Bug: `calc_stats` mutates dataset in loop, causing incorrect `diff_std` computation

[RajdeepKushwaha5]

Description

The calc_stats function in mllam_data_prep/ops/statistics.py mutates the ds variable inside the loop over statistics_config.ops. When a diff_* operation is encountered, the original dataset is overwritten:

ds = ds[vars_to_keep].diff(dim=splitting_dim)

This means any subsequent diff_* operation in the same loop will apply .diff() on an already-diffed dataset instead of the original.

Reproduction

With the default config shipped in all example YAML files:

ops: [mean, std, diff_mean, diff_std]

The loop processes:

1. mean → ds.mean(...) ✅
2. std → ds.std(...) ✅
3. diff_mean → ds = ds.diff(...) then ds.mean(...) ✅ (but mutates ds)
4. diff_std → ds = ds.diff(...) again on already-diffed data → computes std(diff(diff(ds))) ❌

Expected: diff_std should compute std(diff(ds)) (single diff of the original dataset).
Actual: diff_std computes std(diff(diff(ds))) (double diff).

Impact

• Silent data corruption — no error or warning is raised; the output is numerically wrong.
• Affects all default configs — every shipped example (example.danra.yaml, example.era5_cropped.yaml) uses ops: [mean, std, diff_mean, diff_std].
• Incorrect normalization statistics are fed to downstream ML weather models.

Root Cause

Introduced in commit 7784a46 (PR #10 — "Add splitting and calculation of statistics", merged Jul 17 2024). The file has never been modified since.

Suggested Fix

Use a local variable for the per-iteration dataset instead of overwriting ds:

for op_split in statistics_config.ops:
    ds_op = ds  # preserve original dataset for each iteration
    try:
        pre_op, op = op_split.split("_")
    except ValueError:
        op = op_split
        pre_op = None

    if pre_op is not None:
        if pre_op == "diff":
            vars_to_keep = [v for v in ds.data_vars if splitting_dim in ds[v].dims]
            ds_op = ds[vars_to_keep].diff(dim=splitting_dim)
        else:
            raise NotImplementedError(pre_op)

    fn = getattr(ds_op, op)
    stats[op_split] = fn(dim=statistics_config.dims)

Environment

• mllam-data-prep version: v0.7.0 (current main)
• File: mllam_data_prep/ops/statistics.py, lines 34–52