anemoi-datasets 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

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

@@ -0,0 +1,98 @@
+# (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 datetime
+
+
+class Time:
+    @classmethod
+    def from_coordinates(cls, coordinates):
+        time_coordinate = [c for c in coordinates if c.is_time]
+        step_coordinate = [c for c in coordinates if c.is_step]
+        date_coordinate = [c for c in coordinates if c.is_date]
+
+        if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 1:
+            return ForecasstFromValidTimeAndStep(step_coordinate[0])
+
+        if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
+            return Analysis()
+
+        if len(date_coordinate) == 0 and len(time_coordinate) == 0 and len(step_coordinate) == 0:
+            return Constant()
+
+        if len(date_coordinate) == 1 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
+            return ForecastFromValidTimeAndBaseTime(date_coordinate[0])
+
+        if len(date_coordinate) == 1 and len(time_coordinate) == 0 and len(step_coordinate) == 1:
+            return ForecastFromBaseTimeAndDate(date_coordinate[0], step_coordinate[0])
+
+        raise NotImplementedError(f"{date_coordinate=} {time_coordinate=} {step_coordinate=}")
+
+
+class Constant(Time):
+
+    def fill_time_metadata(self, time, metadata):
+        metadata["date"] = time.strftime("%Y%m%d")
+        metadata["time"] = time.strftime("%H%M")
+        metadata["step"] = 0
+
+
+class Analysis(Time):
+
+    def fill_time_metadata(self, time, metadata):
+        metadata["date"] = time.strftime("%Y%m%d")
+        metadata["time"] = time.strftime("%H%M")
+        metadata["step"] = 0
+
+
+class ForecasstFromValidTimeAndStep(Time):
+    def __init__(self, step_coordinate):
+        self.step_name = step_coordinate.variable.name
+
+    def fill_time_metadata(self, time, metadata):
+        step = metadata.pop(self.step_name)
+        assert isinstance(step, datetime.timedelta)
+        base = time - step
+
+        hours = step.total_seconds() / 3600
+        assert int(hours) == hours
+
+        metadata["date"] = base.strftime("%Y%m%d")
+        metadata["time"] = base.strftime("%H%M")
+        metadata["step"] = int(hours)
+
+
+class ForecastFromValidTimeAndBaseTime(Time):
+    def __init__(self, date_coordinate):
+        self.date_coordinate = date_coordinate
+
+    def fill_time_metadata(self, time, metadata):
+
+        step = time - self.date_coordinate
+
+        hours = step.total_seconds() / 3600
+        assert int(hours) == hours
+
+        metadata["date"] = self.date_coordinate.single_value.strftime("%Y%m%d")
+        metadata["time"] = self.date_coordinate.single_value.strftime("%H%M")
+        metadata["step"] = int(hours)
+
+
+class ForecastFromBaseTimeAndDate(Time):
+    def __init__(self, date_coordinate, step_coordinate):
+        self.date_coordinate = date_coordinate
+        self.step_coordinate = step_coordinate
+
+    def fill_time_metadata(self, time, metadata):
+        metadata["date"] = time.strftime("%Y%m%d")
+        metadata["time"] = time.strftime("%H%M")
+        hours = metadata[self.step_coordinate.name].total_seconds() / 3600
+        assert int(hours) == hours
+        metadata["step"] = int(hours)

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

@@ -0,0 +1,198 @@
+# (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 math
+from functools import cached_property
+
+import numpy as np
+from earthkit.data.utils.array import ensure_backend
+
+from anemoi.datasets.create.functions.sources.xarray.metadata import MDMapping
+
+from .field import XArrayField
+
+LOG = logging.getLogger(__name__)
+
+
+class Variable:
+    def __init__(self, *, ds, var, coordinates, grid, time, metadata, mapping=None, array_backend=None):
+        self.ds = ds
+        self.var = var
+
+        self.grid = grid
+        self.coordinates = coordinates
+
+        # print("Variable", var.name)
+        # for c in coordinates:
+        #     print(" ", c)
+
+        self._metadata = metadata.copy()
+        # self._metadata.update(var.attrs)
+        self._metadata.update({"variable": var.name})
+
+        # self._metadata.setdefault("level", None)
+        # self._metadata.setdefault("number", 0)
+        # self._metadata.setdefault("levtype", "sfc")
+        self._mapping = mapping
+
+        self.time = time
+
+        self.shape = tuple(len(c.variable) for c in coordinates if c.is_dim and not c.scalar and not c.is_grid)
+        self.names = {c.variable.name: c for c in coordinates if c.is_dim and not c.scalar and not c.is_grid}
+        self.by_name = {c.variable.name: c for c in coordinates}
+
+        self.length = math.prod(self.shape)
+        self.array_backend = ensure_backend(array_backend)
+
+    def update_metadata_mapping(self, kwargs):
+
+        result = {}
+
+        for k, v in kwargs.items():
+            if k == "param":
+                result[k] = "variable"
+                continue
+
+            for c in self.coordinates:
+                if k in c.mars_names:
+                    for v in c.mars_names:
+                        result[v] = c.variable.name
+                    break
+
+        self._mapping = MDMapping(result)
+
+    @property
+    def name(self):
+        return self.var.name
+
+    def __len__(self):
+        return self.length
+
+    @property
+    def grid_mapping(self):
+        grid_mapping = self.var.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()
+
+    @property
+    def latitudes(self):
+        return self.grid.latitudes
+
+    @property
+    def longitudes(self):
+        return self.grid.longitudes
+
+    def __repr__(self):
+        return "Variable[name=%s,coordinates=%s,metadata=%s]" % (
+            self.var.name,
+            self.coordinates,
+            self._metadata,
+        )
+
+    def __getitem__(self, i):
+        """
+        Get a 2D field from the variable
+        """
+        if i >= self.length:
+            raise IndexError(i)
+
+        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))
+
+    @property
+    def mapping(self):
+        return self._mapping
+
+    def sel(self, missing, **kwargs):
+
+        if not kwargs:
+            return self
+
+        kwargs = self._mapping.from_user(kwargs)
+
+        k, v = kwargs.popitem()
+
+        c = self.by_name.get(k)
+
+        if c is None:
+            missing[k] = v
+            return self.sel(missing, **kwargs)
+
+        i = c.index(v)
+        if i is None:
+            LOG.warning(f"Could not find {k}={v} in {c}")
+            return None
+
+        coordinates = [x.reduced(i) if c is x else x for x in self.coordinates]
+
+        metadata = self._metadata.copy()
+        metadata.update({k: v})
+
+        variable = Variable(
+            ds=self.ds,
+            var=self.var.isel({k: i}),
+            coordinates=coordinates,
+            grid=self.grid,
+            time=self.time,
+            metadata=metadata,
+            mapping=self.mapping,
+        )
+
+        return variable.sel(missing, **kwargs)
+
+    def match(self, **kwargs):
+        kwargs = self._mapping.from_user(kwargs)
+
+        if "variable" in kwargs:
+            name = kwargs.pop("variable")
+            if not isinstance(name, (list, tuple)):
+                name = [name]
+            if self.var.name not in name:
+                return False, None
+            return True, kwargs
+        return True, kwargs
+
+
+class FilteredVariable:
+    def __init__(self, variable, **kwargs):
+        self.variable = variable
+        self.kwargs = kwargs
+
+    @cached_property
+    def fields(self):
+        """Filter the fields of a variable based on metadata.
+
+        Returns
+        -------
+        list
+            A list of fields that match the metadata.
+        """
+        return [
+            field
+            for field in self.variable
+            if all(field.metadata(k, default=None) == v for k, v in self.kwargs.items())
+        ]
+
+    @property
+    def length(self):
+        return len(self.fields)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        if i >= self.length:
+            raise IndexError(i)
+        return self.fields[i]

anemoi/datasets/create/functions/sources/xarray_kerchunk.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.
+#
+
+
+from earthkit.data.core.fieldlist import MultiFieldList
+
+from . import iterate_patterns
+from .xarray import load_one
+
+
+def load_many(emoji, context, dates, pattern, options, **kwargs):
+
+    result = []
+    options = options.copy() if options is not None else {}
+
+    options.setdefault("engine", "zarr")
+    options.setdefault("backend_kwargs", {})
+
+    backend_kwargs = options["backend_kwargs"]
+    backend_kwargs.setdefault("consolidated", False)
+    backend_kwargs.setdefault("storage_options", {})
+
+    storage_options = backend_kwargs["storage_options"]
+    storage_options.setdefault("remote_protocol", "s3")
+    storage_options.setdefault("remote_options", {"anon": True})
+
+    for path, dates in iterate_patterns(pattern, dates, **kwargs):
+        storage_options["fo"] = path
+
+        result.append(load_one(emoji, context, dates, "reference://", options=options, **kwargs))
+
+    return MultiFieldList(result)
+
+
+def execute(context, dates, json, options=None, **kwargs):
+    return load_many("🧱", context, dates, json, options, **kwargs)

anemoi/datasets/create/functions/sources/xarray_zarr.py

@@ -0,0 +1,15 @@
+# (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.
+#
+
+
+from .xarray import load_many
+
+
+def execute(context, dates, url, *args, **kwargs):
+    return load_many("🇿", context, dates, url, *args, **kwargs)

anemoi/datasets/create/functions/sources/zenodo.py

@@ -0,0 +1,40 @@
+# (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.
+#
+
+
+from earthkit.data.core.fieldlist import MultiFieldList
+from earthkit.data.sources.url import download_and_cache
+
+from . import iterate_patterns
+from .xarray import load_one
+
+
+def execute(context, dates, record_id, file_key, *args, **kwargs):
+    import requests
+
+    result = []
+
+    URLPATTERN = "https://zenodo.org/api/records/{record_id}"
+    url = URLPATTERN.format(record_id=record_id)
+    r = requests.get(url)
+    r.raise_for_status()
+    record = r.json()
+
+    urls = {}
+    for file in record["files"]:
+        urls[file["key"]] = file["links"]["self"]
+
+    for url, dates in iterate_patterns(file_key, dates, **kwargs):
+        if url not in urls:
+            continue
+
+        path = download_and_cache(urls[url])
+        result.append(load_one("?", context, dates, path, options={}, flavour=None, **kwargs))
+
+    return MultiFieldList(result)