Skip to content
Navigation Menu
{{ message }}
-
Notifications
You must be signed in to change notification settings - Fork 99
Expand file tree
/
Copy pathtutorial.py
More file actions
280 lines (223 loc) · 10.5 KB
/
tutorial.py
File metadata and controls
280 lines (223 loc) · 10.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
#!/usr/bin/env python
# coding: utf-8
# # Tutorial: Learning Bifurcating PDE Solutions with Physics-Informed Deep Ensembles
#
# [](https://colab.research.google.com/github/mathLab/PINA/blob/master/tutorials/tutorial14/tutorial.ipynb)
#
# This tutorial demonstrates how to use the Deep Ensemble Physics Informed Network (DeepEnsemblePINN) to learn PDEs exhibiting bifurcating behavior, as discussed in [*Learning and Discovering Multiple Solutions Using Physics-Informed Neural Networks with Random Initialization and Deep Ensemble*](https://arxiv.org/abs/2503.06320).
#
# Let’s begin by importing the necessary libraries.
# In[1]:
## 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 lightning.pytorch.callbacks import Callback
from pina import Trainer, Condition, LabelTensor
from pina.solver import DeepEnsemblePINN
from pina.model import FeedForward
from pina.operator import laplacian
from pina.problem import TimeDependentProblem
from pina.domain import CartesianDomain
from pina.equation import Equation
from pina.optim import TorchOptimizer
warnings.filterwarnings("ignore")
# ## Deep Ensemble
#
# Deep Ensemble methods improve model performance by leveraging the diversity of predictions generated by multiple neural networks trained on the same problem. Each network in the ensemble is trained independently—typically with different weight initializations or even slight variations in the architecture or data sampling. By combining their outputs (e.g., via averaging or majority voting), ensembles reduce overfitting, increase robustness, and improve generalization.
#
# This approach allows the ensemble to capture different perspectives of the problem, leading to more accurate and reliable predictions.
#
# <p align="center">
# <img src="http://raw.githubusercontent.com/mathLab/PINA/master/tutorials/static/deep_ensemble.png" alt="Deep ensemble" width="600"/>
# </p>
#
# The image above illustrates a Deep Ensemble setup, where multiple models attempt to predict the text from an image. While individual models may make errors (e.g., predicting "PONY" instead of "PINA"), combining their outputs—such as taking the majority vote—often leads to the correct result. This ensemble effect improves reliability by mitigating the impact of individual model biases.
#
#
# ## Deep Ensemble Physics-Informed Networks
#
# In the context of Physics-Informed Neural Networks (PINNs), Deep Ensembles help the network discover different branches or multiple solutions of a PDE that exhibits bifurcating behavior.
#
# By training a diverse set of models with different initializations, Deep Ensemble methods overcome the limitations of single-initialization models, which may converge to only one of the possible solutions. This approach is particularly useful when the solution space of the problem contains multiple valid physical states or behaviors.
#
#
# ## The Bratu Problem
#
# In this tutorial, we'll train a `DeepEnsemblePINN` solver to solve a bifurcating ODE known as the **Bratu problem**. The ODE is given by:
#
# $$
# \frac{d^2u}{dt^2} + \lambda e^u = 0, \quad t \in (0, 1)
# $$
#
# with boundary conditions:
#
# $$
# u(0) = u(1) = 0,
# $$
#
# where $\lambda > 0$ is a scalar parameter. The analytical solutions to the 1D Bratu problem can be expressed as:
#
# $$
# u(t, \alpha) = 2 \log\left(\frac{\cosh(\alpha)}{\cosh(\alpha(1 - 2t))}\right),
# $$
#
# where $\alpha$ satisfies:
#
# $$
# \cosh(\alpha) - 2\sqrt{2}\alpha = 0.
# $$
#
# When $\lambda < 3.513830719$, the equation admits two solutions $\alpha_1$ and $\alpha_2$, which correspond to two distinct solutions of the original ODE: $u_1$ and $u_2$.
#
# In this tutorial, we set $\lambda = 1$, which leads to:
#
# - $\alpha_1 \approx 0.37929$
# - $\alpha_2 \approx 2.73468$
#
# We first write the problem class, we do not write the boundary conditions as we will hard impose them.
#
# > **👉 We have a dedicated [tutorial](https://mathlab.github.io/PINA/tutorial16/tutorial.html) to teach how to build a Problem — have a look if you're interested!**
#
# > **👉 We have a dedicated [tutorial](https://mathlab.github.io/PINA/tutorial3/tutorial.html) to teach how to impose hard constraints — have a look if you're interested!**
# In[2]:
# define bratu equation
def bratu_eq(input_, output_):
u_tt = laplacian(output_=output_, input_=input_, components=["u"], d=["t"])
return u_tt + torch.exp(output_)
# define true solution
def true_solution(x):
alpha1 = torch.tensor([0.37929])
alpha2 = torch.tensor([2.73468])
u1 = 2 * torch.log(torch.cosh(alpha1) / torch.cosh(alpha1 * (1 - 2 * x)))
u2 = 2 * torch.log(torch.cosh(alpha2) / torch.cosh(alpha2 * (1 - 2 * x)))
return u1, u2
# build problem class
class BratuProblem(TimeDependentProblem):
output_variables = ["u"]
temporal_domain = CartesianDomain({"t": [0, 1]})
domains = {"interior": temporal_domain}
conditions = {
"interior": Condition(domain="interior", equation=Equation(bratu_eq))
}
# define problem and discretise domain
problem = BratuProblem()
problem.discretise_domain(n=101, mode="grid", domains="interior")
# ## Defining the Deep Ensemble Models
#
# Now that the problem setup is complete, we move on to creating an **ensemble of models**. Each ensemble member will be a standard `FeedForward` neural network, wrapped inside a custom `Model` class.
#
# Each model's weights are initialized using a **normal distribution** with mean 0 and standard deviation 2. This random initialization is crucial to promote diversity across the ensemble members, allowing the models to converge to potentially different solutions of the PDE.
#
# The final ensemble is simply a **list of PyTorch models**, which we will later pass to the `DeepEnsemblePINN`
# In[3]:
# define a single model (ensemble member)
class Model(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
self.model = FeedForward(*args, **kwargs)
self.init_weights_gaussian()
def forward(self, x):
return x * (1 - x) * self.model(x)
def init_weights_gaussian(self):
for param in self.model.parameters():
if param.requires_grad:
torch.nn.init.normal_(param, mean=0.0, std=2.0)
# define a list of models with different initializations
models = [Model(1, 1, inner_size=50, n_layers=2) for _ in range(10)]
# Let's visualize the networks output before strated training
# In[4]:
# plot solution
with torch.no_grad():
pts = problem.input_pts["interior"]
for model in models:
plt.plot(pts, model(pts), "--")
plt.plot()
# As you can see we get different output since the neural networks are initialized differently.
#
# ## Training with `DeepEnsemblePINN`
#
# Now that everything is ready, we can train the models using the `DeepEnsemblePINN` solver! 🎯
#
# This solver is constructed by combining multiple neural network models that all aim to solve the same PDE. Each model $\mathcal{M}_{i \in \{1, \dots, 10\}}$ in the ensemble contributes a unique perspective due to different random initializations.
#
# This diversity allows the ensemble to **capture multiple branches or bifurcating solutions** of the problem, making it especially powerful for PDEs like the Bratu problem.
#
# Once the `DeepEnsemblePINN` solver is defined with all the models, we train them using the `Trainer` class, as with any other solver in **PINA**. We also build a callback to store the value of `u(0.5)` during training iterations.
# In[ ]:
# define the optimizers, one per model
optimizers = [TorchOptimizer(torch.optim.Adam, lr=0.006) for _ in range(10)]
# define solver
solver = DeepEnsemblePINN(
problem,
models,
optimizers=optimizers,
)
# callback
class StoreValue(Callback):
def on_train_epoch_start(self, trainer, pl_module):
input = LabelTensor(torch.tensor([[0.5]]), "t")
output = pl_module(input).tensor.flatten()
if trainer.current_epoch == 0:
self.store = [output]
else:
self.store.append(output)
# define trainer
trainer = Trainer(
solver,
max_epochs=500,
accelerator="cpu",
enable_model_summary=False,
callbacks=[StoreValue()],
)
# train
trainer.train()
# The training finished, let's first plot how the value of $u(0.5)$ changed during training
# In[6]:
with torch.no_grad():
metrics = torch.stack(trainer.callbacks[0].store, dim=0)
plt.plot(range(metrics.shape[0]), metrics)
plt.title("Ensemble Convergence")
plt.ylabel(r"$u(0.5)$")
plt.xlabel("epochs")
plt.plot()
# As you can see, different networks in the ensemble converge to different values pf $u(0.5)$ — this means we can actually **spot the bifurcation** in the solution space!
#
# This is a powerful demonstration of how **Deep Ensemble Physics-Informed Neural Networks** are capable of learning **multiple valid solutions** of a PDE that exhibits bifurcating behavior.
#
# We can also visualize the ensemble predictions to better observe the multiple branches:
#
# In[7]:
# plot solution
with torch.no_grad():
pts = problem.input_pts["interior"]
u_ensemble = solver(pts)
u1, u2 = true_solution(pts)
plt.plot(pts, u1, label="Reference solution u1")
plt.plot(pts, u2, label="Reference solution u2")
for idx, sol in enumerate(u_ensemble):
if idx == 0:
plt.plot(pts, sol, "--", label="PINNs", c="k")
else:
plt.plot(pts, sol, "--", c="k")
plt.legend()
plt.plot()
plt.show()
# ## What's Next?
#
# You have completed the tutorial on deep ensemble PINNs for bifurcating PDEs, well don! There are many potential next steps you can explore:
#
# 1. **Train the network longer or with different hyperparameters**: Experiment with different configurations of the single model, you can compose an ensemble by also stacking models with different layers, activation, ... to improve accuracy.
#
# 2. **Solve more complex problems**: The original paper provides very complex problems that can be solved with PINA, we suggest you to try implement and solve them!
#
# 3. **...and many more!**: There are countless directions to further explore, for example, what does it happen when you vary the network initialization hyperparameters?
#
# For more resources and tutorials, check out the [PINA Documentation](https://mathlab.github.io/PINA/).
You can’t perform that action at this time.