From 0bd45c55ed9534765a8c37e28a58cf0e23b9aa77 Mon Sep 17 00:00:00 2001 From: GiGiKoneti Date: Sun, 21 Jun 2026 02:15:43 +0530 Subject: [PATCH 1/6] Add PersistencePredictor baseline model --- CHANGELOG.md | 5 +- neural_lam/models/__init__.py | 2 + .../models/step_predictors/persistence.py | 83 ++++++++++++++ neural_lam/train_model.py | 5 + tests/test_persistence.py | 104 ++++++++++++++++++ 5 files changed, 198 insertions(+), 1 deletion(-) create mode 100644 neural_lam/models/step_predictors/persistence.py create mode 100644 tests/test_persistence.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ccf79fdb..cfc1a74f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Add `PersistencePredictor` baseline model (`--model persistence`) that returns the previous state unchanged, enabling standard evaluation of a persistence baseline through the existing pipeline [\#676](https://github.com/mllam/neural-lam/issues/676) @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`, @@ -112,12 +114,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Explicitly cleanup the temporary directory in `DummyDatastore` by adding a `__del__` method, preventing `ResourceWarning` during garbage collection in tests [\#487](https://github.com/mllam/neural-lam/pull/487) @sohampatil01-svg -- Add comprehensive type hints to all functions and class methods in `utils.py` [\#620](https://github.com/mllam/neural-lam/pull/620) @GiGiKoneti - Add probabilistic objective regression coverage for weighted losses and `pred_std` broadcasting semantics [\#504](https://github.com/mllam/neural-lam/pull/504) @kshirajahere - Add comprehensive type hints to `neural_lam/create_graph.py` [\#618](https://github.com/mllam/neural-lam/pull/618) @GiGiKoneti +- Add comprehensive type hints to all functions and class methods in `utils.py` [\#620](https://github.com/mllam/neural-lam/pull/620) @GiGiKoneti + - Select the torch build via mutually-exclusive `cpu`, `gpu` (CUDA 13.0) and `gpu-cu128` (CUDA 12.8) extras routed through `[tool.uv.sources]`, with torch versions pinned per CUDA build and a committed `uv.lock`. CI now installs and 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/persistence.py b/neural_lam/models/step_predictors/persistence.py new file mode 100644 index 00000000..a66709c3 --- /dev/null +++ b/neural_lam/models/step_predictors/persistence.py @@ -0,0 +1,83 @@ +"""Persistence baseline step predictor.""" + +# Third-party +import torch +from torch import nn + +# 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. + """ + 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 f98065c4..003c9072 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -355,6 +355,11 @@ def main(input_args=None): ), ) args = parser.parse_args(input_args) + if not args.eval and args.model == "persistence": + raise ValueError( + "The persistence model cannot be trained. Run with " + "--eval to evaluate a persistence baseline." + ) args.var_leads_metrics_watch = { int(k): v for k, v in json.loads(args.var_leads_metrics_watch).items() } diff --git a/tests/test_persistence.py b/tests/test_persistence.py new file mode 100644 index 00000000..a11c0e0c --- /dev/null +++ b/tests/test_persistence.py @@ -0,0 +1,104 @@ +"""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) + + 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 + + # Interior nodes should equal init_states[:, 1] (persistence), + # boundary nodes should equal boundary_states + interior_mask = forecaster.interior_mask.squeeze(0).squeeze(-1).bool() + for t in range(pred_steps): + interior_pred = prediction[:, t, interior_mask, :] + # After first step interior is init_states[:, 1], then persists + interior_boundary = boundary_states[:, t, interior_mask, :] + # Verify interior nodes are NOT equal to boundary_states + # (they should be the persisted initial state) + assert interior_pred.shape == interior_boundary.shape + + +def test_persistence_predicts_std_false(): + """PersistencePredictor.predicts_std is always False.""" + datastore = init_datastore_example("mdp") + predictor = PersistencePredictor(datastore=datastore, output_std=True) + assert not predictor.predicts_std + + +def test_persistence_training_error(): + """ValueError must be raised if trying to train the persistence model.""" + from unittest.mock import MagicMock, patch + from neural_lam.train_model import main + + mock_args = MagicMock() + mock_args.eval = None # training mode + mock_args.model = "persistence" + + with patch( + "neural_lam.train_model.ArgumentParser.parse_args", + return_value=mock_args, + ): + with pytest.raises( + ValueError, + match="The persistence model cannot be trained", + ): + getattr(main, "__wrapped__", main)() From e4174ab67533a2c3c991213950a93ebbf2005421 Mon Sep 17 00:00:00 2001 From: GiGiKoneti Date: Sun, 21 Jun 2026 02:20:51 +0530 Subject: [PATCH 2/6] Fix imports and pre-commit checks for persistence predictor --- neural_lam/models/step_predictors/persistence.py | 1 - tests/test_persistence.py | 3 +++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/neural_lam/models/step_predictors/persistence.py b/neural_lam/models/step_predictors/persistence.py index a66709c3..9873e6e8 100644 --- a/neural_lam/models/step_predictors/persistence.py +++ b/neural_lam/models/step_predictors/persistence.py @@ -2,7 +2,6 @@ # Third-party import torch -from torch import nn # Local from ...datastore import BaseDatastore diff --git a/tests/test_persistence.py b/tests/test_persistence.py index a11c0e0c..9a6d3c4e 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -86,7 +86,10 @@ def test_persistence_predicts_std_false(): 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() From 5c5998713ca22cce3124bbe2ba83d2a4fdcfc5e2 Mon Sep 17 00:00:00 2001 From: GiGiKoneti Date: Mon, 22 Jun 2026 15:37:01 +0530 Subject: [PATCH 3/6] Address review feedback --- CHANGELOG.md | 5 ++--- README.md | 3 +++ neural_lam/models/step_predictors/base.py | 3 +++ .../models/step_predictors/persistence.py | 3 +++ neural_lam/train_model.py | 13 ++++++++++--- tests/test_persistence.py | 19 +++++++++++-------- 6 files changed, 32 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cfc1a74f..c4600c7f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add `PersistencePredictor` baseline model (`--model persistence`) that returns the previous state unchanged, enabling standard evaluation of a persistence baseline through the existing pipeline [\#676](https://github.com/mllam/neural-lam/issues/676) @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 @@ -114,13 +114,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Explicitly cleanup the temporary directory in `DummyDatastore` by adding a `__del__` method, preventing `ResourceWarning` during garbage collection in tests [\#487](https://github.com/mllam/neural-lam/pull/487) @sohampatil01-svg +- Add comprehensive type hints to all functions and class methods in `utils.py` [\#620](https://github.com/mllam/neural-lam/pull/620) @GiGiKoneti - Add probabilistic objective regression coverage for weighted losses and `pred_std` broadcasting semantics [\#504](https://github.com/mllam/neural-lam/pull/504) @kshirajahere - Add comprehensive type hints to `neural_lam/create_graph.py` [\#618](https://github.com/mllam/neural-lam/pull/618) @GiGiKoneti -- Add comprehensive type hints to all functions and class methods in `utils.py` [\#620](https://github.com/mllam/neural-lam/pull/620) @GiGiKoneti - - Select the torch build via mutually-exclusive `cpu`, `gpu` (CUDA 13.0) and `gpu-cu128` (CUDA 12.8) extras routed through `[tool.uv.sources]`, with torch versions pinned per CUDA build and a committed `uv.lock`. CI now installs and diff --git a/README.md b/README.md index 68f7faa8..aca32514 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/step_predictors/base.py b/neural_lam/models/step_predictors/base.py index daf8b0e6..dff11b21 100644 --- a/neural_lam/models/step_predictors/base.py +++ b/neural_lam/models/step_predictors/base.py @@ -18,6 +18,9 @@ class StepPredictor(nn.Module, ABC): time steps plus forcing into a prediction of the next state. """ + trainable: bool = True + """Whether the predictor has learnable parameters and can be trained.""" + def __init__( self, datastore: BaseDatastore, diff --git a/neural_lam/models/step_predictors/persistence.py b/neural_lam/models/step_predictors/persistence.py index 9873e6e8..bca317a9 100644 --- a/neural_lam/models/step_predictors/persistence.py +++ b/neural_lam/models/step_predictors/persistence.py @@ -18,6 +18,9 @@ class PersistencePredictor(StepPredictor): ``ForecasterModule`` / ``ARForecaster`` pipeline. """ + trainable: bool = False + """Persistence predictors have no learnable parameters.""" + def __init__( self, datastore: BaseDatastore, diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index 003c9072..b127ec47 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -355,10 +355,10 @@ def main(input_args=None): ), ) args = parser.parse_args(input_args) - if not args.eval and args.model == "persistence": + if not args.eval and not MODELS[args.model].trainable: raise ValueError( - "The persistence model cannot be trained. Run with " - "--eval to evaluate a persistence baseline." + f"The {args.model} model cannot be trained. Run with " + "--eval to evaluate a baseline." ) args.var_leads_metrics_watch = { int(k): v for k, v in json.loads(args.var_leads_metrics_watch).items() @@ -462,6 +462,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 args.output_std and not predictor.predicts_std: + logger.warning( + f"Model '{args.model}' does not support predicting " + "standard deviation. The --output_std flag will be ignored." + ) + forecaster = ARForecaster(predictor, datastore) model = ForecasterModule( diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 9a6d3c4e..bdfafbb0 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -65,16 +65,19 @@ def test_persistence_forecaster_unroll(): assert prediction.shape == (B, pred_steps, num_grid_nodes, d_state) assert pred_std is None - # Interior nodes should equal init_states[:, 1] (persistence), - # boundary nodes should equal boundary_states + # 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): - interior_pred = prediction[:, t, interior_mask, :] - # After first step interior is init_states[:, 1], then persists - interior_boundary = boundary_states[:, t, interior_mask, :] - # Verify interior nodes are NOT equal to boundary_states - # (they should be the persisted initial state) - assert interior_pred.shape == interior_boundary.shape + 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(): From 9a705b0a0003038214b0217370f7f44f1f868485 Mon Sep 17 00:00:00 2001 From: GiGiKoneti Date: Mon, 22 Jun 2026 15:51:26 +0530 Subject: [PATCH 4/6] Fix test_eval_without_load_warning by patching MODELS --- tests/test_train_model_warnings.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index a0b5f92a..454ddbd2 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -24,6 +24,7 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): mock_args.val_steps_to_log = [] mock_args.var_leads_metrics_watch = "{}" mock_args.ar_steps_eval = 10 + mock_args.model = "graph_lam" with patch( "neural_lam.train_model.ArgumentParser.parse_args", @@ -34,13 +35,21 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): side_effect=SystemExit(0), ): with patch("neural_lam.train_model.logger.warning") as mock_warning: - with pytest.raises(SystemExit): - main() - if expect_warning: - mock_warning.assert_called_once() - assert "--load" in mock_warning.call_args[0][0] - else: - mock_warning.assert_not_called() + with patch("neural_lam.train_model.WeatherDataModule"): + # Add this patch to prevent further execution + mock_model_class = MagicMock() + mock_model_class.trainable = True + with patch( + "neural_lam.train_model.MODELS", + {"graph_lam": mock_model_class}, + ): + with pytest.raises(SystemExit): + main() + if expect_warning: + mock_warning.assert_called_once() + assert "--load" in mock_warning.call_args[0][0] + else: + mock_warning.assert_not_called() def test_create_gif_forwarded_to_forecaster_module(): From 2a5defe840eccd9ef58a929321349cfff774a43b Mon Sep 17 00:00:00 2001 From: GiGiKoneti Date: Wed, 24 Jun 2026 01:37:11 +0530 Subject: [PATCH 5/6] Address review feedback --- neural_lam/models/step_predictors/base.py | 14 ++++- .../models/step_predictors/persistence.py | 10 +++- neural_lam/train_model.py | 13 ++--- tests/test_persistence.py | 56 +++++++++++++++---- 4 files changed, 68 insertions(+), 25 deletions(-) diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py index dff11b21..f1404c55 100644 --- a/neural_lam/models/step_predictors/base.py +++ b/neural_lam/models/step_predictors/base.py @@ -18,8 +18,18 @@ class StepPredictor(nn.Module, ABC): time steps plus forcing into a prediction of the next state. """ - trainable: bool = True - """Whether the predictor has learnable parameters and can be trained.""" + @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, diff --git a/neural_lam/models/step_predictors/persistence.py b/neural_lam/models/step_predictors/persistence.py index bca317a9..129d220a 100644 --- a/neural_lam/models/step_predictors/persistence.py +++ b/neural_lam/models/step_predictors/persistence.py @@ -2,6 +2,7 @@ # Third-party import torch +from loguru import logger # Local from ...datastore import BaseDatastore @@ -18,9 +19,6 @@ class PersistencePredictor(StepPredictor): ``ForecasterModule`` / ``ARForecaster`` pipeline. """ - trainable: bool = False - """Persistence predictors have no learnable parameters.""" - def __init__( self, datastore: BaseDatastore, @@ -46,6 +44,12 @@ def __init__( 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, diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index b127ec47..15e0721f 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -355,11 +355,6 @@ def main(input_args=None): ), ) args = parser.parse_args(input_args) - if not args.eval and not MODELS[args.model].trainable: - raise ValueError( - f"The {args.model} model cannot be trained. Run with " - "--eval to evaluate a baseline." - ) args.var_leads_metrics_watch = { int(k): v for k, v in json.loads(args.var_leads_metrics_watch).items() } @@ -463,10 +458,10 @@ def main(input_args=None): mesh_down_gnn_type=args.mesh_down_gnn_type, ) - if args.output_std and not predictor.predicts_std: - logger.warning( - f"Model '{args.model}' does not support predicting " - "standard deviation. The --output_std flag will be ignored." + 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) diff --git a/tests/test_persistence.py b/tests/test_persistence.py index bdfafbb0..b7df6f33 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -81,10 +81,24 @@ def test_persistence_forecaster_unroll(): def test_persistence_predicts_std_false(): - """PersistencePredictor.predicts_std is always 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") - predictor = PersistencePredictor(datastore=datastore, output_std=True) - assert not predictor.predicts_std + 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(): @@ -97,14 +111,34 @@ def test_persistence_training_error(): 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" - - with patch( - "neural_lam.train_model.ArgumentParser.parse_args", - return_value=mock_args, - ): - with pytest.raises( + 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)() + ), + ): + getattr(main, "__wrapped__", main)() From 2f7e8ac4db0c280f62016dd373967af44b37e49b Mon Sep 17 00:00:00 2001 From: GiGiKoneti Date: Mon, 29 Jun 2026 15:41:06 +0530 Subject: [PATCH 6/6] Address review suggestions --- tests/test_persistence.py | 1 + tests/test_train_model_warnings.py | 22 +++++++--------------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/tests/test_persistence.py b/tests/test_persistence.py index b7df6f33..b4258e8d 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -14,6 +14,7 @@ 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 diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index 454ddbd2..fd33e013 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -35,21 +35,13 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): side_effect=SystemExit(0), ): with patch("neural_lam.train_model.logger.warning") as mock_warning: - with patch("neural_lam.train_model.WeatherDataModule"): - # Add this patch to prevent further execution - mock_model_class = MagicMock() - mock_model_class.trainable = True - with patch( - "neural_lam.train_model.MODELS", - {"graph_lam": mock_model_class}, - ): - with pytest.raises(SystemExit): - main() - if expect_warning: - mock_warning.assert_called_once() - assert "--load" in mock_warning.call_args[0][0] - else: - mock_warning.assert_not_called() + with pytest.raises(SystemExit): + main() + if expect_warning: + mock_warning.assert_called_once() + assert "--load" in mock_warning.call_args[0][0] + else: + mock_warning.assert_not_called() def test_create_gif_forwarded_to_forecaster_module():