# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # 2D Advection-Diffusion equation

# in this notebook we provide a simple example of the DeepMoD algorithm and apply it on the 2D advection-diffusion equation. 

# +
# General imports
import numpy as np
import torch
import matplotlib.pylab as plt

# DeepMoD functions

from deepymod import DeepMoD
from deepymod.model.func_approx import NN, Siren
from deepymod.model.library import Library2D_third
from deepymod.model.constraint import LeastSquares
from deepymod.model.sparse_estimators import Threshold,PDEFIND
from deepymod.training import train
from deepymod.training.sparsity_scheduler import TrainTestPeriodic
from scipy.io import loadmat

# Settings for reproducibility
np.random.seed(42)
torch.manual_seed(0)


if torch.cuda.is_available():
    device = 'cuda'
else:
    device = 'cpu'


# %load_ext autoreload
# %autoreload 2
# -

# ## Prepare the data

# Next, we prepare the dataset.

data_pre = np.load('diffusion_advection_29longb.npy').T

data= data_pre[120:-60,:,10:30]

down_data= np.take(np.take(np.take(data,np.arange(0,data.shape[0],5),axis=0),np.arange(0,data.shape[1],5),axis=1),np.arange(0,data.shape[2],1),axis=2)

down_data.shape

steps = down_data.shape[2]
width = down_data.shape[0]
width_2 = down_data.shape[1]

plt.plot(np.sum(np.sum(data_pre,axis=1),axis=0))

x_arr = np.arange(0,width)
y_arr = np.arange(0,width_2)
t_arr = np.arange(0,steps)
x_grid, y_grid, t_grid = np.meshgrid(x_arr, y_arr, t_arr, indexing='ij')
X = np.transpose((t_grid.flatten(), x_grid.flatten(), y_grid.flatten()))

plt.imshow(down_data[:,:,1])

# Next we plot the dataset for three different time-points

# We flatten it to give it the right dimensions for feeding it to the network:

# +
X = np.transpose((t_grid.flatten()/10, x_grid.flatten()/np.max(y_grid), y_grid.flatten()/np.max(y_grid)))
#X = np.transpose((t_grid.flatten(), x_grid.flatten(), y_grid.flatten()))

y = np.float32(down_data.reshape((down_data.size, 1)))
y = 4*y/np.max(y)
# -

len(y)

# +
number_of_samples = 5000

idx = np.random.permutation(y.shape[0])
X_train = torch.tensor(X[idx, :][:number_of_samples], dtype=torch.float32, requires_grad=True).to(device)
y_train = torch.tensor(y[idx, :][:number_of_samples], dtype=torch.float32).to(device)


# -

# ## Configuration of DeepMoD

# Configuration of the function approximator: Here the first argument is the number of input and the last argument the number of output layers.

network = NN(3, [30, 30, 30, 30], 1)

# Configuration of the library function: We select athe library with a 2D spatial input. Note that that the max differential order has been pre-determined here out of convinience. So, for poly_order 1 the library contains the following 12 terms:
# * [$1, u_x, u_y, u_{xx}, u_{yy}, u_{xy}, u, u u_x, u u_y, u u_{xx}, u u_{yy}, u u_{xy}$]

library = Library2D_third(poly_order=0) 

# Configuration of the sparsity estimator and sparsity scheduler used. In this case we use the most basic threshold-based Lasso estimator and a scheduler that asseses the validation loss after a given patience. If that value is smaller than 1e-5, the algorithm is converged.  

estimator = PDEFIND()
estimator = Threshold(0.05) 
sparsity_scheduler = TrainTestPeriodic(periodicity=50, patience=25, delta=1e-5) 

#
# Configuration of the sparsity estimator 

constraint = LeastSquares() 
# Configuration of the sparsity scheduler

# Now we instantiate the model and select the optimizer 

# +
model = DeepMoD(network, library, estimator, constraint).to(device)

# Defining optimizer
optimizer = torch.optim.Adam(model.parameters(), betas=(0.99, 0.99), amsgrad=True, lr=1e-3) 

# -

# ## Run DeepMoD 

# We can now run DeepMoD using all the options we have set and the training data:
# * The directory where the tensorboard file is written (log_dir)
# * The ratio of train/test set used (split)
# * The maximum number of iterations performed (max_iterations)
# * The absolute change in L1 norm considered converged (delta)
# * The amount of epochs over which the absolute change in L1 norm is calculated (patience)

train(model, X_train, y_train, optimizer,sparsity_scheduler, log_dir='runs/news10/', split=0.8, max_iterations=100000, delta=1e-7, patience=200) 

# Sparsity masks provide the active and non-active terms in the PDE:

sol = model(torch.tensor(X, dtype=torch.float32))[0].reshape((width,width_2,steps)).detach().numpy()

ux = model(torch.tensor(X, dtype=torch.float32))[2][0][:,1].reshape((width,width_2,steps)).detach().numpy()
uy = model(torch.tensor(X, dtype=torch.float32))[2][0][:,2].reshape((width,width_2,steps)).detach().numpy()

uxx = model(torch.tensor(X, dtype=torch.float32))[2][0][:,3].reshape((width,width_2,steps)).detach().numpy()
uyy = model(torch.tensor(X, dtype=torch.float32))[2][0][:,4].reshape((width,width_2,steps)).detach().numpy()

import pysindy as ps

fd_spline = ps.SINDyDerivative(kind='spline', s=1e-2)
fd_spectral = ps.SINDyDerivative(kind='spectral')
fd_sg = ps.SINDyDerivative(kind='savitzky_golay', left=0.5, right=0.5, order=3)

y = down_data[2,:,19]
x = x_arr
plt.plot(x,y, 'b--')
plt.plot(x,sol[2,:,19]*np.max(down_data),'b', label='x = 1')
y = down_data[5,:,19]
x = x_arr
plt.plot(x,y, 'g--')
plt.plot(x,sol[5,:,19]*np.max(down_data),'g', label='x = 5')
y = down_data[11,:,19]
x = x_arr
plt.plot(x,y, 'r--')
plt.plot(x,sol[11,:,19]*np.max(down_data),'r', label='x = 10')
plt.legend()

y = down_data[1,:,1]
x = x_arr
plt.plot(x,y, 'b--')
plt.plot(x,sol[1,:,1]*np.max(down_data),'b', label='x = 1')
y = down_data[5,:,1]
x = x_arr
plt.plot(x,y, 'g--')
plt.plot(x,sol[5,:,1]*np.max(down_data),'g', label='x = 5')
y = down_data[11,:,1]
x = x_arr
plt.plot(x,y, 'r--')
plt.plot(x,sol[11,:,1]*np.max(down_data),'r', label='x = 10')
plt.legend()

np.max(down_data)/100

plt.plot(x,fd_sg(y,x), 'ro')

y = down_data[1,:,19]
x = x_arr
plt.plot(x,fd_sg(y,x), 'b--')
plt.plot(x,uy[1,:,19]*np.max(down_data)/100,'b', label='x = 1')
y = down_data[5,:,19]
x = x_arr
plt.plot(x,fd_sg(y,x), 'g--')
plt.plot(x,uy[5,:,19]*np.max(down_data)/100,'g', label='x = 5')
y = down_data[10,:,19]
x = x_arr
plt.plot(x,fd_sg(y,x), 'r--')
plt.plot(x,uy[10,:,19]*np.max(down_data)/100,'r', label='x = 10')
plt.legend()

y = down_data[2,:,19]
x = x_arr
plt.plot(x,fd_sg(fd_sg(y,x)), 'b--')
plt.plot(x,uyy[2,:,19]*np.max(down_data)/(100*100),'b')
y = down_data[5,:,19]
x = x_arr
plt.plot(x,fd_sg(fd_sg(y,x)), 'g--')
plt.plot(x,uyy[5,:,19]*np.max(down_data)/(100*100),'g')
y = down_data[11,:,19]
x = x_arr
plt.plot(x,fd_sg(fd_sg(y,x)), 'r--')
plt.plot(x,uyy[11,:,19]*np.max(down_data)/(100*100),'r')

# +


fig = plt.figure(figsize=(15,5))

plt.subplot(1,3, 1)
y = down_data[2,:,2]
x = x_arr
plt.plot(x,y)
plt.plot(x,sol[2,:,2]*np.max(down_data))
plt.legend()

plt.subplot(1,3, 2)
y = down_data[2,:,2]
x = x_arr
plt.plot(x,y)
plt.plot(x,sol[2,:,2]*np.max(down_data))

plt.subplot(1,3, 3)
y = down_data[2,:,2]
x = x_arr
plt.plot(x,y)
plt.plot(x,sol[2,:,2]*np.max(down_data))
plt.legend()

plt.show()

# +
fig = plt.figure(figsize=(15,5))

plt.subplot(1,3, 1)
plt.imshow(sol[:,:,1], aspect=0.5)

plt.subplot(1,3, 2)
plt.imshow(sol[:,:,19], aspect=0.5)

plt.subplot(1,3, 3)
plt.imshow(sol[:,:,39], aspect=0.5)


plt.savefig('reconstruction.pdf')

# +
fig = plt.figure(figsize=(15,5))

plt.subplot(1,3, 1)
plt.imshow(down_data[:,:,1], aspect=0.5)

plt.subplot(1,3, 2)
plt.imshow(down_data[:,:,19], aspect=0.5)

plt.subplot(1,3, 3)
plt.imshow(down_data[:,:,39], aspect=0.5)

plt.savefig('original_20_20_40.pdf')
# -

np.max(down_data)

plt.plot(x,sol[5,:,10]*np.max(down_data))

noise_level = 0.025
y_noisy = y + noise_level * np.std(y) * np.random.randn(y.size)

plt.plot(x,uy[25,:,10])
plt.plot(x,ux[25,:,10])

# +
fig = plt.figure(figsize=(15,5))

plt.subplot(1,3, 1)
plt.plot(fd_spline(y.reshape(-1,1),x), label='Ground truth',linewidth=3)
plt.plot(fd_spline(y_noisy.reshape(-1,1),x), label='Spline',linewidth=3)
plt.legend()

plt.subplot(1,3, 2)
plt.plot(fd_spline(y.reshape(-1,1),x), label='Ground truth',linewidth=3)
plt.plot(fd_sg(y_noisy.reshape(-1,1),x), label='<NAME>',linewidth=3)
plt.legend()

plt.subplot(1,3, 3)
plt.plot(fd_spline(y.reshape(-1,1),x), label='Ground truth',linewidth=3)
plt.plot(uy[25,:,10],linewidth=3, label='DeepMoD')
plt.legend()

plt.show()
# -

plt.plot(ux[10,:,5])
ax = plt.subplot(1,1,1)
ax.plot(fd(y.reshape(-1,1),x), label='Ground truth')
ax.plot(fd_sline(y_noisy.reshape(-1,1),x), label='Spline')
ax.plot(fd_sg(y_noisy.reshape(-1,1),x), label='Savitzky Golay')
ax.legend()

plt.plot(model(torch.tensor(X, dtype=torch.float32))[2][0].detach().numpy())

sol = model(torch.tensor(X, dtype=torch.float32))[0]

plt.imshow(sol[:,:,4].detach().numpy())

plt.plot(sol[10,:,6].detach().numpy())
plt.plot(down_data[10,:,6]/np.max(down_data))

x = np.arange(0,len(y))

import pysindy as ps
diffs = [
    ('PySINDy Finite Difference', ps.FiniteDifference()),
    ('Smoothed Finite Difference', ps.SmoothedFiniteDifference()),
    ('Savitzky Golay', ps.SINDyDerivative(kind='savitzky_golay', left=0.5, right=0.5, order=3)),
    ('Spline', ps.SINDyDerivative(kind='spline', s=1e-2)),
    ('Trend Filtered', ps.SINDyDerivative(kind='trend_filtered', order=0, alpha=1e-2)),
    ('Spectral', ps.SINDyDerivative(kind='spectral')),
]

fd = ps.SINDyDerivative(kind='spline', s=1e-2)

y = down_data[:,10,9]/np.max(down_data)

x = np.arange(0,len(y))

t = np.linspace(0,1,5)
X = np.vstack((np.sin(t),np.cos(t))).T

plt.plot(y)

plt.plot(fd(y.reshape(-1,1),x))

y.shape

plt.plot(fd._differentiate(y.reshape(-1,1),x))


plt.plot(ux[:,10,6])

plt.plot(sol[:,10,6].detach().numpy())
plt.plot(down_data[:,10,6]/np.max(down_data))

model.sparsity_masks

# estimatior_coeffs gives the magnitude of the active terms:

print(model.estimator_coeffs())





plt.contourf(ux[:,:,10])

plt.plot(ux[25,:,2])

ax = plt.subplot(1,1,1)
ax.plot(fd(y.reshape(-1,1),x), label='Ground truth')
ax.plot(fd_sline(y_noisy.reshape(-1,1),x), label='Spline')
ax.plot(fd_sg(y_noisy.reshape(-1,1),x), label='<NAME>')
ax.legend()

import pysindy as ps

fd_spline = ps.SINDyDerivative(kind='spline', s=1e-2)
fd_spectral = ps.SINDyDerivative(kind='spectral')
fd_sg = ps.SINDyDerivative(kind='savitzky_golay', left=0.5, right=0.5, order=3)

y = u_v[25,:,2]
x = y_v[25,:,2]
plt.scatter(x,y)

y.shape

noise_level = 0.025
y_noisy = y + noise_level * np.std(y) * np.random.randn(y.size)

ax = plt.subplot(1,1,1)
ax.plot(x,y_noisy, label="line 1")
ax.plot(x,y, label="line 2")
ax.legend()

ax = plt.subplot(1,1,1)
ax.plot(fd(y.reshape(-1,1),x), label='Ground truth')
ax.plot(fd_sline(y_noisy.reshape(-1,1),x), label='Spline')
ax.plot(fd_sg(y_noisy.reshape(-1,1),x), label='<NAME>')
ax.legend()


StarcoderdataJupyterScriptsDedupFiltered