fix bug + add training tests

Author: GiovanniCanali

pina/equation/equation_factory.py

@@ -246,12 +246,12 @@ def equation(input_, output_):
                 )
 
             # Repeat c to ensure consistent shape for advection
-            self.c = self.c.repeat(output_.shape[0], 1)
-            if self.c.shape[1] != (len(input_lbl) - 1):
-                self.c = self.c.repeat(1, len(input_lbl) - 1)
+            c = self.c.repeat(output_.shape[0], 1).to(input_.device)
+            if c.shape[1] != (len(input_lbl) - 1):
+                c = c.repeat(1, len(input_lbl) - 1)
 
             # Add a dimension to c for the following operations
-            self.c = self.c.unsqueeze(-1)
+            c = c.unsqueeze(-1)
 
             # Compute the time derivative and the spatial gradient
             time_der = grad(output_, input_, components=None, d="t")
@@ -262,7 +262,7 @@ def equation(input_, output_):
             tmp = tmp.transpose(-1, -2)
 
             # Compute advection term
-            adv = (tmp * self.c).sum(dim=tmp.tensor.ndim - 2)
+            adv = (tmp * c).sum(dim=tmp.tensor.ndim - 2)
 
             return time_der + adv
 

pina/problem/zoo/helmholtz.py

@@ -48,11 +48,10 @@ def __init__(self, alpha=3.0):
         :type alpha: float | int
         """
         super().__init__()
-
-        self.alpha = alpha
         check_consistency(alpha, (int, float))
+        self.alpha = alpha
 
-        def forcing_term(self, input_):
+        def forcing_term(input_):
             """
             Implementation of the forcing term.
             """

tests/test_problem_zoo/test_advection.py

@@ -1,6 +1,9 @@
 import pytest
-from pina.problem.zoo import AdvectionProblem
 from pina.problem import SpatialProblem, TimeDependentProblem
+from pina.problem.zoo import AdvectionProblem
+from pina.model import FeedForward
+from pina.solver import PINN
+from pina import Trainer
 
 
 @pytest.mark.parametrize("c", [1.5, 3])
@@ -17,3 +20,14 @@ def test_constructor(c):
     # Should fail if c is not a float or int
     with pytest.raises(ValueError):
         AdvectionProblem(c="invalid")
+
+
+@pytest.mark.parametrize("c", [1.5, 3])
+def test_train(c):
+
+    problem = AdvectionProblem(c=c)
+    problem.discretise_domain(n=10)
+    model = FeedForward(2, 1, 10, 2)
+    solver = PINN(problem, model)
+    trainer = Trainer(solver, max_epochs=2)
+    trainer.train()

tests/test_problem_zoo/test_allen_cahn.py

@@ -1,6 +1,10 @@
 import pytest
-from pina.problem.zoo import AllenCahnProblem
 from pina.problem import SpatialProblem, TimeDependentProblem
+from pina.problem.zoo import AllenCahnProblem
+from pina.problem import SpatialProblem
+from pina.model import FeedForward
+from pina.solver import PINN
+from pina import Trainer
 
 
 @pytest.mark.parametrize("alpha", [0.1, 1])
@@ -22,3 +26,15 @@ def test_constructor(alpha, beta):
     # Should fail if beta is not a float or int
     with pytest.raises(ValueError):
         AllenCahnProblem(alpha=alpha, beta="invalid")
+
+
+@pytest.mark.parametrize("alpha", [0.1, 1])
+@pytest.mark.parametrize("beta", [0.1, 1])
+def test_train(alpha, beta):
+
+    problem = AllenCahnProblem(alpha=alpha, beta=beta)
+    problem.discretise_domain(n=10)
+    model = FeedForward(2, 1, 10, 2)
+    solver = PINN(problem, model)
+    trainer = Trainer(solver, max_epochs=2)
+    trainer.train()

tests/test_problem_zoo/test_diffusion_reaction.py

@@ -1,6 +1,10 @@
 import pytest
-from pina.problem.zoo import DiffusionReactionProblem
 from pina.problem import TimeDependentProblem, SpatialProblem
+from pina.problem.zoo import DiffusionReactionProblem
+from pina.problem import SpatialProblem
+from pina.model import FeedForward
+from pina.solver import PINN
+from pina import Trainer
 
 
 @pytest.mark.parametrize("alpha", [0.1, 1])
@@ -17,3 +21,14 @@ def test_constructor(alpha):
     # Should fail if alpha is not a float or int
     with pytest.raises(ValueError):
         problem = DiffusionReactionProblem(alpha="invalid")
+
+
+@pytest.mark.parametrize("alpha", [0.1, 1])
+def test_train(alpha):
+
+    problem = DiffusionReactionProblem(alpha=alpha)
+    problem.discretise_domain(n=10)
+    model = FeedForward(2, 1, 10, 2)
+    solver = PINN(problem, model)
+    trainer = Trainer(solver, max_epochs=2)
+    trainer.train()

tests/test_problem_zoo/test_helmholtz.py

@@ -1,6 +1,9 @@
 import pytest
 from pina.problem.zoo import HelmholtzProblem
 from pina.problem import SpatialProblem
+from pina.model import FeedForward
+from pina.solver import PINN
+from pina import Trainer
 
 
 @pytest.mark.parametrize("alpha", [1.5, 3])
@@ -15,3 +18,14 @@ def test_constructor(alpha):
 
     with pytest.raises(ValueError):
         HelmholtzProblem(alpha="invalid")
+
+
+@pytest.mark.parametrize("alpha", [1.5, 3])
+def test_train(alpha):
+
+    problem = HelmholtzProblem(alpha=alpha)
+    problem.discretise_domain(n=10)
+    model = FeedForward(2, 1, 10, 2)
+    solver = PINN(problem, model)
+    trainer = Trainer(solver, max_epochs=2, accelerator="cpu")
+    trainer.train()

tests/test_problem_zoo/test_inverse_poisson_2d_square.py

@@ -1,6 +1,9 @@
 import pytest
 from pina.problem.zoo import InversePoisson2DSquareProblem
 from pina.problem import InverseProblem, SpatialProblem
+from pina.model import FeedForward
+from pina.solver import PINN
+from pina import Trainer
 
 
 @pytest.mark.parametrize("load", [True, False])
@@ -23,3 +26,13 @@ def test_constructor(load, data_size):
     # Should fail if data_size is not in the range [0.0, 1.0]
     with pytest.raises(ValueError):
         problem = InversePoisson2DSquareProblem(load=load, data_size=3.0)
+
+
+def test_train():
+
+    problem = InversePoisson2DSquareProblem()
+    problem.discretise_domain(n=10)
+    model = FeedForward(2, 1, 10, 2)
+    solver = PINN(problem, model)
+    trainer = Trainer(solver, max_epochs=2, accelerator="gpu")
+    trainer.train()

tests/test_problem_zoo/test_poisson_2d_square.py

@@ -1,5 +1,8 @@
 from pina.problem.zoo import Poisson2DSquareProblem
 from pina.problem import SpatialProblem
+from pina.model import FeedForward
+from pina.solver import PINN
+from pina import Trainer
 
 
 def test_constructor():
@@ -10,3 +13,13 @@ def test_constructor():
     assert isinstance(problem, SpatialProblem)
     assert hasattr(problem, "conditions")
     assert isinstance(problem.conditions, dict)
+
+
+def test_train():
+
+    problem = Poisson2DSquareProblem()
+    problem.discretise_domain(n=10)
+    model = FeedForward(2, 1, 10, 2)
+    solver = PINN(problem, model)
+    trainer = Trainer(solver, max_epochs=2)
+    trainer.train()

tests/test_problem_zoo/test_supervised_problem.py

@@ -1,8 +1,11 @@
 import torch
-from pina.problem import AbstractProblem
-from pina.condition import InputTargetCondition
 from pina.problem.zoo.supervised_problem import SupervisedProblem
+from pina.condition import InputTargetCondition
+from pina.problem import AbstractProblem
+from pina.solver import SupervisedSolver
 from pina.graph import RadiusGraph
+from pina.model import FeedForward
+from pina import Trainer
 
 
 def test_constructor():
@@ -32,3 +35,14 @@ def test_constructor_graph():
     assert isinstance(problem.conditions["data"], InputTargetCondition)
     assert isinstance(problem.conditions["data"].input, list)
     assert isinstance(problem.conditions["data"].target, torch.Tensor)
+
+
+def test_train():
+
+    input_ = torch.rand((100, 10))
+    output_ = torch.rand((100, 10))
+    problem = SupervisedProblem(input_=input_, output_=output_)
+    model = FeedForward(10, 10, 10, 2)
+    solver = SupervisedSolver(problem, model, use_lt=False)
+    trainer = Trainer(solver, max_epochs=2, accelerator="cpu")
+    trainer.train()