From e1d9eb01bd77b9acffdf34a751a2f25ce764d588 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 25 Nov 2025 13:51:50 +0100 Subject: [PATCH 1/3] Use CF-compliant coordinate stacking rather than .reset_index() This commit replaces the current stacking method where `xr.Dataset.reset_index()` is used to change the stacked coordinate (for example `grid_index` as a stacking of `x` and `y` coordinates) from a `pd.MultiIndex` to a regular index, and instead uses the CF-compliant "gather compression" of coordinate stacking (see https://cf-xarray.readthedocs.io/en/latest/coding.html). Using `reset_index()` is a lossy operation in that downstream applications will not know wwhich coordinates have been stacked, meaning we must instead rely on this information from elsewhere (e.g. the config) in downstream applications. The reason that stacked coordinates must be handled in a special way is that `pd.MultiIndex` coordinates cannot be directly written to netCDF/zarr datasets, but with the CF-compliant method they can be handled in such a manner that we can completely roundtrip in downstream applications. --- mllam_data_prep/create_dataset.py | 14 ++++++++++++++ mllam_data_prep/ops/mapping.py | 2 +- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py index 3daf321..f8ffe7b 100644 --- a/mllam_data_prep/create_dataset.py +++ b/mllam_data_prep/create_dataset.py @@ -4,7 +4,9 @@ from pathlib import Path from typing import Optional, Union +import cf_xarray as cfxr import numpy as np +import pandas as pd import xarray as xr import yaml import zarr @@ -288,6 +290,18 @@ def create_dataset(config: Config): ) ds["splits"] = da_splits + # We have to deal with the fact that MultiIndex objects (this would + # commonly before example `grid_index` created by stacking the `x` and `y` + # coordinates) can't be written to netcdf/zarr. In cf_xarray this has been + # handled in a cf-compliant manner using so-called "compression by + # gathering" (see + # https://cf-xarray.readthedocs.io/en/latest/generated/cf_xarray.encode_multi_index_as_compress.html#cf_xarray.encode_multi_index_as_compress). + # which allows us to safely roundtrip MultiIndexes through netcdf/zarr, + # using their encode and decode functions. + for idx in ds.indexes: + if isinstance(ds.indexes[idx], pd.MultiIndex): + ds = cfxr.encode_multi_index_as_compress(ds, idxnames=idx) + ds.attrs = {} ds.attrs["schema_version"] = config.schema_version ds.attrs["dataset_version"] = config.dataset_version diff --git a/mllam_data_prep/ops/mapping.py b/mllam_data_prep/ops/mapping.py index 9482ff8..f453699 100644 --- a/mllam_data_prep/ops/mapping.py +++ b/mllam_data_prep/ops/mapping.py @@ -100,7 +100,7 @@ def map_dims_and_variables(ds, dim_mapping, expected_input_var_dims): # in the input dataset that we want to stack to create the architecture # dimension, this is for example used for flatting the spatial dimensions # into a single dimension representing the grid index - ds = ds.stack({arch_dim: source_dims}).reset_index(arch_dim) + ds = ds.stack({arch_dim: source_dims}) else: raise NotImplementedError(method) From cc2ec4fd4aa205089ac7a5196c2caaf027780239 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 25 Nov 2025 14:52:21 +0100 Subject: [PATCH 2/3] add test for cf-compat unstacking --- tests/test_stacking.py | 61 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/tests/test_stacking.py b/tests/test_stacking.py index d8fa859..ebe313c 100644 --- a/tests/test_stacking.py +++ b/tests/test_stacking.py @@ -1,7 +1,9 @@ +import cf_xarray as cfxr import numpy as np import xarray as xr -from mllam_data_prep.config import DimMapping +from mllam_data_prep.config import Config, DimMapping +from mllam_data_prep.create_dataset import create_dataset from mllam_data_prep.ops import mapping as mdp_mapping from mllam_data_prep.ops import stacking as mdp_stacking @@ -85,3 +87,60 @@ def test_stack_xy_coords(): assert set(da_stacked.dims) == set(("grid_index", "feature")) assert da_stacked.coords["grid_index"].shape == (nx * ny,) + + +STACKING_EXAMPLE_YAML = """ +schema_version: v0.6.0 +dataset_version: v0.1.0 + +output: + variables: + static: [grid_index, static_feature] + coord_ranges: + time: + start: 1990-09-03T00:00 + end: 1990-09-06T00:00 + chunking: + time: 1 + +inputs: + danra_height_levels: + path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr + dims: [time, x, y] + variables: [t2m] + dim_mapping: + static_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + dims: [x, y] + + target_output_variable: static +""" + + +def test_unstack_on_processed_dataset(): + config = Config.from_yaml(STACKING_EXAMPLE_YAML) + + ds_stacked = create_dataset(config) + + # check that we can get back the original x, y values from the stacked + # coord. This is possible because we use the CF-compliant "gather + # compression" for encoding the stacked coordinate MultiIndex + # (https://cf-xarray.readthedocs.io/en/latest/coding.html) + ds_multi_index = cfxr.decode_compress_to_multi_index( + ds_stacked, idxnames="grid_index" + ) + # check we can get out the names of the stacked coords + stacked_coords = ds_multi_index.grid_index.to_index().names + assert set(stacked_coords) == set(["x", "y"]) + + ds_unstacked = ds_multi_index.unstack("grid_index") + + # compare unstacked coords to original dataset coords + + ds_orig = xr.open_zarr(config.inputs["danra_height_levels"].path, chunks={}) + + for coord in ["x", "y"]: + xr.testing.assert_equal(ds_unstacked[coord], ds_orig[coord]) From 7945ce3ba859f49bd1b529351967aa06e59ef6a4 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 25 Nov 2025 20:59:13 +0100 Subject: [PATCH 3/3] add cf-xarray dep --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 9778172..257a367 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "psutil>=5.7.2", "packaging>=23.1", "deepdiff>=8.2.0", + "cf-xarray>=0.9.4", ] requires-python = ">=3.9" readme = "README.md"