310 lines
12 KiB
Python
Vendored
310 lines
12 KiB
Python
Vendored
#!/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/).
|