From a38dda9b150361aa4ce82f67c342e164cb101576 Mon Sep 17 00:00:00 2001 From: Nina Raoult Date: Fri, 10 Jul 2026 10:58:44 +0000 Subject: [PATCH] fix(grouping): ignore metadata keys missing from some fields when matching GroupByParam groups fields by their full MARS namespace metadata, so fields from GRIB1 files without ECMWF local definitions (missing class/type/stream/expver, e.g. IFS climate files such as climate.v015/95_4/sfc) could never be matched with GRIB2 fields that carry those keys. This broke MatchingFieldsFilter subclasses like land_parameters when joining such files. Keys that are not present in every candidate field are now ignored for grouping (with a warning), while keys present in all fields still separate groups as before. Co-Authored-By: Claude Opus 4.6 --- src/anemoi/transform/grouping/__init__.py | 37 +++++++++++++++++++ tests/test_grouping.py | 44 +++++++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/src/anemoi/transform/grouping/__init__.py b/src/anemoi/transform/grouping/__init__.py index 5c9e4672..d2a3a514 100644 --- a/src/anemoi/transform/grouping/__init__.py +++ b/src/anemoi/transform/grouping/__init__.py @@ -90,10 +90,40 @@ def _get_grouping_key( raise ValueError(f"Expected {extract_from_grouping_key} keys to extract, got {extracted_keys}") return grouping_key, extracted_keys + @staticmethod + def _restrict_to_common_keys(entries: list[tuple]) -> list[tuple]: + """Restrict grouping keys to the metadata keys present in all fields. + + Some fields (e.g. GRIB1 climate files without ECMWF local definitions) + lack metadata keys such as 'class', 'type', 'stream' or 'expver'. Keys + that are not present in every field are ignored when grouping, + otherwise such fields could never be matched together. + + Parameters + ---------- + entries : list of tuple + Tuples whose first element is the grouping key dict. + + Returns + ------- + list of tuple + Entries with grouping key dicts restricted to the common keys. + """ + if not entries: + return entries + all_keys = [set(key) for key, *_ in entries] + common = set.intersection(*all_keys) + dropped = set.union(*all_keys) - common + if dropped: + LOG.warning(f"Ignoring metadata keys not present in all fields when grouping: {sorted(dropped)}") + entries = [({k: v for k, v in key.items() if k in common}, *rest) for key, *rest in entries] + return entries + def _get_groups(self, data: list[Any], *, other: Callable[[Any], None] = _lost) -> None: assert callable(other), type(other) self.groups: dict[frozenset[Any], dict[str, Any]] = defaultdict(dict) self.groups_params = set() + entries = [] for f in data: key, extras = self._get_grouping_key( f, extract_from_grouping_key=["param"], remove_from_grouping_key=["variable"] @@ -104,6 +134,9 @@ def _get_groups(self, data: list[Any], *, other: Callable[[Any], None] = _lost) other(f) continue + entries.append((key, param, f)) + + for key, param, f in self._restrict_to_common_keys(entries): key = frozenset(key.items()) if param in self.groups[key]: @@ -143,6 +176,7 @@ def _get_groups(self, data: list[Any], *, other: Callable[[Any], None] = _lost) self.groups: dict[frozenset[Any], dict[str, Any]] = defaultdict(dict) self.groups_params = set() levels: dict[str, Any] = defaultdict(list) + entries = [] for f in data: key, extras = self._get_grouping_key( f, extract_from_grouping_key=["param", "levelist"], remove_from_grouping_key=["variable", "levtype"] @@ -154,6 +188,9 @@ def _get_groups(self, data: list[Any], *, other: Callable[[Any], None] = _lost) other(f) continue + entries.append((key, param, level, f)) + + for key, param, level, f in self._restrict_to_common_keys(entries): key = frozenset(key.items()) if level is None: diff --git a/tests/test_grouping.py b/tests/test_grouping.py index d02393fc..22ee76cd 100644 --- a/tests/test_grouping.py +++ b/tests/test_grouping.py @@ -84,6 +84,50 @@ def test_group_by_param(sample_fields): assert field in sample_fields +def test_group_by_param_missing_metadata_keys(): + """Fields missing some MARS keys (e.g. GRIB1 climate files without ECMWF + local definitions, lacking class/type/stream/expver) must still be grouped + with fields that have them. + """ + base = {"domain": "g", "levtype": "sfc", "date": 20200513, "time": 1200, "step": 0} + full = base | {"class": "od", "type": "an", "stream": "oper", "expver": "0001"} + fields = [ + mock_field(param="slt", **full), + mock_field(param="tvh", **base), + mock_field(param="tvl", **base), + ] + + match_params = ["tvh", "tvl", "slt"] + grouper = GroupByParam(params=match_params) + + groups = list(grouper.iterate(fields)) + assert len(groups) == 1 + assert [field.metadata("param") for field in groups[0]] == match_params + + +def test_group_by_param_missing_metadata_keys_still_separates_groups(): + """Ignoring missing keys must not merge fields that differ in a key + present in all fields. + """ + base = {"domain": "g", "levtype": "sfc", "date": 20200513, "time": 1200} + full = base | {"class": "od", "type": "an", "stream": "oper", "expver": "0001"} + fields = [ + mock_field(param="slt", step=0, **full), + mock_field(param="tvh", step=0, **base), + mock_field(param="slt", step=1, **full), + mock_field(param="tvh", step=1, **base), + ] + + match_params = ["tvh", "slt"] + grouper = GroupByParam(params=match_params) + + groups = list(grouper.iterate(fields)) + assert len(groups) == 2 + for group in groups: + assert [field.metadata("param") for field in group] == match_params + assert len({field.metadata("step") for field in group}) == 1 + + @pytest.mark.xfail(reason="vertical grouping not yet implemented") def test_group_by_param_vertical(sample_fields_vertical): from anemoi.transform.grouping import GroupByParamVertical