labelize_forward on ensemble

Author: ndem0

pina/_src/solver/ensemble_simple_solver.py

@@ -1,6 +1,7 @@
 """Module for the DeepEnsemble simple solver."""
 
 from pina._src.solver.multi_model_simple_solver import MultiModelSimpleSolver
+from pina._src.core.utils import check_consistency
 
 
 class DeepEnsembleSimpleSolver(MultiModelSimpleSolver):
@@ -99,5 +100,7 @@ def __init__(
             weighting=weighting,
             loss=loss,
             use_lt=use_lt,
-            ensemble_dim=ensemble_dim,
         )
+
+        check_consistency(ensemble_dim, int)
+        self.num_ensemble = len(models)

pina/_src/solver/multi_model_simple_solver.py

@@ -62,7 +62,6 @@ def __init__(
         weighting=None,
         loss=None,
         use_lt=True,
-        ensemble_dim=0,
     ):
         """
         Initialize the multi-model simple solver.
@@ -90,7 +89,6 @@ def __init__(
             loss = torch.nn.MSELoss()
 
         check_consistency(loss, (LossInterface, _Loss), subclass=False)
-        check_consistency(ensemble_dim, int)
 
         super().__init__(
             problem=problem,
@@ -103,7 +101,6 @@ def __init__(
 
         self._loss_fn = loss
         self._reduction = getattr(loss, "reduction", "mean")
-        self._ensemble_dim = ensemble_dim
 
         if hasattr(self._loss_fn, "reduction"):
             self._loss_fn.reduction = "none"
@@ -194,9 +191,16 @@ def optimization_cycle(self, batch):
                 self.forward = (  # noqa: E731
                     lambda x, _idx=idx: self.models[_idx].forward(x)
                 )
+                from pina._src.core.utils import labelize_forward
+                problem = self.problem
+                self.forward = labelize_forward(
+                    self.forward, 
+                    input_variables=problem.input_variables,
+                    output_variables=problem.output_variables,
+                )
                 loss_tensor = condition.evaluate(
                     condition_data, self, self._loss_fn
-                )
+                ).tensor
                 self.forward = original_forward
                 per_model_losses.append(self._apply_reduction(loss_tensor))
 

tests/test_solver/test_causal_pinn.py tests/test_solver/old_causal_pinn.pytests/test_solver/test_causal_pinn.py renamed to tests/test_solver/old_causal_pinn.py

@@ -9,6 +9,7 @@
 from pina.condition import TimeSeriesCondition
 from pina.problem import AbstractProblem
 from pina.model import FeedForward
+from torch._dynamo import OptimizedModule
 
 
 # Hyperparameters and settings