diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 6026191a32..f56363164d 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -39,7 +39,6 @@ from firedrake.petsc import PETSc, DEFAULT_PARTITIONER from firedrake.adjoint_utils import MeshGeometryMixin from firedrake.exceptions import VertexOnlyMeshMissingPointsError, NonUniqueMeshSequenceError -from pyadjoint import stop_annotating import gem try: @@ -4377,11 +4376,6 @@ def _parent_mesh_embedding( "VertexOnlyMeshes don't have a working locate_cells_ref_coords_and_dists method" ) - import firedrake.functionspace as functionspace - import firedrake.constant as constant - import firedrake.interpolation as interpolation - import firedrake.assemble as assemble - 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. @@ -4433,25 +4427,6 @@ def _parent_mesh_embedding( icomm.Allgatherv( input_ranks_local, (input_ranks_global, ncoords_local_allranks) ) - - # Get parent mesh rank ownership information: - # Interpolating Constant(parent_mesh.comm.rank) into P0DG cleverly creates - # a Function whose dat contains rank ownership information in an ordering - # that is accessible using Firedrake's cell numbering. This is because, on - # each rank, parent_mesh.comm.rank creates a Constant with the local rank - # number, and halo exchange ensures that this information is visible, as - # nessesary, to other processes. - P0DG = functionspace.FunctionSpace(parent_mesh, "DG", 0) - with stop_annotating(): - visible_ranks = interpolation.interpolate( - constant.Constant(parent_mesh.comm.rank), P0DG - ) - visible_ranks = assemble(visible_ranks).dat.data_ro_with_halos.real - - locally_visible = np.full(ncoords_global, False) - # See below for why np.inf is used here. - ranks = np.full(ncoords_global, np.inf) - ( parent_cell_nums, reference_coords, @@ -4466,8 +4441,26 @@ def _parent_mesh_embedding( # which we can safely delete reference_coords = reference_coords[:, : parent_mesh.topological_dimension] - locally_visible[:] = parent_cell_nums != -1 - ranks[locally_visible] = visible_ranks[parent_cell_nums[locally_visible]] + # Get parent mesh rank ownership information. + visible_ranks = np.empty(parent_mesh.cell_set.total_size, dtype=IntType) + 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 + ) + 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] + + ranks = np.full(ncoords_global, np.inf) # See below for why np.inf is used here. + ranks[locally_visible] = visible_ranks[locally_visible_cell_nums] + # see below for why np.inf is used here. ref_cell_dists_l1[~locally_visible] = np.inf @@ -4547,9 +4540,11 @@ def _parent_mesh_embedding( ) changed_ranks_tied &= locally_visible # update the identified rank - ranks[changed_ranks_tied] = visible_ranks[ - parent_cell_nums[changed_ranks_tied] - ] + 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]