Examples update for v0.1 (#206)

* modify examples/problems
* modify tutorials

---------

Co-authored-by: Dario Coscia <dariocoscia@dhcp-235.eduroam.sissa.it>
Co-authored-by: Dario Coscia <dariocoscia@dhcp-015.eduroam.sissa.it>

examples/problems/burgers.py

@@ -1,9 +1,5 @@
-import torch
+""" Burgers' problem. """
 
-from pina.problem import TimeDependentProblem, SpatialProblem
-from pina.operators import grad
-from pina import Condition
-from pina.span import Span
 
 # ===================================================== #
 #                                                       #
@@ -17,12 +13,16 @@ from pina.span import Span
 #                                                       #
 # ===================================================== #
 
-class Burgers1D(TimeDependentProblem, SpatialProblem):
 
-    # assign output/ spatial and temporal variables
+import torch
-    output_variables = ['u']
+from pina.geometry import CartesianDomain
-    spatial_domain = Span({'x': [-1, 1]})
+from pina import Condition
-    temporal_domain = Span({'t': [0, 1]})
+from pina.problem import TimeDependentProblem, SpatialProblem
+from pina.operators import grad
+from pina.equation import FixedValue, Equation
+
+        
+class Burgers1D(TimeDependentProblem, SpatialProblem):
 
     # define the burger equation
     def burger_equation(input_, output_):
@@ -34,20 +34,20 @@ class Burgers1D(TimeDependentProblem, SpatialProblem):
             (0.01/torch.pi)*ddu.extract(['ddudxdx'])
         )
 
-    # define nill dirichlet boundary conditions
-    def nil_dirichlet(input_, output_):
-        u_expected = 0.0
-        return output_.extract(['u']) - u_expected
-
     # define initial condition
     def initial_condition(input_, output_):
         u_expected = -torch.sin(torch.pi*input_.extract(['x']))
         return output_.extract(['u']) - u_expected
 
+    # assign output/ spatial and temporal variables
+    output_variables = ['u']
+    spatial_domain = CartesianDomain({'x': [-1, 1]})
+    temporal_domain = CartesianDomain({'t': [0, 1]})
+
     # problem condition statement
     conditions = {
-        'gamma1': Condition(location=Span({'x': -1, 't': [0, 1]}), function=nil_dirichlet),
+        'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)),
-        'gamma2': Condition(location=Span({'x':  1, 't': [0, 1]}), function=nil_dirichlet),
+        'gamma2': Condition(location=CartesianDomain({'x':  1, 't': [0, 1]}), equation=FixedValue(0.)),
-        't0': Condition(location=Span({'x': [-1, 1], 't': 0}), function=initial_condition),
+        't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)),
-        'D': Condition(location=Span({'x': [-1, 1], 't': [0, 1]}), function=burger_equation),
+        'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Equation(burger_equation)),
     }

examples/problems/elliptic_optimal_control.py

@@ -1,47 +0,0 @@
-# import torch
-# from pina.problem import Problem
-# from pina.segment import Segment
-# from pina.cube import Cube
-# from pina.problem2d import Problem2D
-
-# xmin, xmax, ymin, ymax = -1, 1, -1, 1
-
-# class EllipticOptimalControl(Problem2D):
-
-#     def __init__(self, alpha=1):
-
-#         def term1(input_, output_):
-#             grad_p = self.grad(output_.extract(['p']), input_)
-#             gradgrad_p_x1 = self.grad(grad_p.extract(['x1']), input_)
-#             gradgrad_p_x2 = self.grad(grad_p.extract(['x2']), input_)
-#             yd = 2.0
-#             return output_.extract(['y']) - yd - (gradgrad_p_x1.extract(['x1']) + gradgrad_p_x2.extract(['x2']))
-
-#         def term2(input_, output_):
-#             grad_y = self.grad(output_.extract(['y']), input_)
-#             gradgrad_y_x1 = self.grad(grad_y.extract(['x1']), input_)
-#             gradgrad_y_x2 = self.grad(grad_y.extract(['x2']), input_)
-#             return - (gradgrad_y_x1.extract(['x1']) + gradgrad_y_x2.extract(['x2'])) - output_.extract(['u'])
-
-#         def term3(input_, output_):
-#             return output_.extract(['p']) - output_.extract(['u'])*alpha
-
-
-#         def nil_dirichlet(input_, output_):
-#             y_value = 0.0
-#             p_value = 0.0
-#             return torch.abs(output_.extract(['y']) - y_value) + torch.abs(output_.extract(['p']) - p_value)
-
-#         self.conditions = {
-#             'gamma1': {'location': Segment((xmin, ymin), (xmax, ymin)), 'func': nil_dirichlet},
-#             'gamma2': {'location': Segment((xmax, ymin), (xmax, ymax)), 'func': nil_dirichlet},
-#             'gamma3': {'location': Segment((xmax, ymax), (xmin, ymax)), 'func': nil_dirichlet},
-#             'gamma4': {'location': Segment((xmin, ymax), (xmin, ymin)), 'func': nil_dirichlet},
-#             'D1': {'location': Cube([[xmin, xmax], [ymin, ymax]]), 'func': [term1, term2, term3]},
-#         }
-
-#         self.input_variables = ['x1', 'x2']
-#         self.output_variables = ['u', 'p', 'y']
-#         self.spatial_domain = Cube([[xmin, xmax], [xmin, xmax]])
-
-raise NotImplementedError('not available problem at the moment...')

examples/problems/first_order_ode.py

@@ -1,7 +1,5 @@
-from pina.problem import SpatialProblem
+""" Simple ODE problem. """
-from pina import Condition, Span
+
-from pina.operators import grad
-import torch
 
 # ===================================================== #
 #                                                       #
@@ -11,16 +9,28 @@ import torch
 #           y --> field variable                        #
 #           x --> spatial variable                      #
 #                                                       #
+#  The equation is:                                     #
+#           dy(x)/dx + y(x) = x                         #
+#                                                       #
 # ===================================================== #
 
+
+from pina.problem import SpatialProblem
+from pina import Condition
+from pina.geometry import CartesianDomain
+from pina.operators import grad
+from pina.equation import Equation, FixedValue
+import torch
+
+
 class FirstOrderODE(SpatialProblem):
 
     # variable domain range
-    x_rng = [0, 5]
+    x_rng = [0., 5.]
     # field variable
     output_variables = ['y']
     # create domain
-    spatial_domain = Span({'x': x_rng})
+    spatial_domain = CartesianDomain({'x': x_rng})
 
     # define the ode
     def ode(input_, output_):
@@ -28,11 +38,6 @@ class FirstOrderODE(SpatialProblem):
         x = input_
         return grad(y, x) + y - x
 
-    # define initial conditions
-    def fixed(input_, output_):
-        exp_value = 1.
-        return output_ - exp_value
-
     # define real solution
     def solution(self, input_):
         x = input_
@@ -40,7 +45,8 @@ class FirstOrderODE(SpatialProblem):
 
     # define problem conditions
     conditions = {
-        'bc': Condition(location=Span({'x': x_rng[0]}), function=fixed),
+        'BC': Condition(location=CartesianDomain({'x': x_rng[0]}), equation=FixedValue(1.)),
-        'dd': Condition(location=Span({'x': x_rng}), function=ode),
+        'D': Condition(location=CartesianDomain({'x': x_rng}), equation=Equation(ode)),
     }
+
     truth_solution = solution

examples/problems/parametric_elliptic_optimal_control.py

@@ -1,52 +1,80 @@
-import numpy as np
+""" Poisson OCP problem. """
-import torch
-from pina.segment import Segment
-from pina.cube import Cube
-from pina.problem2d import Problem2D
-
-xmin, xmax, ymin, ymax = -1, 1, -1, 1
-
-class ParametricEllipticOptimalControl(Problem2D):
-
-    def __init__(self, alpha=1):
-
-        def term1(input_, param_, output_):
-            grad_p = self.grad(output_['p'], input_)
-            gradgrad_p_x1 = self.grad(grad_p['x1'], input_)
-            gradgrad_p_x2 = self.grad(grad_p['x2'], input_)
-            return output_['y'] - param_ - (gradgrad_p_x1['x1'] + gradgrad_p_x2['x2'])
-
-        def term2(input_, param_, output_):
-            grad_y = self.grad(output_['y'], input_)
-            gradgrad_y_x1 = self.grad(grad_y['x1'], input_)
-            gradgrad_y_x2 = self.grad(grad_y['x2'], input_)
-            return - (gradgrad_y_x1['x1'] + gradgrad_y_x2['x2']) - output_['u_param']
-
-        def term3(input_, param_, output_):
-            return output_['p'] - output_['u_param']*alpha
 
 
-        def term(input_, param_, output_):
+from pina import Condition
-            return term1( input_, param_, output_)  +term2( input_, param_, output_) + term3( input_, param_, output_)
+from pina.geometry import CartesianDomain
+from pina.equation import SystemEquation, FixedValue
+from pina.problem import SpatialProblem, ParametricProblem
+from pina.operators import laplacian
 
-        def nil_dirichlet(input_, param_, output_):
+# ===================================================== #
-            y_value = 0.0
+#                                                       #
-            p_value = 0.0
+#  This script implements the two dimensional           #
-            return torch.abs(output_['y'] - y_value) + torch.abs(output_['p'] - p_value)
+#  Parametric Elliptic Optimal Control problem.         #
+#  The ParametricEllipticOptimalControl class is        #
+#  inherited from TimeDependentProblem, SpatialProblem  #
+#  and we denote:                                       #
+#           u --> field variable                        #
+#           p --> field variable                        #
+#           y --> field variable                        #
+#           x1, x2 --> spatial variables                #
+#           mu, alpha --> problem parameters            #
+#                                                       #
+#  More info in https://arxiv.org/pdf/2110.13530.pdf    #
+#  Section 4.2 of the article                           #
+# ===================================================== #
 
-        self.conditions = {
-            'gamma1': {'location': Segment((xmin, ymin), (xmax, ymin)), 'func': nil_dirichlet},
-            'gamma2': {'location': Segment((xmax, ymin), (xmax, ymax)), 'func': nil_dirichlet},
-            'gamma3': {'location': Segment((xmax, ymax), (xmin, ymax)), 'func': nil_dirichlet},
-            'gamma4': {'location': Segment((xmin, ymax), (xmin, ymin)), 'func': nil_dirichlet},
-            'D1': {'location': Cube([[xmin, xmax], [ymin, ymax]]), 'func': term},
-            #'D2': {'location': Cube([[0, 1], [0, 1]]), 'func': term2},
-            #'D3': {'location': Cube([[0, 1], [0, 1]]), 'func': term3}
-        }
 
-        self.input_variables = ['x1', 'x2']
+class ParametricEllipticOptimalControl(SpatialProblem, ParametricProblem):
-        self.output_variables = ['u', 'p', 'y']
-        self.parameters = ['mu']
-        self.spatial_domain = Cube([[xmin, xmax], [xmin, xmax]])
-        self.parameter_domain = np.array([[0.5, 3]])
 
+    # setting spatial variables ranges
+    xmin, xmax, ymin, ymax = -1, 1, -1, 1
+    x_range = [xmin, xmax]
+    y_range = [ymin, ymax]
+    # setting parameters range
+    amin, amax = 0.0001, 1
+    mumin, mumax = 0.5, 3
+    mu_range = [mumin, mumax]
+    a_range = [amin, amax]
+    # setting field variables
+    output_variables = ['u', 'p', 'y']
+    # setting spatial and parameter domain
+    spatial_domain = CartesianDomain({'x1': x_range, 'x2': y_range})
+    parameter_domain = CartesianDomain({'mu': mu_range, 'alpha': a_range})
+
+    # equation terms as in https://arxiv.org/pdf/2110.13530.pdf
+    def term1(input_, output_):
+        laplace_p = laplacian(output_, input_, components=['p'], d=['x1', 'x2'])
+        return output_.extract(['y']) - input_.extract(['mu']) - laplace_p
+
+    def term2(input_, output_):
+        laplace_y = laplacian(output_, input_, components=['y'], d=['x1', 'x2'])
+        return - laplace_y - output_.extract(['u'])
+    
+    def fixed_y(input_, output_):
+        return output_.extract(['y'])
+
+    def fixed_p(input_, output_):
+        return output_.extract(['p']) 
+
+    # setting problem condition formulation
+    conditions = {
+        'gamma1': Condition(
+            location=CartesianDomain({'x1': x_range, 'x2':  1, 'mu': mu_range, 'alpha': a_range}),
+            equation=SystemEquation([fixed_y, fixed_p])),
+        'gamma2': Condition(
+            location=CartesianDomain({'x1': x_range, 'x2': -1, 'mu': mu_range, 'alpha': a_range}),
+            equation=SystemEquation([fixed_y, fixed_p])),
+        'gamma3': Condition(
+            location=CartesianDomain({'x1':  1, 'x2': y_range, 'mu': mu_range, 'alpha': a_range}),
+            equation=SystemEquation([fixed_y, fixed_p])),
+        'gamma4': Condition(
+            location=CartesianDomain({'x1': -1, 'x2': y_range, 'mu': mu_range, 'alpha': a_range}),
+            equation=SystemEquation([fixed_y, fixed_p])),
+        'D': Condition(
+            location=CartesianDomain(
+                {'x1': x_range, 'x2': y_range,
+                'mu': mu_range, 'alpha': a_range
+                }),
+            equation=SystemEquation([term1, term2])),
+    }

examples/problems/parametric_elliptic_optimal_control_alpha_variable.py

@@ -1,78 +0,0 @@
-import numpy as np
-import torch
-
-from pina import Span, Condition
-from pina.problem import SpatialProblem, ParametricProblem
-from pina.operators import grad, laplacian
-
-# ===================================================== #
-#                                                       #
-#  This script implements the two dimensional           #
-#  Parametric Elliptic Optimal Control problem.         #
-#  The ParametricEllipticOptimalControl class is        #
-#  inherited from TimeDependentProblem, SpatialProblem  #
-#  and we denote:                                       #
-#           u --> field variable                        #
-#           p --> field variable                        #
-#           y --> field variable                        #
-#           x1, x2 --> spatial variables                #
-#           mu, alpha --> problem parameters            #
-#                                                       #
-#  More info in https://arxiv.org/pdf/2110.13530.pdf    #
-#  Section 4.2 of the article                           #
-# ===================================================== #
-
-
-class ParametricEllipticOptimalControl(SpatialProblem, ParametricProblem):
-
-    # setting spatial variables ranges
-    xmin, xmax, ymin, ymax = -1, 1, -1, 1
-    x_range = [xmin, xmax]
-    y_range = [ymin, ymax]
-    # setting parameters range
-    amin, amax = 0.0001, 1
-    mumin, mumax = 0.5, 3
-    mu_range = [mumin, mumax]
-    a_range = [amin, amax]
-    # setting field variables
-    output_variables = ['u', 'p', 'y']
-    # setting spatial and parameter domain
-    spatial_domain = Span({'x1': x_range, 'x2': y_range})
-    parameter_domain = Span({'mu': mu_range, 'alpha': a_range})
-
-    # equation terms as in https://arxiv.org/pdf/2110.13530.pdf
-    def term1(input_, output_):
-        laplace_p = laplacian(output_, input_, components=['p'], d=['x1', 'x2'])
-        return output_.extract(['y']) - input_.extract(['mu']) - laplace_p
-
-    def term2(input_, output_):
-        laplace_y = laplacian(output_, input_, components=['y'], d=['x1', 'x2'])
-        return - laplace_y - output_.extract(['u_param'])
-
-    def state_dirichlet(input_, output_):
-        y_exp = 0.0
-        return output_.extract(['y']) - y_exp
-
-    def adj_dirichlet(input_, output_):
-        p_exp = 0.0
-        return output_.extract(['p']) - p_exp
-
-    # setting problem condition formulation
-    conditions = {
-        'gamma1': Condition(
-            location=Span({'x1': x_range, 'x2':  1, 'mu': mu_range, 'alpha': a_range}),
-            function=[state_dirichlet, adj_dirichlet]),
-        'gamma2': Condition(
-            location=Span({'x1': x_range, 'x2': -1, 'mu': mu_range, 'alpha': a_range}),
-            function=[state_dirichlet, adj_dirichlet]),
-        'gamma3': Condition(
-            location=Span({'x1':  1, 'x2': y_range, 'mu': mu_range, 'alpha': a_range}),
-            function=[state_dirichlet, adj_dirichlet]),
-        'gamma4': Condition(
-            location=Span({'x1': -1, 'x2': y_range, 'mu': mu_range, 'alpha': a_range}),
-            function=[state_dirichlet, adj_dirichlet]),
-        'D': Condition(
-            location=Span({'x1': x_range, 'x2': y_range,
-                  'mu': mu_range, 'alpha': a_range}),
-            function=[term1, term2]),
-    }

examples/problems/parametric_poisson.py

@@ -1,8 +1,5 @@
-import torch
+""" Parametric Poisson problem. """
 
-from pina.problem import SpatialProblem, ParametricProblem
-from pina.operators import laplacian
-from pina import Span, Condition
 
 # ===================================================== #
 #                                                       #
@@ -16,12 +13,20 @@ from pina import Span, Condition
 #                                                       #
 # ===================================================== #
 
+
+from pina.geometry import CartesianDomain
+from pina.problem import SpatialProblem, ParametricProblem
+from pina.operators import laplacian
+from pina.equation import FixedValue, Equation
+from pina import Condition
+import torch
+
 class ParametricPoisson(SpatialProblem, ParametricProblem):
 
     # assign output/ spatial and parameter variables
     output_variables = ['u']
-    spatial_domain = Span({'x': [-1, 1], 'y': [-1, 1]})
+    spatial_domain = CartesianDomain({'x': [-1, 1], 'y': [-1, 1]})
-    parameter_domain = Span({'mu1': [-1, 1], 'mu2': [-1, 1]})
+    parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]})
 
     # define the laplace equation
     def laplace_equation(input_, output_):
@@ -30,26 +35,21 @@ class ParametricPoisson(SpatialProblem, ParametricProblem):
                 - 2*(input_.extract(['y']) - input_.extract(['mu2']))**2)
         return laplacian(output_.extract(['u']), input_) - force_term
 
-    # define nill dirichlet boundary conditions
-    def nil_dirichlet(input_, output_):
-        value = 0.0
-        return output_.extract(['u']) - value
-
     # problem condition statement
     conditions = {
         'gamma1': Condition(
-            location=Span({'x': [-1, 1], 'y': 1, 'mu1': [-1, 1], 'mu2': [-1, 1]}),
+            location=CartesianDomain({'x': [-1, 1], 'y': 1, 'mu1': [-1, 1], 'mu2': [-1, 1]}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.)),
         'gamma2': Condition(
-            location=Span({'x': [-1, 1], 'y': -1, 'mu1': [-1, 1], 'mu2': [-1, 1]}),
+            location=CartesianDomain({'x': [-1, 1], 'y': -1, 'mu1': [-1, 1], 'mu2': [-1, 1]}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.)),
         'gamma3': Condition(
-            location=Span({'x': 1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}),
+            location=CartesianDomain({'x': 1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.)),
         'gamma4': Condition(
-            location=Span({'x': -1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}),
+            location=CartesianDomain({'x': -1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.)),
         'D': Condition(
-            location=Span({'x': [-1, 1], 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}),
+            location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}),
-            function=laplace_equation),
+            equation=Equation(laplace_equation)),
     }

examples/problems/poisson.py

@@ -1,10 +1,5 @@
-""" Poisson equation example. """
+""" Poisson problem. """
-import numpy as np
-import torch
 
-from pina.problem import SpatialProblem
-from pina.operators import laplacian
-from pina import Condition, Span
 
 # ===================================================== #
 #                                                       #
@@ -17,39 +12,46 @@ from pina import Condition, Span
 # ===================================================== #
 
 
+import torch
+from pina.geometry import CartesianDomain
+from pina import Condition
+from pina.problem import SpatialProblem
+from pina.operators import laplacian
+from pina.equation import FixedValue, Equation
+
+
 class Poisson(SpatialProblem):
-
-    # assign output/ spatial variables
     output_variables = ['u']
-    spatial_domain = Span({'x': [0, 1], 'y': [0, 1]})
+    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
 
-    # define the laplace equation
     def laplace_equation(input_, output_):
         force_term = (torch.sin(input_.extract(['x'])*torch.pi) *
-                      torch.sin(input_.extract(['y'])*torch.pi))
+                        torch.sin(input_.extract(['y'])*torch.pi))
-        delta_u = laplacian(output_.extract(['u']), input_)
+        nabla_u = laplacian(output_.extract(['u']), input_)
-        return delta_u - force_term
+        return nabla_u - force_term
 
-    # define nill dirichlet boundary conditions
-    def nil_dirichlet(input_, output_):
-        value = 0.0
-        return output_.extract(['u']) - value
-
-    # problem condition statement
     conditions = {
-        'gamma1': Condition(location=Span({'x': [0, 1], 'y':  1}), function=nil_dirichlet),
+        'gamma1': Condition(
-        'gamma2': Condition(location=Span({'x': [0, 1], 'y': 0}), function=nil_dirichlet),
+            location=CartesianDomain({'x': [0, 1], 'y':  1}),
-        'gamma3': Condition(location=Span({'x':  1, 'y': [0, 1]}),function=nil_dirichlet),
+            equation=FixedValue(0.0)),
-        'gamma4': Condition(location=Span({'x': 0, 'y': [0, 1]}), function=nil_dirichlet),
+        'gamma2': Condition(
-        'D': Condition(location=Span({'x': [0, 1], 'y': [0, 1]}), function=laplace_equation),
+            location=CartesianDomain({'x': [0, 1], 'y': 0}),
+            equation=FixedValue(0.0)),
+        'gamma3': Condition(
+            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+            equation=FixedValue(0.0)),
+        'gamma4': Condition(
+            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+            equation=FixedValue(0.0)),
+        'D': Condition(
+            location=CartesianDomain({'x': [0, 1], 'y': [0, 1]}),
+            equation=Equation(laplace_equation)),
     }
 
-    # real poisson solution
     def poisson_sol(self, pts):
         return -(
             torch.sin(pts.extract(['x'])*torch.pi) *
             torch.sin(pts.extract(['y'])*torch.pi)
         )/(2*torch.pi**2)
-        # return -(np.sin(x*np.pi)*np.sin(y*np.pi))/(2*np.pi**2)
 
     truth_solution = poisson_sol

examples/problems/stokes.py

@@ -1,9 +1,11 @@
-import numpy as np
+""" Navier Stokes Problem """
-import torch
 
+import torch
 from pina.problem import SpatialProblem
 from pina.operators import laplacian, grad, div
-from pina import Condition, Span, LabelTensor
+from pina import Condition, LabelTensor
+from pina.geometry import CartesianDomain
+from pina.equation import SystemEquation, Equation
 
 # ===================================================== #
 #                                                       #
@@ -21,7 +23,7 @@ class Stokes(SpatialProblem):
 
     # assign output/ spatial variables
     output_variables = ['ux', 'uy', 'p']
-    spatial_domain = Span({'x': [-2, 2], 'y': [-1, 1]})
+    spatial_domain = CartesianDomain({'x': [-2, 2], 'y': [-1, 1]})
 
     # define the momentum equation
     def momentum(input_, output_):
@@ -49,17 +51,9 @@ class Stokes(SpatialProblem):
 
     # problem condition statement
     conditions = {
-        'gamma_top': Condition(location=Span({'x': [-2, 2], 'y':  1}), function=wall),
+        'gamma_top': Condition(location=CartesianDomain({'x': [-2, 2], 'y':  1}), equation=Equation(wall)),
-        'gamma_bot': Condition(location=Span({'x': [-2, 2], 'y': -1}), function=wall),
+        'gamma_bot': Condition(location=CartesianDomain({'x': [-2, 2], 'y': -1}), equation=Equation(wall)),
-        'gamma_out': Condition(location=Span({'x':  2, 'y': [-1, 1]}), function=outlet),
+        'gamma_out': Condition(location=CartesianDomain({'x':  2, 'y': [-1, 1]}), equation=Equation(outlet)),
-        'gamma_in':  Condition(location=Span({'x': -2, 'y': [-1, 1]}), function=inlet),
+        'gamma_in':  Condition(location=CartesianDomain({'x': -2, 'y': [-1, 1]}), equation=Equation(inlet)),
-        'D1': Condition(location=Span({'x': [-2, 2], 'y': [-1, 1]}), function=momentum),
+        'D': Condition(location=CartesianDomain({'x': [-2, 2], 'y': [-1, 1]}), equation=SystemEquation([momentum, continuity]))
-        'D2': Condition(location=Span({'x': [-2, 2], 'y': [-1, 1]}), function=continuity),
     }
-    # conditions = {
-    #     'gamma_top': Condition(location=Span({'x': [-2, 2], 'y':  1}), function=wall),
-    #     'gamma_bot': Condition(location=Span({'x': [-2, 2], 'y': -1}), function=wall),
-    #     'gamma_out': Condition(location=Span({'x':  2, 'y': [-1, 1]}), function=outlet),
-    #     'gamma_in':  Condition(location=Span({'x': -2, 'y': [-1, 1]}), function=inlet),
-    #     'D': Condition(location=Span({'x': [-2, 2], 'y': [-1, 1]}), function=[momentum, continuity]),
-    # }

examples/problems/wave.py

@@ -0,0 +1,57 @@
+""" Wave equation Problem """
+
+
+import torch
+from pina.geometry import CartesianDomain
+from pina import Condition
+from pina.problem import SpatialProblem, TimeDependentProblem
+from pina.operators import laplacian, grad
+from pina.equation import FixedValue, Equation
+
+
+# ===================================================== #
+#                                                       #
+#  This script implements the two dimensional           #
+#  Wave equation. The Wave class is defined inheriting  #
+#  from SpatialProblem and TimeDependentProblem. Let    #
+#           u --> field variable                        #
+#           x,y --> spatial variables                   #
+#           t --> temporal variables                    #
+# the velocity coefficient is set to one.               #
+#                                                       #
+# ===================================================== #
+
+
+
+class Wave(TimeDependentProblem, SpatialProblem):
+    output_variables = ['u']
+    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 = laplacian(output_, input_, components=['u'], d=['x', 'y'])
+        return nabla_u - u_tt
+
+    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=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):
+        sqrt_2 = torch.sqrt(torch.tensor(2.))
+        return (torch.sin(torch.pi*pts.extract(['x'])) *
+                torch.sin(torch.pi*pts.extract(['y'])) *
+                torch.cos(sqrt_2*torch.pi*pts.extract(['t'])))
+
+    truth_solution = wave_sol

examples/run_burgers.py

@@ -1,10 +1,14 @@
-"""Run PINA on Burgers equation"""
+""" Run PINA on Burgers equation. """
+
 import argparse
 import torch
 from torch.nn import Softplus
 
-from pina import PINN, Plotter, LabelTensor
+from pina import LabelTensor
 from pina.model import FeedForward
+from pina.solvers import PINN
+from pina.plotter import Plotter
+from pina.trainer import Trainer
 from problems.burgers import Burgers1D
 
 
@@ -13,9 +17,8 @@ class myFeature(torch.nn.Module):
     Feature: sin(pi*x)
     """
 
-    def __init__(self, idx):
+    def __init__(self):
         super(myFeature, self).__init__()
-        self.idx = idx
 
     def forward(self, x):
         return LabelTensor(torch.sin(torch.pi * x.extract(['x'])), ['sin(x)'])
@@ -24,40 +27,47 @@ class myFeature(torch.nn.Module):
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Run PINA")
-    group = parser.add_mutually_exclusive_group(required=True)
+    parser.add_argument("--load", help="directory to save or load file", type=str)
-    group.add_argument("-s", "-save", action="store_true")
+    parser.add_argument("--features", help="extra features", type=int)
-    group.add_argument("-l", "-load", action="store_true")
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
-    parser.add_argument("id_run", help="number of run", type=int)
-    parser.add_argument("features", help="extra features", type=int)
     args = parser.parse_args()
 
-    feat = [myFeature(0)] if args.features else []
+    if args.features is None:
+        args.features = 0
 
+    # extra features
+    feat = [myFeature()] if args.features else []
+
+    # create problem and discretise domain
     burgers_problem = Burgers1D()
+    burgers_problem.discretise_domain(n=200, mode='grid', variables = 't', locations=['D'])
+    burgers_problem.discretise_domain(n=20, mode='grid', variables = 'x', locations=['D'])
+    burgers_problem.discretise_domain(n=150, mode='random', locations=['gamma1', 'gamma2', 't0'])
+
+    # create model
     model = FeedForward(
         layers=[30, 20, 10, 5],
-        output_variables=burgers_problem.output_variables,
+        output_dimensions=len(burgers_problem.output_variables),
-        input_variables=burgers_problem.input_variables,
+        input_dimensions=len(burgers_problem.input_variables) + len(feat),
-        func=Softplus,
+        func=Softplus
-        extra_features=feat,
     )
 
+    # create solver
     pinn = PINN(
-        burgers_problem,
+        problem=burgers_problem,
-        model,
+        model=model,
-        lr=0.006,
+        extra_features=feat,
-        error_norm='mse',
+        optimizer_kwargs={'lr' : 0.006}
-        regularizer=0)
+    )
 
-    if args.s:
+    # create trainer
-        pinn.span_pts(
+    directory = 'pina.burger_extrafeats_{}'.format(bool(args.features))
-            {'n': 200, 'mode': 'grid', 'variables': 't'},
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
-            {'n': 20, 'mode': 'grid', 'variables': 'x'},
+
-            locations=['D'])
+
-        pinn.span_pts(150, 'random', location=['gamma1', 'gamma2', 't0'])
+    if args.load:
-        pinn.train(5000, 100)
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=burgers_problem, model=model)
-        pinn.save_state('pina.burger.{}.{}'.format(args.id_run, args.features))
-    else:
-        pinn.load_state('pina.burger.{}.{}'.format(args.id_run, args.features))
         plotter = Plotter()
         plotter.plot(pinn)
+    else:
+        trainer.train()

examples/run_first_order_ode.py

@@ -1,67 +1,53 @@
+""" Run PINA on ODE equation. """
 import argparse
-
+import torch
 from torch.nn import Softplus
 
 from pina.model import FeedForward
-from pina import Condition, CartesianDomain, Plotter, PINN
+from pina.solvers import PINN
+from pina.plotter import Plotter
+from pina.trainer import Trainer
+from problems.first_order_ode import FirstOrderODE
 
 
-class FirstOrderODE(SpatialProblem):
-
-    x_rng = [0, 5]
-    output_variables = ['y']
-    spatial_domain = CartesianDomain({'x': x_rng})
-
-    def ode(input_, output_):
-        y = output_
-        x = input_
-        return grad(y, x) + y - x
-
-    def fixed(input_, output_):
-        exp_value = 1.
-        return output_ - exp_value
-
-    def solution(self, input_):
-        x = input_
-        return x - 1.0 + 2*torch.exp(-x)
-
-    conditions = {
-        'bc': Condition(CartesianDomain({'x': x_rng[0]}), fixed),
-        'dd': Condition(CartesianDomain({'x': x_rng}), ode),
-    }
-    truth_solution = solution
-
 
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Run PINA")
-    group = parser.add_mutually_exclusive_group(required=True)
+    parser.add_argument("--load", help="directory to save or load file", type=str)
-    group.add_argument("-s", "-save", action="store_true")
+    parser.add_argument("--epochs", help="extra features", type=int, default=3000)
-    group.add_argument("-l", "-load", action="store_true")
-    parser.add_argument("id_run", help="number of run", type=int)
     args = parser.parse_args()
 
-    # define Problem + Model + PINN
+
+    # create problem and discretise domain
     problem = FirstOrderODE()
+    problem.discretise_domain(n=500, mode='grid', variables = 'x', locations=['D'])
+    problem.discretise_domain(n=1, mode='grid', variables = 'x', locations=['BC'])
+
+    # create model
     model = FeedForward(
-        layers=[4]*2,
+        layers=[10, 10],
-        output_variables=problem.output_variables,
+        output_dimensions=len(problem.output_variables),
-        input_variables=problem.input_variables,
+        input_dimensions=len(problem.input_variables),
-        func=Softplus,
+        func=Softplus
     )
-    pinn = PINN(problem, model, lr=0.03, error_norm='mse', regularizer=0)
 
-    if args.s:
+    # create solver
+    pinn = PINN(
+        problem=problem,
+        model=model,
+        extra_features=None,
+        optimizer_kwargs={'lr' : 0.001}
+    )
 
-        pinn.span_pts(
+    # create trainer
-            {'variables': ['x'], 'mode': 'grid', 'n': 1}, locations=['bc'])
+    directory = 'pina.ode'
-        pinn.span_pts(
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
-            {'variables': ['x'], 'mode': 'grid', 'n': 30}, locations=['dd'])
-        Plotter().plot_samples(pinn, ['x'])
-        pinn.train(1200, 50)
-        pinn.save_state('pina.ode')
 
-    else:
+
-        pinn.load_state('pina.ode')
+    if args.load:
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=problem, model=model)
         plotter = Plotter()
-        plotter.plot(pinn, components=['y'])
+        plotter.plot(pinn)
+    else:
+        trainer.train()

examples/run_parametric_elliptic_optimal.py

@@ -0,0 +1,88 @@
+import argparse
+import numpy as np
+import torch
+from torch.nn import Softplus
+
+from pina import LabelTensor
+from pina.solvers import PINN
+from pina.model import MultiFeedForward
+from pina.plotter import Plotter
+from pina.trainer import Trainer
+from problems.parametric_elliptic_optimal_control import (
+    ParametricEllipticOptimalControl)
+
+
+class myFeature(torch.nn.Module):
+    """
+    Feature: sin(x)
+    """
+
+    def __init__(self):
+        super(myFeature, self).__init__()
+
+    def forward(self, x):
+        t = (-x.extract(['x1'])**2+1) * (-x.extract(['x2'])**2+1)
+        return LabelTensor(t, ['k0'])
+
+
+class CustomMultiDFF(MultiFeedForward):
+
+    def __init__(self, dff_dict):
+        super().__init__(dff_dict)
+
+    def forward(self, x):
+        out = self.uu(x)
+        out.labels = ['u', 'y']
+        p = LabelTensor(
+            (out.extract(['u']) * x.extract(['alpha'])), ['p'])
+        return out.append(p)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Run PINA")
+    parser.add_argument("--load", help="directory to save or load file", type=str)
+    parser.add_argument("--features", help="extra features", type=int)
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
+    args = parser.parse_args()
+
+    if args.features is None:
+        args.features = 0
+
+    # extra features
+    feat = [myFeature()] if args.features else []
+    args = parser.parse_args()
+
+    # create problem and discretise domain
+    opc = ParametricEllipticOptimalControl()
+    opc.discretise_domain(n= 100, mode='random', variables=['x1', 'x2'], locations=['D'])
+    opc.discretise_domain(n= 5, mode='random', variables=['mu', 'alpha'], locations=['D'])
+    opc.discretise_domain(n= 20, mode='random', variables=['x1', 'x2'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
+    opc.discretise_domain(n= 5, mode='random', variables=['mu', 'alpha'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
+
+    # create model
+    model = CustomMultiDFF(
+        {
+            'uu': {
+                'input_dimensions': 4 + len(feat),
+                'output_dimensions': 2,
+                'layers': [40, 40, 20],
+                'func': Softplus,
+            },
+        }
+    )
+
+    # create PINN
+    pinn = PINN(problem=opc, model=model, optimizer_kwargs={'lr' : 0.002}, extra_features=feat)
+
+    # create trainer
+    directory = 'pina.parametric_optimal_control_{}'.format(bool(args.features))
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
+
+
+    if args.load:
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=opc, model=model, extra_features=feat)
+        plotter = Plotter()
+        plotter.plot(pinn, fixed_variables={'mu' : 1 , 'alpha' : 0.001}, components='y')
+    else:
+        trainer.train()

examples/run_parametric_elliptic_optimal_control_alpha.py

@@ -1,87 +0,0 @@
-import argparse
-import numpy as np
-import torch
-from torch.nn import Softplus
-
-from pina import PINN, LabelTensor, Plotter
-from pina.model import MultiFeedForward
-from problems.parametric_elliptic_optimal_control_alpha_variable import (
-    ParametricEllipticOptimalControl)
-
-
-class myFeature(torch.nn.Module):
-    """
-    Feature: sin(x)
-    """
-
-    def __init__(self):
-        super(myFeature, self).__init__()
-
-    def forward(self, x):
-        t = (-x.extract(['x1'])**2+1) * (-x.extract(['x2'])**2+1)
-        return LabelTensor(t, ['k0'])
-
-
-class CustomMultiDFF(MultiFeedForward):
-
-    def __init__(self, dff_dict):
-        super().__init__(dff_dict)
-
-    def forward(self, x):
-        out = self.uu(x)
-        p = LabelTensor(
-            (out.extract(['u_param']) * x.extract(['alpha'])), ['p'])
-        return out.append(p)
-
-
-if __name__ == "__main__":
-
-    parser = argparse.ArgumentParser(description="Run PINA")
-    group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument("-s", "-save", action="store_true")
-    group.add_argument("-l", "-load", action="store_true")
-    args = parser.parse_args()
-
-    opc = ParametricEllipticOptimalControl()
-    model = CustomMultiDFF(
-        {
-            'uu': {
-                'input_variables': ['x1', 'x2', 'mu', 'alpha'],
-                'output_variables': ['u_param', 'y'],
-                'layers': [40, 40, 20],
-                'func': Softplus,
-                'extra_features': [myFeature()],
-            },
-        }
-    )
-
-    pinn = PINN(
-        opc,
-        model,
-        lr=0.002,
-        error_norm='mse',
-        regularizer=1e-8)
-
-    if args.s:
-
-        pinn.span_pts(
-            {'variables': ['x1', 'x2'], 'mode': 'random', 'n': 100},
-            {'variables': ['mu', 'alpha'], 'mode': 'grid', 'n': 5},
-            locations=['D'])
-        pinn.span_pts(
-            {'variables': ['x1', 'x2'], 'mode': 'grid', 'n': 20},
-            {'variables': ['mu', 'alpha'], 'mode': 'grid', 'n': 5},
-            locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
-
-        pinn.train(1000, 20)
-        pinn.save_state('pina.ocp')
-
-    else:
-        pinn.load_state('pina.ocp')
-        plotter = Plotter()
-        plotter.plot(pinn, components='y',
-                     fixed_variables={'alpha': 0.01, 'mu': 1.0})
-        plotter.plot(pinn, components='u_param',
-                     fixed_variables={'alpha': 0.01, 'mu': 1.0})
-        plotter.plot(pinn, components='p', fixed_variables={
-                     'alpha': 0.01, 'mu': 1.0})

examples/run_parametric_poisson.py

@@ -1,7 +1,8 @@
 import argparse
 import torch
 from torch.nn import Softplus
-from pina import Plotter, LabelTensor, PINN
+from pina import Plotter, LabelTensor, Trainer
+from pina.solvers import PINN
 from pina.model import FeedForward
 from problems.parametric_poisson import ParametricPoisson
 
@@ -25,41 +26,48 @@ class myFeature(torch.nn.Module):
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Run PINA")
-    group = parser.add_mutually_exclusive_group(required=True)
+    parser.add_argument("--load", help="directory to save or load file", type=str)
-    group.add_argument("-s", "-save", action="store_true")
+    parser.add_argument("--features", help="extra features", type=int)
-    group.add_argument("-l", "-load", action="store_true")
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
-    parser.add_argument("id_run", help="number of run", type=int)
-    parser.add_argument("features", help="extra features", type=int)
     args = parser.parse_args()
 
+    if args.features is None:
+        args.features = 0
+
+    # extra features
     feat = [myFeature()] if args.features else []
 
-    poisson_problem = ParametricPoisson()
+    # create problem and discretise domain
+    ppoisson_problem = ParametricPoisson()
+    ppoisson_problem.discretise_domain(n=100, mode='random', variables = ['x', 'y'], locations=['D'])
+    ppoisson_problem.discretise_domain(n=100, mode='random', variables = ['mu1', 'mu2'], locations=['D'])
+    ppoisson_problem.discretise_domain(n=20, mode='random', variables = ['x', 'y'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
+    ppoisson_problem.discretise_domain(n=5, mode='random', variables = ['mu1', 'mu2'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
+
+    # create model
     model = FeedForward(
         layers=[10, 10, 10],
-        output_variables=poisson_problem.output_variables,
+        output_dimensions=len(ppoisson_problem.output_variables),
-        input_variables=poisson_problem.input_variables,
+        input_dimensions=len(ppoisson_problem.input_variables) + len(feat),
-        func=Softplus,
+        func=Softplus
-        extra_features=feat
     )
 
-    pinn = PINN(poisson_problem, model, lr=0.006, regularizer=1e-6)
+    # create solver
+    pinn = PINN(
+        problem=ppoisson_problem,
+        model=model,
+        extra_features=feat,
+        optimizer_kwargs={'lr' : 0.006}
+    )
 
-    if args.s:
+    # create trainer
+    directory = 'pina.parametric_poisson_extrafeats_{}'.format(bool(args.features))
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
 
-        pinn.span_pts(
-            {'variables': ['x', 'y'], 'mode': 'random', 'n': 100},
-            {'variables': ['mu1', 'mu2'], 'mode': 'grid', 'n': 5},
-            locations=['D'])
-        pinn.span_pts(
-            {'variables': ['x', 'y'], 'mode': 'grid', 'n': 20},
-            {'variables': ['mu1', 'mu2'], 'mode': 'grid', 'n': 5},
-            locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
-        pinn.train(10000, 100)
-        pinn.save_state('pina.poisson_param')
 
-    else:
+    if args.load:
-        pinn.load_state('pina.poisson_param')
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=ppoisson_problem, model=model, extra_features=feat)
         plotter = Plotter()
-        plotter.plot(pinn, fixed_variables={'mu1': 0, 'mu2': 1}, levels=21)
+        plotter.plot(pinn, fixed_variables={'mu1': 1, 'mu2': -1})
-        plotter.plot(pinn, fixed_variables={'mu1': 1, 'mu2': -1}, levels=21)
+    else:
+        trainer.train()

examples/run_poisson.py

@@ -1,12 +1,13 @@
+""" Run PINA on ODE equation. """
 import argparse
-import sys
-import numpy as np
 import torch
-from torch.nn import ReLU, Tanh, Softplus
+from torch.nn import Softplus
 
-from pina import PINN, LabelTensor, Plotter
+from pina import LabelTensor 
 from pina.model import FeedForward
-from pina.adaptive_functions import AdaptiveSin, AdaptiveCos, AdaptiveTanh
+from pina.solvers import PINN
+from pina.plotter import Plotter
+from pina.trainer import Trainer
 from problems.poisson import Poisson
 
 
@@ -26,39 +27,47 @@ class myFeature(torch.nn.Module):
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Run PINA")
-    group = parser.add_mutually_exclusive_group(required=True)
+    parser.add_argument("--load", help="directory to save or load file", type=str)
-    group.add_argument("-s", "-save", action="store_true")
+    parser.add_argument("--features", help="extra features", type=int)
-    group.add_argument("-l", "-load", action="store_true")
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
-    parser.add_argument("id_run", help="number of run", type=int)
-    parser.add_argument("features", help="extra features", type=int)
     args = parser.parse_args()
 
-    feat = [myFeature()] if args.features else []
+    if args.features is None:
+        args.features = 0
 
-    poisson_problem = Poisson()
+    # extra features
+    feat = [myFeature()] if args.features else []
+    args = parser.parse_args()
+
+    # create problem and discretise domain
+    problem = Poisson()
+    problem.discretise_domain(n=20, mode='grid', locations=['D'])
+    problem.discretise_domain(n=100, mode='random', locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
+
+    # create model
     model = FeedForward(
-        layers=[20, 20],
+        layers=[10, 10],
-        output_variables=poisson_problem.output_variables,
+        output_dimensions=len(problem.output_variables),
-        input_variables=poisson_problem.input_variables,
+        input_dimensions=len(problem.input_variables) + len(feat),
-        func=Softplus,
+        func=Softplus
-        extra_features=feat
     )
 
+    # create solver
     pinn = PINN(
-        poisson_problem,
+        problem=problem,
-        model,
+        model=model,
-        lr=0.03,
+        extra_features=feat,
-        error_norm='mse',
+        optimizer_kwargs={'lr' : 0.001}
-        regularizer=1e-8)
+    )
 
-    if args.s:
+    # create trainer
+    directory = 'pina.parametric_poisson_extrafeats_{}'.format(bool(args.features))
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
 
-        pinn.span_pts(20, 'grid', locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
-        pinn.span_pts(20, 'grid', locations=['D'])
-        pinn.train(5000, 100)
-        pinn.save_state('pina.poisson')
 
-    else:
+    if args.load:
-        pinn.load_state('pina.poisson')
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=problem, model=model, extra_features=feat)
         plotter = Plotter()
         plotter.plot(pinn)
+    else:
+        trainer.train()

examples/run_poisson_deeponet.py

@@ -1,120 +1,75 @@
 import argparse
-import logging
-
 import torch
-from problems.poisson import Poisson
+from pina import Plotter, LabelTensor, Trainer
-
+from pina.solvers import PINN
-from pina import PINN, LabelTensor, Plotter
 from pina.model import DeepONet, FeedForward
+from problems.parametric_poisson import ParametricPoisson
 
 
-class SinFeature(torch.nn.Module):
+class myFeature(torch.nn.Module):
     """
-    Feature: sin(x)
     """
-
-    def __init__(self, label):
-        super().__init__()
-
-        if not isinstance(label, (tuple, list)):
-            label = [label]
-        self._label = label
-
-    def forward(self, x):
-        """
-        Defines the computation performed at every call.
-
-        :param LabelTensor x: the input tensor.
-        :return: the output computed by the model.
-        :rtype: LabelTensor
-        """
-        t = torch.sin(x.extract(self._label) * torch.pi)
-        return LabelTensor(t, [f"sin({self._label})"])
-
-
-class myRBF(torch.nn.Module):
-    def __init__(self, input_):
-
-        super().__init__()
-
-        self.input_variables = [input_]
-        self.a = torch.nn.Parameter(torch.tensor([-.3]))
-        # self.b = torch.nn.Parameter(torch.tensor([0.5]))
-        self.b = torch.tensor([0.5])
-        self.c = torch.nn.Parameter(torch.tensor([.5]))
-
-    def forward(self, x):
-        x = x.extract(self.input_variables)
-        result = self.a * torch.exp(-(x - self.b)**2/(self.c**2))
-        return result
-
-
-class myModel(torch.nn.Module):
-    """ Model for the Poisson equation."""
-
     def __init__(self):
-
+        super(myFeature, self).__init__()
-        super().__init__()
-        self.ffn_x = myRBF('x')
-        self.ffn_y = myRBF('y')
 
     def forward(self, x):
-        result = self.ffn_x(x) * self.ffn_y(x)
+        t = (
-        result.labels = ['u']
+            torch.exp(
-        return result
+                - 2*(x.extract(['x']) - x.extract(['mu1']))**2
+                - 2*(x.extract(['y']) - x.extract(['mu2']))**2
+            )
+        )
+        return LabelTensor(t, ['k0'])
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run PINA")
-    parser.add_argument("-s", "--save", action="store_true")
-    parser.add_argument("-l", "--load", action="store_true")
-    parser.add_argument("id_run", help="Run ID", type=int)
 
-    parser.add_argument("--extra", help="Extra features", action="store_true")
+    parser = argparse.ArgumentParser(description="Run PINA")
+    parser.add_argument("--load", help="directory to save or load file", type=str)
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
     args = parser.parse_args()
 
-    problem = Poisson()
 
-    # ffn_x = FeedForward(
+    # create problem and discretise domain
-    #     input_variables=['x'], layers=[], output_variables=1,
+    ppoisson_problem = ParametricPoisson()
-    #     func=torch.nn.Softplus,
+    ppoisson_problem.discretise_domain(n=100, mode='random', variables = ['x', 'y'], locations=['D'])
-    #     extra_features=[SinFeature('x')]
+    ppoisson_problem.discretise_domain(n=100, mode='random', variables = ['mu1', 'mu2'], locations=['D'])
-    # )
+    ppoisson_problem.discretise_domain(n=20, mode='random', variables = ['x', 'y'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
-    # ffn_y = FeedForward
+    ppoisson_problem.discretise_domain(n=5, mode='random', variables = ['mu1', 'mu2'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
-    #     input_variables=['y'], layers=[], output_variables=1,
-    #     func=torch.nn.Softplus,
-    #     extra_features=[SinFeature('y')]
-    # )
-    model = myModel()
-    test = torch.tensor([[0.0, 0.5]])
-    test.labels = ['x', 'y']
-    pinn = PINN(problem, model, lr=0.0001)
 
-    if args.save:
+    # create model
-        pinn.span_pts(
+    trunck = FeedForward(
-            20, "grid", locations=["gamma1", "gamma2", "gamma3", "gamma4"]
+        layers=[40, 40],
-        )
+        output_dimensions=1,
-        pinn.span_pts(20, "grid", locations=["D"])
+        input_dimensions=2,
-        while True:
+        func=torch.nn.ReLU
-            pinn.train(500, 50)
+    )
-            print(model.ffn_x.a)
+    branch = FeedForward(
-            print(model.ffn_x.b)
+        layers=[40, 40],
-            print(model.ffn_x.c)
+        output_dimensions=1,
+        input_dimensions=2,
+        func=torch.nn.ReLU
+    )
+    model = DeepONet(branch_net=branch,
+                     trunk_net=trunck,
+                     input_indeces_branch_net=['x', 'y'],
+                     input_indeces_trunk_net=['mu1', 'mu2'])
 
-            xi = torch.linspace(0, 1, 64).reshape(-1,
+    # create solver
-                                                  1).as_subclass(LabelTensor)
+    pinn = PINN(
-            xi.labels = ['x']
+        problem=ppoisson_problem,
-            yi = model.ffn_x(xi)
+        model=model,
-            y_truth = -torch.sin(xi*torch.pi)
+        optimizer_kwargs={'lr' : 0.006}
+    )
+
+    # create trainer
+    directory = 'pina.parametric_poisson_deeponet'
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
 
-            import matplotlib.pyplot as plt
-            plt.plot(xi.detach().flatten(), yi.detach().flatten(), 'r-')
-            plt.plot(xi.detach().flatten(), y_truth.detach().flatten(), 'b-')
-            plt.plot(xi.detach().flatten(), -y_truth.detach().flatten(), 'b-')
-            plt.show()
-        pinn.save_state(f"pina.poisson_{args.id_run}")
 
     if args.load:
-        pinn.load_state(f"pina.poisson_{args.id_run}")
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=ppoisson_problem, model=model)
         plotter = Plotter()
-        plotter.plot(pinn)
+        plotter.plot(pinn, fixed_variables={'mu1': 1, 'mu2': -1})
+    else:
+        trainer.train()

examples/run_stokes.py

@@ -1,55 +1,52 @@
 import argparse
-import numpy as np
+from torch.nn import Softplus
-import torch
-from torch.nn import ReLU, Tanh, Softplus
 
-from pina import PINN, Plotter
+from pina import Plotter, Trainer
 from pina.model import FeedForward
-from pina.adaptive_functions import AdaptiveSin, AdaptiveCos, AdaptiveTanh
+from pina.solvers import PINN
 from problems.stokes import Stokes
 
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Run PINA")
-    group = parser.add_mutually_exclusive_group(required=True)
+    parser = argparse.ArgumentParser(description="Run PINA")
-    group.add_argument("-s", "-save", action="store_true")
+    parser.add_argument("--load", help="directory to save or load file", type=str)
-    group.add_argument("-l", "-load", action="store_true")
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
-    parser.add_argument("id_run", help="number of run", type=int)
     args = parser.parse_args()
 
 
+    # create problem and discretise domain
     stokes_problem = Stokes()
+    stokes_problem.discretise_domain(n=1000, locations=['gamma_top', 'gamma_bot', 'gamma_in', 'gamma_out'])
+    stokes_problem.discretise_domain(n=2000, locations=['D'])
+
+    # make the model
     model = FeedForward(
         layers=[10, 10, 10, 10],
-        output_variables=stokes_problem.output_variables,
+        output_dimensions=len(stokes_problem.output_variables),
-        input_variables=stokes_problem.input_variables,
+        input_dimensions=len(stokes_problem.input_variables),
         func=Softplus,
     )
 
+    # make the pinn
     pinn = PINN(
         stokes_problem,
         model,
-        lr=0.006,
+        optimizer_kwargs={'lr' : 0.001}
-        error_norm='mse',
+        )
-        regularizer=1e-8)
 
-    if args.s:
+    # create trainer
+    directory = 'pina.navier_stokes'
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
 
-        pinn.span_pts(200, 'grid', locations=['gamma_top', 'gamma_bot', 'gamma_in', 'gamma_out'])
-        # pinn.span_pts(2000, 'random', locations=['D'])
-        pinn.span_pts(2000, 'random', locations=['D1'])
-        pinn.span_pts(2000, 'random', locations=['D2'])
-        pinn.train(10000, 100)
-        with open('stokes_history_{}.txt'.format(args.id_run), 'w') as file_:
-            for i, losses in pinn.history_loss.items():
-                file_.write('{} {}\n'.format(i, sum(losses)))
-        pinn.save_state('pina.stokes')
 
-    else:
+    if args.load:
-        pinn.load_state('pina.stokes')
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=stokes_problem, model=model)
         plotter = Plotter()
         plotter.plot(pinn, components='ux')
         plotter.plot(pinn, components='uy')
         plotter.plot(pinn, components='p')
+    else:
+        trainer.train()
 
 

examples/run_wave.py

@@ -0,0 +1,64 @@
+""" Run PINA on Burgers equation. """
+
+import argparse
+import torch
+from torch.nn import Softplus
+
+from pina import LabelTensor
+from pina.model import FeedForward
+from pina.solvers import PINN
+from pina.plotter import Plotter
+from pina.trainer import Trainer
+from problems.wave import Wave
+
+class HardMLP(torch.nn.Module):
+
+    def __init__(self, **kwargs):
+        super().__init__()
+
+        self.layers = FeedForward(**kwargs)
+        
+    # here in the foward we implement the hard constraints
+    def forward(self, x):
+        hard_space = x.extract(['x'])*(1-x.extract(['x']))*x.extract(['y'])*(1-x.extract(['y']))
+        hard_t = torch.sin(torch.pi*x.extract(['x'])) * torch.sin(torch.pi*x.extract(['y'])) * torch.cos(torch.sqrt(torch.tensor(2.))*torch.pi*x.extract(['t']))
+        return hard_space * self.layers(x) * x.extract(['t']) + hard_t
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Run PINA")
+    parser.add_argument("--load", help="directory to save or load file", type=str)
+    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
+    args = parser.parse_args()
+
+
+    # create problem and discretise domain
+    wave_problem = Wave()
+    wave_problem.discretise_domain(1000, 'random', locations=['D', 't0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])
+
+    # create model
+    model = HardMLP(
+        layers=[40, 40, 40],
+        output_dimensions=len(wave_problem.output_variables),
+        input_dimensions=len(wave_problem.input_variables),
+        func=Softplus
+    )
+
+    # create solver
+    pinn = PINN(
+        problem=wave_problem,
+        model=model,
+        optimizer_kwargs={'lr' : 0.006}
+    )
+
+    # create trainer
+    directory = 'pina.wave'
+    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
+
+
+    if args.load:
+        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=wave_problem, model=model)
+        plotter = Plotter()
+        plotter.plot(pinn)
+    else:
+        trainer.train()

pina/solvers/solver.py

@@ -69,7 +69,7 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
                              f' {len_optimizer_kwargs} dicitionaries')
 
         # extra features handling
-        if extra_features is None or (len(extra_features)==0):
+        if (extra_features is None) or (len(extra_features)==0):
             extra_features = [None] * len_model
         else:
             # if we only have a list of extra features