implement also concat for trajectories

Author: floriankrb

src/anemoi/datasets/create/gridded/context.py

@@ -67,13 +67,14 @@ def create_result(self, argument: Any, data: Any) -> SimpleGriddedResult:
         """
         return SimpleGriddedResult(self, argument, data)
 
-    def matching_dates(self, filtering_dates: Any, group_of_dates: Any) -> GroupOfDates:
-        """Find dates that match between filtering_dates and group_of_dates.
+    def matching_dates(self, filters: dict, group_of_dates: Any) -> GroupOfDates:
+        """Find dates that match between filters and group_of_dates.
 
         Parameters
         ----------
-        filtering_dates : Any
-            The dates to filter by.
+        filters : dict
+            A dict mapping filter keys to DatesProvider objects.
+            Gridded layouts only support ``'dates'``.
         group_of_dates : Any
             The group of dates to compare against.
 
@@ -82,7 +83,14 @@ def matching_dates(self, filtering_dates: Any, group_of_dates: Any) -> GroupOfDa
         GroupOfDates
             A GroupOfDates object containing the intersection of the two sets.
         """
-
+        unsupported = set(filters) - {"dates"}
+        if unsupported:
+            raise ValueError(
+                f"Gridded layout does not support filtering by {unsupported}. "
+                "Use 'dates' instead."
+            )
+
+        filtering_dates = filters["dates"]
         return GroupOfDates(sorted(set(group_of_dates) & set(filtering_dates)), group_of_dates.provider)
 
     def origin(self, data: Any, action: Any, action_arguments: Any) -> Any:

src/anemoi/datasets/create/input/action.py

@@ -69,6 +69,8 @@ class Concat(Action):
 
     """
 
+    FILTER_KEYS = ("dates", "base_dates", "steps")
+
     def __init__(self, config, *path):
         super().__init__(config, *path, "concat")
 
@@ -78,10 +80,24 @@ def __init__(self, config, *path):
 
         for i, item in enumerate(config):
 
-            dates = item["dates"]
-            filtering_dates = DatesProvider.from_config(**dates)
-            action = action_factory({k: v for k, v in item.items() if k != "dates"}, *self.path, str(i))
-            self.choices.append((filtering_dates, action))
+            filters = {}
+            for key in self.FILTER_KEYS:
+                if key in item:
+                    if key == "steps":
+                        from anemoi.datasets.create.trajectories.context import Steps
+
+                        filters[key] = Steps(item[key])
+                    else:
+                        filters[key] = DatesProvider.from_config(**item[key])
+
+            if not filters:
+                raise ValueError(
+                    f"Concat entry must have at least one of {self.FILTER_KEYS}."
+                )
+
+            action_config = {k: v for k, v in item.items() if k not in self.FILTER_KEYS}
+            action = action_factory(action_config, *self.path, str(i))
+            self.choices.append((filters, action))
 
     def __repr__(self):
         return f"Concat({self.choices})"
@@ -90,8 +106,8 @@ def __call__(self, context, argument):
 
         results = []
 
-        for filtering_dates, action in self.choices:
-            dates = context.matching_dates(filtering_dates, argument)
+        for filters, action in self.choices:
+            dates = context.matching_dates(filters, argument)
             if len(dates) == 0:
                 continue
             results.append(action(context, dates))
@@ -102,7 +118,7 @@ def __call__(self, context, argument):
 
     def dump(self, dumper):
         return dumper.concat(
-            {filtering_dates.dump(dumper): action.dump(dumper) for filtering_dates, action in self.choices}
+            {list(filters.values())[0].dump(dumper): action.dump(dumper) for filters, action in self.choices}
         )
 
 

src/anemoi/datasets/create/trajectories/context.py

@@ -74,26 +74,41 @@ def __init__(self, recipe: Any) -> None:
     def create_result(self, argument: Any, data: Any) -> TrajectoryGriddedResult:
         return TrajectoryGriddedResult(self, argument, data)
 
-    def matching_dates(self, filtering_dates: Any, group_of_dates: Any) -> Any:
-        """Find dates that match between filtering_dates and group_of_dates.
+    def matching_dates(self, filters: dict, group_of_dates: Any) -> Any:
+        """Find dates that match between filters and group_of_dates.
 
         Parameters
         ----------
-        filtering_dates : Any
-            The dates to filter by.
+        filters : dict
+            A dict mapping filter keys to DatesProvider objects.
+            Trajectory layouts support ``'base_dates'`` and ``'steps'``.
         group_of_dates : ForecastDates
             The ``(valid_time, basetime)`` pairs for this group.
 
         Returns
         -------
         ForecastDates
-            A ForecastDates containing only the pairs whose basetime is in
-            ``filtering_dates``.
+            A ForecastDates containing only the matching pairs.
         """
+        unsupported = set(filters) - {"base_dates", "steps"}
+        if unsupported:
+            raise ValueError(
+                f"Trajectory layout does not support filtering by {unsupported}. "
+                "Use 'base_dates' and/or 'steps' instead."
+            )
+
         from anemoi.datasets.create.arguments import ForecastDates
 
-        filter_set = set(filtering_dates)
-        matched = [(vt, bt) for vt, bt in group_of_dates if bt in filter_set]
+        matched = list(group_of_dates)
+
+        if "base_dates" in filters:
+            base_dates_set = set(filters["base_dates"])
+            matched = [(vt, bt) for vt, bt in matched if bt in base_dates_set]
+
+        if "steps" in filters:
+            steps_set = set(filters["steps"])
+            matched = [(vt, bt) for vt, bt in matched if (vt - bt) in steps_set]
+
         return ForecastDates(matched)
 
     def origin(self, data: Any, action: Any, action_arguments: Any) -> Any: