anemoi-datasets 0.4.5__py3-none-any.whl → 0.5.5__py3-none-any.whl

anemoi/datasets/create/input/data_sources.py

@@ -0,0 +1,71 @@
+# (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 .action import Action
+from .action import action_factory
+from .misc import _tidy
+from .result import Result
+
+LOG = logging.getLogger(__name__)
+
+
+class DataSourcesAction(Action):
+    def __init__(self, context, action_path, sources, input):
+        super().__init__(context, ["data_sources"], *sources)
+        if isinstance(sources, dict):
+            configs = [(str(k), c) for k, c in sources.items()]
+        elif isinstance(sources, list):
+            configs = [(str(i), c) for i, c in enumerate(sources)]
+        else:
+            raise ValueError(f"Invalid data_sources, expecting list or dict, got {type(sources)}: {sources}")
+
+        self.sources = [action_factory(config, context, ["data_sources"] + [a_path]) for a_path, config in configs]
+        self.input = action_factory(input, context, ["input"])
+
+    def select(self, group_of_dates):
+        sources_results = [a.select(group_of_dates) for a in self.sources]
+        return DataSourcesResult(
+            self.context,
+            self.action_path,
+            group_of_dates,
+            self.input.select(group_of_dates),
+            sources_results,
+        )
+
+    def __repr__(self):
+        content = "\n".join([str(i) for i in self.sources])
+        return super().__repr__(content)
+
+
+class DataSourcesResult(Result):
+    def __init__(self, context, action_path, dates, input_result, sources_results):
+        super().__init__(context, action_path, dates)
+        # result is the main input result
+        self.input_result = input_result
+        # sources_results is the list of the sources_results
+        self.sources_results = sources_results
+
+    @cached_property
+    def datasource(self):
+        for i in self.sources_results:
+            # for each result trigger the datasource to be computed
+            # and saved in context
+            self.context.notify_result(i.action_path[:-1], i.datasource)
+        # then return the input result
+        # which can use the datasources of the included results
+        return _tidy(self.input_result.datasource)

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)