anemoi-utils 0.4.12__py3-none-any.whl → 0.4.14__py3-none-any.whl

anemoi/utils/grids.py

@@ -8,9 +8,7 @@
 # nor does it submit to any jurisdiction.
 
 
-"""Utilities for working with grids.
-
-"""
+"""Utilities for working with grids."""
 
 import logging
 import os
@@ -27,14 +25,46 @@ LOG = logging.getLogger(__name__)
 GRIDS_URL_PATTERN = "https://get.ecmwf.int/repository/anemoi/grids/grid-{name}.npz"
 
 
-def xyz_to_latlon(x, y, z):
+def xyz_to_latlon(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Convert Cartesian coordinates to latitude and longitude.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        The x coordinates
+    y : np.ndarray
+        The y coordinates
+    z : np.ndarray
+        The z coordinates
+
+    Returns
+    -------
+    tuple[np.ndarray, np.ndarray]
+        The latitude and longitude
+    """
     return (
         np.rad2deg(np.arcsin(np.minimum(1.0, np.maximum(-1.0, z)))),
         np.rad2deg(np.arctan2(y, x)),
     )
 
 
-def latlon_to_xyz(lat, lon, radius=1.0):
+def latlon_to_xyz(lat: np.ndarray, lon: np.ndarray, radius: float = 1.0) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Convert latitude and longitude to Cartesian coordinates.
+
+    Parameters
+    ----------
+    lat : np.ndarray
+        The latitudes
+    lon : np.ndarray
+        The longitudes
+    radius : float, optional
+        The radius of the sphere, by default 1.0
+
+    Returns
+    -------
+    tuple[np.ndarray, np.ndarray, np.ndarray]
+        The x, y, and z coordinates
+    """
     # https://en.wikipedia.org/wiki/Geographic_coordinate_conversion#From_geodetic_to_ECEF_coordinates
     # We assume that the Earth is a sphere of radius 1 so N(phi) = 1
     # We assume h = 0
@@ -54,7 +84,30 @@ def latlon_to_xyz(lat, lon, radius=1.0):
     return x, y, z
 
 
-def nearest_grid_points(source_latitudes, source_longitudes, target_latitudes, target_longitudes):
+def nearest_grid_points(
+    source_latitudes: np.ndarray,
+    source_longitudes: np.ndarray,
+    target_latitudes: np.ndarray,
+    target_longitudes: np.ndarray,
+) -> np.ndarray:
+    """Find the nearest grid points.
+
+    Parameters
+    ----------
+    source_latitudes : np.ndarray
+        The source latitudes
+    source_longitudes : np.ndarray
+        The source longitudes
+    target_latitudes : np.ndarray
+        The target latitudes
+    target_longitudes : np.ndarray
+        The target longitudes
+
+    Returns
+    -------
+    np.ndarray
+        The indices of the nearest grid points
+    """
     from scipy.spatial import cKDTree
 
     source_xyz = latlon_to_xyz(source_latitudes, source_longitudes)
@@ -68,7 +121,19 @@ def nearest_grid_points(source_latitudes, source_longitudes, target_latitudes, t
 
 
 @cached(collection="grids", encoding="npz")
-def _grids(name):
+def _grids(name: str) -> bytes:
+    """Get grid data by name.
+
+    Parameters
+    ----------
+    name : str
+        The name of the grid
+
+    Returns
+    -------
+    bytes
+        The grid data
+    """
     from anemoi.utils.config import load_config
 
     user_path = load_config().get("utils", {}).get("grids_path")
@@ -88,7 +153,19 @@ def _grids(name):
     return response.content
 
 
-def grids(name):
+def grids(name: str) -> dict:
+    """Load grid data by name.
+
+    Parameters
+    ----------
+    name : str
+        The name of the grid
+
+    Returns
+    -------
+    dict
+        The grid data
+    """
     if name.endswith(".npz"):
         return dict(np.load(name))
 

anemoi/utils/hindcasts.py

@@ -9,29 +9,46 @@
 
 
 import datetime
+from typing import Iterator
+from typing import List
+from typing import Tuple
 
 
 class HindcastDatesTimes:
-    """The HindcastDatesTimes class is an iterator that generates datetime objects within a given range."""
+    """The HindcastDatesTimes class is an iterator that generates datetime objects within a given range.
 
-    def __init__(self, reference_dates, years=20):
-        """_summary_
+    Attributes
+    ----------
+    reference_dates : List[datetime.datetime]
+        List of reference dates.
+    years : int
+        Number of years to go back from each reference date.
+    """
+
+    def __init__(self, reference_dates: List[datetime.datetime], years: int = 20):
+        """Initialize the HindcastDatesTimes iterator.
 
         Parameters
         ----------
-        reference_dates : _type_
-            _description_
+        reference_dates : List[datetime.datetime]
+            List of reference dates.
         years : int, optional
-            _description_, by default 20
+            Number of years to go back from each reference date, by default 20.
         """
-
         self.reference_dates = reference_dates
 
         assert isinstance(years, int), f"years must be an integer, got {years}"
         assert years > 0, f"years must be greater than 0, got {years}"
         self.years = years
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Tuple[datetime.datetime, datetime.datetime]]:
+        """Generate tuples of past dates and their corresponding reference dates.
+
+        Yields
+        ------
+        Iterator[Tuple[datetime.datetime, datetime.datetime]]
+            Tuples containing a past date and its corresponding reference date.
+        """
         for reference_date in self.reference_dates:
             year, month, day = reference_date.year, reference_date.month, reference_date.day
             if (month, day) == (2, 29):