Mask boundary nodes in spatial loss maps during test_step

CHANGELOG.md

@@ -23,6 +23,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Mask boundary nodes in spatial loss maps computed during `test_step`, consistent with all other loss calls that use `interior_mask_bool` [\#568](https://github.com/mllam/neural-lam/pull/568) @RajdeepKushwaha5
+
 - Standardize all script references to use `create_graph` instead of the legacy `create_mesh` name in README and `pyproject.toml`, and fix minor README typos [\#426](https://github.com/mllam/neural-lam/pull/426) @GiGiKoneti
 
 - 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

neural_lam/models/ar_model.py

@@ -447,6 +447,8 @@ def test_step(self, batch, batch_idx):
         spatial_loss = self.loss(
             prediction, target, pred_std, average_grid=False
         )  # (B, pred_steps, num_grid_nodes)
+        # Exclude boundary nodes, consistent with training/validation loss
+        spatial_loss[..., ~self.interior_mask_bool] = float("nan")
         log_spatial_losses = spatial_loss[
             :, [step - 1 for step in self.args.val_steps_to_log]
         ]
@@ -753,7 +755,7 @@ def on_test_epoch_end(self):
             torch.cat(self.spatial_loss_maps, dim=0)
         )  # (N_test, N_log, num_grid_nodes)
         if self.trainer.is_global_zero:
-            mean_spatial_loss = torch.mean(
+            mean_spatial_loss = torch.nanmean(
                 spatial_loss_tensor, dim=0
             )  # (N_log, num_grid_nodes)
 

tests/test_training.py

@@ -189,3 +189,50 @@ def all_gather(self, tensor, sync_grads=False):
         "all_gather_cat produced incorrectly ordered/combined values "
         "on multi-device simulation"
     )
+
+
+def test_spatial_loss_maps_exclude_boundary():
+    """
+    Test that spatial loss maps mask out boundary nodes with NaN and that
+    nanmean correctly ignores them during aggregation, matching the masking
+    logic in test_step and on_test_epoch_end.
+    """
+    num_grid_nodes = 10
+    pred_steps = 3
+    batch_size = 2
+
+    # Interior mask: first 6 nodes interior, last 4 boundary
+    interior_mask_bool = torch.tensor(
+        [True] * 6 + [False] * 4
+    )  # (num_grid_nodes,)
+
+    # Simulate spatial loss (all ones for simplicity)
+    spatial_loss = torch.ones(batch_size, pred_steps, num_grid_nodes)
+
+    # Apply the same masking as test_step
+    spatial_loss[..., ~interior_mask_bool] = float("nan")
+
+    # Boundary nodes should be NaN
+    assert torch.all(
+        torch.isnan(spatial_loss[..., ~interior_mask_bool])
+    ), "Boundary nodes should be NaN"
+
+    # Interior nodes should not be NaN
+    assert not torch.any(
+        torch.isnan(spatial_loss[..., interior_mask_bool])
+    ), "Interior nodes should not be NaN"
+
+    # nanmean over batch dim should ignore NaN boundary nodes
+    mean_spatial_loss = torch.nanmean(spatial_loss, dim=0)
+    assert mean_spatial_loss.shape == (pred_steps, num_grid_nodes)
+
+    # Interior nodes: mean of 1.0 = 1.0
+    assert torch.allclose(
+        mean_spatial_loss[:, interior_mask_bool],
+        torch.ones(pred_steps, interior_mask_bool.sum().item()),
+    )
+
+    # Boundary nodes: nanmean of all-NaN = NaN
+    assert torch.all(
+        torch.isnan(mean_spatial_loss[:, ~interior_mask_bool])
+    ), "Boundary nodes should remain NaN after nanmean"