Switch to using CF compliant coord stacking (rather than xr.Dataset.reset_index())

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

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)
 

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"

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])