From ee39b39805ee6517a375918c573160d20747ff6c Mon Sep 17 00:00:00 2001 From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com> Date: Tue, 14 Nov 2023 18:24:07 +0100 Subject: [PATCH] Examples update for v0.1 (#206) * modify examples/problems * modify tutorials --------- Co-authored-by: Dario Coscia Co-authored-by: Dario Coscia --- examples/problems/burgers.py | 40 ++--- examples/problems/elliptic_optimal_control.py | 47 ------ examples/problems/first_order_ode.py | 32 ++-- .../parametric_elliptic_optimal_control.py | 120 ++++++++------ ...elliptic_optimal_control_alpha_variable.py | 78 --------- examples/problems/parametric_poisson.py | 44 +++--- examples/problems/poisson.py | 54 ++++--- examples/problems/stokes.py | 28 ++-- examples/problems/wave.py | 57 +++++++ examples/run_burgers.py | 68 ++++---- examples/run_first_order_ode.py | 82 ++++------ examples/run_parametric_elliptic_optimal.py | 88 +++++++++++ ...rametric_elliptic_optimal_control_alpha.py | 87 ---------- examples/run_parametric_poisson.py | 62 ++++---- examples/run_poisson.py | 67 ++++---- examples/run_poisson_deeponet.py | 149 ++++++------------ examples/run_stokes.py | 49 +++--- examples/run_wave.py | 64 ++++++++ pina/solvers/solver.py | 2 +- 19 files changed, 605 insertions(+), 613 deletions(-) delete mode 100644 examples/problems/elliptic_optimal_control.py delete mode 100644 examples/problems/parametric_elliptic_optimal_control_alpha_variable.py create mode 100644 examples/problems/wave.py create mode 100644 examples/run_parametric_elliptic_optimal.py delete mode 100644 examples/run_parametric_elliptic_optimal_control_alpha.py create mode 100644 examples/run_wave.py diff --git a/examples/problems/burgers.py b/examples/problems/burgers.py index d06fd4e..5da9ccb 100644 --- a/examples/problems/burgers.py +++ b/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 - output_variables = ['u'] - spatial_domain = Span({'x': [-1, 1]}) - temporal_domain = Span({'t': [0, 1]}) +import torch +from pina.geometry import CartesianDomain +from pina import Condition +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), - 'gamma2': Condition(location=Span({'x': 1, 't': [0, 1]}), function=nil_dirichlet), - 't0': Condition(location=Span({'x': [-1, 1], 't': 0}), function=initial_condition), - 'D': Condition(location=Span({'x': [-1, 1], 't': [0, 1]}), function=burger_equation), - } + 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)), + 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)), + 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)), + 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Equation(burger_equation)), + } \ No newline at end of file diff --git a/examples/problems/elliptic_optimal_control.py b/examples/problems/elliptic_optimal_control.py deleted file mode 100644 index 034b09d..0000000 --- a/examples/problems/elliptic_optimal_control.py +++ /dev/null @@ -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...') \ No newline at end of file diff --git a/examples/problems/first_order_ode.py b/examples/problems/first_order_ode.py index eb83941..802b85b 100644 --- a/examples/problems/first_order_ode.py +++ b/examples/problems/first_order_ode.py @@ -1,7 +1,5 @@ -from pina.problem import SpatialProblem -from pina import Condition, Span -from pina.operators import grad -import torch +""" Simple ODE problem. """ + # ===================================================== # # # @@ -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), - 'dd': Condition(location=Span({'x': x_rng}), function=ode), + 'BC': Condition(location=CartesianDomain({'x': x_rng[0]}), equation=FixedValue(1.)), + 'D': Condition(location=CartesianDomain({'x': x_rng}), equation=Equation(ode)), } + truth_solution = solution \ No newline at end of file diff --git a/examples/problems/parametric_elliptic_optimal_control.py b/examples/problems/parametric_elliptic_optimal_control.py index f244d93..3fd26c4 100644 --- a/examples/problems/parametric_elliptic_optimal_control.py +++ b/examples/problems/parametric_elliptic_optimal_control.py @@ -1,52 +1,80 @@ -import numpy as np -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 +""" Poisson OCP problem. """ - def term(input_, param_, output_): - return term1( input_, param_, output_) +term2( input_, param_, output_) + term3( input_, param_, output_) +from pina import Condition +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 - return torch.abs(output_['y'] - y_value) + torch.abs(output_['p'] - p_value) +# ===================================================== # +# # +# 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 # +# ===================================================== # - 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'] - 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]]) +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 = 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])), + } \ No newline at end of file diff --git a/examples/problems/parametric_elliptic_optimal_control_alpha_variable.py b/examples/problems/parametric_elliptic_optimal_control_alpha_variable.py deleted file mode 100644 index 0dabdb4..0000000 --- a/examples/problems/parametric_elliptic_optimal_control_alpha_variable.py +++ /dev/null @@ -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]), - } \ No newline at end of file diff --git a/examples/problems/parametric_poisson.py b/examples/problems/parametric_poisson.py index 82f2113..c6581cb 100644 --- a/examples/problems/parametric_poisson.py +++ b/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]}) - parameter_domain = Span({'mu1': [-1, 1], 'mu2': [-1, 1]}) + spatial_domain = CartesianDomain({'x': [-1, 1], 'y': [-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]}), - function=nil_dirichlet), + location=CartesianDomain({'x': [-1, 1], 'y': 1, 'mu1': [-1, 1], 'mu2': [-1, 1]}), + equation=FixedValue(0.)), 'gamma2': Condition( - location=Span({'x': [-1, 1], 'y': -1, 'mu1': [-1, 1], 'mu2': [-1, 1]}), - function=nil_dirichlet), + location=CartesianDomain({'x': [-1, 1], 'y': -1, 'mu1': [-1, 1], 'mu2': [-1, 1]}), + equation=FixedValue(0.)), 'gamma3': Condition( - location=Span({'x': 1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}), - function=nil_dirichlet), + location=CartesianDomain({'x': 1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}), + equation=FixedValue(0.)), 'gamma4': Condition( - location=Span({'x': -1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}), - function=nil_dirichlet), + location=CartesianDomain({'x': -1, 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}), + equation=FixedValue(0.)), 'D': Condition( - location=Span({'x': [-1, 1], 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}), - function=laplace_equation), - } + location=CartesianDomain({'x': [-1, 1], 'y': [-1, 1], 'mu1': [-1, 1], 'mu2': [-1, 1]}), + equation=Equation(laplace_equation)), + } \ No newline at end of file diff --git a/examples/problems/poisson.py b/examples/problems/poisson.py index c451f53..c817787 100644 --- a/examples/problems/poisson.py +++ b/examples/problems/poisson.py @@ -1,10 +1,5 @@ -""" Poisson equation example. """ -import numpy as np -import torch +""" Poisson problem. """ -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)) - delta_u = laplacian(output_.extract(['u']), input_) - return delta_u - force_term + torch.sin(input_.extract(['y'])*torch.pi)) + nabla_u = laplacian(output_.extract(['u']), input_) + 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), - 'gamma2': Condition(location=Span({'x': [0, 1], 'y': 0}), function=nil_dirichlet), - 'gamma3': Condition(location=Span({'x': 1, 'y': [0, 1]}),function=nil_dirichlet), - 'gamma4': Condition(location=Span({'x': 0, 'y': [0, 1]}), function=nil_dirichlet), - 'D': Condition(location=Span({'x': [0, 1], 'y': [0, 1]}), function=laplace_equation), + 'gamma1': Condition( + location=CartesianDomain({'x': [0, 1], 'y': 1}), + equation=FixedValue(0.0)), + 'gamma2': Condition( + 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 diff --git a/examples/problems/stokes.py b/examples/problems/stokes.py index 3a60498..dbf4015 100644 --- a/examples/problems/stokes.py +++ b/examples/problems/stokes.py @@ -1,9 +1,11 @@ -import numpy as np -import torch +""" Navier Stokes Problem """ +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_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), - 'D1': Condition(location=Span({'x': [-2, 2], 'y': [-1, 1]}), function=momentum), - 'D2': Condition(location=Span({'x': [-2, 2], 'y': [-1, 1]}), function=continuity), + 'gamma_top': Condition(location=CartesianDomain({'x': [-2, 2], 'y': 1}), equation=Equation(wall)), + 'gamma_bot': Condition(location=CartesianDomain({'x': [-2, 2], 'y': -1}), equation=Equation(wall)), + 'gamma_out': Condition(location=CartesianDomain({'x': 2, 'y': [-1, 1]}), equation=Equation(outlet)), + 'gamma_in': Condition(location=CartesianDomain({'x': -2, 'y': [-1, 1]}), equation=Equation(inlet)), + 'D': Condition(location=CartesianDomain({'x': [-2, 2], 'y': [-1, 1]}), equation=SystemEquation([momentum, 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]), - # } diff --git a/examples/problems/wave.py b/examples/problems/wave.py new file mode 100644 index 0000000..cce94da --- /dev/null +++ b/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 \ No newline at end of file diff --git a/examples/run_burgers.py b/examples/run_burgers.py index 8aade04..10f217a 100644 --- a/examples/run_burgers.py +++ b/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) - group.add_argument("-s", "-save", action="store_true") - group.add_argument("-l", "-load", action="store_true") - parser.add_argument("id_run", help="number of run", type=int) - parser.add_argument("features", help="extra features", type=int) + 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() - 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, - input_variables=burgers_problem.input_variables, - func=Softplus, - extra_features=feat, + output_dimensions=len(burgers_problem.output_variables), + input_dimensions=len(burgers_problem.input_variables) + len(feat), + func=Softplus ) + # create solver pinn = PINN( - burgers_problem, - model, - lr=0.006, - error_norm='mse', - regularizer=0) + problem=burgers_problem, + model=model, + extra_features=feat, + optimizer_kwargs={'lr' : 0.006} + ) - if args.s: - pinn.span_pts( - {'n': 200, 'mode': 'grid', 'variables': 't'}, - {'n': 20, 'mode': 'grid', 'variables': 'x'}, - locations=['D']) - pinn.span_pts(150, 'random', location=['gamma1', 'gamma2', 't0']) - pinn.train(5000, 100) - pinn.save_state('pina.burger.{}.{}'.format(args.id_run, args.features)) - else: - pinn.load_state('pina.burger.{}.{}'.format(args.id_run, args.features)) + # create trainer + directory = 'pina.burger_extrafeats_{}'.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=burgers_problem, model=model) plotter = Plotter() plotter.plot(pinn) + else: + trainer.train() diff --git a/examples/run_first_order_ode.py b/examples/run_first_order_ode.py index 57ccfce..b41b470 100644 --- a/examples/run_first_order_ode.py +++ b/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) - group.add_argument("-s", "-save", action="store_true") - group.add_argument("-l", "-load", action="store_true") - parser.add_argument("id_run", help="number of run", type=int) + parser.add_argument("--load", help="directory to save or load file", type=str) + parser.add_argument("--epochs", help="extra features", type=int, default=3000) 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, - output_variables=problem.output_variables, - input_variables=problem.input_variables, - func=Softplus, + layers=[10, 10], + output_dimensions=len(problem.output_variables), + input_dimensions=len(problem.input_variables), + 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( - {'variables': ['x'], 'mode': 'grid', 'n': 1}, locations=['bc']) - pinn.span_pts( - {'variables': ['x'], 'mode': 'grid', 'n': 30}, locations=['dd']) - Plotter().plot_samples(pinn, ['x']) - pinn.train(1200, 50) - pinn.save_state('pina.ode') + # create trainer + directory = 'pina.ode' + trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory) - 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() \ No newline at end of file diff --git a/examples/run_parametric_elliptic_optimal.py b/examples/run_parametric_elliptic_optimal.py new file mode 100644 index 0000000..cf05f7e --- /dev/null +++ b/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() diff --git a/examples/run_parametric_elliptic_optimal_control_alpha.py b/examples/run_parametric_elliptic_optimal_control_alpha.py deleted file mode 100644 index 29964fa..0000000 --- a/examples/run_parametric_elliptic_optimal_control_alpha.py +++ /dev/null @@ -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}) diff --git a/examples/run_parametric_poisson.py b/examples/run_parametric_poisson.py index 0bd5445..1c71366 100644 --- a/examples/run_parametric_poisson.py +++ b/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) - group.add_argument("-s", "-save", action="store_true") - group.add_argument("-l", "-load", action="store_true") - parser.add_argument("id_run", help="number of run", type=int) - parser.add_argument("features", help="extra features", type=int) + 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 [] - 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, - input_variables=poisson_problem.input_variables, - func=Softplus, - extra_features=feat + output_dimensions=len(ppoisson_problem.output_variables), + input_dimensions=len(ppoisson_problem.input_variables) + len(feat), + func=Softplus ) - 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: - pinn.load_state('pina.poisson_param') + if args.load: + 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}, levels=21) + plotter.plot(pinn, fixed_variables={'mu1': 1, 'mu2': -1}) + else: + trainer.train() diff --git a/examples/run_poisson.py b/examples/run_poisson.py index 4b2aad3..390e042 100644 --- a/examples/run_poisson.py +++ b/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) - group.add_argument("-s", "-save", action="store_true") - group.add_argument("-l", "-load", action="store_true") - parser.add_argument("id_run", help="number of run", type=int) - parser.add_argument("features", help="extra features", type=int) + 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() - 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], - output_variables=poisson_problem.output_variables, - input_variables=poisson_problem.input_variables, - func=Softplus, - extra_features=feat + layers=[10, 10], + output_dimensions=len(problem.output_variables), + input_dimensions=len(problem.input_variables) + len(feat), + func=Softplus ) + # create solver pinn = PINN( - poisson_problem, - model, - lr=0.03, - error_norm='mse', - regularizer=1e-8) + problem=problem, + model=model, + extra_features=feat, + optimizer_kwargs={'lr' : 0.001} + ) - 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: - pinn.load_state('pina.poisson') + if args.load: + pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=problem, model=model, extra_features=feat) plotter = Plotter() plotter.plot(pinn) + else: + trainer.train() \ No newline at end of file diff --git a/examples/run_poisson_deeponet.py b/examples/run_poisson_deeponet.py index 2bd566a..3e577a6 100644 --- a/examples/run_poisson_deeponet.py +++ b/examples/run_poisson_deeponet.py @@ -1,120 +1,75 @@ import argparse -import logging - import torch -from problems.poisson import Poisson - -from pina import PINN, LabelTensor, Plotter +from pina import Plotter, LabelTensor, Trainer +from pina.solvers import PINN 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().__init__() - self.ffn_x = myRBF('x') - self.ffn_y = myRBF('y') + super(myFeature, self).__init__() def forward(self, x): - result = self.ffn_x(x) * self.ffn_y(x) - result.labels = ['u'] - return result + t = ( + torch.exp( + - 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( - # input_variables=['x'], layers=[], output_variables=1, - # func=torch.nn.Softplus, - # extra_features=[SinFeature('x')] - # ) - # ffn_y = FeedForward - # 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) + # 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']) - if args.save: - pinn.span_pts( - 20, "grid", locations=["gamma1", "gamma2", "gamma3", "gamma4"] - ) - pinn.span_pts(20, "grid", locations=["D"]) - while True: - pinn.train(500, 50) - print(model.ffn_x.a) - print(model.ffn_x.b) - print(model.ffn_x.c) + # create model + trunck = FeedForward( + layers=[40, 40], + output_dimensions=1, + input_dimensions=2, + func=torch.nn.ReLU + ) + branch = FeedForward( + layers=[40, 40], + 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, - 1).as_subclass(LabelTensor) - xi.labels = ['x'] - yi = model.ffn_x(xi) - y_truth = -torch.sin(xi*torch.pi) + # create solver + pinn = PINN( + problem=ppoisson_problem, + model=model, + 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() diff --git a/examples/run_stokes.py b/examples/run_stokes.py index c0711a9..54b2aec 100644 --- a/examples/run_stokes.py +++ b/examples/run_stokes.py @@ -1,55 +1,52 @@ import argparse -import numpy as np -import torch -from torch.nn import ReLU, Tanh, Softplus +from torch.nn import 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) - group.add_argument("-s", "-save", action="store_true") - group.add_argument("-l", "-load", action="store_true") - parser.add_argument("id_run", help="number of run", type=int) + 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 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, - input_variables=stokes_problem.input_variables, + output_dimensions=len(stokes_problem.output_variables), + input_dimensions=len(stokes_problem.input_variables), func=Softplus, ) + # make the pinn pinn = PINN( stokes_problem, model, - lr=0.006, - error_norm='mse', - regularizer=1e-8) + optimizer_kwargs={'lr' : 0.001} + ) - 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: - pinn.load_state('pina.stokes') + if args.load: + 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() diff --git a/examples/run_wave.py b/examples/run_wave.py new file mode 100644 index 0000000..2d4b4e6 --- /dev/null +++ b/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() diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index c0fade5..93f1bbf 100644 --- a/pina/solvers/solver.py +++ b/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