diff --git a/CHANGELOG.md b/CHANGELOG.md index d90b45498..4db40dcc2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index d6eea4c59..f885aa226 100644 --- a/neural_lam/models/ar_model.py +++ b/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) diff --git a/tests/test_training.py b/tests/test_training.py index 972740695..7b2d6cd02 100644 --- a/tests/test_training.py +++ b/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"