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

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

@@ -9,10 +9,10 @@
 import datetime
 from collections import defaultdict
 
-from climetlab.core.temporary import temp_file
-from climetlab.readers.grib.output import new_grib_output
+from earthkit.data.core.temporary import temp_file
+from earthkit.data.readers.grib.output import new_grib_output
 
-from anemoi.datasets.create.functions import assert_is_fieldset
+from anemoi.datasets.create.functions import assert_is_fieldlist
 from anemoi.datasets.create.utils import to_datetime_list
 
 
@@ -36,7 +36,7 @@ def normalise_time_delta(t):
 def group_by_field(ds):
     d = defaultdict(list)
     for field in ds.order_by("valid_datetime"):
-        m = field.as_mars()
+        m = field.metadata(namespace="mars")
         for k in ("date", "time", "step"):
             m.pop(k, None)
         keys = tuple(m.items())
@@ -103,10 +103,10 @@ def tendencies(dates, time_increment, **kwargs):
 
     out.close()
 
-    from climetlab import load_source
+    from earthkit.data import from_source
 
-    ds = load_source("file", path)
-    assert_is_fieldset(ds)
+    ds = from_source("file", path)
+    assert_is_fieldlist(ds)
     # save a reference to the tmp file so it is deleted
     # only when the dataset is not used anymore
     ds._tmp = tmp

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

@@ -0,0 +1,73 @@
+# (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 earthkit.data.core.fieldlist import MultiFieldList
+
+from anemoi.datasets.data.stores import name_to_zarr_store
+
+from .. import iterate_patterns
+from .fieldlist import XarrayFieldList
+
+LOG = logging.getLogger(__name__)
+
+
+def check(what, ds, paths, **kwargs):
+    count = 1
+    for k, v in kwargs.items():
+        if isinstance(v, (tuple, list)):
+            count *= len(v)
+
+    if len(ds) != count:
+        raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, {what}s={paths})")
+
+
+def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs):
+    import xarray as xr
+
+    """
+    We manage the S3 client ourselve, bypassing fsspec and s3fs layers, because sometimes something on the stack
+    zarr/fsspec/s3fs/boto3 (?) seem to flags files as missing when they actually are not (maybe when S3 reports some sort of
+    connection error). In that case,  Zarr will silently fill the chunks that could not be downloaded with NaNs.
+    See https://github.com/pydata/xarray/issues/8842
+
+    We have seen this bug triggered when we run many clients in parallel, for example, when we create a new dataset using `xarray-zarr`.
+    """
+
+    context.trace(emoji, dataset, options)
+
+    if isinstance(dataset, str) and ".zarr" in dataset:
+        data = xr.open_zarr(name_to_zarr_store(dataset), **options)
+    else:
+        data = xr.open_dataset(dataset, **options)
+
+    fs = XarrayFieldList.from_xarray(data, flavour)
+    result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
+
+    if len(result) == 0:
+        LOG.warning(f"No data found for {dataset} and dates {dates}")
+        LOG.warning(f"Options: {options}")
+        LOG.warning(data)
+
+    return result
+
+
+def load_many(emoji, context, dates, pattern, **kwargs):
+
+    result = []
+
+    for path, dates in iterate_patterns(pattern, dates, **kwargs):
+        result.append(load_one(emoji, context, dates, path, **kwargs))
+
+    return MultiFieldList(result)
+
+
+def execute(context, dates, url, *args, **kwargs):
+    return load_many("🌐", context, dates, url, *args, **kwargs)

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

@@ -0,0 +1,234 @@
+# (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
+import logging
+
+import numpy as np
+from earthkit.data.utils.dates import to_datetime
+
+LOG = logging.getLogger(__name__)
+
+
+def is_scalar(variable):
+    shape = variable.shape
+    if shape == (1,):
+        return True
+    if len(shape) == 0:
+        return True
+    return False
+
+
+def extract_single_value(variable):
+    shape = variable.shape
+    if np.issubdtype(variable.values.dtype, np.datetime64):
+        if len(shape) == 0:
+            return to_datetime(variable.values)  # Convert to python datetime
+        assert False, (shape, variable.values[:2])
+
+    if np.issubdtype(variable.values.dtype, np.timedelta64):
+        if len(shape) == 0:
+            # Convert to python timedelta64
+            return datetime.timedelta(seconds=variable.values.astype("timedelta64[s]").astype(int).item())
+        assert False, (shape, variable.values)
+
+    if shape == (1,):
+        return variable.values[0]
+
+    if len(shape) == 0:
+        return variable.values.item()
+
+    assert False, (shape, variable.values)
+
+
+class Coordinate:
+    is_grid = False
+    is_dim = True
+    is_lat = False
+    is_lon = False
+    is_time = False
+    is_step = False
+    is_date = False
+
+    def __init__(self, variable):
+        self.variable = variable
+        self.scalar = is_scalar(variable)
+        self.kwargs = {}  # Used when creating a new coordinate (reduced method)
+
+    def __len__(self):
+        return 1 if self.scalar else len(self.variable)
+
+    def __repr__(self):
+        return "%s[name=%s,values=%s]" % (
+            self.__class__.__name__,
+            self.variable.name,
+            self.variable.values if self.scalar else len(self),
+        )
+
+    def reduced(self, i):
+        """Create a new coordinate with a single value
+
+        Parameters
+        ----------
+        i : int
+            the index of the value to select
+
+        Returns
+        -------
+        Coordinate
+            the new coordinate
+        """
+        return self.__class__(
+            self.variable.isel({self.variable.dims[0]: i}),
+            **self.kwargs,
+        )
+
+    def index(self, value):
+        """Return the index of the value in the coordinate
+
+        Parameters
+        ----------
+        value : Any
+            The value to search for
+
+        Returns
+        -------
+        int or None
+            The index of the value in the coordinate or None if not found
+        """
+
+        if isinstance(value, (list, tuple)):
+            if len(value) == 1:
+                return self._index_single(value)
+            else:
+                return self._index_multiple(value)
+        return self._index_single(value)
+
+    def _index_single(self, value):
+
+        values = self.variable.values
+
+        # Assume the array is sorted
+        index = np.searchsorted(values, value)
+
+        if index < len(values) and values[index] == value:
+            return index
+
+        # If not found, we need to check if the value is in the array
+
+        index = np.where(values == value)[0]
+        if len(index) > 0:
+            return index[0]
+
+        return None
+
+    def _index_multiple(self, value):
+
+        values = self.variable.values
+
+        # Assume the array is sorted
+
+        index = np.searchsorted(values, value)
+        index = index[index < len(values)]
+
+        if np.all(values[index] == value):
+            return index
+
+        # If not found, we need to check if the value is in the array
+
+        index = np.where(np.isin(values, value))[0]
+
+        # We could also return incomplete matches
+        if len(index) == len(value):
+            return index
+
+        return None
+
+    @property
+    def name(self):
+        return self.variable.name
+
+    def normalise(self, value):
+        # Subclasses to format values that will be added to the field metadata
+        return value
+
+    @property
+    def single_value(self):
+        return extract_single_value(self.variable)
+
+
+class TimeCoordinate(Coordinate):
+    is_time = True
+    mars_names = ("valid_datetime",)
+
+    def index(self, time):
+        return super().index(np.datetime64(time))
+
+
+class DateCoordinate(Coordinate):
+    is_date = True
+    mars_names = ("date",)
+
+    def index(self, date):
+        return super().index(np.datetime64(date))
+
+
+class StepCoordinate(Coordinate):
+    is_step = True
+    mars_names = ("step",)
+
+
+class LevelCoordinate(Coordinate):
+    mars_names = ("level", "levelist")
+
+    def __init__(self, variable, levtype):
+        super().__init__(variable)
+        self.levtype = levtype
+        # kwargs is used when creating a new coordinate (reduced method)
+        self.kwargs = {"levtype": levtype}
+
+    def normalise(self, value):
+        # Some netcdf have pressue levels in float
+        if int(value) == value:
+            return int(value)
+        return value
+
+
+class EnsembleCoordinate(Coordinate):
+    mars_names = ("number",)
+
+
+class LongitudeCoordinate(Coordinate):
+    is_grid = True
+    is_lon = True
+    mars_names = ("longitude",)
+
+
+class LatitudeCoordinate(Coordinate):
+    is_grid = True
+    is_lat = True
+    mars_names = ("latitude",)
+
+
+class XCoordinate(Coordinate):
+    is_grid = True
+    mars_names = ("x",)
+
+
+class YCoordinate(Coordinate):
+    is_grid = True
+    mars_names = ("y",)
+
+
+class ScalarCoordinate(Coordinate):
+    is_grid = False
+
+    @property
+    def mars_names(self):
+        return (self.variable.name,)

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

@@ -0,0 +1,109 @@
+# (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 earthkit.data.core.fieldlist import Field
+from earthkit.data.core.fieldlist import math
+
+from .coordinates import extract_single_value
+from .coordinates import is_scalar
+from .metadata import XArrayMetadata
+
+LOG = logging.getLogger(__name__)
+
+
+class EmptyFieldList:
+    def __len__(self):
+        return 0
+
+    def __getitem__(self, i):
+        raise IndexError(i)
+
+    def __repr__(self) -> str:
+        return "EmptyFieldList()"
+
+
+class XArrayField(Field):
+
+    def __init__(self, owner, selection):
+        """Create a new XArrayField object.
+
+        Parameters
+        ----------
+        owner : Variable
+            The variable that owns this field.
+        selection : XArrayDataArray
+            A 2D sub-selection of the variable's underlying array.
+            This is actually a nD object, but the first dimensions are always 1.
+            The other two dimensions are latitude and longitude.
+        """
+        super().__init__(owner.array_backend)
+
+        self.owner = owner
+        self.selection = selection
+
+        # Copy the metadata from the owner
+        self._md = owner._metadata.copy()
+
+        for coord_name, coord_value in self.selection.coords.items():
+            if is_scalar(coord_value):
+                # Extract the single value from the scalar dimension
+                # and store it in the metadata
+                coordinate = owner.by_name[coord_name]
+                self._md[coord_name] = coordinate.normalise(extract_single_value(coord_value))
+
+        # print(values.ndim, values.shape, selection.dims)
+        # By now, the only dimensions should be latitude and longitude
+        self._shape = tuple(list(self.selection.shape)[-2:])
+        if math.prod(self._shape) != math.prod(self.selection.shape):
+            print(self.selection.ndim, self.selection.shape)
+            print(self.selection)
+            raise ValueError("Invalid shape for selection")
+
+    @property
+    def shape(self):
+        return self._shape
+
+    def to_numpy(self, flatten=False, dtype=None):
+        values = self.selection.values
+
+        assert dtype is None
+        if flatten:
+            return values.flatten()
+        return values.reshape(self.shape)
+
+    def _make_metadata(self):
+        return XArrayMetadata(self, self.owner.mapping)
+
+    def grid_points(self):
+        return self.owner.grid_points()
+
+    @property
+    def resolution(self):
+        return None
+
+    @property
+    def grid_mapping(self):
+        return self.owner.grid_mapping
+
+    @property
+    def latitudes(self):
+        return self.owner.latitudes
+
+    @property
+    def longitudes(self):
+        return self.owner.longitudes
+
+    @property
+    def forecast_reference_time(self):
+        return self.owner.forecast_reference_time
+
+    def __repr__(self):
+        return repr(self._metadata)

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

@@ -0,0 +1,171 @@
+# (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 json
+import logging
+
+import yaml
+from earthkit.data.core.fieldlist import FieldList
+
+from .coordinates import is_scalar as is_scalar
+from .field import EmptyFieldList
+from .flavour import CoordinateGuesser
+from .metadata import XArrayMetadata as XArrayMetadata
+from .time import Time
+from .variable import FilteredVariable
+from .variable import Variable
+
+LOG = logging.getLogger(__name__)
+
+
+class XarrayFieldList(FieldList):
+    def __init__(self, ds, variables):
+        self.ds = ds
+        self.variables = variables.copy()
+        self.total_length = sum(v.length for v in variables)
+
+    def __repr__(self):
+        return f"XarrayFieldList({self.total_length})"
+
+    def __len__(self):
+        return self.total_length
+
+    def __getitem__(self, i):
+        k = i
+
+        if i < 0:
+            i = self.total_length + i
+
+        for v in self.variables:
+            if i < v.length:
+                return v[i]
+            i -= v.length
+
+        raise IndexError(k)
+
+    @classmethod
+    def from_xarray(cls, ds, flavour=None):
+        variables = []
+
+        if isinstance(flavour, str):
+            with open(flavour) as f:
+                if flavour.endswith(".yaml") or flavour.endswith(".yml"):
+                    flavour = yaml.safe_load(f)
+                else:
+                    flavour = json.load(f)
+
+        guess = CoordinateGuesser.from_flavour(ds, flavour)
+
+        skip = set()
+
+        def _skip_attr(v, attr_name):
+            attr_val = getattr(v, attr_name, "")
+            if isinstance(attr_val, str):
+                skip.update(attr_val.split(" "))
+
+        for name in ds.data_vars:
+            v = ds[name]
+            _skip_attr(v, "coordinates")
+            _skip_attr(v, "bounds")
+            _skip_attr(v, "grid_mapping")
+
+        # Select only geographical variables
+        for name in ds.data_vars:
+
+            if name in skip:
+                continue
+
+            v = ds[name]
+            coordinates = []
+
+            for coord in v.coords:
+
+                c = guess.guess(ds[coord], coord)
+                assert c, f"Could not guess coordinate for {coord}"
+                if coord not in v.dims:
+                    c.is_dim = False
+                coordinates.append(c)
+
+            grid_coords = sum(1 for c in coordinates if c.is_grid and c.is_dim)
+            assert grid_coords <= 2
+
+            if grid_coords < 2:
+                continue
+
+            variables.append(
+                Variable(
+                    ds=ds,
+                    var=v,
+                    coordinates=coordinates,
+                    grid=guess.grid(coordinates),
+                    time=Time.from_coordinates(coordinates),
+                    metadata={},
+                )
+            )
+
+        return cls(ds, variables)
+
+    def sel(self, **kwargs):
+        """Override the FieldList's sel method
+
+        Returns
+        -------
+        FieldList
+            The new FieldList
+
+        The algorithm is as follows:
+        1 - Use the kwargs to select the variables that match the selection (`param` or `variable`)
+        2 - For each variable, use the remaining kwargs to select the coordinates (`level`, `number`, ...)
+        3 - Some mars like keys, like `date`, `time`, `step` are not found in the coordinates,
+            but added to the metadata of the selected fields. A example is `step` that is added to the
+            metadata of the field. Step 2 may return a variable that contain all the fields that
+            verify at the same `valid_datetime`, with different base `date` and `time` and a different `step`.
+            So we get an extra chance to filter the fields by the metadata.
+        """
+
+        variables = []
+        count = 0
+
+        for v in self.variables:
+
+            v.update_metadata_mapping(kwargs)
+
+            # First, select matching variables
+            # This will consume 'param' or 'variable' from kwargs
+            # and return the rest
+            match, rest = v.match(**kwargs)
+
+            if match:
+                count += 1
+                missing = {}
+
+                # Select from the variable's coordinates (time, level, number, ....)
+                # This may return a new variable with a isel() slice of the selection
+                # or None if the selection is not possible. In this case, missing is updated
+                # with the values of kwargs (rest) that are not relevant for this variable
+                v = v.sel(missing, **rest)
+                if missing:
+                    if v is not None:
+                        # The remaining kwargs are passed used to create a FilteredVariable
+                        # that will select 2D slices based on their metadata
+                        v = FilteredVariable(v, **missing)
+                    else:
+                        LOG.warning(f"Variable {v} has missing coordinates: {missing}")
+
+                if v is not None:
+                    variables.append(v)
+
+        if count == 0:
+            LOG.warning("No variable found for %s", kwargs)
+            LOG.warning("Variables: %s", sorted([v.name for v in self.variables]))
+
+        if not variables:
+            return EmptyFieldList()
+
+        return self.__class__(self.ds, variables)