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

anemoi/datasets/create/input/empty.py

@@ -0,0 +1,42 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+from functools import cached_property
+
+from anemoi.utils.dates import as_datetime as as_datetime
+from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
+
+from anemoi.datasets.dates import DatesProvider as DatesProvider
+from anemoi.datasets.fields import FieldArray as FieldArray
+from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
+
+from .misc import assert_fieldlist
+from .result import Result
+from .trace import trace_datasource
+
+LOG = logging.getLogger(__name__)
+
+
+class EmptyResult(Result):
+    empty = True
+
+    def __init__(self, context, action_path, dates):
+        super().__init__(context, action_path + ["empty"], dates)
+
+    @cached_property
+    @assert_fieldlist
+    @trace_datasource
+    def datasource(self):
+        from earthkit.data import from_source
+
+        return from_source("empty")
+
+    @property
+    def variables(self):
+        return []

anemoi/datasets/create/input/filter.py

@@ -0,0 +1,76 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+from functools import cached_property
+
+from anemoi.utils.dates import as_datetime as as_datetime
+from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
+
+from anemoi.datasets.dates import DatesProvider as DatesProvider
+from anemoi.datasets.fields import FieldArray as FieldArray
+from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
+
+from ..functions import import_function
+from .function import FunctionContext
+from .misc import _tidy
+from .misc import assert_fieldlist
+from .step import StepAction
+from .step import StepResult
+from .template import notify_result
+from .trace import trace_datasource
+
+LOG = logging.getLogger(__name__)
+
+
+class FilterStepResult(StepResult):
+    @property
+    @notify_result
+    @assert_fieldlist
+    @trace_datasource
+    def datasource(self):
+        ds = self.upstream_result.datasource
+        ds = ds.sel(**self.action.kwargs)
+        return _tidy(ds)
+
+
+class FilterStepAction(StepAction):
+    result_class = FilterStepResult
+
+
+class StepFunctionResult(StepResult):
+    @cached_property
+    @assert_fieldlist
+    @notify_result
+    @trace_datasource
+    def datasource(self):
+        try:
+            return _tidy(
+                self.action.function(
+                    FunctionContext(self),
+                    self.upstream_result.datasource,
+                    *self.action.args[1:],
+                    **self.action.kwargs,
+                )
+            )
+
+        except Exception:
+            LOG.error(f"Error in {self.action.name}", exc_info=True)
+            raise
+
+    def _trace_datasource(self, *args, **kwargs):
+        return f"{self.action.name}({self.group_of_dates})"
+
+
+class FunctionStepAction(StepAction):
+    result_class = StepFunctionResult
+
+    def __init__(self, context, action_path, previous_step, *args, **kwargs):
+        super().__init__(context, action_path, previous_step, *args, **kwargs)
+        self.name = args[0]
+        self.function = import_function(self.name, "filters")

anemoi/datasets/create/input/function.py

@@ -0,0 +1,122 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+from functools import cached_property
+
+from anemoi.utils.dates import as_datetime as as_datetime
+from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
+
+from anemoi.datasets.dates import DatesProvider as DatesProvider
+from anemoi.datasets.fields import FieldArray as FieldArray
+from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
+
+from ..functions import import_function
+from .action import Action
+from .misc import _tidy
+from .misc import assert_fieldlist
+from .result import Result
+from .template import notify_result
+from .template import resolve
+from .template import substitute
+from .trace import trace
+from .trace import trace_datasource
+from .trace import trace_select
+
+LOG = logging.getLogger(__name__)
+
+
+class FunctionContext:
+    """A FunctionContext is passed to all functions, it will be used to pass information
+    to the functions from the other actions and filters and results.
+    """
+
+    def __init__(self, owner):
+        self.owner = owner
+        self.use_grib_paramid = owner.context.use_grib_paramid
+
+    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
+
+    @property
+    def partial_ok(self):
+        return self.owner.group_of_dates.partial_ok
+
+
+class FunctionAction(Action):
+    def __init__(self, context, action_path, _name, **kwargs):
+        super().__init__(context, action_path, **kwargs)
+        self.name = _name
+
+    @trace_select
+    def select(self, group_of_dates):
+        return FunctionResult(self.context, self.action_path, group_of_dates, action=self)
+
+    @property
+    def function(self):
+        # name, delta = parse_function_name(self.name)
+        return import_function(self.name, "sources")
+
+    def __repr__(self):
+        content = ""
+        content += ",".join([self._short_str(a) for a in self.args])
+        content += " ".join([self._short_str(f"{k}={v}") for k, v in self.kwargs.items()])
+        content = self._short_str(content)
+        return super().__repr__(_inline_=content, _indent_=" ")
+
+    def _trace_select(self, group_of_dates):
+        return f"{self.name}({group_of_dates})"
+
+
+class FunctionResult(Result):
+    def __init__(self, context, action_path, group_of_dates, action):
+        super().__init__(context, action_path, group_of_dates)
+        assert isinstance(action, Action), type(action)
+        self.action = action
+
+        self.args, self.kwargs = substitute(context, (self.action.args, self.action.kwargs))
+
+    def _trace_datasource(self, *args, **kwargs):
+        return f"{self.action.name}({self.group_of_dates})"
+
+    @cached_property
+    @assert_fieldlist
+    @notify_result
+    @trace_datasource
+    def datasource(self):
+        args, kwargs = resolve(self.context, (self.args, self.kwargs))
+
+        try:
+            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}({self.group_of_dates})"
+        except Exception:
+            return f"{self.__class__.__name__}(unitialised)"
+
+    @property
+    def function(self):
+        raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")

anemoi/datasets/create/input/join.py

@@ -0,0 +1,57 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+from functools import cached_property
+
+from .action import Action
+from .action import action_factory
+from .empty import EmptyResult
+from .misc import _tidy
+from .misc import assert_fieldlist
+from .result import Result
+from .template import notify_result
+from .trace import trace_datasource
+from .trace import trace_select
+
+LOG = logging.getLogger(__name__)
+
+
+class JoinResult(Result):
+    def __init__(self, context, action_path, group_of_dates, results, **kwargs):
+        super().__init__(context, action_path, group_of_dates)
+        self.results = [r for r in results if not r.empty]
+
+    @cached_property
+    @assert_fieldlist
+    @notify_result
+    @trace_datasource
+    def datasource(self):
+        ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
+        for i in self.results:
+            ds += i.datasource
+        return _tidy(ds)
+
+    def __repr__(self):
+        content = "\n".join([str(i) for i in self.results])
+        return super().__repr__(content)
+
+
+class JoinAction(Action):
+    def __init__(self, context, action_path, *configs):
+        super().__init__(context, action_path, *configs)
+        self.actions = [action_factory(c, context, action_path + [str(i)]) for i, c in enumerate(configs)]
+
+    def __repr__(self):
+        content = "\n".join([str(i) for i in self.actions])
+        return super().__repr__(content)
+
+    @trace_select
+    def select(self, group_of_dates):
+        results = [a.select(group_of_dates) for a in self.actions]
+        return JoinResult(self.context, self.action_path, group_of_dates, results)

anemoi/datasets/create/input/misc.py

@@ -0,0 +1,85 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+from functools import wraps
+
+from anemoi.utils.dates import as_datetime as as_datetime
+from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
+from earthkit.data.core.fieldlist import MultiFieldList
+from earthkit.data.indexing.fieldlist import FieldList
+
+from anemoi.datasets.dates import DatesProvider as DatesProvider
+from anemoi.datasets.fields import FieldArray as FieldArray
+from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
+
+from ..functions import import_function
+
+LOG = logging.getLogger(__name__)
+
+
+def parse_function_name(name):
+
+    if name.endswith("h") and name[:-1].isdigit():
+
+        if "-" in name:
+            name, delta = name.split("-")
+            sign = -1
+
+        elif "+" in name:
+            name, delta = name.split("+")
+            sign = 1
+
+        else:
+            return name, None
+
+        assert delta[-1] == "h", (name, delta)
+        delta = sign * int(delta[:-1])
+        return name, delta
+
+    return name, None
+
+
+def is_function(name, kind):
+    name, _ = parse_function_name(name)
+    try:
+        import_function(name, kind)
+        return True
+    except ImportError as e:
+        print(e)
+        return False
+
+
+def assert_fieldlist(method):
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        result = method(self, *args, **kwargs)
+        assert isinstance(result, FieldList), type(result)
+        return result
+
+    return wrapper
+
+
+def assert_is_fieldlist(obj):
+    assert isinstance(obj, FieldList), type(obj)
+
+
+def _flatten(ds):
+    if isinstance(ds, MultiFieldList):
+        return [_tidy(f) for s in ds._indexes for f in _flatten(s)]
+    return [ds]
+
+
+def _tidy(ds, indent=0):
+    if isinstance(ds, MultiFieldList):
+
+        sources = [s for s in _flatten(ds) if len(s) > 0]
+        if len(sources) == 1:
+            return sources[0]
+        return MultiFieldList(sources)
+    return ds

anemoi/datasets/create/input/pipe.py

@@ -0,0 +1,33 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+
+from .action import Action
+from .action import action_factory
+from .step import step_factory
+from .trace import trace_select
+
+LOG = logging.getLogger(__name__)
+
+
+class PipeAction(Action):
+    def __init__(self, context, action_path, *configs):
+        super().__init__(context, action_path, *configs)
+        assert len(configs) > 1, configs
+        current = action_factory(configs[0], context, action_path + ["0"])
+        for i, c in enumerate(configs[1:]):
+            current = step_factory(c, context, action_path + [str(i + 1)], previous_step=current)
+        self.last_step = current
+
+    @trace_select
+    def select(self, group_of_dates):
+        return self.last_step.select(group_of_dates)
+
+    def __repr__(self):
+        return super().__repr__(self.last_step)

anemoi/datasets/create/input/repeated_dates.py

@@ -0,0 +1,217 @@
+# (C) Copyright 2023 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import logging
+from collections import defaultdict
+
+import numpy as np
+from anemoi.utils.dates import as_datetime
+from anemoi.utils.dates import frequency_to_timedelta
+
+from anemoi.datasets.fields import FieldArray
+from anemoi.datasets.fields import NewValidDateTimeField
+
+from .action import Action
+from .action import action_factory
+from .join import JoinResult
+from .result import Result
+from .trace import trace_select
+
+LOG = logging.getLogger(__name__)
+
+
+class DateMapper:
+
+    @staticmethod
+    def from_mode(mode, source, config):
+
+        MODES = dict(
+            closest=DateMapperClosest,
+            climatology=DateMapperClimatology,
+            constant=DateMapperConstant,
+        )
+
+        if mode not in MODES:
+            raise ValueError(f"Invalid mode for DateMapper: {mode}")
+
+        return MODES[mode](source, **config)
+
+
+class DateMapperClosest(DateMapper):
+    def __init__(self, source, frequency="1h", maximum="30d", skip_all_nans=False):
+        self.source = source
+        self.maximum = frequency_to_timedelta(maximum)
+        self.frequency = frequency_to_timedelta(frequency)
+        self.skip_all_nans = skip_all_nans
+        self.tried = set()
+        self.found = set()
+
+    def transform(self, group_of_dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        asked_dates = list(group_of_dates)
+        if not asked_dates:
+            return []
+
+        to_try = set()
+        for date in asked_dates:
+            start = date
+            while start >= date - self.maximum:
+                to_try.add(start)
+                start -= self.frequency
+
+            end = date
+            while end <= date + self.maximum:
+                to_try.add(end)
+                end += self.frequency
+
+        to_try = sorted(to_try - self.tried)
+
+        if to_try:
+            result = self.source.select(
+                GroupOfDates(
+                    sorted(to_try),
+                    group_of_dates.provider,
+                    partial_ok=True,
+                )
+            )
+
+            for f in result.datasource:
+                # We could keep the fields in a dictionary, but we don't want to keep the fields in memory
+                date = as_datetime(f.metadata("valid_datetime"))
+
+                if self.skip_all_nans:
+                    if np.isnan(f.to_numpy()).all():
+                        LOG.warning(f"Skipping {date} because all values are NaN")
+                        continue
+
+                self.found.add(date)
+
+            self.tried.update(to_try)
+
+        new_dates = defaultdict(list)
+
+        for date in asked_dates:
+            best = None
+            for found_date in sorted(self.found):
+                delta = abs(date - found_date)
+                # With < we prefer the first date
+                # With <= we prefer the last date
+                if best is None or delta <= best[0]:
+                    best = delta, found_date
+            new_dates[best[1]].append(date)
+
+        for date, dates in new_dates.items():
+            yield (
+                GroupOfDates([date], group_of_dates.provider),
+                GroupOfDates(dates, group_of_dates.provider),
+            )
+
+
+class DateMapperClimatology(DateMapper):
+    def __init__(self, source, year, day):
+        self.year = year
+        self.day = day
+
+    def transform(self, group_of_dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        dates = list(group_of_dates)
+        if not dates:
+            return []
+
+        new_dates = defaultdict(list)
+        for date in dates:
+            new_date = date.replace(year=self.year, day=self.day)
+            new_dates[new_date].append(date)
+
+        for date, dates in new_dates.items():
+            yield (
+                GroupOfDates([date], group_of_dates.provider),
+                GroupOfDates(dates, group_of_dates.provider),
+            )
+
+
+class DateMapperConstant(DateMapper):
+    def __init__(self, source, date=None):
+        self.source = source
+        self.date = date
+
+    def transform(self, group_of_dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        if self.date is None:
+            return [
+                (
+                    GroupOfDates([], group_of_dates.provider),
+                    group_of_dates,
+                )
+            ]
+
+        return [
+            (
+                GroupOfDates([self.date], group_of_dates.provider),
+                group_of_dates,
+            )
+        ]
+
+
+class DateMapperResult(Result):
+    def __init__(
+        self,
+        context,
+        action_path,
+        group_of_dates,
+        source_result,
+        mapper,
+        original_group_of_dates,
+    ):
+        super().__init__(context, action_path, group_of_dates)
+
+        self.source_results = source_result
+        self.mapper = mapper
+        self.original_group_of_dates = original_group_of_dates
+
+    @property
+    def datasource(self):
+        result = []
+
+        for field in self.source_results.datasource:
+            for date in self.original_group_of_dates:
+                result.append(NewValidDateTimeField(field, date))
+
+        return FieldArray(result)
+
+
+class RepeatedDatesAction(Action):
+    def __init__(self, context, action_path, source, mode, **kwargs):
+        super().__init__(context, action_path, source, mode, **kwargs)
+
+        self.source = action_factory(source, context, action_path + ["source"])
+        self.mapper = DateMapper.from_mode(mode, self.source, kwargs)
+
+    @trace_select
+    def select(self, group_of_dates):
+        results = []
+        for one_date_group, many_dates_group in self.mapper.transform(group_of_dates):
+            results.append(
+                DateMapperResult(
+                    self.context,
+                    self.action_path,
+                    one_date_group,
+                    self.source.select(one_date_group),
+                    self.mapper,
+                    many_dates_group,
+                )
+            )
+
+        return JoinResult(self.context, self.action_path, group_of_dates, results)
+
+    def __repr__(self):
+        return f"MultiDateMatchAction({self.source}, {self.mapper})"