Tutorials v0.1 (#178)
Tutorial update and small fixes * Tutorials update + Tutorial FNO * Create a metric tracker callback * Update PINN for logging * Update plotter for plotting * Small fix LabelTensor * Small fix FNO --------- Co-authored-by: Dario Coscia <dariocoscia@cli-10-110-13-250.WIFIeduroamSTUD.units.it> Co-authored-by: Dario Coscia <dariocoscia@dhcp-176.eduroam.sissa.it>
This commit is contained in:
Dario Coscia
committed by
Nicola Demo
Nicola Demo
parent
939353f517
commit
a9b1bd2826
@@ -1,12 +1,12 @@
|
||||
Tutorial 3: resolution of wave equation with custom Network
|
||||
===========================================================
|
||||
Tutorial 3: resolution of wave equation with hard constraint PINNs.
|
||||
===================================================================
|
||||
|
||||
The problem solution
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
In this tutorial we present how to solve the wave equation using the
|
||||
``SpatialProblem`` and ``TimeDependentProblem`` class, and the
|
||||
``Network`` class for building custom **torch** networks.
|
||||
In this tutorial we present how to solve the wave equation using hard
|
||||
constraint PINNs. For doing so we will build a costum torch model and
|
||||
pass it to the ``PINN`` solver.
|
||||
|
||||
The problem is written in the following form:
|
||||
|
||||
@@ -29,9 +29,13 @@ First of all, some useful imports.
|
||||
import torch
|
||||
|
||||
from pina.problem import SpatialProblem, TimeDependentProblem
|
||||
from pina.operators import nabla, grad
|
||||
from pina.model import Network
|
||||
from pina import Condition, Span, PINN, Plotter
|
||||
from pina.operators import laplacian, grad
|
||||
from pina.geometry import CartesianDomain
|
||||
from pina.solvers import PINN
|
||||
from pina.trainer import Trainer
|
||||
from pina.equation import Equation
|
||||
from pina.equation.equation_factory import FixedValue
|
||||
from pina import Condition, Plotter
|
||||
|
||||
Now, the wave problem is written in PINA code as a class, inheriting
|
||||
from ``SpatialProblem`` and ``TimeDependentProblem`` since we deal with
|
||||
@@ -44,31 +48,27 @@ predicted one.
|
||||
|
||||
class Wave(TimeDependentProblem, SpatialProblem):
|
||||
output_variables = ['u']
|
||||
spatial_domain = Span({'x': [0, 1], 'y': [0, 1]})
|
||||
temporal_domain = Span({'t': [0, 1]})
|
||||
spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
|
||||
temporal_domain = CartesianDomain({'t': [0, 1]})
|
||||
|
||||
def wave_equation(input_, output_):
|
||||
u_t = grad(output_, input_, components=['u'], d=['t'])
|
||||
u_tt = grad(u_t, input_, components=['dudt'], d=['t'])
|
||||
nabla_u = nabla(output_, input_, components=['u'], d=['x', 'y'])
|
||||
nabla_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])
|
||||
return nabla_u - u_tt
|
||||
|
||||
def nil_dirichlet(input_, output_):
|
||||
value = 0.0
|
||||
return output_.extract(['u']) - value
|
||||
|
||||
def initial_condition(input_, output_):
|
||||
u_expected = (torch.sin(torch.pi*input_.extract(['x'])) *
|
||||
torch.sin(torch.pi*input_.extract(['y'])))
|
||||
return output_.extract(['u']) - u_expected
|
||||
|
||||
conditions = {
|
||||
'gamma1': Condition(location=Span({'x': [0, 1], 'y': 1, 't': [0, 1]}), function=nil_dirichlet),
|
||||
'gamma2': Condition(location=Span({'x': [0, 1], 'y': 0, 't': [0, 1]}), function=nil_dirichlet),
|
||||
'gamma3': Condition(location=Span({'x': 1, 'y': [0, 1], 't': [0, 1]}), function=nil_dirichlet),
|
||||
'gamma4': Condition(location=Span({'x': 0, 'y': [0, 1], 't': [0, 1]}), function=nil_dirichlet),
|
||||
't0': Condition(location=Span({'x': [0, 1], 'y': [0, 1], 't': 0}), function=initial_condition),
|
||||
'D': Condition(location=Span({'x': [0, 1], 'y': [0, 1], 't': [0, 1]}), function=wave_equation),
|
||||
'gamma1': Condition(location=CartesianDomain({'x': [0, 1], 'y': 1, 't': [0, 1]}), equation=FixedValue(0.)),
|
||||
'gamma2': Condition(location=CartesianDomain({'x': [0, 1], 'y': 0, 't': [0, 1]}), equation=FixedValue(0.)),
|
||||
'gamma3': Condition(location=CartesianDomain({'x': 1, 'y': [0, 1], 't': [0, 1]}), equation=FixedValue(0.)),
|
||||
'gamma4': Condition(location=CartesianDomain({'x': 0, 'y': [0, 1], 't': [0, 1]}), equation=FixedValue(0.)),
|
||||
't0': Condition(location=CartesianDomain({'x': [0, 1], 'y': [0, 1], 't': 0}), equation=Equation(initial_condition)),
|
||||
'D': Condition(location=CartesianDomain({'x': [0, 1], 'y': [0, 1], 't': [0, 1]}), equation=Equation(wave_equation)),
|
||||
}
|
||||
|
||||
def wave_sol(self, pts):
|
||||
@@ -80,101 +80,100 @@ predicted one.
|
||||
|
||||
problem = Wave()
|
||||
|
||||
After the problem, a **torch** model is needed to solve the PINN. With
|
||||
the ``Network`` class the users can convert any **torch** model in a
|
||||
**PINA** model which uses label tensors with a single line of code. We
|
||||
will write a simple residual network using linear layers. Here we
|
||||
implement a simple residual network composed by linear torch layers.
|
||||
After the problem, a **torch** model is needed to solve the PINN.
|
||||
Usually many models are already implemented in ``PINA``, but the user
|
||||
has the possibility to build his/her own model in ``pyTorch``. The hard
|
||||
constraint we impose are on the boundary of the spatial domain.
|
||||
Specificly our solution is written as:
|
||||
|
||||
This neural network takes as input the coordinates (in this case
|
||||
:math:`x`, :math:`y` and :math:`t`) and provides the unkwown field of
|
||||
the Wave problem. The residual of the equations are evaluated at several
|
||||
sampling points (which the user can manipulate using the method
|
||||
``span_pts``) and the loss minimized by the neural network is the sum of
|
||||
the residuals.
|
||||
.. math:: u_{\rm{pinn}} = xy(1-x)(1-y)\cdot NN(x, y, t),
|
||||
|
||||
where :math:`NN` is the neural net output. This neural network takes as
|
||||
input the coordinates (in this case :math:`x`, :math:`y` and :math:`t`)
|
||||
and provides the unkwown field of the Wave problem. By construction it
|
||||
is zero on the boundaries. The residual of the equations are evaluated
|
||||
at several sampling points (which the user can manipulate using the
|
||||
method ``discretise_domain``) and the loss minimized by the neural
|
||||
network is the sum of the residuals.
|
||||
|
||||
.. code:: ipython3
|
||||
|
||||
class TorchNet(torch.nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
class HardMLP(torch.nn.Module):
|
||||
|
||||
def __init__(self, input_dim, output_dim):
|
||||
super().__init__()
|
||||
|
||||
self.residual = torch.nn.Sequential(torch.nn.Linear(3, 24),
|
||||
torch.nn.Tanh(),
|
||||
torch.nn.Linear(24, 3))
|
||||
|
||||
self.mlp = torch.nn.Sequential(torch.nn.Linear(3, 64),
|
||||
torch.nn.Tanh(),
|
||||
torch.nn.Linear(64, 1))
|
||||
def forward(self, x):
|
||||
residual_x = self.residual(x)
|
||||
return self.mlp(x + residual_x)
|
||||
|
||||
# model definition
|
||||
model = Network(model = TorchNet(),
|
||||
input_variables=problem.input_variables,
|
||||
output_variables=problem.output_variables,
|
||||
extra_features=None)
|
||||
self.layers = torch.nn.Sequential(torch.nn.Linear(input_dim, 20),
|
||||
torch.nn.Tanh(),
|
||||
torch.nn.Linear(20, 20),
|
||||
torch.nn.Tanh(),
|
||||
torch.nn.Linear(20, output_dim))
|
||||
|
||||
# here in the foward we implement the hard constraints
|
||||
def forward(self, x):
|
||||
hard = x.extract(['x'])*(1-x.extract(['x']))*x.extract(['y'])*(1-x.extract(['y']))
|
||||
return hard*self.layers(x)
|
||||
|
||||
In this tutorial, the neural network is trained for 2000 epochs with a
|
||||
learning rate of 0.001. These parameters can be modified as desired. We
|
||||
highlight that the generation of the sampling points and the train is
|
||||
here encapsulated within the function ``generate_samples_and_train``,
|
||||
but only for saving some lines of code in the next cells; that function
|
||||
is not mandatory in the **PINA** framework. The training takes
|
||||
approximately one minute.
|
||||
In this tutorial, the neural network is trained for 3000 epochs with a
|
||||
learning rate of 0.001 (default in ``PINN``). Training takes
|
||||
approximately 1 minute.
|
||||
|
||||
.. code:: ipython3
|
||||
|
||||
def generate_samples_and_train(model, problem):
|
||||
# generate pinn object
|
||||
pinn = PINN(problem, model, lr=0.001)
|
||||
|
||||
pinn.span_pts(1000, 'random', locations=['D','t0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])
|
||||
pinn.train(1500, 150)
|
||||
return pinn
|
||||
|
||||
|
||||
pinn = generate_samples_and_train(model, problem)
|
||||
pinn = PINN(problem, HardMLP(len(problem.input_variables), len(problem.output_variables)))
|
||||
problem.discretise_domain(1000, 'random', locations=['D','t0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])
|
||||
trainer = Trainer(pinn, max_epochs=3000)
|
||||
trainer.train()
|
||||
|
||||
|
||||
.. parsed-literal::
|
||||
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00000] 1.021557e-01 1.350026e-02 4.368403e-03 6.463497e-03 1.698729e-03 5.513944e-02 2.098533e-02
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00001] 8.096325e-02 7.543423e-03 2.978407e-03 7.128799e-03 2.084145e-03 3.967418e-02 2.155431e-02
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00150] 4.684930e-02 9.609548e-03 3.093602e-03 7.733506e-03 2.570329e-03 1.896760e-02 4.874712e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00300] 3.519089e-02 6.642059e-03 2.865276e-03 6.399740e-03 2.900236e-03 1.244203e-02 3.941551e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00450] 2.766160e-02 5.089254e-03 2.789679e-03 5.370538e-03 3.071685e-03 7.834940e-03 3.505504e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00600] 2.361075e-02 4.279066e-03 2.785937e-03 4.689044e-03 3.101575e-03 5.907214e-03 2.847910e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00750] 8.005206e-02 3.891625e-03 2.690672e-03 3.808867e-03 3.402538e-03 6.042966e-03 6.021538e-02
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 00900] 1.892301e-02 3.592897e-03 2.639081e-03 3.797543e-03 2.988781e-03 3.860098e-03 2.044612e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 01050] 1.739456e-02 3.420912e-03 2.557583e-03 3.532733e-03 2.910482e-03 3.114843e-03 1.858010e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 01200] 1.663617e-02 3.213567e-03 2.571464e-03 3.355495e-03 2.749454e-03 3.247283e-03 1.498912e-03
|
||||
sum gamma1nil_di gamma2nil_di gamma3nil_di gamma4nil_di t0initial_co Dwave_equati
|
||||
[epoch 01350] 1.551488e-02 3.121611e-03 2.481438e-03 3.141828e-03 2.706321e-03 2.636140e-03 1.427544e-03
|
||||
[epoch 01500] 1.497287e-02 2.974171e-03 2.475442e-03 2.979754e-03 2.593079e-03 2.723322e-03 1.227099e-03
|
||||
GPU available: False, used: False
|
||||
TPU available: False, using: 0 TPU cores
|
||||
IPU available: False, using: 0 IPUs
|
||||
HPU available: False, using: 0 HPUs
|
||||
|
||||
| Name | Type | Params
|
||||
----------------------------------------
|
||||
0 | _loss | MSELoss | 0
|
||||
1 | _neural_net | Network | 521
|
||||
----------------------------------------
|
||||
521 Trainable params
|
||||
0 Non-trainable params
|
||||
521 Total params
|
||||
0.002 Total estimated model params size (MB)
|
||||
|
||||
|
||||
After the training is completed one can now plot some results using the
|
||||
``Plotter`` class of **PINA**.
|
||||
.. parsed-literal::
|
||||
|
||||
Epoch 2999: : 1it [00:00, 79.33it/s, v_num=5, mean_loss=0.00119, D_loss=0.00542, t0_loss=0.0017, gamma1_loss=0.000, gamma2_loss=0.000, gamma3_loss=0.000, gamma4_loss=0.000]
|
||||
|
||||
.. parsed-literal::
|
||||
|
||||
`Trainer.fit` stopped: `max_epochs=3000` reached.
|
||||
|
||||
|
||||
.. parsed-literal::
|
||||
|
||||
Epoch 2999: : 1it [00:00, 68.62it/s, v_num=5, mean_loss=0.00119, D_loss=0.00542, t0_loss=0.0017, gamma1_loss=0.000, gamma2_loss=0.000, gamma3_loss=0.000, gamma4_loss=0.000]
|
||||
|
||||
|
||||
Notice that the loss on the boundaries of the spatial domain is exactly
|
||||
zero, as expected! After the training is completed one can now plot some
|
||||
results using the ``Plotter`` class of **PINA**.
|
||||
|
||||
.. code:: ipython3
|
||||
|
||||
plotter = Plotter()
|
||||
|
||||
# plotting at fixed time t = 0.6
|
||||
plotter.plot(pinn, fixed_variables={'t': 0.6})
|
||||
# plotting at fixed time t = 0.0
|
||||
plotter.plot(trainer, fixed_variables={'t': 0.0})
|
||||
|
||||
# plotting at fixed time t = 0.5
|
||||
plotter.plot(trainer, fixed_variables={'t': 0.5})
|
||||
|
||||
# plotting at fixed time t = 1.
|
||||
plotter.plot(trainer, fixed_variables={'t': 1.0})
|
||||
|
||||
|
||||
|
||||
@@ -182,24 +181,10 @@ After the training is completed one can now plot some results using the
|
||||
.. image:: tutorial_files/tutorial_12_0.png
|
||||
|
||||
|
||||
We can also plot the pinn loss during the training to see the decrease.
|
||||
|
||||
.. code:: ipython3
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
plt.figure(figsize=(16, 6))
|
||||
plotter.plot_loss(pinn, label='Loss')
|
||||
|
||||
plt.grid()
|
||||
plt.legend()
|
||||
plt.show()
|
||||
.. image:: tutorial_files/tutorial_12_1.png
|
||||
|
||||
|
||||
|
||||
.. image:: tutorial_files/tutorial_14_0.png
|
||||
.. image:: tutorial_files/tutorial_12_2.png
|
||||
|
||||
|
||||
You can now trying improving the training by changing network, optimizer
|
||||
and its parameters, changin the sampling points,or adding extra
|
||||
features!
|
||||
|
||||
Reference in New Issue
Block a user