Unify assign across Submesh, RestrictedFunctionSpace and RestrictedElement

[pbrubeck]

Stack

This is PR 3 of 3: #5288 (adaptive-multigrid) -> #5270 (mesh-redistribution) -> this PR (assign-submesh-restricted).

Summary

Unifies how Function.assign moves data between non-matching DOF layouts. A Submesh, a RestrictedFunctionSpace and a RestrictedElement all restrict the nodes of a space, and each was handled by its own machinery: entity maps, a cell_node_map permutation, and nothing at all respectively. They are now one mechanism — a PETSc.SF over the two spaces' PETSc.Sections.

• RestrictedElement can be assigned to and from its parent element, which previously raised outright. Compatibility (_compatible_elements) is deduced from the UFL elements and their entity ids alone: the two must 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.
• The cell_node_map permutation is gone. Two distinct spaces on one mesh are related by an identity point SF and their sections, exactly as a submesh and its parent are related by the submesh point SF and theirs (_assign_via_sections, replacing _assign_redistributed).
• assign accepts a subset of the cells, not just of the nodes (_node_subset), so a partial assignment can be expressed directly in terms of a mesh.cell_subset(...).
• FacetSplitPC builds its restriction indices with assign instead of a hand-written C kernel over a PermutedMap, keeping the kernel as a fallback for extruded meshes.
• Combining functions from more than one distinct-but-compatible function space in a single expression now works (e.g. g.assign(f_interior + f_facet) for complementary restrictions of one element): each function is related to the assignee by its own section SF and moved, unweighted, into the assignee's layout before the weighted combination happens there. Previously this raised a confusing tuple-unpack ValueError.

Requires firedrakeproject/fiat#267: a restricted element reported no entity permutations, so a restricted space could not be laid out consistently with the space it restricts, and the nodes on an edge carrying two or more of them came out transposed.

Not implemented

Extruded meshes. The classic ExtrudedMesh has only the flat base DMPlex, so one base point carries a whole column of entities of which a restriction drops only some; a Section counts the nodes on a point without saying which they are. _compatible_elements raises NotImplementedError there, and FacetSplitPC falls back to the kernel.

Testing

tests/firedrake/submesh/test_submesh_assign.py gains tests for an element against its restriction (CG3, RT2 on triangles, Q3 on quadrilaterals), for the multi-space sum above (test_assign_multiple_source_spaces, nprocs=[1, 3]), for the remaining genuinely-incompatible pairs, and for the extruded case being refused. They compare against interpolation into the restricted space, which numbers its nodes without reference to the parent, so a permutation of the nodes within an entity cannot pass unnoticed.

Also included are tests composing all three restrictions at once (a proper-subdomain Submesh, redistributed or not, a RestrictedFunctionSpace over part of the boundary, and a facet-only RestrictedElement), and a cell-subset assignment from a redistributed Submesh.

Ran tests/firedrake/regression/test_assign.py and tests/firedrake/submesh/test_submesh_assign.py at nprocs 1, 2 and 3.

Warning

test_submesh_assign_cell_subset_redistributed[nprocs=3] fails when run as part of the full test_submesh_assign.py file (though it passes in isolation): after several other tests in the file construct a Submesh(..., redistribute=True), the section-SF coverage computed for a later redistributed submesh silently undercounts by exactly nprocs nodes on one rank. This reproduces identically against this PR's base (pbrubeck/mesh-redistribution, unmodified) with no code from this PR involved, so it is rooted in Submesh(redistribute=True)'s point-SF construction in firedrake/mesh.py (from #5270), not in this PR's diff. Needs investigation in #5270 before this stack is CI-clean.

🤖 Generated with Claude Code

Warning

Contains a DROP BEFORE MERGE commit pinning CI to the FIAT branch.