295 lines
11 KiB
Python
Vendored
295 lines
11 KiB
Python
Vendored
#!/usr/bin/env python
|
||
# coding: utf-8
|
||
|
||
# # Tutorial: Learning Multiscale PDEs Using Fourier Feature Networks
|
||
#
|
||
# [](https://colab.research.google.com/github/mathLab/PINA/blob/master/tutorials/tutorial13/tutorial.ipynb)
|
||
#
|
||
# This tutorial demonstrates how to solve a PDE with multiscale behavior using Physics-Informed Neural Networks (PINNs), as discussed in [*On the Eigenvector Bias of Fourier Feature Networks: From Regression to Solving Multi-Scale PDEs with Physics-Informed Neural Networks*](https://doi.org/10.1016/j.cma.2021.113938).
|
||
#
|
||
# Let’s begin by importing the necessary libraries.
|
||
#
|
||
|
||
# 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]"')
|
||
|
||
import torch
|
||
import matplotlib.pyplot as plt
|
||
import warnings
|
||
|
||
from pina import Condition, Trainer
|
||
from pina.problem import SpatialProblem
|
||
from pina.operator import laplacian
|
||
from pina.solver import PINN, SelfAdaptivePINN as SAPINN
|
||
from pina.loss import LpLoss
|
||
from pina.domain import CartesianDomain
|
||
from pina.equation import Equation, FixedValue
|
||
from pina.model import FeedForward
|
||
from pina.model.block import FourierFeatureEmbedding
|
||
|
||
warnings.filterwarnings("ignore")
|
||
|
||
|
||
# ## Multiscale Problem
|
||
#
|
||
# We begin by presenting the problem, which is also discussed in Section 2 of [*On the Eigenvector Bias of Fourier Feature Networks: From Regression to Solving Multi-Scale PDEs with Physics-Informed Neural Networks*](https://doi.org/10.1016/j.cma.2021.113938). The one-dimensional Poisson problem we aim to solve is mathematically defined as:
|
||
#
|
||
# \begin{equation}
|
||
# \begin{cases}
|
||
# \Delta u(x) + f(x) = 0 \quad x \in [0,1], \\
|
||
# u(x) = 0 \quad x \in \partial[0,1],
|
||
# \end{cases}
|
||
# \end{equation}
|
||
#
|
||
# We define the solution as:
|
||
#
|
||
# $$
|
||
# u(x) = \sin(2\pi x) + 0.1 \sin(50\pi x),
|
||
# $$
|
||
#
|
||
# which leads to the corresponding force term:
|
||
#
|
||
# $$
|
||
# f(x) = (2\pi)^2 \sin(2\pi x) + 0.1 (50 \pi)^2 \sin(50\pi x).
|
||
# $$
|
||
#
|
||
# While this example is simple and pedagogical, it's important to note that the solution exhibits low-frequency behavior in the macro-scale and high-frequency behavior in the micro-scale. This characteristic is common in many practical scenarios.
|
||
#
|
||
# Below is the implementation of the `Poisson` problem as described mathematically above.
|
||
# > **👉 We have a dedicated [tutorial](https://mathlab.github.io/PINA/tutorial16/tutorial.html) to teach how to build a Problem from scratch — have a look if you're interested!**
|
||
|
||
# In[2]:
|
||
|
||
|
||
class Poisson(SpatialProblem):
|
||
output_variables = ["u"]
|
||
spatial_domain = CartesianDomain({"x": [0, 1]})
|
||
|
||
def poisson_equation(input_, output_):
|
||
x = input_.extract("x")
|
||
u_xx = laplacian(output_, input_, components=["u"], d=["x"])
|
||
f = ((2 * torch.pi) ** 2) * torch.sin(2 * torch.pi * x) + 0.1 * (
|
||
(50 * torch.pi) ** 2
|
||
) * torch.sin(50 * torch.pi * x)
|
||
return u_xx + f
|
||
|
||
domains = {
|
||
"bound_cond0": CartesianDomain({"x": 0.0}),
|
||
"bound_cond1": CartesianDomain({"x": 1.0}),
|
||
"phys_cond": spatial_domain,
|
||
}
|
||
# here we write the problem conditions
|
||
conditions = {
|
||
"bound_cond0": Condition(
|
||
domain="bound_cond0", equation=FixedValue(0.0)
|
||
),
|
||
"bound_cond1": Condition(
|
||
domain="bound_cond1", equation=FixedValue(0.0)
|
||
),
|
||
"phys_cond": Condition(
|
||
domain="phys_cond", equation=Equation(poisson_equation)
|
||
),
|
||
}
|
||
|
||
def solution(self, x):
|
||
return torch.sin(2 * torch.pi * x) + 0.1 * torch.sin(50 * torch.pi * x)
|
||
|
||
|
||
problem = Poisson()
|
||
|
||
# let's discretise the domain
|
||
problem.discretise_domain(128, "grid", domains=["phys_cond"])
|
||
problem.discretise_domain(1, "grid", domains=["bound_cond0", "bound_cond1"])
|
||
|
||
|
||
# A standard PINN approach would involve fitting the model using a Feed Forward (fully connected) Neural Network. For a conventional fully-connected neural network, it is relatively easy to approximate a function $u$, given sufficient data inside the computational domain.
|
||
#
|
||
# However, solving high-frequency or multi-scale problems presents significant challenges to PINNs, especially when the number of data points is insufficient to capture the different scales effectively.
|
||
#
|
||
# Below, we run a simulation using both the `PINN` solver and the self-adaptive `SAPINN` solver, employing a [`FeedForward`](https://mathlab.github.io/PINA/_modules/pina/model/feed_forward.html#FeedForward) model.
|
||
#
|
||
|
||
# In[3]:
|
||
|
||
|
||
# training with PINN and visualize results
|
||
pinn = PINN(
|
||
problem=problem,
|
||
model=FeedForward(
|
||
input_dimensions=1, output_dimensions=1, layers=[100, 100, 100]
|
||
),
|
||
)
|
||
|
||
trainer = Trainer(
|
||
pinn,
|
||
max_epochs=1500,
|
||
accelerator="cpu",
|
||
enable_model_summary=False,
|
||
val_size=0.0,
|
||
train_size=1.0,
|
||
test_size=0.0,
|
||
)
|
||
trainer.train()
|
||
|
||
# training with PINN and visualize results
|
||
sapinn = SAPINN(
|
||
problem=problem,
|
||
model=FeedForward(
|
||
input_dimensions=1, output_dimensions=1, layers=[100, 100, 100]
|
||
),
|
||
)
|
||
trainer_sapinn = Trainer(
|
||
sapinn,
|
||
max_epochs=1500,
|
||
accelerator="cpu",
|
||
enable_model_summary=False,
|
||
val_size=0.0,
|
||
train_size=1.0,
|
||
test_size=0.0,
|
||
)
|
||
trainer_sapinn.train()
|
||
|
||
|
||
# In[4]:
|
||
|
||
|
||
# define the function to plot the solution obtained using matplotlib
|
||
def plot_solution(pinn_to_use, title):
|
||
pts = pinn_to_use.problem.spatial_domain.sample(256, "grid", variables="x")
|
||
predicted_output = pinn_to_use(pts).extract("u").tensor.detach()
|
||
true_output = pinn_to_use.problem.solution(pts).detach()
|
||
plt.plot(
|
||
pts.extract(["x"]), predicted_output, label="Neural Network solution"
|
||
)
|
||
plt.plot(pts.extract(["x"]), true_output, label="True solution")
|
||
plt.title(title)
|
||
plt.legend()
|
||
|
||
|
||
# plot the solution of the two PINNs
|
||
plot_solution(pinn, "PINN solution")
|
||
plt.figure()
|
||
plot_solution(sapinn, "Self Adaptive PINN solution")
|
||
|
||
|
||
# We can clearly observe that neither of the two solvers has successfully learned the solution.
|
||
# The issue is not with the optimization strategy (i.e., the solver), but rather with the model used to solve the problem.
|
||
# A simple `FeedForward` network struggles to handle multiscale problems, especially when there are not enough collocation points to capture the different scales effectively.
|
||
#
|
||
# Next, let's compute the $l_2$ relative error for both the `PINN` and `SAPINN` solutions:
|
||
|
||
# In[5]:
|
||
|
||
|
||
# l2 loss from PINA losses
|
||
l2_loss = LpLoss(p=2, relative=False)
|
||
|
||
# sample new test points
|
||
pts = pts = problem.spatial_domain.sample(100, "grid")
|
||
print(
|
||
f"Relative l2 error PINN {l2_loss(pinn(pts), problem.solution(pts)).item():.2%}"
|
||
)
|
||
print(
|
||
f"Relative l2 error SAPINN {l2_loss(sapinn(pts), problem.solution(pts)).item():.2%}"
|
||
)
|
||
|
||
|
||
# Which is indeed very high!
|
||
#
|
||
# ## Fourier Feature Embedding in PINA
|
||
# Fourier Feature Embedding is a technique used to transform the input features, aiding the network in learning multiscale variations in the output. It was first introduced in [*On the Eigenvector Bias of Fourier Feature Networks: From Regression to Solving Multi-Scale PDEs with Physics-Informed Neural Networks*](https://doi.org/10.1016/j.cma.2021.113938), where it demonstrated excellent results for multiscale problems.
|
||
#
|
||
# The core idea behind Fourier Feature Embedding is to map the input $\mathbf{x}$ into an embedding $\tilde{\mathbf{x}}$, defined as:
|
||
#
|
||
# $$
|
||
# \tilde{\mathbf{x}} = \left[\cos\left( \mathbf{B} \mathbf{x} \right), \sin\left( \mathbf{B} \mathbf{x} \right)\right],
|
||
# $$
|
||
#
|
||
# where $\mathbf{B}_{ij} \sim \mathcal{N}(0, \sigma^2)$. This simple operation allows the network to learn across multiple scales!
|
||
#
|
||
# In **PINA**, we have already implemented this feature as a `layer` called [`FourierFeatureEmbedding`](https://mathlab.github.io/PINA/_rst/layers/fourier_embedding.html). Below, we will build the *Multi-scale Fourier Feature Architecture*. In this architecture, multiple Fourier feature embeddings (initialized with different $\sigma$ values) are applied to the input coordinates. These embeddings are then passed through the same fully-connected neural network, and the outputs are concatenated with a final linear layer.
|
||
#
|
||
|
||
# In[6]:
|
||
|
||
|
||
class MultiscaleFourierNet(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.embedding1 = FourierFeatureEmbedding(
|
||
input_dimension=1, output_dimension=100, sigma=1
|
||
)
|
||
self.embedding2 = FourierFeatureEmbedding(
|
||
input_dimension=1, output_dimension=100, sigma=10
|
||
)
|
||
self.layers = FeedForward(
|
||
input_dimensions=100, output_dimensions=100, layers=[100]
|
||
)
|
||
self.final_layer = torch.nn.Linear(2 * 100, 1)
|
||
|
||
def forward(self, x):
|
||
e1 = self.layers(self.embedding1(x))
|
||
e2 = self.layers(self.embedding2(x))
|
||
return self.final_layer(torch.cat([e1, e2], dim=-1))
|
||
|
||
|
||
# We will train the `MultiscaleFourierNet` using the `PINN` solver.
|
||
# Feel free to experiment with other PINN variants as well, such as `SAPINN`, `GPINN`, `CompetitivePINN`, and others, to see how they perform on this multiscale problem.
|
||
|
||
# In[7]:
|
||
|
||
|
||
multiscale_pinn = PINN(problem=problem, model=MultiscaleFourierNet())
|
||
trainer = Trainer(
|
||
multiscale_pinn,
|
||
max_epochs=1500,
|
||
accelerator="cpu",
|
||
enable_model_summary=False,
|
||
val_size=0.0,
|
||
train_size=1.0,
|
||
test_size=0.0,
|
||
)
|
||
trainer.train()
|
||
|
||
|
||
# Let us now plot the solution and compute the relative $l_2$ again!
|
||
|
||
# In[8]:
|
||
|
||
|
||
# plot solution obtained
|
||
plot_solution(multiscale_pinn, "Multiscale PINN solution")
|
||
|
||
# sample new test points
|
||
pts = pts = problem.spatial_domain.sample(100, "grid")
|
||
print(
|
||
f"Relative l2 error PINN with MultiscaleFourierNet: {l2_loss(multiscale_pinn(pts), problem.solution(pts)).item():.2%}"
|
||
)
|
||
|
||
|
||
# It is clear that the network has learned the correct solution, with a very low error. Of course, longer training and a more expressive neural network could further improve the results!
|
||
#
|
||
# ## What's Next?
|
||
#
|
||
# Congratulations on completing the one-dimensional Poisson tutorial of **PINA** using `FourierFeatureEmbedding`! There are many potential next steps you can explore:
|
||
#
|
||
# 1. **Train the network longer or with different layer sizes**: Experiment with different configurations to improve accuracy.
|
||
#
|
||
# 2. **Understand the role of `sigma` in `FourierFeatureEmbedding`**: The original paper provides insightful details on the impact of `sigma`. It's a good next step to dive deeper into its effect.
|
||
#
|
||
# 3. **Implement the *Spatio-temporal Multi-scale Fourier Feature Architecture***: Code this architecture for a more complex, time-dependent PDE (refer to Section 3 of the original paper).
|
||
#
|
||
# 4. **...and many more!**: There are countless directions to further explore, from testing on different problems to refining the model architecture.
|
||
#
|
||
# For more resources and tutorials, check out the [PINA Documentation](https://mathlab.github.io/PINA/).
|