diff --git a/firedrake/evaluate.h b/firedrake/evaluate.h index 4253c2174e..5307fb72da 100644 --- a/firedrake/evaluate.h +++ b/firedrake/evaluate.h @@ -48,6 +48,7 @@ typedef PetscReal (*ref_cell_l1_dist_xtr)(void *data_, extern int locate_cell(struct Function *f, double *x, int dim, + int num_owned_cells, ref_cell_l1_dist try_candidate, ref_cell_l1_dist_xtr try_candidate_xtr, void *temp_ref_coords, @@ -56,6 +57,7 @@ extern int locate_cell(struct Function *f, extern int evaluate(struct Function *f, double *x, + int num_owned_cells, PetscScalar *result); #ifdef __cplusplus diff --git a/firedrake/function.py b/firedrake/function.py index ddedebad10..6f5fe118cb 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -576,7 +576,7 @@ def _c_evaluate(self, tolerance=None): return cache[tolerance] except KeyError: result = make_c_evaluate(self, tolerance=tolerance) - result.argtypes = [POINTER(_CFunction), POINTER(c_double), c_void_p] + result.argtypes = [POINTER(_CFunction), POINTER(c_double), c_int, c_void_p] result.restype = c_int return cache.setdefault(tolerance, result) @@ -650,8 +650,17 @@ def at(self, arg, *args, **kwargs): def single_eval(x, buf): r"""Helper function to evaluate at a single point.""" + mesh = self.ufl_domain() + if mesh.extruded: + if mesh.variable_layers: + raise NotImplementedError("Current codegen for extruded meshes " + "assumes constant layers") + num_owned_cells = mesh.cell_set.size * mesh.layers + else: + num_owned_cells = mesh.cell_set.size err = self._c_evaluate(tolerance=tolerance)(self._ctypes, x.ctypes.data_as(POINTER(c_double)), + num_owned_cells, buf.ctypes.data_as(c_void_p)) if err == -1: raise PointNotInDomainError(self.function_space().mesh(), x.reshape(-1)) diff --git a/firedrake/locate.c b/firedrake/locate.c index 2588fa3ef3..8829744956 100644 --- a/firedrake/locate.c +++ b/firedrake/locate.c @@ -4,140 +4,162 @@ #include #include +/** + * Check whether a point is either within, or sufficiently close to the current cell. + * + * To avoid the case where different ranks each claim the other to own the point, + * points are preferentially claimed by owned cells. + */ +bool check_cell(int64_t current_cell, + double current_ref_cell_dist_l1, + int num_owned_cells, + int64_t *found_cell, + bool *found_cell_is_owned, + void *found_ref_coords, + void *temp_ref_coords, + double *found_ref_cell_dist_l1) +{ + if (current_cell < num_owned_cells) { + if (current_ref_cell_dist_l1 <= 0) { + // found the cell + *found_cell = current_cell; + memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); + *found_ref_cell_dist_l1 = current_ref_cell_dist_l1; + return true; + } + else if (current_ref_cell_dist_l1 < tolerance && + current_ref_cell_dist_l1 < *found_ref_cell_dist_l1) + { + /* Close to cell within tolerance so could be this cell */ + *found_cell_is_owned = true; + + *found_cell = current_cell; + memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); + *found_ref_cell_dist_l1 = current_ref_cell_dist_l1; + } + } + else { + // Only consider ghost cells if an owned one is yet to be found + if (!(*found_cell_is_owned) && + current_ref_cell_dist_l1 < tolerance && + current_ref_cell_dist_l1 < *found_ref_cell_dist_l1) + { + *found_cell = current_cell; + memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords)); + *found_ref_cell_dist_l1 = current_ref_cell_dist_l1; + } + } + return false; +} + int locate_cell(struct Function *f, - double *x, - int dim, - ref_cell_l1_dist try_candidate, - ref_cell_l1_dist_xtr try_candidate_xtr, - void *temp_ref_coords, - void *found_ref_coords, - double *found_ref_cell_dist_l1) + double *x, + int dim, + int num_owned_cells, + ref_cell_l1_dist try_candidate, + ref_cell_l1_dist_xtr try_candidate_xtr, + void *temp_ref_coords, + void *found_ref_coords, + double *found_ref_cell_dist_l1) { - RTError err; - int cell = -1; /* NOTE: temp_ref_coords and found_ref_coords are actually of type - struct ReferenceCoords but can't be declared as such in the function - signature because the dimensions of the reference coordinates in the - ReferenceCoords struct are defined by python when the code which - surrounds this is declared in pointquery_utils.py. We cast when we use the - ref_coords_copy function and trust that the underlying memory which the - pointers refer to is updated as necessary. */ - double ref_cell_dist_l1 = DBL_MAX; - double current_ref_cell_dist_l1 = -0.5; + struct ReferenceCoords but can't be declared as such in the function + signature because the dimensions of the reference coordinates in the + ReferenceCoords struct are defined by python when the code which + surrounds this is declared in pointquery_utils.py. We cast when we use the + ref_coords_copy function and trust that the underlying memory which the + pointers refer to is updated as necessary. */ /* NOTE: `tolerance`, which is used throughout this funciton, is a static variable defined outside this function when putting together all the C code that needs to be compiled - see pointquery_utils.py */ + *found_ref_cell_dist_l1 = DBL_MAX; + + RTError err; + int64_t found_cell = -1; + bool found_cell_is_owned = false; + if (f->sidx) { + /* 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. */ 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) { + + if (!f->extruded) { for (uint64_t i = 0; i < nids; i++) { - current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, ids[i], x); - 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; - } - } + int64_t current_cell = ids[i]; + double current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, current_cell, x); + bool found = check_cell(current_cell, + current_ref_cell_dist_l1, + num_owned_cells, + &found_cell, + &found_cell_is_owned, + found_ref_coords, + temp_ref_coords, + found_ref_cell_dist_l1); + if (found) break; } } else { for (uint64_t i = 0; i < nids; i++) { + int64_t current_cell = ids[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); - 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; - } - } + int c = current_cell / nlayers; + int l = current_cell % nlayers; + double current_ref_cell_dist_l1 = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x); + + bool found = check_cell(current_cell, + current_ref_cell_dist_l1, + num_owned_cells, + &found_cell, + &found_cell_is_owned, + found_ref_coords, + temp_ref_coords, + found_ref_cell_dist_l1); + if (found) break; } } free(ids); } else { - if (f->extruded == 0) { + if (!f->extruded) { for (int c = 0; c < f->n_cols; c++) { - current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, c, x); - 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; - } - } + double current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, c, x); + bool found = check_cell(c, + current_ref_cell_dist_l1, + num_owned_cells, + &found_cell, + &found_cell_is_owned, + found_ref_coords, + temp_ref_coords, + found_ref_cell_dist_l1); + if (found) break; } } 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); - 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; + int64_t current_cell = c * f->n_layers + l; + double current_ref_cell_dist_l1 = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x); + bool found = check_cell(current_cell, + current_ref_cell_dist_l1, + num_owned_cells, + &found_cell, + &found_cell_is_owned, + found_ref_coords, + temp_ref_coords, + found_ref_cell_dist_l1); + if (found) break; } + if (found_cell != -1 && found_cell_is_owned) break; } } } - return cell; + return found_cell; } diff --git a/firedrake/mesh.py b/firedrake/mesh.py index acc9031542..e76c217431 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2575,12 +2575,22 @@ def locate_cells_ref_coords_and_dists(self, xs, tolerance=None): ref_cell_dists_l1 = np.empty(npoints, dtype=utils.RealType) cells = np.empty(npoints, dtype=IntType) assert xs.size == npoints * self.geometric_dimension() + + if self.extruded: + if self.variable_layers: + raise NotImplementedError("Current codegen for extruded meshes " + "assumes constant layers") + num_owned_cells = self.cell_set.size * self.layers + else: + num_owned_cells = self.cell_set.size + self._c_locator(tolerance=tolerance)(self.coordinates._ctypes, xs.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), Xs.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), ref_cell_dists_l1.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), cells.ctypes.data_as(ctypes.POINTER(ctypes.c_int)), - npoints) + npoints, + num_owned_cells) return cells, Xs, ref_cell_dists_l1 def _c_locator(self, tolerance=None): @@ -2596,7 +2606,7 @@ def _c_locator(self, tolerance=None): IntTypeC = as_cstr(IntType) src = pq_utils.src_locate_cell(self, tolerance=tolerance) src += dedent(f""" - int locator(struct Function *f, double *x, double *X, double *ref_cell_dists_l1, {IntTypeC} *cells, {IntTypeC} npoints) + int locator(struct Function *f, double *x, double *X, double *ref_cell_dists_l1, {IntTypeC} *cells, {IntTypeC} npoints, int num_owned_cells) {{ {IntTypeC} j = 0; /* index into x and X */ for({IntTypeC} i=0; i