diff --git a/firedrake/cython/rtree.pyx b/firedrake/cython/rtree.pyx new file mode 100644 index 0000000000..6997cc5771 --- /dev/null +++ b/firedrake/cython/rtree.pyx @@ -0,0 +1,83 @@ +# cython: language_level=3 + +cimport numpy as np +import numpy as np +import ctypes +import cython +from libc.stddef cimport size_t +from libc.stdint cimport uintptr_t + +include "rtreeinc.pxi" + +cdef class RTree(object): + """Python class for holding a spatial index.""" + + cdef RTreeH* tree + + def __cinit__(self, uintptr_t tree_handle): + self.tree = 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 build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min, + np.ndarray[np.float64_t, ndim=2, mode="c"] coords_max, + np.ndarray[np.npy_uintp, ndim=1, mode="c"] ids = None): + """Builds rtree from two arrays of shape (n, dim) containing the coordinates + of the lower and upper corners of n axis-aligned bounding boxes, and an + optional array of shape (n,) containing integer ids for each box. + + Parameters + ---------- + coords_min : (n, dim) array + The lower corner coordinates of the bounding boxes. + regions_hi : (n, dim) array + The upper corner coordinates of the bounding boxes. + ids : (n,) array, optional + Integer ids for each box. If not provided, defaults to 0, 1, ..., n-1. + + Returns + ------- + RTree + An RTree object containing the Rtree. + """ + cdef: + RTreeH* rtree + size_t n + size_t dim + RTreeError err + + if coords_min.shape[0] != coords_max.shape[0] or coords_min.shape[1] != coords_max.shape[1]: + raise ValueError("coords_min and coords_max must have the same shape") + + n = coords_min.shape[0] + dim = coords_min.shape[1] + if ids is None: + ids = np.arange(n, dtype=np.uintp) + elif ids.shape[0] != n: + raise ValueError("Mismatch between number of boxes and number of ids") + + err = rtree_bulk_load( + &rtree, + coords_min.data, + coords_max.data, + ids.data, + n, + dim + ) + if err != Success: + raise RuntimeError("rtree_bulk_load failed") + + return RTree(rtree) diff --git a/firedrake/cython/rtreeinc.pxi b/firedrake/cython/rtreeinc.pxi new file mode 100644 index 0000000000..7d1056bdd7 --- /dev/null +++ b/firedrake/cython/rtreeinc.pxi @@ -0,0 +1,30 @@ +from libc.stddef cimport size_t +from libc.stdint cimport uint32_t + + +cdef extern from "rtree-capi.h": + ctypedef enum RTreeError: + Success + NullPointer + InvalidDimension + + ctypedef struct RTreeH: + pass + + RTreeError rtree_bulk_load( + RTreeH **tree, + const double *mins, + const double *maxs, + const size_t *ids, + size_t n, + uint32_t dim + ) + + RTreeError rtree_free(RTreeH *tree) + + RTreeError rtree_locate_all_at_point( + const RTreeH *tree, + const double *point, + size_t **ids_out, + size_t *nids_out + ) 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..5c3497e76c 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -1,5 +1,5 @@ import numpy as np -import rtree +import firedrake_rtree import sys import ufl import warnings @@ -12,7 +12,6 @@ from ctypes import POINTER, c_int, c_double, c_void_p from collections.abc import Collection from numbers import Number -from pathlib import Path from functools import partial, cached_property from typing import Tuple @@ -877,15 +876,13 @@ def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None): if ldargs is None: ldargs = [] - libspatialindex_so = Path(rtree.core.rt._name).absolute() - lsi_runpath = f"-Wl,-rpath,{libspatialindex_so.parent}" - ldargs += [str(libspatialindex_so), lsi_runpath] + ldargs += [firedrake_rtree.get_lib_filename(), f"-Wl,-rpath,{firedrake_rtree.get_lib()}"] dll = compilation.load( src, "c", cppargs=[ f"-I{path.dirname(__file__)}", f"-I{sys.prefix}/include", - f"-I{rtree.finder.get_include()}" + f"-I{firedrake_rtree.get_include()}" ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=ldargs, comm=function.comm diff --git a/firedrake/locate.c b/firedrake/locate.c index a243e2119f..cc9f940d08 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 @@ -31,156 +31,81 @@ 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) { - 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); - return -1; - } - if (f->extruded == 0) { - for (uint64_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++) { - if (ids[i] == cells_ignore[j]) { - cell_ignore_found = 1; - break; - } - } - if (cell_ignore_found) { - cell_ignore_found = 0; - continue; - } - if (current_ref_cell_dist_l1 <= 0.0) { - /* Found cell! */ - cell = ids[i]; - memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); - found_ref_cell_dist_l1[0] = current_ref_cell_dist_l1; + 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); + rtree_free_ids(ids, nids); + return -1; + } + if (f->extruded == 0) { + for (size_t i = 0; i < nids; i++) { + current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, ids[i], x); + for (size_t j = 0; j < ncells_ignore; j++) { + if (ids[i] == cells_ignore[j]) { + cell_ignore_found = 1; break; } - else if (current_ref_cell_dist_l1 < ref_cell_dist_l1) { - /* getting closer... */ - ref_cell_dist_l1 = current_ref_cell_dist_l1; - if (ref_cell_dist_l1 < tolerance) { - /* Close to cell within tolerance so could be this cell */ - cell = ids[i]; - memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); - found_ref_cell_dist_l1[0] = ref_cell_dist_l1; - } - } } - } - else { - for (uint64_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++) { - if (ids[i] == cells_ignore[j]) { - cell_ignore_found = 1; - break; - } - } - if (cell_ignore_found) { - cell_ignore_found = 0; - continue; - } - if (current_ref_cell_dist_l1 <= 0.0) { - /* Found cell! */ + if (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; } } } - free(ids); - } 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; + } + 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; } - 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; + 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] = ref_cell_dist_l1; } } } } + rtree_free_ids(ids, nids); return cell; } diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 59531066c1..ebd9dabc04 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -16,7 +16,7 @@ import enum import numbers import abc -import rtree +import firedrake_rtree from textwrap import dedent from pathlib import Path import typing @@ -34,10 +34,9 @@ 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.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 @@ -2478,7 +2477,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,9 +2485,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, since libspatialindex - does not support 1D. - Notes ----- If we have a higher-order (bendy) mesh we project the mesh coordinates into @@ -2499,11 +2495,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("libspatialindex does not support 1-dimension, falling back on brute force.") - return None - coord_element = self.ufl_coordinate_element() coord_degree = coord_element.degree() if np.all(np.asarray(coord_degree) == 1): @@ -2538,7 +2529,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) @@ -2547,7 +2538,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]) @@ -2574,8 +2565,7 @@ def spatial_index(self): Returns ------- - :class:`~.spatialindex.SpatialIndex` or None if the mesh is - one-dimensional. + :class:`~.rtree.RTree` Notes ----- @@ -2602,20 +2592,19 @@ 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(-1, 1) + 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 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 = rtree.build_from_aabb(coords_min, coords_max) self._saved_coordinate_dat_version = self.coordinates.dat.dat_version return self._spatial_index @@ -2766,20 +2755,18 @@ def _c_locator(self, tolerance=None): }} """) - libspatialindex_so = Path(rtree.core.rt._name).absolute() - lsi_runpath = f"-Wl,-rpath,{libspatialindex_so.parent}" 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{firedrake_rtree.get_include()}" ] + [f"-I{d}/include" for d in get_petsc_dir()], ldargs=[ f"-L{sys.prefix}/lib", - str(libspatialindex_so), + firedrake_rtree.get_lib_filename(), f"-Wl,-rpath,{sys.prefix}/lib", - lsi_runpath + f"-Wl,-rpath,{firedrake_rtree.get_lib()}" ], comm=self.comm ) 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 diff --git a/setup.py b/setup.py index 1d5b588355..76553e372d 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 @@ -120,9 +121,9 @@ 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 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 rtree and libsupermesh 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. @@ -141,6 +142,15 @@ def __getitem__(self, key): ], ) +# firedrake_rtree +rtree_ = ExternalDependency( + include_dirs=[firedrake_rtree.get_include()], + extra_link_args=[firedrake_rtree.get_lib_filename()], + runtime_library_dirs=[ + os.path.join(dir, "firedrake_rtree") for dir in sitepackage_dirs + ], +) + # libsupermesh # example: # gcc -Ipath/to/libsupermesh/include @@ -196,19 +206,19 @@ def extensions(): sources=[os.path.join("firedrake", "cython", "patchimpl.pyx")], **(mpi_ + petsc_ + numpy_) )) - # firedrake/cython/spatialindex.pyx: numpy, spatialindex + # firedrake/cython/rtree.pyx: numpy, rtree-capi cython_list.append(Extension( - name="firedrake.cython.spatialindex", + name="firedrake.cython.rtree", language="c", - sources=[os.path.join("firedrake", "cython", "spatialindex.pyx")], - **(mpi_ + numpy_ + spatialindex_) + sources=[os.path.join("firedrake", "cython", "rtree.pyx")], + **(mpi_ + numpy_ + rtree_) )) # firedrake/cython/supermeshimpl.pyx: petsc, numpy, supermesh cython_list.append(Extension( 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( 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: