anemoi-datasets 0.5.15__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/field.py

@@ -11,9 +11,14 @@
 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.core.fieldlist import Field
+from earthkit.data import Field
 from earthkit.data.core.fieldlist import math
+from numpy.typing import NDArray
 
 from .coordinates import extract_single_value
 from .coordinates import is_scalar
@@ -23,19 +28,41 @@ LOG = logging.getLogger(__name__)
 
 
 class EmptyFieldList:
-    def __len__(self):
+    """A class to represent an empty list of fields."""
+
+    def __len__(self) -> int:
+        """Return the length of the field list."""
         return 0
 
-    def __getitem__(self, i):
+    def __getitem__(self, i: int) -> Any:
+        """Raise an IndexError when trying to access an item.
+
+        Returns
+        -------
+        Any
+            This method does not return anything as it raises an IndexError.
+
+        Raises
+        ------
+        IndexError
+            Always raised to indicate that the list is empty.
+
+        Args
+        ----
+        i : int
+            Index of the item to access.
+        """
         raise IndexError(i)
 
     def __repr__(self) -> str:
+        """Return a string representation of the EmptyFieldList."""
         return "EmptyFieldList()"
 
 
 class XArrayField(Field):
+    """A class to represent a field in an XArray dataset."""
 
-    def __init__(self, owner, selection):
+    def __init__(self, owner: Any, selection: Any) -> None:
         """Create a new XArrayField object.
 
         Parameters
@@ -69,10 +96,29 @@ class XArrayField(Field):
             raise ValueError("Invalid shape for selection")
 
     @property
-    def shape(self):
+    def shape(self) -> Tuple[int, int]:
+        """Return the shape of the field."""
         return self._shape
 
-    def to_numpy(self, flatten=False, dtype=None, index=None):
+    def to_numpy(
+        self, flatten: bool = False, dtype: Optional[type] = None, index: Optional[int] = None
+    ) -> NDArray[Any]:
+        """Convert the selection to a numpy array.
+
+        Returns
+        -------
+        NDArray[Any]
+            The selection converted to a numpy array.
+
+        Args
+        ----
+        flatten : bool, optional
+            Whether to flatten the array, by default False.
+        dtype : Optional[type], optional
+            Data type of the array, by default None.
+        index : Optional[int], optional
+            Index to select a specific element, by default None.
+        """
         if index is not None:
             values = self.selection[index]
         else:
@@ -86,34 +132,53 @@ class XArrayField(Field):
         return values  # .reshape(self.shape)
 
     @cached_property
-    def _metadata(self):
+    def _metadata(self) -> XArrayMetadata:
+        """Return the metadata of the field."""
         return XArrayMetadata(self)
 
-    def grid_points(self):
+    def grid_points(self) -> Any:
+        """Return the grid points of the field."""
         return self.owner.grid_points()
 
-    def to_latlon(self, flatten=True):
+    def to_latlon(self, flatten: bool = True) -> Dict[str, Any]:
+        """Convert the selection to latitude and longitude coordinates.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The latitude and longitude coordinates.
+
+        Args
+        ----
+        flatten : bool, optional
+            Whether to flatten the coordinates, by default True.
+        """
         assert flatten
         return dict(lat=self.latitudes, lon=self.longitudes)
 
     @property
-    def resolution(self):
+    def resolution(self) -> Optional[Any]:
+        """Return the resolution of the field."""
         return None
 
     @property
-    def grid_mapping(self):
+    def grid_mapping(self) -> Any:
+        """Return the grid mapping of the field."""
         return self.owner.grid_mapping
 
     @property
-    def latitudes(self):
+    def latitudes(self) -> Any:
+        """Return the latitudes of the field."""
         return self.owner.latitudes
 
     @property
-    def longitudes(self):
+    def longitudes(self) -> Any:
+        """Return the longitudes of the field."""
         return self.owner.longitudes
 
     @property
-    def forecast_reference_time(self):
+    def forecast_reference_time(self) -> datetime.datetime:
+        """Return the forecast reference time of the field."""
         date, time = self.metadata("date", "time")
         assert len(time) == 4, time
         assert len(date) == 8, date
@@ -121,9 +186,22 @@ class XArrayField(Field):
         time = int(time) // 100
         return datetime.datetime(yyyymmdd // 10000, yyyymmdd // 100 % 100, yyyymmdd % 100, time)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Return a string representation of the field."""
         return repr(self._metadata)
 
-    def _values(self, dtype=None):
+    def _values(self, dtype: Optional[type] = None) -> Any:
+        """Return the values of the selection.
+
+        Returns
+        -------
+        Any
+            The values of the selection.
+
+        Args
+        ----
+        dtype : Optional[type], optional
+            Data type of the values, by default None.
+        """
         # we don't use .values as this will download the data
         return self.selection

anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/fieldlist.py

@@ -10,9 +10,15 @@
 
 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
-from earthkit.data.core.fieldlist import FieldList
+from earthkit.data import FieldList
 
 from .field import EmptyFieldList
 from .flavour import CoordinateGuesser
@@ -25,19 +31,49 @@ 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)
+    """A class to represent a list of fields from an xarray Dataset."""
 
-    def __repr__(self):
+    def __init__(self, ds: xr.Dataset, variables: List[Variable]) -> None:
+        """Initialize the XarrayFieldList.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            The xarray Dataset.
+        variables : List[Variable]
+            The list of variables.
+        """
+        self.ds: xr.Dataset = ds
+        self.variables: List[Variable] = variables.copy()
+        self.total_length: int = sum(v.length for v in variables)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the XarrayFieldList."""
         return f"XarrayFieldList({self.total_length})"
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Return the length of the XarrayFieldList."""
         return self.total_length
 
-    def __getitem__(self, i):
-        k = i
+    def __getitem__(self, i: int) -> Any:
+        """Get an item from the XarrayFieldList by index.
+
+        Parameters
+        ----------
+        i : int
+            The index of the item to get.
+
+        Returns
+        -------
+        Any
+            The item at the specified index.
+
+        Raises
+        ------
+        IndexError
+            If the index is out of range.
+        """
+        k: int = i
 
         if i < 0:
             i = self.total_length + i
@@ -50,12 +86,33 @@ class XarrayFieldList(FieldList):
         raise IndexError(k)
 
     @classmethod
-    def from_xarray(cls, ds, *, flavour=None, patch=None):
+    def from_xarray(
+        cls,
+        ds: xr.Dataset,
+        *,
+        flavour: Optional[Union[str, Dict[str, Any]]] = None,
+        patch: Optional[Dict[str, Any]] = None,
+    ) -> "XarrayFieldList":
+        """Create an XarrayFieldList from an xarray Dataset.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            The xarray Dataset to create the field list from.
+        flavour : Optional[Union[str, Dict[str, Any]]], optional
+            The flavour to use for guessing coordinates.
+        patch : Optional[Dict[str, Any]], optional
+            The patch to apply to the dataset.
 
+        Returns
+        -------
+        XarrayFieldList
+            The created XarrayFieldList.
+        """
         if patch is not None:
             ds = patch_dataset(ds, patch)
 
-        variables = []
+        variables: List[Variable] = []
 
         if isinstance(flavour, str):
             with open(flavour) as f:
@@ -64,9 +121,9 @@ class XarrayFieldList(FieldList):
                 else:
                     flavour = json.load(f)
 
-        if isinstance(flavour, dict):
-            flavour_coords = [coords["name"] for coords in flavour["rules"].values()]
-            ds_dims = [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]})
@@ -75,10 +132,10 @@ class XarrayFieldList(FieldList):
 
         guess = CoordinateGuesser.from_flavour(ds, flavour)
 
-        skip = set()
+        skip: set = set()
 
-        def _skip_attr(v, attr_name):
-            attr_val = getattr(v, attr_name, "")
+        def _skip_attr(v: Any, attr_name: str) -> None:
+            attr_val: str = getattr(v, attr_name, "")
             if isinstance(attr_val, str):
                 skip.update(attr_val.split(" "))
 
@@ -97,7 +154,7 @@ class XarrayFieldList(FieldList):
                 continue
 
             variable = ds[name]
-            coordinates = []
+            coordinates: List[Any] = []
 
             for coord in variable.coords:
 
@@ -108,7 +165,7 @@ class XarrayFieldList(FieldList):
                     c.is_dim = False
                 coordinates.append(c)
 
-            grid_coords = sum(1 for c in coordinates if c.is_grid and c.is_dim)
+            grid_coords: int = sum(1 for c in coordinates if c.is_grid and c.is_dim)
             assert grid_coords <= 2
 
             if grid_coords < 2:
@@ -128,14 +185,21 @@ class XarrayFieldList(FieldList):
 
         return cls(ds, variables)
 
-    def sel(self, **kwargs):
-        """Override the FieldList's sel method
+    def sel(self, **kwargs: Any) -> FieldList:
+        """Select fields from the XarrayFieldList based on criteria.
+
+        Parameters
+        ----------
+        kwargs : dict
+            The selection criteria.
 
         Returns
         -------
         FieldList
-            The new FieldList
+            The new FieldList with selected fields.
+        """
 
+        """
         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`, ...)
@@ -146,19 +210,20 @@ class XarrayFieldList(FieldList):
             So we get an extra chance to filter the fields by the metadata.
         """
 
-        variables = []
-        count = 0
+        variables: List[Variable] = []
+        count: int = 0
 
         for v in self.variables:
 
             # 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 = {}
+                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