anemoi-datasets 0.5.25__py3-none-any.whl → 0.5.27__py3-none-any.whl

anemoi/datasets/create/sources/tendencies.py

@@ -10,9 +10,6 @@
 import datetime
 from collections import defaultdict
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Tuple
 
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
@@ -63,7 +60,7 @@ def normalise_time_delta(t: Any) -> datetime.timedelta:
     return t
 
 
-def group_by_field(ds: Any) -> Dict[Tuple, List[Any]]:
+def group_by_field(ds: Any) -> dict[tuple, list[Any]]:
     """Groups fields by their metadata excluding 'date', 'time', and 'step'.
 
     Parameters
@@ -87,7 +84,7 @@ def group_by_field(ds: Any) -> Dict[Tuple, List[Any]]:
 
 
 @legacy_source(__file__)
-def tendencies(dates: List[datetime.datetime], time_increment: Any, **kwargs: Any) -> Any:
+def tendencies(dates: list[datetime.datetime], time_increment: Any, **kwargs: Any) -> Any:
     """Computes tendencies for the given dates and time increment.
 
     Parameters

anemoi/datasets/create/sources/xarray.py

@@ -8,8 +8,6 @@
 # nor does it submit to any jurisdiction.
 
 from typing import Any
-from typing import Dict
-from typing import Optional
 
 import earthkit.data as ekd
 
@@ -28,11 +26,11 @@ class XarraySourceBase(Source):
 
     emoji = "✖️"  # For tracing
 
-    options: Optional[Dict[str, Any]] = None
-    flavour: Optional[Dict[str, Any]] = None
-    patch: Optional[Dict[str, Any]] = None
+    options: dict[str, Any] | None = None
+    flavour: dict[str, Any] | None = None
+    patch: dict[str, Any] | None = None
 
-    path_or_url: Optional[str] = None
+    path_or_url: str | None = None
 
     def __init__(self, context: Any, path: str = None, url: str = None, *args: Any, **kwargs: Any):
         """Initialise the source.

anemoi/datasets/create/sources/xarray_support/__init__.py

@@ -10,17 +10,12 @@
 import datetime
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Union
 
 import earthkit.data as ekd
 import xarray as xr
 from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.sources.patterns import iterate_patterns
-from anemoi.datasets.data.stores import name_to_zarr_store
 
 from ..legacy import legacy_source
 from .fieldlist import XarrayFieldList
@@ -28,7 +23,7 @@ from .fieldlist import XarrayFieldList
 LOG = logging.getLogger(__name__)
 
 
-def check(what: str, ds: xr.Dataset, paths: List[str], **kwargs: Any) -> None:
+def check(what: str, ds: xr.Dataset, paths: list[str], **kwargs: Any) -> None:
     """Checks if the dataset has the expected number of fields.
 
     Parameters
@@ -54,12 +49,12 @@ def check(what: str, ds: xr.Dataset, paths: List[str], **kwargs: Any) -> None:
 def load_one(
     emoji: str,
     context: Any,
-    dates: List[str],
-    dataset: Union[str, xr.Dataset],
+    dates: list[str],
+    dataset: str | xr.Dataset,
     *,
-    options: Optional[Dict[str, Any]] = None,
-    flavour: Optional[str] = None,
-    patch: Optional[Any] = None,
+    options: dict[str, Any] | None = None,
+    flavour: str | None = None,
+    patch: Any | None = None,
     **kwargs: Any,
 ) -> ekd.FieldList:
     """Loads a single dataset.
@@ -89,28 +84,17 @@ def load_one(
         The loaded dataset.
     """
 
-    """
-    We manage the S3 client ourselves, 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`.
-    """
-
     if options is None:
         options = {}
 
     context.trace(emoji, dataset, options, kwargs)
 
-    if isinstance(dataset, str) and ".zarr" in dataset:
-        data = xr.open_zarr(name_to_zarr_store(dataset), **options)
-    elif "planetarycomputer" in dataset:
-        store = name_to_zarr_store(dataset)
-        if "store" in store:
-            data = xr.open_zarr(**store)
-        if "filename_or_obj" in store:
-            data = xr.open_dataset(**store)
+    if isinstance(dataset, str) and dataset.endswith(".zarr"):
+        # If the dataset is a zarr store, we need to use the zarr engine
+        options["engine"] = "zarr"
+
+    if isinstance(dataset, xr.Dataset):
+        data = dataset
     else:
         data = xr.open_dataset(dataset, **options)
 
@@ -119,7 +103,6 @@ def load_one(
     if len(dates) == 0:
         result = fs.sel(**kwargs)
     else:
-        print("dates", dates, kwargs)
         result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 
     if len(result) == 0:
@@ -130,7 +113,7 @@ def load_one(
             a = ["valid_datetime", k.metadata("valid_datetime", default=None)]
             for n in kwargs.keys():
                 a.extend([n, k.metadata(n, default=None)])
-            print([str(x) for x in a])
+            LOG.warning(f"{[str(x) for x in a]}")
 
             if i > 16:
                 break
@@ -140,7 +123,7 @@ def load_one(
     return result
 
 
-def load_many(emoji: str, context: Any, dates: List[datetime.datetime], pattern: str, **kwargs: Any) -> ekd.FieldList:
+def load_many(emoji: str, context: Any, dates: list[datetime.datetime], pattern: str, **kwargs: Any) -> ekd.FieldList:
     """Loads multiple datasets.
 
     Parameters
@@ -170,7 +153,7 @@ def load_many(emoji: str, context: Any, dates: List[datetime.datetime], pattern:
 
 
 @legacy_source("xarray")
-def execute(context: Any, dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+def execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
     """Executes the loading of datasets.
 
     Parameters

anemoi/datasets/create/sources/xarray_support/coordinates.py

@@ -13,10 +13,6 @@ from __future__ import annotations
 import datetime
 import logging
 from typing import Any
-from typing import Dict
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 import numpy as np
 import xarray as xr
@@ -95,6 +91,7 @@ class Coordinate:
     is_member = False
     is_x = False
     is_y = False
+    is_point = False
 
     def __init__(self, variable: xr.DataArray) -> None:
         """Initialize the coordinate.
@@ -106,7 +103,7 @@ class Coordinate:
         """
         self.variable = variable
         self.scalar = is_scalar(variable)
-        self.kwargs: Dict[str, Any] = {}  # Used when creating a new coordinate (reduced method)
+        self.kwargs: dict[str, Any] = {}  # Used when creating a new coordinate (reduced method)
 
     def __len__(self) -> int:
         """Get the length of the coordinate.
@@ -126,7 +123,7 @@ class Coordinate:
         str
             The string representation of the coordinate.
         """
-        return "%s[name=%s,values=%s,shape=%s]" % (
+        return "{}[name={},values={},shape={}]".format(
             self.__class__.__name__,
             self.variable.name,
             self.variable.values if self.scalar else len(self),
@@ -151,7 +148,7 @@ class Coordinate:
             **self.kwargs,
         )
 
-    def index(self, value: Union[Any, list, tuple]) -> Optional[Union[int, list]]:
+    def index(self, value: Any | list | tuple) -> int | list | None:
         """Return the index of the value in the coordinate.
 
         Parameters
@@ -171,7 +168,7 @@ class Coordinate:
                 return self._index_multiple(value)
         return self._index_single(value)
 
-    def _index_single(self, value: Any) -> Optional[int]:
+    def _index_single(self, value: Any) -> int | None:
         """Return the index of a single value in the coordinate.
 
         Parameters
@@ -204,7 +201,7 @@ class Coordinate:
 
         return None
 
-    def _index_multiple(self, value: list) -> Optional[list]:
+    def _index_multiple(self, value: list) -> list | None:
         """Return the indices of multiple values in the coordinate.
 
         Parameters
@@ -274,7 +271,7 @@ class TimeCoordinate(Coordinate):
     is_time = True
     mars_names = ("valid_datetime",)
 
-    def index(self, time: datetime.datetime) -> Optional[int]:
+    def index(self, time: datetime.datetime) -> int | None:
         """Return the index of the time in the coordinate.
 
         Parameters
@@ -296,7 +293,7 @@ class DateCoordinate(Coordinate):
     is_date = True
     mars_names = ("date",)
 
-    def index(self, date: datetime.datetime) -> Optional[int]:
+    def index(self, date: datetime.datetime) -> int | None:
         """Return the index of the date in the coordinate.
 
         Parameters
@@ -390,6 +387,13 @@ class EnsembleCoordinate(Coordinate):
         return value
 
 
+class PointCoordinate(Coordinate):
+    """Coordinate class for point data."""
+
+    is_point = True
+    mars_names = ("point",)
+
+
 class LongitudeCoordinate(Coordinate):
     """Coordinate class for longitude."""
 
@@ -428,7 +432,7 @@ class ScalarCoordinate(Coordinate):
     is_grid = False
 
     @property
-    def mars_names(self) -> Tuple[str, ...]:
+    def mars_names(self) -> tuple[str, ...]:
         """Get the MARS names for the coordinate."""
         return (self.variable.name,)
 

anemoi/datasets/create/sources/xarray_support/field.py

@@ -12,9 +12,6 @@ import datetime
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import Optional
-from typing import Tuple
 
 from earthkit.data import Field
 from earthkit.data.core.fieldlist import math
@@ -80,29 +77,33 @@ class XArrayField(Field):
         # Copy the metadata from the owner
         self._md = owner._metadata.copy()
 
+        aliases = {}
         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))
+                normalised = coordinate.normalise(extract_single_value(coord_value))
+                self._md[coord_name] = normalised
+                for alias in coordinate.mars_names:
+                    aliases[alias] = normalised
+
+        # Add metadata aliases (e.g. levelist == level) only if they are not already present
+        for alias, value in aliases.items():
+            if alias not in self._md:
+                self._md[alias] = 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")
+            raise ValueError(f"Invalid shape for selection {self._shape=}, {self.selection.shape=} {self.selection=}")
 
     @property
-    def shape(self) -> Tuple[int, int]:
+    def shape(self) -> tuple[int, int]:
         """Return the shape of the field."""
         return self._shape
 
-    def to_numpy(
-        self, flatten: bool = False, dtype: Optional[type] = None, index: Optional[int] = None
-    ) -> NDArray[Any]:
+    def to_numpy(self, flatten: bool = False, dtype: type | None = None, index: int | None = None) -> NDArray[Any]:
         """Convert the selection to a numpy array.
 
         Returns
@@ -140,7 +141,7 @@ class XArrayField(Field):
         """Return the grid points of the field."""
         return self.owner.grid_points()
 
-    def to_latlon(self, flatten: bool = True) -> Dict[str, Any]:
+    def to_latlon(self, flatten: bool = True) -> dict[str, Any]:
         """Convert the selection to latitude and longitude coordinates.
 
         Returns
@@ -157,7 +158,7 @@ class XArrayField(Field):
         return dict(lat=self.latitudes, lon=self.longitudes)
 
     @property
-    def resolution(self) -> Optional[Any]:
+    def resolution(self) -> Any | None:
         """Return the resolution of the field."""
         return None
 
@@ -188,9 +189,9 @@ class XArrayField(Field):
 
     def __repr__(self) -> str:
         """Return a string representation of the field."""
-        return repr(self._metadata)
+        return f"XArrayField({self._metadata})"
 
-    def _values(self, dtype: Optional[type] = None) -> Any:
+    def _values(self, dtype: type | None = None) -> Any:
         """Return the values of the selection.
 
         Returns

anemoi/datasets/create/sources/xarray_support/fieldlist.py

@@ -11,10 +11,6 @@
 import json
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Union
 
 import xarray as xr
 import yaml
@@ -33,7 +29,7 @@ LOG = logging.getLogger(__name__)
 class XarrayFieldList(FieldList):
     """A class to represent a list of fields from an xarray Dataset."""
 
-    def __init__(self, ds: xr.Dataset, variables: List[Variable]) -> None:
+    def __init__(self, ds: xr.Dataset, variables: list[Variable]) -> None:
         """Initialize the XarrayFieldList.
 
         Parameters
@@ -44,7 +40,7 @@ class XarrayFieldList(FieldList):
             The list of variables.
         """
         self.ds: xr.Dataset = ds
-        self.variables: List[Variable] = variables.copy()
+        self.variables: list[Variable] = variables.copy()
         self.total_length: int = sum(v.length for v in variables)
 
     def __repr__(self) -> str:
@@ -90,8 +86,8 @@ class XarrayFieldList(FieldList):
         cls,
         ds: xr.Dataset,
         *,
-        flavour: Optional[Union[str, Dict[str, Any]]] = None,
-        patch: Optional[Dict[str, Any]] = None,
+        flavour: str | dict[str, Any] | None = None,
+        patch: dict[str, Any] | None = None,
     ) -> "XarrayFieldList":
         """Create an XarrayFieldList from an xarray Dataset.
 
@@ -112,7 +108,7 @@ class XarrayFieldList(FieldList):
         if patch is not None:
             ds = patch_dataset(ds, patch)
 
-        variables: List[Variable] = []
+        variables: list[Variable] = []
 
         if isinstance(flavour, str):
             with open(flavour) as f:
@@ -121,9 +117,9 @@ class XarrayFieldList(FieldList):
                 else:
                     flavour = json.load(f)
 
-        if isinstance(flavour, Dict):
-            flavour_coords: List[str] = [coords["name"] for coords in flavour["rules"].values()]
-            ds_dims: List[str] = [dim for dim in ds._dims]
+        if isinstance(flavour, dict):
+            flavour_coords: list[str] = [coords["name"] for coords in flavour["rules"].values()]
+            ds_dims: list[str] = [dim for dim in ds._dims]
             for dim in ds_dims:
                 if dim in flavour_coords and dim not in ds._coord_names:
                     ds = ds.assign_coords({dim: ds[dim]})
@@ -154,7 +150,7 @@ class XarrayFieldList(FieldList):
                 continue
 
             variable = ds[name]
-            coordinates: List[Any] = []
+            coordinates: list[Any] = []
 
             for coord in variable.coords:
 
@@ -210,7 +206,7 @@ class XarrayFieldList(FieldList):
             So we get an extra chance to filter the fields by the metadata.
         """
 
-        variables: List[Variable] = []
+        variables: list[Variable] = []
         count: int = 0
 
         for v in self.variables:
@@ -223,7 +219,7 @@ class XarrayFieldList(FieldList):
 
             if match:
                 count += 1
-                missing: Dict[str, Any] = {}
+                missing: dict[str, Any] = {}
 
                 # Select from the variable's coordinates (time, level, number, ....)
                 # This may return a new variable with a isel() slice of the selection