fix issues + speedup test

Author: dario-coscia

pina/solver/__init__.py

@@ -11,6 +11,7 @@
     "CompetitivePINN",
     "SelfAdaptivePINN",
     "RBAPINN",
+    "SupervisedSolverInterface",
     "SupervisedSolver",
     "ReducedOrderModelSolver",
     "DeepEnsembleSolverInterface",
@@ -20,7 +21,23 @@
 ]
 
 from .solver import SolverInterface, SingleSolverInterface, MultiSolverInterface
-from .physics_informed_solver import *
-from .supervised_solver import *
-from .ensemble_solver import *
+from .physics_informed_solver import (
+    PINNInterface,
+    PINN,
+    GradientPINN,
+    CausalPINN,
+    CompetitivePINN,
+    SelfAdaptivePINN,
+    RBAPINN,
+)
+from .supervised_solver import (
+    SupervisedSolverInterface,
+    SupervisedSolver,
+    ReducedOrderModelSolver,
+)
+from .ensemble_solver import (
+    DeepEnsembleSolverInterface,
+    DeepEnsembleSupervisedSolver,
+    DeepEnsemblePINN,
+)
 from .garom import GAROM

pina/solver/ensemble_solver/ensemble_pinn.py

@@ -98,10 +98,11 @@ def __init__(
             If ``None``, no weighting schema is used. Default is ``None``.
         :param int ensemble_dim: The dimension along which the ensemble
             outputs are stacked. Default is 0.
+        :raises NotImplementedError: If an inverse problem is passed.
         """
         if isinstance(problem, InverseProblem):
             raise NotImplementedError(
-                "DeepEnsemblePINN does not work on inverse problems."
+                "DeepEnsemblePINN can not be used to solve inverse problems."
             )
         super().__init__(
             problem=problem,
@@ -126,11 +127,12 @@ def loss_data(self, input, target):
         :return: The supervised loss, averaged over the number of observations.
         :rtype: torch.Tensor
         """
+        predictions = self.forward(input)
         loss = sum(
-            self._loss_fn(self.forward(input, idx), target)
-            for idx in range(self.num_ensembles)
+            self._loss_fn(predictions[idx], target)
+            for idx in range(self.num_ensemble)
         )
-        return loss / self.num_ensembles
+        return loss / self.num_ensemble
 
     def loss_phys(self, samples, equation):
         """
@@ -160,7 +162,9 @@ def _residual_loss(self, samples, equation):
         :rtype: torch.Tensor
         """
         loss = 0
-        for idx in range(self.num_ensembles):
-            residuals = equation.residual(samples, self.forward(samples, idx))
-            loss = loss + self._loss_fn(residuals, torch.zeros_like(residuals))
-        return loss / self.num_ensembles
+        predictions = self.forward(samples)
+        for idx in range(self.num_ensemble):
+            residuals = equation.residual(samples, predictions[idx])
+            target = torch.zeros_like(residuals, requires_grad=True)
+            loss = loss + self._loss_fn(residuals, target)
+        return loss / self.num_ensemble

pina/solver/ensemble_solver/ensemble_solver_interface.py

@@ -61,8 +61,6 @@ def __init__(
 
         :param AbstractProblem problem: The problem to be solved.
         :param torch.nn.Module models: The neural network models to be used.
-        :param int ensemble_dim: The dimension along which the ensemble
-            outputs are stacked. Default is 0.
         :param Optimizer optimizer: The optimizer to be used.
             If ``None``, the :class:`torch.optim.Adam` optimizer is used.
             Default is ``None``.
@@ -73,6 +71,8 @@ def __init__(
             If ``None``, no weighting schema is used. Default is ``None``.
         :param bool use_lt: If ``True``, the solver uses LabelTensors as input.
             Default is ``True``.
+        :param int ensemble_dim: The dimension along which the ensemble
+            outputs are stacked. Default is 0.
         """
         super().__init__(
             problem, models, optimizers, schedulers, weighting, use_lt
@@ -90,7 +90,8 @@ def forward(self, x, ensemble_idx=None):
 
         :param LabelTensor x: The input tensor to the models.
         :param int ensemble_idx: Optional index to select a specific
-            model from the ensemble.
+            model from the ensemble. If ``None`` results for all models are
+            stacked in ``ensemble_dim`` dimension. Default is ``None``.
         :return: The output of the selected model or the stacked
             outputs from all models.
         :rtype: LabelTensor
@@ -100,7 +101,7 @@ def forward(self, x, ensemble_idx=None):
             return self.models[ensemble_idx].forward(x)
         # otherwise return the stacked output
         return torch.stack(
-            [self.forward(x, idx) for idx in range(self.num_ensembles)],
+            [self.forward(x, idx) for idx in range(self.num_ensemble)],
             dim=self.ensemble_dim,
         )
 
@@ -125,8 +126,9 @@ def training_step(self, batch):
         # perform backpropagation
         self.manual_backward(loss)
         # optimize
-        for opt in self.optimizers:
+        for opt, sched in zip(self.optimizers, self.schedulers):
             opt.instance.step()
+            sched.instance.step()
         return loss
 
     @property
@@ -140,7 +142,7 @@ def ensemble_dim(self):
         return self._ensemble_dim
 
     @property
-    def num_ensembles(self):
+    def num_ensemble(self):
         """
         The number of models in the ensemble.
 

pina/solver/ensemble_solver/ensemble_supervised.py

@@ -114,8 +114,9 @@ def loss_data(self, input, target):
         :return: The supervised loss, averaged over the number of observations.
         :rtype: torch.Tensor
         """
+        predictions = self.forward(input)
         loss = sum(
-            self._loss_fn(self.forward(input, idx), target)
-            for idx in range(self.num_ensembles)
+            self._loss_fn(predictions[idx], target)
+            for idx in range(self.num_ensemble)
         )
-        return loss / self.num_ensembles
+        return loss / self.num_ensemble

pina/solver/physics_informed_solver/pinn_interface.py

@@ -142,10 +142,9 @@ def loss_data(self, input, target):
         between the network's output and the true solution. This method should
         be overridden by the derived class.
 
-        :param input: The input to the neural network.
-        :type input: LabelTensor
-        :param target: The target to compare with the network's output.
-        :type target: LabelTensor
+        :param LabelTensor input: The input to the neural network.
+        :param LabelTensor target: The target to compare with the
+            network's output.
         :return: The supervised loss, averaged over the number of observations.
         :rtype: LabelTensor
         """

tests/test_solver/test_causal_pinn.py

@@ -27,12 +27,12 @@ class DummySpatialProblem(SpatialProblem):
 
 # define problems
 problem = DiffusionReactionProblem()
-problem.discretise_domain(50)
+problem.discretise_domain(10)
 
 # add input-output condition to test supervised learning
-input_pts = torch.rand(50, len(problem.input_variables))
+input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
-output_pts = torch.rand(50, len(problem.output_variables))
+output_pts = torch.rand(10, len(problem.output_variables))
 output_pts = LabelTensor(output_pts, problem.output_variables)
 problem.conditions["data"] = Condition(input=input_pts, target=output_pts)
 

tests/test_solver/test_competitive_pinn.py

@@ -19,19 +19,19 @@
 
 # define problems
 problem = Poisson()
-problem.discretise_domain(50)
+problem.discretise_domain(10)
 inverse_problem = InversePoisson()
-inverse_problem.discretise_domain(50)
+inverse_problem.discretise_domain(10)
 
 # reduce the number of data points to speed up testing
 data_condition = inverse_problem.conditions["data"]
 data_condition.input = data_condition.input[:10]
 data_condition.target = data_condition.target[:10]
 
 # add input-output condition to test supervised learning
-input_pts = torch.rand(50, len(problem.input_variables))
+input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
-output_pts = torch.rand(50, len(problem.output_variables))
+output_pts = torch.rand(10, len(problem.output_variables))
 output_pts = LabelTensor(output_pts, problem.output_variables)
 problem.conditions["data"] = Condition(input=input_pts, target=output_pts)
 

tests/test_solver/test_ensemble_pinn.py

@@ -19,9 +19,9 @@
 problem.discretise_domain(10)
 
 # add input-output condition to test supervised learning
-input_pts = torch.rand(50, len(problem.input_variables))
+input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
-output_pts = torch.rand(50, len(problem.output_variables))
+output_pts = torch.rand(10, len(problem.output_variables))
 output_pts = LabelTensor(output_pts, problem.output_variables)
 problem.conditions["data"] = Condition(input=input_pts, target=output_pts)
 
@@ -42,7 +42,7 @@ def test_constructor():
         InputEquationCondition,
         DomainEquationCondition,
     )
-    assert solver.num_ensembles == 5
+    assert solver.num_ensemble == 5
 
 
 @pytest.mark.parametrize("batch_size", [None, 1, 5, 20])

tests/test_solver/test_ensemble_supervised_solver.py

@@ -30,9 +30,9 @@ class TensorProblem(AbstractProblem):
     }
 
 
-x = torch.rand((100, 20, 5))
-pos = torch.rand((100, 20, 2))
-output_ = torch.rand((100, 20, 1))
+x = torch.rand((15, 20, 5))
+pos = torch.rand((15, 20, 2))
+output_ = torch.rand((15, 20, 1))
 input_ = [
     KNNGraph(x=x_, pos=pos_, neighbours=3, edge_attr=True)
     for x_, pos_ in zip(x, pos)
@@ -44,9 +44,9 @@ class GraphProblem(AbstractProblem):
     conditions = {"data": Condition(input=input_, target=output_)}
 
 
-x = LabelTensor(torch.rand((100, 20, 5)), ["a", "b", "c", "d", "e"])
-pos = LabelTensor(torch.rand((100, 20, 2)), ["x", "y"])
-output_ = LabelTensor(torch.rand((100, 20, 1)), ["u"])
+x = LabelTensor(torch.rand((15, 20, 5)), ["a", "b", "c", "d", "e"])
+pos = LabelTensor(torch.rand((15, 20, 2)), ["x", "y"])
+output_ = LabelTensor(torch.rand((15, 20, 1)), ["u"])
 input_ = [
     KNNGraph(x=x[i], pos=pos[i], neighbours=3, edge_attr=True)
     for i in range(len(x))
@@ -96,7 +96,7 @@ def test_constructor():
     assert DeepEnsembleSupervisedSolver.accepted_conditions_types == (
         InputTargetCondition
     )
-    assert solver.num_ensembles == 10
+    assert solver.num_ensemble == 10
 
 
 @pytest.mark.parametrize("batch_size", [None, 1, 5, 20])

tests/test_solver/test_garom.py

@@ -15,7 +15,7 @@ class TensorProblem(AbstractProblem):
     input_variables = ["u_0", "u_1"]
     output_variables = ["u"]
     conditions = {
-        "data": Condition(target=torch.randn(50, 2), input=torch.randn(50, 1))
+        "data": Condition(target=torch.randn(10, 2), input=torch.randn(10, 1))
     }