diff --git a/CHANGELOG.md b/CHANGELOG.md index 62dcb764..8fbe4baa 100644 --- a/CHANGELOG.md +++ b/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`, diff --git a/README.md b/README.md index 35d029a6..1101cf01 100644 --- a/README.md +++ b/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 diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..d2bb3485 100644 --- a/neural_lam/models/__init__.py +++ b/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, } diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py index daf8b0e6..f1404c55 100644 --- a/neural_lam/models/step_predictors/base.py +++ b/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, diff --git a/neural_lam/models/step_predictors/persistence.py b/neural_lam/models/step_predictors/persistence.py new file mode 100644 index 00000000..129d220a --- /dev/null +++ b/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 diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index 2a9775c9..851de698 100644 --- a/neural_lam/train_model.py +++ b/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 to evaluate a baseline." + ) + forecaster = ARForecaster(predictor, datastore) model = ForecasterModule( diff --git a/tests/test_persistence.py b/tests/test_persistence.py new file mode 100644 index 00000000..b4258e8d --- /dev/null +++ b/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)() diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index eb4c2bc0..838b820e 100644 --- a/tests/test_train_model_warnings.py +++ b/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(