rtree bulk query

firedrake/cython/rtree.pyx

@@ -5,7 +5,7 @@ import numpy as np
 import ctypes
 import cython
 from libc.stddef cimport size_t
-from libc.stdint cimport uintptr_t, uint32_t
+from libc.stdint cimport uintptr_t, uint32_t, int64_t
 
 cdef extern from "rtree-capi.h":
     ctypedef enum RTreeError:
@@ -20,7 +20,7 @@ cdef extern from "rtree-capi.h":
         RTreeH **tree,
         const double *mins,
         const double *maxs,
-        const size_t *ids,
+        const int64_t *ids,
         size_t n,
         uint32_t dim
     )
@@ -53,7 +53,7 @@ cdef class RTree(object):
 @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):
+                    np.ndarray[np.int64_t, 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.
@@ -76,24 +76,24 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min,
     cdef:
         RTreeH* rtree
         size_t n
-        size_t dim
+        uint32_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]
+    n = coords_min.shape[0]
+    dim = coords_min.shape[1]
     if ids is None:
-        ids = np.arange(n, dtype=np.uintp)
-    elif ids.shape[0] != n:
+        ids = np.arange(n, dtype=np.int64)
+    elif <size_t>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,
+        <const int64_t*>ids.data,
         n,
         dim
     )

firedrake/evaluate.h

@@ -2,6 +2,7 @@
 #define _EVALUATE_H
 
 #include <petsc.h>
+#include <stdint.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -45,14 +46,15 @@ typedef PetscReal (*ref_cell_l1_dist_xtr)(void *data_,
 				int layer,
 				double *x);
 
-extern int locate_cell(struct Function *f,
+extern int locate_cell_from_candidates(struct Function *f,
 		       double *x,
-		       int dim,
 		       ref_cell_l1_dist try_candidate,
 		       ref_cell_l1_dist_xtr try_candidate_xtr,
 		       void *temp_ref_coords,
 		       void *found_ref_coords,
 		       double *found_ref_cell_dist_l1,
+			   size_t nids,
+			   int64_t *ids,
 			   size_t ncells_ignore,
 			   int* cells_ignore);
 

firedrake/locate.c

@@ -4,18 +4,18 @@
 #include <float.h>
 #include <evaluate.h>
 
-int locate_cell(struct Function *f,
+int locate_cell_from_candidates(struct Function *f,
         double *x,
-        int dim,
         ref_cell_l1_dist try_candidate,
         ref_cell_l1_dist_xtr try_candidate_xtr,
         void *temp_ref_coords,
         void *found_ref_coords,
         double *found_ref_cell_dist_l1,
+        size_t nids,
+        int64_t *ids,
         size_t ncells_ignore,
         int* cells_ignore)
 {
-    RTreeError err;
     int cell = -1;
     int cell_ignore_found = 0;
     /* NOTE: temp_ref_coords and found_ref_coords are actually of type
@@ -31,81 +31,45 @@ int locate_cell(struct Function *f,
        variable defined outside this function when putting together all the C
        code that needs to be compiled - see pointquery_utils.py */
 
-    size_t *ids = NULL;
-    size_t nids = 0;
-    err = rtree_locate_all_at_point((const struct RTreeH *)f->rtree, x, &ids, &nids);
-    if (err != Success) {
-        fputs("ERROR: rtree_locate_all_at_point failed.\n", stderr);
-        rtree_free_ids(ids, nids);
-        return -1;
-    }
-    if (f->extruded == 0) {
-        for (size_t i = 0; i < nids; i++) {
-            current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, ids[i], x);
-            for (size_t j = 0; j < ncells_ignore; j++) {
-                if (ids[i] == cells_ignore[j]) {
-                    cell_ignore_found = 1;
-                    break;
-                }
-            }
-            if (cell_ignore_found) {
-                cell_ignore_found = 0;
-                continue;
-            }
-            if (current_ref_cell_dist_l1 <= 0.0) {
-                /* Found cell! */
-                cell = ids[i];
-                memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
-                found_ref_cell_dist_l1[0] = current_ref_cell_dist_l1;
+    for (size_t i = 0; i < nids; ++i) {
+        for (size_t j = 0; j < ncells_ignore; j++) {
+            if (ids[i] == cells_ignore[j]) {
+                cell_ignore_found = 1;
                 break;
             }
-            else if (current_ref_cell_dist_l1 < ref_cell_dist_l1) {
-                /* getting closer... */
-                ref_cell_dist_l1 = current_ref_cell_dist_l1;
-                if (ref_cell_dist_l1 < tolerance) {
-                    /* Close to cell within tolerance so could be this cell */
-                    cell = ids[i];
-                    memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
-                    found_ref_cell_dist_l1[0] = ref_cell_dist_l1;
-                }
-            }
         }
-    }
-    else {
-        for (size_t i = 0; i < nids; i++) {
+        if (cell_ignore_found) {
+            cell_ignore_found = 0;
+            continue;
+        }
+
+        if (f->extruded == 0) {
+            current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, ids[i], x);
+        }
+        else {
             int nlayers = f->n_layers;
             int c = ids[i] / nlayers;
             int l = ids[i] % nlayers;
             current_ref_cell_dist_l1 = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x);
-            for (size_t j = 0; j < ncells_ignore; j++) {
-                if (ids[i] == cells_ignore[j]) {
-                    cell_ignore_found = 1;
-                    break;
-                }
-            }
-            if (cell_ignore_found) {
-                cell_ignore_found = 0;
-                continue;
-            }
-            if (current_ref_cell_dist_l1 <= 0.0) {
-                /* Found cell! */
+        }
+
+        if (current_ref_cell_dist_l1 <= 0.0) {
+            /* Found cell! */
+            cell = ids[i];
+            memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
+            found_ref_cell_dist_l1[0] = current_ref_cell_dist_l1;
+            break;
+        }
+        else if (current_ref_cell_dist_l1 < ref_cell_dist_l1) {
+            /* getting closer... */
+            ref_cell_dist_l1 = current_ref_cell_dist_l1;
+            if (ref_cell_dist_l1 < tolerance) {
+                /* Close to cell within tolerance so could be this cell */
                 cell = ids[i];
                 memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
-                found_ref_cell_dist_l1[0] = current_ref_cell_dist_l1;
-                break;
-            }
-            else if (current_ref_cell_dist_l1 < ref_cell_dist_l1) {
-                /* getting closer... */
-                ref_cell_dist_l1 = current_ref_cell_dist_l1;
-                if (ref_cell_dist_l1 < tolerance) {
-                    /* Close to cell within tolerance so could be this cell */
-                    cell = ids[i];
-                    memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
-                    found_ref_cell_dist_l1[0] = ref_cell_dist_l1;
-                }
+                found_ref_cell_dist_l1[0] = ref_cell_dist_l1;
             }
         }
     }
-    rtree_free_ids(ids, nids);
     return cell;
 }

firedrake/mesh.py

@@ -2748,6 +2748,15 @@ def _c_locator(self, tolerance=None):
             src += dedent(f"""
                 int locator(struct Function *f, double *x, double *X, double *ref_cell_dists_l1, {IntTypeC} *cells, {IntTypeC} npoints, size_t ncells_ignore, int* cells_ignore)
                 {{
+                    int64_t *candidate_ids = NULL;
+                    size_t *candidate_offsets = NULL;
+                    RTreeError rtree_err = rtree_locate_all_at_points(
+                        (const struct RTreeH *)f->rtree, x, npoints, &candidate_ids, &candidate_offsets);
+                    if (rtree_err != Success) {{
+                        fputs("ERROR: rtree_locate_all_at_points failed.\\n", stderr);
+                        return -1;
+                    }}
+
                     {IntTypeC} j = 0;  /* index into x and X */
                     for({IntTypeC} i=0; i<npoints; i++) {{
                         /* i is the index into cells and ref_cell_dists_l1 */
@@ -2756,18 +2765,23 @@ def _c_locator(self, tolerance=None):
                         statics in pointquery_utils.py */
                         struct ReferenceCoords temp_reference_coords, found_reference_coords;
 
+                        size_t nids_i = candidate_offsets[i + 1] - candidate_offsets[i];
+                        int64_t *ids_i = candidate_ids + candidate_offsets[i];
+
                         /* to_reference_coords and to_reference_coords_xtr are defined in
                         pointquery_utils.py. If they contain python calls, this loop will
                         not run at c-loop speed. */
                         /* cells_ignore has shape (npoints, ncells_ignore) - find the ith row */
                         int *cells_ignore_i = cells_ignore + i*ncells_ignore;
-                        cells[i] = locate_cell(f, &x[j], {self.geometric_dimension}, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &ref_cell_dists_l1[i], ncells_ignore, cells_ignore_i);
+                        cells[i] = locate_cell_from_candidates(f, &x[j], &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &ref_cell_dists_l1[i], nids_i, ids_i, ncells_ignore, cells_ignore_i);
 
                         for (int k = 0; k < {self.geometric_dimension}; k++) {{
                             X[j] = found_reference_coords.X[k];
                             j++;
                         }}
                     }}
+                    rtree_free_ids(candidate_ids, candidate_offsets[npoints]);
+                    rtree_free_offsets(candidate_offsets, npoints + 1);
                     return 0;
                 }}
             """)

firedrake/pointeval_utils.py

@@ -180,7 +180,19 @@ def predicate(index):
     double found_ref_cell_dist_l1 = DBL_MAX;
     struct ReferenceCoords temp_reference_coords, found_reference_coords;
     int cells_ignore[1] = {-1};
-    %(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &found_ref_cell_dist_l1, 1, cells_ignore);
+    RTreeError err;
+    int64_t *ids = NULL;
+    size_t nids = 0;
+    err = rtree_locate_all_at_point((const struct RTreeH *)f->rtree, x, &ids, &nids);
+    if (err != Success) {
+        fputs("ERROR: rtree_locate_all_at_point failed.\\n", stderr);
+        rtree_free_ids(ids, nids);
+        return -1;
+    }
+    %(IntType)s cell = locate_cell_from_candidates(f, x, &to_reference_coords, &to_reference_coords_xtr,
+            &temp_reference_coords, &found_reference_coords, &found_ref_cell_dist_l1,
+            nids, ids, 1, cells_ignore);
+    rtree_free_ids(ids, nids);
     if (cell == -1) {
         return -1;
     }

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
   # TODO RELEASE
   "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main",
   "h5py>3.12.1",
-  "firedrake-rtree",
+  "firedrake-rtree @ git+https://github.com/firedrakeproject/firedrake-rtree.git@main",
   "immutabledict",
   "libsupermesh>=2026.0",
   "loopy>2024.1",
@@ -156,7 +156,7 @@ docker = [  # Used in firedrake-vanilla container
 [build-system]
 requires = [
   "Cython>=3.0",
-  "firedrake-rtree",
+  "firedrake-rtree @ git+https://github.com/firedrakeproject/firedrake-rtree.git@main",
   "libsupermesh>=2026.0",
   "mpi4py>3; python_version >= '3.13'",
   "mpi4py; python_version < '3.13'",

requirements-build.txt

@@ -1,6 +1,6 @@
 # Core build dependencies (adapted from pyproject.toml)
 Cython>=3.0
-firedrake-rtree
+firedrake-rtree @ git+https://github.com/firedrakeproject/firedrake-rtree.git@main
 libsupermesh>=2026.0
 mpi4py>3; python_version >= '3.13'
 mpi4py; python_version < '3.13'
@@ -14,3 +14,5 @@ setuptools>=77.0.3
 hatchling
 meson-python
 scikit_build_core
+maturin>=1.12,<2.0
+puccinialin