<jupyter_start><jupyter_text>Example Burgers' equation In this notebook we provide a simple example of the DeepMoD algorithm by applying it on the Burgers' equation. We start by importing the required libraries and setting the plotting style:<jupyter_code>import numpy as np
import torch
import matplotlib.pyplot as plt
plt.style.use('seaborn-notebook')

# DeepMoD stuff
from deepymod_torch.DeepMod import DeepMod
from deepymod_torch.library_functions import library_basic
from deepymod_torch.utilities import create_deriv_data, terms_definition

# Setting cuda if GPU is available
if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

# Settings for reproducibility
np.random.seed(42)
torch.manual_seed(0)<jupyter_output><empty_output><jupyter_text>Next, we prepare the dataset.<jupyter_code>data = np.load('data/burgers.npy', allow_pickle=True).item()
print('Shape of grid:', data['x'].shape)<jupyter_output>Shape of grid: (256, 101)<jupyter_text>Let's plot it to get an idea of the data:<jupyter_code>fig, ax = plt.subplots()
im = ax.contourf(data['x'], data['t'], np.real(data['u']))
ax.set_xlabel('x')
ax.set_ylabel('t')
fig.colorbar(mappable=im)

plt.show()
X = np.transpose((data['t'].flatten(), data['x'].flatten()))
y = np.real(data['u']).reshape((data['u'].size, 1))

print(X.shape, y.shape)
X<jupyter_output><empty_output><jupyter_text>As we can see, $X$ has 2 dimensions, $\{x, t\}$, while $y$ has only one, $\{u\}$. Always explicity set the shape (i.e. $N\times 1$, not $N$) or you'll get errors. This dataset is noiseless, so let's add $5\%$ noise:<jupyter_code>noise_level = 0.05
y_noisy = y + noise_level * np.std(y) * np.random.randn(y.size, 1)<jupyter_output><empty_output><jupyter_text>The dataset is also much larger than needed, so let's hussle it and pick out a 1000 samples:<jupyter_code>number_of_samples = 1000

idx = np.random.permutation(y.shape[0])
X_train = torch.tensor(X[idx, :][:number_of_samples], dtype=torch.float32, requires_grad=True)
y_train = torch.tensor(y_noisy[idx, :][:number_of_samples], dtype=torch.float32)
print(X_train.shape, y_train.shape)<jupyter_output>torch.Size([1000, 2]) torch.Size([1000, 1])<jupyter_text>We now have a dataset which we can use. Let's plot, for a final time, the original dataset, the noisy set and the samples points:<jupyter_code>fig, axes = plt.subplots(ncols=3, figsize=(15, 4))

im0 = axes[0].contourf(data['x'], data['t'], np.real(data['u']), cmap='coolwarm')
axes[0].set_xlabel('x')
axes[0].set_ylabel('t')
axes[0].set_title('Ground truth')

im1 = axes[1].contourf(data['x'], data['t'], y_noisy.reshape(data['x'].shape), cmap='coolwarm')
axes[1].set_xlabel('x')
axes[1].set_title('Noisy')

sampled = np.array([y_noisy[index, 0] if index in idx[:number_of_samples] else np.nan for index in np.arange(data['x'].size)])
sampled = np.rot90(sampled.reshape(data['x'].shape)) #array needs to be rotated because of imshow

im2 = axes[2].imshow(sampled, aspect='auto', cmap='coolwarm')
axes[2].set_xlabel('x')
axes[2].set_title('Sampled')

fig.colorbar(im1, ax=axes.ravel().tolist())

plt.show()<jupyter_output><empty_output><jupyter_text>Configuring DeepMoD We now setup the options for DeepMoD. The setup requires the dimensions of the neural network, a library function and some args for the library function:<jupyter_code>config = {'input_dim': 2, 'hidden_dim': 20, 'layers': 5, 'output_dim': 1, 'library_function': library_basic, 'library_args':{'poly_order': 2, 'diff_order': 2}}<jupyter_output><empty_output><jupyter_text>Now we instantiate the model:<jupyter_code>model = DeepMod(config)<jupyter_output><empty_output><jupyter_text>Run DeepMoD  We can now run DeepMoD using all the options we have set and the training data. We need to slightly preprocess the input data for the derivatives:<jupyter_code>optimizer = torch.optim.Adam(model.parameters(), lr=0.002)
X_input = create_deriv_data(X_train, config['library_args']['diff_order'])
model.train(X_input, y_train, optimizer, 5000, type='deepmod')<jupyter_output>| Iteration | Progress | Time remaining |     Cost |      MSE |      Reg |       L1 |
       5000    100.00%               0s   9.38e-05   7.68e-05   2.42e-06   1.46e-05 
[Parameter containing:
tensor([[ 0.0952],
        [-0.9262]], requires_grad=True)]
[tensor([2, 4])]

| Iteration | Progress | Time remaining |     Cost |      MSE |      Reg |       L1 |
       5000    100.00%               0s   1.13e-04   7.68e-05   3.57e-05   0.00e+00<jupyter_text>Now that DeepMoD has converged, it has found the following numbers:<jupyter_code>print('Final result:')
print(model.coeff_vector_list[0], model.sparsity_mask_list)<jupyter_output>Final result:
Parameter containing:
tensor([[ 0.0952],
        [-0.9262]], requires_grad=True) [tensor([2, 4])]<jupyter_text>We can put that into a slightly friendlier format by defining which terms our library used:<jupyter_code>deriv_list = [['', 'u_x', 'u_xx']]
poly_list = [['', 'u', 'u^2']]
terms = terms_definition(poly_list, deriv_list)
terms
def print_eq(coeff_vector, sparsity_mask, terms):
    for coeff, term in zip(coeff_vector, sparsity_mask):
        print(coeff.item(), terms[term])
print_eq(model.coeff_vector_list[0].detach().cpu().numpy(), model.sparsity_mask_list[0].cpu().numpy(), terms)<jupyter_output>0.09520050138235092 u_xx
-0.9261836409568787 uu_x<jupyter_text>Which is the Burgers equation :-).  Other options We have a few other options, most notably training the mse or a single deepmod cycle (i.e. mse + reg + l1, but not thresholding):<jupyter_code>model.train(X_input, y_train, optimizer, 1000, type='mse')
model.train(X_input, y_train, optimizer, 1000, type='single_cycle')<jupyter_output>| Iteration | Progress | Time remaining |     Cost |      MSE |      Reg |       L1 |
       1000    100.00%               0s   1.26e-04   7.68e-05   3.57e-05   1.35e-05
{'fraction_non_alphanumeric': 0.1442964480874317, 'fraction_numerical': 0.04115437158469945, 'mean_word_length': 5.387131952017448, 'pattern_counts': {'":': 0, '<': 43, '<?xml version=': 0, '>': 43, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '220388', 'license': 'MIT', 'n_tokens_mistral': 2289, 'n_tokens_neox': 2110, 'n_words': 509, 'path': 'examples/Example.ipynb', 'repo_name': 'GJBoth/DeePyMoD_torch', 'size': '5856'}