From 8fcc2e5c315880d2188c7463eb78ee70050b0c6b Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Sat, 6 Jun 2026 00:59:20 -0700 Subject: [PATCH 1/3] prototype for spatial indexing for fast polygon contains operation --- CMakeLists.txt | 2 + docs/api.rst | 12 +++ docs/api_hidden.rst | 3 + spatial_index_plan.md | 178 +++++++++++++++++++++++++++++++ src/predicates.hpp | 30 ++++++ src/predicates_common.cpp | 70 +++++++++++++ src/spatial_index.cpp | 203 ++++++++++++++++++++++++++++++++++++ src/spherely.cpp | 2 + src/spherely.pyi | 16 +++ tests/test_spatial_index.py | 112 ++++++++++++++++++++ 10 files changed, 628 insertions(+) create mode 100644 spatial_index_plan.md create mode 100644 src/predicates.hpp create mode 100644 src/predicates_common.cpp create mode 100644 src/spatial_index.cpp create mode 100644 tests/test_spatial_index.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 3cd1d11..91df11b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -77,6 +77,8 @@ set(CPP_SOURCES src/geography.cpp src/io.cpp src/predicates.cpp + src/predicates_common.cpp + src/spatial_index.cpp src/spherely.cpp) if(${s2geography_VERSION} VERSION_GREATER_EQUAL "0.2.0") diff --git a/docs/api.rst b/docs/api.rst index 1c906da..b2e4e93 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -99,6 +99,18 @@ between two geographies. covers covered_by +.. _api_spatial_index: + +Spatial indexing +---------------- + +Index a collection of geographies for fast spatial queries. + +.. autosummary:: + :toctree: _api_generated/ + + SpatialIndex + .. _api_overlays: Overlays (boolean operations) diff --git a/docs/api_hidden.rst b/docs/api_hidden.rst index 7862012..9241a7d 100644 --- a/docs/api_hidden.rst +++ b/docs/api_hidden.rst @@ -6,6 +6,9 @@ :toctree: _api_generated/ Geography + SpatialIndex + SpatialIndex.geometries + SpatialIndex.query Projection Projection.lnglat Projection.pseudo_mercator diff --git a/spatial_index_plan.md b/spatial_index_plan.md new file mode 100644 index 0000000..6d89622 --- /dev/null +++ b/spatial_index_plan.md @@ -0,0 +1,178 @@ +# Plan: `SpatialIndex` for spherely (issue #72) + +## Context + +spherely (S2-backed analog of shapely) currently has no way to spatially index a +large collection of `Geography` objects. Every pairwise predicate (`intersects`, +`contains`, …) is computed brute-force via `py::vectorize`, which is O(N·M) for a +spatial join. Issue #72 asks for an STRtree-like index so candidate pairs can be +pre-filtered cheaply, as shapely's `STRtree` does. + +**Feasibility: confirmed — all primitives already exist in the pinned +`s2geography 0.2.0` (conda-forge), so no upstream changes are needed.** Key facts +from exploration: + +- `s2geography::GeographyIndex` (header-only, `s2geography/index.h`) wraps a + `MutableS2ShapeIndex` and is purpose-built as "a GEOSSTRTree index" over a + *vector* of geographies. Public API (verified at the `0.2.0` tag): + - `void Add(const Geography& geog, int value)` — `value` is the caller's array index. + - `int value(int shape_id) const` — maps an internal shape id back to that array index. + - `const MutableS2ShapeIndex& ShapeIndex() const`. + - nested `Iterator` with `void Query(const std::vector& covering, std::unordered_set* out)` + which fills `out` with matching array indices (cell-overlap candidates). +- `s2geography::Geography::Region() -> std::unique_ptr` and + `GetCellUnionBound(std::vector*)` exist, so we can compute a query + covering via `S2RegionCoverer`. +- The R `s2` package implements exactly this pattern (`s2-matrix.cpp`): build + `GeographyIndex`, then per query feature `coverer.GetCovering(*Geog().Region(), &cells)` + → `iterator.Query(cells, &indices)`. +- spherely's existing predicate functions (`s2geog::s2_intersects`, `s2_contains`, + etc., used in `src/predicates.cpp`) operate on `ShapeIndexGeography` and can be + reused directly for the refinement step, via the lazy + `Geography::geog_index()` accessor (`src/geography.hpp:82`). + +**Decided scope (per user):** ship the index + `query()` with optional predicate +refinement. Distance/nearest (`S2ClosestEdgeQuery`) and serialization are +explicitly deferred to a follow-up PR. Class name: **`SpatialIndex`**. + +## Target Python API + +```python +tree = spherely.SpatialIndex(geographies) # iterable / object-ndarray of Geography +len(tree) # number of indexed geographies +tree.geometries # object-ndarray of the inputs (same order) + +# scalar query -> sorted np.intp array of tree indices +idx = tree.query(geom) # coarse: cells overlap (candidate set) +idx = tree.query(geom, predicate="intersects") # refined + +# array query -> shape (2, K): row 0 = input index, row 1 = tree index (shapely-compatible) +pairs = tree.query(np.array([...]), predicate="intersects") +``` + +- `predicate=None` → coarse cell-overlap candidates (superset of true intersections). +- Supported `predicate` values (all are subsets of the intersecting-candidate set, + so refinement is correct): `"intersects"`, `"within"`, `"contains"`, `"covers"`, + `"covered_by"`, `"touches"`, `"equals"`. Semantics match + shapely: a tree geom `t` matches when `predicate(query_geom, t)` is True. +- `"disjoint"` is **rejected** in v1 (it is *not* a subset of the overlap candidate + set; computing it via the index needs the complement — defer). Raise `ValueError`. + +## Implementation + +### 1. New C++ source: `src/spatial_index.cpp` +Mirror the structure of `src/predicates.cpp` (a `void init_spatial_index(py::module&)` +entry point, `namespace s2geog = s2geography`, reuse `PyObjectGeography`). + +Includes: ``, ``, ``, +``, `"geography.hpp"`, `"pybind11.hpp"`. + +Define `class SpatialIndex`: +- **Members:** `std::unique_ptr m_index;` and + `py::array m_geographies;` (an object-dtype ndarray held to (a) keep the indexed + `Geography` C++ objects alive — `GeographyIndex` only borrows their `S2Shape`s — + and (b) back the `geometries` property). +- **Constructor `SpatialIndex(py::array_t geographies)`:** + store the array; `m_index = std::make_unique();` then loop + `i` over the flat array, `auto* g = geographies.data()[i]->as_geog_ptr();` + `m_index->Add(g->geog(), static_cast(i));`. Empty geographies contribute no + shapes (documented: they are never returned). Validate 1-D input; reject `None` + entries with a clear `TypeError` (reuse `PyObjectGeography::check_type`). +- **`__len__`** → number of input geographies. +- **`geometries` property** → return `m_geographies`. +- **Private helper `query_one(Geography* query_geog, predicate) -> std::vector`:** + 1. `auto region = query_geog->geog().Region();` + 2. `S2RegionCoverer coverer;` (default options for v1) → + `std::vector covering; coverer.GetCovering(*region, &covering);` + 3. `std::unordered_set hits; GeographyIndex::Iterator it(m_index.get); it.Query(covering, &hits);` + 4. If `predicate != None`: keep only `c` where + `predfn(query_geog->geog_index(), candidate->geog_index(), options)` is true, + where `candidate = m_geographies.data()[c]->as_geog_ptr()`. Map predicate name → + callable over two `ShapeIndexGeography` (see step 2 reuse note below). + 5. copy set → `std::vector`, `std::sort`. +- **`query(scalar Geography)`** → `py::array_t` (sorted). +- **`query(array of Geography)`** → build two `std::vector` (input idx, tree idx) + and return a `(2, K)` `py::array_t` (shapely's `query` bulk layout). + +### 2. Reuse predicate functions (small shared helper) +The predicate lambdas currently live *inside* `init_predicates` in +`src/predicates.cpp` and aren't reachable. Add a tiny shared mapping so both files +use one source of truth. Lightest option: + +- Add `src/predicates.hpp` declaring a factory, e.g. + `using PredicateFunc = std::function;` + and `PredicateFunc get_predicate(const std::string& name);` returning a closure + over the right `s2geog::s2_intersects`/`s2_contains` call + the appropriate + `S2BooleanOperation::Options` (e.g. closed model for `covers`/`covered_by`, + the dual-options trick for `touches`, argument-swap for `within`/`covered_by`). +- Implement it in `predicates.cpp` (or a new `predicates_common.cpp`) and have + `init_predicates` register the public vectorized predicates by reusing the same + closures, avoiding duplication. `SpatialIndex::query_one` calls `get_predicate`. +- Unknown / unsupported name (incl. `"disjoint"`) → throw `py::value_error`. + +### 3. Wire-up +- `src/spherely.cpp`: declare `void init_spatial_index(py::module&);` and call it in + `PYBIND11_MODULE` (next to `init_predicates`). +- `CMakeLists.txt`: add `src/spatial_index.cpp` (and `predicates_common.cpp` if + split out) to `CPP_SOURCES`. No new link deps — `GeographyIndex` is header-only and + `s2geography`/`s2::s2` are already linked. Gate on s2geography ≥ 0.2.0 if needed + (already the project floor: `s2geography >=0.2.0,<0.3`). + +### 4. Type stubs — `src/spherely.pyi` +Add (near the predicates section): +```python +class SpatialIndex: + def __init__(self, geographies: Iterable[Geography]) -> None: ... + def __len__(self) -> int: ... + @property + def geometries(self) -> T_NDArray_Geography: ... + @overload + def query(self, geography: Geography, predicate: str | None = None) -> npt.NDArray[np.intp]: ... + @overload + def query(self, geography: Iterable[Geography], predicate: str | None = None) -> npt.NDArray[np.intp]: ... +``` + +### 5. Docs +- `docs/api.rst`: new "Spatial indexing" section listing `SpatialIndex`. +- `docs/api_hidden.rst`: add the class' members so autosummary picks them up + (match the existing hidden-entry pattern). + +### 6. Tests — `tests/test_spatial_index.py` +Mirror `tests/test_predicates.py` conventions (typed, numpy assertions). Cover: +- build + `len()` + `geometries` round-trips and dtype/order. +- coarse `query(scalar)` returns candidate indices (sorted, np.intp). +- `query(scalar, predicate="intersects")` refines vs coarse (strict subset case). +- each supported predicate gives shapely-equivalent results on a small hand-checked + fixture (points/lines/polygons spanning a couple of cells). +- array query returns `(2, K)` `(input_idx, tree_idx)` layout. +- empty geography in the tree is never returned; empty query returns empty. +- `predicate="disjoint"` and unknown predicate raise; `None` entry raises `TypeError`. + +## Verification + +```bash +pixi run -e test compile # build the extension (configure+compile tasks) +pixi run -e test tests # run pytest, incl. new test_spatial_index.py +pixi run -e lint mypy # type-check stubs against tests +``` +Manual smoke test in `pixi shell -e test`: +```python +import numpy as np, spherely +geoms = spherely.points([0,1,2,50], [0,1,2,50]) +t = spherely.SpatialIndex(geoms) +assert len(t) == 4 +q = spherely.create_polygon([(-1,-1),(3,-1),(3,3),(-1,3)]) +print(t.query(q)) # -> [0 1 2] +print(t.query(q, predicate="contains")) # refined subset +``` +Cross-check a handful of results against shapely's `STRtree` on the planar +equivalents to confirm the join semantics line up. + +## Deferred to follow-up (out of scope here) +- `query_nearest` / `nearest` / distance filtering via `S2ClosestEdgeQuery` + (+ `S2ClosestEdgeQuery::ShapeIndexTarget`, header ``). +- Index serialization / pickling (S2's `EncodedS2ShapeIndex` enables mmap'd, + shippable indexes — noted as high value in issue #72). +- Configurable coverer (`max_cells`) and a dedicated point-only fast path + (`S2PointIndex`) for large point sets. +``` diff --git a/src/predicates.hpp b/src/predicates.hpp new file mode 100644 index 0000000..81bc480 --- /dev/null +++ b/src/predicates.hpp @@ -0,0 +1,30 @@ +#ifndef SPHERELY_PREDICATES_H_ +#define SPHERELY_PREDICATES_H_ + +#include + +#include +#include + +namespace spherely { + +namespace s2geog = s2geography; + +// A binary spatial predicate operating on two indexed geographies, i.e. +// ``predicate(a, b)`` where ``a`` and ``b`` are ``ShapeIndexGeography``. +using PredicateFunc = + std::function; + +// Returns the predicate closure matching ``name`` (one of "intersects", +// "within", "contains", "covers", "covered_by", "touches", "equals"). +// +// Throws ``pybind11::value_error`` for unknown names and for "disjoint" (which +// cannot be evaluated as a refinement of a spatial-index candidate set). +// +// The closures mirror the semantics of the vectorized predicates registered in +// ``predicates.cpp``; that file remains the reference implementation. +PredicateFunc get_predicate(const std::string& name); + +} // namespace spherely + +#endif // SPHERELY_PREDICATES_H_ diff --git a/src/predicates_common.cpp b/src/predicates_common.cpp new file mode 100644 index 0000000..f457577 --- /dev/null +++ b/src/predicates_common.cpp @@ -0,0 +1,70 @@ +#include +#include +#include + +#include "predicates.hpp" + +namespace py = pybind11; +namespace s2geog = s2geography; + +namespace spherely { + +namespace { + +S2BooleanOperation::Options closed_options() { + S2BooleanOperation::Options options; + options.set_polyline_model(S2BooleanOperation::PolylineModel::CLOSED); + options.set_polygon_model(S2BooleanOperation::PolygonModel::CLOSED); + return options; +} + +S2BooleanOperation::Options open_options() { + S2BooleanOperation::Options options; + options.set_polyline_model(S2BooleanOperation::PolylineModel::OPEN); + options.set_polygon_model(S2BooleanOperation::PolygonModel::OPEN); + return options; +} + +} // namespace + +PredicateFunc get_predicate(const std::string& name) { + using Index = s2geog::ShapeIndexGeography; + + if (name == "intersects") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_intersects(a, b, options); + }; + } else if (name == "within") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_contains(b, a, options); + }; + } else if (name == "contains") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_contains(a, b, options); + }; + } else if (name == "equals") { + return [options = S2BooleanOperation::Options()](const Index& a, const Index& b) { + return s2geog::s2_equals(a, b, options); + }; + } else if (name == "covers") { + return [options = closed_options()](const Index& a, const Index& b) { + return s2geog::s2_contains(a, b, options); + }; + } else if (name == "covered_by") { + return [options = closed_options()](const Index& a, const Index& b) { + return s2geog::s2_contains(b, a, options); + }; + } else if (name == "touches") { + return [closed = closed_options(), open = open_options()](const Index& a, const Index& b) { + return s2geog::s2_intersects(a, b, closed) && !s2geog::s2_intersects(a, b, open); + }; + } else if (name == "disjoint") { + throw py::value_error( + "the 'disjoint' predicate is not supported by SpatialIndex.query: it is not a " + "refinement of the spatial-index candidate set"); + } else { + throw py::value_error("invalid predicate: '" + name + "'"); + } +} + +} // namespace spherely diff --git a/src/spatial_index.cpp b/src/spatial_index.cpp new file mode 100644 index 0000000..ef75fe2 --- /dev/null +++ b/src/spatial_index.cpp @@ -0,0 +1,203 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "geography.hpp" +#include "predicates.hpp" +#include "pybind11.hpp" + +namespace py = pybind11; +namespace s2geog = s2geography; +using namespace spherely; + +/* +** A spatial index over a collection of Geography objects, backed by +** s2geography::GeographyIndex (a MutableS2ShapeIndex). Provides fast candidate +** lookup similar to shapely's STRtree. +** +** The input geographies are held as a numpy object-dtype array, which both +** keeps the underlying C++ Geography objects (whose S2Shapes are borrowed by +** the index) alive and backs the ``geometries`` property. +*/ +class SpatialIndex { +public: + SpatialIndex(py::object geographies) { + auto arr = py::array_t::ensure(geographies); + if (!arr) { + throw py::type_error("geographies must be an array-like of Geography objects"); + } + if (arr.ndim() != 1) { + throw py::type_error("geographies must be a 1-dimensional array"); + } + + m_geographies = arr; + m_index = std::make_unique(); + + auto n = arr.size(); + auto* data = static_cast(arr.request().ptr); + for (py::ssize_t i = 0; i < n; i++) { + auto* geog_ptr = data[i].as_geog_ptr(); + m_index->Add(geog_ptr->geog(), static_cast(i)); + } + } + + std::size_t size() const { + return static_cast(m_geographies.size()); + } + + py::array geometries() const { + return m_geographies; + } + + // Dispatch between scalar (single Geography -> 1-d index array) and + // array-like (-> (2, K) array of (input index, tree index) pairs) queries. + py::object query(py::object geography, std::optional predicate) const { + PredicateFunc pred; + bool has_pred = predicate.has_value(); + if (has_pred) { + pred = get_predicate(*predicate); + } + + // PyObjectGeography is layout-compatible with py::object (see + // PyObjectGeography::from_geog), so a reference cast is safe here. + auto& maybe_geog = static_cast(geography); + if (maybe_geog.is_geog_ptr()) { + auto results = query_one(maybe_geog.as_geog_ptr(), has_pred ? &pred : nullptr); + return to_index_array(results); + } + + auto arr = py::array_t::ensure(geography); + if (!arr || arr.ndim() != 1) { + throw py::type_error( + "query geography must be a Geography or a 1-dimensional array of Geography"); + } + + auto n = arr.size(); + auto* data = static_cast(arr.request().ptr); + std::vector input_idx; + std::vector tree_idx; + for (py::ssize_t i = 0; i < n; i++) { + auto results = query_one(data[i].as_geog_ptr(), has_pred ? &pred : nullptr); + for (int t : results) { + input_idx.push_back(i); + tree_idx.push_back(static_cast(t)); + } + } + + auto k = input_idx.size(); + py::array_t out({static_cast(2), static_cast(k)}); + auto out_data = out.mutable_unchecked<2>(); + for (std::size_t j = 0; j < k; j++) { + out_data(0, static_cast(j)) = input_idx[j]; + out_data(1, static_cast(j)) = tree_idx[j]; + } + return out; + } + +private: + std::unique_ptr m_index; + py::array_t m_geographies; + + // Return the sorted tree indices whose cells overlap the query geography, + // optionally refined by ``pred`` (predicate(query, candidate)). + std::vector query_one(Geography* query_geog, const PredicateFunc* pred) const { + std::unordered_set hits; + + auto region = query_geog->geog().Region(); + S2RegionCoverer coverer; + std::vector covering; + coverer.GetCovering(*region, &covering); + + s2geog::GeographyIndex::Iterator iter(m_index.get()); + iter.Query(covering, &hits); + + std::vector results; + if (pred == nullptr) { + results.assign(hits.begin(), hits.end()); + } else { + const auto& query_index = query_geog->geog_index(); + auto* data = static_cast(m_geographies.request().ptr); + for (int candidate : hits) { + auto* cand_geog = data[candidate].as_geog_ptr(); + if ((*pred)(query_index, cand_geog->geog_index())) { + results.push_back(candidate); + } + } + } + + std::sort(results.begin(), results.end()); + return results; + } + + static py::array_t to_index_array(const std::vector& indices) { + auto n = indices.size(); + py::array_t out(static_cast(n)); + auto out_data = out.mutable_unchecked<1>(); + for (std::size_t i = 0; i < n; i++) { + out_data(static_cast(i)) = static_cast(indices[i]); + } + return out; + } +}; + +void init_spatial_index(py::module& m) { + py::class_(m, "SpatialIndex", R"pbdoc( + A spatial index for a collection of geographies. + + The index allows fast retrieval of the geographies that may interact + with a given geography, e.g., for accelerating spatial joins. It is + built on top of s2geometry's shape index and is conceptually similar to + shapely's ``STRtree``. + + Parameters + ---------- + geographies : array_like + An array (or other sequence) of :py:class:`Geography` objects. Empty + geographies are indexed but never returned by queries. + + )pbdoc") + .def(py::init(), py::arg("geographies")) + .def("__len__", &SpatialIndex::size) + .def_property_readonly("geometries", + &SpatialIndex::geometries, + "The array of geographies in the index (in input order).") + .def("query", + &SpatialIndex::query, + py::arg("geography"), + py::arg("predicate") = py::none(), + R"pbdoc(query(geography, predicate=None) + + Return the integer indices of all geographies in the index whose cells + overlap the given geography (a coarse candidate set), optionally refined + by a spatial predicate. + + Parameters + ---------- + geography : :py:class:`Geography` or array_like + The geography or geographies to query. + predicate : str, optional + If provided, only return candidates for which + ``predicate(geography, indexed_geography)`` is True. One of + "intersects", "within", "contains", "covers", "covered_by", + "touches" or "equals". + + Returns + ------- + ndarray + If ``geography`` is a scalar, a 1-d array of (sorted) indices into + the index. If ``geography`` is an array, a ``(2, N)`` array where the + first row holds the input geography indices and the second row the + matching index (tree) indices. + + )pbdoc"); +} diff --git a/src/spherely.cpp b/src/spherely.cpp index 1f4dce7..42571db 100644 --- a/src/spherely.cpp +++ b/src/spherely.cpp @@ -8,6 +8,7 @@ namespace py = pybind11; void init_geography(py::module&); void init_creation(py::module&); void init_predicates(py::module&); +void init_spatial_index(py::module&); void init_boolean_operations(py::module&); void init_accessors(py::module&); void init_io(py::module&); @@ -29,6 +30,7 @@ PYBIND11_MODULE(spherely, m) { init_geography(m); init_creation(m); init_predicates(m); + init_spatial_index(m); init_boolean_operations(m); init_accessors(m); init_io(m); diff --git a/src/spherely.pyi b/src/spherely.pyi index e9f03d0..146a43c 100644 --- a/src/spherely.pyi +++ b/src/spherely.pyi @@ -232,6 +232,22 @@ touches: _VFunc_Nin2_Nout1[Literal["touches"], bool, bool] covers: _VFunc_Nin2_Nout1[Literal["covers"], bool, bool] covered_by: _VFunc_Nin2_Nout1[Literal["covered_by"], bool, bool] +# spatial index + +class SpatialIndex: + def __init__(self, geographies: Iterable[Geography]) -> None: ... + def __len__(self) -> int: ... + @property + def geometries(self) -> T_NDArray_Geography: ... + @overload + def query( + self, geography: Geography, predicate: str | None = None + ) -> npt.NDArray[np.intp]: ... + @overload + def query( + self, geography: Iterable[Geography], predicate: str | None = None + ) -> npt.NDArray[np.intp]: ... + # boolean operations union: _VFunc_Nin2_Nout1[Literal["union"], Geography, Geography] diff --git a/tests/test_spatial_index.py b/tests/test_spatial_index.py new file mode 100644 index 0000000..94cb484 --- /dev/null +++ b/tests/test_spatial_index.py @@ -0,0 +1,112 @@ +from typing import Any + +import numpy as np +import numpy.typing as npt +import pytest + +import spherely + + +@pytest.fixture +def geographies() -> npt.NDArray[Any]: + # three points near (0, 0)..(2, 2) and one far-away point + return np.array( + [ + spherely.create_point(0, 0), + spherely.create_point(1, 1), + spherely.create_point(2, 2), + spherely.create_point(50, 50), + ] + ) + + +def test_spatial_index_len(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + assert len(tree) == 4 + + +def test_spatial_index_geometries(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + geoms = tree.geometries + assert isinstance(geoms, np.ndarray) + assert geoms.shape == (4,) + assert all(spherely.equals(a, b) for a, b in zip(geoms, geographies)) + + +def test_spatial_index_from_list() -> None: + tree = spherely.SpatialIndex( + [spherely.create_point(0, 0), spherely.create_point(1, 1)] + ) + assert len(tree) == 2 + + +def test_query_scalar(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + result = tree.query(poly) + assert result.dtype == np.intp + # the three nearby points are candidates; the far point is not + np.testing.assert_array_equal(result, [0, 1, 2]) + + +def test_query_predicate_refines(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + poly = spherely.create_polygon([(-1, -1), (3, -1), (3, 3), (-1, 3), (-1, -1)]) + + coarse = tree.query(poly) + refined = tree.query(poly, predicate="contains") + # refinement is a subset of the coarse candidate set + assert {int(x) for x in refined}.issubset({int(x) for x in coarse}) + # all three nearby points are actually inside the polygon + np.testing.assert_array_equal(refined, [0, 1, 2]) + + +def test_query_predicate_intersects(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + result = tree.query(point, predicate="intersects") + np.testing.assert_array_equal(result, [1]) + + +def test_query_array(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + queries = np.array( + [ + spherely.create_point(1, 1), + spherely.create_point(50, 50), + ] + ) + result = tree.query(queries, predicate="intersects") + assert result.shape[0] == 2 + # (input_index, tree_index) pairs + pairs = {(int(a), int(b)) for a, b in zip(result[0], result[1])} + assert pairs == {(0, 1), (1, 3)} + + +def test_query_empty_geography_never_returned() -> None: + geoms = np.array( + [ + spherely.create_point(0, 0), + spherely.create_polygon(None), # empty + ] + ) + tree = spherely.SpatialIndex(geoms) + assert len(tree) == 2 + poly = spherely.create_polygon([(-1, -1), (1, -1), (1, 1), (-1, 1), (-1, -1)]) + result = tree.query(poly) + assert 1 not in {int(x) for x in result} + + +def test_query_disjoint_rejected(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + with pytest.raises(ValueError, match="disjoint"): + tree.query(point, predicate="disjoint") + + +def test_query_invalid_predicate(geographies: npt.NDArray[Any]) -> None: + tree = spherely.SpatialIndex(geographies) + point = spherely.create_point(1, 1) + with pytest.raises(ValueError, match="invalid predicate"): + tree.query(point, predicate="not_a_predicate") From 48949482ea74410d9201daa1f9e993a7a02a8803 Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Sat, 6 Jun 2026 13:13:20 -0700 Subject: [PATCH 2/3] culling planning doc --- spatial_index_plan.md | 178 ------------------------------------------ 1 file changed, 178 deletions(-) delete mode 100644 spatial_index_plan.md diff --git a/spatial_index_plan.md b/spatial_index_plan.md deleted file mode 100644 index 6d89622..0000000 --- a/spatial_index_plan.md +++ /dev/null @@ -1,178 +0,0 @@ -# Plan: `SpatialIndex` for spherely (issue #72) - -## Context - -spherely (S2-backed analog of shapely) currently has no way to spatially index a -large collection of `Geography` objects. Every pairwise predicate (`intersects`, -`contains`, …) is computed brute-force via `py::vectorize`, which is O(N·M) for a -spatial join. Issue #72 asks for an STRtree-like index so candidate pairs can be -pre-filtered cheaply, as shapely's `STRtree` does. - -**Feasibility: confirmed — all primitives already exist in the pinned -`s2geography 0.2.0` (conda-forge), so no upstream changes are needed.** Key facts -from exploration: - -- `s2geography::GeographyIndex` (header-only, `s2geography/index.h`) wraps a - `MutableS2ShapeIndex` and is purpose-built as "a GEOSSTRTree index" over a - *vector* of geographies. Public API (verified at the `0.2.0` tag): - - `void Add(const Geography& geog, int value)` — `value` is the caller's array index. - - `int value(int shape_id) const` — maps an internal shape id back to that array index. - - `const MutableS2ShapeIndex& ShapeIndex() const`. - - nested `Iterator` with `void Query(const std::vector& covering, std::unordered_set* out)` - which fills `out` with matching array indices (cell-overlap candidates). -- `s2geography::Geography::Region() -> std::unique_ptr` and - `GetCellUnionBound(std::vector*)` exist, so we can compute a query - covering via `S2RegionCoverer`. -- The R `s2` package implements exactly this pattern (`s2-matrix.cpp`): build - `GeographyIndex`, then per query feature `coverer.GetCovering(*Geog().Region(), &cells)` - → `iterator.Query(cells, &indices)`. -- spherely's existing predicate functions (`s2geog::s2_intersects`, `s2_contains`, - etc., used in `src/predicates.cpp`) operate on `ShapeIndexGeography` and can be - reused directly for the refinement step, via the lazy - `Geography::geog_index()` accessor (`src/geography.hpp:82`). - -**Decided scope (per user):** ship the index + `query()` with optional predicate -refinement. Distance/nearest (`S2ClosestEdgeQuery`) and serialization are -explicitly deferred to a follow-up PR. Class name: **`SpatialIndex`**. - -## Target Python API - -```python -tree = spherely.SpatialIndex(geographies) # iterable / object-ndarray of Geography -len(tree) # number of indexed geographies -tree.geometries # object-ndarray of the inputs (same order) - -# scalar query -> sorted np.intp array of tree indices -idx = tree.query(geom) # coarse: cells overlap (candidate set) -idx = tree.query(geom, predicate="intersects") # refined - -# array query -> shape (2, K): row 0 = input index, row 1 = tree index (shapely-compatible) -pairs = tree.query(np.array([...]), predicate="intersects") -``` - -- `predicate=None` → coarse cell-overlap candidates (superset of true intersections). -- Supported `predicate` values (all are subsets of the intersecting-candidate set, - so refinement is correct): `"intersects"`, `"within"`, `"contains"`, `"covers"`, - `"covered_by"`, `"touches"`, `"equals"`. Semantics match - shapely: a tree geom `t` matches when `predicate(query_geom, t)` is True. -- `"disjoint"` is **rejected** in v1 (it is *not* a subset of the overlap candidate - set; computing it via the index needs the complement — defer). Raise `ValueError`. - -## Implementation - -### 1. New C++ source: `src/spatial_index.cpp` -Mirror the structure of `src/predicates.cpp` (a `void init_spatial_index(py::module&)` -entry point, `namespace s2geog = s2geography`, reuse `PyObjectGeography`). - -Includes: ``, ``, ``, -``, `"geography.hpp"`, `"pybind11.hpp"`. - -Define `class SpatialIndex`: -- **Members:** `std::unique_ptr m_index;` and - `py::array m_geographies;` (an object-dtype ndarray held to (a) keep the indexed - `Geography` C++ objects alive — `GeographyIndex` only borrows their `S2Shape`s — - and (b) back the `geometries` property). -- **Constructor `SpatialIndex(py::array_t geographies)`:** - store the array; `m_index = std::make_unique();` then loop - `i` over the flat array, `auto* g = geographies.data()[i]->as_geog_ptr();` - `m_index->Add(g->geog(), static_cast(i));`. Empty geographies contribute no - shapes (documented: they are never returned). Validate 1-D input; reject `None` - entries with a clear `TypeError` (reuse `PyObjectGeography::check_type`). -- **`__len__`** → number of input geographies. -- **`geometries` property** → return `m_geographies`. -- **Private helper `query_one(Geography* query_geog, predicate) -> std::vector`:** - 1. `auto region = query_geog->geog().Region();` - 2. `S2RegionCoverer coverer;` (default options for v1) → - `std::vector covering; coverer.GetCovering(*region, &covering);` - 3. `std::unordered_set hits; GeographyIndex::Iterator it(m_index.get); it.Query(covering, &hits);` - 4. If `predicate != None`: keep only `c` where - `predfn(query_geog->geog_index(), candidate->geog_index(), options)` is true, - where `candidate = m_geographies.data()[c]->as_geog_ptr()`. Map predicate name → - callable over two `ShapeIndexGeography` (see step 2 reuse note below). - 5. copy set → `std::vector`, `std::sort`. -- **`query(scalar Geography)`** → `py::array_t` (sorted). -- **`query(array of Geography)`** → build two `std::vector` (input idx, tree idx) - and return a `(2, K)` `py::array_t` (shapely's `query` bulk layout). - -### 2. Reuse predicate functions (small shared helper) -The predicate lambdas currently live *inside* `init_predicates` in -`src/predicates.cpp` and aren't reachable. Add a tiny shared mapping so both files -use one source of truth. Lightest option: - -- Add `src/predicates.hpp` declaring a factory, e.g. - `using PredicateFunc = std::function;` - and `PredicateFunc get_predicate(const std::string& name);` returning a closure - over the right `s2geog::s2_intersects`/`s2_contains` call + the appropriate - `S2BooleanOperation::Options` (e.g. closed model for `covers`/`covered_by`, - the dual-options trick for `touches`, argument-swap for `within`/`covered_by`). -- Implement it in `predicates.cpp` (or a new `predicates_common.cpp`) and have - `init_predicates` register the public vectorized predicates by reusing the same - closures, avoiding duplication. `SpatialIndex::query_one` calls `get_predicate`. -- Unknown / unsupported name (incl. `"disjoint"`) → throw `py::value_error`. - -### 3. Wire-up -- `src/spherely.cpp`: declare `void init_spatial_index(py::module&);` and call it in - `PYBIND11_MODULE` (next to `init_predicates`). -- `CMakeLists.txt`: add `src/spatial_index.cpp` (and `predicates_common.cpp` if - split out) to `CPP_SOURCES`. No new link deps — `GeographyIndex` is header-only and - `s2geography`/`s2::s2` are already linked. Gate on s2geography ≥ 0.2.0 if needed - (already the project floor: `s2geography >=0.2.0,<0.3`). - -### 4. Type stubs — `src/spherely.pyi` -Add (near the predicates section): -```python -class SpatialIndex: - def __init__(self, geographies: Iterable[Geography]) -> None: ... - def __len__(self) -> int: ... - @property - def geometries(self) -> T_NDArray_Geography: ... - @overload - def query(self, geography: Geography, predicate: str | None = None) -> npt.NDArray[np.intp]: ... - @overload - def query(self, geography: Iterable[Geography], predicate: str | None = None) -> npt.NDArray[np.intp]: ... -``` - -### 5. Docs -- `docs/api.rst`: new "Spatial indexing" section listing `SpatialIndex`. -- `docs/api_hidden.rst`: add the class' members so autosummary picks them up - (match the existing hidden-entry pattern). - -### 6. Tests — `tests/test_spatial_index.py` -Mirror `tests/test_predicates.py` conventions (typed, numpy assertions). Cover: -- build + `len()` + `geometries` round-trips and dtype/order. -- coarse `query(scalar)` returns candidate indices (sorted, np.intp). -- `query(scalar, predicate="intersects")` refines vs coarse (strict subset case). -- each supported predicate gives shapely-equivalent results on a small hand-checked - fixture (points/lines/polygons spanning a couple of cells). -- array query returns `(2, K)` `(input_idx, tree_idx)` layout. -- empty geography in the tree is never returned; empty query returns empty. -- `predicate="disjoint"` and unknown predicate raise; `None` entry raises `TypeError`. - -## Verification - -```bash -pixi run -e test compile # build the extension (configure+compile tasks) -pixi run -e test tests # run pytest, incl. new test_spatial_index.py -pixi run -e lint mypy # type-check stubs against tests -``` -Manual smoke test in `pixi shell -e test`: -```python -import numpy as np, spherely -geoms = spherely.points([0,1,2,50], [0,1,2,50]) -t = spherely.SpatialIndex(geoms) -assert len(t) == 4 -q = spherely.create_polygon([(-1,-1),(3,-1),(3,3),(-1,3)]) -print(t.query(q)) # -> [0 1 2] -print(t.query(q, predicate="contains")) # refined subset -``` -Cross-check a handful of results against shapely's `STRtree` on the planar -equivalents to confirm the join semantics line up. - -## Deferred to follow-up (out of scope here) -- `query_nearest` / `nearest` / distance filtering via `S2ClosestEdgeQuery` - (+ `S2ClosestEdgeQuery::ShapeIndexTarget`, header ``). -- Index serialization / pickling (S2's `EncodedS2ShapeIndex` enables mmap'd, - shippable indexes — noted as high value in issue #72). -- Configurable coverer (`max_cells`) and a dedicated point-only fast path - (`S2PointIndex`) for large point sets. -``` From 194c6bda8dc750dcb08bb8806ac7b068edf3284c Mon Sep 17 00:00:00 2001 From: Shane Date: Tue, 9 Jun 2026 18:15:52 -0700 Subject: [PATCH 3/3] Bump version to 0.1.1+spatialindex for fork wheel builds --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 1606a70..164a31b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build" [project] name = "spherely" -version = "0.1.1" +version = "0.1.1+spatialindex" description = "Manipulation and analysis of geometric objects on the sphere" keywords = ["gis", "geometry", "s2geometry", "shapely"] readme = "README.md"