Dario Coscia
151
tutorials/tutorial10/tutorial.ipynb
vendored
151
tutorials/tutorial10/tutorial.ipynb
vendored
File diff suppressed because one or more lines are too long
314
tutorials/tutorial10/tutorial.py
vendored
314
tutorials/tutorial10/tutorial.py
vendored
@@ -1,314 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# # Tutorial: Averaging Neural Operator for solving Kuramoto Sivashinsky equation
|
||||
#
|
||||
# [](https://colab.research.google.com/github/mathLab/PINA/blob/master/tutorials/tutorial10/tutorial.ipynb)
|
||||
#
|
||||
# In this tutorial we will build a Neural Operator using the
|
||||
# `AveragingNeuralOperator` model and the `SupervisedSolver`. At the end of the
|
||||
# tutorial you will be able to train a Neural Operator for learning
|
||||
# the operator of time dependent PDEs.
|
||||
#
|
||||
#
|
||||
# First of all, some useful imports. Note we use `scipy` for i/o operations.
|
||||
#
|
||||
|
||||
# In[ ]:
|
||||
|
||||
|
||||
## routine needed to run the notebook on Google Colab
|
||||
try:
|
||||
import google.colab
|
||||
|
||||
IN_COLAB = True
|
||||
except:
|
||||
IN_COLAB = False
|
||||
if IN_COLAB:
|
||||
get_ipython().system('pip install "pina-mathlab"')
|
||||
# get the data
|
||||
get_ipython().system('mkdir "data"')
|
||||
get_ipython().system('wget "https://github.com/mathLab/PINA/raw/refs/heads/master/tutorials/tutorial10/data/Data_KS.mat" -O "data/Data_KS.mat"')
|
||||
get_ipython().system('wget "https://github.com/mathLab/PINA/raw/refs/heads/master/tutorials/tutorial10/data/Data_KS2.mat" -O "data/Data_KS2.mat"')
|
||||
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
import warnings
|
||||
|
||||
from scipy import io
|
||||
from pina import Condition, Trainer, LabelTensor
|
||||
from pina.model import AveragingNeuralOperator
|
||||
from pina.solver import SupervisedSolver
|
||||
from pina.problem.zoo import SupervisedProblem
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
|
||||
# ## Data Generation
|
||||
#
|
||||
# We will focus on solving a specific PDE, the **Kuramoto Sivashinsky** (KS) equation.
|
||||
# The KS PDE is a fourth-order nonlinear PDE with the following form:
|
||||
#
|
||||
# $$
|
||||
# \frac{\partial u}{\partial t}(x,t) = -u(x,t)\frac{\partial u}{\partial x}(x,t)- \frac{\partial^{4}u}{\partial x^{4}}(x,t) - \frac{\partial^{2}u}{\partial x^{2}}(x,t).
|
||||
# $$
|
||||
#
|
||||
# In the above $x\in \Omega=[0, 64]$ represents a spatial location, $t\in\mathbb{T}=[0,50]$ the time and $u(x, t)$ is the value of the function $u:\Omega \times\mathbb{T}\in\mathbb{R}$. We indicate with $\mathbb{U}$ a suitable space for $u$, i.e. we have that the solution $u\in\mathbb{U}$.
|
||||
#
|
||||
#
|
||||
# We impose Dirichlet boundary conditions on the derivative of $u$ on the border of the domain $\partial \Omega$
|
||||
# $$
|
||||
# \frac{\partial u}{\partial x}(x,t)=0 \quad \forall (x,t)\in \partial \Omega\times\mathbb{T}.
|
||||
# $$
|
||||
#
|
||||
# Initial conditions are sampled from a distribution over truncated Fourier series with random coefficients
|
||||
# $\{A_k, \ell_k, \phi_k\}_k$ as
|
||||
# $$
|
||||
# u(x,0) = \sum_{k=1}^N A_k \sin(2 \pi \ell_k x / L + \phi_k) \ ,
|
||||
# $$
|
||||
#
|
||||
# where $A_k \in [-0.4, -0.3]$, $\ell_k = 2$, $\phi_k = 2\pi \quad \forall k=1,\dots,N$.
|
||||
#
|
||||
#
|
||||
# We have already generated some data for differenti initial conditions, and our objective will
|
||||
# be to build a Neural Operator that, given $u(x, t)$ will output $u(x, t+\delta)$, where
|
||||
# $\delta$ is a fixed time step. We will come back on the Neural Operator architecture, for now
|
||||
# we first need to import the data.
|
||||
#
|
||||
# **Note:**
|
||||
# *The numerical integration is obtained by using pseudospectral method for spatial derivative discratization and
|
||||
# implicit Runge Kutta 5 for temporal dynamics.*
|
||||
#
|
||||
|
||||
# In[2]:
|
||||
|
||||
|
||||
# load data
|
||||
data = io.loadmat("data/Data_KS.mat")
|
||||
|
||||
# converting to label tensor
|
||||
initial_cond_train = LabelTensor(
|
||||
torch.tensor(data["initial_cond_train"], dtype=torch.float),
|
||||
["t", "x", "u0"],
|
||||
)
|
||||
initial_cond_test = LabelTensor(
|
||||
torch.tensor(data["initial_cond_test"], dtype=torch.float), ["t", "x", "u0"]
|
||||
)
|
||||
sol_train = LabelTensor(
|
||||
torch.tensor(data["sol_train"], dtype=torch.float), ["u"]
|
||||
)
|
||||
sol_test = LabelTensor(torch.tensor(data["sol_test"], dtype=torch.float), ["u"])
|
||||
|
||||
print("Data Loaded")
|
||||
print(f" shape initial condition: {initial_cond_train.shape}")
|
||||
print(f" shape solution: {sol_train.shape}")
|
||||
|
||||
|
||||
# The data are saved in the form `B \times N \times D`, where `B` is the batch_size
|
||||
# (basically how many initial conditions we sample), `N` the number of points in the mesh
|
||||
# (which is the product of the discretization in `x` timese the one in `t`), and
|
||||
# `D` the dimension of the problem (in this case we have three variables `[u, t, x]`).
|
||||
#
|
||||
# We are now going to plot some trajectories!
|
||||
|
||||
# In[3]:
|
||||
|
||||
|
||||
# helper function
|
||||
def plot_trajectory(coords, real, no_sol=None):
|
||||
# find the x-t shapes
|
||||
dim_x = len(torch.unique(coords.extract("x")))
|
||||
dim_t = len(torch.unique(coords.extract("t")))
|
||||
# if we don't have the Neural Operator solution we simply plot the real one
|
||||
if no_sol is None:
|
||||
fig, axs = plt.subplots(1, 1, figsize=(15, 5), sharex=True, sharey=True)
|
||||
c = axs.imshow(
|
||||
real.reshape(dim_t, dim_x).T.detach(),
|
||||
extent=[0, 50, 0, 64],
|
||||
cmap="PuOr_r",
|
||||
aspect="auto",
|
||||
)
|
||||
axs.set_title("Real solution")
|
||||
fig.colorbar(c, ax=axs)
|
||||
axs.set_xlabel("t")
|
||||
axs.set_ylabel("x")
|
||||
# otherwise we plot the real one, the Neural Operator one, and their difference
|
||||
else:
|
||||
fig, axs = plt.subplots(1, 3, figsize=(15, 5), sharex=True, sharey=True)
|
||||
axs[0].imshow(
|
||||
real.reshape(dim_t, dim_x).T.detach(),
|
||||
extent=[0, 50, 0, 64],
|
||||
cmap="PuOr_r",
|
||||
aspect="auto",
|
||||
)
|
||||
axs[0].set_title("Real solution")
|
||||
axs[1].imshow(
|
||||
no_sol.reshape(dim_t, dim_x).T.detach(),
|
||||
extent=[0, 50, 0, 64],
|
||||
cmap="PuOr_r",
|
||||
aspect="auto",
|
||||
)
|
||||
axs[1].set_title("NO solution")
|
||||
c = axs[2].imshow(
|
||||
(real - no_sol).abs().reshape(dim_t, dim_x).T.detach(),
|
||||
extent=[0, 50, 0, 64],
|
||||
cmap="PuOr_r",
|
||||
aspect="auto",
|
||||
)
|
||||
axs[2].set_title("Absolute difference")
|
||||
fig.colorbar(c, ax=axs.ravel().tolist())
|
||||
for ax in axs:
|
||||
ax.set_xlabel("t")
|
||||
ax.set_ylabel("x")
|
||||
plt.show()
|
||||
|
||||
|
||||
# a sample trajectory (we use the sample 5, feel free to change)
|
||||
sample_number = 20
|
||||
plot_trajectory(
|
||||
coords=initial_cond_train[sample_number].extract(["x", "t"]),
|
||||
real=sol_train[sample_number].extract("u"),
|
||||
)
|
||||
|
||||
|
||||
# As we can see, as the time progresses the solution becomes chaotic, which makes
|
||||
# it really hard to learn! We will now focus on building a Neural Operator using the
|
||||
# `SupervisedSolver` class to tackle the problem.
|
||||
#
|
||||
# ## Averaging Neural Operator
|
||||
#
|
||||
# We will build a neural operator $\texttt{NO}$ which takes the solution at time $t=0$ for any $x\in\Omega$,
|
||||
# the time $(t)$ at which we want to compute the solution, and gives back the solution to the KS equation $u(x, t)$, mathematically:
|
||||
# $$
|
||||
# \texttt{NO}_\theta : \mathbb{U} \rightarrow \mathbb{U},
|
||||
# $$
|
||||
# such that
|
||||
# $$
|
||||
# \texttt{NO}_\theta[u(t=0)](x, t) \rightarrow u(x, t).
|
||||
# $$
|
||||
#
|
||||
# There are many ways on approximating the following operator, e.g. by 2D [FNO](https://mathlab.github.io/PINA/_rst/models/fno.html) (for regular meshes),
|
||||
# a [DeepOnet](https://mathlab.github.io/PINA/_rst/models/deeponet.html), [Continuous Convolutional Neural Operator](https://mathlab.github.io/PINA/_rst/layers/convolution.html),
|
||||
# [MIONet](https://mathlab.github.io/PINA/_rst/models/mionet.html).
|
||||
# In this tutorial we will use the *Averaging Neural Operator* presented in [*The Nonlocal Neural Operator: Universal Approximation*](https://arxiv.org/abs/2304.13221)
|
||||
# which is a [Kernel Neural Operator](https://mathlab.github.io/PINA/_rst/models/base_no.html) with integral kernel:
|
||||
#
|
||||
# $$
|
||||
# K(v) = \sigma\left(Wv(x) + b + \frac{1}{|\Omega|}\int_\Omega v(y)dy\right)
|
||||
# $$
|
||||
#
|
||||
# where:
|
||||
#
|
||||
# * $v(x)\in\mathbb{R}^{\rm{emb}}$ is the update for a function $v$ with $\mathbb{R}^{\rm{emb}}$ the embedding (hidden) size
|
||||
# * $\sigma$ is a non-linear activation
|
||||
# * $W\in\mathbb{R}^{\rm{emb}\times\rm{emb}}$ is a tunable matrix.
|
||||
# * $b\in\mathbb{R}^{\rm{emb}}$ is a tunable bias.
|
||||
#
|
||||
# If PINA many Kernel Neural Operators are already implemented, and the modular componets of the [Kernel Neural Operator](https://mathlab.github.io/PINA/_rst/models/base_no.html) class permits to create new ones by composing base kernel layers.
|
||||
#
|
||||
# **Note:*** We will use the already built class* `AveragingNeuralOperator`, *as constructive excercise try to use the* [KernelNeuralOperator](https://mathlab.github.io/PINA/_rst/models/base_no.html) *class for building a kernel neural operator from scratch. You might employ the different layers that we have in pina, e.g.* [FeedForward](https://mathlab.github.io/PINA/_rst/models/fnn.html), *and* [AveragingNeuralOperator](https://mathlab.github.io/PINA/_rst/layers/avno_layer.html) *layers*.
|
||||
|
||||
# In[4]:
|
||||
|
||||
|
||||
class SIREN(torch.nn.Module):
|
||||
def forward(self, x):
|
||||
return torch.sin(x)
|
||||
|
||||
|
||||
embedding_dimesion = 40 # hyperparameter embedding dimension
|
||||
input_dimension = 3 # ['u', 'x', 't']
|
||||
number_of_coordinates = 2 # ['x', 't']
|
||||
lifting_net = torch.nn.Linear(input_dimension, embedding_dimesion)
|
||||
projecting_net = torch.nn.Linear(embedding_dimesion + number_of_coordinates, 1)
|
||||
model = AveragingNeuralOperator(
|
||||
lifting_net=lifting_net,
|
||||
projecting_net=projecting_net,
|
||||
coordinates_indices=["x", "t"],
|
||||
field_indices=["u0"],
|
||||
n_layers=4,
|
||||
func=SIREN,
|
||||
)
|
||||
|
||||
|
||||
# Super easy! Notice that we use the `SIREN` activation function, more on [Implicit Neural Representations with Periodic Activation Functions](https://arxiv.org/abs/2006.09661).
|
||||
#
|
||||
# ## Solving the KS problem
|
||||
#
|
||||
# We will now focus on solving the KS equation using the `SupervisedSolver` class
|
||||
# and the `AveragingNeuralOperator` model. As done in the [FNO tutorial](https://github.com/mathLab/PINA/blob/master/tutorials/tutorial5/tutorial.ipynb) we now create the Neural Operator problem class with `SupervisedProblem`.
|
||||
|
||||
# In[5]:
|
||||
|
||||
|
||||
# initialize problem
|
||||
problem = SupervisedProblem(
|
||||
initial_cond_train,
|
||||
sol_train,
|
||||
input_variables=initial_cond_train.labels,
|
||||
output_variables=sol_train.labels,
|
||||
)
|
||||
# initialize solver
|
||||
solver = SupervisedSolver(problem=problem, model=model)
|
||||
# train, only CPU and avoid model summary at beginning of training (optional)
|
||||
trainer = Trainer(
|
||||
solver=solver,
|
||||
max_epochs=40,
|
||||
accelerator="cpu",
|
||||
enable_model_summary=False,
|
||||
batch_size=5, # we train on CPU and avoid model summary at beginning of training (optional)
|
||||
train_size=1.0,
|
||||
val_size=0.0,
|
||||
test_size=0.0,
|
||||
)
|
||||
trainer.train()
|
||||
|
||||
|
||||
# We can now see some plots for the solutions
|
||||
|
||||
# In[6]:
|
||||
|
||||
|
||||
sample_number = 2
|
||||
no_sol = solver(initial_cond_test)
|
||||
plot_trajectory(
|
||||
coords=initial_cond_test[sample_number].extract(["x", "t"]),
|
||||
real=sol_test[sample_number].extract("u"),
|
||||
no_sol=no_sol[5],
|
||||
)
|
||||
|
||||
|
||||
# As we can see we can obtain nice result considering the small training time and the difficulty of the problem!
|
||||
# Let's take a look at the training and testing error:
|
||||
|
||||
# In[7]:
|
||||
|
||||
|
||||
from pina.loss import PowerLoss
|
||||
|
||||
error_metric = PowerLoss(p=2) # we use the MSE loss
|
||||
|
||||
with torch.no_grad():
|
||||
no_sol_train = solver(initial_cond_train)
|
||||
err_train = error_metric(
|
||||
sol_train.extract("u"), no_sol_train
|
||||
).mean() # we average the error over trajectories
|
||||
no_sol_test = solver(initial_cond_test)
|
||||
err_test = error_metric(
|
||||
sol_test.extract("u"), no_sol_test
|
||||
).mean() # we average the error over trajectories
|
||||
print(f"Training error: {float(err_train):.3f}")
|
||||
print(f"Testing error: {float(err_test):.3f}")
|
||||
|
||||
|
||||
# As we can see the error is pretty small, which agrees with what we can see from the previous plots.
|
||||
|
||||
# ## What's next?
|
||||
#
|
||||
# Now you know how to solve a time dependent neural operator problem in **PINA**! There are multiple directions you can go now:
|
||||
#
|
||||
# 1. Train the network for longer or with different layer sizes and assert the final accuracy
|
||||
#
|
||||
# 2. We left a more challenging dataset [Data_KS2.mat](dat/Data_KS2.mat) where $A_k \in [-0.5, 0.5]$, $\ell_k \in [1, 2, 3]$, $\phi_k \in [0, 2\pi]$ for longer training
|
||||
#
|
||||
# 3. Compare the performance between the different neural operators (you can even try to implement your favourite one!)
|
||||
Reference in New Issue
Block a user