feat: Graph-EFM ensemble forecasting model

.gitignore

@@ -77,6 +77,7 @@ tags
 # Coc configuration directory
 .vim
 .vscode
+.idea
 
 # macos
 .DS_Store

CHANGELOG.md

@@ -12,6 +12,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add `--train_steps_to_log` CLI option to log training loss for individual unroll steps, and deduplicate common prediction and loss computation steps across loops [\#674](https://github.com/mllam/neural-lam/issues/674) @GiGiKoneti
 
+- 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/config.py

@@ -3,7 +3,7 @@
 # Standard library
 import dataclasses
 from pathlib import Path
-from typing import Dict, Union
+from typing import Dict, Optional, Union
 
 # Third-party
 import dataclass_wizard
@@ -116,6 +116,49 @@ class TrainingConfig:
     )
 
 
+@dataclasses.dataclass
+class ProbabilisticConfig:
+    """
+    Configuration for the probabilistic (Graph-EFM) models.
+
+    Only used when training or evaluating a probabilistic model
+    (``--model graph_efm``/``graph_efm_ms``); ignored otherwise. Every field
+    has a default, so the ``probabilistic`` config section may be omitted
+    entirely for deterministic models.
+
+    Attributes
+    ----------
+    latent_dim : int, optional
+        Dimensionality of the latent variable at each latent-carrying mesh
+        node. Defaults to the model's ``hidden_dim`` when None.
+    prior_layers : int
+        Number of on-mesh GNN layers in the prior.
+    encoder_layers : int
+        Number of on-mesh GNN layers in the variational encoder.
+    decoder_layers : int
+        Number of on-mesh GNN layers in the latent decoder.
+    learn_prior : bool
+        If True, the prior is a learned encoder conditioned on the previous
+        state; if False, a constant ``Normal(0, 1)`` prior is used.
+    prior_dist : str
+        Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``.
+    kl_beta : float
+        Weight of the KL term in the ELBO. When 0, the prior and KL are not
+        computed (pure auto-encoder training).
+    eval_ensemble_size : int
+        Number of ensemble members sampled during validation and testing.
+    """
+
+    latent_dim: Optional[int] = None
+    prior_layers: int = 2
+    encoder_layers: int = 2
+    decoder_layers: int = 4
+    learn_prior: bool = True
+    prior_dist: str = "isotropic"
+    kl_beta: float = 1.0
+    eval_ensemble_size: int = 5
+
+
 @dataclasses.dataclass
 class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
     """
@@ -133,10 +176,16 @@ class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
     training : TrainingConfig
         Configuration for training the model, including loss function and
         feature-weighting strategy. Defaults to ``TrainingConfig()``.
+    probabilistic : ProbabilisticConfig
+        Configuration for the probabilistic (Graph-EFM) models. Defaults to
+        ``ProbabilisticConfig()`` and is ignored by deterministic models.
     """
 
     datastore: DatastoreSelection
     training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig)
+    probabilistic: ProbabilisticConfig = dataclasses.field(
+        default_factory=ProbabilisticConfig
+    )
 
     class _(dataclass_wizard.JSONWizard.Meta):
         """

neural_lam/datastore/mdp.py

@@ -18,7 +18,7 @@
 from numpy import ndarray
 
 # Local
-from ..utils import log_on_rank_zero
+from ..utils.logging import log_on_rank_zero
 from .base import BaseRegularGridDatastore, CartesianGridShape
 
 

neural_lam/models/__init__.py

@@ -3,16 +3,32 @@
 # Local
 from .forecasters.autoregressive import ARForecaster
 from .forecasters.base import Forecaster
-from .module import ForecasterModule
+from .forecasters.graph_efm import GraphEFMForecaster
+from .forecasters.probabilistic import (
+    ProbabilisticARForecaster,
+    ProbabilisticForecaster,
+)
+from .modules.base import BaseForecasterModule
+from .modules.deterministic import DeterministicForecasterModule
+from .modules.probabilistic import ProbabilisticForecasterModule
 from .step_predictors.base import StepPredictor
 from .step_predictors.graph.base import BaseGraphModel
+from .step_predictors.graph.graph_efm import GraphEFM, GraphEFMMultiScale
 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
 
+# Graph-EFM models are probabilistic: train_model.py builds them with a
+# config-aware, probabilistic assembly path (GraphEFMForecaster wrapped in a
+# ProbabilisticForecasterModule), distinct from the deterministic models
+# above. ``PROBABILISTIC_MODELS`` marks which entries take that path.
 MODELS = {
     "graph_lam": GraphLAM,
     "hi_lam": HiLAM,
     "hi_lam_parallel": HiLAMParallel,
+    "graph_efm": GraphEFM,
+    "graph_efm_ms": GraphEFMMultiScale,
 }
+
+PROBABILISTIC_MODELS = {"graph_efm", "graph_efm_ms"}

neural_lam/models/forecasters/__init__.py

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