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>

docs/source/_rst/tutorial3/tutorial.rst

@@ -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!