From 5559682c3cea646b30fa28aa1c152ec2ca6c11ce Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 18 Feb 2026 16:47:58 +0000 Subject: [PATCH 1/6] initial --- firedrake/mesh.py | 47 +++++++++++++++++++++-------------------------- 1 file changed, 21 insertions(+), 26 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 6026191a32..96e8695f3e 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4377,11 +4377,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 +4428,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 +4442,27 @@ 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. + with PETSc.Log.Event("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 + # These entries will be overwritten by the halo exchange below. + visible_ranks[parent_mesh.cell_set.size :] = -1 + dmcommon.exchange_cell_orientations( + parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks + ) + locally_visible = np.full(ncoords_global, False) + # See below for why np.inf is used here. + ranks = np.full(ncoords_global, np.inf) + locally_visible[:] = parent_cell_nums != -1 + if parent_mesh.extruded: + base_parent_cell_nums, _ = _parent_extrusion_numbering( + parent_cell_nums[locally_visible], parent_mesh.layers + ) + ranks[locally_visible] = visible_ranks[base_parent_cell_nums] + else: + ranks[locally_visible] = visible_ranks[parent_cell_nums[locally_visible]] + # see below for why np.inf is used here. ref_cell_dists_l1[~locally_visible] = np.inf From 37d63d8816519ed90b1535ba78ebd8441191d361 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 18 Feb 2026 17:32:27 +0000 Subject: [PATCH 2/6] comment --- firedrake/mesh.py | 1 - 1 file changed, 1 deletion(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 96e8695f3e..f16310407c 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4446,7 +4446,6 @@ def _parent_mesh_embedding( with PETSc.Log.Event("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 - # These entries will be overwritten by the halo exchange below. visible_ranks[parent_mesh.cell_set.size :] = -1 dmcommon.exchange_cell_orientations( parent_mesh.topology.topology_dm, parent_mesh.topology._cell_numbering, visible_ranks From 95dc00e40b9419850d50cae48d747f5891453d68 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 19 Feb 2026 12:03:37 +0000 Subject: [PATCH 3/6] add comment and lint --- firedrake/mesh.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index f16310407c..b56880cc75 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: @@ -4445,8 +4444,9 @@ def _parent_mesh_embedding( # Get parent mesh rank ownership information. with PETSc.Log.Event("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 + 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 ) From fce076f6728c69ff2bcd3711169d4781cf94f70a Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 19 Feb 2026 13:22:55 +0000 Subject: [PATCH 4/6] extruded comment --- firedrake/mesh.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index b56880cc75..9db182bcca 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4454,13 +4454,15 @@ def _parent_mesh_embedding( # See below for why np.inf is used here. ranks = np.full(ncoords_global, np.inf) locally_visible[:] = parent_cell_nums != -1 + if parent_mesh.extruded: - base_parent_cell_nums, _ = _parent_extrusion_numbering( - parent_cell_nums[locally_visible], parent_mesh.layers - ) - ranks[locally_visible] = visible_ranks[base_parent_cell_nums] + # 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: - ranks[locally_visible] = visible_ranks[parent_cell_nums[locally_visible]] + locally_visible_cell_nums = parent_cell_nums[locally_visible] + + 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 From 3bbb78ecaf6a1a9abd991264f031a6d6776515f7 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 19 Feb 2026 17:35:02 +0000 Subject: [PATCH 5/6] missed one --- firedrake/mesh.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 9db182bcca..ce0affb54f 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4461,7 +4461,6 @@ def _parent_mesh_embedding( locally_visible_cell_nums = parent_cell_nums[locally_visible] // (parent_mesh.layers - 1) else: locally_visible_cell_nums = parent_cell_nums[locally_visible] - ranks[locally_visible] = visible_ranks[locally_visible_cell_nums] # see below for why np.inf is used here. @@ -4543,9 +4542,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] From fb27252ceae80eedcc8f5123950b4bb300b59ea6 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 20 Feb 2026 12:09:31 +0000 Subject: [PATCH 6/6] review suggestion & remove log context --- firedrake/mesh.py | 38 ++++++++++++++++++-------------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index ce0affb54f..f56363164d 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4442,26 +4442,24 @@ def _parent_mesh_embedding( reference_coords = reference_coords[:, : parent_mesh.topological_dimension] # Get parent mesh rank ownership information. - with PETSc.Log.Event("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 = np.full(ncoords_global, False) - # See below for why np.inf is used here. - ranks = np.full(ncoords_global, np.inf) - 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[locally_visible] = visible_ranks[locally_visible_cell_nums] + 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