anemoi-datasets 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

anemoi/datasets/create/input.py

@@ -23,7 +23,7 @@ from earthkit.data.core.fieldlist import FieldList
 from earthkit.data.core.fieldlist import MultiFieldList
 from earthkit.data.core.order import build_remapping
 
-from anemoi.datasets.dates import Dates
+from anemoi.datasets.dates import DatesProvider
 
 from .functions import import_function
 from .template import Context
@@ -75,7 +75,7 @@ def time_delta_to_string(delta):
 
 
 def is_function(name, kind):
-    name, delta = parse_function_name(name)  # noqa
+    name, _ = parse_function_name(name)
     try:
         import_function(name, kind)
         return True
@@ -106,30 +106,32 @@ def _data_request(data):
     area = grid = None
 
     for field in data:
-        if not hasattr(field, "as_mars"):
-            continue
-
-        if date is None:
-            date = field.datetime()["valid_time"]
-
-        if field.datetime()["valid_time"] != date:
-            continue
+        try:
+            if date is None:
+                date = field.datetime()["valid_time"]
 
-        as_mars = field.metadata(namespace="mars")
-        step = as_mars.get("step")
-        levtype = as_mars.get("levtype", "sfc")
-        param = as_mars["param"]
-        levelist = as_mars.get("levelist", None)
-        area = field.mars_area
-        grid = field.mars_grid
+            if field.datetime()["valid_time"] != date:
+                continue
 
-        if levelist is None:
-            params_levels[levtype].add(param)
-        else:
-            params_levels[levtype].add((param, levelist))
+            as_mars = field.metadata(namespace="mars")
+            if not as_mars:
+                continue
+            step = as_mars.get("step")
+            levtype = as_mars.get("levtype", "sfc")
+            param = as_mars["param"]
+            levelist = as_mars.get("levelist", None)
+            area = field.mars_area
+            grid = field.mars_grid
+
+            if levelist is None:
+                params_levels[levtype].add(param)
+            else:
+                params_levels[levtype].add((param, levelist))
 
-        if step:
-            params_steps[levtype].add((param, step))
+            if step:
+                params_steps[levtype].add((param, step))
+        except Exception:
+            LOG.error(f"Error in retrieving metadata (cannot build data request info) for {field}", exc_info=True)
 
     def sort(old_dic):
         new_dic = {}
@@ -202,11 +204,15 @@ class Result:
     _coords_already_built = False
 
     def __init__(self, context, action_path, dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        assert isinstance(dates, GroupOfDates), dates
+
         assert isinstance(context, ActionContext), type(context)
         assert isinstance(action_path, list), action_path
 
         self.context = context
-        self.dates = dates
+        self.group_of_dates = dates
         self.action_path = action_path
 
     @property
@@ -288,7 +294,6 @@ class Result:
                 names += list(a.keys())
 
         print(f"Building a {len(names)}D hypercube using", names)
-
         ds = ds.order_by(*args, remapping=remapping, patches=patches)
         user_coords = ds.unique_values(*names, remapping=remapping, patches=patches, progress_bar=False)
 
@@ -404,10 +409,10 @@ class Result:
         more += ",".join([f"{k}={v}"[:5000] for k, v in kwargs.items()])
 
         dates = " no-dates"
-        if self.dates is not None:
-            dates = f" {len(self.dates)} dates"
+        if self.group_of_dates is not None:
+            dates = f" {len(self.group_of_dates)} dates"
             dates += " ("
-            dates += "/".join(d.strftime("%Y-%m-%d:%H") for d in self.dates)
+            dates += "/".join(d.strftime("%Y-%m-%d:%H") for d in self.group_of_dates)
             if len(dates) > 100:
                 dates = dates[:100] + "..."
             dates += ")"
@@ -422,7 +427,7 @@ class Result:
         raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
 
     def _trace_datasource(self, *args, **kwargs):
-        return f"{self.__class__.__name__}({shorten(self.dates)})"
+        return f"{self.__class__.__name__}({self.group_of_dates})"
 
     def build_coords(self):
         if self._coords_already_built:
@@ -512,7 +517,7 @@ class Result:
     @cached_property
     def shape(self):
         return [
-            len(self.dates),
+            len(self.group_of_dates),
             len(self.variables),
             len(self.ensembles),
             len(self.grid_values),
@@ -521,7 +526,7 @@ class Result:
     @cached_property
     def coords(self):
         return {
-            "dates": self.dates,
+            "dates": list(self.group_of_dates),
             "variables": self.variables,
             "ensembles": self.ensembles,
             "values": self.grid_values,
@@ -572,7 +577,7 @@ class FunctionResult(Result):
         self.args, self.kwargs = substitute(context, (self.action.args, self.action.kwargs))
 
     def _trace_datasource(self, *args, **kwargs):
-        return f"{self.action.name}({shorten(self.dates)})"
+        return f"{self.action.name}({self.group_of_dates})"
 
     @cached_property
     @assert_fieldlist
@@ -582,14 +587,21 @@ class FunctionResult(Result):
         args, kwargs = resolve(self.context, (self.args, self.kwargs))
 
         try:
-            return _tidy(self.action.function(FunctionContext(self), self.dates, *args, **kwargs))
+            return _tidy(
+                self.action.function(
+                    FunctionContext(self),
+                    list(self.group_of_dates),  # Will provide a list of datetime objects
+                    *args,
+                    **kwargs,
+                )
+            )
         except Exception:
             LOG.error(f"Error in {self.action.function.__name__}", exc_info=True)
             raise
 
     def __repr__(self):
         try:
-            return f"{self.action.name}({shorten(self.dates)})"
+            return f"{self.action.name}({self.group_of_dates})"
         except Exception:
             return f"{self.__class__.__name__}(unitialised)"
 
@@ -608,7 +620,7 @@ class JoinResult(Result):
     @notify_result
     @trace_datasource
     def datasource(self):
-        ds = EmptyResult(self.context, self.action_path, self.dates).datasource
+        ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
         for i in self.results:
             ds += i.datasource
         return _tidy(ds)
@@ -823,7 +835,7 @@ class ConcatResult(Result):
     @notify_result
     @trace_datasource
     def datasource(self):
-        ds = EmptyResult(self.context, self.action_path, self.dates).datasource
+        ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
         for i in self.results:
             ds += i.datasource
         return _tidy(ds)
@@ -903,7 +915,7 @@ class ConcatAction(Action):
             cfg = deepcopy(cfg)
             dates_cfg = cfg.pop("dates")
             assert isinstance(dates_cfg, dict), dates_cfg
-            filtering_dates = Dates.from_config(**dates_cfg)
+            filtering_dates = DatesProvider.from_config(**dates_cfg)
             action = action_factory(cfg, context, action_path + [str(i)])
             parts.append((filtering_dates, action))
         self.parts = parts
@@ -914,9 +926,11 @@ class ConcatAction(Action):
 
     @trace_select
     def select(self, dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
         results = []
         for filtering_dates, action in self.parts:
-            newdates = sorted(set(dates) & set(filtering_dates))
+            newdates = GroupOfDates(sorted(set(dates) & set(filtering_dates)), dates.provider)
             if newdates:
                 results.append(action.select(newdates))
         if not results:
@@ -952,8 +966,10 @@ def action_factory(config, context, action_path):
 
     if isinstance(config[key], list):
         args, kwargs = config[key], {}
-    if isinstance(config[key], dict):
+    elif isinstance(config[key], dict):
         args, kwargs = [], config[key]
+    else:
+        raise ValueError(f"Invalid input config {config[key]} ({type(config[key])}")
 
     cls = {
         # "date_shift": DateShiftAction,
@@ -1020,6 +1036,13 @@ class FunctionContext:
     def trace(self, emoji, *args):
         trace(emoji, *args)
 
+    def info(self, *args, **kwargs):
+        LOG.info(*args, **kwargs)
+
+    @property
+    def dates_provider(self):
+        return self.owner.group_of_dates.provider
+
 
 class ActionContext(Context):
     def __init__(self, /, order_by, flatten_grid, remapping, use_grib_paramid):

anemoi/datasets/create/persistent.py

@@ -68,7 +68,7 @@ class PersistentDict:
         path = os.path.join(self.dirname, f"{h}.pickle")
 
         if os.path.exists(path):
-            LOG.warn(f"{path} already exists")
+            LOG.warning(f"{path} already exists")
 
         tmp_path = path + f".tmp-{os.getpid()}-on-{socket.gethostname()}"
         with open(tmp_path, "wb") as f:

anemoi/datasets/create/statistics/__init__.py

@@ -79,6 +79,37 @@ def to_datetimes(dates):
     return [to_datetime(d) for d in dates]
 
 
+def fix_variance(x, name, count, sums, squares):
+    assert count.shape == sums.shape == squares.shape
+    assert isinstance(x, float)
+
+    mean = sums / count
+    assert mean.shape == count.shape
+
+    if x >= 0:
+        return x
+
+    LOG.warning(f"Negative variance for {name=}, variance={x}")
+    magnitude = np.sqrt((squares / count + mean * mean) / 2)
+    LOG.warning(f"square / count - mean * mean =  {squares/count} - {mean*mean} = {squares/count - mean*mean}")
+    LOG.warning(f"Variable span order of magnitude is {magnitude}.")
+    LOG.warning(f"Count is {count}.")
+
+    variances = squares / count - mean * mean
+    assert variances.shape == squares.shape == mean.shape
+    if all(variances >= 0):
+        LOG.warning(f"All individual variances for {name} are positive, setting variance to 0.")
+        return 0
+
+    # if abs(x) < magnitude * 1e-6 and abs(x) < range * 1e-6:
+    #     LOG.warning("Variance is negative but very small.")
+    #     variances = squares / count - mean * mean
+    #     return 0
+
+    LOG.warning(f"ERROR at least one individual variance is negative ({np.nanmin(variances)}).")
+    return x
+
+
 def check_variance(x, variables_names, minimum, maximum, mean, count, sums, squares):
     if (x >= 0).all():
         return
@@ -292,39 +323,24 @@ class StatAggregator:
     def aggregate(self):
         minimum = np.nanmin(self.minimum, axis=0)
         maximum = np.nanmax(self.maximum, axis=0)
+
         sums = np.nansum(self.sums, axis=0)
         squares = np.nansum(self.squares, axis=0)
         count = np.nansum(self.count, axis=0)
         has_nans = np.any(self.has_nans, axis=0)
-        mean = sums / count
+        assert sums.shape == count.shape == squares.shape == minimum.shape == maximum.shape
 
-        assert sums.shape == count.shape == squares.shape == mean.shape == minimum.shape == maximum.shape
+        mean = sums / count
+        assert mean.shape == minimum.shape
 
         x = squares / count - mean * mean
-
-        # def fix_variance(x, name, minimum, maximum, mean, count, sums, squares):
-        #     assert x.shape == minimum.shape == maximum.shape == mean.shape == count.shape == sums.shape == squares.shape
-        #     assert x.shape == (1,)
-        #     x, minimum, maximum, mean, count, sums, squares = x[0], minimum[0], maximum[0], mean[0], count[0], sums[0], squares[0]
-        #     if x >= 0:
-        #         return x
-        #
-        #     order = np.sqrt((squares / count + mean * mean)/2)
-        #     range = maximum - minimum
-        #     LOG.warning(f"Negative variance for {name=}, variance={x}")
-        #     LOG.warning(f"square / count - mean * mean =  {squares / count} - {mean * mean} = {squares / count - mean * mean}")
-        #     LOG.warning(f"Variable order of magnitude is {order}.")
-        #     LOG.warning(f"Range is {range} ({maximum=} - {minimum=}).")
-        #     LOG.warning(f"Count is {count}.")
-        #     if abs(x) < order * 1e-6 and abs(x) < range * 1e-6:
-        #         LOG.warning(f"Variance is negative but very small, setting to 0.")
-        #         return x*0
-        #     return x
+        assert x.shape == minimum.shape
 
         for i, name in enumerate(self.variables_names):
             # remove negative variance due to numerical errors
-            # Not needed for now, fix_variance is disabled
-            # x[i] = fix_variance(x[i:i+1], name, minimum[i:i+1], maximum[i:i+1], mean[i:i+1], count[i:i+1], sums[i:i+1], squares[i:i+1])
+            x[i] = fix_variance(x[i], name, self.count[i : i + 1], self.sums[i : i + 1], self.squares[i : i + 1])
+
+        for i, name in enumerate(self.variables_names):
             check_variance(
                 x[i : i + 1],
                 [name],

anemoi/datasets/create/utils.py

@@ -7,6 +7,7 @@
 # nor does it submit to any jurisdiction.
 #
 
+import datetime
 import os
 from contextlib import contextmanager
 
@@ -61,10 +62,13 @@ def make_list_int(value):
 
 
 def normalize_and_check_dates(dates, start, end, frequency, dtype="datetime64[s]"):
-    assert isinstance(frequency, int), frequency
+
+    dates = [d.hdate if hasattr(d, "hdate") else d for d in dates]
+
+    assert isinstance(frequency, datetime.timedelta), frequency
     start = np.datetime64(start)
     end = np.datetime64(end)
-    delta = np.timedelta64(frequency, "h")
+    delta = np.timedelta64(frequency)
 
     res = []
     while start <= end:

anemoi/datasets/data/__init__.py

@@ -27,6 +27,7 @@ class MissingDateError(Exception):
 
 def open_dataset(*args, **kwargs):
     ds = _open_dataset(*args, **kwargs)
+    ds = ds.mutate()
     ds.arguments = {"args": args, "kwargs": kwargs}
     ds._check()
     return ds

anemoi/datasets/data/concat.py

@@ -9,6 +9,7 @@ import logging
 from functools import cached_property
 
 import numpy as np
+from anemoi.utils.dates import frequency_to_timedelta
 
 from .debug import Node
 from .debug import debug_indexing
@@ -102,20 +103,63 @@ class Concat(ConcatMixin, Combined):
     def tree(self):
         return Node(self, [d.tree() for d in self.datasets])
 
+    @classmethod
+    def check_dataset_compatibility(cls, datasets, fill_missing_gaps=False):
+        # Study the dates
+        ranges = [(d.dates[0].astype(object), d.dates[-1].astype(object)) for d in datasets]
 
-def concat_factory(args, kwargs, zarr_root):
+        # Make sure the dates are disjoint
+        for i in range(len(ranges)):
+            r = ranges[i]
+            for j in range(i + 1, len(ranges)):
+                s = ranges[j]
+                if r[0] <= s[0] <= r[1] or r[0] <= s[1] <= r[1]:
+                    raise ValueError(f"Overlapping dates: {r} and {s} ({datasets[i]} {datasets[j]})")
+
+        # For now we should have the datasets in order with no gaps
+
+        frequency = frequency_to_timedelta(datasets[0].frequency)
+        result = []
+
+        for i in range(len(ranges) - 1):
+            result.append(datasets[i])
+            r = ranges[i]
+            s = ranges[i + 1]
+            if r[1] + frequency != s[0]:
+                if fill_missing_gaps:
+                    from .missing import MissingDataset
+
+                    result.append(MissingDataset(datasets[i], r[1] + frequency, s[0] - frequency))
+                else:
+                    r = [str(e) for e in r]
+                    s = [str(e) for e in s]
+                    raise ValueError(
+                        "Datasets must be sorted by dates, with no gaps: "
+                        f"{r} and {s} ({datasets[i]} {datasets[i+1]})"
+                    )
+
+        result.append(datasets[-1])
+        assert len(result) >= len(datasets), (len(result), len(datasets))
+
+        return result
+
+
+def concat_factory(args, kwargs):
 
     datasets = kwargs.pop("concat")
+    fill_missing_gaps = kwargs.pop("fill_missing_gaps", False)
     assert isinstance(datasets, (list, tuple))
     assert len(args) == 0
 
     assert isinstance(datasets, (list, tuple))
 
-    datasets = [_open(e, zarr_root) for e in datasets]
+    datasets = [_open(e) for e in datasets]
 
     if len(datasets) == 1:
         return datasets[0]._subset(**kwargs)
 
     datasets, kwargs = _auto_adjust(datasets, kwargs)
 
+    datasets = Concat.check_dataset_compatibility(datasets, fill_missing_gaps)
+
     return Concat(datasets)._subset(**kwargs)