Use firedrake-rtree for spatial indexing

firedrake/cython/rtree.pyx

@@ -0,0 +1,83 @@
+# cython: language_level=3
+
+cimport numpy as np
+import numpy as np
+import ctypes
+import cython
+from libc.stddef cimport size_t
+from libc.stdint cimport uintptr_t
+
+include "rtreeinc.pxi"
+
+cdef class RTree(object):
+    """Python class for holding a spatial index."""
+
+    cdef RTreeH* tree
+
+    def __cinit__(self, uintptr_t tree_handle):
+        self.tree = <RTreeH*>0
+        if tree_handle == 0:
+            raise RuntimeError("invalid tree handle")
+        self.tree = <RTreeH*>tree_handle
+
+    def __dealloc__(self):
+        if self.tree != <RTreeH*>0:
+            rtree_free(self.tree)
+
+    @property
+    def ctypes(self):
+        """Returns a ctypes pointer to the native spatial index."""
+        return ctypes.c_void_p(<uintptr_t> self.tree)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min,
+                    np.ndarray[np.float64_t, ndim=2, mode="c"] coords_max,
+                    np.ndarray[np.npy_uintp, ndim=1, mode="c"] ids = None):
+    """Builds rtree from two arrays of shape (n, dim) containing the coordinates
+    of the lower and upper corners of n axis-aligned bounding boxes, and an
+    optional array of shape (n,) containing integer ids for each box.
+
+    Parameters
+    ----------
+    coords_min : (n, dim) array
+        The lower corner coordinates of the bounding boxes.
+    regions_hi : (n, dim) array
+        The upper corner coordinates of the bounding boxes.
+    ids : (n,) array, optional
+        Integer ids for each box. If not provided, defaults to 0, 1, ..., n-1.
+
+    Returns
+    -------
+    RTree
+        An RTree object containing the Rtree.
+    """    
+    cdef:
+        RTreeH* rtree
+        size_t n
+        size_t dim
+        RTreeError err
+
+    if coords_min.shape[0] != coords_max.shape[0] or coords_min.shape[1] != coords_max.shape[1]:
+        raise ValueError("coords_min and coords_max must have the same shape")
+
+    n = <size_t>coords_min.shape[0]
+    dim = <size_t>coords_min.shape[1]
+    if ids is None:
+        ids = np.arange(n, dtype=np.uintp)
+    elif ids.shape[0] != n:
+        raise ValueError("Mismatch between number of boxes and number of ids")
+
+    err = rtree_bulk_load(
+        &rtree,
+        <const double*>coords_min.data,
+        <const double*>coords_max.data,
+        <const size_t*>ids.data,
+        n,
+        dim
+    )
+    if err != Success:
+        raise RuntimeError("rtree_bulk_load failed")
+
+    return RTree(<uintptr_t>rtree)

firedrake/cython/rtreeinc.pxi

@@ -0,0 +1,30 @@
+from libc.stddef cimport size_t
+from libc.stdint cimport uint32_t
+
+
+cdef extern from "rtree-capi.h":
+    ctypedef enum RTreeError:
+        Success
+        NullPointer
+        InvalidDimension
+
+    ctypedef struct RTreeH:
+        pass
+
+    RTreeError rtree_bulk_load(
+        RTreeH **tree,
+        const double *mins,
+        const double *maxs,
+        const size_t *ids,
+        size_t n,
+        uint32_t dim
+    )
+
+    RTreeError rtree_free(RTreeH *tree)
+
+    RTreeError rtree_locate_all_at_point(
+        const RTreeH *tree,
+        const double *point,
+        size_t **ids_out,
+        size_t *nids_out
+    )

firedrake/function.py

@@ -1,5 +1,5 @@
 import numpy as np
-import rtree
+import firedrake_rtree
 import sys
 import ufl
 import warnings
@@ -12,7 +12,6 @@
 from ctypes import POINTER, c_int, c_double, c_void_p
 from collections.abc import Collection
 from numbers import Number
-from pathlib import Path
 from functools import partial, cached_property
 from typing import Tuple
 
@@ -877,15 +876,13 @@ def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None):
 
     if ldargs is None:
         ldargs = []
-    libspatialindex_so = Path(rtree.core.rt._name).absolute()
-    lsi_runpath = f"-Wl,-rpath,{libspatialindex_so.parent}"
-    ldargs += [str(libspatialindex_so), lsi_runpath]
+    ldargs += [firedrake_rtree.get_lib_filename(), f"-Wl,-rpath,{firedrake_rtree.get_lib()}"]
     dll = compilation.load(
         src, "c",
         cppargs=[
             f"-I{path.dirname(__file__)}",
             f"-I{sys.prefix}/include",
-            f"-I{rtree.finder.get_include()}"
+            f"-I{firedrake_rtree.get_include()}"
         ] + [f"-I{d}/include" for d in get_petsc_dir()],
         ldargs=ldargs,
         comm=function.comm