Document saved statistics variable names

README.md

@@ -283,7 +283,9 @@ The `output` section defines three things:
 1. `variables`: what input variables the model architecture you are targeting expects, and what the dimensions are for each of these variables.
 2. `coord_ranges`: the range of values for each of the dimensions that the model architecture expects as input. These are optional, but allows you to ensure that the training dataset is created with the correct range of values for each dimension.
 3. `chunking`: the chunk sizes to use when writing the training dataset to zarr. This is optional, but can be used to optimise the performance of the zarr dataset. By default the chunk sizes are set to the size of the dimension, but this can be overridden by setting the chunk size in the configuration file. A common choice is to set the dimension along which you are batching to align with the of each training item (e.g. if you are training a model with time-step roll-out of 10 timesteps, you might choose a chunksize of 10 along the time dimension).
-4. Splitting and calculation of statistics of the output variables, using the `splitting` section. The `output.splitting.splits` attribute defines the individual splits to create (for example `train`, `val` and `test`) and `output.splitting.dim` defines the dimension to split along. The `compute_statistics` can be optionally set for a given split to calculate the statistical properties requested (for example `mean`, `std`) any method available on `xarray.Dataset.{op}` can be used. In addition methods prefixed by `diff_` (so the operational would be listed as `diff_{op}`) to compute a statistic based on difference of consecutive time-steps, e.g. `diff_mean` to compute the `mean` of the difference between consecutive timesteps (these are used for normalisating increments). The `dims` attribute defines the dimensions to calculate the statistics over (for example `grid_index` and `time`).
+4. Splitting and calculation of statistics for the output variables, using the `splitting` section. The `output.splitting.splits` attribute defines the individual splits to create (for example `train`, `val` and `test`) and `output.splitting.dim` defines the dimension to split along. The `compute_statistics` attribute can be set for a given split to calculate the requested statistical properties (for example `mean` and `std`). Any method available on `xarray.Dataset.{op}` can be used. In addition, methods prefixed with `diff_` (so the operation is listed as `diff_{op}`) compute statistics from differences between consecutive time steps. For example, `diff_mean` computes the `mean` of those differences, which can be used for normalising increments. The `dims` attribute defines the dimensions to reduce over (for example `grid_index` and `time`).
+
+   The calculated statistics are saved as additional variables in the output dataset. Their names follow the pattern `{output_variable}__{split_name}__{operation}`. For the example above, statistics for the `state` output variable in the `train` split are saved as variables such as `state__train__mean`, `state__train__std`, `state__train__diff_mean` and `state__train__diff_std`. Each statistic keeps the dimensions that were not listed in `compute_statistics.dims`, so reducing over `grid_index` and `time` for a `state(time, grid_index, state_feature)` variable leaves a result indexed by `state_feature`. Operations prefixed with `diff_` are only calculated for output variables that span the splitting dimension, since the time-step difference is taken along that dimension before applying the statistic.
 
 In addition the `output` section can also contain a configuration for cropping the output dataset using the convex hull of coordinates from a different dataset. This is used for example when creating training datasets from limited area modelling (LAM) setups, where a separate dataset is used for the boundary data. The example above doesn't include this section, but [see below](#cropping-the-output-dataset-using-convex-hull-of-another-dataset) for an example of how to use this feature.
 

mllam_data_prep/ops/statistics.py

@@ -38,15 +38,16 @@ def calc_stats(
             op = op_split
             pre_op = None
 
+        ds_for_op = ds
         if pre_op is not None:
             if pre_op == "diff":
                 # subset to select only the variable which have the splitting_dim
                 vars_to_keep = [v for v in ds.data_vars if splitting_dim in ds[v].dims]
-                ds = ds[vars_to_keep].diff(dim=splitting_dim)
+                ds_for_op = ds[vars_to_keep].diff(dim=splitting_dim)
             else:
                 raise NotImplementedError(pre_op)
 
-        fn = getattr(ds, op)
+        fn = getattr(ds_for_op, op)
         stats[op_split] = fn(dim=statistics_config.dims)
 
     return stats

tests/test_statistics.py

@@ -0,0 +1,37 @@
+import numpy as np
+import xarray as xr
+
+from mllam_data_prep.config import Statistics
+from mllam_data_prep.ops.statistics import calc_stats
+
+
+def test_diff_statistics_start_from_original_dataset_for_each_operation():
+    ds = xr.Dataset(
+        {
+            "state": xr.DataArray(
+                np.array([[0.0], [1.0], [3.0], [6.0]]),
+                dims=("time", "feature"),
+                coords={"time": [0, 1, 2, 3], "feature": ["a"]},
+            )
+        }
+    )
+    statistics_config = Statistics(
+        ops=["diff_mean", "diff_std"],
+        dims=["time"],
+    )
+
+    stats = calc_stats(
+        ds=ds,
+        statistics_config=statistics_config,
+        splitting_dim="time",
+    )
+
+    expected_diffs = ds["state"].diff(dim="time")
+    xr.testing.assert_allclose(
+        stats["diff_mean"]["state"],
+        expected_diffs.mean(dim="time"),
+    )
+    xr.testing.assert_allclose(
+        stats["diff_std"]["state"],
+        expected_diffs.std(dim="time"),
+    )