fix(grouping): ignore metadata keys missing from some fields when matching

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:

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