From f6acb2dad041d17c7e08dd3b0542471a346726f2 Mon Sep 17 00:00:00 2001 From: zweihuehner <2huehner@gmail.com> Date: Wed, 26 Nov 2025 09:52:50 +0000 Subject: [PATCH 1/8] Problem with output.coord_ranges slicing for coordinates not present in all of state/static/forcing #81 --- mllam_data_prep/create_dataset.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py index 3daf321..6ed6542 100644 --- a/mllam_data_prep/create_dataset.py +++ b/mllam_data_prep/create_dataset.py @@ -236,10 +236,12 @@ def create_dataset(config: Config): # only need to do selection for the coordinates that the input dataset actually has if output_coord_ranges is not None: - output_coord_ranges = { + # Use a temporary dict to avoid modifying the original ranges. + # Static features have no time dimension, so they would return an empty dict. + output_coord_ranges_tmp = { k: w for k, w in output_coord_ranges.items() if k in output_dims } - da_target = select_by_kwargs(da_target, **output_coord_ranges) + da_target = select_by_kwargs(da_target, **output_coord_ranges_tmp) dataarrays_by_target[target_output_var].append(da_target) From 5ec9174181c6b272886990b6b369e913b9a5922f Mon Sep 17 00:00:00 2001 From: zweihuehner <2huehner@gmail.com> Date: Wed, 26 Nov 2025 12:03:58 +0000 Subject: [PATCH 2/8] adapt changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 105c9d4..78dd7f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ This release contains bugfixes to update tests to use newer version of pre-commi ### Fixes - use old union typing notation compatible with all required python versions [\#77](https://github.com/mllam/mllam-data-prep/pull/77) @SimonKamuk +- fix bug that would overwrite the output_coord_ranges dictionary if feature order in datastore.yaml is incorrect [\#87](https://github.com/mllam/mllam-data-prep/pull/87) @zweihuehner ### Maintenance - update pre-commit action to v3.0.1 [\#77](https://github.com/mllam/mllam-data-prep/pull/77) @SimonKamuk From 536c3178e887ef1147fec7254b96c0c530bcc388 Mon Sep 17 00:00:00 2001 From: zweihuehner <2huehner@gmail.com> Date: Wed, 26 Nov 2025 12:55:06 +0000 Subject: [PATCH 3/8] refined changelog --- CHANGELOG.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 78dd7f6..79fc1c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [unreleased](https://github.com/mllam/mllam-data-prep/compare/v0.6.1...HEAD) + +### Fixes +- fix bug where coordinate selection of an unshared dimension isn't applied to subsequent ouput variables when an output variable without this dimension is processed before the others [\#87](https://github.com/mllam/mllam-data-prep/pull/87) @zweihuehner + ## [v0.6.1](https://github.com/mllam/mllam-data-prep/release/tag/v0.6.1) [All changes](https://github.com/mllam/mllam-data-prep/compare/v0.6.1...v0.6.0) @@ -13,7 +18,6 @@ This release contains bugfixes to update tests to use newer version of pre-commi ### Fixes - use old union typing notation compatible with all required python versions [\#77](https://github.com/mllam/mllam-data-prep/pull/77) @SimonKamuk -- fix bug that would overwrite the output_coord_ranges dictionary if feature order in datastore.yaml is incorrect [\#87](https://github.com/mllam/mllam-data-prep/pull/87) @zweihuehner ### Maintenance - update pre-commit action to v3.0.1 [\#77](https://github.com/mllam/mllam-data-prep/pull/77) @SimonKamuk From fbd66a10e17c616e8bacd4221b90f6e761e2fa25 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Faro=20Sch=C3=A4fer?= <58438715+zweihuehner@users.noreply.github.com> Date: Thu, 27 Nov 2025 10:03:07 +0100 Subject: [PATCH 4/8] Improve comments on coordinate range selection logic Clarified comments regarding coordinate range selection and static features. --- mllam_data_prep/create_dataset.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py index 6ed6542..faef71a 100644 --- a/mllam_data_prep/create_dataset.py +++ b/mllam_data_prep/create_dataset.py @@ -236,8 +236,11 @@ def create_dataset(config: Config): # only need to do selection for the coordinates that the input dataset actually has if output_coord_ranges is not None: - # Use a temporary dict to avoid modifying the original ranges. - # Static features have no time dimension, so they would return an empty dict. + # Use a temporary dict to apply selection on coordinate ranges to avoid + # modifying the original ranges given in the config. This is needed because + # static features, for example, do not have a time dimension. Hence, the time + # based selection returns an empty dictionary, which should not overwrite the + # selection for the other variables. output_coord_ranges_tmp = { k: w for k, w in output_coord_ranges.items() if k in output_dims } From dd70f5ccddf516fbab466531a63bbde07332ccf4 Mon Sep 17 00:00:00 2001 From: zweihuehner <2huehner@gmail.com> Date: Thu, 27 Nov 2025 15:53:08 +0000 Subject: [PATCH 5/8] run pre-commit hook --- mllam_data_prep/create_dataset.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py index faef71a..c51eef6 100644 --- a/mllam_data_prep/create_dataset.py +++ b/mllam_data_prep/create_dataset.py @@ -236,10 +236,10 @@ def create_dataset(config: Config): # only need to do selection for the coordinates that the input dataset actually has if output_coord_ranges is not None: - # Use a temporary dict to apply selection on coordinate ranges to avoid - # modifying the original ranges given in the config. This is needed because - # static features, for example, do not have a time dimension. Hence, the time - # based selection returns an empty dictionary, which should not overwrite the + # Use a temporary dict to apply selection on coordinate ranges to avoid + # modifying the original ranges given in the config. This is needed because + # static features, for example, do not have a time dimension. Hence, the time + # based selection returns an empty dictionary, which should not overwrite the # selection for the other variables. output_coord_ranges_tmp = { k: w for k, w in output_coord_ranges.items() if k in output_dims From 05e8ad68063c5e9769fa053f2bb1e547969c6820 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Mon, 9 Feb 2026 13:45:58 +0100 Subject: [PATCH 6/8] add test that exposes input ordering bug --- tests/test_output_coord_ranges_slicing.py | 124 ++++++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 tests/test_output_coord_ranges_slicing.py diff --git a/tests/test_output_coord_ranges_slicing.py b/tests/test_output_coord_ranges_slicing.py new file mode 100644 index 0000000..17f840a --- /dev/null +++ b/tests/test_output_coord_ranges_slicing.py @@ -0,0 +1,124 @@ +import itertools + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +import mllam_data_prep as mdp +from mllam_data_prep.config import DimMapping, InputDataset, Output, Range + + +def _write_zarr(ds: xr.Dataset, path): + ds.to_zarr(path, mode="w") + + +@pytest.mark.parametrize( + "input_order", + list(itertools.permutations(["state", "static", "forcing"])), +) +def test_output_coord_ranges_not_dropped_between_inputs(tmp_path, input_order): + """ + Ensure output coord range slicing is applied per-input without being + affected by the order of inputs. This guards against mutating the shared + output_coord_ranges dict when an input lacks a dimension (e.g. `static` + without `time`), which would otherwise remove slicing for later inputs. + + See https://github.com/mllam/mllam-data-prep/issues/81 for bug report. + """ + time = pd.date_range("2000-01-01", "2000-01-05", freq="1D") + x = np.arange(2) + + state_ds = xr.Dataset( + {"s": (("time", "x"), np.zeros((len(time), len(x))))}, + coords={"time": time, "x": x}, + ) + forcing_ds = xr.Dataset( + {"f": (("time", "x"), np.ones((len(time), len(x))))}, + coords={"time": time, "x": x}, + ) + static_ds = xr.Dataset( + {"static_feature": (("x",), np.array([10.0, 20.0]))}, + coords={"x": x}, + ) + + state_path = tmp_path / "state.zarr" + forcing_path = tmp_path / "forcing.zarr" + static_path = tmp_path / "static.zarr" + + _write_zarr(state_ds, state_path) + _write_zarr(forcing_ds, forcing_path) + _write_zarr(static_ds, static_path) + + inputs_by_name = { + "state": InputDataset( + path=str(state_path), + dims=["time", "x"], + variables=["s"], + target_output_variable="state", + dim_mapping={ + "time": DimMapping(method="rename", dim="time"), + "grid_index": DimMapping(method="stack", dims=["x"]), + "state_feature": DimMapping( + method="stack_variables_by_var_name", + name_format="{var_name}", + ), + }, + ), + "static": InputDataset( + path=str(static_path), + dims=["x"], + variables=["static_feature"], + target_output_variable="static", + dim_mapping={ + "grid_index": DimMapping(method="stack", dims=["x"]), + "static_feature": DimMapping( + method="stack_variables_by_var_name", + name_format="{var_name}", + ), + }, + ), + "forcing": InputDataset( + path=str(forcing_path), + dims=["time", "x"], + variables=["f"], + target_output_variable="forcing", + dim_mapping={ + "time": DimMapping(method="rename", dim="time"), + "grid_index": DimMapping(method="stack", dims=["x"]), + "forcing_feature": DimMapping( + method="stack_variables_by_var_name", + name_format="{var_name}", + ), + }, + ), + } + + ordered_inputs = {name: inputs_by_name[name] for name in input_order} + + config = mdp.Config( + schema_version="v0.6.0", + dataset_version="v0.0.0", + output=Output( + variables={ + "state": ["time", "grid_index", "state_feature"], + "forcing": ["time", "grid_index", "forcing_feature"], + "static": ["grid_index", "static_feature"], + }, + coord_ranges={ + "time": Range( + start="2000-01-01T00:00", + end="2000-01-03T00:00", + step="PT24H", + ) + }, + ), + inputs=ordered_inputs, + ) + + ds = mdp.create_dataset(config=config) + + expected_len = 3 + assert ds["state"].sizes["time"] == expected_len + assert ds["forcing"].sizes["time"] == expected_len + assert "time" not in ds["static"].dims From 3b2071725f0cf084c7b080eae0d102b7e2ab467e Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Mon, 9 Feb 2026 13:46:29 +0100 Subject: [PATCH 7/8] add more helpful input merge exception message --- mllam_data_prep/create_dataset.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py index 5f12612..052eeef 100644 --- a/mllam_data_prep/create_dataset.py +++ b/mllam_data_prep/create_dataset.py @@ -105,10 +105,17 @@ def _merge_dataarrays_by_target(dataarrays_by_target): ds = xr.merge(dataarrays, join="exact") except ValueError as ex: if ex.args[0].startswith("cannot align objects with join='exact'"): + + def _summarize(da): + dims = ", ".join([f"{k}: {v}" for k, v in da.sizes.items()]) + return f"{da.name} ({dims})\n{da.coords}" + + coord_summaries = "\n".join([_summarize(da) for da in dataarrays]) raise InvalidConfigException( - f"Couldn't merge together the dataarrays for all targets ({', '.join(dataarrays_by_target.keys())})" - f" This is likely because the dataarrays have different dimensions or coordinates." - " Maybe you need to give the 'feature' dimension a unique name for each target variable?" + f"Couldn't merge together the dataarrays for all targets ({', '.join(dataarrays_by_target.keys())}). " + "This is likely because the dataarrays have different dimensions or coordinates. " + f"Dataarray coords:\n{coord_summaries}" + "Maybe you need to give the 'feature' dimension a unique name for each target variable?" ) from ex else: raise ex From 4985d63b7225fbb639554fcde731f7290f9369be Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Mon, 9 Feb 2026 14:01:45 +0100 Subject: [PATCH 8/8] changelog --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4595e43..846a62f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,10 +5,10 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [unreleased](https://github.com/mllam/mllam-data-prep/compare/v0.6.1...HEAD) +## [unreleased](https://github.com/mllam/mllam-data-prep/compare/v0.7.0...HEAD) ### Fixes -- fix bug where coordinate selection of an unshared dimension isn't applied to subsequent ouput variables when an output variable without this dimension is processed before the others [\#87](https://github.com/mllam/mllam-data-prep/pull/87) @zweihuehner +- fix bug where coordinate selection of an unshared dimension isn't applied to subsequent ouput variables when an output variable without this dimension is processed before the others [\#90](https://github.com/mllam/mllam-data-prep/pull/90) @zweihuehner & @leifdenby ## [v0.7.0](https://github.com/mllam/mllam-data-prep/release/tag/v0.7.0)