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/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" 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")