Feature/merge (#126)

* save masks to checkpoint
* force np.datetime64 is seconds
* Call filters from anemoi-transform
* when merging datasets, consider missing dates
* add gcd for frequency

---------

Co-authored-by: Florian Pinault <[email protected]>

Author: b8raoult

CHANGELOG.md

@@ -10,9 +10,13 @@ Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-datasets/compare/0.5.8...HEAD)
 
-### Changed
+
+### Added
+
+- Call filters from anemoi-transform
 - make test optional when adls is not installed Pull request #110
 
+
 ## [0.5.8](https://github.com/ecmwf/anemoi-datasets/compare/0.5.7...0.5.8) - 2024-10-26
 
 ### Changed

src/anemoi/datasets/create/functions/__init__.py

@@ -21,6 +21,7 @@ def assert_is_fieldlist(obj):
 
 def import_function(name, kind):
 
+    from anemoi.transform.filters import filter_registry
     from anemoi.transforms import Transform as Transform
 
     name = name.replace("-", "_")
@@ -32,8 +33,21 @@ def import_function(name, kind):
     if name in plugins:
         return plugins[name].load()
 
-    module = importlib.import_module(
-        f".{kind}.{name}",
-        package=__name__,
-    )
-    return module.execute
+    try:
+        module = importlib.import_module(
+            f".{kind}.{name}",
+            package=__name__,
+        )
+        return module.execute
+    except ModuleNotFoundError:
+        pass
+
+    if kind == "filters":
+        if filter_registry.lookup(name, return_none=True):
+
+            def proc(context, data, *args, **kwargs):
+                return filter_registry.create(name, *args, **kwargs)(data)
+
+            return proc
+
+    raise ValueError(f"Unknown {kind} '{name}'")

src/anemoi/datasets/create/functions/filters/rename.py

@@ -25,7 +25,9 @@ class RenamedFieldMapping:
     def __init__(self, field, what, renaming):
         self.field = field
         self.what = what
-        self.renaming = renaming
+        self.renaming = {}
+        for k, v in renaming.items():
+            self.renaming[k] = {str(a): str(b) for a, b in v.items()}
 
     def metadata(self, key=None, **kwargs):
         if key is None:

src/anemoi/datasets/data/merge.py

@@ -40,16 +40,29 @@ def __init__(self, datasets, allow_gaps_in_dates=False):
 
         self.allow_gaps_in_dates = allow_gaps_in_dates
 
-        dates = dict()
+        dates = dict()  # date -> (dataset_index, date_index)
 
         for i, d in enumerate(datasets):
             for j, date in enumerate(d.dates):
                 date = date.astype(object)
                 if date in dates:
-                    d1 = datasets[dates[date][0]]
-                    d2 = datasets[i]
+
+                    d1 = datasets[dates[date][0]]  # Selected
+                    d2 = datasets[i]  # The new one
+
+                    if j in d2.missing:
+                        # LOG.warning(f"Duplicate date {date} found in datasets {d1} and {d2}, but {date} is missing in {d}, ignoring")
+                        continue
+
+                    k = dates[date][1]
+                    if k in d1.missing:
+                        # LOG.warning(f"Duplicate date {date} found in datasets {d1} and {d2}, but {date} is missing in {d}, ignoring")
+                        dates[date] = (i, j)  # Replace the missing date with the new one
+                        continue
+
                     raise ValueError(f"Duplicate date {date} found in datasets {d1} and {d2}")
-                dates[date] = (i, j)
+                else:
+                    dates[date] = (i, j)
 
         all_dates = sorted(dates)
         start = all_dates[0]

src/anemoi/datasets/dates/__init__.py

@@ -10,6 +10,8 @@
 
 import datetime
 import warnings
+from functools import reduce
+from math import gcd
 
 # from anemoi.utils.dates import as_datetime
 from anemoi.utils.dates import DateTimes
@@ -195,18 +197,16 @@ def __init__(self, start, end, steps=[0], years=20, **kwargs):
 
         dates = sorted(dates)
 
-        mindelta = None
+        deltas = set()
         for a, b in zip(dates, dates[1:]):
             delta = b - a
             assert isinstance(delta, datetime.timedelta), delta
-            if mindelta is None:
-                mindelta = delta
-            else:
-                mindelta = min(mindelta, delta)
+            deltas.add(delta)
 
+        mindelta_seconds = reduce(gcd, [int(delta.total_seconds()) for delta in deltas])
+        mindelta = datetime.timedelta(seconds=mindelta_seconds)
         self.frequency = mindelta
         assert mindelta.total_seconds() > 0, mindelta
-
         print("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", dates[0], dates[-1], mindelta)
 
         # Use all values between start and end by frequency, and set the ones that are missing