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

anemoi/datasets/create/functions/sources/xarray/grid.py

@@ -8,39 +8,128 @@
 #
 
 
+import logging
+from functools import cached_property
+
 import numpy as np
 
+LOG = logging.getLogger(__name__)
+
 
 class Grid:
-    def __init__(self, lat, lon):
-        self.lat = lat
-        self.lon = lon
+
+    def __init__(self):
+        pass
 
     @property
     def latitudes(self):
-        return self.grid_points()[0]
+        return self.grid_points[0]
 
     @property
     def longitudes(self):
-        return self.grid_points()[1]
+        return self.grid_points[1]
+
+
+class LatLonGrid(Grid):
+    def __init__(self, lat, lon, variable_dims):
+        super().__init__()
+        self.lat = lat
+        self.lon = lon
+
 
+class XYGrid(Grid):
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
 
-class MeshedGrid(Grid):
-    _cache = None
 
+class MeshedGrid(LatLonGrid):
+
+    @cached_property
     def grid_points(self):
-        if self._cache is not None:
-            return self._cache
-        lat = self.lat.variable.values
-        lon = self.lon.variable.values
 
-        lat, lon = np.meshgrid(lat, lon)
-        self._cache = (lat.flatten(), lon.flatten())
-        return self._cache
+        lat, lon = np.meshgrid(
+            self.lat.variable.values,
+            self.lon.variable.values,
+        )
+
+        return lat.flatten(), lon.flatten()
+
 
+class UnstructuredGrid(LatLonGrid):
 
-class UnstructuredGrid(Grid):
+    def __init__(self, lat, lon, variable_dims):
+        super().__init__(lat, lon, variable_dims)
+        assert len(lat) == len(lon), (len(lat), len(lon))
+        self.variable_dims = variable_dims
+        self.grid_dims = lat.variable.dims
+        assert lon.variable.dims == self.grid_dims, (lon.variable.dims, self.grid_dims)
+        assert set(self.variable_dims) == set(self.grid_dims), (self.variable_dims, self.grid_dims)
+
+    @cached_property
     def grid_points(self):
-        lat = self.lat.variable.values.flatten()
-        lon = self.lon.variable.values.flatten()
+
+        assert 1 <= len(self.variable_dims) <= 2
+
+        if len(self.variable_dims) == 1:
+            return self.lat.variable.values.flatten(), self.lon.variable.values.flatten()
+
+        if len(self.variable_dims) == 2 and self.variable_dims == self.grid_dims:
+            return self.lat.variable.values.flatten(), self.lon.variable.values.flatten()
+
+        LOG.warning(
+            "UnstructuredGrid: variable indexing %s does not match grid indexing %s", self.variable_dims, self.grid_dims
+        )
+
+        lat = self.lat.variable.values.transpose().flatten()
+        lon = self.lon.variable.values.transpose().flatten()
+
         return lat, lon
+
+
+class ProjectionGrid(XYGrid):
+    def __init__(self, x, y, projection):
+        super().__init__(x, y)
+        self.projection = projection
+
+    def transformer(self):
+        from pyproj import CRS
+        from pyproj import Transformer
+
+        if isinstance(self.projection, dict):
+            data_crs = CRS.from_cf(self.projection)
+        else:
+            data_crs = self.projection
+        wgs84_crs = CRS.from_epsg(4326)  # WGS84
+
+        return Transformer.from_crs(data_crs, wgs84_crs, always_xy=True)
+
+
+class MeshProjectionGrid(ProjectionGrid):
+
+    @cached_property
+    def grid_points(self):
+
+        transformer = self.transformer()
+        xv, yv = np.meshgrid(self.x.variable.values, self.y.variable.values)  # , indexing="ij")
+        lon, lat = transformer.transform(xv, yv)
+        return lat.flatten(), lon.flatten()
+
+
+class UnstructuredProjectionGrid(XYGrid):
+    @cached_property
+    def grid_points(self):
+        assert False, "Not implemented"
+
+        # lat, lon = transformer.transform(
+        #      self.y.variable.values.flatten(),
+        #     self.x.variable.values.flatten(),
+
+        # )
+
+        # lat = lat[::len(lat)//100]
+        # lon = lon[::len(lon)//100]
+
+        # print(len(lat), len(lon))
+
+        # return np.meshgrid(lat, lon)

anemoi/datasets/create/functions/sources/xarray/metadata.py

@@ -40,7 +40,9 @@ class _MDMapping:
         return f"MDMapping({self.mapping})"
 
     def fill_time_metadata(self, field, md):
-        md["valid_datetime"] = as_datetime(self.variable.time.fill_time_metadata(field._md, md)).isoformat()
+        valid_datetime = self.variable.time.fill_time_metadata(field._md, md)
+        if valid_datetime is not None:
+            md["valid_datetime"] = as_datetime(valid_datetime).isoformat()
 
 
 class XArrayMetadata(RawMetadata):
@@ -71,16 +73,6 @@ class XArrayMetadata(RawMetadata):
 
     def _as_mars(self):
         return {}
-        # p =  dict(
-        #     param=self.get("variable", self.get("param")),
-        #     step=self.get("step"),
-        #     levelist=self.get("levelist", self.get("level")),
-        #     levtype=self.get("levtype"),
-        #     number=self.get("number"),
-        #     date=self.get("date"),
-        #     time=self.get("time"),
-        # )
-        # return {k: v for k, v in p.items() if v is not None}
 
     def _base_datetime(self):
         return self._field.forecast_reference_time

anemoi/datasets/create/functions/sources/xarray/time.py

@@ -42,11 +42,7 @@ class Time:
 class Constant(Time):
 
     def fill_time_metadata(self, coords_values, metadata):
-        raise NotImplementedError("Constant time not implemented")
-        # print("Constant", coords_values, metadata)
-        # metadata["date"] = time.strftime("%Y%m%d")
-        # metadata["time"] = time.strftime("%H%M")
-        # metadata["step"] = 0
+        return None
 
 
 class Analysis(Time):

anemoi/datasets/create/functions/sources/xarray/variable.py

@@ -24,7 +24,7 @@ class Variable:
         self,
         *,
         ds,
-        var,
+        variable,
         coordinates,
         grid,
         time,
@@ -32,13 +32,13 @@ class Variable:
         array_backend=None,
     ):
         self.ds = ds
-        self.var = var
+        self.variable = variable
 
         self.grid = grid
         self.coordinates = coordinates
 
         self._metadata = metadata.copy()
-        self._metadata.update({"variable": var.name})
+        self._metadata.update({"variable": variable.name})
 
         self.time = time
 
@@ -51,20 +51,20 @@ class Variable:
 
     @property
     def name(self):
-        return self.var.name
+        return self.variable.name
 
     def __len__(self):
         return self.length
 
     @property
     def grid_mapping(self):
-        grid_mapping = self.var.attrs.get("grid_mapping", None)
+        grid_mapping = self.variable.attrs.get("grid_mapping", None)
         if grid_mapping is None:
             return None
         return self.ds[grid_mapping].attrs
 
     def grid_points(self):
-        return self.grid.grid_points()
+        return self.grid.grid_points
 
     @property
     def latitudes(self):
@@ -76,7 +76,7 @@ class Variable:
 
     def __repr__(self):
         return "Variable[name=%s,coordinates=%s,metadata=%s]" % (
-            self.var.name,
+            self.variable.name,
             self.coordinates,
             self._metadata,
         )
@@ -90,7 +90,7 @@ class Variable:
 
         coords = np.unravel_index(i, self.shape)
         kwargs = {k: v for k, v in zip(self.names, coords)}
-        return XArrayField(self, self.var.isel(kwargs))
+        return XArrayField(self, self.variable.isel(kwargs))
 
     def sel(self, missing, **kwargs):
 
@@ -117,7 +117,7 @@ class Variable:
 
         variable = Variable(
             ds=self.ds,
-            var=self.var.isel({k: i}),
+            var=self.variable.isel({k: i}),
             coordinates=coordinates,
             grid=self.grid,
             time=self.time,
@@ -136,7 +136,7 @@ class Variable:
             name = kwargs.pop("variable")
             if not isinstance(name, (list, tuple)):
                 name = [name]
-            if self.var.name not in name:
+            if self.variable.name not in name:
                 return False, None
             return True, kwargs
         return True, kwargs

anemoi/datasets/create/input/__init__.py

@@ -0,0 +1,69 @@
+# (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 datetime
+import itertools
+import logging
+import math
+import time
+from collections import defaultdict
+from copy import deepcopy
+from functools import cached_property
+from functools import wraps
+
+import numpy as np
+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 .trace import trace_select
+
+LOG = logging.getLogger(__name__)
+
+
+class InputBuilder:
+    def __init__(self, config, data_sources, **kwargs):
+        self.kwargs = kwargs
+
+        config = deepcopy(config)
+        if data_sources:
+            config = dict(
+                data_sources=dict(
+                    sources=data_sources,
+                    input=config,
+                )
+            )
+        self.config = config
+        self.action_path = ["input"]
+
+    @trace_select
+    def select(self, group_of_dates):
+        from .action import ActionContext
+        from .action import action_factory
+
+        """This changes the context."""
+        context = ActionContext(**self.kwargs)
+        action = action_factory(self.config, context, self.action_path)
+        return action.select(group_of_dates)
+
+    def __repr__(self):
+        from .action import ActionContext
+        from .action import action_factory
+
+        context = ActionContext(**self.kwargs)
+        a = action_factory(self.config, context, self.action_path)
+        return repr(a)
+
+    def _trace_select(self, group_of_dates):
+        return f"InputBuilder({group_of_dates})"
+
+
+build_input = InputBuilder

anemoi/datasets/create/input/action.py

@@ -0,0 +1,123 @@
+# (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 copy import deepcopy
+
+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.order import build_remapping
+
+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 .context import Context
+from .misc import is_function
+
+LOG = logging.getLogger(__name__)
+
+
+class Action:
+    def __init__(self, context, action_path, /, *args, **kwargs):
+        if "args" in kwargs and "kwargs" in kwargs:
+            """We have:
+               args = []
+               kwargs = {args: [...], kwargs: {...}}
+            move the content of kwargs to args and kwargs.
+            """
+            assert len(kwargs) == 2, (args, kwargs)
+            assert not args, (args, kwargs)
+            args = kwargs.pop("args")
+            kwargs = kwargs.pop("kwargs")
+
+        assert isinstance(context, ActionContext), type(context)
+        self.context = context
+        self.kwargs = kwargs
+        self.args = args
+        self.action_path = action_path
+
+    @classmethod
+    def _short_str(cls, x):
+        x = str(x)
+        if len(x) < 1000:
+            return x
+        return x[:1000] + "..."
+
+    def __repr__(self, *args, _indent_="\n", _inline_="", **kwargs):
+        more = ",".join([str(a)[:5000] for a in args])
+        more += ",".join([f"{k}={v}"[:5000] for k, v in kwargs.items()])
+
+        more = more[:5000]
+        txt = f"{self.__class__.__name__}: {_inline_}{_indent_}{more}"
+        if _indent_:
+            txt = txt.replace("\n", "\n  ")
+        return txt
+
+    def select(self, dates, **kwargs):
+        self._raise_not_implemented()
+
+    def _raise_not_implemented(self):
+        raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
+
+    def _trace_select(self, group_of_dates):
+        return f"{self.__class__.__name__}({group_of_dates})"
+
+
+class ActionContext(Context):
+    def __init__(self, /, order_by, flatten_grid, remapping, use_grib_paramid):
+        super().__init__()
+        self.order_by = order_by
+        self.flatten_grid = flatten_grid
+        self.remapping = build_remapping(remapping)
+        self.use_grib_paramid = use_grib_paramid
+
+
+def action_factory(config, context, action_path):
+
+    from .concat import ConcatAction
+    from .data_sources import DataSourcesAction
+    from .function import FunctionAction
+    from .join import JoinAction
+    from .pipe import PipeAction
+    from .repeated_dates import RepeatedDatesAction
+
+    # from .data_sources import DataSourcesAction
+
+    assert isinstance(context, Context), (type, context)
+    if not isinstance(config, dict):
+        raise ValueError(f"Invalid input config {config}")
+    if len(config) != 1:
+        raise ValueError(f"Invalid input config. Expecting dict with only one key, got {list(config.keys())}")
+
+    config = deepcopy(config)
+    key = list(config.keys())[0]
+
+    if isinstance(config[key], list):
+        args, kwargs = config[key], {}
+    elif isinstance(config[key], dict):
+        args, kwargs = [], config[key]
+    else:
+        raise ValueError(f"Invalid input config {config[key]} ({type(config[key])}")
+
+    cls = {
+        "data_sources": DataSourcesAction,
+        "concat": ConcatAction,
+        "join": JoinAction,
+        "pipe": PipeAction,
+        "function": FunctionAction,
+        "repeated_dates": RepeatedDatesAction,
+    }.get(key)
+
+    if cls is None:
+        if not is_function(key, "sources"):
+            raise ValueError(f"Unknown action '{key}' in {config}")
+        cls = FunctionAction
+        args = [key] + args
+
+    return cls(context, action_path + [key], *args, **kwargs)

anemoi/datasets/create/input/concat.py

@@ -0,0 +1,92 @@
+# (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 copy import deepcopy
+from functools import cached_property
+
+from anemoi.datasets.dates import DatesProvider
+
+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 ConcatResult(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)
+
+    @property
+    def variables(self):
+        """Check that all the results objects have the same variables."""
+        variables = None
+        for f in self.results:
+            if f.empty:
+                continue
+            if variables is None:
+                variables = f.variables
+            assert variables == f.variables, (variables, f.variables)
+        assert variables is not None, self.results
+        return variables
+
+    def __repr__(self):
+        content = "\n".join([str(i) for i in self.results])
+        return super().__repr__(content)
+
+
+class ConcatAction(Action):
+    def __init__(self, context, action_path, *configs):
+        super().__init__(context, action_path, *configs)
+        parts = []
+        for i, cfg in enumerate(configs):
+            if "dates" not in cfg:
+                raise ValueError(f"Missing 'dates' in {cfg}")
+            cfg = deepcopy(cfg)
+            dates_cfg = cfg.pop("dates")
+            assert isinstance(dates_cfg, dict), 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
+
+    def __repr__(self):
+        content = "\n".join([str(i) for i in self.parts])
+        return super().__repr__(content)
+
+    @trace_select
+    def select(self, group_of_dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        results = []
+        for filtering_dates, action in self.parts:
+            newdates = GroupOfDates(sorted(set(group_of_dates) & set(filtering_dates)), group_of_dates.provider)
+            if newdates:
+                results.append(action.select(newdates))
+        if not results:
+            return EmptyResult(self.context, self.action_path, group_of_dates)
+
+        return ConcatResult(self.context, self.action_path, group_of_dates, results)

anemoi/datasets/create/input/context.py

@@ -0,0 +1,59 @@
+# (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
+import textwrap
+
+from anemoi.utils.dates import as_datetime as as_datetime
+from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
+from anemoi.utils.humanize import plural
+
+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 .trace import step
+from .trace import trace
+
+LOG = logging.getLogger(__name__)
+
+
+class Context:
+    def __init__(self):
+        # used_references is a set of reference paths that will be needed
+        self.used_references = set()
+        # results is a dictionary of reference path -> obj
+        self.results = {}
+
+    def will_need_reference(self, key):
+        assert isinstance(key, (list, tuple)), key
+        key = tuple(key)
+        self.used_references.add(key)
+
+    def notify_result(self, key, result):
+        trace(
+            "🎯",
+            step(key),
+            "notify result",
+            textwrap.shorten(repr(result).replace(",", ", "), width=40),
+            plural(len(result), "field"),
+        )
+        assert isinstance(key, (list, tuple)), key
+        key = tuple(key)
+        if key in self.used_references:
+            if key in self.results:
+                raise ValueError(f"Duplicate result {key}")
+            self.results[key] = result
+
+    def get_result(self, key):
+        assert isinstance(key, (list, tuple)), key
+        key = tuple(key)
+        if key in self.results:
+            return self.results[key]
+        all_keys = sorted(list(self.results.keys()))
+        raise ValueError(f"Cannot find result {key} in {all_keys}")

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)