Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
Need to see if I'm correctly dealing with the case of two ranks thinking
the other owns the point.
  • Loading branch information
ReubenHill committed Dec 14, 2023
1 parent df8eb23 commit d33cca9
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 28 deletions.
3 changes: 2 additions & 1 deletion firedrake/evaluate.h
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,8 @@ extern int locate_cell(struct Function *f,
void *temp_ref_coords,
void *found_ref_coords,
double *found_ref_cell_dist_l1,
int cell_ignore);
size_t ncells_ignore,
int* cells_ignore);

extern int evaluate(struct Function *f,
double *x,
Expand Down
40 changes: 35 additions & 5 deletions firedrake/locate.c
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,12 @@ int locate_cell(struct Function *f,
void *temp_ref_coords,
void *found_ref_coords,
double *found_ref_cell_dist_l1,
int cell_ignore)
size_t ncells_ignore,
int* cells_ignore)
{
RTError err;
int cell = -1;
int cell_ignore_found = 0;
/* 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
Expand Down Expand Up @@ -43,7 +45,14 @@ int locate_cell(struct Function *f,
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);
if (ids[i] == cell_ignore) {
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) {
Expand Down Expand Up @@ -71,7 +80,14 @@ int locate_cell(struct Function *f,
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 (ids[i] == cell_ignore) {
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) {
Expand All @@ -98,7 +114,14 @@ int locate_cell(struct Function *f,
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);
if (c == cell_ignore) {
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) {
Expand All @@ -124,7 +147,14 @@ int locate_cell(struct Function *f,
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 (l == cell_ignore) {
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) {
Expand Down
66 changes: 46 additions & 20 deletions firedrake/mesh.py
Original file line number Diff line number Diff line change
Expand Up @@ -2191,7 +2191,7 @@ def locate_cell_and_reference_coordinate(self, x, tolerance=None, cell_ignore=No
if x.size != self.geometric_dimension():
raise ValueError("Point must have the same geometric dimension as the mesh")
x = x.reshape((1, self.geometric_dimension()))
cells, ref_coords, _ = self.locate_cells_ref_coords_and_dists(x, tolerance=tolerance, cells_ignore=[cell_ignore])
cells, ref_coords, _ = self.locate_cells_ref_coords_and_dists(x, tolerance=tolerance, cells_ignore=[[cell_ignore]])
if cells[0] == -1:
return None, None
return cells[0], ref_coords[0]
Expand All @@ -2206,9 +2206,10 @@ def locate_cells_ref_coords_and_dists(self, xs, tolerance=None, cells_ignore=Non
this from default will cause the spatial index to be rebuilt which
can take some time.
:kwarg cells_ignore: Cell numbers to ignore in the search for each
point in xs. Each entry corresponds to a single coordinate in xs.
To not ignore any cells, pass None. To ensure a full cell search
for any given point, set the corresponding list entry to -1.
point in xs. Shape should be (n_ignore_pts, npoints). Each column
corresponds to a single coordinate in xs. To not ignore any cells,
pass None. To ensure a full cell search for any given point, set
the corresponding entries to -1.
:returns: tuple either
(cell numbers array, reference coordinates array, ref_cell_dists_l1 array)
of type
Expand All @@ -2229,12 +2230,13 @@ def locate_cells_ref_coords_and_dists(self, xs, tolerance=None, cells_ignore=Non
raise ValueError("Point coordinate dimension does not match mesh geometric dimension")
Xs = np.empty_like(xs)
npoints = len(xs)
if cells_ignore is None or cells_ignore[0] is None:
cells_ignore = np.full(npoints, -1, dtype=IntType)
if cells_ignore is None or cells_ignore[0][0] is None:
cells_ignore = np.full((1, npoints), -1, dtype=IntType)
else:
cells_ignore = np.asarray(cells_ignore, dtype=IntType)
if len(cells_ignore) != npoints:
if cells_ignore.shape[1] != npoints:
raise ValueError("Number of cells to ignore does not match number of points")
assert cells_ignore.shape == (cells_ignore.shape[0], npoints)
ref_cell_dists_l1 = np.empty(npoints, dtype=utils.RealType)
cells = np.empty(npoints, dtype=IntType)
assert xs.size == npoints * self.geometric_dimension()
Expand All @@ -2244,6 +2246,7 @@ def locate_cells_ref_coords_and_dists(self, xs, tolerance=None, cells_ignore=Non
ref_cell_dists_l1.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
cells.ctypes.data_as(ctypes.POINTER(ctypes.c_int)),
npoints,
cells_ignore.shape[0],
cells_ignore.ctypes.data_as(ctypes.POINTER(ctypes.c_int)))
return cells, Xs, ref_cell_dists_l1

Expand All @@ -2259,7 +2262,7 @@ def _c_locator(self, tolerance=None):
except KeyError:
src = pq_utils.src_locate_cell(self, tolerance=tolerance)
src += """
int locator(struct Function *f, double *x, double *X, double *ref_cell_dists_l1, int *cells, size_t npoints, int* cells_ignore)
int locator(struct Function *f, double *x, double *X, double *ref_cell_dists_l1, int *cells, size_t npoints, size_t ncells_ignore, int* cells_ignore)
{
size_t j = 0; /* index into x and X */
for(size_t i=0; i<npoints; i++) {
Expand All @@ -2272,7 +2275,10 @@ def _c_locator(self, tolerance=None):
/* to_reference_coords and to_reference_coords_xtr are defined in
pointquery_utils.py. If they contain python calls, this loop will
not run at c-loop speed. */
cells[i] = locate_cell(f, &x[j], %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &ref_cell_dists_l1[i], cells_ignore[i]);
/* cells_ignore has shape (ncells_ignore, npoints) - find the ith column */
int *cells_ignore_i = cells_ignore + i*ncells_ignore;
cells[i] = locate_cell(f, &x[j], %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &ref_cell_dists_l1[i], ncells_ignore, cells_ignore_i);
for (int k = 0; k < %(geometric_dimension)d; k++) {
X[j] = found_reference_coords.X[k];
Expand All @@ -2297,6 +2303,7 @@ def _c_locator(self, tolerance=None):
ctypes.POINTER(ctypes.c_double),
ctypes.POINTER(ctypes.c_int),
ctypes.c_size_t,
ctypes.c_size_t,
ctypes.POINTER(ctypes.c_int)]
locator.restype = ctypes.c_int
return cache.setdefault(tolerance, locator)
Expand Down Expand Up @@ -3748,24 +3755,32 @@ def _parent_mesh_embedding(
changed_ranks = owned_ranks != ranks

# If distance has changed the the point is not in local mesh partition
# since some other cell is closer.
# 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
# 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. This should only happen for halo points.
# point.
changed_ranks_tied = changed_ranks & ~changed_ref_cell_dists_l1

# parent_mesh.comm.rank
# parent_mesh.comm.rank
# locally_visible[9846]
# ranks[9846]
# owned_ranks[9846]
# ref_cell_dists_l1[9846]
# parent_cell_nums[9846]
# for debugging: find index into owned ranks where changed ranks ties equals
idxs = []
for i in range(len(owned_ranks)):
if (owned_ranks[i] == parent_mesh.comm.rank) & changed_ranks_tied[i]:
idxs.append(i)
import pdb; pdb.set_trace()
if any(changed_ranks_tied):
# check that only halos are affected
if any(owned_ranks[changed_ranks_tied] == parent_mesh.comm.rank):
raise RuntimeError(
"A tie break has occurred in the voting algorithm which has "
"resulted in a non-halo point changing MPI rank. This should "
"not happen."
)
old_parent_cell_nums = np.asarray([np.copy(parent_cell_nums)])
while any(changed_ranks_tied):
(
parent_cell_nums[changed_ranks_tied],
Expand All @@ -3774,25 +3789,36 @@ def _parent_mesh_embedding(
) = parent_mesh.locate_cells_ref_coords_and_dists(
coords_global[changed_ranks_tied],
tolerance,
cells_ignore=parent_cell_nums[changed_ranks_tied],
cells_ignore=old_parent_cell_nums[:, changed_ranks_tied],
)
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
# 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
ranks[changed_ranks_tied] = visible_ranks[
parent_cell_nums[changed_ranks_tied]
]
# 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]
)
changed_ranks_tied &= locally_visible
import pdb; pdb.set_trace()
# add more cells to ignore
old_parent_cell_nums = np.vstack((
old_parent_cell_nums,
parent_cell_nums)
)

import pdb; pdb.set_trace()
# Any ranks which are still np.inf are not in the mesh
missing_global_idxs = np.where(owned_ranks == np.inf)[0]

Expand Down
3 changes: 2 additions & 1 deletion firedrake/pointeval_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,8 @@ def predicate(index):
/* The type definitions and arguments used here are defined as statics in pointquery_utils.py */
double found_ref_cell_dist_l1 = DBL_MAX;
struct ReferenceCoords temp_reference_coords, found_reference_coords;
%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &found_ref_cell_dist_l1, -1);
int cells_ignore[1] = {-1};
%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &temp_reference_coords, &found_reference_coords, &found_ref_cell_dist_l1, 1, cells_ignore);
if (cell == -1) {
return -1;
}
Expand Down
2 changes: 1 addition & 1 deletion tests/regression/test_locate_cell.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def value_at(p, cell_ignore=None):

def value_at_and_dist(p, cell_ignore=None):
if cell_ignore is not None:
cell_ignore = [cell_ignore]
cell_ignore = [[cell_ignore]]
cells, _, l1_dists = m.locate_cells_ref_coords_and_dists([p], cells_ignore=cell_ignore)
return f.dat.data[cells[0]], l1_dists[0]

Expand Down

0 comments on commit d33cca9

Please sign in to comment.