Use firedrake-rtree for spatial indexing

firedrake/cython/rtree.pyx

@@ -0,0 +1,103 @@
+# 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, 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)
+
+
+cdef class RTree(object):
+    """Python class for holding an Rtree."""
+
+    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 rtree."""
+        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 : numpy.ndarray
+        Lower corner coordinates of the bounding boxes, with shape `(n, dim)`.
+    coords_max : numpy.ndarray
+        Upper corner coordinates of the bounding boxes, with shape `(n, dim)`.
+    ids : numpy.ndarray
+        Optional integer ids for each box, with shape `(n,)`. 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/evaluate.h

@@ -25,8 +25,8 @@ struct Function {
 	PetscScalar *f;
 	PetscInt *f_map;
 
-	/* Spatial index */
-	void *sidx;
+	/* rtree */
+	void *rtree;
 
 	/*
 	 * TODO:

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
 
@@ -44,7 +43,7 @@ class _CFunction(ctypes.Structure):
                 ("coords_map", POINTER(as_ctypes(IntType))),
                 ("f", c_void_p),
                 ("f_map", POINTER(as_ctypes(IntType))),
-                ("sidx", c_void_p)]
+                ("rtree", c_void_p)]
 
 
 class CoordinatelessFunction(ufl.Coefficient):
@@ -545,7 +544,7 @@ def _constant_ctypes(self):
     def _ctypes(self):
         mesh = extract_unique_domain(self)
         c_function = self._constant_ctypes
-        c_function.sidx = mesh.spatial_index and mesh.spatial_index.ctypes
+        c_function.rtree = mesh.rtree and mesh.rtree.ctypes
 
         # Return pointer
         return ctypes.pointer(c_function)
@@ -588,7 +587,7 @@ def _at(self, arg, *args, **kwargs):
         :kwarg tolerance: Tolerence to use when checking if a point is
             in a cell. Default is the ``tolerance`` provided when
             creating the :func:`~.Mesh` the function is defined on.
-            Changing this from default will cause the spatial index to
+            Changing this from default will cause the rtree to
             be rebuilt which can take some time.
         """
         # Shortcut if function space is the R-space
@@ -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()}",
             *petsctools.get_petsc_dirs(prefix="-I", subdir="include"),
         ],
         ldargs=ldargs,