update tests for cleaner version

Author: Davide-Miotti

pina/_src/solver/autoregressive_solver/autoregressive_solver.py

@@ -42,7 +42,6 @@ def __init__(
         self,
         problem,
         model,
-        eps=None,
         loss=None,
         optimizer=None,
         scheduler=None,
@@ -57,8 +56,6 @@ def __init__(
             the time series data conditions.
         :param torch.nn.Module model: Neural network that predicts the
             next state given the current state.
-        :param float eps: If provided, applies exponential weighting to the per-step losses.
-            If ``None``, uniform weights are used. Default is ``None``.
         :param torch.nn.Module loss: Loss function to minimize.
             If ``None``, :class:`torch.nn.MSELoss` is used.
             Default is ``None``.
@@ -207,15 +204,13 @@ def loss_data(
 
         :param torch.Tensor unroll: Batch of unroll windows with shape
             ``[B1, B2, Twin, *state_shape]`` where ``Twin = unroll_length + 1``.
+        :param str condition_name: Name of the condition associated with this data.
         :param float eps: If provided, applies step weighting through
             :meth:`weighting_strategy`. If ``None``, uniform normalized weights are used.
         :param callable aggregation_strategy: Reduction applied to the weighted per-step
             losses. If ``None``, :func:`torch.sum` is used.
-        :param str condition_name: Name of the condition associated with this data.
-        :kwargs: Additional keyword arguments forwarded to
-            :meth:`_step_kwargs` and subsequently to :meth:`preprocess_step`.
+        :kwargs: Additional keyword arguments.
         :return: Scalar loss value for the given batch.
-        :rtype: torch.Tensor
         """
         if unroll.dim() < 4:
             raise ValueError(
@@ -225,23 +220,19 @@ def loss_data(
         B1, B2, Twin = unroll.shape[0], unroll.shape[1], unroll.shape[2]
         state_shape = unroll.shape[3:]
 
-        # current_value = unroll[:, :, 0, ...]  # first time step of each batch
         losses = []
-
         for step in range(1, Twin):
 
             model_input = self.preprocess_step(unroll, step=step, **kwargs)
             model_output = self.model_forward(model_input)
-            predicted_state = self.postprocess_step(model_output, unroll=unroll, step=step, **kwargs)
+            predicted_state = self.postprocess_step(
+                model_output, unroll=unroll, step=step, **kwargs
+            )
 
             target_state = unroll[:, :, step, ...]
-            step_loss = self._loss_fn(
-                predicted_state, target_state, **kwargs
-            )
+            step_loss = self._loss_fn(predicted_state, target_state, **kwargs)
             losses.append(step_loss)
 
-            # current_value = predicted_state
-
         step_losses = torch.stack(losses)  # [unroll_length]
 
         with torch.no_grad():
@@ -260,7 +251,7 @@ def loss_data(
             aggregation_strategy = torch.sum
 
         return aggregation_strategy(step_losses * weights)
-    
+
     def preprocess_step(self, unroll, step=None, **kwargs):
         """
         Pre-process the input unroll for the current step before feeding it to the model.
@@ -276,7 +267,6 @@ def preprocess_step(self, unroll, step=None, **kwargs):
             return unroll
         else:
             return unroll[:, :, step - 1, ...]
-       
 
     def model_forward(self, model_input, **kwargs):
         """
@@ -292,8 +282,8 @@ def model_forward(self, model_input, **kwargs):
         """
 
         return self.model(model_input)
-    
-    def postprocess_step(self, model_output, unroll=None, step=None, **kwargs):
+
+    def postprocess_step(self, model_output, unroll, step=None, **kwargs):
         """
         Post-process the predicted state after obtaining it from the model.
         This method can be overridden by subclasses to implement specific post-processing logic.
@@ -303,13 +293,17 @@ def postprocess_step(self, model_output, unroll=None, step=None, **kwargs):
         or incorporating additional context from the unroll.
 
         :param torch.Tensor model_output: The output of the model.
+        :param torch.Tensor unroll: The original unroll tensor, which can be used for context.
+        :param int step: The current step index within the unroll.
+            By default is ``None``, which is meant to be used in inference.
+        :kwargs: Additional keyword arguments for post-processing.
         :return: The post-processed model output.
         :rtype: torch.Tensor
         """
-        if unroll is not None and step is not None:
-            #reshape model output to match target shape if needed
-            if model_output.shape != unroll[:, :, step, ...].shape:
-                model_output = model_output.view_as(unroll[:, :, step, ...])
+        if step is not None:
+            # do the logic for the training phase, also involving unroll if needed
+            return model_output
+
         return model_output
 
     def get_weights(self, condition_name, step_losses, eps):
@@ -401,12 +395,12 @@ def predict(self, initial_input, num_steps, **kwargs):
 
         with torch.no_grad():
             for step in range(1, num_steps + 1):
-                model_input = self.preprocess_state(
-                    predictions[-1], **kwargs
+                model_input = self.preprocess_step(
+                    predictions[-1], step=None, **kwargs
                 )
                 next_state = self.model_forward(model_input)
-                next_state = self.post_process_state(
-                    next_state, **kwargs
+                next_state = self.postprocess_step(
+                    next_state, unroll=predictions[-1], step=None, **kwargs
                 )
                 predictions.append(next_state)
 

pina/_src/solver/autoregressive_solver/autoregressive_solver_interface.py

@@ -79,11 +79,19 @@ def optimization_cycle(self, batch):
 
         condition_loss = {}
         for condition_name, points in batch:
+            if hasattr(self.problem.conditions[condition_name], "settings"):
+                settings = self.problem.conditions[condition_name].settings
+                eps = settings.get("eps", None)
+                kwargs = settings.get("kwargs", {})
+            else:
+                eps = None
+                kwargs = {}
+
             loss = self.loss_data(
-                points["input"]["unroll"],
+                points["input"],
                 condition_name=condition_name,
-                eps=points["input"].get("eps", None),
-                **points["input"].get("kwargs", {})
+                eps=eps,
+                **kwargs,
             )
             condition_loss[condition_name] = loss
         return condition_loss

tests/test_solver/test_autoregressive_solver.py

@@ -53,26 +53,25 @@ def test_end_to_end(y_data_large):
     the AutoregressiveSolver with curriculum learning
     """
 
-    # AbstratProblem with empty conditions and conditions_settings to be filled later
+    # AbstratProblem with empty conditions to be filled later
     class Problem(AbstractProblem):
         output_variables = None
         input_variables = None
         conditions = {}
-        conditions_settings = {}
 
     problem = Problem()
 
     solver = AutoregressiveSolver(
         problem=problem,
         model=MinimalModel(),
-        optimizer=TorchOptimizer(torch.optim.AdamW, lr=0.008),
+        optimizer=TorchOptimizer(torch.optim.AdamW, lr=0.015),
     )
     # PHASE1: train with 'short' condition only
     y_short = AutoregressiveSolver.unroll(
         y_data_large, unroll_length=4, num_unrolls=20, randomize=False
     )
     problem.conditions["short"] = DataCondition(input=y_short)
-    problem.conditions_settings["short"] = {"eps": 0.1}
+    problem.conditions["short"].settings = {"eps": 0.1}
     trainer1 = Trainer(
         solver, max_epochs=300, accelerator="cpu", enable_model_summary=False
     )
@@ -84,8 +83,7 @@ class Problem(AbstractProblem):
     )
     problem.conditions.clear()
     problem.conditions["medium"] = DataCondition(input=y_medium)
-    problem.conditions_settings.clear()
-    problem.conditions_settings["medium"] = {"eps": 0.2}
+    problem.conditions["medium"].settings = {"eps": 0.2}
     trainer2 = Trainer(
         solver, max_epochs=1500, accelerator="cpu", enable_model_summary=False
     )
@@ -97,8 +95,7 @@ class Problem(AbstractProblem):
     )
     problem.conditions.clear()
     problem.conditions["long"] = DataCondition(input=y_long)
-    problem.conditions_settings.clear()
-    problem.conditions_settings["long"] = {"eps": 0.2}
+    problem.conditions["long"].settings = {"eps": 0.25}
     trainer3 = Trainer(
         solver, max_epochs=4000, accelerator="cpu", enable_model_summary=False
     )
@@ -117,10 +114,10 @@ class Problem(AbstractProblem):
     total_mse = torch.nn.functional.mse_loss(
         prediction.squeeze(1)[:, 1:, :], ground_truth[:, 1:, :]
     )
-    assert total_mse < 1e-6
+    assert total_mse < 1e-5
 
 
-### UNIT TESTS #############################################################################
+# ### UNIT TESTS #############################################################################
 
 NUM_TIMESTEPS = 10
 NUM_FEATURES = 3