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