typing

Author: benbovy

xdggs/accessor.py

@@ -1,31 +1,64 @@
+import numpy.typing as npt
 import xarray as xr
 
 from xdggs.index import DGGSIndex
 
 
 @xr.register_dataset_accessor("dggs")
+@xr.register_dataarray_accessor("dggs")
 class DGGSAccessor:
-    def __init__(self, obj):
+    _obj: xr.Dataset | xr.DataArray
+    _index: DGGSIndex | None
+    _name: str
+
+    def __init__(self, obj: xr.Dataset | xr.DataArray):
         self._obj = obj
 
-        indexes = {k: idx for k, idx in obj.xindexes.items() if isinstance(idx, DGGSIndex)}
-        if len(indexes) > 1:
-            raise ValueError("Only one DGGSIndex per object is supported")
+        index = None
+        name = ""
+        for k, idx in obj.xindexes.items():
+            if isinstance(idx, DGGSIndex):
+                if index is not None:
+                    raise ValueError("Only one DGGSIndex per object is supported")
+                index = idx
+                name = k
+        self._name = name
+        self._index = index
+
+    @property
+    def index(self) -> DGGSIndex:
+        if self._index is None:
+            raise ValueError("no DGGSIndex found on this Dataset or DataArray")
+        return self._index
+
+    def sel_latlon(
+        self, latitude: npt.ArrayLike, longitude: npt.ArrayLike
+    ) -> xr.Dataset | xr.DataArray:
+        """Select grid cells from latitude/longitude data.
 
-        self._name, self._index = next(iter(indexes.items()))
+        Parameters
+        ----------
+        latitude : array-like
+            Latitude coordinates (degrees).
+        longitude : array-like
+            Longitude coordinates (degrees).
 
-    def sel_latlon(self, lat, lon):
-        """Point-wise, nearest-neighbor selection from lat/lon data."""
+        Returns
+        -------
+        subset
+            A new :py:class:`xarray.Dataset` or :py:class:`xarray.DataArray`
+            with all cells that contain the input latitude/longitude data points.
 
-        cell_indexers = {self._name: self._index._latlon2cellid(lat, lon)}
+        """
+        cell_indexers = {self._name: self.index._latlon2cellid(latitude, longitude)}
         return self._obj.sel(cell_indexers)
 
-    def assign_latlon_coords(self):
-        """Return a new object with latitude and longitude coordinates
-        of the cell centers."""
+    def assign_latlon_coords(self) -> xr.Dataset | xr.DataArray:
+        """Return a new Dataset or DataArray with new "latitude" and "longitude"
+        coordinates representing the grid cell centers."""
 
-        lat_data, lon_data = self._index.cell_centers
+        lat_data, lon_data = self.index.cell_centers
         return self._obj.assign_coords(
-            latitude=(self._index._dim, lat_data),
-            longitude=(self._index._dim, lon_data),
+            latitude=(self.index._dim, lat_data),
+            longitude=(self.index._dim, lon_data),
         )

xdggs/healpix.py

@@ -1,39 +1,57 @@
+from collections.abc import Mapping
+from typing import Any
+
 import healpy
+import numpy as np
+import xarray as xr
+from xarray.indexes import PandasIndex
 
 from xdggs.index import DGGSIndex
 from xdggs.utils import _extract_cell_id_variable, register_dggs
 
 
 @register_dggs("healpix")
 class HealpixIndex(DGGSIndex):
-    def __init__(self, cell_ids, dim, nside, nest, rot_latlon):
+    def __init__(
+        self,
+        cell_ids: Any | PandasIndex,
+        dim: str,
+        nside: int,
+        nest: bool,
+        rot_latlon: tuple[float, float],
+    ):
         super().__init__(cell_ids, dim)
 
         self._nside = nside
         self._nest = nest
         self._rot_latlon = rot_latlon
 
     @classmethod
-    def from_variables(cls, variables, *, options):
-        name, var, dim = _extract_cell_id_variable(variables)
+    def from_variables(
+        cls: type["HealpixIndex"],
+        variables: Mapping[Any, xr.Variable],
+        *,
+        options: Mapping[str, Any],
+    ) -> "HealpixIndex":
+        _, var, dim = _extract_cell_id_variable(variables)
 
         nside = var.attrs.get("nside", options.get("nside"))
         nest = var.attrs.get("nest", options.get("nest", False))
         rot_latlon = var.attrs.get("rot_latlon", options.get("rot_latlon", (0.0, 0.0)))
 
         return cls(var.data, dim, nside, nest, rot_latlon)
 
-    def _replace(self, new_pd_index):
+    def _replace(self, new_pd_index: PandasIndex):
         return type(self)(new_pd_index, self._dim, self._nside, self._nest, self._rot_latlon)
 
-    def _latlon2cellid(self, lat, lon):
+    def _latlon2cellid(self, lat: Any, lon: Any) -> np.ndarray:
         return healpy.ang2pix(self._nside, -lon, lat, lonlat=True, nest=self._nest)
 
-    def _cellid2latlon(self, cell_ids):
+    def _cellid2latlon(self, cell_ids: Any) -> tuple[np.ndarray, np.ndarray]:
         lon, lat = healpy.pix2ang(self._nside, cell_ids, nest=self._nest, lonlat=True)
         return lat, -lon
 
-    def _repr_inline_(self, max_width):
+    def _repr_inline_(self, max_width: int):
         return (
             f"HealpixIndex(nside={self._nside}, nest={self._nest}, rot_latlon={self._rot_latlon!r})"
         )

xdggs/index.py

@@ -1,10 +1,18 @@
+from collections.abc import Hashable, Mapping
+from typing import Any, Union
+
+import numpy as np
+import xarray as xr
 from xarray.indexes import Index, PandasIndex
 
 from .utils import GRID_REGISTRY, _extract_cell_id_variable
 
 
 class DGGSIndex(Index):
-    def __init__(self, cell_ids, dim):
+    _dim: str
+    _pd_index: PandasIndex
+
+    def __init__(self, cell_ids: Any | PandasIndex, dim: str):
         self._dim = dim
 
         if isinstance(cell_ids, PandasIndex):
@@ -13,18 +21,27 @@ def __init__(self, cell_ids, dim):
             self._pd_index = PandasIndex(cell_ids, dim)
 
     @classmethod
-    def from_variables(cls, variables, *, options):
-        name, var, dim = _extract_cell_id_variable(variables)
+    def from_variables(
+        cls: type["DGGSIndex"],
+        variables: Mapping[Any, xr.Variable],
+        *,
+        options: Mapping[str, Any],
+    ) -> "DGGSIndex":
+        _, var, _ = _extract_cell_id_variable(variables)
 
         grid_name = var.attrs["grid_name"]
         cls = GRID_REGISTRY[grid_name]
 
         return cls.from_variables(variables, options=options)
 
-    def create_variables(self, variables=None):
+    def create_variables(
+        self, variables: Mapping[Any, xr.Variable] | None = None
+    ) -> dict[Hashable, xr.Variable]:
         return self._pd_index.create_variables(variables)
 
-    def isel(self, indexers):
+    def isel(
+        self: "DGGSIndex", indexers: Mapping[Any, int | np.ndarray | xr.Variable]
+    ) -> Union["DGGSIndex", None]:
         new_pd_index = self._pd_index.isel(indexers)
         if new_pd_index is not None:
             return self._replace(new_pd_index)
@@ -36,17 +53,17 @@ def sel(self, labels, method=None, tolerance=None):
             raise ValueError("finding nearest grid cell has no meaning")
         return self._pd_index.sel(labels, method=method, tolerance=tolerance)
 
-    def _replace(self, new_pd_index):
+    def _replace(self, new_pd_index: PandasIndex):
         raise NotImplementedError()
 
-    def _latlon2cellid(self, lat, lon):
+    def _latlon2cellid(self, lat: Any, lon: Any) -> np.ndarray:
         """convert latitude / longitude points to cell ids."""
         raise NotImplementedError()
 
-    def _cellid2latlon(self, cell_ids):
-        """convert cell centers to latitude / longitude."""
+    def _cellid2latlon(self, cell_ids: Any) -> tuple[np.ndarray, np.ndarray]:
+        """convert cell ids to latitude / longitude (cell centers)."""
         raise NotImplementedError()
 
     @property
-    def cell_centers(self):
+    def cell_centers(self) -> tuple[np.ndarray, np.ndarray]:
         return self._cellid2latlon(self._pd_index.index.values)