feat : a general probabilistic forecasting interface

CHANGELOG.md

@@ -9,6 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add a general probabilistic forecasting interface (`ProbabilisticForecaster`,
+  `ProbabilisticARForecaster`, `ProbabilisticForecasterModule`) and move
+  ownership of the training objective, scoring rule and per-variable std
+  from `ForecasterModule` onto the `Forecaster`.
+  [\#685](https://github.com/mllam/neural-lam/issues/685)
+  @Sir-Sloth-The-Lazy
+
 - Add `PropagationNet` GNN layer that incentivises directional message
   propagation from sender to receiver nodes, and expose it alongside
   `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`,

neural_lam/models/__init__.py

@@ -3,7 +3,12 @@
 # Local
 from .forecasters.autoregressive import ARForecaster
 from .forecasters.base import Forecaster
+from .forecasters.probabilistic import (
+    ProbabilisticARForecaster,
+    ProbabilisticForecaster,
+)
 from .module import ForecasterModule
+from .probabilistic_module import ProbabilisticForecasterModule
 from .step_predictors.base import StepPredictor
 from .step_predictors.graph.base import BaseGraphModel
 from .step_predictors.graph.graph_lam import GraphLAM

neural_lam/models/forecasters/__init__.py

@@ -5,3 +5,4 @@
 # Local
 from .autoregressive import ARForecaster
 from .base import Forecaster
+from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster

neural_lam/models/forecasters/autoregressive.py

@@ -1,12 +1,13 @@
 """Forecaster that uses an auto-regressive strategy to unroll a forecast."""
 
-# Standard library
-
 # Third-party
 import torch
 
 # Local
+from ... import metrics
+from ...config import NeuralLAMConfig
 from ...datastore import BaseDatastore
+from ...loss_weighting import get_state_feature_weighting
 from ..step_predictors.base import StepPredictor
 from .base import Forecaster
 
@@ -18,7 +19,11 @@ class ARForecaster(Forecaster):
     """
 
     def __init__(
-        self, predictor: StepPredictor, datastore: BaseDatastore
+        self,
+        predictor: StepPredictor,
+        datastore: BaseDatastore,
+        config: NeuralLAMConfig | None = None,
+        loss: str = "wmse",
     ) -> None:
         """
         Initialize the ARForecaster.
@@ -29,6 +34,14 @@ def __init__(
             The predictor to use for each step.
         datastore : BaseDatastore
             The datastore providing grid metadata and boundary masks.
+        config : NeuralLAMConfig or None
+            Configuration used to compute the constant per-variable std
+            substituted for ``pred_std`` when ``predictor`` does not output
+            its own (see ``per_var_std``). Only required for that case;
+            forecasters used purely for inference can omit it.
+        loss : str, default "wmse"
+            The scoring rule (from ``neural_lam.metrics``) used by
+            ``compute_training_loss`` and stored as ``self.loss``.
         """
         super().__init__()
         self.predictor = predictor
@@ -44,6 +57,31 @@ def __init__(
             "interior_mask", 1.0 - self.boundary_mask, persistent=False
         )
 
+        self.loss = metrics.get_metric(loss)
+
+        # Store per_var_std here if the predictor does not output its own std
+        if not self.predicts_std and config is not None:
+            da_state_stats = datastore.get_standardization_dataarray(
+                category="state"
+            )
+            state_feature_weights = get_state_feature_weighting(
+                config=config, datastore=datastore
+            )
+            diff_std = torch.tensor(
+                da_state_stats.state_diff_std_standardized.values,
+                dtype=torch.float32,
+            )
+            feature_weights_t = torch.tensor(
+                state_feature_weights, dtype=torch.float32
+            )
+            self.register_buffer(
+                "per_var_std",
+                diff_std / torch.sqrt(feature_weights_t),
+                persistent=False,
+            )
+        else:
+            self.per_var_std = None
+
     @property
     def predicts_std(self) -> bool:
         """
@@ -137,10 +175,73 @@ def forward(
 
         prediction = torch.stack(prediction_list, dim=1)
         # If predictor outputs std, stack it; otherwise return None so
-        # ForecasterModule can substitute the constant per_var_std
+        # callers can substitute the constant per_var_std
         if pred_std_list:
             pred_std = torch.stack(pred_std_list, dim=1)
         else:
             pred_std = None
 
         return prediction, pred_std
+
+    def compute_training_loss(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        target_states: torch.Tensor,
+        interior_mask_bool: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        Score the deterministic rollout with ``self.loss``.
+
+        Unrolls a single forecast over the full rollout, scores it against
+        the target states on interior nodes and averages over batch and
+        time.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the rollout length,
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_forcing_vars`` is the forcing feature dimension (already
+            concatenated past/current/future windows).
+        target_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            states at each predicted step, used both as the prediction
+            targets and to overwrite boundary nodes during the rollout.
+            Dims: same as the prediction.
+        interior_mask_bool : torch.Tensor
+            Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior
+            nodes; passed as ``mask`` to ``self.loss`` so that only interior
+            nodes are scored.
+
+        Returns
+        -------
+        batch_loss : torch.Tensor
+            Scalar. The scoring rule applied to the rollout, averaged over
+            batch and time.
+        loss_components : dict of {str: torch.Tensor}
+            Empty; the deterministic objective has no separate components.
+        """
+        prediction, pred_std = self(
+            init_states, forcing_features, target_states
+        )
+        if pred_std is None:
+            pred_std = self.per_var_std
+
+        batch_loss = torch.mean(
+            self.loss(
+                prediction,
+                target_states,
+                pred_std,
+                mask=interior_mask_bool,
+            )
+        )
+        return batch_loss, {}

neural_lam/models/forecasters/base.py

@@ -79,3 +79,58 @@ def forward(
             per-variable std is substituted upstream by
             ``ForecasterModule``. Dims: same as ``prediction``.
         """
+
+    @abstractmethod
+    def compute_training_loss(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        target_states: torch.Tensor,
+        interior_mask_bool: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        Compute the training objective for one batch.
+
+        The forecaster owns its complete training objective: which forecasts
+        to produce from the batch, which loss terms to compute from them and
+        how to combine those terms into a single scalar, using its own
+        ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The
+        wrapping ``ForecasterModule`` only injects the interior mask, logs
+        the returned components and optimizes the returned loss.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the rollout length,
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_forcing_vars`` is the forcing feature dimension (already
+            concatenated past/current/future windows).
+        target_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            states at each predicted step, used both as the prediction
+            targets and to overwrite boundary nodes during forecasting.
+            Dims: same as the prediction.
+        interior_mask_bool : torch.Tensor
+            Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior
+            nodes; passed as ``mask`` to ``self.loss`` so that only interior
+            nodes are scored.
+
+        Returns
+        -------
+        batch_loss : torch.Tensor
+            Scalar. The full training loss for the batch, to take gradients
+            of.
+        loss_components : dict of {str: torch.Tensor}
+            Scalar loss-related quantities to log alongside the loss, keyed
+            by component name. The wrapping module prefixes the names with
+            the training phase. Empty when the objective has no separate
+            components worth logging.
+        """