Skip to content

Commit

Permalink
Added eq. printing function
Browse files Browse the repository at this point in the history
  • Loading branch information
Gert-Jan committed Jan 24, 2020
1 parent 11e6b76 commit bb2cd18
Show file tree
Hide file tree
Showing 4 changed files with 112 additions and 99 deletions.
120 changes: 101 additions & 19 deletions examples/Example.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,16 +18,16 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 2,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<torch._C.Generator at 0x7fa64a6c9f70>"
"<torch._C.Generator at 0x7f5e10766f70>"
]
},
"execution_count": 4,
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -41,7 +41,7 @@
"# DeepMoD stuff\n",
"from deepymod_torch.DeepMod import DeepMod\n",
"from deepymod_torch.library_functions import library_basic\n",
"from deepymod_torch.utilities import create_deriv_data\n",
"from deepymod_torch.utilities import create_deriv_data, terms_definition\n",
"\n",
"# Setting cuda if GPU is available\n",
"if torch.cuda.is_available():\n",
Expand All @@ -61,7 +61,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 3,
"metadata": {},
"outputs": [
{
Expand All @@ -86,7 +86,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 4,
"metadata": {},
"outputs": [
{
Expand Down Expand Up @@ -114,7 +114,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 5,
"metadata": {},
"outputs": [
{
Expand All @@ -134,7 +134,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": 6,
"metadata": {},
"outputs": [
{
Expand All @@ -149,7 +149,7 @@
" [10. , 7.9375]])"
]
},
"execution_count": 8,
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -167,7 +167,7 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -184,7 +184,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -197,7 +197,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 9,
"metadata": {},
"outputs": [
{
Expand All @@ -221,7 +221,7 @@
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": 10,
"metadata": {},
"outputs": [
{
Expand Down Expand Up @@ -277,7 +277,7 @@
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -293,7 +293,7 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -316,7 +316,7 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -326,7 +326,7 @@
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": 14,
"metadata": {},
"outputs": [
{
Expand All @@ -353,12 +353,12 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that DeepMoD has converged, it has found the following equation:"
"Now that DeepMoD has converged, it has found the following numbers:"
]
},
{
"cell_type": "code",
"execution_count": 21,
"execution_count": 15,
"metadata": {},
"outputs": [
{
Expand All @@ -377,6 +377,88 @@
"print(model.coeff_vector_list[0], model.sparsity_mask_list)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can put that into a slightly friendlier format by defining which terms our library used:"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": [
"deriv_list = [['', 'u_x', 'u_xx']]\n",
"poly_list = [['', 'u', 'u^2']]"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
"terms = terms_definition(poly_list, deriv_list)"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"['', 'u_x', 'u_xx', 'u', 'uu_x', 'uu_xx', 'u^2', 'u^2u_x', 'u^2u_xx']"
]
},
"execution_count": 23,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"terms"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {},
"outputs": [],
"source": [
"def print_eq(coeff_vector, sparsity_mask, terms):\n",
" for coeff, term in zip(coeff_vector, sparsity_mask):\n",
" print(coeff.item(), terms[term])"
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"0.09520050138235092 u_xx\n",
"-0.9261836409568787 uu_x\n"
]
}
],
"source": [
"print_eq(model.coeff_vector_list[0].detach().cpu().numpy(), model.sparsity_mask_list[0].cpu().numpy(), terms)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Which is the Burgers equation :-)."
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down
75 changes: 2 additions & 73 deletions src/deepymod_torch/library_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,35 +3,14 @@
from itertools import combinations, product


class library_1D(Library):
def __init__(self, input_dim, output_dim, diff_order, poly_order):
self.poly_order = poly_order
super().__init__(input_dim, output_dim, diff_order)

def theta(self, input):
# Works for all N-D inputs, 1-D outputs
X, dX = input
dt = dX[:, 0, 0:1, 0]
dx = dX[:, :, 1:, 0]

# Calculate the polynomes of u
u = torch.ones_like(X)
for order in torch.arange(1, self.poly_order+1):
u = torch.cat((u, u[:, order-1:order] * X), dim=1)

# Calculate derivs
du = torch.cat((torch.ones((dx.shape[0], 1)), dx.reshape(dx.shape[0], -1)), dim=1) # u_x, u_y, u_xx, u_yy etc

theta = (u[:, :, None] @ du[:, None, :]).reshape(u.shape[0], -1)

return [dt], theta

class library_basic(Library):
'''Implementation of library layer. Inherets from Library layer.'''
def __init__(self, input_dim, output_dim, diff_order, poly_order):
self.poly_order = poly_order
super().__init__(input_dim, output_dim, diff_order)

def theta(self, input):
'''Calculates the library and time deriv from NN output'''
X, dX = input
samples = X.shape[0]

Expand Down Expand Up @@ -60,53 +39,3 @@ def theta(self, input):

return time_deriv_list, theta

def library_1D_in(data, prediction, library_config):
'''
Calculates a library function for a 1D+1 input for M coupled equations consisting of all polynomials up to order K and derivatives up to order
L and all possible combinations (i.e. combining two terms) of these.
Parameters
----------
data : tensor of size (N x 2)
coordinates to whose respect the derivatives of prediction are calculated. First column is time, space second column.
prediction : tensor of size (N x M)
dataset from which the library is constructed.
library_config : dict
dictionary containing options for the library function.
Returns
-------
time_deriv_list : tensor list of length of M
list containing the time derivatives, each entry corresponds to an equation.
theta : tensor
library matrix tensor.
'''
poly_list = []
deriv_list = []
time_deriv_list = []

# Creating lists for all outputs
for output in torch.arange(prediction.shape[1]):
time_deriv, du = library_deriv(data, prediction[:, output:output+1], library_config)
u = library_poly(prediction[:, output:output+1], library_config)

poly_list.append(u)
deriv_list.append(du)
time_deriv_list.append(time_deriv)

samples = time_deriv_list[0].shape[0]
total_terms = poly_list[0].shape[1] * deriv_list[0].shape[1]

# Calculating theta
if len(poly_list) == 1:
theta = torch.matmul(poly_list[0][:, :, None], deriv_list[0][:, None, :]).view(samples, total_terms) # If we have a single output, we simply calculate and flatten matrix product between polynomials and derivatives to get library
else:
theta_uv = torch.cat([torch.matmul(u[:, :, None], v[:, None, :]).view(samples, total_terms) for u, v in combinations(poly_list, 2)], 1) # calculate all unique combinations between polynomials
theta_dudv = torch.cat([torch.matmul(du[:, :, None], dv[:, None, :]).view(samples, total_terms)[:, 1:] for du, dv in combinations(deriv_list, 2)], 1) # calculate all unique combinations of derivatives
theta_udu = torch.cat([torch.matmul(u[:, 1:, None], du[:, None, 1:]).view(samples, (poly_list[0].shape[1]-1) * (deriv_list[0].shape[1]-1)) for u, du in product(poly_list, deriv_list)], 1) # calculate all unique products of polynomials and derivatives
theta = torch.cat([theta_uv, theta_dudv, theta_udu], dim=1)

return time_deriv_list, theta



14 changes: 9 additions & 5 deletions src/deepymod_torch/utilities.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,20 @@
import torch

def string_matmul(list_1, list_2):
''' Matrix multiplication with strings.'''
prod = [element[0] + element[1] for element in product(list_1, list_2)]
return prod


def terms_definition(poly_list, deriv_list):
theta_uv = [string_matmul(u, v) for u, v in combinations(poly_list, 2)]
theta_dudv = [string_matmul(du, dv)[1:] for du, dv in combinations(deriv_list, 2)]
theta_udv = [string_matmul(u[1:], dv[1:]) for u, dv in product(poly_list, deriv_list)]
theta = [element for theta_specific in (theta_uv + theta_dudv + theta_udv) for element in theta_specific]

''' Calculates which terms are in the library.'''
if len(poly_list) == 1:
theta = string_matmul(poly_list[0], deriv_list[0]) # If we have a single output, we simply calculate and flatten matrix product between polynomials and derivatives to get library
else:
theta_uv = list(chain.from_iterable([string_matmul(u, v) for u, v in combinations(poly_list, 2)])) # calculate all unique combinations between polynomials
theta_dudv = list(chain.from_iterable([string_matmul(du, dv)[1:] for du, dv in combinations(deriv_list, 2)])) # calculate all unique combinations of derivatives
theta_udu = list(chain.from_iterable([string_matmul(u[1:], du[1:]) for u, du in product(poly_list, deriv_list)])) # calculate all unique combinations of derivatives
theta = theta_uv + theta_dudv + theta_udu
return theta

def create_deriv_data(X, max_order):
Expand Down
2 changes: 0 additions & 2 deletions tests/keller_segel.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,6 @@
config = {'input_dim': 2, 'hidden_dim': 20, 'layers': 5, 'output_dim': 2, 'library_function': library_basic, 'library_args':{'poly_order': 1, 'diff_order': 2}}

X_input = create_deriv_data(X_train, config['library_args']['diff_order'])

model = DeepMod(config)
optimizer = torch.optim.Adam(model.parameters())
model.train(X_input, y_train, optimizer, 100000, type='deepmod')

Expand Down

0 comments on commit bb2cd18

Please sign in to comment.