new solvers structure

pina/__init__.py

@@ -12,13 +12,10 @@
     "Condition",
     "PinaDataModule",
     "Graph",
-    "SolverInterface",
-    "MultiSolverInterface",
 ]
 
 from pina._src.core.label_tensor import LabelTensor
 from pina._src.core.graph import Graph
-from pina._src.solver.solver import SolverInterface, MultiSolverInterface
 from pina._src.core.trainer import Trainer
 from pina._src.condition.condition import Condition
 from pina._src.data.data_module import PinaDataModule

pina/_src/callback/refinement/r3_refinement.py

@@ -6,7 +6,7 @@
 )
 from pina._src.core.label_tensor import LabelTensor
 from pina._src.core.utils import check_consistency
-from pina._src.loss.loss_interface import LossInterface
+from pina._src.loss.loss_interface import DualLossInterface
 
 
 class R3Refinement(RefinementInterface):
@@ -44,7 +44,7 @@ def __init__(
         :param int sample_every: The sampling frequency.
         :param loss: The loss function to compute the residuals.
             Default is :class:`~torch.nn.L1Loss`.
-        :type loss: LossInterface | :class:`~torch.nn.modules.loss._Loss`
+        :type loss: DualLossInterface | :class:`~torch.nn.modules.loss._Loss`
         :param condition_to_update: The conditions to update during the
             refinement process. If None, all conditions will be updated.
             Default is None.
@@ -59,7 +59,7 @@ def __init__(
         # Check consistency
         check_consistency(
             residual_loss,
-            (LossInterface, torch.nn.modules.loss._Loss),
+            (DualLossInterface, torch.nn.modules.loss._Loss),
             subclass=True,
         )
 

pina/_src/callback/refinement/refinement_interface.py

@@ -6,9 +6,7 @@
 from abc import ABCMeta, abstractmethod
 from lightning.pytorch import Callback
 from pina._src.core.utils import check_consistency
-from pina._src.solver.physics_informed_solver.pinn_interface import (
-    PINNInterface,
-)
+from pina._src.solver.pinn import PINN as PINNInterface
 
 
 class RefinementInterface(Callback, metaclass=ABCMeta):

pina/_src/condition/equation_condition_base.py

@@ -0,0 +1,49 @@
+"""Module for the EquationConditionBase class."""
+
+from pina._src.condition.base_condition import BaseCondition
+
+
+class EquationConditionBase(BaseCondition):
+    """
+    Base class for conditions that involve an equation.
+
+    This class provides the :meth:`evaluate` method, which computes the
+    non-aggregated residual of the equation given the input samples and a
+    solver. It is intended to be subclassed by conditions that define an
+    ``equation`` attribute, such as
+    :class:`~pina.condition.DomainEquationCondition` and
+    :class:`~pina.condition.InputEquationCondition`.
+    """
+
+    def evaluate(self, batch, solver, loss):
+        """
+        Evaluate the equation residual on the given batch using the solver.
+
+        This method computes the non-aggregated, element-wise residual of the
+        equation. It performs a forward pass of the solver's model on the
+        input samples and then evaluates the equation residual. The returned
+        tensor is **not** reduced (i.e., no mean, sum, etc.), preserving the
+        per-sample residual values.
+
+        :param batch: The batch containing the ``input`` entry.
+        :type batch: dict | _DataManager
+        :param solver: The solver containing the model and any additional
+            parameters (e.g., unknown parameters for inverse problems).
+        :type solver: ~pina.solver.solver.SolverInterface
+        :param loss: The non-aggregating loss function to apply to the
+            computed residual against zero.
+        :type loss: torch.nn.Module
+        :return: The non-aggregated loss tensor.
+        :rtype: ~pina.label_tensor.LabelTensor
+
+        :Example:
+
+            >>> residuals = condition.evaluate(
+            ...     {"input": input_samples}, solver, loss
+            ... )
+            >>> # residuals is a non-reduced tensor of shape (n_samples, ...)
+        """
+        samples = batch["input"].requires_grad_(True)
+        return self.equation.residual(
+            samples, solver.forward(samples), solver._params
+        )

pina/_src/condition/input_equation_condition.py

@@ -107,3 +107,36 @@ def equation(self, value):
         # Check consistency
         check_consistency(value, self._avail_equation_cls)
         self._equation = value
+
+    def evaluate(self, batch, solver, loss):
+        """
+        Evaluate the equation residual on the given batch using the solver.
+
+        This method computes the non-aggregated, element-wise residual of the
+        equation. It performs a forward pass of the solver's model on the
+        input samples and then evaluates the equation residual. The returned
+        tensor is **not** reduced (i.e., no mean, sum, etc.), preserving the
+        per-sample residual values.
+
+        :param batch: The batch containing the ``input`` entry.
+        :type batch: dict | _DataManager
+        :param solver: The solver containing the model and any additional
+            parameters (e.g., unknown parameters for inverse problems).
+        :type solver: ~pina.solver.solver.SolverInterface
+        :param loss: The non-aggregating loss function to apply to the
+            computed residual against zero.
+        :type loss: torch.nn.Module
+        :return: The non-aggregated loss tensor.
+        :rtype: ~pina.label_tensor.LabelTensor
+
+        :Example:
+
+            >>> residuals = condition.evaluate(
+            ...     {"input": input_samples}, solver, loss
+            ... )
+            >>> # residuals is a non-reduced tensor of shape (n_samples, ...)
+        """
+        samples = batch["input"].requires_grad_(True)
+        return self.equation.residual(
+            samples, solver.forward(samples), solver._params
+        )

pina/_src/condition/input_target_condition.py

@@ -104,3 +104,21 @@ def target(self):
             list[Data] | tuple[Graph] | tuple[Data]
         """
         return self.data.target
+
+    def evaluate(self, batch, solver, loss):
+        """
+        Evaluate the supervised condition on the given batch using the solver.
+
+        This method computes the element-wise loss associated with the
+        condition using the input and target stored in the provided batch.
+
+        :param batch: The batch containing ``input`` and ``target`` entries.
+        :type batch: dict | _DataManager
+        :param solver: The solver containing the model.
+        :type solver: ~pina.solver.solver.SolverInterface
+        :param loss: The non-aggregating loss function to apply.
+        :type loss: torch.nn.Module
+        :return: The non-aggregated loss tensor.
+        :rtype: LabelTensor | torch.Tensor | Graph | Data
+        """
+        return loss(solver.forward(batch["input"]), batch["target"])

pina/_src/condition/time_series_condition.py

@@ -0,0 +1,168 @@
+"""Module for the TimeSeriesCondition class."""
+
+import torch
+
+from pina._src.data.manager.data_manager import _DataManager
+from pina._src.core.label_tensor import LabelTensor
+from pina._src.condition.base_condition import BaseCondition
+
+
+class TimeSeriesCondition(BaseCondition):
+    """
+    Condition for autoregressive time-series training.
+
+    The condition stores an input tensor containing unroll windows with shape
+    ``[trajectories, windows, time_steps, *features]`` and computes the
+    autoregressive non-aggregated/aggregated temporal loss inside
+    :meth:`evaluate` by recursively applying the solver model over time.
+    """
+
+    __fields__ = ["input", "eps", "aggregation_strategy", "kwargs"]
+    _avail_input_cls = (torch.Tensor, LabelTensor)
+
+    def __new__(cls, input, eps=None, aggregation_strategy=None, kwargs=None):
+        if cls != TimeSeriesCondition:
+            return super().__new__(cls)
+
+        if not isinstance(input, cls._avail_input_cls):
+            raise ValueError(
+                "Invalid input type. Expected one of the following: "
+                "torch.Tensor, LabelTensor."
+            )
+
+        return super().__new__(cls)
+
+    def store_data(self, **kwargs):
+        return _DataManager(input=kwargs.get("input"))
+
+    @property
+    def input(self):
+        return self.data.input
+
+    @property
+    def settings(self):
+        return {
+            "eps": getattr(self, "_eps", None),
+            "aggregation_strategy": getattr(
+                self, "_aggregation_strategy", None
+            ),
+            "kwargs": getattr(self, "_kwargs", {}),
+        }
+
+    def __init__(self, input, eps=None, aggregation_strategy=None, kwargs=None):
+        super().__init__(input=input)
+        self._eps = eps
+        self._aggregation_strategy = aggregation_strategy
+        self._kwargs = kwargs or {}
+
+    def evaluate(self, batch, solver, loss, condition_name=None):
+        input_tensor = batch["input"]
+
+        if input_tensor.dim() < 4:
+            raise ValueError(
+                "The provided input tensor must have at least 4 dimensions:"
+                " [trajectories, windows, time_steps, *features]."
+                f" Got shape {input_tensor.shape}."
+            )
+
+        current_state = input_tensor[:, :, 0]
+        losses = []
+        step_kwargs = self._kwargs.copy()
+
+        for step in range(1, input_tensor.shape[2]):
+            processed_input = solver.preprocess_step(
+                current_state, **step_kwargs
+            )
+            output = solver.forward(processed_input)
+            predicted_state = solver.postprocess_step(output, **step_kwargs)
+
+            target_state = input_tensor[:, :, step]
+            step_loss = loss(predicted_state, target_state, **step_kwargs)
+            losses.append(step_loss)
+            current_state = predicted_state
+
+        step_losses = torch.stack(losses).as_subclass(torch.Tensor)
+
+        with torch.no_grad():
+            name = condition_name or getattr(self, "name", None) or "default"
+            # weights = solver._get_weights(name, step_losses, self._eps)
+
+        aggregation_strategy = self._aggregation_strategy or torch.mean
+        return aggregation_strategy(step_losses)  # * weights)
+
+    @staticmethod
+    def unroll(data, unroll_length, n_unrolls=None, randomize=True):
+        """
+        Create unrolling time windows from temporal data.
+
+        This function takes as input a tensor of shape
+        ``[trajectories, time_steps, *features]`` and produces a tensor of
+        shape ``[trajectories, windows, unroll_length, *features]``.
+        Each window contains a sequence of subsequent states used for
+        computing the multi-step loss during training.
+
+        :param data: The temporal data tensor to be unrolled.
+        :type data: torch.Tensor | LabelTensor
+        :param int unroll_length: The number of time steps in each window.
+        :param int n_unrolls: The maximum number of windows to return.
+            If ``None``, all valid windows are returned. Default is ``None``.
+        :param bool randomize: If ``True``, starting indices are randomly
+            permuted before applying ``n_unrolls``. Default is ``True``.
+        :raise ValueError: If the input ``data`` has less than 3 dimensions.
+        :raise ValueError: If ``unroll_length`` is greater or equal to the
+            number of time steps in ``data``.
+        :return: A tensor of unrolled windows.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        if data.dim() < 3:
+            raise ValueError(
+                "The provided data tensor must have at least 3 dimensions:"
+                " [trajectories, time_steps, *features]."
+                f" Got shape {data.shape}."
+            )
+
+        start_idx = TimeSeriesCondition._get_start_idx(
+            n_steps=data.shape[1],
+            unroll_length=unroll_length,
+            n_unrolls=n_unrolls,
+            randomize=randomize,
+        )
+
+        windows = [data[:, s : s + unroll_length] for s in start_idx]
+        return torch.stack(windows, dim=1)
+
+    @staticmethod
+    def _get_start_idx(n_steps, unroll_length, n_unrolls=None, randomize=True):
+        """
+        Determine starting indices for unroll windows.
+
+        :param int n_steps: The total number of time steps in the data.
+        :param int unroll_length: The number of time steps in each window.
+        :param int n_unrolls: The maximum number of windows to return.
+            If ``None``, all valid windows are returned. Default is ``None``.
+        :param bool randomize: If ``True``, starting indices are randomly
+            permuted before applying ``n_unrolls``. Default is ``True``.
+        :raise ValueError: If ``unroll_length`` is greater or equal to the
+            number of time steps in ``data``.
+        :return: A tensor of starting indices for unroll windows.
+        :rtype: torch.Tensor
+        """
+        last_idx = n_steps - unroll_length
+
+        if last_idx < 0:
+            raise ValueError(
+                "Cannot create unroll windows: "
+                f"unroll_length ({unroll_length})"
+                " cannot be greater or equal to the number of time_steps"
+                f" ({n_steps})."
+            )
+
+        indices = torch.arange(last_idx + 1)
+
+        if randomize:
+            indices = indices[torch.randperm(len(indices))]
+
+        if n_unrolls is not None and n_unrolls < len(indices):
+            indices = indices[:n_unrolls]
+
+        return indices

pina/_src/core/trainer.py

@@ -6,12 +6,13 @@
 import lightning
 from pina._src.core.utils import check_consistency, custom_warning_format
 from pina._src.data.data_module import PinaDataModule
-from pina._src.solver.supervised_solver.supervised_solver_interface import (
+from pina._src.solver.solver_interface import (
     SolverInterface,
 )
-from pina._src.solver.physics_informed_solver.pinn_interface import (
-    PINNInterface,
-)
+
+# from pina._src.solver.physics_informed_solver.pinn_interface import (
+#     PINNInterface,
+# )
 
 # set the warning for compile options
 warnings.formatwarning = custom_warning_format
@@ -99,8 +100,8 @@ def __init__(
 
         # inference mode set to false when validating/testing PINNs otherwise
         # gradient is not tracked and optimization_cycle fails
-        if isinstance(solver, PINNInterface):
-            kwargs["inference_mode"] = False
+        # if isinstance(solver, PINNInterface):
+        kwargs["inference_mode"] = False
 
         # Logging depends on the batch size, when batch_size is None then
         # log_every_n_steps should be zero

pina/_src/loss/base_loss.py

@@ -1,10 +1,10 @@
 """Module for the BaseLoss class."""
 
 import torch
-from pina._src.loss.loss_interface import LossInterface
+from pina._src.loss.loss_interface import DualLossInterface
 
 
-class BaseLoss(LossInterface):
+class BaseLoss(DualLossInterface):
     """
     Base class for all losses, implementing common functionality.
 

pina/_src/loss/loss_interface.py

@@ -4,7 +4,7 @@
 from torch.nn.modules.loss import _Loss
 
 
-class LossInterface(_Loss, metaclass=ABCMeta):
+class DualLossInterface(_Loss, metaclass=ABCMeta):
     """
     Abstract interface for all losses.
     """