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

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

@@ -7,7 +7,6 @@
 # nor does it submit to any jurisdiction.
 #
 import datetime
-from copy import deepcopy
 
 from anemoi.utils.humanize import did_you_mean
 from earthkit.data import from_source
@@ -43,25 +42,27 @@ def normalise_time_delta(t):
     return t
 
 
-def _expand_mars_request(request, date, date_key="date"):
+def _expand_mars_request(request, date, request_already_using_valid_datetime=False, date_key="date"):
     requests = []
     step = to_list(request.get("step", [0]))
     for s in step:
-        r = deepcopy(request)
-
-        if isinstance(s, str) and "-" in s:
-            assert s.count("-") == 1, s
-        # this takes care of the cases where the step is a period such as 0-24 or 12-24
-        hours = int(str(s).split("-")[-1])
-
-        base = date - datetime.timedelta(hours=hours)
-        r.update(
-            {
-                date_key: base.strftime("%Y%m%d"),
-                "time": base.strftime("%H%M"),
-                "step": s,
-            }
-        )
+        r = request.copy()
+
+        if not request_already_using_valid_datetime:
+
+            if isinstance(s, str) and "-" in s:
+                assert s.count("-") == 1, s
+            # this takes care of the cases where the step is a period such as 0-24 or 12-24
+            hours = int(str(s).split("-")[-1])
+
+            base = date - datetime.timedelta(hours=hours)
+            r.update(
+                {
+                    date_key: base.strftime("%Y%m%d"),
+                    "time": base.strftime("%H%M"),
+                    "step": s,
+                }
+            )
 
         for pproc in ("grid", "rotation", "frame", "area", "bitmap", "resol"):
             if pproc in r:
@@ -73,13 +74,18 @@ def _expand_mars_request(request, date, date_key="date"):
     return requests
 
 
-def factorise_requests(dates, *requests, date_key="date"):
+def factorise_requests(dates, *requests, request_already_using_valid_datetime=False, date_key="date"):
     updates = []
     for req in requests:
         # req = normalise_request(req)
 
         for d in dates:
-            updates += _expand_mars_request(req, date=d, date_key=date_key)
+            updates += _expand_mars_request(
+                req,
+                date=d,
+                request_already_using_valid_datetime=request_already_using_valid_datetime,
+                date_key=date_key,
+            )
 
     compressed = Availability(updates)
     for r in compressed.iterate():
@@ -171,11 +177,32 @@ MARS_KEYS = [
 ]
 
 
-def mars(context, dates, *requests, date_key="date", **kwargs):
+def mars(context, dates, *requests, request_already_using_valid_datetime=False, date_key="date", **kwargs):
     if not requests:
         requests = [kwargs]
 
-    requests = factorise_requests(dates, *requests, date_key=date_key)
+    for r in requests:
+        # check for "Norway bug" where yaml transforms 'no' into False, etc.
+        for p in r.get("param", []):
+            if p is False:
+                raise ValueError(
+                    "'param' cannot be 'False'. If you wrote 'param: no' or 'param: off' in yaml, you may want to use quotes?"
+                )
+            if p is None:
+                raise ValueError(
+                    "'param' cannot be 'None'. If you wrote 'param: no' in yaml, you may want to use quotes?"
+                )
+            if p is True:
+                raise ValueError(
+                    "'param' cannot be 'True'. If you wrote 'param: on' in yaml, you may want to use quotes?"
+                )
+
+    requests = factorise_requests(
+        dates,
+        *requests,
+        request_already_using_valid_datetime=request_already_using_valid_datetime,
+        date_key=date_key,
+    )
     ds = from_source("empty")
     for r in requests:
         r = {k: v for k, v in r.items() if v != ("-",)}
@@ -191,7 +218,11 @@ def mars(context, dates, *requests, date_key="date", **kwargs):
                 raise ValueError(
                     f"⚠️ Unknown key {k}={v} in MARS request. Did you mean '{did_you_mean(k, MARS_KEYS)}' ?"
                 )
-        ds = ds + from_source("mars", **r)
+        try:
+            ds = ds + from_source("mars", **r)
+        except Exception as e:
+            if "File is empty:" not in str(e):
+                raise
     return ds
 
 

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

@@ -7,66 +7,8 @@
 # nor does it submit to any jurisdiction.
 #
 
-import glob
-
-from earthkit.data import from_source
-from earthkit.data.utils.patterns import Pattern
-
-
-def _expand(paths):
-    for path in paths:
-        if path.startswith("file://"):
-            path = path[7:]
-
-        if path.startswith("http://"):
-            yield path
-            continue
-
-        if path.startswith("https://"):
-            yield path
-            continue
-
-        for p in glob.glob(path):
-            yield p
-
-
-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_netcdfs(emoji, what, context, dates, path, *args, **kwargs):
-    given_paths = path if isinstance(path, list) else [path]
-
-    dates = [d.isoformat() for d in dates]
-    ds = from_source("empty")
-
-    for path in given_paths:
-        paths = Pattern(path, ignore_missing_keys=True).substitute(*args, date=dates, **kwargs)
-
-        levels = kwargs.get("level", kwargs.get("levelist"))
-
-        for path in _expand(paths):
-            context.trace(emoji, what.upper(), path)
-            s = from_source("opendap", path)
-            s = s.sel(
-                valid_datetime=dates,
-                param=kwargs["param"],
-                step=kwargs.get("step", 0),
-            )
-            if levels:
-                s = s.sel(levelist=levels)
-            ds = ds + s
-
-    check(what, ds, given_paths, valid_datetime=dates, **kwargs)
-
-    return ds
+from .xarray import load_many
 
 
 def execute(context, dates, path, *args, **kwargs):
-    return load_netcdfs("📁", "path", context, dates, path, *args, **kwargs)
+    return load_many("📁", context, dates, path, *args, **kwargs)

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

@@ -7,8 +7,9 @@
 # nor does it submit to any jurisdiction.
 #
 
-from .netcdf import load_netcdfs
+
+from .xarray import load_many
 
 
 def execute(context, dates, url, *args, **kwargs):
-    return load_netcdfs("🌐", "url", context, dates, url, *args, **kwargs)
+    return load_many("🌐", context, dates, url, *args, **kwargs)

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)