From 7d1fe4b745131944553f9856308bb4bba51b4213 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Tue, 27 Oct 2026 09:40:02 -0600 Subject: [PATCH 1/2] Add RasterIndex.get_indexer for vectorized point lookup Return the int64 position of the cell that contains each point, or -1 for a point outside the extent, NaN or inf. Wraps on periodic axes; uses the full inverse affine for rotated/skewed transforms. Co-Authored-By: Claude Opus 5.5 --- src/rasterix/raster_index.py | 96 ++++++++++++++++++ tests/test_get_indexer.py | 187 +++++++++++++++++++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 tests/test_get_indexer.py diff --git a/src/rasterix/raster_index.py b/src/rasterix/raster_index.py index 917ccc1..fc21ab2 100644 --- a/src/rasterix/raster_index.py +++ b/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,17 @@ def _assert_same_frame(ours: BoundingBox, theirs: BoundingBox, index: RasterInde ) +def _containing_cells(positions: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """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.""" @@ -396,6 +408,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: Any) -> np.ndarray: + """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 +674,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 +1227,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, np.ndarray]: + """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 = {} diff --git a/tests/test_get_indexer.py b/tests/test_get_indexer.py new file mode 100644 index 0000000..3ba8607 --- /dev/null +++ b/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 From b278efe20956c44ca53a0d5699c261ebf73bc1a3 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Tue, 27 Oct 2026 09:45:22 -0600 Subject: [PATCH 2/2] Type get_indexer results and cell positions Co-Authored-By: Claude Opus 5.5 --- src/rasterix/raster_index.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/rasterix/raster_index.py b/src/rasterix/raster_index.py index fc21ab2..8f307cb 100644 --- a/src/rasterix/raster_index.py +++ b/src/rasterix/raster_index.py @@ -210,7 +210,9 @@ def _assert_same_frame(ours: BoundingBox, theirs: BoundingBox, index: RasterInde ) -def _containing_cells(positions: np.ndarray) -> tuple[np.ndarray, np.ndarray]: +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. @@ -323,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: @@ -332,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) @@ -408,7 +410,7 @@ 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: Any) -> np.ndarray: + 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)) @@ -1227,7 +1229,7 @@ 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, np.ndarray]: + 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