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

anemoi/datasets/data/subset.py

@@ -10,12 +10,9 @@
 
 import datetime
 import logging
+from collections.abc import Sequence
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Sequence
-from typing import Set
 from typing import Union
 
 import numpy as np
@@ -95,7 +92,7 @@ def _end(a: int, b: int, dates: NDArray[np.datetime64]) -> int:
         return b
 
 
-def _combine_reasons(reason1: Dict[str, Any], reason2: Dict[str, Any], dates: NDArray[np.datetime64]) -> Dict[str, Any]:
+def _combine_reasons(reason1: dict[str, Any], reason2: dict[str, Any], dates: NDArray[np.datetime64]) -> dict[str, Any]:
     """Combine two reason dictionaries.
 
     Parameters:
@@ -126,7 +123,7 @@ class Subset(Forwards):
     reason (Dict[str, Any]): Dictionary of reasons.
     """
 
-    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: Dict[str, Any]) -> None:
+    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: dict[str, Any]) -> None:
         """Initialize the Subset.
 
         Parameters:
@@ -140,8 +137,8 @@ class Subset(Forwards):
             dataset = dataset.dataset
 
         self.dataset: Dataset = dataset
-        self.indices: List[int] = list(indices)
-        self.reason: Dict[str, Any] = {k: v for k, v in reason.items() if v is not None}
+        self.indices: list[int] = list(indices)
+        self.reason: dict[str, Any] = {k: v for k, v in reason.items() if v is not None}
 
         # Forward other properties to the super dataset
         super().__init__(dataset)
@@ -274,10 +271,10 @@ class Subset(Forwards):
         return f"Subset({self.dataset},{self.dates[0]}...{self.dates[-1]}/{self.frequency})"
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing indices of the subset."""
         missing = self.dataset.missing
-        result: Set[int] = set()
+        result: set[int] = set()
         for j, i in enumerate(self.indices):
             if i in missing:
                 result.add(j)
@@ -291,7 +288,7 @@ class Subset(Forwards):
         """
         return Node(self, [self.dataset.tree()], **self.reason)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the forwards subclass.
 
         Returns:

anemoi/datasets/data/unchecked.py

@@ -10,14 +10,10 @@
 
 import datetime
 import logging
+from collections.abc import Callable
 from functools import cached_property
 from functools import wraps
 from typing import Any
-from typing import Callable
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Set
 
 import numpy as np
 from numpy.typing import NDArray
@@ -163,13 +159,13 @@ class Unchecked(Combined):
 
     @check("check_same_variables")
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Get the mapping of variable names to their indices."""
         raise NotImplementedError()
 
     @check("check_same_variables")
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Get the list of variables in the dataset."""
         raise NotImplementedError()
 
@@ -181,12 +177,12 @@ class Unchecked(Combined):
 
     @check("check_same_variables")
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Get the statistics of the dataset."""
         raise NotImplementedError()
 
     @check("check_same_variables")
-    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         """Get the statistics tendencies of the dataset.
 
         Parameters
@@ -207,9 +203,9 @@ class Unchecked(Combined):
         raise NotImplementedError()
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing data indices."""
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result

anemoi/datasets/data/xy.py

@@ -11,10 +11,6 @@
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Set
-from typing import Tuple
 
 from .dataset import Dataset
 from .dataset import FullIndex
@@ -29,7 +25,7 @@ LOG = logging.getLogger(__name__)
 class ZipBase(Combined):
     """Base class for handling zipped datasets."""
 
-    def __init__(self, datasets: List[Any], check_compatibility: bool = True) -> None:
+    def __init__(self, datasets: list[Any], check_compatibility: bool = True) -> None:
         """Initialize ZipBase with a list of datasets.
 
         Parameters
@@ -58,7 +54,7 @@ class ZipBase(Combined):
         new_parents = [parent.clone(ds) for ds in self.datasets]
         return self.clone(new_parents)
 
-    def clone(self, datasets: List[Any]) -> "ZipBase":
+    def clone(self, datasets: list[Any]) -> "ZipBase":
         """Clone the ZipBase with new datasets.
 
         Parameters
@@ -81,7 +77,11 @@ class ZipBase(Combined):
         Node
             Tree representation of the datasets.
         """
-        return Node(self, [d.tree() for d in self.datasets], check_compatibility=self._check_compatibility)
+        return Node(
+            self,
+            [d.tree() for d in self.datasets],
+            check_compatibility=self._check_compatibility,
+        )
 
     def __len__(self) -> int:
         """Get the length of the smallest dataset.
@@ -93,7 +93,7 @@ class ZipBase(Combined):
         """
         return min(len(d) for d in self.datasets)
 
-    def __getitem__(self, n: FullIndex) -> Tuple[Any, ...]:
+    def __getitem__(self, n: FullIndex) -> tuple[Any, ...]:
         """Get the item at the specified index from all datasets.
 
         Parameters
@@ -145,55 +145,55 @@ class ZipBase(Combined):
         pass
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the set of missing indices from all datasets."""
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result
 
     @property
-    def shape(self) -> Tuple[Any, ...]:
+    def shape(self) -> tuple[Any, ...]:
         """Get the shape of all datasets."""
         return tuple(d.shape for d in self.datasets)
 
     @property
-    def field_shape(self) -> Tuple[Any, ...]:
+    def field_shape(self) -> tuple[Any, ...]:
         """Get the field shape of all datasets."""
         return tuple(d.shape for d in self.datasets)
 
     @property
-    def latitudes(self) -> Tuple[Any, ...]:
+    def latitudes(self) -> tuple[Any, ...]:
         """Get the latitudes of all datasets."""
         return tuple(d.latitudes for d in self.datasets)
 
     @property
-    def longitudes(self) -> Tuple[Any, ...]:
+    def longitudes(self) -> tuple[Any, ...]:
         """Get the longitudes of all datasets."""
         return tuple(d.longitudes for d in self.datasets)
 
     @property
-    def dtype(self) -> Tuple[Any, ...]:
+    def dtype(self) -> tuple[Any, ...]:
         """Get the data types of all datasets."""
         return tuple(d.dtype for d in self.datasets)
 
     @property
-    def grids(self) -> Tuple[Any, ...]:
+    def grids(self) -> tuple[Any, ...]:
         """Get the grids of all datasets."""
         return tuple(d.grids for d in self.datasets)
 
     @property
-    def statistics(self) -> Tuple[Any, ...]:
+    def statistics(self) -> tuple[Any, ...]:
         """Get the statistics of all datasets."""
         return tuple(d.statistics for d in self.datasets)
 
     @property
-    def resolution(self) -> Tuple[Any, ...]:
+    def resolution(self) -> tuple[Any, ...]:
         """Get the resolution of all datasets."""
         return tuple(d.resolution for d in self.datasets)
 
     @property
-    def name_to_index(self) -> Tuple[Any, ...]:
+    def name_to_index(self) -> tuple[Any, ...]:
         """Get the name to index mapping of all datasets."""
         return tuple(d.name_to_index for d in self.datasets)
 
@@ -210,6 +210,10 @@ class ZipBase(Combined):
         if self._check_compatibility:
             super().check_compatibility(d1, d2)
 
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
+        """Get the metadata specific to the subclass."""
+        return {}
+
 
 class Zip(ZipBase):
     """Class for handling zipped datasets."""
@@ -223,7 +227,7 @@ class XY(ZipBase):
     pass
 
 
-def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
+def xy_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> XY:
     """Factory function to create an XY instance.
 
     Parameters
@@ -256,7 +260,7 @@ def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
     return XY(datasets, check_compatibility=check_compatibility)._subset(**kwargs)
 
 
-def zip_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Zip:
+def zip_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Zip:
     """Factory function to create a Zip instance.
 
     Parameters

anemoi/datasets/dates/__init__.py

@@ -10,15 +10,10 @@
 
 import datetime
 import warnings
+from collections.abc import Iterator
 from functools import reduce
 from math import gcd
 from typing import Any
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 # from anemoi.utils.dates import as_datetime
 from anemoi.utils.dates import DateTimes
@@ -29,7 +24,7 @@ from anemoi.utils.hindcasts import HindcastDatesTimes
 from anemoi.utils.humanize import print_dates
 
 
-def extend(x: Union[str, List[Any], Tuple[Any, ...]]) -> Iterator[datetime.datetime]:
+def extend(x: str | list[Any] | tuple[Any, ...]) -> Iterator[datetime.datetime]:
     """Extend a date range or list of dates into individual datetime objects.
 
     Args:
@@ -86,7 +81,7 @@ class DatesProvider:
     3
     """
 
-    def __init__(self, missing: Optional[List[Union[str, datetime.datetime]]] = None) -> None:
+    def __init__(self, missing: list[str | datetime.datetime] | None = None) -> None:
         """Initialize the DatesProvider with optional missing dates.
 
         Parameters
@@ -168,7 +163,7 @@ class ValuesDates(DatesProvider):
         **kwargs (Any): Additional arguments.
     """
 
-    def __init__(self, values: List[Union[str, datetime.datetime]], **kwargs: Any) -> None:
+    def __init__(self, values: list[str | datetime.datetime], **kwargs: Any) -> None:
         """Initialize ValuesDates with a list of values.
 
         Args:
@@ -188,7 +183,7 @@ class ValuesDates(DatesProvider):
         """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the ValuesDates instance to a dictionary.
 
         Returns
@@ -215,9 +210,9 @@ class StartEndDates(DatesProvider):
 
     def __init__(
         self,
-        start: Union[str, datetime.datetime],
-        end: Union[str, datetime.datetime],
-        frequency: Union[int, str] = 1,
+        start: str | datetime.datetime,
+        end: str | datetime.datetime,
+        frequency: int | str = 1,
         **kwargs: Any,
     ) -> None:
         """Initialize StartEndDates with start, end, and frequency.
@@ -259,7 +254,7 @@ class StartEndDates(DatesProvider):
 
         super().__init__(missing=missing)
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the StartEndDates instance to a dictionary.
 
         Returns
@@ -314,9 +309,9 @@ class HindcastsDates(DatesProvider):
 
     def __init__(
         self,
-        start: Union[str, List[str]],
-        end: Union[str, List[str]],
-        steps: List[int] = [0],
+        start: str | list[str],
+        end: str | list[str],
+        steps: list[int] = [0],
         years: int = 20,
         **kwargs: Any,
     ) -> None:
@@ -403,7 +398,7 @@ class HindcastsDates(DatesProvider):
         """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the HindcastsDates instance to a dictionary.
 
         Returns

anemoi/datasets/dates/groups.py

@@ -12,19 +12,16 @@ import datetime
 import itertools
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Callable
+from collections.abc import Iterator
 from functools import cached_property
 from typing import Any
-from typing import Callable
-from typing import Iterator
-from typing import List
-from typing import Tuple
-from typing import Union
 
 from anemoi.datasets.dates import DatesProvider
 from anemoi.datasets.dates import as_datetime
 
 
-def _shorten(dates: Union[List[datetime.datetime], Tuple[datetime.datetime, ...]]) -> Union[str, List[str]]:
+def _shorten(dates: list[datetime.datetime] | tuple[datetime.datetime, ...]) -> str | list[str]:
     """Shorten the list of dates for display.
 
     Args:
@@ -43,7 +40,7 @@ def _shorten(dates: Union[List[datetime.datetime], Tuple[datetime.datetime, ...]
 class GroupOfDates:
     """A class to represent a group of dates."""
 
-    def __init__(self, dates: List[datetime.datetime], provider: DatesProvider, partial_ok: bool = False) -> None:
+    def __init__(self, dates: list[datetime.datetime], provider: DatesProvider, partial_ok: bool = False) -> None:
         assert isinstance(provider, DatesProvider), type(provider)
         assert isinstance(dates, list)
 
@@ -197,10 +194,10 @@ class Groups:
 class Filter:
     """A class to filter out missing dates."""
 
-    def __init__(self, missing: List[datetime.datetime]) -> None:
-        self.missing = set(as_datetime(m) for m in missing)
+    def __init__(self, missing: list[datetime.datetime]) -> None:
+        self.missing = {as_datetime(m) for m in missing}
 
-    def __call__(self, dates: List[datetime.datetime]) -> List[datetime.datetime]:
+    def __call__(self, dates: list[datetime.datetime]) -> list[datetime.datetime]:
         """Filter out missing dates from the list of dates.
 
         Args:

anemoi/datasets/grids.py

@@ -11,10 +11,6 @@
 import base64
 import logging
 from typing import Any
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -94,7 +90,7 @@ def plot_mask(
 
 # TODO: Use the one from anemoi.utils.grids instead
 # from anemoi.utils.grids import ...
-def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> Tuple[NDArray[Any], NDArray[Any]]:
+def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> tuple[NDArray[Any], NDArray[Any]]:
     """Convert Cartesian coordinates to latitude and longitude.
 
     Parameters
@@ -121,7 +117,7 @@ def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> Tuple[ND
 # from anemoi.utils.grids import ...
 def latlon_to_xyz(
     lat: NDArray[Any], lon: NDArray[Any], radius: float = 1.0
-) -> Tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
+) -> tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
     """Convert latitude and longitude to Cartesian coordinates.
 
     Parameters
@@ -272,8 +268,8 @@ def cutout_mask(
     global_lons: NDArray[Any],
     cropping_distance: float = 2.0,
     neighbours: int = 5,
-    min_distance_km: Optional[Union[int, float]] = None,
-    plot: Optional[str] = None,
+    min_distance_km: int | float | None = None,
+    plot: str | None = None,
 ) -> NDArray[Any]:
     """Return a mask for the points in [global_lats, global_lons] that are inside of [lats, lons].
 
@@ -465,7 +461,7 @@ def thinning_mask(
     return np.array([i for i in indices])
 
 
-def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> List[int]:
+def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> list[int]:
     """Find the outline of the grid points.
 
     Parameters
@@ -605,6 +601,7 @@ def nearest_grid_points(
     target_latitudes: NDArray[Any],
     target_longitudes: NDArray[Any],
     max_distance: float = None,
+    k: int = 1,
 ) -> NDArray[Any]:
     """Find the nearest grid points from source to target coordinates.
 
@@ -621,6 +618,8 @@ def nearest_grid_points(
     max_distance: float, optional
         Maximum distance between nearest point and point to interpolate. Defaults to None.
         For example, 1e-3 is 1 km.
+    k : int, optional
+        The number of k closest neighbors to consider for interpolation
 
     Returns
     -------
@@ -637,10 +636,10 @@ def nearest_grid_points(
     target_xyz = latlon_to_xyz(target_latitudes, target_longitudes)
     target_points = np.array(target_xyz).transpose()
     if max_distance is None:
-        _, indices = cKDTree(source_points).query(target_points, k=1)
+        distances, indices = cKDTree(source_points).query(target_points, k=k)
     else:
-        _, indices = cKDTree(source_points).query(target_points, k=1, distance_upper_bound=max_distance)
-    return indices
+        distances, indices = cKDTree(source_points).query(target_points, k=k, distance_upper_bound=max_distance)
+    return distances, indices
 
 
 if __name__ == "__main__":

anemoi/datasets/testing.py

@@ -12,19 +12,17 @@
 
 import logging
 from typing import Any
-from typing import List
-from typing import Optional
 
 LOG = logging.getLogger(__name__)
 
 
 def assert_field_list(
-    fs: List[Any],
-    size: Optional[int] = None,
-    start: Optional[Any] = None,
-    end: Optional[Any] = None,
+    fs: list[Any],
+    size: int | None = None,
+    start: Any | None = None,
+    end: Any | None = None,
     constant: bool = False,
-    skip: Optional[Any] = None,
+    skip: Any | None = None,
 ) -> None:
     """Asserts various properties of a list of fields.
 
@@ -85,3 +83,91 @@ def assert_field_list(
     assert south >= -90, south
     assert east <= 360, east
     assert west >= -180, west
+
+
+class IndexTester:
+    """Class to test indexing of datasets."""
+
+    def __init__(self, ds: Any) -> None:
+        """Initialise the IndexTester.
+
+        Parameters
+        ----------
+        ds : Any
+            Dataset.
+        """
+        self.ds = ds
+        self.np = ds[:]  # Numpy array
+
+        assert self.ds.shape == self.np.shape, (self.ds.shape, self.np.shape)
+        assert (self.ds == self.np).all()
+
+    def __getitem__(self, index: Any) -> None:
+        """Test indexing.
+
+        Parameters
+        ----------
+        index : Any
+            Index.
+        """
+        LOG.info("IndexTester: %s", index)
+        if self.ds[index] is None:
+            assert False, (self.ds, index)
+
+        if not (self.ds[index] == self.np[index]).all():
+            assert (self.ds[index] == self.np[index]).all()
+
+
+def default_test_indexing(ds):
+
+    t = IndexTester(ds)
+
+    t[0:10, :, 0]
+    t[:, 0:3, 0]
+    # t[:, :, 0]
+    t[0:10, 0:3, 0]
+    t[:, :, :]
+
+    if ds.shape[1] > 2:  # Variable dimension
+        t[:, (1, 2), :]
+        t[:, (1, 2)]
+
+    t[0]
+    t[0, :]
+    t[0, 0, :]
+    t[0, 0, 0, :]
+
+    if ds.shape[2] > 1:  # Ensemble dimension
+        t[0:10, :, (0, 1)]
+
+    for i in range(3):
+        t[i]
+        start = 5 * i
+        end = len(ds) - 5 * i
+        step = len(ds) // 10
+
+        t[start:end:step]
+        t[start:end]
+        t[start:]
+        t[:end]
+        t[::step]
+
+
+class Trace:
+
+    def __init__(self, ds):
+        self.ds = ds
+        self.f = open("trace.txt", "a")
+
+    def __getattr__(self, name: str) -> Any:
+
+        print(name, file=self.f, flush=True)
+        return getattr(self.ds, name)
+
+    def __len__(self) -> int:
+        print("__len__", file=self.f, flush=True)
+        return len(self.ds)
+
+    def __getitem__(self, index: Any) -> Any:
+        print("__getitem__", file=self.f, flush=True)
+        return self.ds[index]