implement autoregressive condition, time_weighting, solver

Author: Davide-Miotti

pina/condition/__init__.py

@@ -15,6 +15,7 @@
     "DataCondition",
     "GraphDataCondition",
     "TensorDataCondition",
+    "AutoregressiveCondition",
 ]
 
 from .condition_interface import ConditionInterface
@@ -37,3 +38,5 @@
     GraphDataCondition,
     TensorDataCondition,
 )
+
+from .autoregressive_condition import AutoregressiveCondition

pina/condition/autoregressive_condition.py

@@ -0,0 +1,91 @@
+import torch
+from .condition_interface import ConditionInterface
+from ..loss import TimeWeightingInterface, ConstantTimeWeighting
+from ..utils import check_consistency
+
+
+class AutoregressiveCondition(ConditionInterface):
+    """
+    A specialized condition for autoregressive tasks.
+    It generates input/unroll pairs from a single time-series tensor.
+    """
+
+    __slots__ = ["input", "unroll"]
+
+    def __init__(
+        self,
+        data,
+        unroll_length,
+        num_unrolls=None,
+        randomize=True,
+        time_weighting=None,
+    ):
+        """
+        Create an AutoregressiveCondition.
+        """
+        super().__init__()
+
+        self._n_timesteps, n_features = data.shape
+        self._unroll_length = unroll_length
+        self._requested_num_unrolls = num_unrolls
+        self._randomize = randomize
+
+        # time weighting: weight the loss differently along the unroll
+        if time_weighting is None:
+            self._time_weighting = ConstantTimeWeighting()
+        else:
+            check_consistency(time_weighting, TimeWeightingInterface)
+            self._time_weighting = time_weighting
+
+        # windows creation
+        initial_data = []
+        unroll_data = []
+
+        for starting_index in self.starting_indices:
+            initial_data.append(data[starting_index])
+            target_start = starting_index + 1
+            unroll_data.append(
+                data[target_start : target_start + self._unroll_length, :]
+            )
+
+        self.input = torch.stack(initial_data)  # [num_unrolls, features]
+        self.unroll = torch.stack(
+            unroll_data
+        )  # [num_unrolls, unroll_length, features]
+
+    @property
+    def unroll_length(self):
+        return self._unroll_length
+
+    @property
+    def time_weighting(self):
+        return self._time_weighting
+
+    @property
+    def max_start_idx(self):
+        max_start_idx = self._n_timesteps - self._unroll_length
+        assert max_start_idx > 0, "Provided data sequence too short"
+        return max_start_idx
+
+    @property
+    def num_unrolls(self):
+        if self._requested_num_unrolls is None:
+            return self.max_start_idx
+        else:
+            assert (
+                self._requested_num_unrolls < self.max_start_idx
+            ), "too many samples requested"
+            return self._requested_num_unrolls
+
+    @property
+    def starting_indices(self):
+        all_starting_indices = torch.arange(self.max_start_idx)
+
+        if self._randomize:
+            perm = torch.randperm(len(all_starting_indices))
+            return all_starting_indices[perm[: self.num_unrolls]]
+        else:
+            selected_indices = torch.linspace(
+                0, len(all_starting_indices) - 1, self.num_unrolls
+            ).long()
+            return all_starting_indices[selected_indices]

pina/loss/__init__.py

@@ -9,6 +9,10 @@
     "NeuralTangentKernelWeighting",
     "SelfAdaptiveWeighting",
     "LinearWeighting",
+    "TimeWeightingInterface",
+    "ConstantTimeWeighting",
+    "ExponentialTimeWeighting",
+    "LinearTimeWeighting",
 ]
 
 from .loss_interface import LossInterface
@@ -19,3 +23,9 @@
 from .ntk_weighting import NeuralTangentKernelWeighting
 from .self_adaptive_weighting import SelfAdaptiveWeighting
 from .linear_weighting import LinearWeighting
+from .time_weighting_interface import TimeWeightingInterface
+from .time_weighting import (
+    ConstantTimeWeighting,
+    ExponentialTimeWeighting,
+    LinearTimeWeighting,
+)

pina/loss/time_weighting.py

@@ -0,0 +1,57 @@
+"""Module for the Time Weighting."""
+
+import torch
+from .time_weighting_interface import TimeWeightingInterface
+
+
+class ConstantTimeWeighting(TimeWeightingInterface):
+    """
+    Weighting scheme that assigns equal weight to all time steps.
+    """
+
+    def __call__(self, num_steps, device):
+        return torch.ones(num_steps, device=device) / num_steps
+
+
+class ExponentialTimeWeighting(TimeWeightingInterface):
+    """
+    Weighting scheme change exponentially with time.
+    gamma > 1.0: increasing weights
+    0 < gamma < 1.0: decreasing weights
+    weight at time t is gamma^t
+    """
+
+    def __init__(self, gamma=0.9):
+        """
+        Initialization of the :class:`ExponentialTimeWeighting` class.
+        :param float gamma: The decay factor. Default is 0.9.
+        """
+        self.gamma = gamma
+
+    def __call__(self, num_steps, device):
+        steps = torch.arange(num_steps, device=device, dtype=torch.float32)
+        weights = self.gamma**steps
+        return weights / weights.sum()
+
+
+class LinearTimeWeighting(TimeWeightingInterface):
+    """
+    Weighting scheme that changes linearly from a start weight to an end weight.
+    """
+
+    def __init__(self, start=0.1, end=1.0):
+        """
+        Initialization of the :class:`LinearDecayTimeWeighting` class.
+
+        :param float start: The starting weight. Default is 0.1.
+        :param float end: The ending weight. Default is 1.0.
+        """
+        self.start = start
+        self.end = end
+
+    def __call__(self, num_steps, device):
+        if num_steps == 1:
+            return torch.ones(1, device=device)
+
+        weights = torch.linspace(self.start, self.end, num_steps, device=device)
+        return weights / weights.sum()

pina/loss/time_weighting_interface.py

@@ -0,0 +1,24 @@
+"""Module for the Time Weighting Interface."""
+
+from abc import ABCMeta, abstractmethod
+import torch
+
+
+class TimeWeightingInterface(metaclass=ABCMeta):
+    """
+    Abstract base class for all time weighting schemas. All time weighting
+    schemas should inherit from this class.
+    """
+
+    @abstractmethod
+    def __call__(self, num_steps, device):
+        """
+        Compute the weights for the time steps.
+
+        :param int num_steps: The number of time steps.
+        :param torch.device device: The device on which the weights should be
+            created.
+        :return: The weights for the time steps.
+        :rtype: torch.Tensor
+        """
+        pass

pina/solver/__init__.py

@@ -18,6 +18,7 @@
     "DeepEnsembleSupervisedSolver",
     "DeepEnsemblePINN",
     "GAROM",
+    "AutoregressiveSolver",
 ]
 
 from .solver import SolverInterface, SingleSolverInterface, MultiSolverInterface
@@ -41,3 +42,7 @@
     DeepEnsemblePINN,
 )
 from .garom import GAROM
+from .autoregressive_solver import (
+    AutoregressiveSolver,
+    AutoregressiveSolverInterface,
+)

pina/solver/autoregressive_solver/__init__.py

@@ -0,0 +1,4 @@
+__all__ = ["AutoregressiveSolver", "AutoregressiveSolverInterface"]
+
+from .autoregressive_solver import AutoregressiveSolver
+from .autoregressive_solver_interface import AutoregressiveSolverInterface

pina/solver/autoregressive_solver/autoregressive_solver.py

@@ -0,0 +1,88 @@
+import torch
+from torch.nn.modules.loss import _Loss
+
+from pina.utils import check_consistency
+from pina.solver.solver import SingleSolverInterface
+from pina.condition import AutoregressiveCondition
+from pina.loss import (
+    LossInterface,
+    TimeWeightingInterface,
+    ConstantTimeWeighting,
+)
+from .autoregressive_solver_interface import AutoregressiveSolverInterface
+
+
+class AutoregressiveSolver(
+    AutoregressiveSolverInterface, SingleSolverInterface
+):
+    """
+    Autoregressive Solver class.
+    """
+
+    accepted_conditions_types = AutoregressiveCondition
+
+    def __init__(
+        self,
+        problem,
+        model,
+        loss=None,
+        optimizer=None,
+        scheduler=None,
+        weighting=None,
+        use_lt=False,
+    ):
+        """
+        Initialization of the :class:`AutoregressiveSolver` class.
+        """
+        super().__init__(
+            problem=problem,
+            model=model,
+            loss=loss,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            weighting=weighting,
+            use_lt=use_lt,
+        )
+
+    def loss_data(self, input, target, unroll_length, time_weighting):
+        """
+        Compute the data loss for the recursive autoregressive solver.
+        This will be applied to each condition individually.
+        """
+        steps_to_predict = unroll_length - 1
+        # weights are passed from the condition
+        weights = time_weighting(steps_to_predict, device=input.device)
+
+        total_loss = 0.0
+        current_state = input
+
+        for step in range(steps_to_predict):
+
+            predicted_next_state = self.forward(
+                current_state
+            )  # [batch_size, features]
+            actual_next_state = target[:, step, :]  # [batch_size, features]
+
+            step_loss = self.loss(predicted_next_state, actual_next_state)
+
+            total_loss += step_loss * weights[step]
+
+            current_state = predicted_next_state.detach()
+
+        return total_loss
+
+    def predict(self, initial_state, num_steps):
+        """
+        Make recursive predictions starting from an initial state.
+        """
+        self.eval()  # Set model to evaluation mode
+
+        current_state = initial_state
+        predictions = [current_state]  # Store initial state without batch dim
+        with torch.no_grad():
+            for step in range(num_steps):
+                next_state = self.forward(current_state)
+                predictions.append(next_state)  # Keep batch dim for storage
+                current_state = next_state
+
+        return torch.stack(predictions)