separate residuals and losses logic

Author: GiovanniCanali

pina/_src/condition/condition_interface.py

@@ -57,7 +57,7 @@ def create_dataloader(
         """
 
     @abstractmethod
-    def evaluate(self, batch, solver, loss):
+    def evaluate(self, batch, solver):
         """
         Evaluate the residual of the condition on the given batch using the
         solver.
@@ -75,8 +75,6 @@ def evaluate(self, batch, solver, loss):
             pass and compute the residual. The solver provides access to the
             model and its parameters, which may be necessary for evaluating the
             condition residual.
-        :param torch.nn.Module loss: The non-aggregating loss function used to
-            compare the condition residual against its reference value.
         :return: The non-aggregated residual tensor.
         :rtype: torch.Tensor | LabelTensor
         """

pina/_src/condition/data_condition.py

@@ -85,7 +85,7 @@ def store_data(self, **kwargs):
 
         return _DataManager(**data_dict)
 
-    def evaluate(self, batch, solver, loss):
+    def evaluate(self, batch, solver):
         """
         Evaluate the residual of the condition on the given batch using the
         solver.
@@ -103,13 +103,10 @@ def evaluate(self, batch, solver, loss):
             pass and compute the residual. The solver provides access to the
             model and its parameters, which may be necessary for evaluating the
             condition residual.
-        :param torch.nn.Module loss: The non-aggregating loss function used to
-            compare the condition residual against its reference value.
         :return: The non-aggregated residual tensor.
         :rtype: torch.Tensor | LabelTensor
         """
-        output_ = solver.forward(batch["input"])
-        return loss(output_, torch.zeros_like(output_))
+        return solver.forward(batch["input"])
 
     @property
     def conditional_variables(self):

pina/_src/condition/domain_equation_condition.py

@@ -75,7 +75,7 @@ def store_data(self, **kwargs):
         setattr(self, "domain", kwargs.get("domain"))
         setattr(self, "equation", kwargs.get("equation"))
 
-    def evaluate(self, batch, solver, loss):
+    def evaluate(self, batch, solver):
         """
         Evaluate the residual of the condition on the given batch using the
         solver.
@@ -93,8 +93,6 @@ def evaluate(self, batch, solver, loss):
             pass and compute the residual. The solver provides access to the
             model and its parameters, which may be necessary for evaluating the
             condition residual.
-        :param torch.nn.Module loss: The non-aggregating loss function used to
-            compare the condition residual against its reference value.
         :raises NotImplementedError: Always raised since any domain-equation
             condition is transformed into an input-equation condition before
             evaluation, and the residual is computed using the input-equation

pina/_src/condition/input_equation_condition.py

@@ -109,7 +109,7 @@ def equation(self, value):
         check_consistency(value, self._avail_equation_cls)
         self._equation = value
 
-    def evaluate(self, batch, solver, loss):
+    def evaluate(self, batch, solver):
         """
         Evaluate the residual of the condition on the given batch using the
         solver.
@@ -127,15 +127,11 @@ def evaluate(self, batch, solver, loss):
             pass and compute the residual. The solver provides access to the
             model and its parameters, which may be necessary for evaluating the
             condition residual.
-        :param torch.nn.Module loss: The non-aggregating loss function used to
-            compare the condition residual against its reference value.
         :return: The non-aggregated residual tensor.
         :rtype: LabelTensor
         """
         # Compute residuals
         samples = batch["input"].requires_grad_(True)
-        residual = self.equation.residual(
+        return self.equation.residual(
             samples, solver.forward(samples), solver._params
         )
-
-        return loss(residual, torch.zeros_like(residual))

pina/_src/condition/input_target_condition.py

@@ -83,7 +83,7 @@ def store_data(self, **kwargs):
         """
         return _DataManager(**kwargs)
 
-    def evaluate(self, batch, solver, loss):
+    def evaluate(self, batch, solver):
         """
         Evaluate the residual of the condition on the given batch using the
         solver.
@@ -101,12 +101,10 @@ def evaluate(self, batch, solver, loss):
             pass and compute the residual. The solver provides access to the
             model and its parameters, which may be necessary for evaluating the
             condition residual.
-        :param torch.nn.Module loss: The non-aggregating loss function used to
-            compare the condition residual against its reference value.
         :return: The non-aggregated residual tensor.
         :rtype: torch.Tensor | LabelTensor
         """
-        return loss(solver.forward(batch["input"]), batch["target"])
+        return solver.forward(batch["input"]) - batch["target"]
 
     @property
     def input(self):

pina/_src/condition/time_series_condition.py

@@ -170,32 +170,32 @@ def _unroll(self, data, n_windows, unroll_length, randomize):
 
         return torch.stack(windows, dim=1)
 
-    def evaluate(self, batch, solver, loss):
+    def evaluate(self, batch, solver):
         """
         Evaluate the residual of the condition on the given batch using the
         solver.
 
-        This method computes the non-aggregated, element-wise residual of the
-        condition. A forward pass of the solver's model is performed on the
-        input samples, and the condition residual is evaluated accordingly.
+        This method computes the per-step residuals through autoregressive
+        unrolling. A forward pass of the solver's model is performed at each
+        time step, and the per-step residuals (predicted - target) are
+        returned as a stacked tensor.
 
-        The returned tensor is not reduced, preserving the per-sample residual
-        values.
+        The returned tensor preserves all per-step residual values without
+        reduction or loss aggregation.
 
         :param dict batch: The batch containing the data required by the
             condition evaluation.
         :param SolverInterface solver: The solver used to perform the forward
             pass and compute the residual. The solver provides access to the
             model and its parameters, which may be necessary for evaluating the
             condition residual.
-        :param torch.nn.Module loss: The non-aggregating loss function used to
-            compare the condition residual against its reference value.
         :raises ValueError: If the input tensor in the batch has less than 4
             dimensions.
-        :return: The non-aggregated residual tensor.
+        :return: The stacked per-step residual tensor of shape
+            [time_steps - 1, trajectories, windows, *features].
         :rtype: torch.Tensor | LabelTensor
         """
-        # Raise error if input tensor does not have at least4 dimensions
+        # Raise error if input tensor does not have at least 4 dimensions
         if batch["input"].dim() < 4:
             raise ValueError(
                 "The provided input tensor must have at least 4 dimensions:"
@@ -206,9 +206,9 @@ def evaluate(self, batch, solver, loss):
         # Copy the kwargs to avoid modifying the original settings
         kwargs = solver._kwargs.copy()
 
-        # Extract the initial state and initialize the list of step-wise losses
+        # Extract the initial state and initialize the step-wise residuals list
         current_state = batch["input"][:, :, 0]
-        losses = []
+        residuals = []
 
         # Iterate over the time steps
         for step in range(1, batch["input"].shape[2]):
@@ -218,23 +218,16 @@ def evaluate(self, batch, solver, loss):
             output = solver.forward(processed_input)
             predicted_state = solver.postprocess_step(output, **kwargs)
 
-            # Retrieve the target and compute the step-wise loss
+            # Retrieve the target and compute the step-wise residual
             target_state = batch["input"][:, :, step]
-            step_loss = loss(predicted_state, target_state, **kwargs)
-            losses.append(step_loss)
+            step_residual = predicted_state - target_state
+            residuals.append(step_residual)
 
             # Update the current state for the next iteration
             current_state = predicted_state
 
-        # Stack the step-wise losses
-        step_losses = torch.stack(losses).as_subclass(torch.Tensor)
-
-        # Compute adaptive weights and aggregate the step-wise losses
-        with torch.no_grad():
-            name = getattr(self, "name", None) or "default"
-            weights = solver._get_weights(name, step_losses)
-
-        return solver.aggregation_strategy(step_losses * weights)
+        # Stack the step-wise residuals
+        return torch.stack(residuals).as_subclass(torch.Tensor)
 
     @property
     def input(self):

pina/_src/solver/autoregressive_solver.py

@@ -216,6 +216,45 @@ def postprocess_step(self, predicted_state, **kwargs):
         """
         return predicted_state
 
+    def optimization_cycle(self, batch):
+        """
+        Compute one reduced, aggregated loss per condition in the batch.
+
+        For TimeSeriesCondition, this method evaluates the condition to obtain
+        per-step residuals, applies the pointwise loss function to each step,
+        computes adaptive weights based on the step-wise losses, and returns
+        the aggregated weighted loss.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :return: The reduced, aggregated losses for all conditions.
+        :rtype: dict[str, torch.Tensor]
+        """
+        condition_losses = {}
+
+        for condition_name, data in batch:
+            condition = self.problem.conditions[condition_name]
+            condition_data = dict(data)
+
+            # Evaluate condition to get per-step residuals
+            self.step_residuals = condition.evaluate(condition_data, self)
+
+            # Apply the loss function to each step-wise residual
+            step_losses = self._loss_fn(
+                self.step_residuals, torch.zeros_like(self.step_residuals)
+            )
+
+            # Compute adaptive weights and aggregate the step-wise losses
+            with torch.no_grad():
+                name = condition_name or "default"
+                weights = self._get_weights(name, step_losses)
+
+            # Aggregate using the configured strategy
+            aggregated_loss = self.aggregation_strategy(step_losses * weights)
+            condition_losses[condition_name] = aggregated_loss
+
+        return condition_losses
+
     def predict(self, initial_state, n_steps, **kwargs):
         """
         Generate predictions by recursively calling the model's forward.

pina/_src/solver/multi_model_simple_solver.py

@@ -30,22 +30,24 @@ class MultiModelSimpleSolver(BaseSolver):
         \quad i = 1, \dots, N_{\rm ensemble}
 
     During the optimization cycle each model's prediction is evaluated against
-    the condition independently, and the resulting per-model losses are
-    averaged to form the aggregated condition loss:
+    the condition independently, the residual is converted into a pointwise
+    loss, and the resulting per-model losses are averaged to form the
+    aggregated condition loss:
 
     .. math::
         \mathcal{L}_{\rm condition} = \frac{1}{N_{\rm ensemble}}
         \sum_{i=1}^{N_{\rm ensemble}} \mathcal{L}_i
 
     The per-condition workflow is:
 
-         1. evaluate the condition for each model and obtain non-aggregated
-            loss tensors;
-         2. apply the configured reduction to each per-model tensor;
-         3. average the reduced per-model losses into a single scalar for
-            the condition;
-         4. return the per-condition losses, which are aggregated by the
-            inherited solver machinery through the configured weighting.
+     1. evaluate the condition for each model and obtain non-aggregated
+         residual tensors;
+     2. apply the pointwise loss and the configured reduction to each
+         per-model tensor;
+     3. average the reduced per-model losses into a single scalar for the
+         condition;
+     4. return the per-condition losses, which are aggregated by the
+         inherited solver machinery through the configured weighting.
     """
 
     accepted_conditions_types = (
@@ -191,7 +193,7 @@ def forward(self, x, model_idx=None):
         :rtype: LabelTensor | torch.Tensor
         """
         if model_idx is not None:
-            return self.models[model_idx].forward(x)
+            return self.models[model_idx](x)
         return torch.stack(
             [self.forward(x, idx) for idx in range(self.num_models)],
         )
@@ -259,10 +261,16 @@ def optimization_cycle(self, batch):
                         self.problem.output_variables,
                     )
 
-                loss_tensor = condition.evaluate(
-                    condition_data, self, self._loss_fn
+                # Store the residual tensor
+                self.residual_tensor = condition.evaluate(condition_data, self)
+
+                # Compute the per-sample loss tensor
+                loss_tensor = self._loss_fn(
+                    self.residual_tensor, torch.zeros_like(self.residual_tensor)
                 )
                 self.forward = original_forward
+
+                # Apply reduction and store the result
                 per_model_losses.append(self._apply_reduction(loss_tensor))
 
             condition_losses[condition_name] = torch.stack(

pina/_src/solver/single_model_simple_solver.py

@@ -22,8 +22,9 @@ class SingleModelSimpleSolver(BaseSolver):
 
     The solver orchestrates a uniform workflow for all conditions in the batch:
 
-         1. evaluate the condition and obtain a non-aggregated loss tensor;
-         2. apply a reduction to obtain a scalar loss for that condition;
+     1. evaluate the condition and obtain a non-aggregated residual tensor;
+     2. apply the pointwise loss and a reduction to obtain a scalar loss for
+         that condition;
      4. return the per-condition losses, which are aggregated by the inherited
        solver machinery through the configured weighting.
     """
@@ -99,12 +100,19 @@ def optimization_cycle(self, batch):
             condition = self.problem.conditions[condition_name]
             condition_data = dict(data)
 
-            condition_loss_tensor = condition.evaluate(
-                condition_data, self, self._loss_fn
+            # Store the residual tensor
+            self.residual_tensor = condition.evaluate(condition_data, self)
+
+            # Compute the per-sample loss tensor
+            loss_tensor = self._loss_fn(
+                self.residual_tensor, torch.zeros_like(self.residual_tensor)
             )
+
+            # Apply reduction and store the result
             condition_losses[condition_name] = self._apply_reduction(
-                condition_loss_tensor
+                loss_tensor
             )
+
         return condition_losses
 
     def _apply_reduction(self, value):

tests/test_condition/test_data_condition.py

@@ -382,9 +382,8 @@ def test_evaluate(case, use_lt, conditional_variables):
         }
 
     # Evaluate the condition and compute the expected value
-    loss = condition.evaluate(batch, solver, loss_fn)
-    output_ = solver.forward(batch["input"])
-    expected = loss_fn(output_, torch.zeros_like(output_))
+    loss = condition.evaluate(batch, solver)
+    expected = solver.forward(batch["input"])
 
     # Assert that the evaluated loss is correct
     assert torch.allclose(loss, expected)