From 5323ae8346859d4fc3dc7b84166882d8833fde8a Mon Sep 17 00:00:00 2001 From: EddyCMWF Date: Wed, 19 Jun 2024 16:58:49 +0100 Subject: [PATCH 1/3] preserve coordinate encoding --- cfgrib/xarray_store.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/cfgrib/xarray_store.py b/cfgrib/xarray_store.py index d0a32cd9..8a5e0200 100644 --- a/cfgrib/xarray_store.py +++ b/cfgrib/xarray_store.py @@ -42,6 +42,7 @@ def open_dataset(path, **kwargs): def merge_datasets(datasets, **kwargs): # type: (T.Sequence[xr.Dataset], T.Any) -> T.List[xr.Dataset] merged = [] # type: T.List[xr.Dataset] + first = [] # type: T.List[xr.Dataset] for ds in datasets: ds.attrs.pop("history", None) for i, o in enumerate(merged): @@ -55,6 +56,19 @@ def merge_datasets(datasets, **kwargs): pass else: merged.append(ds) + first.append(ds) + + # Add the important coordinate encoding fields from the first found, to the merged: + preserve_encoding_fields = [ + 'source', 'units', 'calendar', 'dtype' + ] + for i, o in enumerate(first): + for var in o.coords: + out_encoding = { + key: o[var].encoding[key] for key in preserve_encoding_fields if key in o[var].encoding + } + merged[i][var].encoding.update(out_encoding) + return merged From e64f500e9cb3d0d4231971c7a4a1fd788ed43e2c Mon Sep 17 00:00:00 2001 From: EddyCMWF Date: Wed, 19 Jun 2024 17:07:04 +0100 Subject: [PATCH 2/3] test for encoding --- tests/test_40_xarray_store.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_40_xarray_store.py b/tests/test_40_xarray_store.py index 541c8562..a224caf4 100644 --- a/tests/test_40_xarray_store.py +++ b/tests/test_40_xarray_store.py @@ -159,6 +159,18 @@ def test_open_datasets_differet_step_types_zeros() -> None: assert res[1].cfrzr.attrs["GRIB_stepType"] == "avg" +# ensure that the encoding of the coordinates is preserved +def test_open_datasets_differet_preserve_coordinate_encoding() -> None: + res = xarray_store.open_datasets(TEST_DATA_DIFFERENT_STEP_TYPES) + assert len(res) == 2 + assert "units" in res[0].valid_time.encoding + assert "units" in res[1].valid_time.encoding + + res = xarray_store.open_datasets(TEST_DATA_DIFFERENT_STEP_TYPES_ZEROS) + assert len(res) == 2 + assert "units" in res[0].valid_time.encoding + assert "units" in res[1].valid_time.encoding + def test_open_dataset_steps_in_minutes() -> None: res = xarray_store.open_dataset(TEST_DATA_STEPS_IN_MINUTES) From 96f5c9d361efcf73c34a0d03611082f5334f92f2 Mon Sep 17 00:00:00 2001 From: EddyCMWF Date: Thu, 20 Jun 2024 11:44:36 +0100 Subject: [PATCH 3/3] QA --- cfgrib/xarray_store.py | 10 +++++----- tests/test_40_xarray_store.py | 1 + 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/cfgrib/xarray_store.py b/cfgrib/xarray_store.py index 8a5e0200..b2b9a47b 100644 --- a/cfgrib/xarray_store.py +++ b/cfgrib/xarray_store.py @@ -57,15 +57,15 @@ def merge_datasets(datasets, **kwargs): else: merged.append(ds) first.append(ds) - + # Add the important coordinate encoding fields from the first found, to the merged: - preserve_encoding_fields = [ - 'source', 'units', 'calendar', 'dtype' - ] + preserve_encoding_fields = ["source", "units", "calendar", "dtype"] for i, o in enumerate(first): for var in o.coords: out_encoding = { - key: o[var].encoding[key] for key in preserve_encoding_fields if key in o[var].encoding + key: o[var].encoding[key] + for key in preserve_encoding_fields + if key in o[var].encoding } merged[i][var].encoding.update(out_encoding) diff --git a/tests/test_40_xarray_store.py b/tests/test_40_xarray_store.py index a224caf4..f14d70ae 100644 --- a/tests/test_40_xarray_store.py +++ b/tests/test_40_xarray_store.py @@ -171,6 +171,7 @@ def test_open_datasets_differet_preserve_coordinate_encoding() -> None: assert "units" in res[0].valid_time.encoding assert "units" in res[1].valid_time.encoding + def test_open_dataset_steps_in_minutes() -> None: res = xarray_store.open_dataset(TEST_DATA_STEPS_IN_MINUTES)