export tutorials changed in dd88513 (#559)
Co-authored-by: dario-coscia <dario-coscia@users.noreply.github.com>
This commit is contained in:
github-actions[bot]
committed by
Dario Coscia
Dario Coscia
parent
228f807d92
commit
d10c525e74
305
tutorials/tutorial10/tutorial.py
vendored
Normal file
305
tutorials/tutorial10/tutorial.py
vendored
Normal file
@@ -0,0 +1,305 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# # Tutorial: Solving the Kuramoto–Sivashinsky Equation with Averaging Neural Operator
|
||||
#
|
||||
# [](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`**. By the end of this tutorial, you will be able to train a Neural Operator to learn the operator for time-dependent PDEs.
|
||||
#
|
||||
# Let's start by importing the necessary modules.
|
||||
|
||||
# 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[tutorial]"')
|
||||
# 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 Trainer, LabelTensor
|
||||
from pina.model import AveragingNeuralOperator
|
||||
from pina.solver import SupervisedSolver
|
||||
from pina.problem.zoo import SupervisedProblem
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
|
||||
# ## Data Generation
|
||||
#
|
||||
# In this tutorial, we will focus on solving the **Kuramoto-Sivashinsky (KS)** equation, a fourth-order nonlinear PDE. The equation is given by:
|
||||
#
|
||||
# $$
|
||||
# \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 this equation, $x \in \Omega = [0, 64]$ represents a spatial location, and $t \in \mathbb{T} = [0, 50]$ represents time. The function $u(x, t)$ is the value of the function at each point in space and time, with $u(x, t) \in \mathbb{R}$. We denote the solution space as $\mathbb{U}$, where $u \in \mathbb{U}$.
|
||||
#
|
||||
# We impose Dirichlet boundary conditions on the derivative of $u$ at the boundary of the domain $\partial \Omega$:
|
||||
#
|
||||
# $$
|
||||
# \frac{\partial u}{\partial x}(x,t) = 0 \quad \forall (x,t) \in \partial \Omega \times \mathbb{T}.
|
||||
# $$
|
||||
#
|
||||
# The initial conditions are sampled from a distribution over truncated Fourier series with random coefficients $\{A_k, \ell_k, \phi_k\}_k$, as follows:
|
||||
#
|
||||
# $$
|
||||
# u(x,0) = \sum_{k=1}^N A_k \sin\left(2 \pi \frac{\ell_k x}{L} + \phi_k\right),
|
||||
# $$
|
||||
#
|
||||
# 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 data for different initial conditions. The goal is to build a Neural Operator that, given $u(x,t)$, outputs $u(x,t+\delta)$, where $\delta$ is a fixed time step.
|
||||
#
|
||||
# We will cover the Neural Operator architecture later, but for now, let’s start by importing the data.
|
||||
#
|
||||
# **Note:**
|
||||
# The numerical integration is obtained using a pseudospectral method for spatial derivative discretization and implicit Runge-Kutta 5 for temporal dynamics.
|
||||
|
||||
# In[3]:
|
||||
|
||||
|
||||
# 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 is saved in the form `[B, N, D]`, where:
|
||||
# - `B` is the batch size (i.e., how many initial conditions we sample),
|
||||
# - `N` is the number of points in the mesh (which is the product of the discretization in $x$ times the one in $t$),
|
||||
# - `D` is the dimension of the problem (in this case, we have three variables: $[u, t, x]$).
|
||||
#
|
||||
# We are now going to plot some trajectories!
|
||||
|
||||
# In[4]:
|
||||
|
||||
|
||||
# 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 time progresses, the solution becomes chaotic, making it very difficult to learn! We will now focus on building a Neural Operator using the `SupervisedSolver` class to tackle this 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 to approximate the following operator, for example, by using a 2D [FNO](https://mathlab.github.io/PINA/_rst/model/fourier_neural_operator.html) (for regular meshes), a [DeepOnet](https://mathlab.github.io/PINA/_rst/model/deeponet.html), [Continuous Convolutional Neural Operator](https://mathlab.github.io/PINA/_rst/model/block/convolution.html), or [MIONet](https://mathlab.github.io/PINA/_rst/model/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/model/kernel_neural_operator.html) with an 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}}$ being the embedding (hidden) size.
|
||||
# * $\sigma$ is a non-linear activation function.
|
||||
# * $W \in \mathbb{R}^{\rm{emb} \times \rm{emb}}$ is a tunable matrix.
|
||||
# * $b \in \mathbb{R}^{\rm{emb}}$ is a tunable bias.
|
||||
#
|
||||
# In PINA, many Kernel Neural Operators are already implemented. The modular components of the [Kernel Neural Operator](https://mathlab.github.io/PINA/_rst/model/kernel_neural_operator.html) class allow you to create new ones by composing base kernel layers.
|
||||
#
|
||||
# **Note:** We will use the already built class `AveragingNeuralOperator`. As a constructive exercise, try to use the [KernelNeuralOperator](https://mathlab.github.io/PINA/_rst/model/kernel_neural_operator.html) class to build a kernel neural operator from scratch. You might employ the different layers that we have in PINA, such as [FeedForward](https://mathlab.github.io/PINA/_rst/model/feed_forward.html) and [AveragingNeuralOperator](https://mathlab.github.io/PINA/_rst/model/average_neural_operator.html) layers.
|
||||
|
||||
# In[5]:
|
||||
|
||||
|
||||
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, which is discussed in more detail in the paper [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[6]:
|
||||
|
||||
|
||||
# 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 visualize some plots for the solutions!
|
||||
|
||||
# In[7]:
|
||||
|
||||
|
||||
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 results considering the small training time and the difficulty of the problem!
|
||||
# Let's take a look at the training and testing error:
|
||||
|
||||
# In[8]:
|
||||
|
||||
|
||||
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 aligns with the observations from the previous plots.
|
||||
|
||||
# ## What's Next?
|
||||
#
|
||||
# You have completed the tutorial on solving time-dependent PDEs using Neural Operators in **PINA**. Great job! Here are some potential next steps you can explore:
|
||||
#
|
||||
# 1. **Train the network for longer or with different layer sizes**: Experiment with various configurations, such as adjusting the number of layers or hidden dimensions, to further improve accuracy and observe the impact on performance.
|
||||
#
|
||||
# 2. **Use a more challenging dataset**: Try using the more complex dataset [Data_KS2.mat](dat/Data_KS2.mat) where $A_k \in [-0.5, 0.5]$, $\ell_k \in [1, 2, 3]$, and $\phi_k \in [0, 2\pi]$ for a more difficult task. This dataset may require longer training and testing.
|
||||
#
|
||||
# 3. **... and many more...**: Explore other models, such as the [FNO](https://mathlab.github.io/PINA/_rst/models/fno.html), [DeepOnet](https://mathlab.github.io/PINA/_rst/models/deeponet.html), or implement your own operator using the [KernelNeuralOperator](https://mathlab.github.io/PINA/_rst/models/base_no.html) class to compare performance and find the best model for your task.
|
||||
#
|
||||
# For more resources and tutorials, check out the [PINA Documentation](https://mathlab.github.io/PINA/).
|
||||
Reference in New Issue
Block a user