diff --git a/firedrake/cython/rtree.pyx b/firedrake/cython/rtree.pyx index 346e03bb4b..44a8409c0f 100644 --- a/firedrake/cython/rtree.pyx +++ b/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 = coords_min.shape[0] - dim = 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 ids.shape[0] != n: raise ValueError("Mismatch between number of boxes and number of ids") err = rtree_bulk_load( &rtree, coords_min.data, coords_max.data, - ids.data, + ids.data, n, dim ) diff --git a/firedrake/evaluate.h b/firedrake/evaluate.h index 47bf93c23f..33c07eba5c 100644 --- a/firedrake/evaluate.h +++ b/firedrake/evaluate.h @@ -2,6 +2,7 @@ #define _EVALUATE_H #include +#include #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); diff --git a/firedrake/locate.c b/firedrake/locate.c index f1fa6856a3..9327ec8805 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -4,18 +4,18 @@ #include #include -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; } diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 8e6ebed99e..6e71046c99 100644 --- a/firedrake/mesh.py +++ b/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; irtree, 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; } diff --git a/pyproject.toml b/pyproject.toml index 74d38ccf8e..dedb532f9e 100644 --- a/pyproject.toml +++ b/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'", diff --git a/requirements-build.txt b/requirements-build.txt index 7916a00db7..6c8f71e67f 100644 --- a/requirements-build.txt +++ b/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