From 49517b4cc0f59dc5dd6c37aca3b5c2f4a8125067 Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 11:38:45 +0530 Subject: [PATCH 1/7] fix: Handle single-device runs in all_gather_cat --- neural_lam/models/ar_model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index f1bcb461d..e8f7ca8eb 100644 --- a/neural_lam/models/ar_model.py +++ b/neural_lam/models/ar_model.py @@ -327,7 +327,10 @@ def all_gather_cat(self, tensor_to_gather): returns: (K*d1, d2, ...) """ - return self.all_gather(tensor_to_gather).flatten(0, 1) + gathered = self.all_gather(tensor_to_gather) + if gathered.dim() > tensor_to_gather.dim(): + return gathered.flatten(0, 1) + return gathered # newer lightning versions requires batch_idx argument, even if unused # pylint: disable-next=unused-argument From a4a977a0246e5defbeb20ece8e49dfe731a0f060 Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 11:34:15 +0530 Subject: [PATCH 2/7] feat: Add GlobalDummyDatastore for global domain testing --- tests/conftest.py | 4 +++- tests/dummy_datastore.py | 15 +++++++++++++++ tests/test_datastores.py | 12 +++++++++--- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 47237ed55..c3f6a1445 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,7 @@ ) # Local -from .dummy_datastore import DummyDatastore +from .dummy_datastore import DummyDatastore, GlobalDummyDatastore # Disable weights and biases to avoid unnecessary logging # and to avoid having to deal with authentication @@ -103,9 +103,11 @@ def download_meps_example_reduced_dataset(): ), npyfilesmeps=None, dummydata=None, + dummydata_global=None, ) DATASTORES[DummyDatastore.SHORT_NAME] = DummyDatastore +DATASTORES[GlobalDummyDatastore.SHORT_NAME] = GlobalDummyDatastore def init_datastore_example(datastore_kind): diff --git a/tests/dummy_datastore.py b/tests/dummy_datastore.py index e3de83ca4..3108bb0cd 100644 --- a/tests/dummy_datastore.py +++ b/tests/dummy_datastore.py @@ -466,3 +466,18 @@ def grid_shape_state(self) -> CartesianGridShape: n_points_1d = int(np.sqrt(self.num_grid_points)) return CartesianGridShape(x=n_points_1d, y=n_points_1d) + + +class GlobalDummyDatastore(DummyDatastore): + """ + Dummy datastore that represents a global domain (no lateral boundaries). + """ + + SHORT_NAME = "dummydata_global" + + @cached_property + def boundary_mask(self) -> xr.DataArray: + """ + Return an all-zero boundary mask. + """ + return xr.zeros_like(self.ds["boundary_mask"]) diff --git a/tests/test_datastores.py b/tests/test_datastores.py index c719c69df..885d8f2a8 100644 --- a/tests/test_datastores.py +++ b/tests/test_datastores.py @@ -229,9 +229,15 @@ def test_boundary_mask(datastore_name): assert isinstance(da_mask, xr.DataArray) assert set(da_mask.dims) == {"grid_index"} assert da_mask.dtype == "int" - assert set(da_mask.values) == {0, 1} - assert da_mask.sum() > 0 - assert da_mask.sum() < da_mask.size + assert set(da_mask.values).issubset({0, 1}) + + if "global" in datastore_name: + # global datastores have no lateral boundaries + assert da_mask.sum() == 0 + else: + # LAM datastores must have a boundary + assert da_mask.sum() > 0 + assert da_mask.sum() < da_mask.size if isinstance(datastore, BaseRegularGridDatastore): grid_shape = datastore.grid_shape_state From 94d412e9fce083b48d4767704ba36082416e11d7 Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 11:40:32 +0530 Subject: [PATCH 3/7] test: Add global training test and update config --- tests/test_training.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/test_training.py b/tests/test_training.py index e8c131572..930742c41 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -41,10 +41,7 @@ def run_simple_training(datastore, set_output_std): max_epochs=1, deterministic=True, accelerator=device_name, - # XXX: `devices` has to be set to 2 otherwise - # neural_lam.models.ar_model.ARModel.aggregate_and_plot_metrics fails - # because it expects to aggregate over multiple devices - devices=2, + devices=1, log_every_n_steps=1, # use `detect_anomaly` to ensure that we don't have NaNs popping up # during training @@ -124,3 +121,11 @@ def test_training(datastore_name): def test_training_output_std(): datastore = init_datastore_example("mdp") # Test only with mdp datastore run_simple_training(datastore, set_output_std=True) + + +def test_training_global(): + """ + Test training with a global datastore (no boundaries). + """ + datastore = init_datastore_example("dummydata_global") + run_simple_training(datastore, set_output_std=False) From ef00869f23c065bc135c1d244b32004f7c593722 Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 11:43:42 +0530 Subject: [PATCH 4/7] docs: Fix docstring formatting in ARModel --- neural_lam/models/ar_model.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index f1bcb461d..aeed3720f 100644 --- a/neural_lam/models/ar_model.py +++ b/neural_lam/models/ar_model.py @@ -278,11 +278,14 @@ def unroll_prediction(self, init_states, forcing_features, true_states): def common_step(self, batch): """ - Predict on single batch batch consists of: init_states: (B, 2, - num_grid_nodes, d_features) target_states: (B, pred_steps, - num_grid_nodes, d_features) forcing_features: (B, pred_steps, - num_grid_nodes, d_forcing), - where index 0 corresponds to index 1 of init_states + Predict on single batch. + + Batch consists of: + init_states: (B, 2, num_grid_nodes, d_features) + target_states: (B, pred_steps, num_grid_nodes, d_features) + forcing_features: (B, pred_steps, num_grid_nodes, d_forcing) + + Note: index 0 corresponds to index 1 of init_states. """ (init_states, target_states, forcing_features, batch_times) = batch From 8ace28f13f908f5814ffcb6ffaeea7a5adad6dbe Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 12:47:56 +0530 Subject: [PATCH 5/7] docs: Update CHANGELOG.md for PR #452 --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d976ad4b9..82a77a669 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Handle single-device runs in all_gather_cat by skipping all_gather and returning the concatenated input tensor directly. [\#452](https://github.com/mllam/neural-lam/pull/452) @sohampatil01-svg + - Initialize `da_forcing_mean` and `da_forcing_std` to `None` when forcing data is absent, fixing `AttributeError` in `WeatherDataset` with `standardize=True` [\#369](https://github.com/mllam/neural-lam/issues/369) @Sir-Sloth-The-Lazy - Ensure proper sorting of `analysis_time` in `NpyFilesDatastoreMEPS._get_analysis_times` independent of the order in which files are processed with glob [\#386](https://github.com/mllam/neural-lam/pull/386) @Gopisokk From 6a791e761e65d074a57ceca2b9d09550554057bc Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 12:48:35 +0530 Subject: [PATCH 6/7] feat: Update CHANGELOG.md for PR #453 --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d976ad4b9..41266d5f0 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 `GlobalDummyDatastore` to support testing and development of global domain models without requiring large datasets. [\#453](https://github.com/mllam/neural-lam/pull/453) @sohampatil01-svg + - Enable `pin_memory` in DataLoaders when GPU is available for faster async CPU-to-GPU data transfers [\#236](https://github.com/mllam/neural-lam/pull/236) @abhaygoudannavar ### Changed From d2c27ef08baa9509ec916141e74f16ad6cb16a80 Mon Sep 17 00:00:00 2001 From: Soham_patil Date: Fri, 20 Mar 2026 12:49:01 +0530 Subject: [PATCH 7/7] docs: Update CHANGELOG.md for PR #458 --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d976ad4b9..6e3c06d8d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,6 +57,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Maintenance +- Fix reST formatting errors in `ARModel` methods to ensure documentation builds correctly. [\#458](https://github.com/mllam/neural-lam/pull/458) @sohampatil01-svg + - Update PR template to clarify milestone/roadmap requirement and maintenance changes [\#186](https://github.com/mllam/neural-lam/pull/186) @joeloskarsson - Update CI/CD to use python 3.13 for testing and full range of current python versions for linting (3.10 - 3.14) [\#173](https://github.com/mllam/neural-lam/pull/173) @observingClouds