From 058953077307c5510359418592870b814030a35c Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 17 Feb 2026 15:12:41 +0000 Subject: [PATCH 01/27] use rstar capi for cell location --- firedrake/cython/rstar.pyx | 69 +++++++++++++++++ firedrake/cython/rstarinc.pxi | 28 +++++++ firedrake/cython/spatialindex.pyx | 110 --------------------------- firedrake/cython/spatialindexinc.pxi | 47 ------------ firedrake/function.py | 10 +-- firedrake/locate.c | 26 +++---- firedrake/mesh.py | 24 +++--- setup.py | 34 ++++----- 8 files changed, 142 insertions(+), 206 deletions(-) create mode 100644 firedrake/cython/rstar.pyx create mode 100644 firedrake/cython/rstarinc.pxi delete mode 100644 firedrake/cython/spatialindex.pyx delete mode 100644 firedrake/cython/spatialindexinc.pxi diff --git a/firedrake/cython/rstar.pyx b/firedrake/cython/rstar.pyx new file mode 100644 index 0000000000..4717f943ab --- /dev/null +++ b/firedrake/cython/rstar.pyx @@ -0,0 +1,69 @@ +# 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 "rstarinc.pxi" + +cdef class RStarTree(object): + """Python class for holding a native spatial index object.""" + + cdef RStar_RTree* tree + + def __cinit__(self, uintptr_t tree_handle): + self.tree = 0 + if tree_handle == 0: + raise RuntimeError("invalid tree handle") + self.tree = tree_handle + + def __dealloc__(self): + if self.tree != 0: + RTree_Free(self.tree) + + @property + def ctypes(self): + """Returns a ctypes pointer to the native spatial index.""" + return ctypes.c_void_p( self.tree) + + +@cython.boundscheck(False) +@cython.wraparound(False) +def from_regions(np.ndarray[np.float64_t, ndim=2, mode="c"] regions_lo, + np.ndarray[np.float64_t, ndim=2, mode="c"] regions_hi): + """Builds a spatial index from a set of maximum bounding regions (MBRs). + + regions_lo and regions_hi must have the same size. + regions_lo[i] and regions_hi[i] contain the coordinates of the diagonally + opposite lower and higher corners of the i-th MBR, respectively. + """ + cdef: + RStarTree rstar_tree + np.ndarray[np.npy_uintp, ndim=1, mode="c"] ids + RStar_RTree* rtree + size_t n + size_t dim + RTreeError err + + assert regions_lo.shape[0] == regions_hi.shape[0] + assert regions_lo.shape[1] == regions_hi.shape[1] + n = regions_lo.shape[0] + dim = regions_lo.shape[1] + ids = np.arange(n, dtype=np.uintp) + + err = RTree_FromArray( + regions_lo.data, + regions_hi.data, + ids.data, + n, + dim, + &rtree + ) + if err != Ok: + PrintRTreeError(err) + raise RuntimeError("RTree_FromArray failed") + rstar_tree = RStarTree(rtree) + return rstar_tree diff --git a/firedrake/cython/rstarinc.pxi b/firedrake/cython/rstarinc.pxi new file mode 100644 index 0000000000..85f4d7a829 --- /dev/null +++ b/firedrake/cython/rstarinc.pxi @@ -0,0 +1,28 @@ +from libc.stddef cimport size_t + + +cdef extern from "rstar_capi.h": + ctypedef enum RTreeError: + Ok + NullPointer + DimensionNotImplemented + SizeOverflow + OutputTooSmall + Panic + + ctypedef struct RStar_RTree: + pass + + void PrintRTreeError(RTreeError error) + + RTreeError RTree_FromArray(const double *mins, + const double *maxs, + const size_t *ids, + size_t len, + size_t dim, + RStar_RTree **out_tree) + RTreeError RTree_Free(RStar_RTree *tree) + RTreeError RTree_LocateAllAtPoint(const RStar_RTree *tree, + const double *point, + size_t **ids, + size_t *nids) diff --git a/firedrake/cython/spatialindex.pyx b/firedrake/cython/spatialindex.pyx deleted file mode 100644 index 88466650e1..0000000000 --- a/firedrake/cython/spatialindex.pyx +++ /dev/null @@ -1,110 +0,0 @@ -# cython: language_level=3 - -cimport numpy as np -import numpy as np -import ctypes -import cython -from libc.stdint cimport uintptr_t - -include "spatialindexinc.pxi" - -cdef class SpatialIndex(object): - """Python class for holding a native spatial index object.""" - - cdef IndexH index - - def __cinit__(self, uintptr_t handle): - self.index = NULL - if handle == 0: - raise ValueError("SpatialIndex handle must be nonzero") - self.index = handle - - def __dealloc__(self): - Index_Destroy(self.index) - - @property - def ctypes(self): - """Returns a ctypes pointer to the native spatial index.""" - return ctypes.c_void_p( self.index) - - -cdef IndexPropertyH _make_index_properties(uint32_t dim) except *: - cdef IndexPropertyH ps = NULL - cdef RTError err = RT_None - - ps = IndexProperty_Create() - if ps == NULL: - raise RuntimeError("failed to create index properties") - - err = IndexProperty_SetIndexType(ps, RT_RTree) - if err != RT_None: - IndexProperty_Destroy(ps) - raise RuntimeError("failed to set index type") - - err = IndexProperty_SetDimension(ps, dim) - if err != RT_None: - IndexProperty_Destroy(ps) - raise RuntimeError("failed to set dimension") - - err = IndexProperty_SetIndexStorage(ps, RT_Memory) - if err != RT_None: - IndexProperty_Destroy(ps) - raise RuntimeError("failed to set index storage") - - return ps - - -@cython.boundscheck(False) -@cython.wraparound(False) -def from_regions(np.ndarray[np.float64_t, ndim=2, mode="c"] regions_lo, - np.ndarray[np.float64_t, ndim=2, mode="c"] regions_hi): - """Builds a spatial index from a set of maximum bounding regions (MBRs). - - regions_lo and regions_hi must have the same size. - regions_lo[i] and regions_hi[i] contain the coordinates of the diagonally - opposite lower and higher corners of the i-th MBR, respectively. - """ - cdef: - SpatialIndex spatial_index - np.ndarray[np.int64_t, ndim=1, mode="c"] ids - IndexPropertyH ps - IndexH index - uint64_t n - uint32_t dim - uint64_t i_stri - uint64_t d_i_stri - uint64_t d_j_stri - - assert regions_lo.shape[0] == regions_hi.shape[0] - assert regions_lo.shape[1] == regions_hi.shape[1] - n = regions_lo.shape[0] - dim = regions_lo.shape[1] - - ps = NULL - index = NULL - try: - ps = _make_index_properties(dim) - if n == 0: - # Index_CreateWithArray will fail for n=0, so create an empty index instead. - index = Index_Create(ps) - else: - ids = np.arange(n, dtype=np.int64) - - # Calculate the strides - i_stri = (ids.strides[0] // ids.itemsize) - d_i_stri = (regions_lo.strides[0] // regions_lo.itemsize) - d_j_stri = (regions_lo.strides[1] // regions_lo.itemsize) - - index = Index_CreateWithArray(ps, n, dim, - i_stri, d_i_stri, d_j_stri, - ids.data, - regions_lo.data, - regions_hi.data) - if index == NULL: - raise RuntimeError("failed to create index") - - spatial_index = SpatialIndex(index) - finally: - IndexProperty_Destroy(ps) - - return spatial_index diff --git a/firedrake/cython/spatialindexinc.pxi b/firedrake/cython/spatialindexinc.pxi deleted file mode 100644 index 8c0fb573af..0000000000 --- a/firedrake/cython/spatialindexinc.pxi +++ /dev/null @@ -1,47 +0,0 @@ -from libc.stdint cimport int64_t, uint32_t, uint64_t - -cdef extern from "spatialindex/capi/sidx_api.h": - ctypedef enum RTError: - RT_None = 0, - RT_Debug = 1, - RT_Warning = 2, - RT_Failure = 3, - RT_Fatal = 4 - - ctypedef enum RTIndexType: - RT_RTree = 0, - RT_MVRTree = 1, - RT_TPRTree = 2, - RT_InvalidIndexType = -99 - - ctypedef enum RTStorageType: - RT_Memory = 0, - RT_Disk = 1, - RT_Custom = 2, - RT_InvalidStorageType = -99 - - ctypedef enum RTIndexVariant: - RT_Linear = 0, - RT_Quadratic = 1, - RT_Star = 2, - RT_InvalidIndexVariant = -99 - - struct Index - struct Tools_PropertySet - - ctypedef Index *IndexH - ctypedef Tools_PropertySet *IndexPropertyH - - IndexPropertyH IndexProperty_Create() - RTError IndexProperty_SetIndexType(IndexPropertyH hProp, RTIndexType value) - RTError IndexProperty_SetDimension(IndexPropertyH hProp, uint32_t value) - RTError IndexProperty_SetIndexVariant(IndexPropertyH hProp, RTIndexVariant value) - RTError IndexProperty_SetIndexStorage(IndexPropertyH hProp, RTStorageType value) - void IndexProperty_Destroy(IndexPropertyH hProp) - IndexH Index_Create(IndexPropertyH hProp) - IndexH Index_CreateWithArray(IndexPropertyH hProp, uint64_t n, uint32_t dimension, - uint64_t i_stri, uint64_t d_i_stri, uint64_t d_j_stri, - int64_t *ids, double *mins, double *maxs) - RTError Index_Intersects_id(IndexH index, double* pdMin, double* pdMax, uint32_t nDimension, - int64_t** ids, uint64_t* nResults) - void Index_Destroy(IndexH index) diff --git a/firedrake/function.py b/firedrake/function.py index 5bec04b799..b7e5b37550 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -1,5 +1,4 @@ import numpy as np -import rtree import sys import ufl import warnings @@ -877,15 +876,16 @@ 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] + rstar_root = Path(__file__).resolve().parents[2] / "rstar" + rstar_include = rstar_root / "rstar-capi" / "include" + rstar_lib = rstar_root / "target" / "release" + ldargs += [f"-L{rstar_lib}", "-lrstar_capi", f"-Wl,-rpath,{rstar_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{rstar_include}" ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=ldargs, comm=function.comm diff --git a/firedrake/locate.c b/firedrake/locate.c index a243e2119f..801665b5ce 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -1,6 +1,6 @@ #include #include -#include +#include #include #include @@ -15,7 +15,7 @@ int locate_cell(struct Function *f, size_t ncells_ignore, int* cells_ignore) { - RTError err; + RTreeError err; int cell = -1; int cell_ignore_found = 0; /* NOTE: temp_ref_coords and found_ref_coords are actually of type @@ -32,20 +32,18 @@ int locate_cell(struct Function *f, code that needs to be compiled - see pointquery_utils.py */ if (f->sidx) { - int64_t *ids = NULL; - uint64_t nids = 0; - /* We treat our list of candidate cells (ids) from libspatialindex's - Index_Intersects_id as our source of truth: the point must be in - one of the cells. */ - err = Index_Intersects_id(f->sidx, x, x, dim, &ids, &nids); - if (err != RT_None) { - fputs("ERROR: Index_Intersects_id failed in libspatialindex!", stderr); + size_t *ids = NULL; + size_t nids = 0; + err = RTree_LocateAllAtPoint((const struct RStar_RTree *)f->sidx, x, &ids, &nids); + if (err != Ok) { + fputs("ERROR: RTree_LocateAllAtPoint failed in rstar-capi!\n", stderr); + PrintRTreeError(err); return -1; } if (f->extruded == 0) { - for (uint64_t i = 0; i < nids; i++) { + for (size_t i = 0; i < nids; i++) { current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, ids[i], x); - for (uint64_t j = 0; j < ncells_ignore; j++) { + for (size_t j = 0; j < ncells_ignore; j++) { if (ids[i] == cells_ignore[j]) { cell_ignore_found = 1; break; @@ -75,12 +73,12 @@ int locate_cell(struct Function *f, } } else { - for (uint64_t i = 0; i < nids; i++) { + for (size_t i = 0; i < nids; i++) { 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 (uint64_t j = 0; j < ncells_ignore; j++) { + for (size_t j = 0; j < ncells_ignore; j++) { if (ids[i] == cells_ignore[j]) { cell_ignore_found = 1; break; diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 59531066c1..4dd8f227ae 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -16,7 +16,6 @@ import enum import numbers import abc -import rtree from textwrap import dedent from pathlib import Path import typing @@ -34,7 +33,7 @@ from firedrake.cython.dmcommon import DistributedMeshOverlapType import firedrake.cython.extrusion_numbering as extnum import firedrake.extrusion_utils as eutils -import firedrake.cython.spatialindex as spatialindex +import firedrake.cython.rstar as rstar import firedrake.utils as utils from firedrake.utils import as_cstr, IntType, RealType from firedrake.logging import info_red @@ -2486,8 +2485,8 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: Tuple of arrays of shape (num_cells, gdim) containing the minimum and maximum coordinates of each cell's bounding box. - None if the geometric dimension is 1, since libspatialindex - does not support 1D. + None if the geometric dimension is 1, since point location falls + back to brute force in that case. Notes ----- @@ -2501,7 +2500,7 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: gdim = self.geometric_dimension if gdim <= 1: - info_red("libspatialindex does not support 1-dimension, falling back on brute force.") + info_red("Spatial indexing is skipped for 1-dimensional meshes, falling back on brute force.") return None coord_element = self.ufl_coordinate_element() @@ -2614,8 +2613,7 @@ def spatial_index(self): coords_max = coords_mid + (tolerance + 0.5)*d # Build spatial index - with PETSc.Log.Event("spatial_index_build"): - self._spatial_index = spatialindex.from_regions(coords_min, coords_max) + self._spatial_index = rstar.from_regions(coords_min, coords_max) self._saved_coordinate_dat_version = self.coordinates.dat.dat_version return self._spatial_index @@ -2766,20 +2764,22 @@ def _c_locator(self, tolerance=None): }} """) - libspatialindex_so = Path(rtree.core.rt._name).absolute() - lsi_runpath = f"-Wl,-rpath,{libspatialindex_so.parent}" + rstar_root = Path(__file__).resolve().parents[2] / "rstar" + rstar_include = rstar_root / "rstar-capi" / "include" + rstar_lib = rstar_root / "target" / "release" dll = compilation.load( src, "c", cppargs=[ f"-I{os.path.dirname(__file__)}", f"-I{sys.prefix}/include", - f"-I{rtree.finder.get_include()}" + f"-I{rstar_include}" ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=[ f"-L{sys.prefix}/lib", - str(libspatialindex_so), + f"-L{rstar_lib}", + "-lrstar_capi", f"-Wl,-rpath,{sys.prefix}/lib", - lsi_runpath + f"-Wl,-rpath,{rstar_lib}" ], comm=self.comm ) diff --git a/setup.py b/setup.py index 1d5b588355..d786b256c9 100644 --- a/setup.py +++ b/setup.py @@ -11,7 +11,6 @@ import numpy as np import pybind11 import petsctools -import rtree from Cython.Build import cythonize from setuptools import setup, find_packages, Extension from setuptools.command.editable_wheel import editable_wheel as _editable_wheel @@ -120,25 +119,24 @@ def __getitem__(self, key): # Set the library name and hope for the best hdf5_ = ExternalDependency(libraries=["hdf5"]) -# When we link against spatialindex or libsupermesh we need to know where +# When we link against rstar-capi or libsupermesh we need to know where # the '.so' files end up. Since installation happens in an isolated -# environment we cannot simply query rtree and libsupermesh for the +# environment we cannot simply query the currently loaded libraries for the # current paths as they will not be valid once the installation is complete. # Therefore we set the runtime library search path to all the different # possible site package locations we can think of. sitepackage_dirs = site.getsitepackages() + [site.getusersitepackages()] -# libspatialindex -# example: -# gcc -I/rtree/include -# gcc /rtree.libs/libspatialindex.so -Wl,-rpath,$ORIGIN/../../Rtree.libs -libspatialindex_so = Path(rtree.core.rt._name).absolute() -spatialindex_ = ExternalDependency( - include_dirs=[rtree.finder.get_include()], - extra_link_args=[str(libspatialindex_so)], - runtime_library_dirs=[ - os.path.join(dir, "Rtree.libs") for dir in sitepackage_dirs - ], +# rstar-capi +rstar_root = Path(__file__).resolve().parents[1] / "rstar" +rstar_capi_include = rstar_root / "rstar-capi" / "include" +rstar_capi_libdir = rstar_root / "target" / "release" +rstar_ = ExternalDependency( + include_dirs=[str(rstar_capi_include)], + library_dirs=[str(rstar_capi_libdir)], + libraries=["rstar_capi"], + runtime_library_dirs=[str(rstar_capi_libdir)], + extra_link_args=[f"-Wl,-rpath,{rstar_capi_libdir}"], ) # libsupermesh @@ -196,12 +194,12 @@ def extensions(): sources=[os.path.join("firedrake", "cython", "patchimpl.pyx")], **(mpi_ + petsc_ + numpy_) )) - # firedrake/cython/spatialindex.pyx: numpy, spatialindex + # firedrake/cython/rstar.pyx: numpy, rstar-capi cython_list.append(Extension( - name="firedrake.cython.spatialindex", + name="firedrake.cython.rstar", language="c", - sources=[os.path.join("firedrake", "cython", "spatialindex.pyx")], - **(mpi_ + numpy_ + spatialindex_) + sources=[os.path.join("firedrake", "cython", "rstar.pyx")], + **(mpi_ + numpy_ + rstar_) )) # firedrake/cython/supermeshimpl.pyx: petsc, numpy, supermesh cython_list.append(Extension( From 5b68b0e9c346fc25586a5575f5e3d5cbc47ac14d Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 17 Feb 2026 15:13:50 +0000 Subject: [PATCH 02/27] add logging --- firedrake/mesh.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 4dd8f227ae..767e9bf9cf 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2475,7 +2475,7 @@ def clear_spatial_index(self): the coordinate field).""" self._spatial_index = None - @cached_property + @utils.cached_property @PETSc.Log.EventDecorator() def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: """Calculates bounding boxes for spatial indexing. @@ -2607,13 +2607,15 @@ def spatial_index(self): else: coords_min, coords_max = self.bounding_box_coords tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 + coords_mid = (coords_max + coords_min)/2 d = np.max(coords_max - coords_min, axis=1)[:, None] coords_min = coords_mid - (tolerance + 0.5)*d coords_max = coords_mid + (tolerance + 0.5)*d # Build spatial index - self._spatial_index = rstar.from_regions(coords_min, coords_max) + with PETSc.Log.Event("spatial_index_build"): + self._spatial_index = rstar.from_regions(coords_min, coords_max) self._saved_coordinate_dat_version = self.coordinates.dat.dat_version return self._spatial_index @@ -4276,7 +4278,8 @@ def _parent_mesh_embedding( # it. if redundant: # rank 0 broadcasts coords to all ranks - coords_local = icomm.bcast(coords, root=0) + with PETSc.Log.Event("pm_embed_coords_bcast"): + coords_local = icomm.bcast(coords, root=0) ncoords_local = coords_local.shape[0] coords_global = coords_local ncoords_global = coords_global.shape[0] @@ -4294,17 +4297,20 @@ def _parent_mesh_embedding( # 30-34. coords_local = coords ncoords_local = coords.shape[0] - ncoords_local_allranks = icomm.allgather(ncoords_local) + with PETSc.Log.Event("pm_embed_ncoords_allgather"): + ncoords_local_allranks = icomm.allgather(ncoords_local) ncoords_global = sum(ncoords_local_allranks) # The below code looks complicated but it's just an allgather of the # (variable length) coords_local array such that they are concatenated. coords_local_size = np.array(coords_local.size) coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) - icomm.Allgatherv(coords_local_size, coords_local_sizes) + with PETSc.Log.Event("pm_embed_coord_sizes_allgatherv"): + icomm.Allgatherv(coords_local_size, coords_local_sizes) coords_global = np.empty( (ncoords_global, coords.shape[1]), dtype=coords_local.dtype ) - icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) + with PETSc.Log.Event("pm_embed_coords_allgatherv"): + icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) # # ncoords_local_allranks is in rank order so we can just sum up the # # previous ranks to get the starting index for the global numbering. # # For rank 0 we make use of the fact that sum([]) = 0. @@ -4314,9 +4320,11 @@ def _parent_mesh_embedding( global_idxs_global = np.arange(coords_global.shape[0]) input_coords_idxs_local = np.arange(ncoords_local) input_coords_idxs_global = np.empty(ncoords_global, dtype=int) - icomm.Allgatherv( - input_coords_idxs_local, (input_coords_idxs_global, ncoords_local_allranks) - ) + with PETSc.Log.Event("pm_embed_inputidx_allgatherv"): + icomm.Allgatherv( + input_coords_idxs_local, + (input_coords_idxs_global, ncoords_local_allranks), + ) input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) input_ranks_global = np.empty(ncoords_global, dtype=int) icomm.Allgatherv( From 8e48eea220aa57af6cb7f8be51c6829bc9a97e2b Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 17 Feb 2026 15:19:56 +0000 Subject: [PATCH 03/27] remove allgather logs --- firedrake/mesh.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 767e9bf9cf..d254e53c45 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4297,20 +4297,17 @@ def _parent_mesh_embedding( # 30-34. coords_local = coords ncoords_local = coords.shape[0] - with PETSc.Log.Event("pm_embed_ncoords_allgather"): - ncoords_local_allranks = icomm.allgather(ncoords_local) + ncoords_local_allranks = icomm.allgather(ncoords_local) ncoords_global = sum(ncoords_local_allranks) # The below code looks complicated but it's just an allgather of the # (variable length) coords_local array such that they are concatenated. coords_local_size = np.array(coords_local.size) coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) - with PETSc.Log.Event("pm_embed_coord_sizes_allgatherv"): - icomm.Allgatherv(coords_local_size, coords_local_sizes) + icomm.Allgatherv(coords_local_size, coords_local_sizes) coords_global = np.empty( (ncoords_global, coords.shape[1]), dtype=coords_local.dtype ) - with PETSc.Log.Event("pm_embed_coords_allgatherv"): - icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) + icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) # # ncoords_local_allranks is in rank order so we can just sum up the # # previous ranks to get the starting index for the global numbering. # # For rank 0 we make use of the fact that sum([]) = 0. @@ -4320,11 +4317,10 @@ def _parent_mesh_embedding( global_idxs_global = np.arange(coords_global.shape[0]) input_coords_idxs_local = np.arange(ncoords_local) input_coords_idxs_global = np.empty(ncoords_global, dtype=int) - with PETSc.Log.Event("pm_embed_inputidx_allgatherv"): - icomm.Allgatherv( - input_coords_idxs_local, - (input_coords_idxs_global, ncoords_local_allranks), - ) + icomm.Allgatherv( + input_coords_idxs_local, + (input_coords_idxs_global, ncoords_local_allranks), + ) input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) input_ranks_global = np.empty(ncoords_global, dtype=int) icomm.Allgatherv( From 6e79d56527cc4b519794dc6d76d9216e2e70d8ff Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 23 Feb 2026 15:48:20 +0000 Subject: [PATCH 04/27] update comments --- firedrake/mesh.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index d254e53c45..bc7218d1f1 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2485,8 +2485,9 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: Tuple of arrays of shape (num_cells, gdim) containing the minimum and maximum coordinates of each cell's bounding box. - None if the geometric dimension is 1, since point location falls - back to brute force in that case. + None if the geometric dimension is 1. + TODO: rstar supports 1D, we just need to add a case for it in the + C api. Notes ----- @@ -2500,7 +2501,7 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: gdim = self.geometric_dimension if gdim <= 1: - info_red("Spatial indexing is skipped for 1-dimensional meshes, falling back on brute force.") + info_red("TODO: add 1D case to rstar-capi") return None coord_element = self.ufl_coordinate_element() @@ -2606,8 +2607,8 @@ def spatial_index(self): return None else: coords_min, coords_max = self.bounding_box_coords - tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 + tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 coords_mid = (coords_max + coords_min)/2 d = np.max(coords_max - coords_min, axis=1)[:, None] coords_min = coords_mid - (tolerance + 0.5)*d From c1be0f44d5a68ff5c4b60c6fa9166b83fa1432ad Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 23 Feb 2026 16:11:15 +0000 Subject: [PATCH 05/27] fix cached property --- firedrake/mesh.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index bc7218d1f1..65203d646d 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2475,7 +2475,7 @@ def clear_spatial_index(self): the coordinate field).""" self._spatial_index = None - @utils.cached_property + @cached_property @PETSc.Log.EventDecorator() def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: """Calculates bounding boxes for spatial indexing. From 0c7a6939ca6e72bc55a8e74225ac090ccafcc7e7 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 24 Feb 2026 15:25:02 +0000 Subject: [PATCH 06/27] merge logging more logging --- firedrake/mesh.py | 223 +++++++++++++++++++++++++--------------------- 1 file changed, 123 insertions(+), 100 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 65203d646d..efc3328bda 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2267,6 +2267,7 @@ def input_ordering(self): ) @staticmethod + @PETSc.Log.EventDecorator() def _make_input_ordering_sf(swarm, nroots, ilocal): # ilocal = None -> leaves are swarm points [0, 1, 2, ...). # ilocal can also be Firedrake cell numbers. @@ -2609,6 +2610,7 @@ def spatial_index(self): coords_min, coords_max = self.bounding_box_coords tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 + coords_mid = (coords_max + coords_min)/2 d = np.max(coords_max - coords_min, axis=1)[:, None] coords_min = coords_mid - (tolerance + 0.5)*d @@ -4055,7 +4057,8 @@ def _dmswarm_create( _, coordsdim = coords.shape # Create a DMSWARM - swarm = FiredrakeDMSwarm().create(comm=comm) + with PETSc.Log.Event("FiredrakeDMSwarmCreate"): + swarm = FiredrakeDMSwarm().create(comm=comm) # save the fields on the swarm swarm.fields = default_fields + default_extra_fields + other_fields @@ -4108,33 +4111,34 @@ def _dmswarm_create( # topological dimension of the base mesh. # NOTE ensure that swarm.restoreField is called for each field too! - swarm_coords = swarm.getField("DMSwarmPIC_coor").reshape((num_vertices, gdim)) - cell_id_name = swarm.getCellDMActive().getCellID() - swarm_parent_cell_nums = swarm.getField(cell_id_name).ravel() - field_parent_cell_nums = swarm.getField("parentcellnum").ravel() - field_reference_coords = swarm.getField("refcoord").reshape((num_vertices, tdim)) - field_global_index = swarm.getField("globalindex").ravel() - field_rank = swarm.getField("DMSwarm_rank").ravel() - field_input_rank = swarm.getField("inputrank").ravel() - field_input_index = swarm.getField("inputindex").ravel() - swarm_coords[...] = coords - swarm_parent_cell_nums[...] = plex_parent_cell_nums - field_parent_cell_nums[...] = parent_cell_nums - field_reference_coords[...] = reference_coords - field_global_index[...] = coords_idxs - field_rank[...] = ranks - field_input_rank[...] = input_ranks - field_input_index[...] = input_coords_idxs - - # have to restore fields once accessed to allow access again - swarm.restoreField("inputindex") - swarm.restoreField("inputrank") - swarm.restoreField("DMSwarm_rank") - swarm.restoreField("globalindex") - swarm.restoreField("refcoord") - swarm.restoreField("parentcellnum") - swarm.restoreField("DMSwarmPIC_coor") - swarm.restoreField(cell_id_name) + with PETSc.Log.Event("FiredrakeDMSwarmSetFields"): + swarm_coords = swarm.getField("DMSwarmPIC_coor").reshape((num_vertices, gdim)) + cell_id_name = swarm.getCellDMActive().getCellID() + swarm_parent_cell_nums = swarm.getField(cell_id_name).ravel() + field_parent_cell_nums = swarm.getField("parentcellnum").ravel() + field_reference_coords = swarm.getField("refcoord").reshape((num_vertices, tdim)) + field_global_index = swarm.getField("globalindex").ravel() + field_rank = swarm.getField("DMSwarm_rank").ravel() + field_input_rank = swarm.getField("inputrank").ravel() + field_input_index = swarm.getField("inputindex").ravel() + swarm_coords[...] = coords + swarm_parent_cell_nums[...] = plex_parent_cell_nums + field_parent_cell_nums[...] = parent_cell_nums + field_reference_coords[...] = reference_coords + field_global_index[...] = coords_idxs + field_rank[...] = ranks + field_input_rank[...] = input_ranks + field_input_index[...] = input_coords_idxs + + # have to restore fields once accessed to allow access again + swarm.restoreField("inputindex") + swarm.restoreField("inputrank") + swarm.restoreField("DMSwarm_rank") + swarm.restoreField("globalindex") + swarm.restoreField("refcoord") + swarm.restoreField("parentcellnum") + swarm.restoreField("DMSwarmPIC_coor") + swarm.restoreField(cell_id_name) if extruded: field_base_parent_cell_nums = swarm.getField("parentcellbasenum").ravel() @@ -4298,17 +4302,20 @@ def _parent_mesh_embedding( # 30-34. coords_local = coords ncoords_local = coords.shape[0] - ncoords_local_allranks = icomm.allgather(ncoords_local) + with PETSc.Log.Event("pm_embed_ncoords_allgather"): + ncoords_local_allranks = icomm.allgather(ncoords_local) ncoords_global = sum(ncoords_local_allranks) # The below code looks complicated but it's just an allgather of the # (variable length) coords_local array such that they are concatenated. coords_local_size = np.array(coords_local.size) coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) - icomm.Allgatherv(coords_local_size, coords_local_sizes) + with PETSc.Log.Event("pm_embed_coord_sizes_allgatherv"): + icomm.Allgatherv(coords_local_size, coords_local_sizes) coords_global = np.empty( (ncoords_global, coords.shape[1]), dtype=coords_local.dtype ) - icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) + with PETSc.Log.Event("pm_embed_coords_allgatherv"): + icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) # # ncoords_local_allranks is in rank order so we can just sum up the # # previous ranks to get the starting index for the global numbering. # # For rank 0 we make use of the fact that sum([]) = 0. @@ -4318,15 +4325,17 @@ def _parent_mesh_embedding( global_idxs_global = np.arange(coords_global.shape[0]) input_coords_idxs_local = np.arange(ncoords_local) input_coords_idxs_global = np.empty(ncoords_global, dtype=int) - icomm.Allgatherv( - input_coords_idxs_local, - (input_coords_idxs_global, ncoords_local_allranks), - ) + with PETSc.Log.Event("pm_embed_inputidx_allgatherv"): + icomm.Allgatherv( + input_coords_idxs_local, + (input_coords_idxs_global, ncoords_local_allranks), + ) input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) input_ranks_global = np.empty(ncoords_global, dtype=int) - icomm.Allgatherv( - input_ranks_local, (input_ranks_global, ncoords_local_allranks) - ) + with PETSc.Log.Event("pm_embed_inputranks_allgatherv"): + icomm.Allgatherv( + input_ranks_local, (input_ranks_global, ncoords_local_allranks) + ) ( parent_cell_nums, reference_coords, @@ -4346,9 +4355,10 @@ def _parent_mesh_embedding( visible_ranks[:parent_mesh.cell_set.size] = parent_mesh.comm.rank visible_ranks[parent_mesh.cell_set.size:] = -1 # Halo exchange the visible ranks so that each rank knows which ranks can see each cell. - dmcommon.exchange_cell_orientations( - parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks - ) + with PETSc.Log.Event("pm_embed_exchange_cell_orientations"): + dmcommon.exchange_cell_orientations( + parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks + ) locally_visible = parent_cell_nums != -1 if parent_mesh.extruded: @@ -4504,50 +4514,60 @@ def _swarm_original_ordering_preserve( # Gather everything except original_ordering_coords_local from all mpi # ranks - ncoords_local_allranks = comm.allgather(ncoords_local) + with PETSc.Log.Event("original_ordering_allgather_ncoords_local"): + ncoords_local_allranks = comm.allgather(ncoords_local) ncoords_global = sum(ncoords_local_allranks) parent_cell_nums_global = np.empty( ncoords_global, dtype=parent_cell_nums_local.dtype ) - comm.Allgatherv( - parent_cell_nums_local, (parent_cell_nums_global, ncoords_local_allranks) - ) + with PETSc.Log.Event("original_ordering_allgatherv_parent_cell_nums"): + comm.Allgatherv( + parent_cell_nums_local, (parent_cell_nums_global, ncoords_local_allranks) + ) plex_parent_cell_nums_global = np.empty( ncoords_global, dtype=plex_parent_cell_nums_local.dtype ) - comm.Allgatherv( - plex_parent_cell_nums_local, - (plex_parent_cell_nums_global, ncoords_local_allranks), - ) + with PETSc.Log.Event("original_ordering_allgatherv_plex_parent_cell_nums"): + comm.Allgatherv( + plex_parent_cell_nums_local, + (plex_parent_cell_nums_global, ncoords_local_allranks), + ) reference_coords_local_size = np.array(reference_coords_local.size) reference_coords_local_sizes = np.empty(comm.size, dtype=int) - comm.Allgatherv(reference_coords_local_size, reference_coords_local_sizes) + with PETSc.Log.Event("original_ordering_allgatherv_reference_coords_sizes"): + comm.Allgatherv(reference_coords_local_size, reference_coords_local_sizes) reference_coords_global = np.empty( (ncoords_global, reference_coords_local.shape[1]), dtype=reference_coords_local.dtype, ) - comm.Allgatherv( - reference_coords_local, (reference_coords_global, reference_coords_local_sizes) - ) + with PETSc.Log.Event("original_ordering_allgatherv_reference_coords"): + comm.Allgatherv( + reference_coords_local, (reference_coords_global, reference_coords_local_sizes) + ) global_idxs_global = np.empty(ncoords_global, dtype=global_idxs_local.dtype) - comm.Allgatherv(global_idxs_local, (global_idxs_global, ncoords_local_allranks)) + with PETSc.Log.Event("original_ordering_allgatherv_global_idxs"): + comm.Allgatherv(global_idxs_local, (global_idxs_global, ncoords_local_allranks)) ranks_global = np.empty(ncoords_global, dtype=ranks_local.dtype) - comm.Allgatherv(ranks_local, (ranks_global, ncoords_local_allranks)) + with PETSc.Log.Event("original_ordering_allgatherv_ranks"): + comm.Allgatherv(ranks_local, (ranks_global, ncoords_local_allranks)) input_ranks_global = np.empty(ncoords_global, dtype=input_ranks_local.dtype) - comm.Allgatherv(input_ranks_local, (input_ranks_global, ncoords_local_allranks)) + with PETSc.Log.Event("original_ordering_allgatherv_input_ranks"): + comm.Allgatherv(input_ranks_local, (input_ranks_global, ncoords_local_allranks)) input_idxs_global = np.empty(ncoords_global, dtype=input_idxs_local.dtype) - comm.Allgatherv(input_idxs_local, (input_idxs_global, ncoords_local_allranks)) + with PETSc.Log.Event("original_ordering_allgatherv_input_idxs"): + comm.Allgatherv(input_idxs_local, (input_idxs_global, ncoords_local_allranks)) # Sort by global index, which will be in rank order (they probably already # are but we can't rely on that) - global_idxs_global_order = np.argsort(global_idxs_global) + with PETSc.Log.Event("original_ordering_sort_by_global_idx"): + global_idxs_global_order = np.argsort(global_idxs_global) sorted_parent_cell_nums_global = parent_cell_nums_global[global_idxs_global_order] sorted_plex_parent_cell_nums_global = plex_parent_cell_nums_global[ global_idxs_global_order @@ -4561,13 +4581,15 @@ def _swarm_original_ordering_preserve( sorted_input_idxs_global = input_idxs_global[global_idxs_global_order] # Check order is correct - we can probably remove this eventually since it's # quite expensive - if not np.all(sorted_input_ranks_global[1:] >= sorted_input_ranks_global[:-1]): - raise ValueError("Global indexing has not ordered the ranks as expected") + with PETSc.Log.Event("original_ordering_check_global_idx_order"): + if not np.all(sorted_input_ranks_global[1:] >= sorted_input_ranks_global[:-1]): + raise ValueError("Global indexing has not ordered the ranks as expected") # get rid of any duplicated global indices (i.e. points in halos) - unique_global_idxs, unique_idxs = np.unique( - sorted_global_idxs_global, return_index=True - ) + with PETSc.Log.Event("original_ordering_remove_duplicate_global_idxs"): + unique_global_idxs, unique_idxs = np.unique( + sorted_global_idxs_global, return_index=True + ) unique_parent_cell_nums_global = sorted_parent_cell_nums_global[unique_idxs] unique_plex_parent_cell_nums_global = sorted_plex_parent_cell_nums_global[ unique_idxs @@ -4599,47 +4621,48 @@ def _swarm_original_ordering_preserve( # check if the input indices are in order from zero - this can also probably # be removed eventually because, again, it's expensive. - if not np.array_equal(output_input_idxs, np.arange(output_input_idxs.size)): - raise ValueError( - "Global indexing has not ordered the input indices as expected." - ) - if len(output_global_idxs) != len(original_ordering_coords_local): - raise ValueError( - "The number of local global indices which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_parent_cell_nums) != len(original_ordering_coords_local): - raise ValueError( - "The number of local parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_plex_parent_cell_nums) != len(original_ordering_coords_local): - raise ValueError( - "The number of local plex parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_reference_coords) != len(original_ordering_coords_local): - raise ValueError( - "The number of local reference coordinates which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_ranks) != len(original_ordering_coords_local): - raise ValueError( - "The number of local rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_input_ranks) != len(original_ordering_coords_local): - raise ValueError( - "The number of local input rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_input_idxs) != len(original_ordering_coords_local): - raise ValueError( - "The number of local input indices which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if extruded: - if len(output_base_parent_cell_nums) != len(original_ordering_coords_local): + with PETSc.Log.Event("original_ordering_check_input_idx_order"): + if not np.array_equal(output_input_idxs, np.arange(output_input_idxs.size)): + raise ValueError( + "Global indexing has not ordered the input indices as expected." + ) + if len(output_global_idxs) != len(original_ordering_coords_local): + raise ValueError( + "The number of local global indices which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_parent_cell_nums) != len(original_ordering_coords_local): raise ValueError( - "The number of local base parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + "The number of local parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." ) - if len(output_extrusion_heights) != len(original_ordering_coords_local): + if len(output_plex_parent_cell_nums) != len(original_ordering_coords_local): raise ValueError( - "The number of local extrusion heights which will be used to make the swarm do not match the input number of original ordering coordinates." + "The number of local plex parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." ) + if len(output_reference_coords) != len(original_ordering_coords_local): + raise ValueError( + "The number of local reference coordinates which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_ranks) != len(original_ordering_coords_local): + raise ValueError( + "The number of local rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_input_ranks) != len(original_ordering_coords_local): + raise ValueError( + "The number of local input rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_input_idxs) != len(original_ordering_coords_local): + raise ValueError( + "The number of local input indices which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if extruded: + if len(output_base_parent_cell_nums) != len(original_ordering_coords_local): + raise ValueError( + "The number of local base parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_extrusion_heights) != len(original_ordering_coords_local): + raise ValueError( + "The number of local extrusion heights which will be used to make the swarm do not match the input number of original ordering coordinates." + ) return _dmswarm_create( [], From d2fb722348badbf957f8911e4e3fe9caf9200ef4 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 25 Feb 2026 17:35:26 +0000 Subject: [PATCH 07/27] more logs --- firedrake/mesh.py | 124 ++++++++++++++++++++++++---------------------- 1 file changed, 64 insertions(+), 60 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index efc3328bda..a9079fff04 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4402,52 +4402,54 @@ def _parent_mesh_embedding( # the distance is the same then we have found the correct cell. If we # cannot make a match to owned_rank and distance then we can't see the # point. - changed_ranks_tied = changed_ranks & ~changed_ref_cell_dists_l1 - if any(changed_ranks_tied): - cells_ignore_T = np.asarray([np.copy(parent_cell_nums)]) - while any(changed_ranks_tied): - ( - parent_cell_nums[changed_ranks_tied], - new_reference_coords, - ref_cell_dists_l1[changed_ranks_tied], - ) = parent_mesh.locate_cells_ref_coords_and_dists( - coords_global[changed_ranks_tied], - tolerance, - cells_ignore=cells_ignore_T.T[changed_ranks_tied, :], - ) - # delete extra dimension if necessary - if parent_mesh.geometric_dimension > parent_mesh.topological_dimension: - new_reference_coords = new_reference_coords[:, : parent_mesh.topological_dimension] - reference_coords[changed_ranks_tied, :] = new_reference_coords - # remove newly lost points - locally_visible[changed_ranks_tied] = ( - parent_cell_nums[changed_ranks_tied] != -1 - ) - changed_ranks_tied &= locally_visible - # if new ref_cell_dists_l1 > owned_ref_cell_dists_l1 then we should - # disregard the point. - locally_visible[changed_ranks_tied] &= ( - ref_cell_dists_l1[changed_ranks_tied] - <= owned_ref_cell_dists_l1[changed_ranks_tied] - ) - changed_ranks_tied &= locally_visible - # update the identified rank - if parent_mesh.extruded: - _retry_cell_nums = parent_cell_nums[changed_ranks_tied] // (parent_mesh.layers - 1) - else: - _retry_cell_nums = parent_cell_nums[changed_ranks_tied] - ranks[changed_ranks_tied] = visible_ranks[_retry_cell_nums] - # if the rank now matches then we have found the correct cell - locally_visible[changed_ranks_tied] &= ( - owned_ranks[changed_ranks_tied] == ranks[changed_ranks_tied] - ) - # remove these rank matches from changed_ranks_tied - changed_ranks_tied &= ~locally_visible - # add more cells to ignore - cells_ignore_T = np.vstack(( - cells_ignore_T, - parent_cell_nums) - ) + with PETSc.Log.Event("pm_embed_tied_ranks"): + changed_ranks_tied = changed_ranks & ~changed_ref_cell_dists_l1 + if any(changed_ranks_tied): + cells_ignore_T = np.asarray([np.copy(parent_cell_nums)]) + while any(changed_ranks_tied): + with PETSc.Log.Event("pm_embed_tied_rank_locate_cells"): + ( + parent_cell_nums[changed_ranks_tied], + new_reference_coords, + ref_cell_dists_l1[changed_ranks_tied], + ) = parent_mesh.locate_cells_ref_coords_and_dists( + coords_global[changed_ranks_tied], + tolerance, + cells_ignore=cells_ignore_T.T[changed_ranks_tied, :], + ) + # delete extra dimension if necessary + if parent_mesh.geometric_dimension > parent_mesh.topological_dimension: + new_reference_coords = new_reference_coords[:, : parent_mesh.topological_dimension] + reference_coords[changed_ranks_tied, :] = new_reference_coords + # remove newly lost points + locally_visible[changed_ranks_tied] = ( + parent_cell_nums[changed_ranks_tied] != -1 + ) + changed_ranks_tied &= locally_visible + # if new ref_cell_dists_l1 > owned_ref_cell_dists_l1 then we should + # disregard the point. + locally_visible[changed_ranks_tied] &= ( + ref_cell_dists_l1[changed_ranks_tied] + <= owned_ref_cell_dists_l1[changed_ranks_tied] + ) + changed_ranks_tied &= locally_visible + # update the identified rank + if parent_mesh.extruded: + _retry_cell_nums = parent_cell_nums[changed_ranks_tied] // (parent_mesh.layers - 1) + else: + _retry_cell_nums = parent_cell_nums[changed_ranks_tied] + ranks[changed_ranks_tied] = visible_ranks[_retry_cell_nums] + # if the rank now matches then we have found the correct cell + locally_visible[changed_ranks_tied] &= ( + owned_ranks[changed_ranks_tied] == ranks[changed_ranks_tied] + ) + # remove these rank matches from changed_ranks_tied + changed_ranks_tied &= ~locally_visible + # add more cells to ignore + cells_ignore_T = np.vstack(( + cells_ignore_T, + parent_cell_nums) + ) # Any ranks which are still -np.inf are not in the mesh missing_global_idxs = np.where(owned_ranks == -np.inf)[0] @@ -4461,20 +4463,22 @@ def _parent_mesh_embedding( reference_coords[missing_coords_idxs_on_rank, :] = np.nan owned_ranks[missing_coords_idxs_on_rank] = parent_mesh.comm.size + 1 - if exclude_halos and parent_mesh.comm.size > 1: - off_rank_coords_idxs = np.where( - (owned_ranks != parent_mesh.comm.rank) - & (owned_ranks != parent_mesh.comm.size + 1) - )[0] - locally_visible[off_rank_coords_idxs] = False - - coords_embedded = np.compress(locally_visible, coords_global, axis=0) - global_idxs = np.compress(locally_visible, global_idxs_global, axis=0) - reference_coords = np.compress(locally_visible, reference_coords, axis=0) - parent_cell_nums = np.compress(locally_visible, parent_cell_nums, axis=0) - owned_ranks = np.compress(locally_visible, owned_ranks, axis=0).astype(int) - input_ranks = np.compress(locally_visible, input_ranks_global, axis=0) - input_coords_idxs = np.compress(locally_visible, input_coords_idxs_global, axis=0) + with PETSc.Log.Event("pm_embed_exclude_halos_parallel"): + if exclude_halos and parent_mesh.comm.size > 1: + off_rank_coords_idxs = np.where( + (owned_ranks != parent_mesh.comm.rank) + & (owned_ranks != parent_mesh.comm.size + 1) + )[0] + locally_visible[off_rank_coords_idxs] = False + + with PETSc.Log.Event("pm_embed_compress_arrays"): + coords_embedded = np.compress(locally_visible, coords_global, axis=0) + global_idxs = np.compress(locally_visible, global_idxs_global, axis=0) + reference_coords = np.compress(locally_visible, reference_coords, axis=0) + parent_cell_nums = np.compress(locally_visible, parent_cell_nums, axis=0) + owned_ranks = np.compress(locally_visible, owned_ranks, axis=0).astype(int) + input_ranks = np.compress(locally_visible, input_ranks_global, axis=0) + input_coords_idxs = np.compress(locally_visible, input_coords_idxs_global, axis=0) return ( coords_embedded, From a81f19d71e444b406ccb865481af56c752064bd5 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 25 Feb 2026 18:07:22 +0000 Subject: [PATCH 08/27] more logging --- firedrake/mesh.py | 222 +++++++++++++++++++++++----------------------- 1 file changed, 113 insertions(+), 109 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index a9079fff04..2ce6f9b4bc 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4278,64 +4278,65 @@ def _parent_mesh_embedding( "VertexOnlyMeshes don't have a working locate_cells_ref_coords_and_dists method" ) - with temp_internal_comm(parent_mesh.comm) as icomm: - # In parallel, we need to make sure we know which point is which and save - # it. - if redundant: - # rank 0 broadcasts coords to all ranks - with PETSc.Log.Event("pm_embed_coords_bcast"): - coords_local = icomm.bcast(coords, root=0) - ncoords_local = coords_local.shape[0] - coords_global = coords_local - ncoords_global = coords_global.shape[0] - global_idxs_global = np.arange(coords_global.shape[0]) - input_coords_idxs_local = np.arange(ncoords_local) - input_coords_idxs_global = input_coords_idxs_local - input_ranks_local = np.zeros(ncoords_local, dtype=int) - input_ranks_global = input_ranks_local - else: - # Here, we have to assume that all points we can see are unique. - # We therefore gather all points on all ranks in rank order: if rank 0 - # has 10 points, rank 1 has 20 points, and rank 3 has 5 points, then - # rank 0's points have global numbering 0-9, rank 1's points have - # global numbering 10-29, and rank 3's points have global numbering - # 30-34. - coords_local = coords - ncoords_local = coords.shape[0] - with PETSc.Log.Event("pm_embed_ncoords_allgather"): - ncoords_local_allranks = icomm.allgather(ncoords_local) - ncoords_global = sum(ncoords_local_allranks) - # The below code looks complicated but it's just an allgather of the - # (variable length) coords_local array such that they are concatenated. - coords_local_size = np.array(coords_local.size) - coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) - with PETSc.Log.Event("pm_embed_coord_sizes_allgatherv"): - icomm.Allgatherv(coords_local_size, coords_local_sizes) - coords_global = np.empty( - (ncoords_global, coords.shape[1]), dtype=coords_local.dtype - ) - with PETSc.Log.Event("pm_embed_coords_allgatherv"): - icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) - # # ncoords_local_allranks is in rank order so we can just sum up the - # # previous ranks to get the starting index for the global numbering. - # # For rank 0 we make use of the fact that sum([]) = 0. - # startidx = sum(ncoords_local_allranks[:parent_mesh.comm.rank]) - # endidx = startidx + ncoords_local - # global_idxs_global = np.arange(startidx, endidx) - global_idxs_global = np.arange(coords_global.shape[0]) - input_coords_idxs_local = np.arange(ncoords_local) - input_coords_idxs_global = np.empty(ncoords_global, dtype=int) - with PETSc.Log.Event("pm_embed_inputidx_allgatherv"): - icomm.Allgatherv( - input_coords_idxs_local, - (input_coords_idxs_global, ncoords_local_allranks), - ) - input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) - input_ranks_global = np.empty(ncoords_global, dtype=int) - with PETSc.Log.Event("pm_embed_inputranks_allgatherv"): - icomm.Allgatherv( - input_ranks_local, (input_ranks_global, ncoords_local_allranks) + with PETSc.Log.Event("pm_embed_initial"): + with temp_internal_comm(parent_mesh.comm) as icomm: + # In parallel, we need to make sure we know which point is which and save + # it. + if redundant: + # rank 0 broadcasts coords to all ranks + with PETSc.Log.Event("pm_embed_coords_bcast"): + coords_local = icomm.bcast(coords, root=0) + ncoords_local = coords_local.shape[0] + coords_global = coords_local + ncoords_global = coords_global.shape[0] + global_idxs_global = np.arange(coords_global.shape[0]) + input_coords_idxs_local = np.arange(ncoords_local) + input_coords_idxs_global = input_coords_idxs_local + input_ranks_local = np.zeros(ncoords_local, dtype=int) + input_ranks_global = input_ranks_local + else: + # Here, we have to assume that all points we can see are unique. + # We therefore gather all points on all ranks in rank order: if rank 0 + # has 10 points, rank 1 has 20 points, and rank 3 has 5 points, then + # rank 0's points have global numbering 0-9, rank 1's points have + # global numbering 10-29, and rank 3's points have global numbering + # 30-34. + coords_local = coords + ncoords_local = coords.shape[0] + with PETSc.Log.Event("pm_embed_ncoords_allgather"): + ncoords_local_allranks = icomm.allgather(ncoords_local) + ncoords_global = sum(ncoords_local_allranks) + # The below code looks complicated but it's just an allgather of the + # (variable length) coords_local array such that they are concatenated. + coords_local_size = np.array(coords_local.size) + coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) + with PETSc.Log.Event("pm_embed_coord_sizes_allgatherv"): + icomm.Allgatherv(coords_local_size, coords_local_sizes) + coords_global = np.empty( + (ncoords_global, coords.shape[1]), dtype=coords_local.dtype ) + with PETSc.Log.Event("pm_embed_coords_allgatherv"): + icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) + # # ncoords_local_allranks is in rank order so we can just sum up the + # # previous ranks to get the starting index for the global numbering. + # # For rank 0 we make use of the fact that sum([]) = 0. + # startidx = sum(ncoords_local_allranks[:parent_mesh.comm.rank]) + # endidx = startidx + ncoords_local + # global_idxs_global = np.arange(startidx, endidx) + global_idxs_global = np.arange(coords_global.shape[0]) + input_coords_idxs_local = np.arange(ncoords_local) + input_coords_idxs_global = np.empty(ncoords_global, dtype=int) + with PETSc.Log.Event("pm_embed_inputidx_allgatherv"): + icomm.Allgatherv( + input_coords_idxs_local, + (input_coords_idxs_global, ncoords_local_allranks), + ) + input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) + input_ranks_global = np.empty(ncoords_global, dtype=int) + with PETSc.Log.Event("pm_embed_inputranks_allgatherv"): + icomm.Allgatherv( + input_ranks_local, (input_ranks_global, ncoords_local_allranks) + ) ( parent_cell_nums, reference_coords, @@ -4359,14 +4360,15 @@ def _parent_mesh_embedding( dmcommon.exchange_cell_orientations( parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks ) - locally_visible = parent_cell_nums != -1 + with PETSc.Log.Event("pm_embed_build_rank_arrays"): + locally_visible = parent_cell_nums != -1 - if parent_mesh.extruded: - # Halo exchange of visible_ranks is over the base mesh topology and cell numbering, - # so we need to map back to extruded cell numbering after indexing parent_cell_nums. - locally_visible_cell_nums = parent_cell_nums[locally_visible] // (parent_mesh.layers - 1) - else: - locally_visible_cell_nums = parent_cell_nums[locally_visible] + if parent_mesh.extruded: + # Halo exchange of visible_ranks is over the base mesh topology and cell numbering, + # so we need to map back to extruded cell numbering after indexing parent_cell_nums. + locally_visible_cell_nums = parent_cell_nums[locally_visible] // (parent_mesh.layers - 1) + else: + locally_visible_cell_nums = parent_cell_nums[locally_visible] # In parallel there will regularly be disagreements about which cell owns a # point when those points are close to mesh partition boundaries. @@ -4389,13 +4391,13 @@ def _parent_mesh_embedding( owned_ranks = np.empty_like(rank_candidates) parent_mesh.comm.Allreduce(rank_candidates, owned_ranks, op=MPI.MAX) - changed_ref_cell_dists_l1 = owned_ref_cell_dists_l1 != ref_cell_dists_l1 - changed_ranks = owned_ranks != ranks + changed_ref_cell_dists_l1 = owned_ref_cell_dists_l1 != ref_cell_dists_l1 + changed_ranks = owned_ranks != ranks - # If distance has changed the the point is not in local mesh partition - # since some other cell on another rank is closer. - locally_visible[changed_ref_cell_dists_l1] = False - parent_cell_nums[changed_ref_cell_dists_l1] = -1 + # If distance has changed the the point is not in local mesh partition + # since some other cell on another rank is closer. + locally_visible[changed_ref_cell_dists_l1] = False + parent_cell_nums[changed_ref_cell_dists_l1] = -1 # If the rank has changed but the distance hasn't then there was a tie # break and we need to search for the point again, this time disallowing # the previously identified cell: if we match the identified owned_rank AND @@ -4572,17 +4574,18 @@ def _swarm_original_ordering_preserve( # are but we can't rely on that) with PETSc.Log.Event("original_ordering_sort_by_global_idx"): global_idxs_global_order = np.argsort(global_idxs_global) - sorted_parent_cell_nums_global = parent_cell_nums_global[global_idxs_global_order] - sorted_plex_parent_cell_nums_global = plex_parent_cell_nums_global[ - global_idxs_global_order - ] - sorted_reference_coords_global = reference_coords_global[ - global_idxs_global_order, : - ] - sorted_global_idxs_global = global_idxs_global[global_idxs_global_order] - sorted_ranks_global = ranks_global[global_idxs_global_order] - sorted_input_ranks_global = input_ranks_global[global_idxs_global_order] - sorted_input_idxs_global = input_idxs_global[global_idxs_global_order] + with PETSc.Log.Event("original_ordering_reindex_sorted"): + sorted_parent_cell_nums_global = parent_cell_nums_global[global_idxs_global_order] + sorted_plex_parent_cell_nums_global = plex_parent_cell_nums_global[ + global_idxs_global_order + ] + sorted_reference_coords_global = reference_coords_global[ + global_idxs_global_order, : + ] + sorted_global_idxs_global = global_idxs_global[global_idxs_global_order] + sorted_ranks_global = ranks_global[global_idxs_global_order] + sorted_input_ranks_global = input_ranks_global[global_idxs_global_order] + sorted_input_idxs_global = input_idxs_global[global_idxs_global_order] # Check order is correct - we can probably remove this eventually since it's # quite expensive with PETSc.Log.Event("original_ordering_check_global_idx_order"): @@ -4594,34 +4597,35 @@ def _swarm_original_ordering_preserve( unique_global_idxs, unique_idxs = np.unique( sorted_global_idxs_global, return_index=True ) - unique_parent_cell_nums_global = sorted_parent_cell_nums_global[unique_idxs] - unique_plex_parent_cell_nums_global = sorted_plex_parent_cell_nums_global[ - unique_idxs - ] - unique_reference_coords_global = sorted_reference_coords_global[unique_idxs, :] - unique_ranks_global = sorted_ranks_global[unique_idxs] - unique_input_ranks_global = sorted_input_ranks_global[unique_idxs] - unique_input_idxs_global = sorted_input_idxs_global[unique_idxs] - - # save the points on this rank which match the input rank ready for output - input_ranks_match = unique_input_ranks_global == comm.rank - output_global_idxs = unique_global_idxs[input_ranks_match] - output_parent_cell_nums = unique_parent_cell_nums_global[input_ranks_match] - output_plex_parent_cell_nums = unique_plex_parent_cell_nums_global[ - input_ranks_match - ] - output_reference_coords = unique_reference_coords_global[input_ranks_match, :] - output_ranks = unique_ranks_global[input_ranks_match] - output_input_ranks = unique_input_ranks_global[input_ranks_match] - output_input_idxs = unique_input_idxs_global[input_ranks_match] - if extruded: - ( - output_base_parent_cell_nums, - output_extrusion_heights, - ) = _parent_extrusion_numbering(output_parent_cell_nums, layers) - else: - output_base_parent_cell_nums = None - output_extrusion_heights = None + with PETSc.Log.Event("original_ordering_filter_and_select_outputs"): + unique_parent_cell_nums_global = sorted_parent_cell_nums_global[unique_idxs] + unique_plex_parent_cell_nums_global = sorted_plex_parent_cell_nums_global[ + unique_idxs + ] + unique_reference_coords_global = sorted_reference_coords_global[unique_idxs, :] + unique_ranks_global = sorted_ranks_global[unique_idxs] + unique_input_ranks_global = sorted_input_ranks_global[unique_idxs] + unique_input_idxs_global = sorted_input_idxs_global[unique_idxs] + + # save the points on this rank which match the input rank ready for output + input_ranks_match = unique_input_ranks_global == comm.rank + output_global_idxs = unique_global_idxs[input_ranks_match] + output_parent_cell_nums = unique_parent_cell_nums_global[input_ranks_match] + output_plex_parent_cell_nums = unique_plex_parent_cell_nums_global[ + input_ranks_match + ] + output_reference_coords = unique_reference_coords_global[input_ranks_match, :] + output_ranks = unique_ranks_global[input_ranks_match] + output_input_ranks = unique_input_ranks_global[input_ranks_match] + output_input_idxs = unique_input_idxs_global[input_ranks_match] + if extruded: + ( + output_base_parent_cell_nums, + output_extrusion_heights, + ) = _parent_extrusion_numbering(output_parent_cell_nums, layers) + else: + output_base_parent_cell_nums = None + output_extrusion_heights = None # check if the input indices are in order from zero - this can also probably # be removed eventually because, again, it's expensive. From 92e0aec1dd6230a2189e41ae79fcb108778c2cea Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 4 Mar 2026 17:15:43 +0000 Subject: [PATCH 09/27] adapt to new API --- firedrake/cython/rstar.pyx | 21 ++++++----- firedrake/cython/rstarinc.pxi | 68 ++++++++++++++++++++++++----------- firedrake/locate.c | 9 +++-- 3 files changed, 62 insertions(+), 36 deletions(-) diff --git a/firedrake/cython/rstar.pyx b/firedrake/cython/rstar.pyx index 4717f943ab..b3b6a67321 100644 --- a/firedrake/cython/rstar.pyx +++ b/firedrake/cython/rstar.pyx @@ -12,17 +12,17 @@ include "rstarinc.pxi" cdef class RStarTree(object): """Python class for holding a native spatial index object.""" - cdef RStar_RTree* tree + cdef RTreeH* tree def __cinit__(self, uintptr_t tree_handle): - self.tree = 0 + self.tree = 0 if tree_handle == 0: raise RuntimeError("invalid tree handle") - self.tree = tree_handle + self.tree = tree_handle def __dealloc__(self): - if self.tree != 0: - RTree_Free(self.tree) + if self.tree != 0: + rtree_free(self.tree) @property def ctypes(self): @@ -43,7 +43,7 @@ def from_regions(np.ndarray[np.float64_t, ndim=2, mode="c"] regions_lo, cdef: RStarTree rstar_tree np.ndarray[np.npy_uintp, ndim=1, mode="c"] ids - RStar_RTree* rtree + RTreeH* rtree size_t n size_t dim RTreeError err @@ -54,16 +54,15 @@ def from_regions(np.ndarray[np.float64_t, ndim=2, mode="c"] regions_lo, dim = regions_lo.shape[1] ids = np.arange(n, dtype=np.uintp) - err = RTree_FromArray( + err = rtree_bulk_load( + &rtree, regions_lo.data, regions_hi.data, ids.data, n, - dim, - &rtree + dim ) - if err != Ok: - PrintRTreeError(err) + if err != Success: raise RuntimeError("RTree_FromArray failed") rstar_tree = RStarTree(rtree) return rstar_tree diff --git a/firedrake/cython/rstarinc.pxi b/firedrake/cython/rstarinc.pxi index 85f4d7a829..c5f30af919 100644 --- a/firedrake/cython/rstarinc.pxi +++ b/firedrake/cython/rstarinc.pxi @@ -1,28 +1,56 @@ from libc.stddef cimport size_t +from libc.stdint cimport uint32_t -cdef extern from "rstar_capi.h": +cdef extern from "rstar-capi.h": ctypedef enum RTreeError: - Ok + Success NullPointer - DimensionNotImplemented - SizeOverflow - OutputTooSmall - Panic + InvalidDimension + NodeNotLeaf - ctypedef struct RStar_RTree: + ctypedef struct RTreeH: pass - void PrintRTreeError(RTreeError error) - - RTreeError RTree_FromArray(const double *mins, - const double *maxs, - const size_t *ids, - size_t len, - size_t dim, - RStar_RTree **out_tree) - RTreeError RTree_Free(RStar_RTree *tree) - RTreeError RTree_LocateAllAtPoint(const RStar_RTree *tree, - const double *point, - size_t **ids, - size_t *nids) + ctypedef struct RTreeNodeH: + 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 + ) + + RTreeError rtree_root_node( + const RTreeH *tree, + RTreeNodeH **node + ) + + RTreeError rtree_node_children( + const RTreeNodeH *node, + RTreeNodeH ***children_out, + size_t *nchildren_out + ) + + RTreeError rtree_node_id( + const RTreeNodeH *node, + size_t *id_out + ) + + RTreeError rtree_node_envelope( + const RTreeNodeH *node, + double *mins_out, + double *maxs_out + ) diff --git a/firedrake/locate.c b/firedrake/locate.c index 801665b5ce..80c9fe06d6 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -1,6 +1,6 @@ #include #include -#include +#include #include #include @@ -34,10 +34,9 @@ int locate_cell(struct Function *f, if (f->sidx) { size_t *ids = NULL; size_t nids = 0; - err = RTree_LocateAllAtPoint((const struct RStar_RTree *)f->sidx, x, &ids, &nids); - if (err != Ok) { + err = rtree_locate_all_at_point((const struct RTreeH *)f->sidx, x, &ids, &nids); + if (err != Success) { fputs("ERROR: RTree_LocateAllAtPoint failed in rstar-capi!\n", stderr); - PrintRTreeError(err); return -1; } if (f->extruded == 0) { @@ -107,7 +106,7 @@ int locate_cell(struct Function *f, } } } - free(ids); + rtree_free_ids(ids, nids); } else { if (f->extruded == 0) { for (int c = 0; c < f->n_cols; c++) { From d0b3a5324110eb43fe5a01b35029c346088df873 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 10 Mar 2026 11:43:29 +0000 Subject: [PATCH 10/27] rebase --- firedrake/mesh.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 2ce6f9b4bc..d610ccf2e6 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4390,14 +4390,13 @@ def _parent_mesh_embedding( rank_candidates = np.where(ref_cell_dists_l1 == owned_ref_cell_dists_l1, ranks, -np.inf) owned_ranks = np.empty_like(rank_candidates) parent_mesh.comm.Allreduce(rank_candidates, owned_ranks, op=MPI.MAX) + changed_ref_cell_dists_l1 = owned_ref_cell_dists_l1 != ref_cell_dists_l1 + changed_ranks = owned_ranks != ranks - changed_ref_cell_dists_l1 = owned_ref_cell_dists_l1 != ref_cell_dists_l1 - changed_ranks = owned_ranks != ranks - - # If distance has changed the the point is not in local mesh partition - # since some other cell on another rank is closer. - locally_visible[changed_ref_cell_dists_l1] = False - parent_cell_nums[changed_ref_cell_dists_l1] = -1 + # If distance has changed the the point is not in local mesh partition + # since some other cell on another rank is closer. + locally_visible[changed_ref_cell_dists_l1] = False + parent_cell_nums[changed_ref_cell_dists_l1] = -1 # If the rank has changed but the distance hasn't then there was a tie # break and we need to search for the point again, this time disallowing # the previously identified cell: if we match the identified owned_rank AND From 6e5d9eb79f117df20b9ac21deb1f2dfcf6ae2b1e Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 9 Mar 2026 19:22:19 +0000 Subject: [PATCH 11/27] spatial index in 1D --- firedrake/locate.c | 71 ---------------------------------------------- firedrake/mesh.py | 25 +++++----------- 2 files changed, 7 insertions(+), 89 deletions(-) diff --git a/firedrake/locate.c b/firedrake/locate.c index 80c9fe06d6..f8003458bd 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -107,77 +107,6 @@ int locate_cell(struct Function *f, } } rtree_free_ids(ids, nids); - } else { - if (f->extruded == 0) { - for (int c = 0; c < f->n_cols; c++) { - current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, c, x); - for (uint64_t j = 0; j < ncells_ignore; j++) { - if (c == 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 = c; - 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 = c; - memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); - found_ref_cell_dist_l1[0] = ref_cell_dist_l1; - } - } - } - } - else { - for (int c = 0; c < f->n_cols; c++) { - for (int l = 0; l < f->n_layers; l++) { - current_ref_cell_dist_l1 = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x); - for (uint64_t j = 0; j < ncells_ignore; j++) { - if (l == 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 = l; - 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 = l; - memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); - found_ref_cell_dist_l1[0] = ref_cell_dist_l1; - } - } - } - if (cell != -1) { - cell = c * f->n_layers + cell; - break; - } - } - } } return cell; } diff --git a/firedrake/mesh.py b/firedrake/mesh.py index d610ccf2e6..16ff37cd01 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2478,7 +2478,7 @@ def clear_spatial_index(self): @cached_property @PETSc.Log.EventDecorator() - def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: + def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray]: """Calculates bounding boxes for spatial indexing. Returns @@ -2486,10 +2486,6 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: Tuple of arrays of shape (num_cells, gdim) containing the minimum and maximum coordinates of each cell's bounding box. - None if the geometric dimension is 1. - TODO: rstar supports 1D, we just need to add a case for it in the - C api. - Notes ----- If we have a higher-order (bendy) mesh we project the mesh coordinates into @@ -2500,11 +2496,6 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: from firedrake import function, functionspace from firedrake.parloops import par_loop, READ, MIN, MAX - gdim = self.geometric_dimension - if gdim <= 1: - info_red("TODO: add 1D case to rstar-capi") - return None - coord_element = self.ufl_coordinate_element() coord_degree = coord_element.degree() if np.all(np.asarray(coord_degree) == 1): @@ -2539,7 +2530,7 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: return np.min(all_coords, axis=1), np.max(all_coords, axis=1) # Extruded case: calculate the bounding boxes for all cells by running a kernel - V = functionspace.VectorFunctionSpace(mesh, "DG", 0, dim=gdim) + V = functionspace.VectorFunctionSpace(mesh, "DG", 0, dim=self.geometric_dimension) coords_min = function.Function(V, dtype=RealType) coords_max = function.Function(V, dtype=RealType) @@ -2548,7 +2539,7 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None: _, nodes_per_cell = cell_node_list.shape - domain = f"{{[d, i]: 0 <= d < {gdim} and 0 <= i < {nodes_per_cell}}}" + domain = f"{{[d, i]: 0 <= d < {self.geometric_dimension} and 0 <= i < {nodes_per_cell}}}" instructions = """ for d, i f_min[0, d] = fmin(f_min[0, d], f[i, d]) @@ -2603,11 +2594,10 @@ def spatial_index(self): # within the mesh tolerance. # NOTE: getattr doesn't work here due to the inheritance games that are # going on in getattr. - if self.bounding_box_coords is None: - # This happens in 1D meshes - return None - else: - coords_min, coords_max = self.bounding_box_coords + coords_min, coords_max = self.bounding_box_coords + if self.geometric_dimension == 1: + coords_min = coords_min.reshape(len(coords_min), -1) + coords_max = coords_max.reshape(len(coords_max), -1) tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 @@ -2616,7 +2606,6 @@ def spatial_index(self): coords_min = coords_mid - (tolerance + 0.5)*d coords_max = coords_mid + (tolerance + 0.5)*d - # Build spatial index with PETSc.Log.Event("spatial_index_build"): self._spatial_index = rstar.from_regions(coords_min, coords_max) self._saved_coordinate_dat_version = self.coordinates.dat.dat_version From 50d18bef496dab7bcd06737d7ae568b7d9dd3920 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 10 Mar 2026 11:31:05 +0000 Subject: [PATCH 12/27] partial fix --- firedrake/locate.c | 126 ++++++++++++++++++++++----------------------- firedrake/mesh.py | 4 +- 2 files changed, 64 insertions(+), 66 deletions(-) diff --git a/firedrake/locate.c b/firedrake/locate.c index f8003458bd..4575ceb954 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -31,82 +31,80 @@ 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 */ - if (f->sidx) { - size_t *ids = NULL; - size_t nids = 0; - err = rtree_locate_all_at_point((const struct RTreeH *)f->sidx, x, &ids, &nids); - if (err != Success) { - fputs("ERROR: RTree_LocateAllAtPoint failed in rstar-capi!\n", stderr); - 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; + size_t *ids = NULL; + size_t nids = 0; + err = rtree_locate_all_at_point((const struct RTreeH *)f->sidx, x, &ids, &nids); + if (err != Success) { + fputs("ERROR: RTree_LocateAllAtPoint failed in rstar-capi!\n", stderr); + 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 (current_ref_cell_dist_l1 <= 0.0) { - /* Found cell! */ + } + 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; + 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; } } } - else { - for (size_t i = 0; i < nids; i++) { - 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; + } + else { + for (size_t i = 0; i < nids; i++) { + 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 (current_ref_cell_dist_l1 <= 0.0) { - /* Found cell! */ + } + 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; + 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); } + rtree_free_ids(ids, nids); return cell; } diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 16ff37cd01..1efc10f8a0 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2596,8 +2596,8 @@ def spatial_index(self): # going on in getattr. coords_min, coords_max = self.bounding_box_coords if self.geometric_dimension == 1: - coords_min = coords_min.reshape(len(coords_min), -1) - coords_max = coords_max.reshape(len(coords_max), -1) + coords_min = coords_min.reshape(-1, 1) + coords_max = coords_max.reshape(-1, 1) tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 From 863fc75705cbc6b2259594ed22edd92d95f4bb84 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 10 Mar 2026 15:39:14 +0000 Subject: [PATCH 13/27] fix test This test depended on the ordering of the candidate cells returned by the rtree query. Instead we check the missing coordinates correctly as suggested by the comment --- .../test_vertex_only_mesh_generation.py | 34 +++++++------------ 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/tests/firedrake/vertexonly/test_vertex_only_mesh_generation.py b/tests/firedrake/vertexonly/test_vertex_only_mesh_generation.py index 2ca7acda3f..344e9e2702 100644 --- a/tests/firedrake/vertexonly/test_vertex_only_mesh_generation.py +++ b/tests/firedrake/vertexonly/test_vertex_only_mesh_generation.py @@ -3,6 +3,7 @@ import pytest import numpy as np from mpi4py import MPI +from pytest_mpi.parallel_assert import parallel_assert # Utility Functions @@ -124,39 +125,30 @@ def verify_vertexonly_mesh(m, vm, inputvertexcoords, name): assert vm.name == name # Find in-bounds and non-halo-region input coordinates in_bounds = [] - ref_cell_dists_l1 = [] + ref_cell_dists_l1 = {} # this method of getting owned cells works for all mesh types owned_cells = len(Function(FunctionSpace(m, "DG", 0)).dat.data_ro) for i in range(len(inputvertexcoords)): cell_num, _, ref_cell_dist_l1 = m.locate_cells_ref_coords_and_dists(inputvertexcoords[i].reshape(1, gdim)) if cell_num != -1 and cell_num < owned_cells: in_bounds.append(i) - ref_cell_dists_l1.append(ref_cell_dist_l1) + ref_cell_dists_l1[i] = ref_cell_dist_l1[0] # In parallel locate_cells_ref_coords_and_dists might give point # duplication: the voting algorithm in VertexOnlyMesh should remove these. # We can check that this is the case by seeing if all the missing - # coordinates have positive distances from the reference cell but this is a - # faff. For now just check that, where point duplication occurs, we have - # some ref_cell_dists_l1 over half the mesh tolerance (i.e. where I'd - # expect this to start ocurring). - total_cells = MPI.COMM_WORLD.allreduce(len(vm.coordinates.dat.data_ro), op=MPI.SUM) - total_in_bounds = MPI.COMM_WORLD.allreduce(len(in_bounds), op=MPI.SUM) + # coordinates have positive distances from the reference cell. skip_in_bounds_checks = False - local_cells = len(vm.coordinates.dat.data_ro) - if total_cells != total_in_bounds: - assert MPI.COMM_WORLD.size > 1 # i.e. we're in parallel - assert total_cells < total_in_bounds # i.e. some points are duplicated - local_in_bounds = len(in_bounds) - if not local_cells == local_in_bounds and local_in_bounds > 0: - # This assertion needs to happen in parallel! - assertion = (max(ref_cell_dists_l1) > 0.5*m.tolerance) - skip_in_bounds_checks = True - else: - assertion = True + final_input_indices = np.copy(vm.topology_dm.getField("inputindex").ravel()) + vm.topology_dm.restoreField("inputindex") + final_input_indices = np.sort(final_input_indices) + local_in_bounds = np.array(sorted(in_bounds), dtype=int) + if not np.array_equal(final_input_indices, local_in_bounds): + missing_coordinate_indices = np.setdiff1d(local_in_bounds, final_input_indices) + assertion = np.all([ref_cell_dists_l1[i] > 0.0 for i in missing_coordinate_indices]) + skip_in_bounds_checks = True else: assertion = True - # FIXME: Replace with parallel assert when it's merged into pytest-mpi - assert min(MPI.COMM_WORLD.allgather([assertion])) + parallel_assert([assertion]) # Correct local coordinates (though not guaranteed to be in same order) if not skip_in_bounds_checks: From e7b45af872de20ddad16b24abb2c18cb4e8cc306 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 10 Mar 2026 16:03:50 +0000 Subject: [PATCH 14/27] tidy up cython --- firedrake/cython/rstar.pyx | 55 ++++++++++++++++++++++++-------------- firedrake/mesh.py | 2 +- 2 files changed, 36 insertions(+), 21 deletions(-) diff --git a/firedrake/cython/rstar.pyx b/firedrake/cython/rstar.pyx index b3b6a67321..55f280b17e 100644 --- a/firedrake/cython/rstar.pyx +++ b/firedrake/cython/rstar.pyx @@ -9,7 +9,7 @@ from libc.stdint cimport uintptr_t include "rstarinc.pxi" -cdef class RStarTree(object): +cdef class RTree(object): """Python class for holding a native spatial index object.""" cdef RTreeH* tree @@ -32,37 +32,52 @@ cdef class RStarTree(object): @cython.boundscheck(False) @cython.wraparound(False) -def from_regions(np.ndarray[np.float64_t, ndim=2, mode="c"] regions_lo, - np.ndarray[np.float64_t, ndim=2, mode="c"] regions_hi): - """Builds a spatial index from a set of maximum bounding regions (MBRs). +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. - regions_lo and regions_hi must have the same size. - regions_lo[i] and regions_hi[i] contain the coordinates of the diagonally - opposite lower and higher corners of the i-th MBR, respectively. - """ + 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 built R*-tree. + """ cdef: - RStarTree rstar_tree - np.ndarray[np.npy_uintp, ndim=1, mode="c"] ids RTreeH* rtree size_t n size_t dim - RTreeError err + 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") - assert regions_lo.shape[0] == regions_hi.shape[0] - assert regions_lo.shape[1] == regions_hi.shape[1] - n = regions_lo.shape[0] - dim = regions_lo.shape[1] - ids = np.arange(n, dtype=np.uintp) + n = coords_min.shape[0] + dim = coords_min.shape[1] + if ids is not None and ids.shape[0] != n: + raise ValueError("Mismatch between number of boxes and number of ids") + else: + ids = np.arange(n, dtype=np.uintp) err = rtree_bulk_load( &rtree, - regions_lo.data, - regions_hi.data, + coords_min.data, + coords_max.data, ids.data, n, dim ) if err != Success: raise RuntimeError("RTree_FromArray failed") - rstar_tree = RStarTree(rtree) - return rstar_tree + + return RTree(rtree) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 1efc10f8a0..5d4bda5657 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2607,7 +2607,7 @@ def spatial_index(self): coords_max = coords_mid + (tolerance + 0.5)*d with PETSc.Log.Event("spatial_index_build"): - self._spatial_index = rstar.from_regions(coords_min, coords_max) + self._spatial_index = rstar.build_from_aabb(coords_min, coords_max) self._saved_coordinate_dat_version = self.coordinates.dat.dat_version return self._spatial_index From 9b1d409a20c70bb1dccf25eff226749adea135d4 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 12 Mar 2026 18:16:21 +0000 Subject: [PATCH 15/27] fixes --- firedrake/cython/petschdr.pxi | 12 ++ firedrake/cython/rstar.pyx | 237 +++++++++++++++++++++++++++++++++- 2 files changed, 245 insertions(+), 4 deletions(-) diff --git a/firedrake/cython/petschdr.pxi b/firedrake/cython/petschdr.pxi index d7b1478ca0..1026473594 100644 --- a/firedrake/cython/petschdr.pxi +++ b/firedrake/cython/petschdr.pxi @@ -8,6 +8,7 @@ cdef extern from "mpi-compat.h": cdef extern from "petsc.h": ctypedef long PetscInt + ctypedef int PetscMPIInt ctypedef double PetscReal ctypedef double PetscScalar ctypedef enum PetscBool: @@ -32,6 +33,17 @@ cdef extern from "petscsys.h" nogil: PetscErrorCode PetscFree(void*) PetscErrorCode PetscFree2(void*,void*) PetscErrorCode PetscSortIntWithArray(PetscInt,PetscInt[],PetscInt[]) + PetscErrorCode PetscCommBuildTwoSided( + MPI.MPI_Comm comm, + PetscMPIInt count, + MPI.MPI_Datatype dtype, + PetscMPIInt nto, + const PetscMPIInt *toranks, + const void *todata, + PetscMPIInt *nfrom, + PetscMPIInt **fromranks, + void *fromdata + ) cdef extern from "petscdmtypes.h" nogil: ctypedef enum PetscDMPolytopeType "DMPolytopeType": diff --git a/firedrake/cython/rstar.pyx b/firedrake/cython/rstar.pyx index 55f280b17e..3d4d7045fb 100644 --- a/firedrake/cython/rstar.pyx +++ b/firedrake/cython/rstar.pyx @@ -7,10 +7,15 @@ import cython from libc.stddef cimport size_t from libc.stdint cimport uintptr_t +cimport mpi4py.MPI as MPI +from mpi4py.libmpi cimport MPI_INT +from petsc4py.PETSc cimport CHKERR + include "rstarinc.pxi" +include "petschdr.pxi" cdef class RTree(object): - """Python class for holding a native spatial index object.""" + """Python class for holding a spatial index.""" cdef RTreeH* tree @@ -64,10 +69,10 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min, n = coords_min.shape[0] dim = coords_min.shape[1] - if ids is not None and ids.shape[0] != n: - raise ValueError("Mismatch between number of boxes and number of ids") - else: + 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, @@ -81,3 +86,227 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min, raise RuntimeError("RTree_FromArray failed") return RTree(rtree) + +@cython.boundscheck(False) +@cython.wraparound(False) +def locate_all_at_point( + RTree rtree, + np.ndarray[np.float64_t, ndim=1, mode="c"] point): + """Return the ids of all leaves whose bounding box contains ``point``. + + Parameters + ---------- + rtree : RTree + The R*-tree to query. + point : (dim,) float64 array + The query point. Must have the same dimensionality as the tree. + + Returns + ------- + An array of integer ids corresponding to the leaves whose bounding boxes contain the query point. + The array must be freed by the caller using rtree_locate_all_at_point_free. + """ + cdef: + size_t *ids_out = NULL + size_t nids_out = 0 + RTreeError err + np.ndarray[np.npy_uintp, ndim=1, mode="c"] result + + err = rtree_locate_all_at_point( + rtree.tree, + point.data, + &ids_out, + &nids_out, + ) + if err != Success: + raise RuntimeError("rtree_locate_all_at_point failed") + + result = np.empty(nids_out, dtype=np.uintp) + for i in range(nids_out): + result[i] = ids_out[i] + rtree_free_ids(ids_out, nids_out) + return result + + +@cython.boundscheck(False) +@cython.wraparound(False) +def discover_ranks( + RTree rtree, + np.ndarray[np.float64_t, ndim=2, mode="c"] points, + MPI.Comm comm): + """Query a distributed Rtree and discover which ranks to send points, + and which ranks are going to send points to us. + + For each point in `points`, we find all candidate ranks whose bounding + boxes contain that point. We then use `PetscCommBuildTwoSided` to discover + which ranks will be sending points to us. + + Parameters + ---------- + rtree : RTree + The distributed Rtree built by :func:`build_from_aabb` with + rank numbers as leaf ids. + points : (n_points, gdim) float64 array + The local points to send to remote ranks. + comm : mpi4py.MPI.Comm + The MPI communicator. + + Returns + ------- + toranks : (nto,) int32 array + Target ranks to send points to. + send_offsets : (nto + 1,) int32 array + Points destined for `toranks[i]` are + `point_indices[send_offsets[i]:send_offsets[i+1]]`. + point_indices : (total_sends,) int32 array + Indices into `points` determining which points to send. + fromranks : (nfrom,) int32 array + Ranks that will send points to us. + recv_counts : (nfrom,) int32 array + Number of points we will receive from each rank in `fromranks`. + """ + cdef: + size_t *ids_out = NULL + size_t nids_out = 0 + RTreeError err + size_t i, j + size_t n_points = points.shape[0] + MPI.MPI_Comm mpi_comm = comm.ob_mpi + PetscMPIInt nto + PetscMPIInt nfrom = 0 + PetscMPIInt *fromranks = NULL + void *fromdata = NULL + np.ndarray[np.int32_t, ndim=1, mode="c"] toranks + np.ndarray[np.int32_t, ndim=1, mode="c"] send_counts + np.ndarray[np.int32_t, ndim=1, mode="c"] point_indices + np.ndarray[np.int32_t, ndim=1, mode="c"] send_offsets + np.ndarray[np.int32_t, ndim=1, mode="c"] fromranks_out + np.ndarray[np.int32_t, ndim=1, mode="c"] recv_counts_out + PetscMPIInt k + + # map dest_rank -> list of point indices to send there + rank_to_indices: dict[int, list[int]] = {} + for i in range(n_points): + err = rtree_locate_all_at_point( + rtree.tree, + &points[i, 0], + &ids_out, + &nids_out, + ) + if err != Success: + raise RuntimeError("rtree_locate_all_at_point failed") + + # Points may lie in multiple bounding boxes owned by the same rank + seen_ranks: set[int] = set() + for j in range(nids_out): + seen_ranks.add(ids_out[j]) + rtree_free_ids(ids_out, nids_out) + + ids_out = NULL + for dest_rank in seen_ranks: + if dest_rank in rank_to_indices: + rank_to_indices[dest_rank].append(i) + else: + rank_to_indices[dest_rank] = [i] + + nto = len(rank_to_indices) + toranks = np.empty(nto, dtype=np.int32) + send_counts = np.empty(nto, dtype=np.int32) + send_offsets = np.empty(nto + 1, dtype=np.int32) + all_indices: list[int] = [] + send_offsets[0] = 0 + for i, (rank, idx_list) in enumerate(rank_to_indices.items()): + toranks[i] = rank + send_counts[i] = len(idx_list) + send_offsets[i + 1] = send_offsets[i] + len(idx_list) + all_indices.extend(idx_list) + point_indices = np.array(all_indices, dtype=np.int32) + + # Routine that discovers communicating ranks given one-sided information + CHKERR(PetscCommBuildTwoSided( + mpi_comm, + 1, # sending/receiving one entry (the number of points) + MPI_INT, + nto, # number of ranks to send data to + toranks.data, # ranks to send to (array of length nto) + send_counts.data, # data to send to each rank (array of length nto) + &nfrom, # number of ranks we're receiving messages from + &fromranks, # ranks we're receiving messages from (array of length nfrom) + &fromdata, # data we're receiving from each rank (array of length nfrom) + )) + + # Copy petsc-allocated results into numpy arrays then free them + fromranks_out = np.empty(nfrom, dtype=np.int32) + recv_counts_out = np.empty(nfrom, dtype=np.int32) + for k in range(nfrom): + fromranks_out[k] = fromranks[k] + recv_counts_out[k] = (fromdata)[k] + CHKERR(PetscFree(fromranks)) + CHKERR(PetscFree(fromdata)) + + return toranks, send_offsets, point_indices, fromranks_out, recv_counts_out + + +cdef class RTreeNode(object): + """Python class for holding a spatial index node.""" + + cdef RTreeNodeH* node + + def __cinit__(self, uintptr_t node_handle): + self.node = 0 + if node_handle == 0: + raise RuntimeError("invalid node handle") + self.node = node_handle + + def __dealloc__(self): + if self.node != 0: + rtree_node_free(self.node) + self.node = 0 + + +def root_node(RTree rtree): + """Return the root node of the R*-tree.""" + cdef: + RTreeNodeH* node + RTreeError err + err = rtree_root_node(rtree.tree, &node) + if err != Success: + raise RuntimeError("rtree_root_node failed") + return RTreeNode(node) + + +def node_children(RTreeNode node): + """Return the children of an R*-tree node as a list of RStarTreeNode.""" + cdef: + RTreeNodeH** children + size_t nchildren + RTreeError err + err = rtree_node_children(node.node, &children, &nchildren) + if err != Success: + raise RuntimeError("rtree_node_children failed") + result = [RTreeNode(children[i]) for i in range(nchildren)] + rtree_node_children_free(children, nchildren) + return result + + +def node_id(RTreeNode node): + """Return the id of a leaf node.""" + cdef: + size_t id_out + RTreeError err + err = rtree_node_id(node.node, &id_out) + if err != Success: + raise RuntimeError("rtree_node_id failed (node may not be a leaf)") + return id_out + + +def node_envelope(RTreeNode node, size_t dim): + """Return the (mins, maxs) bounding envelope of an R*-tree node.""" + cdef: + np.ndarray[np.float64_t, ndim=1, mode="c"] mins = np.empty(dim, dtype=np.float64) + np.ndarray[np.float64_t, ndim=1, mode="c"] maxs = np.empty(dim, dtype=np.float64) + RTreeError err + err = rtree_node_envelope(node.node, mins.data, maxs.data) + if err != Success: + raise RuntimeError("rtree_node_envelope failed") + return mins, maxs From 12de8f60fb7f48fd3274b61465b70d3134e9a52d Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 14:22:46 +0000 Subject: [PATCH 16/27] renaming and use `firedrake_rtree` package --- firedrake/cython/{rstar.pyx => rtree.pyx} | 2 +- .../cython/{rstarinc.pxi => rtreeinc.pxi} | 9 +++++-- firedrake/mesh.py | 16 +++++------ setup.py | 27 ++++++++++--------- 4 files changed, 29 insertions(+), 25 deletions(-) rename firedrake/cython/{rstar.pyx => rtree.pyx} (99%) rename firedrake/cython/{rstarinc.pxi => rtreeinc.pxi} (83%) diff --git a/firedrake/cython/rstar.pyx b/firedrake/cython/rtree.pyx similarity index 99% rename from firedrake/cython/rstar.pyx rename to firedrake/cython/rtree.pyx index 3d4d7045fb..e9121a13a9 100644 --- a/firedrake/cython/rstar.pyx +++ b/firedrake/cython/rtree.pyx @@ -11,7 +11,7 @@ cimport mpi4py.MPI as MPI from mpi4py.libmpi cimport MPI_INT from petsc4py.PETSc cimport CHKERR -include "rstarinc.pxi" +include "rtreeinc.pxi" include "petschdr.pxi" cdef class RTree(object): diff --git a/firedrake/cython/rstarinc.pxi b/firedrake/cython/rtreeinc.pxi similarity index 83% rename from firedrake/cython/rstarinc.pxi rename to firedrake/cython/rtreeinc.pxi index c5f30af919..dd8edeb50d 100644 --- a/firedrake/cython/rstarinc.pxi +++ b/firedrake/cython/rtreeinc.pxi @@ -2,12 +2,11 @@ from libc.stddef cimport size_t from libc.stdint cimport uint32_t -cdef extern from "rstar-capi.h": +cdef extern from "rtree-capi.h": ctypedef enum RTreeError: Success NullPointer InvalidDimension - NodeNotLeaf ctypedef struct RTreeH: pass @@ -26,6 +25,8 @@ cdef extern from "rstar-capi.h": RTreeError rtree_free(RTreeH *tree) + RTreeError rtree_free_ids(size_t *ids, size_t n) + RTreeError rtree_locate_all_at_point( const RTreeH *tree, const double *point, @@ -44,6 +45,10 @@ cdef extern from "rstar-capi.h": size_t *nchildren_out ) + RTreeError rtree_node_children_free(RTreeNodeH **children, size_t n) + + RTreeError rtree_node_free(RTreeNodeH *node) + RTreeError rtree_node_id( const RTreeNodeH *node, size_t *id_out diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 5d4bda5657..f06213cbb8 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -33,7 +33,7 @@ from firedrake.cython.dmcommon import DistributedMeshOverlapType import firedrake.cython.extrusion_numbering as extnum import firedrake.extrusion_utils as eutils -import firedrake.cython.rstar as rstar +import firedrake.cython.rtree as rtree import firedrake.utils as utils from firedrake.utils import as_cstr, IntType, RealType from firedrake.logging import info_red @@ -42,6 +42,7 @@ from firedrake.adjoint_utils import MeshGeometryMixin from firedrake.exceptions import VertexOnlyMeshMissingPointsError, NonUniqueMeshSequenceError import gem +import firedrake_rtree try: import netgen @@ -2607,7 +2608,7 @@ def spatial_index(self): coords_max = coords_mid + (tolerance + 0.5)*d with PETSc.Log.Event("spatial_index_build"): - self._spatial_index = rstar.build_from_aabb(coords_min, coords_max) + self._spatial_index = rtree.build_from_aabb(coords_min, coords_max) self._saved_coordinate_dat_version = self.coordinates.dat.dat_version return self._spatial_index @@ -2758,22 +2759,19 @@ def _c_locator(self, tolerance=None): }} """) - rstar_root = Path(__file__).resolve().parents[2] / "rstar" - rstar_include = rstar_root / "rstar-capi" / "include" - rstar_lib = rstar_root / "target" / "release" dll = compilation.load( src, "c", cppargs=[ f"-I{os.path.dirname(__file__)}", f"-I{sys.prefix}/include", - f"-I{rstar_include}" + f"-I{firedrake_rtree.get_include()}" ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=[ f"-L{sys.prefix}/lib", - f"-L{rstar_lib}", - "-lrstar_capi", + f"-L{firedrake_rtree.get_lib_filename()}", + "-lrtree_capi", f"-Wl,-rpath,{sys.prefix}/lib", - f"-Wl,-rpath,{rstar_lib}" + f"-Wl,-rpath,{firedrake_rtree.get_lib_filename()}" ], comm=self.comm ) diff --git a/setup.py b/setup.py index d786b256c9..0beddf972f 100644 --- a/setup.py +++ b/setup.py @@ -8,6 +8,7 @@ from pathlib import Path import libsupermesh +import firedrake_rtree import numpy as np import pybind11 import petsctools @@ -127,16 +128,16 @@ def __getitem__(self, key): # possible site package locations we can think of. sitepackage_dirs = site.getsitepackages() + [site.getusersitepackages()] -# rstar-capi -rstar_root = Path(__file__).resolve().parents[1] / "rstar" -rstar_capi_include = rstar_root / "rstar-capi" / "include" -rstar_capi_libdir = rstar_root / "target" / "release" -rstar_ = ExternalDependency( - include_dirs=[str(rstar_capi_include)], - library_dirs=[str(rstar_capi_libdir)], - libraries=["rstar_capi"], - runtime_library_dirs=[str(rstar_capi_libdir)], - extra_link_args=[f"-Wl,-rpath,{rstar_capi_libdir}"], +# firedrake_rtree +# The installed .so is named rtree_capi.cpython-XYZ-darwin.so (a maturin cdylib), +# so -lrtree_capi won't find it. Link by full path via extra_link_args instead. +rtree_ = ExternalDependency( + include_dirs=[firedrake_rtree.get_include()], + extra_link_args=[ + firedrake_rtree.get_lib_filename(), + f"-Wl,-rpath,{firedrake_rtree.get_lib()}", + ], + runtime_library_dirs=[firedrake_rtree.get_lib()], ) # libsupermesh @@ -196,10 +197,10 @@ def extensions(): )) # firedrake/cython/rstar.pyx: numpy, rstar-capi cython_list.append(Extension( - name="firedrake.cython.rstar", + name="firedrake.cython.rtree", language="c", - sources=[os.path.join("firedrake", "cython", "rstar.pyx")], - **(mpi_ + numpy_ + rstar_) + sources=[os.path.join("firedrake", "cython", "rtree.pyx")], + **(mpi_ + numpy_ + rtree_) )) # firedrake/cython/supermeshimpl.pyx: petsc, numpy, supermesh cython_list.append(Extension( From d46eac61fa79ae62fad6236bf5d9f813c35f234e Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 15:12:51 +0000 Subject: [PATCH 17/27] firedrake-rtree fixes --- firedrake/locate.c | 2 +- firedrake/mesh.py | 5 ++--- setup.py | 9 +++------ 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/firedrake/locate.c b/firedrake/locate.c index 4575ceb954..0ee7ac8639 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -1,6 +1,6 @@ #include #include -#include +#include #include #include diff --git a/firedrake/mesh.py b/firedrake/mesh.py index f06213cbb8..8fff3dfdce 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2768,10 +2768,9 @@ def _c_locator(self, tolerance=None): ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=[ f"-L{sys.prefix}/lib", - f"-L{firedrake_rtree.get_lib_filename()}", - "-lrtree_capi", + firedrake_rtree.get_lib_filename(), f"-Wl,-rpath,{sys.prefix}/lib", - f"-Wl,-rpath,{firedrake_rtree.get_lib_filename()}" + f"-Wl,-rpath,{firedrake_rtree.get_lib()}" ], comm=self.comm ) diff --git a/setup.py b/setup.py index 0beddf972f..2ba4f41a15 100644 --- a/setup.py +++ b/setup.py @@ -129,15 +129,12 @@ def __getitem__(self, key): sitepackage_dirs = site.getsitepackages() + [site.getusersitepackages()] # firedrake_rtree -# The installed .so is named rtree_capi.cpython-XYZ-darwin.so (a maturin cdylib), -# so -lrtree_capi won't find it. Link by full path via extra_link_args instead. rtree_ = ExternalDependency( include_dirs=[firedrake_rtree.get_include()], - extra_link_args=[ - firedrake_rtree.get_lib_filename(), - f"-Wl,-rpath,{firedrake_rtree.get_lib()}", + extra_link_args=[firedrake_rtree.get_lib_filename()], + runtime_library_dirs=[ + os.path.join(dir, "firedrake_rtree") for dir in sitepackage_dirs ], - runtime_library_dirs=[firedrake_rtree.get_lib()], ) # libsupermesh From 1c20e4552ffa62dff1f3bad6aa2ef8d2b3b4c7a6 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 15:13:08 +0000 Subject: [PATCH 18/27] remove unused stuff from rtree cython --- firedrake/cython/rtree.pyx | 229 ---------------------------------- firedrake/cython/rtreeinc.pxi | 31 ----- 2 files changed, 260 deletions(-) diff --git a/firedrake/cython/rtree.pyx b/firedrake/cython/rtree.pyx index e9121a13a9..7568a6a822 100644 --- a/firedrake/cython/rtree.pyx +++ b/firedrake/cython/rtree.pyx @@ -7,12 +7,7 @@ import cython from libc.stddef cimport size_t from libc.stdint cimport uintptr_t -cimport mpi4py.MPI as MPI -from mpi4py.libmpi cimport MPI_INT -from petsc4py.PETSc cimport CHKERR - include "rtreeinc.pxi" -include "petschdr.pxi" cdef class RTree(object): """Python class for holding a spatial index.""" @@ -86,227 +81,3 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min, raise RuntimeError("RTree_FromArray failed") return RTree(rtree) - -@cython.boundscheck(False) -@cython.wraparound(False) -def locate_all_at_point( - RTree rtree, - np.ndarray[np.float64_t, ndim=1, mode="c"] point): - """Return the ids of all leaves whose bounding box contains ``point``. - - Parameters - ---------- - rtree : RTree - The R*-tree to query. - point : (dim,) float64 array - The query point. Must have the same dimensionality as the tree. - - Returns - ------- - An array of integer ids corresponding to the leaves whose bounding boxes contain the query point. - The array must be freed by the caller using rtree_locate_all_at_point_free. - """ - cdef: - size_t *ids_out = NULL - size_t nids_out = 0 - RTreeError err - np.ndarray[np.npy_uintp, ndim=1, mode="c"] result - - err = rtree_locate_all_at_point( - rtree.tree, - point.data, - &ids_out, - &nids_out, - ) - if err != Success: - raise RuntimeError("rtree_locate_all_at_point failed") - - result = np.empty(nids_out, dtype=np.uintp) - for i in range(nids_out): - result[i] = ids_out[i] - rtree_free_ids(ids_out, nids_out) - return result - - -@cython.boundscheck(False) -@cython.wraparound(False) -def discover_ranks( - RTree rtree, - np.ndarray[np.float64_t, ndim=2, mode="c"] points, - MPI.Comm comm): - """Query a distributed Rtree and discover which ranks to send points, - and which ranks are going to send points to us. - - For each point in `points`, we find all candidate ranks whose bounding - boxes contain that point. We then use `PetscCommBuildTwoSided` to discover - which ranks will be sending points to us. - - Parameters - ---------- - rtree : RTree - The distributed Rtree built by :func:`build_from_aabb` with - rank numbers as leaf ids. - points : (n_points, gdim) float64 array - The local points to send to remote ranks. - comm : mpi4py.MPI.Comm - The MPI communicator. - - Returns - ------- - toranks : (nto,) int32 array - Target ranks to send points to. - send_offsets : (nto + 1,) int32 array - Points destined for `toranks[i]` are - `point_indices[send_offsets[i]:send_offsets[i+1]]`. - point_indices : (total_sends,) int32 array - Indices into `points` determining which points to send. - fromranks : (nfrom,) int32 array - Ranks that will send points to us. - recv_counts : (nfrom,) int32 array - Number of points we will receive from each rank in `fromranks`. - """ - cdef: - size_t *ids_out = NULL - size_t nids_out = 0 - RTreeError err - size_t i, j - size_t n_points = points.shape[0] - MPI.MPI_Comm mpi_comm = comm.ob_mpi - PetscMPIInt nto - PetscMPIInt nfrom = 0 - PetscMPIInt *fromranks = NULL - void *fromdata = NULL - np.ndarray[np.int32_t, ndim=1, mode="c"] toranks - np.ndarray[np.int32_t, ndim=1, mode="c"] send_counts - np.ndarray[np.int32_t, ndim=1, mode="c"] point_indices - np.ndarray[np.int32_t, ndim=1, mode="c"] send_offsets - np.ndarray[np.int32_t, ndim=1, mode="c"] fromranks_out - np.ndarray[np.int32_t, ndim=1, mode="c"] recv_counts_out - PetscMPIInt k - - # map dest_rank -> list of point indices to send there - rank_to_indices: dict[int, list[int]] = {} - for i in range(n_points): - err = rtree_locate_all_at_point( - rtree.tree, - &points[i, 0], - &ids_out, - &nids_out, - ) - if err != Success: - raise RuntimeError("rtree_locate_all_at_point failed") - - # Points may lie in multiple bounding boxes owned by the same rank - seen_ranks: set[int] = set() - for j in range(nids_out): - seen_ranks.add(ids_out[j]) - rtree_free_ids(ids_out, nids_out) - - ids_out = NULL - for dest_rank in seen_ranks: - if dest_rank in rank_to_indices: - rank_to_indices[dest_rank].append(i) - else: - rank_to_indices[dest_rank] = [i] - - nto = len(rank_to_indices) - toranks = np.empty(nto, dtype=np.int32) - send_counts = np.empty(nto, dtype=np.int32) - send_offsets = np.empty(nto + 1, dtype=np.int32) - all_indices: list[int] = [] - send_offsets[0] = 0 - for i, (rank, idx_list) in enumerate(rank_to_indices.items()): - toranks[i] = rank - send_counts[i] = len(idx_list) - send_offsets[i + 1] = send_offsets[i] + len(idx_list) - all_indices.extend(idx_list) - point_indices = np.array(all_indices, dtype=np.int32) - - # Routine that discovers communicating ranks given one-sided information - CHKERR(PetscCommBuildTwoSided( - mpi_comm, - 1, # sending/receiving one entry (the number of points) - MPI_INT, - nto, # number of ranks to send data to - toranks.data, # ranks to send to (array of length nto) - send_counts.data, # data to send to each rank (array of length nto) - &nfrom, # number of ranks we're receiving messages from - &fromranks, # ranks we're receiving messages from (array of length nfrom) - &fromdata, # data we're receiving from each rank (array of length nfrom) - )) - - # Copy petsc-allocated results into numpy arrays then free them - fromranks_out = np.empty(nfrom, dtype=np.int32) - recv_counts_out = np.empty(nfrom, dtype=np.int32) - for k in range(nfrom): - fromranks_out[k] = fromranks[k] - recv_counts_out[k] = (fromdata)[k] - CHKERR(PetscFree(fromranks)) - CHKERR(PetscFree(fromdata)) - - return toranks, send_offsets, point_indices, fromranks_out, recv_counts_out - - -cdef class RTreeNode(object): - """Python class for holding a spatial index node.""" - - cdef RTreeNodeH* node - - def __cinit__(self, uintptr_t node_handle): - self.node = 0 - if node_handle == 0: - raise RuntimeError("invalid node handle") - self.node = node_handle - - def __dealloc__(self): - if self.node != 0: - rtree_node_free(self.node) - self.node = 0 - - -def root_node(RTree rtree): - """Return the root node of the R*-tree.""" - cdef: - RTreeNodeH* node - RTreeError err - err = rtree_root_node(rtree.tree, &node) - if err != Success: - raise RuntimeError("rtree_root_node failed") - return RTreeNode(node) - - -def node_children(RTreeNode node): - """Return the children of an R*-tree node as a list of RStarTreeNode.""" - cdef: - RTreeNodeH** children - size_t nchildren - RTreeError err - err = rtree_node_children(node.node, &children, &nchildren) - if err != Success: - raise RuntimeError("rtree_node_children failed") - result = [RTreeNode(children[i]) for i in range(nchildren)] - rtree_node_children_free(children, nchildren) - return result - - -def node_id(RTreeNode node): - """Return the id of a leaf node.""" - cdef: - size_t id_out - RTreeError err - err = rtree_node_id(node.node, &id_out) - if err != Success: - raise RuntimeError("rtree_node_id failed (node may not be a leaf)") - return id_out - - -def node_envelope(RTreeNode node, size_t dim): - """Return the (mins, maxs) bounding envelope of an R*-tree node.""" - cdef: - np.ndarray[np.float64_t, ndim=1, mode="c"] mins = np.empty(dim, dtype=np.float64) - np.ndarray[np.float64_t, ndim=1, mode="c"] maxs = np.empty(dim, dtype=np.float64) - RTreeError err - err = rtree_node_envelope(node.node, mins.data, maxs.data) - if err != Success: - raise RuntimeError("rtree_node_envelope failed") - return mins, maxs diff --git a/firedrake/cython/rtreeinc.pxi b/firedrake/cython/rtreeinc.pxi index dd8edeb50d..7d1056bdd7 100644 --- a/firedrake/cython/rtreeinc.pxi +++ b/firedrake/cython/rtreeinc.pxi @@ -11,9 +11,6 @@ cdef extern from "rtree-capi.h": ctypedef struct RTreeH: pass - ctypedef struct RTreeNodeH: - pass - RTreeError rtree_bulk_load( RTreeH **tree, const double *mins, @@ -25,37 +22,9 @@ cdef extern from "rtree-capi.h": RTreeError rtree_free(RTreeH *tree) - RTreeError rtree_free_ids(size_t *ids, size_t n) - RTreeError rtree_locate_all_at_point( const RTreeH *tree, const double *point, size_t **ids_out, size_t *nids_out ) - - RTreeError rtree_root_node( - const RTreeH *tree, - RTreeNodeH **node - ) - - RTreeError rtree_node_children( - const RTreeNodeH *node, - RTreeNodeH ***children_out, - size_t *nchildren_out - ) - - RTreeError rtree_node_children_free(RTreeNodeH **children, size_t n) - - RTreeError rtree_node_free(RTreeNodeH *node) - - RTreeError rtree_node_id( - const RTreeNodeH *node, - size_t *id_out - ) - - RTreeError rtree_node_envelope( - const RTreeNodeH *node, - double *mins_out, - double *maxs_out - ) From 10dd698ae01657c5045eb07db767a20e6967532f Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 15:20:01 +0000 Subject: [PATCH 19/27] free ids in error case --- firedrake/locate.c | 1 + 1 file changed, 1 insertion(+) diff --git a/firedrake/locate.c b/firedrake/locate.c index 0ee7ac8639..cc9f940d08 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -36,6 +36,7 @@ int locate_cell(struct Function *f, err = rtree_locate_all_at_point((const struct RTreeH *)f->sidx, x, &ids, &nids); if (err != Success) { fputs("ERROR: RTree_LocateAllAtPoint failed in rstar-capi!\n", stderr); + rtree_free_ids(ids, nids); return -1; } if (f->extruded == 0) { From ef6dc6cd6943568482350eed574f856df4caed35 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 16:35:07 +0000 Subject: [PATCH 20/27] add firedrake-rtree dependency add to build-requirements.txt --- pyproject.toml | 2 ++ requirements-build.txt | 1 + 2 files changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 43796aba86..118edbec9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ # TODO RELEASE "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main", "h5py>3.12.1", + "firedrake-rtree", "immutabledict", "libsupermesh", "loopy>2024.1", @@ -154,6 +155,7 @@ docker = [ # Used in firedrake-vanilla container [build-system] requires = [ "Cython>=3.0", + "firedrake-rtree", "libsupermesh", "mpi4py>3; python_version >= '3.13'", "mpi4py; python_version < '3.13'", diff --git a/requirements-build.txt b/requirements-build.txt index 716b5e0651..2ae541122f 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -1,6 +1,7 @@ # Core build dependencies (adapted from pyproject.toml) Cython>=3.0 libsupermesh +firedrake-rtree mpi4py>3; python_version >= '3.13' mpi4py; python_version < '3.13' numpy From 796b2e948be730100a1a2bd5734f8e25968bdc34 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 17:13:27 +0000 Subject: [PATCH 21/27] remove unnecessary stuff from petschdr --- firedrake/cython/petschdr.pxi | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/firedrake/cython/petschdr.pxi b/firedrake/cython/petschdr.pxi index 1026473594..d7b1478ca0 100644 --- a/firedrake/cython/petschdr.pxi +++ b/firedrake/cython/petschdr.pxi @@ -8,7 +8,6 @@ cdef extern from "mpi-compat.h": cdef extern from "petsc.h": ctypedef long PetscInt - ctypedef int PetscMPIInt ctypedef double PetscReal ctypedef double PetscScalar ctypedef enum PetscBool: @@ -33,17 +32,6 @@ cdef extern from "petscsys.h" nogil: PetscErrorCode PetscFree(void*) PetscErrorCode PetscFree2(void*,void*) PetscErrorCode PetscSortIntWithArray(PetscInt,PetscInt[],PetscInt[]) - PetscErrorCode PetscCommBuildTwoSided( - MPI.MPI_Comm comm, - PetscMPIInt count, - MPI.MPI_Datatype dtype, - PetscMPIInt nto, - const PetscMPIInt *toranks, - const void *todata, - PetscMPIInt *nfrom, - PetscMPIInt **fromranks, - void *fromdata - ) cdef extern from "petscdmtypes.h" nogil: ctypedef enum PetscDMPolytopeType "DMPolytopeType": From c2ba66c38d3ba0061ab4ca2abdd03cb36e389970 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 17:31:29 +0000 Subject: [PATCH 22/27] fix function at --- firedrake/function.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/firedrake/function.py b/firedrake/function.py index b7e5b37550..3be9a18d66 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -1,4 +1,5 @@ import numpy as np +import firedrake_rtree import sys import ufl import warnings @@ -876,16 +877,13 @@ def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None): if ldargs is None: ldargs = [] - rstar_root = Path(__file__).resolve().parents[2] / "rstar" - rstar_include = rstar_root / "rstar-capi" / "include" - rstar_lib = rstar_root / "target" / "release" - ldargs += [f"-L{rstar_lib}", "-lrstar_capi", f"-Wl,-rpath,{rstar_lib}"] + 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{rstar_include}" + f"-I{firedrake_rtree.get_include()}" ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=ldargs, comm=function.comm From 35fc5631c9fbff9d59f71df3aaee7b27b274aef7 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 17:32:47 +0000 Subject: [PATCH 23/27] misc tidying one more --- firedrake/function.py | 1 - firedrake/mesh.py | 531 ++++++++++++++++++++---------------------- setup.py | 6 +- 3 files changed, 253 insertions(+), 285 deletions(-) diff --git a/firedrake/function.py b/firedrake/function.py index 3be9a18d66..5c3497e76c 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -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 diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 8fff3dfdce..4cef0ac979 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -16,6 +16,7 @@ import enum import numbers import abc +import firedrake_rtree from textwrap import dedent from pathlib import Path import typing @@ -36,13 +37,11 @@ import firedrake.cython.rtree as rtree import firedrake.utils as utils from firedrake.utils import as_cstr, IntType, RealType -from firedrake.logging import info_red from firedrake.parameters import parameters from firedrake.petsc import PETSc, DEFAULT_PARTITIONER from firedrake.adjoint_utils import MeshGeometryMixin from firedrake.exceptions import VertexOnlyMeshMissingPointsError, NonUniqueMeshSequenceError import gem -import firedrake_rtree try: import netgen @@ -2601,7 +2600,6 @@ def spatial_index(self): coords_max = coords_max.reshape(-1, 1) tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0 - coords_mid = (coords_max + coords_min)/2 d = np.max(coords_max - coords_min, axis=1)[:, None] coords_min = coords_mid - (tolerance + 0.5)*d @@ -4043,8 +4041,7 @@ def _dmswarm_create( _, coordsdim = coords.shape # Create a DMSWARM - with PETSc.Log.Event("FiredrakeDMSwarmCreate"): - swarm = FiredrakeDMSwarm().create(comm=comm) + swarm = FiredrakeDMSwarm().create(comm=comm) # save the fields on the swarm swarm.fields = default_fields + default_extra_fields + other_fields @@ -4097,34 +4094,33 @@ def _dmswarm_create( # topological dimension of the base mesh. # NOTE ensure that swarm.restoreField is called for each field too! - with PETSc.Log.Event("FiredrakeDMSwarmSetFields"): - swarm_coords = swarm.getField("DMSwarmPIC_coor").reshape((num_vertices, gdim)) - cell_id_name = swarm.getCellDMActive().getCellID() - swarm_parent_cell_nums = swarm.getField(cell_id_name).ravel() - field_parent_cell_nums = swarm.getField("parentcellnum").ravel() - field_reference_coords = swarm.getField("refcoord").reshape((num_vertices, tdim)) - field_global_index = swarm.getField("globalindex").ravel() - field_rank = swarm.getField("DMSwarm_rank").ravel() - field_input_rank = swarm.getField("inputrank").ravel() - field_input_index = swarm.getField("inputindex").ravel() - swarm_coords[...] = coords - swarm_parent_cell_nums[...] = plex_parent_cell_nums - field_parent_cell_nums[...] = parent_cell_nums - field_reference_coords[...] = reference_coords - field_global_index[...] = coords_idxs - field_rank[...] = ranks - field_input_rank[...] = input_ranks - field_input_index[...] = input_coords_idxs - - # have to restore fields once accessed to allow access again - swarm.restoreField("inputindex") - swarm.restoreField("inputrank") - swarm.restoreField("DMSwarm_rank") - swarm.restoreField("globalindex") - swarm.restoreField("refcoord") - swarm.restoreField("parentcellnum") - swarm.restoreField("DMSwarmPIC_coor") - swarm.restoreField(cell_id_name) + swarm_coords = swarm.getField("DMSwarmPIC_coor").reshape((num_vertices, gdim)) + cell_id_name = swarm.getCellDMActive().getCellID() + swarm_parent_cell_nums = swarm.getField(cell_id_name).ravel() + field_parent_cell_nums = swarm.getField("parentcellnum").ravel() + field_reference_coords = swarm.getField("refcoord").reshape((num_vertices, tdim)) + field_global_index = swarm.getField("globalindex").ravel() + field_rank = swarm.getField("DMSwarm_rank").ravel() + field_input_rank = swarm.getField("inputrank").ravel() + field_input_index = swarm.getField("inputindex").ravel() + swarm_coords[...] = coords + swarm_parent_cell_nums[...] = plex_parent_cell_nums + field_parent_cell_nums[...] = parent_cell_nums + field_reference_coords[...] = reference_coords + field_global_index[...] = coords_idxs + field_rank[...] = ranks + field_input_rank[...] = input_ranks + field_input_index[...] = input_coords_idxs + + # have to restore fields once accessed to allow access again + swarm.restoreField("inputindex") + swarm.restoreField("inputrank") + swarm.restoreField("DMSwarm_rank") + swarm.restoreField("globalindex") + swarm.restoreField("refcoord") + swarm.restoreField("parentcellnum") + swarm.restoreField("DMSwarmPIC_coor") + swarm.restoreField(cell_id_name) if extruded: field_base_parent_cell_nums = swarm.getField("parentcellbasenum").ravel() @@ -4264,65 +4260,58 @@ def _parent_mesh_embedding( "VertexOnlyMeshes don't have a working locate_cells_ref_coords_and_dists method" ) - with PETSc.Log.Event("pm_embed_initial"): - with temp_internal_comm(parent_mesh.comm) as icomm: - # In parallel, we need to make sure we know which point is which and save - # it. - if redundant: - # rank 0 broadcasts coords to all ranks - with PETSc.Log.Event("pm_embed_coords_bcast"): - coords_local = icomm.bcast(coords, root=0) - ncoords_local = coords_local.shape[0] - coords_global = coords_local - ncoords_global = coords_global.shape[0] - global_idxs_global = np.arange(coords_global.shape[0]) - input_coords_idxs_local = np.arange(ncoords_local) - input_coords_idxs_global = input_coords_idxs_local - input_ranks_local = np.zeros(ncoords_local, dtype=int) - input_ranks_global = input_ranks_local - else: - # Here, we have to assume that all points we can see are unique. - # We therefore gather all points on all ranks in rank order: if rank 0 - # has 10 points, rank 1 has 20 points, and rank 3 has 5 points, then - # rank 0's points have global numbering 0-9, rank 1's points have - # global numbering 10-29, and rank 3's points have global numbering - # 30-34. - coords_local = coords - ncoords_local = coords.shape[0] - with PETSc.Log.Event("pm_embed_ncoords_allgather"): - ncoords_local_allranks = icomm.allgather(ncoords_local) - ncoords_global = sum(ncoords_local_allranks) - # The below code looks complicated but it's just an allgather of the - # (variable length) coords_local array such that they are concatenated. - coords_local_size = np.array(coords_local.size) - coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) - with PETSc.Log.Event("pm_embed_coord_sizes_allgatherv"): - icomm.Allgatherv(coords_local_size, coords_local_sizes) - coords_global = np.empty( - (ncoords_global, coords.shape[1]), dtype=coords_local.dtype - ) - with PETSc.Log.Event("pm_embed_coords_allgatherv"): - icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) - # # ncoords_local_allranks is in rank order so we can just sum up the - # # previous ranks to get the starting index for the global numbering. - # # For rank 0 we make use of the fact that sum([]) = 0. - # startidx = sum(ncoords_local_allranks[:parent_mesh.comm.rank]) - # endidx = startidx + ncoords_local - # global_idxs_global = np.arange(startidx, endidx) - global_idxs_global = np.arange(coords_global.shape[0]) - input_coords_idxs_local = np.arange(ncoords_local) - input_coords_idxs_global = np.empty(ncoords_global, dtype=int) - with PETSc.Log.Event("pm_embed_inputidx_allgatherv"): - icomm.Allgatherv( - input_coords_idxs_local, - (input_coords_idxs_global, ncoords_local_allranks), - ) - input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) - input_ranks_global = np.empty(ncoords_global, dtype=int) - with PETSc.Log.Event("pm_embed_inputranks_allgatherv"): - icomm.Allgatherv( - input_ranks_local, (input_ranks_global, ncoords_local_allranks) - ) + with temp_internal_comm(parent_mesh.comm) as icomm: + # In parallel, we need to make sure we know which point is which and save + # it. + if redundant: + # rank 0 broadcasts coords to all ranks + coords_local = icomm.bcast(coords, root=0) + ncoords_local = coords_local.shape[0] + coords_global = coords_local + ncoords_global = coords_global.shape[0] + global_idxs_global = np.arange(coords_global.shape[0]) + input_coords_idxs_local = np.arange(ncoords_local) + input_coords_idxs_global = input_coords_idxs_local + input_ranks_local = np.zeros(ncoords_local, dtype=int) + input_ranks_global = input_ranks_local + else: + # Here, we have to assume that all points we can see are unique. + # We therefore gather all points on all ranks in rank order: if rank 0 + # has 10 points, rank 1 has 20 points, and rank 3 has 5 points, then + # rank 0's points have global numbering 0-9, rank 1's points have + # global numbering 10-29, and rank 3's points have global numbering + # 30-34. + coords_local = coords + ncoords_local = coords.shape[0] + ncoords_local_allranks = icomm.allgather(ncoords_local) + ncoords_global = sum(ncoords_local_allranks) + # The below code looks complicated but it's just an allgather of the + # (variable length) coords_local array such that they are concatenated. + coords_local_size = np.array(coords_local.size) + coords_local_sizes = np.empty(parent_mesh.comm.size, dtype=int) + icomm.Allgatherv(coords_local_size, coords_local_sizes) + coords_global = np.empty( + (ncoords_global, coords.shape[1]), dtype=coords_local.dtype + ) + icomm.Allgatherv(coords_local, (coords_global, coords_local_sizes)) + # # ncoords_local_allranks is in rank order so we can just sum up the + # # previous ranks to get the starting index for the global numbering. + # # For rank 0 we make use of the fact that sum([]) = 0. + # startidx = sum(ncoords_local_allranks[:parent_mesh.comm.rank]) + # endidx = startidx + ncoords_local + # global_idxs_global = np.arange(startidx, endidx) + global_idxs_global = np.arange(coords_global.shape[0]) + input_coords_idxs_local = np.arange(ncoords_local) + input_coords_idxs_global = np.empty(ncoords_global, dtype=int) + icomm.Allgatherv( + input_coords_idxs_local, + (input_coords_idxs_global, ncoords_local_allranks), + ) + input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) + input_ranks_global = np.empty(ncoords_global, dtype=int) + icomm.Allgatherv( + input_ranks_local, (input_ranks_global, ncoords_local_allranks) + ) ( parent_cell_nums, reference_coords, @@ -4342,19 +4331,17 @@ def _parent_mesh_embedding( visible_ranks[:parent_mesh.cell_set.size] = parent_mesh.comm.rank visible_ranks[parent_mesh.cell_set.size:] = -1 # Halo exchange the visible ranks so that each rank knows which ranks can see each cell. - with PETSc.Log.Event("pm_embed_exchange_cell_orientations"): - dmcommon.exchange_cell_orientations( - parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks - ) - with PETSc.Log.Event("pm_embed_build_rank_arrays"): - locally_visible = parent_cell_nums != -1 + dmcommon.exchange_cell_orientations( + parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks + ) + locally_visible = parent_cell_nums != -1 - if parent_mesh.extruded: - # Halo exchange of visible_ranks is over the base mesh topology and cell numbering, - # so we need to map back to extruded cell numbering after indexing parent_cell_nums. - locally_visible_cell_nums = parent_cell_nums[locally_visible] // (parent_mesh.layers - 1) - else: - locally_visible_cell_nums = parent_cell_nums[locally_visible] + if parent_mesh.extruded: + # Halo exchange of visible_ranks is over the base mesh topology and cell numbering, + # so we need to map back to extruded cell numbering after indexing parent_cell_nums. + locally_visible_cell_nums = parent_cell_nums[locally_visible] // (parent_mesh.layers - 1) + else: + locally_visible_cell_nums = parent_cell_nums[locally_visible] # In parallel there will regularly be disagreements about which cell owns a # point when those points are close to mesh partition boundaries. @@ -4376,6 +4363,7 @@ def _parent_mesh_embedding( rank_candidates = np.where(ref_cell_dists_l1 == owned_ref_cell_dists_l1, ranks, -np.inf) owned_ranks = np.empty_like(rank_candidates) parent_mesh.comm.Allreduce(rank_candidates, owned_ranks, op=MPI.MAX) + changed_ref_cell_dists_l1 = owned_ref_cell_dists_l1 != ref_cell_dists_l1 changed_ranks = owned_ranks != ranks @@ -4389,54 +4377,52 @@ def _parent_mesh_embedding( # the distance is the same then we have found the correct cell. If we # cannot make a match to owned_rank and distance then we can't see the # point. - with PETSc.Log.Event("pm_embed_tied_ranks"): - changed_ranks_tied = changed_ranks & ~changed_ref_cell_dists_l1 - if any(changed_ranks_tied): - cells_ignore_T = np.asarray([np.copy(parent_cell_nums)]) - while any(changed_ranks_tied): - with PETSc.Log.Event("pm_embed_tied_rank_locate_cells"): - ( - parent_cell_nums[changed_ranks_tied], - new_reference_coords, - ref_cell_dists_l1[changed_ranks_tied], - ) = parent_mesh.locate_cells_ref_coords_and_dists( - coords_global[changed_ranks_tied], - tolerance, - cells_ignore=cells_ignore_T.T[changed_ranks_tied, :], - ) - # delete extra dimension if necessary - if parent_mesh.geometric_dimension > parent_mesh.topological_dimension: - new_reference_coords = new_reference_coords[:, : parent_mesh.topological_dimension] - reference_coords[changed_ranks_tied, :] = new_reference_coords - # remove newly lost points - locally_visible[changed_ranks_tied] = ( - parent_cell_nums[changed_ranks_tied] != -1 - ) - changed_ranks_tied &= locally_visible - # if new ref_cell_dists_l1 > owned_ref_cell_dists_l1 then we should - # disregard the point. - locally_visible[changed_ranks_tied] &= ( - ref_cell_dists_l1[changed_ranks_tied] - <= owned_ref_cell_dists_l1[changed_ranks_tied] - ) - changed_ranks_tied &= locally_visible - # update the identified rank - if parent_mesh.extruded: - _retry_cell_nums = parent_cell_nums[changed_ranks_tied] // (parent_mesh.layers - 1) - else: - _retry_cell_nums = parent_cell_nums[changed_ranks_tied] - ranks[changed_ranks_tied] = visible_ranks[_retry_cell_nums] - # if the rank now matches then we have found the correct cell - locally_visible[changed_ranks_tied] &= ( - owned_ranks[changed_ranks_tied] == ranks[changed_ranks_tied] - ) - # remove these rank matches from changed_ranks_tied - changed_ranks_tied &= ~locally_visible - # add more cells to ignore - cells_ignore_T = np.vstack(( - cells_ignore_T, - parent_cell_nums) - ) + changed_ranks_tied = changed_ranks & ~changed_ref_cell_dists_l1 + if any(changed_ranks_tied): + cells_ignore_T = np.asarray([np.copy(parent_cell_nums)]) + while any(changed_ranks_tied): + ( + parent_cell_nums[changed_ranks_tied], + new_reference_coords, + ref_cell_dists_l1[changed_ranks_tied], + ) = parent_mesh.locate_cells_ref_coords_and_dists( + coords_global[changed_ranks_tied], + tolerance, + cells_ignore=cells_ignore_T.T[changed_ranks_tied, :], + ) + # delete extra dimension if necessary + if parent_mesh.geometric_dimension > parent_mesh.topological_dimension: + new_reference_coords = new_reference_coords[:, : parent_mesh.topological_dimension] + reference_coords[changed_ranks_tied, :] = new_reference_coords + # remove newly lost points + locally_visible[changed_ranks_tied] = ( + parent_cell_nums[changed_ranks_tied] != -1 + ) + changed_ranks_tied &= locally_visible + # if new ref_cell_dists_l1 > owned_ref_cell_dists_l1 then we should + # disregard the point. + locally_visible[changed_ranks_tied] &= ( + ref_cell_dists_l1[changed_ranks_tied] + <= owned_ref_cell_dists_l1[changed_ranks_tied] + ) + changed_ranks_tied &= locally_visible + # update the identified rank + if parent_mesh.extruded: + _retry_cell_nums = parent_cell_nums[changed_ranks_tied] // (parent_mesh.layers - 1) + else: + _retry_cell_nums = parent_cell_nums[changed_ranks_tied] + ranks[changed_ranks_tied] = visible_ranks[_retry_cell_nums] + # if the rank now matches then we have found the correct cell + locally_visible[changed_ranks_tied] &= ( + owned_ranks[changed_ranks_tied] == ranks[changed_ranks_tied] + ) + # remove these rank matches from changed_ranks_tied + changed_ranks_tied &= ~locally_visible + # add more cells to ignore + cells_ignore_T = np.vstack(( + cells_ignore_T, + parent_cell_nums) + ) # Any ranks which are still -np.inf are not in the mesh missing_global_idxs = np.where(owned_ranks == -np.inf)[0] @@ -4450,22 +4436,20 @@ def _parent_mesh_embedding( reference_coords[missing_coords_idxs_on_rank, :] = np.nan owned_ranks[missing_coords_idxs_on_rank] = parent_mesh.comm.size + 1 - with PETSc.Log.Event("pm_embed_exclude_halos_parallel"): - if exclude_halos and parent_mesh.comm.size > 1: - off_rank_coords_idxs = np.where( - (owned_ranks != parent_mesh.comm.rank) - & (owned_ranks != parent_mesh.comm.size + 1) - )[0] - locally_visible[off_rank_coords_idxs] = False - - with PETSc.Log.Event("pm_embed_compress_arrays"): - coords_embedded = np.compress(locally_visible, coords_global, axis=0) - global_idxs = np.compress(locally_visible, global_idxs_global, axis=0) - reference_coords = np.compress(locally_visible, reference_coords, axis=0) - parent_cell_nums = np.compress(locally_visible, parent_cell_nums, axis=0) - owned_ranks = np.compress(locally_visible, owned_ranks, axis=0).astype(int) - input_ranks = np.compress(locally_visible, input_ranks_global, axis=0) - input_coords_idxs = np.compress(locally_visible, input_coords_idxs_global, axis=0) + if exclude_halos and parent_mesh.comm.size > 1: + off_rank_coords_idxs = np.where( + (owned_ranks != parent_mesh.comm.rank) + & (owned_ranks != parent_mesh.comm.size + 1) + )[0] + locally_visible[off_rank_coords_idxs] = False + + coords_embedded = np.compress(locally_visible, coords_global, axis=0) + global_idxs = np.compress(locally_visible, global_idxs_global, axis=0) + reference_coords = np.compress(locally_visible, reference_coords, axis=0) + parent_cell_nums = np.compress(locally_visible, parent_cell_nums, axis=0) + owned_ranks = np.compress(locally_visible, owned_ranks, axis=0).astype(int) + input_ranks = np.compress(locally_visible, input_ranks_global, axis=0) + input_coords_idxs = np.compress(locally_visible, input_coords_idxs_global, axis=0) return ( coords_embedded, @@ -4505,157 +4489,142 @@ def _swarm_original_ordering_preserve( # Gather everything except original_ordering_coords_local from all mpi # ranks - with PETSc.Log.Event("original_ordering_allgather_ncoords_local"): - ncoords_local_allranks = comm.allgather(ncoords_local) + ncoords_local_allranks = comm.allgather(ncoords_local) ncoords_global = sum(ncoords_local_allranks) parent_cell_nums_global = np.empty( ncoords_global, dtype=parent_cell_nums_local.dtype ) - with PETSc.Log.Event("original_ordering_allgatherv_parent_cell_nums"): - comm.Allgatherv( - parent_cell_nums_local, (parent_cell_nums_global, ncoords_local_allranks) - ) + comm.Allgatherv( + parent_cell_nums_local, (parent_cell_nums_global, ncoords_local_allranks) + ) plex_parent_cell_nums_global = np.empty( ncoords_global, dtype=plex_parent_cell_nums_local.dtype ) - with PETSc.Log.Event("original_ordering_allgatherv_plex_parent_cell_nums"): - comm.Allgatherv( - plex_parent_cell_nums_local, - (plex_parent_cell_nums_global, ncoords_local_allranks), - ) + comm.Allgatherv( + plex_parent_cell_nums_local, + (plex_parent_cell_nums_global, ncoords_local_allranks), + ) reference_coords_local_size = np.array(reference_coords_local.size) reference_coords_local_sizes = np.empty(comm.size, dtype=int) - with PETSc.Log.Event("original_ordering_allgatherv_reference_coords_sizes"): - comm.Allgatherv(reference_coords_local_size, reference_coords_local_sizes) + comm.Allgatherv(reference_coords_local_size, reference_coords_local_sizes) reference_coords_global = np.empty( (ncoords_global, reference_coords_local.shape[1]), dtype=reference_coords_local.dtype, ) - with PETSc.Log.Event("original_ordering_allgatherv_reference_coords"): - comm.Allgatherv( - reference_coords_local, (reference_coords_global, reference_coords_local_sizes) - ) + comm.Allgatherv( + reference_coords_local, (reference_coords_global, reference_coords_local_sizes) + ) global_idxs_global = np.empty(ncoords_global, dtype=global_idxs_local.dtype) - with PETSc.Log.Event("original_ordering_allgatherv_global_idxs"): - comm.Allgatherv(global_idxs_local, (global_idxs_global, ncoords_local_allranks)) + comm.Allgatherv(global_idxs_local, (global_idxs_global, ncoords_local_allranks)) ranks_global = np.empty(ncoords_global, dtype=ranks_local.dtype) - with PETSc.Log.Event("original_ordering_allgatherv_ranks"): - comm.Allgatherv(ranks_local, (ranks_global, ncoords_local_allranks)) + comm.Allgatherv(ranks_local, (ranks_global, ncoords_local_allranks)) input_ranks_global = np.empty(ncoords_global, dtype=input_ranks_local.dtype) - with PETSc.Log.Event("original_ordering_allgatherv_input_ranks"): - comm.Allgatherv(input_ranks_local, (input_ranks_global, ncoords_local_allranks)) + comm.Allgatherv(input_ranks_local, (input_ranks_global, ncoords_local_allranks)) input_idxs_global = np.empty(ncoords_global, dtype=input_idxs_local.dtype) - with PETSc.Log.Event("original_ordering_allgatherv_input_idxs"): - comm.Allgatherv(input_idxs_local, (input_idxs_global, ncoords_local_allranks)) + comm.Allgatherv(input_idxs_local, (input_idxs_global, ncoords_local_allranks)) # Sort by global index, which will be in rank order (they probably already # are but we can't rely on that) - with PETSc.Log.Event("original_ordering_sort_by_global_idx"): - global_idxs_global_order = np.argsort(global_idxs_global) - with PETSc.Log.Event("original_ordering_reindex_sorted"): - sorted_parent_cell_nums_global = parent_cell_nums_global[global_idxs_global_order] - sorted_plex_parent_cell_nums_global = plex_parent_cell_nums_global[ - global_idxs_global_order - ] - sorted_reference_coords_global = reference_coords_global[ - global_idxs_global_order, : - ] - sorted_global_idxs_global = global_idxs_global[global_idxs_global_order] - sorted_ranks_global = ranks_global[global_idxs_global_order] - sorted_input_ranks_global = input_ranks_global[global_idxs_global_order] - sorted_input_idxs_global = input_idxs_global[global_idxs_global_order] + global_idxs_global_order = np.argsort(global_idxs_global) + sorted_parent_cell_nums_global = parent_cell_nums_global[global_idxs_global_order] + sorted_plex_parent_cell_nums_global = plex_parent_cell_nums_global[ + global_idxs_global_order + ] + sorted_reference_coords_global = reference_coords_global[ + global_idxs_global_order, : + ] + sorted_global_idxs_global = global_idxs_global[global_idxs_global_order] + sorted_ranks_global = ranks_global[global_idxs_global_order] + sorted_input_ranks_global = input_ranks_global[global_idxs_global_order] + sorted_input_idxs_global = input_idxs_global[global_idxs_global_order] # Check order is correct - we can probably remove this eventually since it's # quite expensive - with PETSc.Log.Event("original_ordering_check_global_idx_order"): - if not np.all(sorted_input_ranks_global[1:] >= sorted_input_ranks_global[:-1]): - raise ValueError("Global indexing has not ordered the ranks as expected") + if not np.all(sorted_input_ranks_global[1:] >= sorted_input_ranks_global[:-1]): + raise ValueError("Global indexing has not ordered the ranks as expected") # get rid of any duplicated global indices (i.e. points in halos) - with PETSc.Log.Event("original_ordering_remove_duplicate_global_idxs"): - unique_global_idxs, unique_idxs = np.unique( - sorted_global_idxs_global, return_index=True - ) - with PETSc.Log.Event("original_ordering_filter_and_select_outputs"): - unique_parent_cell_nums_global = sorted_parent_cell_nums_global[unique_idxs] - unique_plex_parent_cell_nums_global = sorted_plex_parent_cell_nums_global[ - unique_idxs - ] - unique_reference_coords_global = sorted_reference_coords_global[unique_idxs, :] - unique_ranks_global = sorted_ranks_global[unique_idxs] - unique_input_ranks_global = sorted_input_ranks_global[unique_idxs] - unique_input_idxs_global = sorted_input_idxs_global[unique_idxs] - - # save the points on this rank which match the input rank ready for output - input_ranks_match = unique_input_ranks_global == comm.rank - output_global_idxs = unique_global_idxs[input_ranks_match] - output_parent_cell_nums = unique_parent_cell_nums_global[input_ranks_match] - output_plex_parent_cell_nums = unique_plex_parent_cell_nums_global[ - input_ranks_match - ] - output_reference_coords = unique_reference_coords_global[input_ranks_match, :] - output_ranks = unique_ranks_global[input_ranks_match] - output_input_ranks = unique_input_ranks_global[input_ranks_match] - output_input_idxs = unique_input_idxs_global[input_ranks_match] - if extruded: - ( - output_base_parent_cell_nums, - output_extrusion_heights, - ) = _parent_extrusion_numbering(output_parent_cell_nums, layers) - else: - output_base_parent_cell_nums = None - output_extrusion_heights = None + unique_global_idxs, unique_idxs = np.unique( + sorted_global_idxs_global, return_index=True + ) + unique_parent_cell_nums_global = sorted_parent_cell_nums_global[unique_idxs] + unique_plex_parent_cell_nums_global = sorted_plex_parent_cell_nums_global[ + unique_idxs + ] + unique_reference_coords_global = sorted_reference_coords_global[unique_idxs, :] + unique_ranks_global = sorted_ranks_global[unique_idxs] + unique_input_ranks_global = sorted_input_ranks_global[unique_idxs] + unique_input_idxs_global = sorted_input_idxs_global[unique_idxs] + + # save the points on this rank which match the input rank ready for output + input_ranks_match = unique_input_ranks_global == comm.rank + output_global_idxs = unique_global_idxs[input_ranks_match] + output_parent_cell_nums = unique_parent_cell_nums_global[input_ranks_match] + output_plex_parent_cell_nums = unique_plex_parent_cell_nums_global[ + input_ranks_match + ] + output_reference_coords = unique_reference_coords_global[input_ranks_match, :] + output_ranks = unique_ranks_global[input_ranks_match] + output_input_ranks = unique_input_ranks_global[input_ranks_match] + output_input_idxs = unique_input_idxs_global[input_ranks_match] + if extruded: + ( + output_base_parent_cell_nums, + output_extrusion_heights, + ) = _parent_extrusion_numbering(output_parent_cell_nums, layers) + else: + output_base_parent_cell_nums = None + output_extrusion_heights = None # check if the input indices are in order from zero - this can also probably # be removed eventually because, again, it's expensive. - with PETSc.Log.Event("original_ordering_check_input_idx_order"): - if not np.array_equal(output_input_idxs, np.arange(output_input_idxs.size)): - raise ValueError( - "Global indexing has not ordered the input indices as expected." - ) - if len(output_global_idxs) != len(original_ordering_coords_local): - raise ValueError( - "The number of local global indices which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_parent_cell_nums) != len(original_ordering_coords_local): - raise ValueError( - "The number of local parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_plex_parent_cell_nums) != len(original_ordering_coords_local): - raise ValueError( - "The number of local plex parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_reference_coords) != len(original_ordering_coords_local): - raise ValueError( - "The number of local reference coordinates which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_ranks) != len(original_ordering_coords_local): - raise ValueError( - "The number of local rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_input_ranks) != len(original_ordering_coords_local): + if not np.array_equal(output_input_idxs, np.arange(output_input_idxs.size)): + raise ValueError( + "Global indexing has not ordered the input indices as expected." + ) + if len(output_global_idxs) != len(original_ordering_coords_local): + raise ValueError( + "The number of local global indices which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_parent_cell_nums) != len(original_ordering_coords_local): + raise ValueError( + "The number of local parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_plex_parent_cell_nums) != len(original_ordering_coords_local): + raise ValueError( + "The number of local plex parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_reference_coords) != len(original_ordering_coords_local): + raise ValueError( + "The number of local reference coordinates which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_ranks) != len(original_ordering_coords_local): + raise ValueError( + "The number of local rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_input_ranks) != len(original_ordering_coords_local): + raise ValueError( + "The number of local input rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if len(output_input_idxs) != len(original_ordering_coords_local): + raise ValueError( + "The number of local input indices which will be used to make the swarm do not match the input number of original ordering coordinates." + ) + if extruded: + if len(output_base_parent_cell_nums) != len(original_ordering_coords_local): raise ValueError( - "The number of local input rank numbers which will be used to make the swarm do not match the input number of original ordering coordinates." + "The number of local base parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." ) - if len(output_input_idxs) != len(original_ordering_coords_local): + if len(output_extrusion_heights) != len(original_ordering_coords_local): raise ValueError( - "The number of local input indices which will be used to make the swarm do not match the input number of original ordering coordinates." + "The number of local extrusion heights which will be used to make the swarm do not match the input number of original ordering coordinates." ) - if extruded: - if len(output_base_parent_cell_nums) != len(original_ordering_coords_local): - raise ValueError( - "The number of local base parent cell numbers which will be used to make the swarm do not match the input number of original ordering coordinates." - ) - if len(output_extrusion_heights) != len(original_ordering_coords_local): - raise ValueError( - "The number of local extrusion heights which will be used to make the swarm do not match the input number of original ordering coordinates." - ) return _dmswarm_create( [], diff --git a/setup.py b/setup.py index 2ba4f41a15..ac1d006056 100644 --- a/setup.py +++ b/setup.py @@ -120,9 +120,9 @@ def __getitem__(self, key): # Set the library name and hope for the best hdf5_ = ExternalDependency(libraries=["hdf5"]) -# When we link against rstar-capi or libsupermesh we need to know where +# When we link against firedrake-rtree or libsupermesh we need to know where # the '.so' files end up. Since installation happens in an isolated -# environment we cannot simply query the currently loaded libraries for the +# environment we cannot simply query firedrake-rtree and libsupermesh for the # current paths as they will not be valid once the installation is complete. # Therefore we set the runtime library search path to all the different # possible site package locations we can think of. @@ -192,7 +192,7 @@ def extensions(): sources=[os.path.join("firedrake", "cython", "patchimpl.pyx")], **(mpi_ + petsc_ + numpy_) )) - # firedrake/cython/rstar.pyx: numpy, rstar-capi + # firedrake/cython/rstar.pyx: numpy, rtree-capi cython_list.append(Extension( name="firedrake.cython.rtree", language="c", From fcc7a2f0e0011850212f5f32f5846a188e21db59 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 18:29:27 +0000 Subject: [PATCH 24/27] fixes --- firedrake/cython/rtree.pyx | 4 ++-- firedrake/mesh.py | 3 +-- setup.py | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/firedrake/cython/rtree.pyx b/firedrake/cython/rtree.pyx index 7568a6a822..6997cc5771 100644 --- a/firedrake/cython/rtree.pyx +++ b/firedrake/cython/rtree.pyx @@ -51,7 +51,7 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min, Returns ------- RTree - An RTree object containing the built R*-tree. + An RTree object containing the Rtree. """ cdef: RTreeH* rtree @@ -78,6 +78,6 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min, dim ) if err != Success: - raise RuntimeError("RTree_FromArray failed") + raise RuntimeError("rtree_bulk_load failed") return RTree(rtree) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 4cef0ac979..ed57bf3b36 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4304,8 +4304,7 @@ def _parent_mesh_embedding( input_coords_idxs_local = np.arange(ncoords_local) input_coords_idxs_global = np.empty(ncoords_global, dtype=int) icomm.Allgatherv( - input_coords_idxs_local, - (input_coords_idxs_global, ncoords_local_allranks), + input_coords_idxs_local, (input_coords_idxs_global, ncoords_local_allranks), ) input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) input_ranks_global = np.empty(ncoords_global, dtype=int) diff --git a/setup.py b/setup.py index ac1d006056..773b2e6af5 100644 --- a/setup.py +++ b/setup.py @@ -192,7 +192,7 @@ def extensions(): sources=[os.path.join("firedrake", "cython", "patchimpl.pyx")], **(mpi_ + petsc_ + numpy_) )) - # firedrake/cython/rstar.pyx: numpy, rtree-capi + # firedrake/cython/rtree.pyx: numpy, rtree-capi cython_list.append(Extension( name="firedrake.cython.rtree", language="c", From 72948ac397d23f109d9ff96b60e7995a11d12860 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 19:53:24 +0000 Subject: [PATCH 25/27] fixup mesh.py --- firedrake/mesh.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index ed57bf3b36..ebd9dabc04 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2267,7 +2267,6 @@ def input_ordering(self): ) @staticmethod - @PETSc.Log.EventDecorator() def _make_input_ordering_sf(swarm, nroots, ilocal): # ilocal = None -> leaves are swarm points [0, 1, 2, ...). # ilocal can also be Firedrake cell numbers. @@ -2566,8 +2565,7 @@ def spatial_index(self): Returns ------- - :class:`~.spatialindex.SpatialIndex` or None if the mesh is - one-dimensional. + :class:`~.rtree.RTree` Notes ----- @@ -4304,7 +4302,7 @@ def _parent_mesh_embedding( input_coords_idxs_local = np.arange(ncoords_local) input_coords_idxs_global = np.empty(ncoords_global, dtype=int) icomm.Allgatherv( - input_coords_idxs_local, (input_coords_idxs_global, ncoords_local_allranks), + input_coords_idxs_local, (input_coords_idxs_global, ncoords_local_allranks) ) input_ranks_local = np.full(ncoords_local, icomm.rank, dtype=int) input_ranks_global = np.empty(ncoords_global, dtype=int) From a29c6612f9c0f7960934eeff91469bcb91a70435 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Mar 2026 19:54:08 +0000 Subject: [PATCH 26/27] add rtree libspatialindex back in --- setup.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/setup.py b/setup.py index 773b2e6af5..635b6fcecb 100644 --- a/setup.py +++ b/setup.py @@ -12,6 +12,7 @@ import numpy as np import pybind11 import petsctools +import rtree from Cython.Build import cythonize from setuptools import setup, find_packages, Extension from setuptools.command.editable_wheel import editable_wheel as _editable_wheel @@ -128,6 +129,19 @@ def __getitem__(self, key): # possible site package locations we can think of. sitepackage_dirs = site.getsitepackages() + [site.getusersitepackages()] +# libspatialindex +# example: +# gcc -I/rtree/include +# gcc /rtree.libs/libspatialindex.so -Wl,-rpath,$ORIGIN/../../Rtree.libs +libspatialindex_so = Path(rtree.core.rt._name).absolute() +spatialindex_ = ExternalDependency( + include_dirs=[rtree.finder.get_include()], + extra_link_args=[str(libspatialindex_so)], + runtime_library_dirs=[ + os.path.join(dir, "Rtree.libs") for dir in sitepackage_dirs + ], +) + # firedrake_rtree rtree_ = ExternalDependency( include_dirs=[firedrake_rtree.get_include()], From 90ae828649dab09425ba17cda67fad775a8a74a0 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 24 Mar 2026 13:42:53 +0000 Subject: [PATCH 27/27] add spatialindex to supermesh extension --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 635b6fcecb..76553e372d 100644 --- a/setup.py +++ b/setup.py @@ -218,7 +218,7 @@ def extensions(): name="firedrake.cython.supermeshimpl", language="c", sources=[os.path.join("firedrake", "cython", "supermeshimpl.pyx")], - **(mpi_ + petsc_ + numpy_ + libsupermesh_) + **(mpi_ + petsc_ + numpy_ + libsupermesh_ + spatialindex_) )) # pyop2/sparsity.pyx: petsc, numpy, cython_list.append(Extension(