Add PersistencePredictor baseline model

CHANGELOG.md

@@ -12,6 +12,8 @@ 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 `PersistencePredictor` baseline model (`--model persistence`) that returns the previous state unchanged, enabling standard evaluation of a persistence baseline through the existing pipeline [\#677](https://github.com/mllam/neural-lam/pull/677) @GiGiKoneti
+
 - 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`,

README.md

@@ -481,6 +481,9 @@ To train Hi-LAM-Parallel use
 python -m neural_lam.train_model --model hi_lam_parallel --graph hierarchical ...
 ```
 
+### Persistence
+A trivial baseline model that returns the previous state unchanged at each step (`--model persistence`). Since it has no learnable parameters, it must be evaluated directly without training (e.g. `--eval val` or `--eval test`).
+
 Checkpoint files for our models trained on the MEPS data are available upon request.
 
 ### High Performance Computing

neural_lam/models/__init__.py

@@ -10,9 +10,11 @@
 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.persistence import PersistencePredictor
 
 MODELS = {
     "graph_lam": GraphLAM,
     "hi_lam": HiLAM,
     "hi_lam_parallel": HiLAMParallel,
+    "persistence": PersistencePredictor,
 }

neural_lam/models/step_predictors/base.py

@@ -18,6 +18,19 @@ class StepPredictor(nn.Module, ABC):
     time steps plus forcing into a prediction of the next state.
     """
 
+    @property
+    def trainable(self) -> bool:
+        """
+        Whether the predictor has learnable parameters and can be trained.
+
+        Returns
+        -------
+        bool
+            ``True`` if the predictor has trainable parameters,
+            ``False`` otherwise.
+        """
+        return any(p.requires_grad for p in self.parameters())
+
     def __init__(
         self,
         datastore: BaseDatastore,

neural_lam/models/step_predictors/persistence.py

@@ -0,0 +1,89 @@
+"""Persistence baseline step predictor."""
+
+# Third-party
+import torch
+from loguru import logger
+
+# Local
+from ...datastore import BaseDatastore
+from .base import StepPredictor
+
+
+class PersistencePredictor(StepPredictor):
+    """
+    Trivial baseline predictor that returns the previous state unchanged.
+
+    At each AR step the predicted next state is simply the current state,
+    i.e. ``X_{t+1} = X_t``.  This provides a persistence (climatological
+    no-change) baseline that can be evaluated through the standard
+    ``ForecasterModule`` / ``ARForecaster`` pipeline.
+    """
+
+    def __init__(
+        self,
+        datastore: BaseDatastore,
+        output_std: bool = False,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize the PersistencePredictor.
+
+        Parameters
+        ----------
+        datastore : BaseDatastore
+            The datastore providing grid metadata and data access.
+        output_std : bool, default False
+            Ignored — persistence never predicts uncertainty.
+        output_clamping_lower : dict, optional
+            Ignored — persistence returns the raw previous state.
+        output_clamping_upper : dict, optional
+            Ignored — persistence returns the raw previous state.
+        **kwargs
+            Absorbed so that the standard CLI kwargs (``graph_name``,
+            ``hidden_dim``, etc.) do not cause errors.
+        """
+        if output_std:
+            logger.warning(
+                "Persistence predictor does not support predicting "
+                "standard deviation. The output_std parameter will be ignored."
+            )
+
+        super().__init__(
+            datastore=datastore,
+            output_std=False,
+            output_clamping_lower=None,
+            output_clamping_upper=None,
+        )
+
+    def forward(
+        self,
+        prev_state: torch.Tensor,
+        prev_prev_state: torch.Tensor,
+        forcing: torch.Tensor,
+    ) -> tuple[torch.Tensor, None]:
+        """
+        Return the previous state as the prediction.
+
+        Parameters
+        ----------
+        prev_state : torch.Tensor
+            Shape ``(B, num_grid_nodes, num_state_vars)``.
+            The current state ``X_t``.
+        prev_prev_state : torch.Tensor
+            Shape ``(B, num_grid_nodes, num_state_vars)``.
+            The previous state ``X_{t-1}`` (unused).
+        forcing : torch.Tensor
+            Shape ``(B, num_grid_nodes, num_forcing_vars)``.
+            External forcings (unused).
+
+        Returns
+        -------
+        pred_state : torch.Tensor
+            Shape ``(B, num_grid_nodes, num_state_vars)``.
+            Equal to ``prev_state``.
+        pred_std : None
+            Always ``None`` — persistence does not predict uncertainty.
+        """
+        return prev_state, None

neural_lam/train_model.py

@@ -474,6 +474,13 @@ def main(input_args=None):
         mesh_up_gnn_type=args.mesh_up_gnn_type,
         mesh_down_gnn_type=args.mesh_down_gnn_type,
     )
+
+    if not args.eval and not predictor.trainable:
+        raise ValueError(
+            f"The {args.model} model cannot be trained. Run with "
+            "--eval <val/test> to evaluate a baseline."
+        )
+
     forecaster = ARForecaster(predictor, datastore)
 
     model = ForecasterModule(

tests/test_persistence.py

@@ -0,0 +1,145 @@
+"""Tests for the PersistencePredictor baseline model."""
+
+# Third-party
+import pytest
+import torch
+
+# First-party
+from neural_lam.models import ARForecaster
+from neural_lam.models.step_predictors.persistence import PersistencePredictor
+from tests.conftest import init_datastore_example
+
+
+def test_persistence_predictor_returns_prev_state():
+    """PersistencePredictor.forward returns prev_state unchanged."""
+    datastore = init_datastore_example("mdp")
+    predictor = PersistencePredictor(datastore=datastore)
+    assert predictor.trainable is False
+
+    B = 2
+    num_grid_nodes = predictor.num_grid_nodes
+    d_state = datastore.get_num_data_vars(category="state")
+    d_forcing = datastore.get_num_data_vars(category="forcing")
+
+    prev_state = torch.randn(B, num_grid_nodes, d_state)
+    prev_prev_state = torch.randn(B, num_grid_nodes, d_state)
+    forcing = torch.randn(B, num_grid_nodes, d_forcing)
+
+    pred_state, pred_std = predictor(prev_state, prev_prev_state, forcing)
+
+    assert pred_std is None
+    assert torch.equal(pred_state, prev_state)
+
+
+def test_persistence_predictor_ignores_kwargs():
+    """Extra graph-specific kwargs are silently absorbed."""
+    datastore = init_datastore_example("mdp")
+    predictor = PersistencePredictor(
+        datastore=datastore,
+        graph_name="multiscale",
+        hidden_dim=64,
+        hidden_layers=1,
+        processor_layers=4,
+        mesh_aggr="sum",
+    )
+    assert isinstance(predictor, PersistencePredictor)
+
+
+def test_persistence_forecaster_unroll():
+    """ARForecaster with PersistencePredictor reproduces initial state."""
+    datastore = init_datastore_example("mdp")
+    predictor = PersistencePredictor(datastore=datastore)
+    forecaster = ARForecaster(predictor, datastore)
+
+    B = 2
+    num_grid_nodes = predictor.num_grid_nodes
+    d_state = datastore.get_num_data_vars(category="state")
+    d_forcing = datastore.get_num_data_vars(category="forcing") * 3
+    pred_steps = 4
+
+    init_states = torch.randn(B, 2, num_grid_nodes, d_state)
+    forcing = torch.randn(B, pred_steps, num_grid_nodes, d_forcing)
+    boundary_states = torch.randn(B, pred_steps, num_grid_nodes, d_state)
+
+    prediction, pred_std = forecaster(init_states, forcing, boundary_states)
+
+    assert prediction.shape == (B, pred_steps, num_grid_nodes, d_state)
+    assert pred_std is None
+
+    # Persistence keeps every interior node at init_states[:, 1] for all
+    # steps, while boundary nodes are overwritten with boundary_states.
+    interior_mask = forecaster.interior_mask.squeeze(0).squeeze(-1).bool()
+    boundary_mask = forecaster.boundary_mask.squeeze(0).squeeze(-1).bool()
+    for t in range(pred_steps):
+        assert torch.equal(
+            prediction[:, t, interior_mask, :],
+            init_states[:, 1, interior_mask, :],
+        )
+        assert torch.equal(
+            prediction[:, t, boundary_mask, :],
+            boundary_states[:, t, boundary_mask, :],
+        )
+
+
+def test_persistence_predicts_std_false():
+    """
+    PersistencePredictor.predicts_std is always False and logs warning
+    when output_std is True.
+    """
+    # Standard library
+    from unittest.mock import patch
+
+    datastore = init_datastore_example("mdp")
+    target_patch = (
+        "neural_lam.models.step_predictors.persistence.logger.warning"
+    )
+    with patch(target_patch) as mock_warn:
+        predictor = PersistencePredictor(datastore=datastore, output_std=True)
+        assert not predictor.predicts_std
+        mock_warn.assert_called_once_with(
+            "Persistence predictor does not support predicting "
+            "standard deviation. The output_std parameter will be ignored."
+        )
+
+
+def test_persistence_training_error():
+    """ValueError must be raised if trying to train the persistence model."""
+    # Standard library
+    from unittest.mock import MagicMock, patch
+
+    # First-party
+    from neural_lam.train_model import main
+
+    mock_args = MagicMock()
+    mock_args.eval = None  # training mode
+    mock_args.load = None
+    mock_args.config_path = "dummy.yaml"
+    mock_args.val_steps_to_log = []
+    mock_args.var_leads_metrics_watch = "{}"
+    mock_args.ar_steps_eval = 10
+    mock_args.model = "persistence"
+    mock_args.devices = ["auto"]
+
+    mock_predictor = MagicMock()
+    mock_predictor.trainable = False
+
+    with (
+        patch(
+            "neural_lam.train_model.ArgumentParser.parse_args",
+            return_value=mock_args,
+        ),
+        patch(
+            "neural_lam.train_model.load_config_and_datastore",
+            return_value=(MagicMock(), MagicMock()),
+        ),
+        patch("neural_lam.train_model.WeatherDataModule"),
+        patch(
+            "neural_lam.train_model.MODELS",
+            {"persistence": MagicMock(return_value=mock_predictor)},
+        ),
+        pytest.raises(
+            ValueError,
+            match="The persistence model cannot be trained",
+        ),
+    ):
+        getattr(main, "__wrapped__", main)()

tests/test_train_model_warnings.py

@@ -29,6 +29,7 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning):
     mock_args.train_steps_to_log = []
     mock_args.var_leads_metrics_watch = "{}"
     mock_args.ar_steps_eval = 10
+    mock_args.model = "graph_lam"
     mock_args.ar_steps_train = 10
 
     with patch(