Fix silent stat corruption in distributed standardization-stats depadding

[nikhil3495]

Describe your changes

compute_standardization_stats.py --distributed pads WeatherDataset (via PaddedWeatherDataset) so its length divides evenly across ranks, then after the multi-rank gather "depads" by taking gathered[:total_samples] (equivalently gathered[i] for i in original_indices]). That assumes padded rows land at the tail of the gathered, rank-major-concatenated tensor.

They don't: DistributedSampler(shuffle=False) stripes dataset indices across ranks (rank r gets r, r+world_size, r+2*world_size, ...), so whenever total_samples % world_size != 0, the padded rows are scattered across several ranks' local tails rather than one global tail. The prefix-selection then silently keeps a few padded (duplicated last-sample) rows while dropping an equal number of real ones — corrupting the saved parameter_mean/std.pt and diff_mean/std.pt tensors used to standardize the whole dataset.

I confirmed this with a standalone simulation using this repo's actual DistributedSampler: for total_samples=101, world_size=4, batch_size=8, the old logic swaps real samples 95 and 99 for padded samples 101 and 102.

Fix: identify real vs. padded rows locally, per rank, before the gather, instead of trying to reconstruct sample identity after it. real_sample_mask_per_batch() derives — from pure index arithmetic on the sampler's already-deterministic iteration order — which rows of each minibatch are real, with no dependency on gather order. Both the parameter-stats and diff-stats accumulators are filtered by this mask as soon as each batch is computed, so the post-gather concatenation already contains exactly the real samples, in whatever order (order doesn't matter for a mean/std reduction). This also removes the need for the previous get_original_indices() / n_original_windows positional bookkeeping in main().

Not a duplicate of #409/#412/#413 (all closed via #411): those fixed a shape/IndexError bug in the same file's gather logic, not this sample-identity mismatch. The tests #411 added only check that slicing preserves shape/values on synthetic tensors — they never simulate an actual multi-rank DistributedSampler, which is why this gap went uncaught.

No new dependencies required.

Issue Link

closes #749

Type of change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form.
• I have requested a reviewer and an assignee. I don't have write access to the repo — happy to have a maintainer add these.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, under Fixed.

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
• Once the PR is ready to be merged, squash commits and merge the PR.

---

Generated with Claude Code

[sadamov]

Nice catch, reproduced with a real DataLoader (101/4/8 loses samples 95 and 99 and counts sample 100 three times). Suggestions below, please rebase first, the CHANGELOG conflicts with #743.

Outside the diff so not suggestible: get_original_indices, original_indices and padded_indices have no callers left after this PR (__getitem__ can return self.base_dataset[min(idx, self.total_samples - 1)]), I would drop them and test_original_indices too.

Below I am mostly talking about shorter comments and docstrings.

[nikhil3495]

Thanks for the detailed review! I've addressed all points:

• Converted real_sample_mask_per_batch into the real_sample_masks method
  on PaddedWeatherDataset
• Removed the now-dead original_indices/padded_indices/get_original_indices
  and test_original_indices
• Fixed the reversed/misleading comments
• Rewrote the Bug 3 tests to use a real DataLoader with your parametrized cases
• Shrunk the CHANGELOG entry to one line
• Rebased onto main to resolve the CHANGELOG conflict

All pushed. Let me know if anything else needs adjusting!