From cedcf43c71c515b93f93449872de55cbf3141ab5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Jul 2026 14:39:12 +0100 Subject: [PATCH 1/6] Unify assign across Submesh, RestrictedFunctionSpace and RestrictedElement Assigner special-cased Submesh, and RestrictedFunctionSpace went through a separate cell-node-map-based index computation in facet_split.get_restriction_indices. Both relate a source space's nodes to a subset of a target space's nodes, so route them through the same node-matching machinery: get_restriction_indices now walks node numbers through assign for the non-extruded case, and Assigner grows the general handling needed to cover RestrictedFunctionSpace, Submesh, and combinations of the two, including the redistributed Submesh case from the previous commit. --- firedrake/assign.py | 479 +++++++++++++++--- firedrake/preconditioners/facet_split.py | 22 + tests/firedrake/regression/test_assign.py | 28 + .../firedrake/submesh/test_submesh_assign.py | 179 +++++++ 4 files changed, 632 insertions(+), 76 deletions(-) diff --git a/firedrake/assign.py b/firedrake/assign.py index 4a837bd2a0..e87cf0bcde 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -9,6 +9,7 @@ import pytools import finat.ufl from ufl.algorithms import extract_coefficients +from ufl.cell import TensorProductCell from ufl.constantvalue import as_ufl from ufl.corealg.map_dag import map_expr_dag from ufl.corealg.multifunction import MultiFunction @@ -19,13 +20,13 @@ from firedrake.function import Function from firedrake.halo import _get_mtype from firedrake.petsc import PETSc -from firedrake.utils import ScalarType, split_by +from firedrake.utils import IntType, ScalarType, split_by from mpi4py import MPI def _submesh_point_sf(target_mesh, source_mesh): - """Find the point SF relating two meshes with different distributions. + """Find the point SF relating a submesh to its parent. Parameters ---------- @@ -39,16 +40,26 @@ def _submesh_point_sf(target_mesh, source_mesh): tuple The `PETSc.SF` mapping the points of the parent mesh (roots) to the points of the submesh (leaves), and whether ``target_mesh`` is the - submesh. Both are `None` if the two meshes share their distribution, - in which case they are related by entity maps instead. + submesh. Both are `None` if neither mesh is a submesh of the other. + When the two share their distribution, the SF is the local subpoint + map rather than one of the meshes' own ``submesh_point_sf``. """ + from firedrake.mesh import _make_submesh_point_sf + if target_mesh.submesh_parent is source_mesh: - return target_mesh.submesh_point_sf, True + submesh, target_is_submesh = target_mesh, True elif source_mesh.submesh_parent is target_mesh: - return source_mesh.submesh_point_sf, False + submesh, target_is_submesh = source_mesh, False else: return None, None + point_sf = submesh.submesh_point_sf + if point_sf is None: + # The two share their distribution, so their points are related + # locally by the subpoint IS. + point_sf = _make_submesh_point_sf(submesh.submesh_parent.topology_dm, + submesh.topology_dm) + return point_sf, target_is_submesh def _make_section_sf(point_sf, root_V, leaf_V): @@ -90,6 +101,227 @@ def _make_section_sf(point_sf, root_V, leaf_V): return cache.setdefault(key, (section_sf, section_sf.computeDegree() > 0)) +def _identity_point_sf(mesh): + """Create the `PETSc.SF` mapping the points of a mesh onto themselves. + + Two function spaces on the same mesh are related by their `PETSc.Section` + alone, which the section SF machinery expresses as the identity on points. + + Parameters + ---------- + mesh : firedrake.mesh.AbstractMeshTopology + The mesh. + + Returns + ------- + PETSc.SF + SF whose roots and leaves are both the points of ``mesh``. + + """ + plex = mesh.topology_dm + pStart, pEnd = plex.getChart() + remote = np.empty((pEnd - pStart, 2), dtype=IntType) + remote[:, 0] = plex.comm.rank + remote[:, 1] = np.arange(pEnd - pStart, dtype=IntType) + point_sf = PETSc.SF().create(comm=plex.comm) + point_sf.setGraph(pEnd - pStart, None, remote) + return point_sf + + +def _entity_dof_counts(element): + """Count the nodes an element places on each entity of the reference cell. + + Parameters + ---------- + element : finat.ufl.finiteelementbase.FiniteElementBase + The UFL element. + + Returns + ------- + dict + The number of nodes on each ``(dimension, entity)`` of the reference cell. + + """ + from firedrake.functionspacedata import create_element + + entity_dofs = create_element(element).entity_dofs() + return {(dim, entity): len(dofs) + for dim, entities in entity_dofs.items() + for entity, dofs in entities.items()} + + +def _unrestricted(element): + """Strip the topological restrictions off an element. + + Parameters + ---------- + element : finat.ufl.finiteelementbase.FiniteElementBase + The UFL element. + + Returns + ------- + finat.ufl.finiteelementbase.FiniteElementBase + The element whose nodes ``element`` selects from. + + """ + while isinstance(element, finat.ufl.RestrictedElement): + element = element.sub_element() + return element + + +def _compatible_elements(target, source): + """Whether functions in two elements share a node layout. + + Two elements are compatible if they restrict a common element, and place + the same number of nodes on every entity of the reference cell up to the + entities one of them drops. Their nodes are then the same functionals + entity by entity, so the `PETSc.Section` of either function space + describes both and data moves between them without reference to the cell + node maps. + + Parameters + ---------- + target : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned to. + source : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned from. + + Returns + ------- + bool + Whether the two elements share a node layout. + + """ + if target == source: + return True + blocked_types = (finat.ufl.VectorElement, finat.ufl.TensorElement) + if isinstance(target, blocked_types) or isinstance(source, blocked_types): + return (type(target) is type(source) + and target.num_sub_elements == source.num_sub_elements + and target.reference_value_shape == source.reference_value_shape + and _compatible_elements(target.sub_elements[0], source.sub_elements[0])) + # Equal node counts on an entity do not by themselves make two nodes the + # same functional, so the elements must restrict a common element. + if _unrestricted(target) != _unrestricted(source): + return False + if isinstance(target.cell, TensorProductCell): + # A base mesh point of an extruded mesh carries the nodes of a whole + # column of entities, of which a restriction may drop only some, and + # the Section counts the nodes on a point without saying which. + raise NotImplementedError( + "Assigning between an element and its restriction is not " + "implemented on extruded meshes" + ) + target_counts = _entity_dof_counts(target) + source_counts = _entity_dof_counts(source) + if target_counts.keys() != source_counts.keys(): + return False + # An entity either carries the same nodes in both elements, or is + # dropped by one of them; a partial overlap has no entity-wise + # correspondence and so cannot be expressed by the two Sections. + return all(nt == source_counts[entity] or nt == 0 or source_counts[entity] == 0 + for entity, nt in target_counts.items()) + + +def _node_subset(V, cell_subset): + """Find the nodes of the cells of a subset. + + A node on the boundary of the subset belongs to cells outside it too, and + is included: the subset selects the cells whose nodes are assigned, not + the nodes that no other cell shares. + + Parameters + ---------- + V : firedrake.functionspaceimpl.WithGeometry + Function space whose nodes are selected. + cell_subset : pyop2.types.set.Subset + Subset of the cells of the mesh of ``V``. + + Returns + ------- + pyop2.types.set.Subset + The nodes of ``V`` on the cells of ``cell_subset``. + + """ + if V.extruded: + raise NotImplementedError( + "Assigning over a subset of the cells is not implemented on " + "extruded meshes" + ) + node_map = V.cell_node_map() + if node_map is None: + raise ValueError(f"Function space ({V}) has no nodes on the cells") + # A node is on the subset for every rank that shares it as soon as it is on + # the subset for one of them, which the cells known to a single rank do not + # say: a rank owning the node need not own, or even halo, a cell of the + # subset that carries it. + marker = op2.Dat(V.node_set, dtype=IntType) + marker.data_wo_with_halos[np.unique(node_map.values_with_halo[cell_subset.indices])] = 1 + marker.local_to_global_begin(op2.MAX) + marker.local_to_global_end(op2.MAX) + marker.global_to_local_begin(op2.READ) + marker.global_to_local_end(op2.READ) + nodes, = np.nonzero(marker.data_ro_with_halos) + return op2.Subset(V.node_set, nodes) + + +def _target_is_leaf(target, source): + """Whether the target element's nodes are the subset of the two. + + Data travels root to leaf by broadcast, which requires every leaf node to + have a counterpart, and leaf to root by reduction, which does not. The + target can therefore be the leaf unless it carries nodes on an entity + that the source drops, and taking it to be the leaf whenever possible + keeps the halo of the assignee up to date. + + Parameters + ---------- + target : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned to. + source : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned from. + + Returns + ------- + bool + Whether every node of ``target`` has a counterpart in ``source``. + + """ + target_counts = _entity_dof_counts(target) + source_counts = _entity_dof_counts(source) + return not any(source_counts[e] == 0 < target_counts[e] for e in target_counts) + + +def _relate_to_target(target_mesh, target_element, source_V): + """Find the section SF relating a source function space to the assignee. + + Parameters + ---------- + target_mesh : AbstractMeshTopology + Mesh of the function being assigned to. + target_element : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned to. + source_V : firedrake.functionspaceimpl.WithGeometry + Function space of the function being assigned from. + + Returns + ------- + tuple + The `PETSc.SF` relating the points of ``target_mesh`` to those of + ``source_V``'s mesh, and whether the assignee is the leaf. + + """ + source_mesh = source_V.mesh().topology + if target_mesh is source_mesh: + return _identity_point_sf(target_mesh), _target_is_leaf(target_element, source_V.ufl_element()) + point_sf, target_is_leaf = _submesh_point_sf(target_mesh, source_mesh) + if point_sf is None: + raise NotImplementedError( + "Can only assign between a redistributed mesh and its parent" + ) + return point_sf, target_is_leaf + + def _isconstant(expr): return isinstance(expr, Constant) or \ (isinstance(expr, (Function, Cofunction)) and expr.ufl_element().family() == "Real") @@ -226,9 +458,9 @@ def __init__(self, assignee, expression, subset=None): source_meshes = set() for coeff in extract_coefficients(expression): if isinstance(coeff, (Function, Cofunction)) and coeff.ufl_element().family() != "Real": - if coeff.ufl_element() != assignee.ufl_element(): - raise ValueError("All functions in the expression must have the same " - "element as the assignee") + if not _compatible_elements(assignee.ufl_element(), coeff.ufl_element()): + raise ValueError("All functions in the expression must have an " + "element compatible with that of the assignee") source_meshes.add(extract_unique_domain(coeff, expand_mesh_sequence=False)) if len(source_meshes) == 0: pass @@ -303,14 +535,21 @@ def assign(self, allow_missing_dofs=False): target_V = lhs_func.function_space() # Validate / Process subset. if subset is not None: - if subset is target_V.node_set: - # The whole set. + cell_set = target_V.mesh().cell_set + superset = getattr(subset, "superset", None) + if subset is target_V.node_set or subset is cell_set: + # The whole set, as `cell_subset("everywhere")` gives. subset = None - elif subset.superset is target_V.node_set: + elif superset is target_V.node_set: # op2.Subset of target_V.node_set pass + elif superset is cell_set: + # op2.Subset of the cells, e.g. mesh.cell_subset(id) + subset = _node_subset(target_V, subset) else: - raise ValueError(f"subset ({subset}) not a subset of target_V.node_set ({target_V.node_set})") + raise ValueError(f"subset ({subset}) is neither a subset of " + f"target_V.node_set ({target_V.node_set}) nor " + f"of the cells of its mesh ({cell_set})") source_meshes = set(extract_unique_domain(f) for f in funcs) if len(source_meshes) == 0: # Assign constants only. @@ -318,8 +557,12 @@ def assign(self, allow_missing_dofs=False): elif len(source_meshes) == 1: source_mesh, = source_meshes if target_mesh is source_mesh: - # Assign (co)functions from one mesh to the same mesh. - single_mesh_assign = True + # Assign (co)functions from one mesh to the same mesh. Two + # distinct spaces on it lay their nodes out differently, + # even when they share an element, so those are related by + # their sections like spaces on different meshes are. + single_mesh_assign = all(f.function_space() == lhs_func.function_space() + for f in funcs) else: # Assign (co)functions between a submesh and the parent or between two submeshes. single_mesh_assign = False @@ -335,30 +578,12 @@ def _assign_single_mesh(self, lhs_func, subset, funcs, operator): if assign_to_halos: indices = operator.attrgetter("indices") data_ro = operator.attrgetter("data_ro_with_halos") - values = operator.attrgetter("values_with_halo") else: indices = operator.attrgetter("owned_indices") data_ro = operator.attrgetter("data_ro") - values = operator.attrgetter("values") subset_indices = Ellipsis if subset is None else indices(subset) - def source_indices(f): - target_space = lhs_func.function_space() - target_map = target_space.cell_node_map() - source_map = f.function_space().cell_node_map() - if source_map is target_map: - # Source and target spaces have the same DoF ordering. - return subset_indices - else: - # Permute source indices into the target ordering. - size = target_space.dof_dset.total_size - perm = np.empty((size,), dtype=source_map.values.dtype) - np.put(perm, values(target_map), values(source_map)) - if not assign_to_halos: - perm = perm[:target_space.dof_dset.size] - return perm[subset_indices] - - func_data = np.array([data_ro(f.dat)[source_indices(f)] for f in funcs]) + func_data = np.array([data_ro(f.dat)[subset_indices] for f in funcs]) rvalue = self._compute_rvalue(func_data) self._assign_single_dat(lhs_func.dat, subset_indices, rvalue, assign_to_halos) if assign_to_halos: @@ -366,66 +591,168 @@ def source_indices(f): def _assign_multi_mesh(self, lhs_func, subset, funcs, operator, allow_missing_dofs): target_mesh = extract_unique_domain(lhs_func).topology - source_V, = set(f.function_space() for f in funcs) + source_spaces = set(f.function_space() for f in funcs) + if len(source_spaces) > 1: + # Every function is compatible with the assignee (checked at + # construction time) but not necessarily with one another, so + # each is related to the assignee by its own section SF and + # mapped into its layout before they are combined. + self._assign_multi_space(lhs_func, subset, funcs, allow_missing_dofs) + return + source_V, = source_spaces source_mesh = source_V.mesh().topology - if target_mesh.submesh_shares_distribution(source_mesh): + # Spaces on meshes that share their distribution are related by their + # entity maps, unless their elements lay the nodes out differently, in + # which case only their Sections relate them. + same_element = source_V.ufl_element() == lhs_func.ufl_element() + if target_mesh is not source_mesh and same_element and target_mesh.submesh_shares_distribution(source_mesh): self._assign_submesh(lhs_func, subset, funcs, operator, allow_missing_dofs) return - point_sf, target_is_submesh = _submesh_point_sf(target_mesh, source_mesh) - if point_sf is None: - raise NotImplementedError( - "Can only assign between a redistributed mesh and its parent" - ) - self._assign_redistributed(lhs_func, subset, funcs, point_sf, - target_is_submesh, allow_missing_dofs) + point_sf, target_is_leaf = _relate_to_target(target_mesh, lhs_func.ufl_element(), source_V) + self._assign_via_sections(lhs_func, subset, funcs, point_sf, + target_is_leaf, allow_missing_dofs) - def _assign_redistributed(self, lhs_func, subset, funcs, point_sf, - target_is_submesh, allow_missing_dofs): - """Assign between (co)functions on a redistributed submesh and its parent. + def _assign_via_sections(self, lhs_func, subset, funcs, point_sf, + target_is_leaf, allow_missing_dofs): + """Assign between (co)functions whose nodes are related by a section SF. - The nodes of the two spaces correspond one to one. The expression is - evaluated in the source layout. One communication then moves the - result into the target layout. + The nodes the two spaces have in common correspond one to one. The + expression is evaluated in the source layout. One communication then + moves the result into the target layout. + + ``target_is_leaf`` says which of the two carries the subset of the + nodes. Data travels from root to leaf by broadcast, and from leaf to + root by reduction. Only a reduction can miss nodes. """ target_V = lhs_func.function_space() source_V, = set(f.function_space() for f in funcs) - if target_is_submesh: + func_data = np.array([f.dat.data_ro_with_halos for f in funcs]) + source_data = self._compute_rvalue(func_data) + target_data, covered, assign_to_halos = self._transfer_via_section( + lhs_func, source_V, source_data, point_sf, target_is_leaf) + indices = self._covered_indices(target_V, subset, covered, assign_to_halos, allow_missing_dofs) + self._assign_single_dat(lhs_func.dat, indices, target_data[indices], assign_to_halos) + lhs_func.dat.halo_valid = assign_to_halos + + def _assign_multi_space(self, lhs_func, subset, funcs, allow_missing_dofs): + """Assign an expression combining functions from more than one + function space, each compatible with the assignee's element but not + necessarily with one another's. + + Every function is moved, unweighted, into the assignee's node + layout by its own section SF; the weighted combination then happens + in that common layout exactly as it would on a single mesh. + """ + target_mesh = extract_unique_domain(lhs_func).topology + target_V = lhs_func.function_space() + rows = [] + covered = None + assign_to_halos = True + for f in funcs: + source_V = f.function_space() + point_sf, target_is_leaf = _relate_to_target(target_mesh, lhs_func.ufl_element(), source_V) + data, cov, halo_ok = self._transfer_via_section( + lhs_func, source_V, f.dat.data_ro_with_halos, point_sf, target_is_leaf) + rows.append(data) + covered = cov if covered is None else (covered | cov) + assign_to_halos = assign_to_halos and halo_ok + target_data = self._compute_rvalue(np.array(rows)) + indices = self._covered_indices(target_V, subset, covered, assign_to_halos, allow_missing_dofs) + self._assign_single_dat(lhs_func.dat, indices, target_data[indices], assign_to_halos) + lhs_func.dat.halo_valid = assign_to_halos + + def _transfer_via_section(self, lhs_func, source_V, source_data, point_sf, target_is_leaf): + """Move data from a source layout into the assignee's, via a section SF. + + Parameters + ---------- + lhs_func : firedrake.function.Function or firedrake.cofunction.Cofunction + The function being assigned to; only its type and function space + are used. + source_V : firedrake.functionspaceimpl.WithGeometry + Function space ``source_data`` is laid out in. + source_data : numpy.ndarray + Data in ``source_V``'s with-halo layout, already combined if it + comes from more than one function. + point_sf : PETSc.SF + SF relating the assignee's mesh to ``source_V``'s, as returned by + `_relate_to_target`. + target_is_leaf : bool + Whether the assignee carries the subset of the nodes. + + Returns + ------- + tuple + The data moved into the assignee's with-halo layout, the boolean + array of which of those nodes received data, and whether the + halo of the result can be trusted. + + """ + target_V = lhs_func.function_space() + if target_is_leaf: root_V, leaf_V = source_V, target_V else: root_V, leaf_V = target_V, source_V section_sf, covered_roots = _make_section_sf(point_sf, root_V, leaf_V) - source_buffer = Function(source_V) - target_buffer = Function(target_V) - func_data = np.array([f.dat.data_ro_with_halos for f in funcs]) - source_buffer.dat.data_wo_with_halos[...] = self._compute_rvalue(func_data) + source_buffer = type(lhs_func)(source_V) + target_buffer = type(lhs_func)(target_V) + source_buffer.dat.data_wo_with_halos[...] = source_data mtype, _ = _get_mtype(source_buffer.dat) - source_data = source_buffer.dat.data_ro_with_halos - target_data = target_buffer.dat.data_wo_with_halos - if target_is_submesh: - section_sf.bcastBegin(mtype, source_data, target_data, MPI.REPLACE) - section_sf.bcastEnd(mtype, source_data, target_data, MPI.REPLACE) + source_buffer_data = source_buffer.dat.data_ro_with_halos + target_buffer_data = target_buffer.dat.data_wo_with_halos + if target_is_leaf: + section_sf.bcastBegin(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) + section_sf.bcastEnd(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) # Every node of a submesh, including its halo, has a counterpart # in the parent. - indices = Ellipsis if subset is None else subset.indices - assign_to_halos = True + covered = np.ones(target_buffer_data.shape[0], dtype=bool) + halos_valid = True else: - section_sf.reduceBegin(mtype, source_data, target_data, MPI.REPLACE) - section_sf.reduceEnd(mtype, source_data, target_data, MPI.REPLACE) - # Only the owned parent nodes that the submesh covers have been + section_sf.reduceBegin(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) + section_sf.reduceEnd(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) + # Only the owned parent nodes that the source covers have been # reduced into; the parent halo never is. - owned = covered_roots[:target_V.dof_dset.size] - comm = target_V.mesh().comm - if not comm.allreduce(owned.all(), op=MPI.LAND) and not allow_missing_dofs: - raise ValueError("Found assignee nodes with no matching assigner " - "nodes: run with `allow_missing_dofs=True`") - indices, = np.nonzero(owned) - if subset is not None: - indices = np.intersect1d(indices, subset.owned_indices) - target_data = target_buffer.dat.data_ro - assign_to_halos = False - self._assign_single_dat(lhs_func.dat, indices, target_data[indices], assign_to_halos) - lhs_func.dat.halo_valid = assign_to_halos + covered = np.zeros(target_buffer_data.shape[0], dtype=bool) + covered[:target_V.dof_dset.size] = covered_roots[:target_V.dof_dset.size] + halos_valid = False + return target_buffer.dat.data_ro_with_halos, covered, halos_valid + + def _covered_indices(self, target_V, subset, covered, assign_to_halos, allow_missing_dofs): + """Find the indices of the assignee's nodes that received data. + + Parameters + ---------- + target_V : firedrake.functionspaceimpl.WithGeometry + Function space of the function being assigned to. + subset : pyop2.types.set.Subset or None + Subset of the assignee's node set to restrict the assignment to. + covered : numpy.ndarray + Boolean array, in the assignee's with-halo layout, of which + nodes received data. + assign_to_halos : bool + Whether ``covered`` (and the halo of the assignee) can be trusted. + allow_missing_dofs : bool + Permit assignee nodes with no matching data, subject to + ``subset``, rather than raising. + + Returns + ------- + numpy.ndarray or Ellipsis + The indices, in the assignee's with-halo layout, to assign to. + + """ + if assign_to_halos: + return Ellipsis if subset is None else subset.indices + owned = covered[:target_V.dof_dset.size] + comm = target_V.mesh().comm + if not comm.allreduce(bool(owned.all()), op=MPI.LAND) and not allow_missing_dofs: + raise ValueError("Found assignee nodes with no matching assigner " + "nodes: run with `allow_missing_dofs=True`") + indices, = np.nonzero(owned) + if subset is not None: + indices = np.intersect1d(indices, subset.owned_indices) + return indices def _assign_submesh(self, lhs_func, subset, funcs, operator, allow_missing_dofs): target_mesh = extract_unique_domain(lhs_func) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 08ac21aeb0..6067244dad 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -243,9 +243,31 @@ def restricted_dofs(celem, felem): def get_restriction_indices(V, W): """Return the list of dofs in the space V such that W = V[indices]. """ + from firedrake.function import Function + if V.cell_node_map() is W.cell_node_map(): return numpy.arange(V.dof_dset.layout_vec.getSizes()[0], dtype=PETSc.IntType) + if V.extruded: + return _get_restriction_indices_extruded(V, W) + + # The restriction of an element keeps the nodes of the entities it does not + # drop, so numbering the nodes of V and assigning them into W leaves each + # node of W holding the number of the node of V it was taken from. + numbering = Function(V) + numbering.dat.data_wo_with_halos.flat[...] = numpy.arange(V.dof_count) + indices = numpy.concatenate([numpy.reshape(Function(Wsub).assign(numbering).dat.data_ro, -1) + for Wsub in W]) + return indices.astype(PETSc.IntType) + +def _get_restriction_indices_extruded(V, W): + """Return the list of dofs in the space V such that W = V[indices]. + + A base mesh point of an extruded mesh carries the nodes of a whole column + of entities, of which the restriction drops only some, so the two spaces + are not related by their `PETSc.Section` and the nodes are matched through + the cell node maps instead. + """ vdat = V.make_dat(val=numpy.arange(V.dof_count, dtype=PETSc.IntType)) wdats = [Wsub.make_dat(val=numpy.full((Wsub.dof_count,), -1, dtype=PETSc.IntType)) for Wsub in W] wdat = wdats[0] if len(W) == 1 else op2.MixedDat(wdats) diff --git a/tests/firedrake/regression/test_assign.py b/tests/firedrake/regression/test_assign.py index fd51a4fb77..702e7e0f3e 100644 --- a/tests/firedrake/regression/test_assign.py +++ b/tests/firedrake/regression/test_assign.py @@ -1,3 +1,4 @@ +import pytest from firedrake import * import numpy as np @@ -18,3 +19,30 @@ def test_single_mesh_mixed_assign(): assert np.allclose(w.subfunctions[0].dat.data_ro, [1.0, 2.0]) assert np.allclose(w.subfunctions[1].dat.data_ro, 3.0) + + +@pytest.mark.parallel(nprocs=[1, 3]) +def test_assign_cell_subset(): + """Assigning over a subset of the cells writes the nodes of those cells.""" + sentinel = -1.0 + mesh = UnitSquareMesh(6, 6) + x, y = SpatialCoordinate(mesh) + marker = Function(FunctionSpace(mesh, "DG", 0)).interpolate(conditional(x < 0.5, 1, 0)) + mesh.mark_entities(marker, 7) + + V = FunctionSpace(mesh, "CG", 2) + source = Function(V).interpolate(sin(3 * x)) + target = Function(V).assign(sentinel) + target.assign(source, subset=mesh.cell_subset(7)) + + written = target.dat.data_ro != sentinel + assert np.allclose(target.dat.data_ro[written], source.dat.data_ro[written]) + assert np.all(target.dat.data_ro[~written] == sentinel) + # The nodes of the unmarked cells are left alone, except where they sit on + # a cell of the subset too. + coords = Function(VectorFunctionSpace(mesh, "CG", 2)).interpolate(SpatialCoordinate(mesh)) + assert not written[coords.dat.data_ro[:, 0] > 0.5 + 1e-12].any() + # Which nodes the subset holds cannot depend on how the mesh is + # partitioned, so the count is the same however many ranks there are. + nwritten = written[:V.dof_dset.size] + assert mesh.comm.allreduce(int(nwritten.sum())) == 91 diff --git a/tests/firedrake/submesh/test_submesh_assign.py b/tests/firedrake/submesh/test_submesh_assign.py index ef29c538e7..fe79ee8030 100644 --- a/tests/firedrake/submesh/test_submesh_assign.py +++ b/tests/firedrake/submesh/test_submesh_assign.py @@ -2,6 +2,7 @@ import numpy as np from firedrake import * import finat +from functools import partial from os.path import abspath, dirname, join @@ -374,3 +375,181 @@ def test_submesh_assign_function_redistributed_subdomain(family, degree): g = Function(V).interpolate(expr) g.assign(f_sub, allow_missing_dofs=True) assert np.allclose(norm(g - f), 0) + + +@pytest.mark.parallel(nprocs=[1, 3]) +def test_submesh_assign_cell_subset_redistributed(): + # A cell subset of the parent mesh assigned from a redistributed submesh + # that does not share the parent's parallel distribution. + sentinel = -1.0 + mesh = UnitSquareMesh(6, 6) + x, y = SpatialCoordinate(mesh) + marker = Function(FunctionSpace(mesh, "DG", 0)).interpolate(conditional(x < 0.5, 1, 0)) + mesh.mark_entities(marker, 7) + + V = FunctionSpace(mesh, "CG", 2) + source = Function(V).interpolate(sin(3 * x)) + submesh = Submesh(mesh, subdomain_id=7, redistribute=True) + x_sub, _ = SpatialCoordinate(submesh) + f_sub = Function(FunctionSpace(submesh, "CG", 2)).interpolate(sin(3 * x_sub)) + composed = Function(V).assign(sentinel) + composed.assign(f_sub, subset=mesh.cell_subset(7), allow_missing_dofs=True) + written = composed.dat.data_ro != sentinel + assert np.allclose(composed.dat.data_ro[written], source.dat.data_ro[written]) + assert mesh.comm.allreduce(int(written[:V.dof_dset.size].sum())) == 91 + + +@pytest.mark.parametrize("family,degree", [("CG", 3), ("RT", 2)]) +@pytest.mark.parametrize("redistribute", [False, True]) +@pytest.mark.parallel(nprocs=[1, 3]) +def test_submesh_assign_composed_restrictions(family, degree, redistribute): + # Compose all three restrictions of the node layout at once, each of them + # applying to one side of the assignment only: the source is the whole of + # an element on the whole of the parent mesh, while the target restricts + # the mesh to a proper subdomain (optionally redistributing it), the + # boundary, and the element to the facets of the reference cell. + sentinel = -12345.0 + mesh = UnitSquareMesh(6, 6) + x, y = SpatialCoordinate(mesh) + DG0 = FunctionSpace(mesh, "DG", 0) + mesh.mark_entities(Function(DG0).interpolate(conditional(x < 0.5, 1, 0)), 111) + submesh = Submesh(mesh, subdomain_id=111, redistribute=redistribute) + + elem = finat.ufl.FiniteElement(family, mesh.ufl_cell(), degree) + V = FunctionSpace(mesh, elem) + V_sub = RestrictedFunctionSpace( + FunctionSpace(submesh, finat.ufl.RestrictedElement(elem, restriction_domain="facet")), + boundary_set={1}, + ) + + f = Function(V) + f.dat.data_wo[:] = 1.0 + np.arange(f.dat.data_ro.size) + + f_sub = Function(V_sub).assign(f) + g = Function(V).assign(sentinel) + g.assign(f_sub, allow_missing_dofs=True) + + covered = g.dat.data_ro != sentinel + assert mesh.comm.allreduce(int(covered.sum())) > 0 + assert np.allclose(g.dat.data_ro[covered], f.dat.data_ro[covered]) + + +@pytest.mark.parametrize("cell,family,degree", [ + ("triangle", "CG", 3), + ("triangle", "RT", 2), + ("quadrilateral", "Q", 3), +]) +@pytest.mark.parallel(nprocs=[1, 3]) +def test_assign_restricted_element_same_mesh(cell, family, degree): + # An element and its restriction lay their nodes out differently on the + # very same mesh; the two are related by their Sections alone. + sentinel = -12345.0 + mesh = UnitSquareMesh(6, 6, quadrilateral=(cell == "quadrilateral")) + elem = finat.ufl.FiniteElement(family, mesh.ufl_cell(), degree) + V = FunctionSpace(mesh, elem) + V_facet = FunctionSpace(mesh, finat.ufl.RestrictedElement(elem, restriction_domain="facet")) + + x, y = SpatialCoordinate(mesh) + expr = sin(3 * x) + 2 * cos(5 * y) + if V.value_shape: + expr = as_vector([sin(3 * x), cos(5 * y)]) + f = Function(V).interpolate(expr) + + # -- the restriction keeps the facet nodes of the full element. Compare + # against interpolation, which numbers the nodes of the restricted space + # without reference to the parent, so that a permutation of the nodes + # within an entity cannot go unnoticed. + f_facet = Function(V_facet).assign(f) + assert np.allclose(f_facet.dat.data_ro, Function(V_facet).interpolate(expr).dat.data_ro) + + # -- and back: the full element has interior nodes with no counterpart + with pytest.raises(ValueError): + Function(V).assign(f_facet) + g = Function(V).assign(sentinel) + g.assign(f_facet, allow_missing_dofs=True) + + covered = g.dat.data_ro != sentinel + assert mesh.comm.allreduce(int(covered[:V.dof_dset.size].sum())) == V_facet.dim() + assert np.allclose(g.dat.data_ro[covered], f.dat.data_ro[covered]) + + +@pytest.mark.parametrize("shape", ["vector", "symmetric"]) +@pytest.mark.parallel(nprocs=[1, 3]) +def test_assign_restricted_element_blocked(shape): + # A vector or tensor element holds a block of the nodes of a single + # sub-element, and UFL pushes the restriction inside the block. + mesh = UnitSquareMesh(4, 4) + x, y = SpatialCoordinate(mesh) + elem = finat.ufl.FiniteElement("CG", mesh.ufl_cell(), 3) + restricted = finat.ufl.RestrictedElement(elem, restriction_domain="facet") + if shape == "vector": + block, expr = finat.ufl.VectorElement, as_vector([sin(3 * x), cos(5 * y)]) + else: + block = partial(finat.ufl.TensorElement, symmetry=True) + expr = as_tensor([[sin(3 * x), cos(5 * y)], [cos(5 * y), x - y]]) + + V_facet = FunctionSpace(mesh, block(restricted)) + f = Function(FunctionSpace(mesh, block(elem))).interpolate(expr) + assert np.allclose(Function(V_facet).assign(f).dat.data_ro, + Function(V_facet).interpolate(expr).dat.data_ro) + + +@pytest.mark.parallel(nprocs=1) +def test_assign_incompatible_elements(): + mesh = UnitSquareMesh(2, 2) + elem = finat.ufl.FiniteElement("CG", mesh.ufl_cell(), 3) + + def facet_space(mesh, element): + return FunctionSpace(mesh, finat.ufl.RestrictedElement(element, restriction_domain="facet")) + + f = Function(FunctionSpace(mesh, "CG", 1)) + with pytest.raises(ValueError): + f.assign(Function(FunctionSpace(mesh, "RT", 1))) + # CG1 carries the nodes CG2 puts on the vertices, and drops the rest, but + # they are not the nodes of a common element, so this is not an assignment. + with pytest.raises(ValueError): + f.assign(Function(FunctionSpace(mesh, "CG", 2))) + + # Restricting to the facets of an interval leaves one node per vertex + # whatever the degree, so there the node counts alone cannot tell the two + # elements apart and only the element they restrict does. + interval = UnitIntervalMesh(4) + cells = [finat.ufl.FiniteElement("CG", interval.ufl_cell(), d) for d in (2, 3)] + with pytest.raises(ValueError): + Function(facet_space(interval, cells[1])).assign(Function(facet_space(interval, cells[0]))) + + # Blocks of different shape hold different numbers of nodes. + vector = Function(FunctionSpace(mesh, finat.ufl.VectorElement(elem, dim=2))) + for other in (finat.ufl.VectorElement(elem, dim=3), finat.ufl.TensorElement(elem), elem): + with pytest.raises(ValueError): + vector.assign(Function(FunctionSpace(mesh, other))) + + # A base mesh point of an extruded mesh carries a whole column of + # entities, of which the restriction drops only some. + extruded = ExtrudedMesh(UnitSquareMesh(2, 2, quadrilateral=True), 2) + q = finat.ufl.FiniteElement("Q", extruded.ufl_cell(), 3) + with pytest.raises(NotImplementedError): + Function(facet_space(extruded, q)).assign(Function(FunctionSpace(extruded, q))) + + +@pytest.mark.parallel(nprocs=[1, 3]) +def test_assign_multiple_source_spaces(): + # An interior and a facet restriction of one element are each compatible + # with the parent but not with one another, so each term of the sum is + # related to the assignee by its own section SF and moved into its + # layout before the two are added. + mesh = UnitSquareMesh(6, 6) + elem = finat.ufl.FiniteElement("CG", mesh.ufl_cell(), 3) + V = FunctionSpace(mesh, elem) + V_interior = FunctionSpace(mesh, finat.ufl.RestrictedElement(elem, restriction_domain="interior")) + V_facet = FunctionSpace(mesh, finat.ufl.RestrictedElement(elem, restriction_domain="facet")) + + x, y = SpatialCoordinate(mesh) + expr = sin(3 * x) + 2 * cos(5 * y) + f = Function(V).interpolate(expr) + f_interior = Function(V_interior).assign(f) + f_facet = Function(V_facet).assign(f) + + g = Function(V) + g.assign(f_interior + f_facet) + assert np.allclose(g.dat.data_ro, f.dat.data_ro) From d752060d2a6b51a05d6cee4d1e46c61008871c9d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Jul 2026 15:04:54 +0100 Subject: [PATCH 2/6] Apply suggestion from @pbrubeck --- firedrake/assign.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/assign.py b/firedrake/assign.py index e87cf0bcde..ab3be231b7 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -695,8 +695,8 @@ def _transfer_via_section(self, lhs_func, source_V, source_data, point_sf, targe root_V, leaf_V = target_V, source_V section_sf, covered_roots = _make_section_sf(point_sf, root_V, leaf_V) - source_buffer = type(lhs_func)(source_V) - target_buffer = type(lhs_func)(target_V) + source_buffer = Function(source_V) + target_buffer = Function(target_V) source_buffer.dat.data_wo_with_halos[...] = source_data mtype, _ = _get_mtype(source_buffer.dat) source_buffer_data = source_buffer.dat.data_ro_with_halos From f11d6f714ca0a0c4210e18d5cb5c421627793e0d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 8 Aug 2026 16:11:31 +0100 Subject: [PATCH 3/6] Apply the ASD-STE100 prose rules to the assign comments Split the sentences that stack clauses, in the section-SF assignment path and in the tests that cover it. --- firedrake/assign.py | 59 ++++++++++--------- firedrake/preconditioners/facet_split.py | 10 ++-- .../firedrake/submesh/test_submesh_assign.py | 30 +++++----- 3 files changed, 50 insertions(+), 49 deletions(-) diff --git a/firedrake/assign.py b/firedrake/assign.py index ab3be231b7..edbd91ca89 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -172,12 +172,13 @@ def _unrestricted(element): def _compatible_elements(target, source): """Whether functions in two elements share a node layout. - Two elements are compatible if they restrict a common element, and place - the same number of nodes on every entity of the reference cell up to the - entities one of them drops. Their nodes are then the same functionals - entity by entity, so the `PETSc.Section` of either function space - describes both and data moves between them without reference to the cell - node maps. + Two elements are compatible under two conditions. They must restrict a + common element. They must also place the same number of nodes on every + entity of the reference cell, up to the entities one of them drops. + + Their nodes are then the same functionals, entity by entity. The + `PETSc.Section` of either function space therefore describes both, and + data moves between them without reference to the cell node maps. Parameters ---------- @@ -206,8 +207,8 @@ def _compatible_elements(target, source): return False if isinstance(target.cell, TensorProductCell): # A base mesh point of an extruded mesh carries the nodes of a whole - # column of entities, of which a restriction may drop only some, and - # the Section counts the nodes on a point without saying which. + # column of entities. A restriction may drop only some of them. The + # Section counts the nodes on a point without saying which. raise NotImplementedError( "Assigning between an element and its restriction is not " "implemented on extruded meshes" @@ -216,9 +217,9 @@ def _compatible_elements(target, source): source_counts = _entity_dof_counts(source) if target_counts.keys() != source_counts.keys(): return False - # An entity either carries the same nodes in both elements, or is - # dropped by one of them; a partial overlap has no entity-wise - # correspondence and so cannot be expressed by the two Sections. + # An entity either carries the same nodes in both elements, or is dropped + # by one of them. A partial overlap has no entity-wise correspondence. + # The two Sections therefore cannot express it. return all(nt == source_counts[entity] or nt == 0 or source_counts[entity] == 0 for entity, nt in target_counts.items()) @@ -227,8 +228,8 @@ def _node_subset(V, cell_subset): """Find the nodes of the cells of a subset. A node on the boundary of the subset belongs to cells outside it too, and - is included: the subset selects the cells whose nodes are assigned, not - the nodes that no other cell shares. + is included. The subset selects the cells whose nodes are assigned. It + does not select the nodes that no other cell shares. Parameters ---------- @@ -251,10 +252,10 @@ def _node_subset(V, cell_subset): node_map = V.cell_node_map() if node_map is None: raise ValueError(f"Function space ({V}) has no nodes on the cells") - # A node is on the subset for every rank that shares it as soon as it is on - # the subset for one of them, which the cells known to a single rank do not - # say: a rank owning the node need not own, or even halo, a cell of the - # subset that carries it. + # A node is on the subset for every rank that shares it, as soon as it is + # on the subset for one of them. The cells known to a single rank do not + # say this. A rank owning the node need not own, or even halo, a cell of + # the subset that carries it. marker = op2.Dat(V.node_set, dtype=IntType) marker.data_wo_with_halos[np.unique(node_map.values_with_halo[cell_subset.indices])] = 1 marker.local_to_global_begin(op2.MAX) @@ -270,8 +271,8 @@ def _target_is_leaf(target, source): Data travels root to leaf by broadcast, which requires every leaf node to have a counterpart, and leaf to root by reduction, which does not. The - target can therefore be the leaf unless it carries nodes on an entity - that the source drops, and taking it to be the leaf whenever possible + target can therefore be the leaf, unless it carries nodes on an entity + that the source drops. Take it to be the leaf whenever possible, which keeps the halo of the assignee up to date. Parameters @@ -559,8 +560,8 @@ def assign(self, allow_missing_dofs=False): if target_mesh is source_mesh: # Assign (co)functions from one mesh to the same mesh. Two # distinct spaces on it lay their nodes out differently, - # even when they share an element, so those are related by - # their sections like spaces on different meshes are. + # even when they share an element. Their sections relate + # them, as they relate spaces on different meshes. single_mesh_assign = all(f.function_space() == lhs_func.function_space() for f in funcs) else: @@ -593,17 +594,17 @@ def _assign_multi_mesh(self, lhs_func, subset, funcs, operator, allow_missing_do target_mesh = extract_unique_domain(lhs_func).topology source_spaces = set(f.function_space() for f in funcs) if len(source_spaces) > 1: - # Every function is compatible with the assignee (checked at - # construction time) but not necessarily with one another, so - # each is related to the assignee by its own section SF and + # Every function is compatible with the assignee, which is checked + # at construction time, but not necessarily with one another. Each + # is therefore related to the assignee by its own section SF, and # mapped into its layout before they are combined. self._assign_multi_space(lhs_func, subset, funcs, allow_missing_dofs) return source_V, = source_spaces source_mesh = source_V.mesh().topology # Spaces on meshes that share their distribution are related by their - # entity maps, unless their elements lay the nodes out differently, in - # which case only their Sections relate them. + # entity maps. Elements that lay the nodes out differently are the + # exception: there, only the Sections relate them. same_element = source_V.ufl_element() == lhs_func.ufl_element() if target_mesh is not source_mesh and same_element and target_mesh.submesh_shares_distribution(source_mesh): self._assign_submesh(lhs_func, subset, funcs, operator, allow_missing_dofs) @@ -639,9 +640,9 @@ def _assign_multi_space(self, lhs_func, subset, funcs, allow_missing_dofs): function space, each compatible with the assignee's element but not necessarily with one another's. - Every function is moved, unweighted, into the assignee's node - layout by its own section SF; the weighted combination then happens - in that common layout exactly as it would on a single mesh. + Every function is moved, unweighted, into the assignee's node layout + by its own section SF. The weighted combination then happens in that + common layout, exactly as it would on a single mesh. """ target_mesh = extract_unique_domain(lhs_func).topology target_V = lhs_func.function_space() diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 6067244dad..4071803b79 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -251,8 +251,8 @@ def get_restriction_indices(V, W): return _get_restriction_indices_extruded(V, W) # The restriction of an element keeps the nodes of the entities it does not - # drop, so numbering the nodes of V and assigning them into W leaves each - # node of W holding the number of the node of V it was taken from. + # drop. Number the nodes of V and assign them into W. Each node of W then + # holds the number of the node of V it came from. numbering = Function(V) numbering.dat.data_wo_with_halos.flat[...] = numpy.arange(V.dof_count) indices = numpy.concatenate([numpy.reshape(Function(Wsub).assign(numbering).dat.data_ro, -1) @@ -264,9 +264,9 @@ def _get_restriction_indices_extruded(V, W): """Return the list of dofs in the space V such that W = V[indices]. A base mesh point of an extruded mesh carries the nodes of a whole column - of entities, of which the restriction drops only some, so the two spaces - are not related by their `PETSc.Section` and the nodes are matched through - the cell node maps instead. + of entities. The restriction drops only some of them. The two spaces are + therefore not related by their `PETSc.Section`, and the nodes are matched + through the cell node maps instead. """ vdat = V.make_dat(val=numpy.arange(V.dof_count, dtype=PETSc.IntType)) wdats = [Wsub.make_dat(val=numpy.full((Wsub.dof_count,), -1, dtype=PETSc.IntType)) for Wsub in W] diff --git a/tests/firedrake/submesh/test_submesh_assign.py b/tests/firedrake/submesh/test_submesh_assign.py index fe79ee8030..0d315f1195 100644 --- a/tests/firedrake/submesh/test_submesh_assign.py +++ b/tests/firedrake/submesh/test_submesh_assign.py @@ -403,11 +403,11 @@ def test_submesh_assign_cell_subset_redistributed(): @pytest.mark.parametrize("redistribute", [False, True]) @pytest.mark.parallel(nprocs=[1, 3]) def test_submesh_assign_composed_restrictions(family, degree, redistribute): - # Compose all three restrictions of the node layout at once, each of them - # applying to one side of the assignment only: the source is the whole of - # an element on the whole of the parent mesh, while the target restricts - # the mesh to a proper subdomain (optionally redistributing it), the - # boundary, and the element to the facets of the reference cell. + # Compose all three restrictions of the node layout at once. Each of them + # applies to one side of the assignment only. The source is the whole of + # an element on the whole of the parent mesh. The target restricts the mesh + # to a proper subdomain, optionally redistributing it. It then restricts to + # the boundary, and the element to the facets of the reference cell. sentinel = -12345.0 mesh = UnitSquareMesh(6, 6) x, y = SpatialCoordinate(mesh) @@ -457,8 +457,8 @@ def test_assign_restricted_element_same_mesh(cell, family, degree): # -- the restriction keeps the facet nodes of the full element. Compare # against interpolation, which numbers the nodes of the restricted space - # without reference to the parent, so that a permutation of the nodes - # within an entity cannot go unnoticed. + # without reference to the parent. A permutation of the nodes within an + # entity therefore cannot go unnoticed. f_facet = Function(V_facet).assign(f) assert np.allclose(f_facet.dat.data_ro, Function(V_facet).interpolate(expr).dat.data_ro) @@ -505,14 +505,14 @@ def facet_space(mesh, element): f = Function(FunctionSpace(mesh, "CG", 1)) with pytest.raises(ValueError): f.assign(Function(FunctionSpace(mesh, "RT", 1))) - # CG1 carries the nodes CG2 puts on the vertices, and drops the rest, but - # they are not the nodes of a common element, so this is not an assignment. + # CG1 carries the nodes CG2 puts on the vertices, and drops the rest. They + # are not the nodes of a common element, so this is not an assignment. with pytest.raises(ValueError): f.assign(Function(FunctionSpace(mesh, "CG", 2))) - # Restricting to the facets of an interval leaves one node per vertex - # whatever the degree, so there the node counts alone cannot tell the two - # elements apart and only the element they restrict does. + # Restricting to the facets of an interval leaves one node per vertex, + # whatever the degree. The node counts alone therefore cannot tell the two + # elements apart. Only the element they restrict can. interval = UnitIntervalMesh(4) cells = [finat.ufl.FiniteElement("CG", interval.ufl_cell(), d) for d in (2, 3)] with pytest.raises(ValueError): @@ -535,9 +535,9 @@ def facet_space(mesh, element): @pytest.mark.parallel(nprocs=[1, 3]) def test_assign_multiple_source_spaces(): # An interior and a facet restriction of one element are each compatible - # with the parent but not with one another, so each term of the sum is - # related to the assignee by its own section SF and moved into its - # layout before the two are added. + # with the parent, but not with one another. Each term of the sum is + # therefore related to the assignee by its own section SF. Both are moved + # into that layout before they are added. mesh = UnitSquareMesh(6, 6) elem = finat.ufl.FiniteElement("CG", mesh.ufl_cell(), 3) V = FunctionSpace(mesh, elem) From 13064c79ca1c434ded20a82c0994de5e3b5a3adc Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 8 Aug 2026 19:43:29 +0100 Subject: [PATCH 4/6] Assign between elements that differ only in their cell A submesh of lower dimension holds the entities of its parent up to its own dimension, and the point SF pairs them off by dimension. Compare the node counts by dimension there, because the two cells number their entities differently. Submesh() can then carry the coordinates of a curved or periodic parent onto a submesh of lower dimension. Measured on a unit square warped so that the boundary bends and the area changes: the perimeter of the submesh matches the ds integral of the parent to machine precision, at degree 2 and at degree 3. Degree 3 puts two nodes on an edge, which comes out reversed if the submesh orients that edge differently from its parent. --- firedrake/assign.py | 64 ++++++++++++++++++- firedrake/mesh.py | 5 -- .../firedrake/submesh/test_submesh_basics.py | 30 ++++++--- 3 files changed, 85 insertions(+), 14 deletions(-) diff --git a/firedrake/assign.py b/firedrake/assign.py index edbd91ca89..ef6345a619 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -169,6 +169,57 @@ def _unrestricted(element): return element +def _same_element_on_either_cell(target, source): + """Whether two elements differ only in the cell they are defined on. + + Parameters + ---------- + target : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned to. + source : finat.ufl.finiteelementbase.FiniteElementBase + Element of the function being assigned from. + + Returns + ------- + bool + Whether one element is the other, carried onto a different cell. + + """ + if target.cell == source.cell: + return target == source + # Carry the element of lower dimension up onto the other cell. A family + # that the lower cell does not name, such as "Q" on an interval, raises + # rather than compares, so the direction matters. + low, high = sorted((target, source), key=lambda e: e.cell.topological_dimension) + try: + return low.reconstruct(cell=high.cell) == high + except (ValueError, KeyError): + return False + + +def _dimension_dof_counts(element): + """Count the nodes an element places on each dimension of the reference cell. + + Parameters + ---------- + element : finat.ufl.finiteelementbase.FiniteElementBase + The UFL element. + + Returns + ------- + dict or None + The number of nodes on an entity of each dimension. `None` if some + dimension holds entities that carry different numbers of nodes, which + leaves the dimension alone unable to say how many a point carries. + + """ + counts = {} + for (dim, _), nodes in _entity_dof_counts(element).items(): + if counts.setdefault(dim, nodes) != nodes: + return None + return counts + + def _compatible_elements(target, source): """Whether functions in two elements share a node layout. @@ -203,8 +254,19 @@ def _compatible_elements(target, source): and _compatible_elements(target.sub_elements[0], source.sub_elements[0])) # Equal node counts on an entity do not by themselves make two nodes the # same functional, so the elements must restrict a common element. - if _unrestricted(target) != _unrestricted(source): + if not _same_element_on_either_cell(_unrestricted(target), _unrestricted(source)): return False + if target.cell != source.cell: + # The two cells have different entities, so the entity numbers can not + # be compared. A submesh of lower dimension holds the entities of its + # parent up to its own dimension, and the point SF pairs them off by + # dimension, so the counts must agree dimension by dimension. + target_counts = _dimension_dof_counts(target) + source_counts = _dimension_dof_counts(source) + if target_counts is None or source_counts is None: + return False + shared = min(len(target_counts), len(source_counts)) + return all(target_counts[dim] == source_counts[dim] for dim in range(shared)) if isinstance(target.cell, TensorProductCell): # A base mesh point of an extruded mesh carries the nodes of a whole # column of entities. A restriction may drop only some of them. The diff --git a/firedrake/mesh.py b/firedrake/mesh.py index ba3ddaff5d..f1355be504 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -5233,11 +5233,6 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig if mesh.coordinates.ufl_element() != plex_element: # The parent coordinates are not carried by the plex (e.g. the parent # is curved or periodic), so they must be transferred onto the submesh. - if submesh.ufl_cell() != mesh.ufl_cell(): - raise NotImplementedError( - "Can only transfer the coordinates of a curved or periodic mesh " - "onto a submesh of the same dimension" - ) V = mesh.coordinates.function_space().reconstruct(mesh=submesh) coordinates = function.Function(V).assign(mesh.coordinates) submesh = Mesh(coordinates, name=name) diff --git a/tests/firedrake/submesh/test_submesh_basics.py b/tests/firedrake/submesh/test_submesh_basics.py index cea5acaef4..4054a3786e 100644 --- a/tests/firedrake/submesh/test_submesh_basics.py +++ b/tests/firedrake/submesh/test_submesh_basics.py @@ -32,11 +32,17 @@ def test_submesh_redistribute_codim(): def _curved_mesh(nx=4, degree=2): - """Build a unit square whose curved coordinates the plex can not carry.""" + """Build a unit square whose curved coordinates the plex can not carry. + + The warp bends the boundary as well as the interior, and it does not + preserve area. A warp that vanishes on the boundary would leave the + perimeter at four, and a shear would leave every area unchanged, so + either one would pass these tests without carrying any coordinates. + """ mesh = UnitSquareMesh(nx, nx) V = VectorFunctionSpace(mesh, "CG", degree) x, y = SpatialCoordinate(mesh) - coordinates = Function(V).interpolate(as_vector([x + 0.15 * sin(pi * y) * x * (1 - x), y])) + coordinates = Function(V).interpolate(as_vector([x, y * (1 + 0.3 * sin(2 * pi * x))])) return Mesh(coordinates) @@ -69,12 +75,20 @@ def test_submesh_affine_coordinates(): @pytest.mark.parallel([1, 2, 3]) -def test_submesh_curved_codim(): - # A parent and a submesh of lower dimension share only some of the nodes - # on a parent cell. The cell node maps can not select those nodes. - mesh = _curved_mesh() - with pytest.raises(NotImplementedError): - Submesh(mesh, 1, "on_boundary", label_name="exterior_facets") +@pytest.mark.parametrize("degree", [2, 3]) +def test_submesh_curved_codim(degree): + # A submesh of lower dimension takes the curved coordinates of its parent + # entity by entity. Degree 3 puts two nodes on an edge, so it fails if the + # submesh orients that edge differently from the parent. + mesh = _curved_mesh(degree=degree) + submesh = Submesh(mesh, 1, "on_boundary", label_name="exterior_facets") + assert submesh.coordinates.ufl_element() == \ + mesh.coordinates.ufl_element().reconstruct(cell=submesh.ufl_cell()) + + perimeter = assemble(Constant(1.0) * dx(submesh)) + assert np.isclose(perimeter, assemble(Constant(1.0) * ds(mesh))) + # The boundary of an affine parent would measure exactly four. + assert not np.isclose(perimeter, 4.0) @pytest.mark.parallel([1, 2]) From f129f16c993073b6264179763eee7cee071900ff Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 8 Aug 2026 20:13:50 +0100 Subject: [PATCH 5/6] Name the sets that an assignment subset can be The subset an assignment is given is the node set, the cell set, or a subset of either, and each case reaches the assignment differently. Give that classification a name of its own, so that the branches say which set they match rather than carrying it in a comment beside them. Name the node numbers that get_restriction_indices walks through assign, and the count of dimensions that two cells share. Co-Authored-By: Claude Opus 5 --- firedrake/assign.py | 83 ++++++++++++------- firedrake/preconditioners/facet_split.py | 9 +- .../firedrake/submesh/test_submesh_basics.py | 7 +- 3 files changed, 61 insertions(+), 38 deletions(-) diff --git a/firedrake/assign.py b/firedrake/assign.py index ef6345a619..b9b33acd75 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -187,9 +187,9 @@ def _same_element_on_either_cell(target, source): """ if target.cell == source.cell: return target == source - # Carry the element of lower dimension up onto the other cell. A family - # that the lower cell does not name, such as "Q" on an interval, raises - # rather than compares, so the direction matters. + # A family that the lower cell does not name, such as "Q" on an interval, + # raises rather than compares. The reconstruction must therefore go up, + # onto the cell of higher dimension. low, high = sorted((target, source), key=lambda e: e.cell.topological_dimension) try: return low.reconstruct(cell=high.cell) == high @@ -259,14 +259,15 @@ def _compatible_elements(target, source): if target.cell != source.cell: # The two cells have different entities, so the entity numbers can not # be compared. A submesh of lower dimension holds the entities of its - # parent up to its own dimension, and the point SF pairs them off by - # dimension, so the counts must agree dimension by dimension. + # parent up to its own dimension. The point SF pairs them off by + # dimension. The counts must therefore agree dimension by dimension. target_counts = _dimension_dof_counts(target) source_counts = _dimension_dof_counts(source) if target_counts is None or source_counts is None: return False - shared = min(len(target_counts), len(source_counts)) - return all(target_counts[dim] == source_counts[dim] for dim in range(shared)) + shared_dimensions = min(len(target_counts), len(source_counts)) + return all(target_counts[dim] == source_counts[dim] + for dim in range(shared_dimensions)) if isinstance(target.cell, TensorProductCell): # A base mesh point of an extruded mesh carries the nodes of a whole # column of entities. A restriction may drop only some of them. The @@ -282,8 +283,9 @@ def _compatible_elements(target, source): # An entity either carries the same nodes in both elements, or is dropped # by one of them. A partial overlap has no entity-wise correspondence. # The two Sections therefore cannot express it. - return all(nt == source_counts[entity] or nt == 0 or source_counts[entity] == 0 - for entity, nt in target_counts.items()) + return all(target_nodes == source_counts[entity] + or target_nodes == 0 or source_counts[entity] == 0 + for entity, target_nodes in target_counts.items()) def _node_subset(V, cell_subset): @@ -328,6 +330,43 @@ def _node_subset(V, cell_subset): return op2.Subset(V.node_set, nodes) +def _assigned_nodes(V, subset): + """Find the nodes of a space that an assignment writes to. + + Parameters + ---------- + V : firedrake.functionspaceimpl.WithGeometry + Function space of the function being assigned to. + subset : pyop2.types.set.Set or pyop2.types.set.Subset or None + The set to assign over. This is the node set of ``V``, the cell set of + the mesh of ``V``, a `pyop2.types.set.Subset` of either, or `None`. + + Returns + ------- + pyop2.types.set.Subset or None + The nodes of ``V`` to assign, or `None` for all of them. + + Raises + ------ + ValueError + If the subset belongs to neither the nodes of ``V`` nor the cells of + the mesh of ``V``. + + """ + all_nodes = V.node_set + all_cells = V.mesh().cell_set + if subset is None or subset is all_nodes or subset is all_cells: + return None + superset = getattr(subset, "superset", None) + if superset is all_nodes: + return subset + if superset is all_cells: + return _node_subset(V, subset) + raise ValueError(f"subset ({subset}) is neither a subset of the nodes of " + f"the function space ({all_nodes}) nor of the cells of " + f"its mesh ({all_cells})") + + def _target_is_leaf(target, source): """Whether the target element's nodes are the subset of the two. @@ -596,23 +635,7 @@ def assign(self, allow_missing_dofs=False): for lhs_func, subset, *funcs in zip(self._assignee.subfunctions, self._subset, *(f.subfunctions for f in self._functions)): target_mesh = extract_unique_domain(lhs_func) target_V = lhs_func.function_space() - # Validate / Process subset. - if subset is not None: - cell_set = target_V.mesh().cell_set - superset = getattr(subset, "superset", None) - if subset is target_V.node_set or subset is cell_set: - # The whole set, as `cell_subset("everywhere")` gives. - subset = None - elif superset is target_V.node_set: - # op2.Subset of target_V.node_set - pass - elif superset is cell_set: - # op2.Subset of the cells, e.g. mesh.cell_subset(id) - subset = _node_subset(target_V, subset) - else: - raise ValueError(f"subset ({subset}) is neither a subset of " - f"target_V.node_set ({target_V.node_set}) nor " - f"of the cells of its mesh ({cell_set})") + subset = _assigned_nodes(target_V, subset) source_meshes = set(extract_unique_domain(f) for f in funcs) if len(source_meshes) == 0: # Assign constants only. @@ -620,10 +643,10 @@ def assign(self, allow_missing_dofs=False): elif len(source_meshes) == 1: source_mesh, = source_meshes if target_mesh is source_mesh: - # Assign (co)functions from one mesh to the same mesh. Two - # distinct spaces on it lay their nodes out differently, - # even when they share an element. Their sections relate - # them, as they relate spaces on different meshes. + # Two distinct spaces on one mesh lay their nodes out + # differently, even when they share an element. Their + # Sections relate them, as they relate spaces on different + # meshes. single_mesh_assign = all(f.function_space() == lhs_func.function_space() for f in funcs) else: diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 4071803b79..551698d886 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -251,11 +251,10 @@ def get_restriction_indices(V, W): return _get_restriction_indices_extruded(V, W) # The restriction of an element keeps the nodes of the entities it does not - # drop. Number the nodes of V and assign them into W. Each node of W then - # holds the number of the node of V it came from. - numbering = Function(V) - numbering.dat.data_wo_with_halos.flat[...] = numpy.arange(V.dof_count) - indices = numpy.concatenate([numpy.reshape(Function(Wsub).assign(numbering).dat.data_ro, -1) + # drop. An assignment from V to W therefore pairs those nodes off. + node_numbers = Function(V) + node_numbers.dat.data_wo_with_halos.flat[...] = numpy.arange(V.dof_count) + indices = numpy.concatenate([numpy.reshape(Function(Wsub).assign(node_numbers).dat.data_ro, -1) for Wsub in W]) return indices.astype(PETSc.IntType) diff --git a/tests/firedrake/submesh/test_submesh_basics.py b/tests/firedrake/submesh/test_submesh_basics.py index 4054a3786e..1743ac2d01 100644 --- a/tests/firedrake/submesh/test_submesh_basics.py +++ b/tests/firedrake/submesh/test_submesh_basics.py @@ -35,9 +35,10 @@ def _curved_mesh(nx=4, degree=2): """Build a unit square whose curved coordinates the plex can not carry. The warp bends the boundary as well as the interior, and it does not - preserve area. A warp that vanishes on the boundary would leave the - perimeter at four, and a shear would leave every area unchanged, so - either one would pass these tests without carrying any coordinates. + preserve area. Both properties are necessary. A warp that vanishes on + the boundary leaves the perimeter at four, and a shear leaves every + area unchanged. Either one passes these tests on a plex that carries + no coordinates at all. """ mesh = UnitSquareMesh(nx, nx) V = VectorFunctionSpace(mesh, "CG", degree) From 2b661167905d0478a72be642560bf54df9cd71e0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 22:14:22 +0100 Subject: [PATCH 6/6] Derive the assigned node counts, and keep them real The covered mask no longer needs its halo sliced off: the point SF addresses each root on its owner, so a covered root is one this rank owns. Replace the magic node counts in the subset tests with the quantity they stand for, and take the real part of the node numbers that FacetSplitPC routes through a ScalarType Function, which is complex in complex mode. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01GRZG3iuNzgXQWGiuzrjKQo --- firedrake/assign.py | 13 +++++-------- firedrake/preconditioners/facet_split.py | 6 ++++-- tests/firedrake/regression/test_assign.py | 10 +++------- tests/firedrake/submesh/test_submesh_assign.py | 4 +++- 4 files changed, 15 insertions(+), 18 deletions(-) diff --git a/firedrake/assign.py b/firedrake/assign.py index b9b33acd75..84bc7a9aad 100644 --- a/firedrake/assign.py +++ b/firedrake/assign.py @@ -707,8 +707,10 @@ def _assign_via_sections(self, lhs_func, subset, funcs, point_sf, moves the result into the target layout. ``target_is_leaf`` says which of the two carries the subset of the - nodes. Data travels from root to leaf by broadcast, and from leaf to - root by reduction. Only a reduction can miss nodes. + nodes. Data travels from root to leaf by broadcast, which reaches + every leaf, and from leaf to root by reduction, which reaches only + the owned roots that a leaf points at. Only a reduction can therefore + miss nodes, and only a reduction leaves the halo stale. """ target_V = lhs_func.function_space() source_V, = set(f.function_space() for f in funcs) @@ -790,17 +792,12 @@ def _transfer_via_section(self, lhs_func, source_V, source_data, point_sf, targe if target_is_leaf: section_sf.bcastBegin(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) section_sf.bcastEnd(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) - # Every node of a submesh, including its halo, has a counterpart - # in the parent. covered = np.ones(target_buffer_data.shape[0], dtype=bool) halos_valid = True else: section_sf.reduceBegin(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) section_sf.reduceEnd(mtype, source_buffer_data, target_buffer_data, MPI.REPLACE) - # Only the owned parent nodes that the source covers have been - # reduced into; the parent halo never is. - covered = np.zeros(target_buffer_data.shape[0], dtype=bool) - covered[:target_V.dof_dset.size] = covered_roots[:target_V.dof_dset.size] + covered = covered_roots halos_valid = False return target_buffer.dat.data_ro_with_halos, covered, halos_valid diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 551698d886..11d4a33e6e 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -251,12 +251,14 @@ def get_restriction_indices(V, W): return _get_restriction_indices_extruded(V, W) # The restriction of an element keeps the nodes of the entities it does not - # drop. An assignment from V to W therefore pairs those nodes off. + # drop. An assignment from V to W therefore pairs those nodes off. The + # numbers ride in a ScalarType Function, which holds every integer below + # 2**53 exactly and, in complex mode, holds it in the real part. node_numbers = Function(V) node_numbers.dat.data_wo_with_halos.flat[...] = numpy.arange(V.dof_count) indices = numpy.concatenate([numpy.reshape(Function(Wsub).assign(node_numbers).dat.data_ro, -1) for Wsub in W]) - return indices.astype(PETSc.IntType) + return indices.real.astype(PETSc.IntType) def _get_restriction_indices_extruded(V, W): diff --git a/tests/firedrake/regression/test_assign.py b/tests/firedrake/regression/test_assign.py index 702e7e0f3e..b759848d93 100644 --- a/tests/firedrake/regression/test_assign.py +++ b/tests/firedrake/regression/test_assign.py @@ -38,11 +38,7 @@ def test_assign_cell_subset(): written = target.dat.data_ro != sentinel assert np.allclose(target.dat.data_ro[written], source.dat.data_ro[written]) assert np.all(target.dat.data_ro[~written] == sentinel) - # The nodes of the unmarked cells are left alone, except where they sit on - # a cell of the subset too. + # The marked cells cover the closed left half, so a node is written + # exactly when it lies there. That does not depend on the partition. coords = Function(VectorFunctionSpace(mesh, "CG", 2)).interpolate(SpatialCoordinate(mesh)) - assert not written[coords.dat.data_ro[:, 0] > 0.5 + 1e-12].any() - # Which nodes the subset holds cannot depend on how the mesh is - # partitioned, so the count is the same however many ranks there are. - nwritten = written[:V.dof_dset.size] - assert mesh.comm.allreduce(int(nwritten.sum())) == 91 + assert np.array_equal(written, coords.dat.data_ro[:, 0] <= 0.5 + 1e-12) diff --git a/tests/firedrake/submesh/test_submesh_assign.py b/tests/firedrake/submesh/test_submesh_assign.py index b19447e0c1..b8c23e84db 100644 --- a/tests/firedrake/submesh/test_submesh_assign.py +++ b/tests/firedrake/submesh/test_submesh_assign.py @@ -396,7 +396,9 @@ def test_submesh_assign_cell_subset_redistributed(): composed.assign(f_sub, subset=mesh.cell_subset(7), allow_missing_dofs=True) written = composed.dat.data_ro != sentinel assert np.allclose(composed.dat.data_ro[written], source.dat.data_ro[written]) - assert mesh.comm.allreduce(int(written[:V.dof_dset.size].sum())) == 91 + # The submesh covers the marked cells, so its nodes are exactly the ones + # written. Which they are cannot depend on how the mesh is partitioned. + assert mesh.comm.allreduce(int(written[:V.dof_dset.size].sum())) == f_sub.function_space().dim() @pytest.mark.parametrize("family,degree", [("CG", 3), ("RT", 2)])