Tutorial: Learning Multiscale PDEs Using Fourier Feature NetworksΒΆ
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.
Letβs begin by importing the necessary libraries.
## routine needed to run the notebook on Google Colab
try:
import google.colab
IN_COLAB = True
except:
IN_COLAB = False
if IN_COLAB:
!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. 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 to teach how to build a Problem from scratch β have a look if you're interested!
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
model.
# 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()
You are using the plain ModelCheckpoint callback. Consider using LitModelCheckpoint which with seamless uploading to Model registry.
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
`Trainer.fit` stopped: `max_epochs=1500` reached.
You are using the plain ModelCheckpoint callback. Consider using LitModelCheckpoint which with seamless uploading to Model registry.
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
`Trainer.fit` stopped: `max_epochs=1500` reached.
# 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:
# 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%}"
)
Relative l2 error PINN 2912.82% Relative l2 error SAPINN 2788.09%
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, 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
. 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.
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.
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()
You are using the plain ModelCheckpoint callback. Consider using LitModelCheckpoint which with seamless uploading to Model registry.
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
`Trainer.fit` stopped: `max_epochs=1500` reached.
Let us now plot the solution and compute the relative $l_2$ again!
# 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%}"
)
Relative l2 error PINN with MultiscaleFourierNet: 14.61%
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:
Train the network longer or with different layer sizes: Experiment with different configurations to improve accuracy.
Understand the role of
sigma
inFourierFeatureEmbedding
: The original paper provides insightful details on the impact ofsigma
. It's a good next step to dive deeper into its effect.**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).
...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.