Add RasterIndex.get_indexer for vectorized point lookup

src/rasterix/raster_index.py

@@ -6,6 +6,7 @@
 from typing import Any, Self, TypeVar, cast
 
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 import xproj
 from affine import Affine
@@ -209,6 +210,19 @@ def _assert_same_frame(ours: BoundingBox, theirs: BoundingBox, index: RasterInde
                 )
 
 
+def _containing_cells(
+    positions: npt.NDArray[np.float64],
+) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.bool_]]:
+    """Cell ``i`` whose cell-center positions are ``[i - 0.5, i + 0.5)``, and a mask of usable positions.
+
+    NaN, inf and positions too large for exact int64 cells are not usable; their cell is 0. No bounds check.
+    """
+    positions = np.asarray(positions, dtype=np.float64)
+    valid = np.abs(positions) < 2.0**53
+    cells = np.floor(np.where(valid, positions, 0.0) + 0.5).astype(np.int64)
+    return cells, valid
+
+
 class AffineTransform(CoordinateTransform):
     """Affine 2D transform wrapper."""
 
@@ -311,7 +325,7 @@ def forward(self, dim_positions: dict[str, Any]) -> dict[Hashable, Any]:
 
         return {self.coord_name: labels}
 
-    def _positions(self, labels: Any) -> np.ndarray:
+    def _positions(self, labels: npt.ArrayLike) -> npt.NDArray[np.float64]:
         """Unwrapped (float) cell positions of ``labels``."""
         labels = np.asarray(labels)
         if self.is_xaxis:
@@ -320,7 +334,7 @@ def _positions(self, labels: Any) -> np.ndarray:
             _, positions = ~self.affine * (np.zeros_like(labels), labels)
         return positions
 
-    def _wrap(self, positions: np.ndarray) -> np.ndarray:
+    def _wrap(self, positions: npt.ArrayLike) -> npt.NDArray[np.float64]:
         """Exact position wins; else wrap into [-0.5, P - 0.5) so that rounding gives [0, P)."""
         assert self.period_cells is not None
         positions = np.asarray(positions, dtype=np.float64)
@@ -396,6 +410,16 @@ def periodic_slice_indexer(self, labels: slice) -> slice | np.ndarray:
             return slice(int(positions[0]), int(positions[-1]) + 1)
         return positions
 
+    def get_indexer(self, labels: npt.ArrayLike) -> npt.NDArray[np.int64]:
+        """Index of the cell that contains each label, or -1; see :py:meth:`RasterIndex.get_indexer`."""
+        with np.errstate(invalid="ignore"):  # inf * 0 for the other axis
+            positions = self._positions(np.asarray(labels, dtype=np.float64))
+        cells, valid = _containing_cells(positions)
+        if self.period_cells is not None:
+            # wrap the integer cell, not the float position: no round-off at the seam
+            cells = np.mod(cells, self.period_cells)
+        return np.where(valid & (cells >= 0) & (cells < self.size), cells, -1)
+
     def reverse(self, coord_labels: dict[Hashable, Any]) -> dict[str, Any]:
         positions = self._positions(coord_labels[self.coord_name])
         if self.period_cells is not None:
@@ -652,6 +676,8 @@ class RasterIndex(Index, xproj.ProjIndexMixin):
         Get affine transform for pixel top-left corners.
     center_transform
         Get affine transform for pixel centers.
+    get_indexer
+        Find the cell that contains each point, or -1.
 
     Examples
     --------
@@ -1203,6 +1229,78 @@ def sel(self, labels: dict[Any, Any], method=None, tolerance=None) -> IndexSelRe
 
             return merge_sel_results(results)
 
+    def get_indexer(self, labels: Mapping[Hashable, npt.ArrayLike]) -> dict[str, npt.NDArray[np.int64]]:
+        """Find the cell that contains each point.
+
+        This follows :py:meth:`pandas.Index.get_indexer`: a point that is not in any cell
+        gives ``-1``. Unlike ``sel(..., method="nearest")``, a point outside the extent is
+        not snapped to the nearest edge cell.
+
+        Parameters
+        ----------
+        labels : mapping of coordinate name to array-like
+            Point coordinates, keyed by coordinate name (as for :py:meth:`sel`). For a
+            rectilinear transform, give the X labels, the Y labels or both; each axis is
+            looked up independently. For a rotated or skewed transform, give both, with the
+            same shape.
+
+        Returns
+        -------
+        dict of str to numpy.ndarray
+            ``int64`` cell positions, keyed by dimension name, with the shape of the input
+            labels. ``-1`` for a point outside the extent, and for NaN or inf.
+
+        Notes
+        -----
+        Let ``p`` be the position of a point in the corner-based pixel space of
+        :py:meth:`transform`. The point is in cell ``i`` if ``i <= p < i + 1``. That is, a
+        cell includes the edge at its lower position and excludes the edge at its higher
+        position. For a north-up raster (negative ``e``), a cell includes its top edge and
+        excludes its bottom edge, and the bottom edge of the raster is outside it. Away
+        from cell edges, the result is the same as ``sel(..., method="nearest")``.
+
+        On a periodic axis (``x_period`` or ``y_period``), a point is first wrapped by a
+        whole number of periods; a point outside one period is not a miss.
+
+        Examples
+        --------
+        >>> from affine import Affine
+        >>> index = RasterIndex.from_transform(Affine(1, 0, -10, 0, -1, 10), width=20, height=20)
+        >>> index.get_indexer({"x": [-9.5, 9.99, 10.0, 180.0]})
+        {'x': array([ 0, 19, -1, -1])}
+        """
+        if self._axis_independent:
+            unknown = set(labels) - set(self.xy_coord_names)
+            if unknown:
+                raise KeyError(f"No coordinate(s) {sorted(map(str, unknown))} in this RasterIndex.")
+            return {
+                index.dim: index.axis_transform.get_indexer(labels[index.axis_transform.coord_name])
+                for index in self._xy_indexes
+                if index.axis_transform.coord_name in labels
+            }
+
+        x_name, y_name = self.xy_coord_names
+        if set(labels) != {x_name, y_name}:
+            raise ValueError(
+                f"A rotated or skewed RasterIndex needs both coordinates {x_name!r} and {y_name!r}, "
+                f"got {sorted(map(str, labels))}."
+            )
+        x = np.asarray(labels[x_name], dtype=np.float64)
+        y = np.asarray(labels[y_name], dtype=np.float64)
+        if x.shape != y.shape:
+            raise ValueError(
+                f"Labels for {x_name!r} and {y_name!r} must have the same shape, got {x.shape} and {y.shape}."
+            )
+        transform = cast(AffineTransform, self._xxyy_index.transform)
+        with np.errstate(invalid="ignore"):
+            positions = transform.reverse({x_name: x, y_name: y})
+        x_dim, y_dim = self.xy_dims
+        x_cells, x_valid = _containing_cells(positions[x_dim])
+        y_cells, y_valid = _containing_cells(positions[y_dim])
+        width, height = self.xy_shape
+        hit = x_valid & y_valid & (x_cells >= 0) & (x_cells < width) & (y_cells >= 0) & (y_cells < height)
+        return {x_dim: np.where(hit, x_cells, -1), y_dim: np.where(hit, y_cells, -1)}
+
     def equals(self, other: Index, *, exclude=None) -> bool:
         if exclude is None:
             exclude = {}

tests/test_get_indexer.py

@@ -0,0 +1,187 @@
+import numpy as np
+import pytest
+import xarray as xr
+from affine import Affine
+from hypothesis import given, settings
+from hypothesis import strategies as st
+
+from rasterix import RasterIndex
+
+
+def global_index(**kwargs) -> RasterIndex:
+    """1-degree global grid, north-up, corners at -180..180 and 90..-90."""
+    return RasterIndex.from_transform(Affine(1, 0, -180, 0, -1, 90), width=360, height=180, **kwargs)
+
+
+def regional_index() -> RasterIndex:
+    """20 x 20 cells of 1 degree, corners at -10..10, north-up."""
+    return RasterIndex.from_transform(Affine(1, 0, -10, 0, -1, 10), width=20, height=20)
+
+
+def test_get_indexer_cell_centers_and_interior():
+    index = regional_index()
+    xc = -10 + np.arange(20) + 0.5
+    yc = 10 - np.arange(20) - 0.5
+    result = index.get_indexer({"x": xc, "y": yc})
+    np.testing.assert_array_equal(result["x"], np.arange(20))
+    np.testing.assert_array_equal(result["y"], np.arange(20))
+    assert result["x"].dtype == np.int64
+    assert result["y"].dtype == np.int64
+
+    result = index.get_indexer({"x": [-9.99, -9.01, 9.99], "y": [9.99, 9.01, -9.99]})
+    np.testing.assert_array_equal(result["x"], [0, 0, 19])
+    np.testing.assert_array_equal(result["y"], [0, 0, 19])
+
+
+def test_get_indexer_outside_extent():
+    index = regional_index()
+    result = index.get_indexer({"x": [-100.0, -10.01, 10.01, 180.5], "y": [100.0, 10.01, -10.01, -100.0]})
+    np.testing.assert_array_equal(result["x"], [-1, -1, -1, -1])
+    np.testing.assert_array_equal(result["y"], [-1, -1, -1, -1])
+
+
+def test_get_indexer_edges_are_half_open_in_position():
+    # increasing x: left edge included, right edge excluded
+    # decreasing y: top edge included, bottom edge excluded
+    index = regional_index()
+    result = index.get_indexer({"x": [-10.0, -9.0, 10.0], "y": [10.0, 9.0, -10.0]})
+    np.testing.assert_array_equal(result["x"], [0, 1, -1])
+    np.testing.assert_array_equal(result["y"], [0, 1, -1])
+
+
+def test_get_indexer_increasing_y():
+    index = RasterIndex.from_transform(Affine(1, 0, -10, 0, 1, -10), width=20, height=20)
+    result = index.get_indexer({"y": [-10.0, -9.5, -9.0, 9.5, 10.0]})
+    np.testing.assert_array_equal(result["y"], [0, 0, 1, 19, -1])
+
+
+def test_get_indexer_single_axis():
+    index = regional_index()
+    result = index.get_indexer({"x": [-9.5, 0.5]})
+    assert set(result) == {"x"}
+    np.testing.assert_array_equal(result["x"], [0, 10])
+    result = index.get_indexer({"y": [9.5]})
+    assert set(result) == {"y"}
+    np.testing.assert_array_equal(result["y"], [0])
+
+
+def test_get_indexer_keeps_shape():
+    index = regional_index()
+    x = np.full((2, 3), 0.5)
+    result = index.get_indexer({"x": x, "y": 0.5})
+    assert result["x"].shape == (2, 3)
+    assert result["y"].shape == ()
+    assert result["y"] == 9
+
+
+def test_get_indexer_unknown_coordinate_raises():
+    with pytest.raises(KeyError, match="foo"):
+        regional_index().get_indexer({"foo": [0.0]})
+
+
+def test_get_indexer_periodic_wraps():
+    index = global_index(x_period=360)
+    result = index.get_indexer({"x": [180.5, 539.5, -180.5]})
+    np.testing.assert_array_equal(result["x"], [0, 359, 359])
+
+    x = np.array([-179.5, -0.25, 0.25, 179.9])
+    expected = index.get_indexer({"x": x})["x"]
+    for k in (-2, -1, 1, 2):
+        np.testing.assert_array_equal(index.get_indexer({"x": x + k * 360})["x"], expected)
+
+
+def test_get_indexer_periodic_subset():
+    # 100 of 360 cells; a wrapped label can still miss the subset
+    index = RasterIndex.from_transform(Affine(1, 0, 100, 0, -1, 90), width=100, height=180, x_period=360)
+    result = index.get_indexer({"x": [100.5, 199.5, 200.5, 460.5, -259.5, 99.5]})
+    np.testing.assert_array_equal(result["x"], [0, 99, -1, 0, 0, -1])
+
+
+def test_get_indexer_no_period_no_wrap():
+    index = regional_index()
+    result = index.get_indexer({"x": [0.5, 0.5 + 360, 0.5 - 360]})
+    np.testing.assert_array_equal(result["x"], [10, -1, -1])
+
+
+@pytest.mark.parametrize("x_period", [None, 360])
+def test_get_indexer_non_finite(x_period):
+    index = global_index(x_period=x_period)
+    result = index.get_indexer({"x": [np.nan, np.inf, -np.inf, 0.5], "y": [np.nan, np.inf, -np.inf, 0.5]})
+    np.testing.assert_array_equal(result["x"], [-1, -1, -1, 180])
+    np.testing.assert_array_equal(result["y"], [-1, -1, -1, 89])
+
+
+@pytest.mark.parametrize("x_period", [None, 360])
+def test_get_indexer_matches_sel_nearest(x_period):
+    index = global_index(x_period=x_period)
+    rng = np.random.default_rng(0)
+    x = rng.uniform(-179.99, 179.99, 1000)
+    y = rng.uniform(-89.99, 89.99, 1000)
+    result = index.get_indexer({"x": x, "y": y})
+    expected = index.sel({"x": x, "y": y}, method="nearest").dim_indexers
+    np.testing.assert_array_equal(result["x"], expected["x"])
+    np.testing.assert_array_equal(result["y"], expected["y"])
+
+
+def rotated_index() -> RasterIndex:
+    return RasterIndex.from_transform(
+        Affine.translation(5, 5) * Affine.rotation(30) * Affine.scale(2, -2),
+        width=10,
+        height=8,
+        x_coord_name="xc",
+        y_coord_name="yc",
+    )
+
+
+def test_get_indexer_rotated():
+    index = rotated_index()
+    i, j = np.meshgrid(np.arange(10), np.arange(8))
+    xc, yc = index.center_transform() * (i.ravel(), j.ravel())
+    result = index.get_indexer({"xc": xc, "yc": yc})
+    np.testing.assert_array_equal(result["x"], i.ravel())
+    np.testing.assert_array_equal(result["y"], j.ravel())
+
+    labels = {"xc": xr.Variable("pt", xc), "yc": xr.Variable("pt", yc)}
+    expected = index.sel(labels, method="nearest").dim_indexers
+    np.testing.assert_array_equal(result["x"], expected["x"].values)
+    np.testing.assert_array_equal(result["y"], expected["y"].values)
+
+
+def test_get_indexer_rotated_miss_on_both_dims():
+    index = rotated_index()
+    # corner positions (-0.25, 2.5) miss in x only, (3.5, 8.5) miss in y only
+    xc, yc = index.transform() * (np.array([-0.25, 3.5, 3.5]), np.array([2.5, 8.5, np.nan]))
+    result = index.get_indexer({"xc": xc, "yc": yc})
+    np.testing.assert_array_equal(result["x"], [-1, -1, -1])
+    np.testing.assert_array_equal(result["y"], [-1, -1, -1])
+
+
+def test_get_indexer_rotated_requires_both_coords():
+    index = rotated_index()
+    with pytest.raises(ValueError, match="both"):
+        index.get_indexer({"xc": [5.0]})
+    with pytest.raises(ValueError, match="shape"):
+        index.get_indexer({"xc": [5.0, 6.0], "yc": [5.0]})
+
+
+@settings(deadline=None)  # first call is slow (cold imports)
+@given(
+    dx=st.floats(1e-3, 1e3).flatmap(lambda v: st.sampled_from([v, -v])),
+    dy=st.floats(1e-3, 1e3).flatmap(lambda v: st.sampled_from([v, -v])),
+    x0=st.floats(-1e6, 1e6),
+    y0=st.floats(-1e6, 1e6),
+    width=st.integers(1, 1000),
+    height=st.integers(1, 1000),
+    data=st.data(),
+)
+def test_get_indexer_cell_center_plus_offset(dx, dy, x0, y0, width, height, data):
+    index = RasterIndex.from_transform(Affine(dx, 0, x0, 0, dy, y0), width=width, height=height)
+    i = data.draw(st.integers(0, width - 1))
+    j = data.draw(st.integers(0, height - 1))
+    # stay clear of the cell edges, where round-off can move the point to a neighbour
+    fx = data.draw(st.floats(-0.49, 0.49))
+    fy = data.draw(st.floats(-0.49, 0.49))
+    x, y = index.center_transform() * (i + fx, j + fy)
+    result = index.get_indexer({"x": x, "y": y})
+    assert result["x"] == i
+    assert result["y"] == j