feat: add probabilistic forecaster interface (decoupled from Graph-EFM)

neural_lam/metrics.py

@@ -387,6 +387,80 @@ def crps_gauss(
     )
 
 
+def crps_ens(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: Optional[torch.Tensor] = None,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+    ens_dim: int = 1,
+) -> torch.Tensor:
+    """
+    (Negative) Continuous Ranked Probability Score from ensemble samples.
+
+    Unbiased ("fair") estimator of the CRPS computed directly from a set of
+    ensemble members, following Zamo and Naveau / WeatherBench 2.
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Shape ``(..., M, ..., N, num_variables)``. Ensemble predictions, with
+        the ``M`` ensemble members laid out along ``ens_dim``.
+    target : torch.Tensor
+        Shape ``(..., N, num_variables)``. Ground-truth target (no ensemble
+        dimension).
+    pred_std : torch.Tensor or None, optional
+        Unused; accepted so the signature matches the other metrics.
+    mask : torch.Tensor or None, optional
+        Shape ``(N,)``. Boolean mask selecting grid nodes. ``None`` uses all.
+    average_grid : bool, optional
+        If True, average over the grid dimension ``N``.
+    sum_vars : bool, optional
+        If True, sum over the variable dimension ``num_variables``.
+    ens_dim : int, optional
+        Batch dimension along which the ensemble members are laid out and
+        reduced. Default ``1``.
+
+    Returns
+    -------
+    torch.Tensor
+        Reduced CRPS values; shape depends on ``average_grid`` and
+        ``sum_vars``.
+    """
+    del pred_std  # unused; present for a uniform metric signature
+    num_ens = pred.shape[ens_dim]
+    if num_ens == 1:
+        # With a single sample the CRPS reduces to the MAE
+        return mae(
+            pred.squeeze(ens_dim),
+            target,
+            None,
+            mask=mask,
+            average_grid=average_grid,
+            sum_vars=sum_vars,
+        )
+
+    mean_mae = torch.mean(
+        torch.abs(pred - target.unsqueeze(ens_dim)), dim=ens_dim
+    )  # (..., N, num_variables)
+    if num_ens == 2:
+        # Cheap exact estimator for a two-member ensemble
+        pair_diffs_term = -0.5 * torch.abs(
+            pred.select(ens_dim, 0) - pred.select(ens_dim, 1)
+        )
+    else:
+        # Rank-based estimator, O(M log M) compute and O(M) memory.
+        # Ranks start at 1; two argsorts give the per-entry ranks.
+        ranks = pred.argsort(dim=ens_dim).argsort(ens_dim) + 1
+        pair_diffs_term = (1 / (num_ens - 1)) * torch.mean(
+            (num_ens + 1 - 2 * ranks) * pred, dim=ens_dim
+        )
+    crps_estimator = mean_mae + pair_diffs_term  # (..., N, num_variables)
+
+    return mask_and_reduce_metric(crps_estimator, mask, average_grid, sum_vars)
+
+
 DEFINED_METRICS = {
     "mse": mse,
     "mae": mae,

neural_lam/models/__init__.py

@@ -3,14 +3,24 @@
 # Local
 from .forecasters.autoregressive import ARForecaster
 from .forecasters.base import Forecaster
+from .forecasters.probabilistic import ProbabilisticARForecaster
 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_efm import GraphEFM, GraphEFMMS
 from .step_predictors.graph.graph_lam import GraphLAM
 from .step_predictors.graph.hi_lam import HiLAM
 from .step_predictors.graph.hi_lam_parallel import HiLAMParallel
 from .step_predictors.graph.hierarchical import BaseHiGraphModel
+from .step_predictors.probabilistic import LatentStepPredictor
 
+# NOTE: GraphEFM/GraphEFMMS are intentionally NOT registered in MODELS yet.
+# The shared construction call in train_model.py instantiates the chosen
+# model with a fixed deterministic kwarg set -- datastore-first, no
+# ``config``, and with ``mesh_aggr`` -- whereas the Graph-EFM models require
+# ``config`` (for their per_var_std weighting) and take no ``mesh_aggr``.
+# Registering them requires config-aware model assembly in train_model.py.
 MODELS = {
     "graph_lam": GraphLAM,
     "hi_lam": HiLAM,

neural_lam/models/forecasters/probabilistic.py

@@ -0,0 +1,156 @@
+"""Autoregressive forecaster for Graph-EFM latent-variable predictors."""
+
+# Standard library
+from typing import Optional
+
+# Third-party
+import torch
+
+# Local
+from ..step_predictors.probabilistic import LatentStepPredictor
+from .autoregressive import ARForecaster
+
+
+class ProbabilisticARForecaster(ARForecaster):
+    """
+    Autoregressive forecaster for latent-variable step predictors.
+
+    The inherited :meth:`ARForecaster.forward` is the evaluation rollout: each
+    step calls the predictor (which samples its prior), conditions
+    autoregressively on the predicted state, and overwrites boundary nodes
+    with the true value. This subclass adds :meth:`training_rollout`, the
+    posterior-conditioned rollout used for training: at each step it asks the
+    predictor for the prior, the variational posterior (conditioned on the
+    true target) and the decoded reconstruction, autoregresses on the
+    predicted mean, and reduces the per-step KL between posterior and prior.
+    The KL is purely between the model's own distributions, so it is reduced
+    here; the likelihood and the ELBO assembly are left to the module.
+    """
+
+    # The wrapped predictor must expose per-step latent distributions
+    predictor: LatentStepPredictor
+
+    def training_rollout(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        target_states: torch.Tensor,
+        compute_kl: bool = True,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        Posterior-conditioned rollout producing the per-step ELBO components.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, d_state)``. ``[X_{t-1}, X_t]``.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``.
+        target_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, d_state)``. The true next
+            states, used to condition the posterior and to overwrite the
+            boundary nodes.
+        compute_kl : bool, optional
+            If ``False`` (KL weight zero), skip the prior and return
+            ``kl_terms`` as ``None``.
+
+        Returns
+        -------
+        pred_means : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, d_state)``. Raw decoder
+            means (not boundary-overwritten), for the likelihood.
+        pred_stds : torch.Tensor or None
+            Shape ``(B, pred_steps, num_grid_nodes, d_state)`` when the
+            predictor emits its own std, otherwise ``None``.
+        kl_terms : torch.Tensor or None
+            Shape ``(B, pred_steps)``. Per-step KL of posterior from prior,
+            or ``None`` when ``compute_kl`` is ``False``.
+        """
+        prev_prev_state = init_states[:, 0]
+        prev_state = init_states[:, 1]
+        mean_list = []
+        std_list = []
+        kl_list = []
+
+        for i in range(forcing_features.shape[1]):
+            target = target_states[:, i]
+            prior_dist, posterior_dist, pred_mean, pred_std = (
+                self.predictor.step_distributions(
+                    prev_state,
+                    prev_prev_state,
+                    forcing_features[:, i],
+                    target=target,
+                    compute_prior=compute_kl,
+                )
+            )
+
+            mean_list.append(pred_mean)  # raw, masked to interior later
+            if pred_std is not None:
+                std_list.append(pred_std)
+            if posterior_dist is not None and prior_dist is not None:
+                kl = torch.distributions.kl_divergence(
+                    posterior_dist, prior_dist
+                )
+                # Sum over all latent dimensions, leaving the batch dim
+                kl_list.append(kl.sum(dim=tuple(range(1, kl.ndim))))
+
+            # Autoregress on the predicted mean; boundary overwritten by truth
+            new_state = (
+                self.boundary_mask * target + self.interior_mask * pred_mean
+            )
+            prev_prev_state = prev_state
+            prev_state = new_state
+
+        pred_means = torch.stack(mean_list, dim=1)
+        pred_stds = torch.stack(std_list, dim=1) if std_list else None
+        kl_terms = torch.stack(kl_list, dim=1) if kl_list else None
+        return pred_means, pred_stds, kl_terms
+
+    def sample_trajectories(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        boundary_states: torch.Tensor,
+        num_traj: int,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Sample an ensemble of prior trajectories.
+
+        Runs the evaluation rollout (:meth:`ARForecaster.forward`, which
+        samples the prior each step) ``num_traj`` times and stacks the
+        members along a new ensemble dimension. Used for the CRPS term and
+        for ensemble evaluation.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, d_state)``. ``[X_{t-1}, X_t]``.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``.
+        boundary_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, d_state)``. True states
+            used to overwrite boundary nodes at each step.
+        num_traj : int
+            Number of trajectories ``S`` to sample.
+
+        Returns
+        -------
+        traj_means : torch.Tensor
+            Shape ``(B, S, pred_steps, num_grid_nodes, d_state)``.
+        traj_stds : torch.Tensor or None
+            Shape ``(B, S, pred_steps, num_grid_nodes, d_state)`` when the
+            predictor emits its own std, otherwise ``None``.
+        """
+        mean_list = []
+        std_list = []
+        for _ in range(num_traj):
+            prediction, pred_std = self(
+                init_states, forcing_features, boundary_states
+            )
+            mean_list.append(prediction)
+            if pred_std is not None:
+                std_list.append(pred_std)
+
+        traj_means = torch.stack(mean_list, dim=1)  # (B, S, T, N, d_state)
+        traj_stds = torch.stack(std_list, dim=1) if std_list else None
+        return traj_means, traj_stds

neural_lam/models/latent/__init__.py

@@ -0,0 +1,22 @@
+"""Latent encoder and decoder modules for latent-variable models such as
+GraphEFM, mapping between grid representations and distributions over
+latent variables on mesh nodes."""
+
+# Local
+from .base_decoder import BaseGraphLatentDecoder
+from .base_encoder import BaseLatentEncoder
+from .constant_encoder import ConstantLatentEncoder
+from .graph_decoder import GraphLatentDecoder
+from .graph_encoder import GraphLatentEncoder
+from .hi_graph_decoder import HiGraphLatentDecoder
+from .hi_graph_encoder import HiGraphLatentEncoder
+
+__all__ = [
+    "BaseGraphLatentDecoder",
+    "BaseLatentEncoder",
+    "ConstantLatentEncoder",
+    "GraphLatentDecoder",
+    "GraphLatentEncoder",
+    "HiGraphLatentDecoder",
+    "HiGraphLatentEncoder",
+]