From b77709d0629e24e8eae853caaf9c2b1c9deb8047 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 4 Dec 2025 16:59:24 +0000 Subject: [PATCH 1/8] BUG: Fix symmetry in supermeshing --- firedrake/supermeshing.py | 42 +++++++++++++++++++-------------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index ec037aa283..e09f4b22f0 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -68,14 +68,6 @@ def assemble_mixed_mass_matrix(V_A, V_B): if len(V_A) > 1 or len(V_B) > 1: raise NotImplementedError("Sorry, only implemented for non-mixed spaces") - if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping() != "identity": - msg = """ -Sorry, only implemented for affine maps for now. To do non-affine, we'd need to -import much more of the assembly engine of UFL/TSFC/etc to do the assembly on -each supermesh cell. -""" - raise NotImplementedError(msg) - mesh_A = V_A.mesh() mesh_B = V_B.mesh() @@ -116,15 +108,23 @@ def likely(cell_A): def likely(cell_A): return cell_map[cell_A] - assert V_A.value_size == V_B.value_size - orig_value_size = V_A.value_size - if V_A.value_size > 1: + assert V_A.block_size == V_B.block_size + orig_block_size = V_A.block_size + if V_A.block_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements[0]) - if V_B.value_size > 1: + if V_B.block_size > 1: V_B = firedrake.FunctionSpace(mesh_B, V_B.ufl_element().sub_elements[0]) - assert V_A.value_size == 1 - assert V_B.value_size == 1 + if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping() != "identity": + msg = """ +Sorry, only implemented for affine maps for now. To do non-affine, we'd need to +import much more of the assembly engine of UFL/TSFC/etc to do the assembly on +each supermesh cell. +""" + raise NotImplementedError(msg) + + assert V_A.block_size == 1 + assert V_B.block_size == 1 preallocator = PETSc.Mat().create(comm=mesh_A._comm) preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) @@ -155,7 +155,7 @@ def likely(cell_A): onnz = numpy.repeat(onnz, cset.cdim) preallocator.destroy() - assert V_A.value_size == V_B.value_size + assert V_A.block_size == V_B.block_size rdim = V_B.dof_dset.cdim cdim = V_A.dof_dset.cdim @@ -445,16 +445,16 @@ def likely(cell_A): lib.restype = ctypes.c_int ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat) - if orig_value_size == 1: + if orig_block_size == 1: return mat else: (lrows, grows), (lcols, gcols) = mat.getSizes() - lrows *= orig_value_size - grows *= orig_value_size - lcols *= orig_value_size - gcols *= orig_value_size + lrows *= orig_block_size + grows *= orig_block_size + lcols *= orig_block_size + gcols *= orig_block_size size = ((lrows, grows), (lcols, gcols)) - context = BlockMatrix(mat, orig_value_size) + context = BlockMatrix(mat, orig_block_size) blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm) blockmat.setUp() return blockmat From 8526e6689dfca6ba690a9ec5aeb94d01021fddbe Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 5 Dec 2025 12:42:04 +0000 Subject: [PATCH 2/8] Fix cross mesh interpolation for symmetry=True --- firedrake/interpolation.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 7423eb548e..f94e2ad7eb 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -472,7 +472,11 @@ def __init__( raise NotImplementedError("freeze_expr not implemented") if bcs: raise NotImplementedError("bcs not implemented") - if V.ufl_element().mapping() != "identity": + + target_element = V.ufl_element() + if not ((isinstance(target_element, finat.ufl.MixedElement) + and all(sub.mapping() == "identity" for sub in target_element.sub_elements)) + or target_element.mapping() == "identity"): # Identity mapping between reference cell and physical coordinates # implies point evaluation nodes. A more general version would # require finding the global coordinates of all quadrature points @@ -551,7 +555,8 @@ def __init__( elif len(shape) == 1: fs_type = partial(firedrake.VectorFunctionSpace, dim=shape[0]) else: - fs_type = partial(firedrake.TensorFunctionSpace, shape=shape) + symmetry = V_dest.ufl_element().symmetry() + fs_type = partial(firedrake.TensorFunctionSpace, shape=shape, symmetry=symmetry) P0DG_vom = fs_type(self.vom_dest_node_coords_in_src_mesh, "DG", 0) self.point_eval_interpolate = Interpolate(self.expr_renumbered, P0DG_vom) # The parallel decomposition of the nodes of V_dest in the DESTINATION From 54e8fc5be1528878781cf366b8f4421da6651ace Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 5 Dec 2025 15:41:39 +0000 Subject: [PATCH 3/8] test --- firedrake/supermeshing.py | 21 +++++++++++++++++-- .../regression/test_interpolate_cross_mesh.py | 7 ++++--- .../supermesh/test_galerkin_projection.py | 5 ++++- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index e09f4b22f0..a6f03e3596 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -20,15 +20,17 @@ from pyop2.compilation import load from pyop2.mpi import COMM_SELF from pyop2.utils import get_petsc_dir +from collections import defaultdict __all__ = ["assemble_mixed_mass_matrix", "intersection_finder"] class BlockMatrix(object): - def __init__(self, mat, dimension): + def __init__(self, mat, dimension, symmetry=None): self.mat = mat self.dimension = dimension + self.symmetry = symmetry or {} def mult(self, mat, x, y): sizes = self.mat.getSizes() @@ -41,6 +43,8 @@ def mult(self, mat, x, y): xi = PETSc.Vec().createWithArray(xa, size=sizes[1], comm=x.comm) yi = PETSc.Vec().createWithArray(ya, size=sizes[0], comm=y.comm) self.mat.mult(xi, yi) + if i in self.symmetry: + yi.array[:] *= self.symmetry[i] y.array[start::stride] = yi.array_r def multTranspose(self, mat, x, y): @@ -54,6 +58,8 @@ def multTranspose(self, mat, x, y): xi = PETSc.Vec().createWithArray(xa, size=sizes[0], comm=x.comm) yi = PETSc.Vec().createWithArray(ya, size=sizes[1], comm=y.comm) self.mat.multTranspose(xi, yi) + if i in self.symmetry: + yi.array[:] *= self.symmetry[i] y.array[start::stride] = yi.array_r @@ -110,6 +116,17 @@ def likely(cell_A): assert V_A.block_size == V_B.block_size orig_block_size = V_A.block_size + + symmetry = V_A.ufl_element().symmetry() + assert symmetry == V_B.ufl_element().symmetry() + if symmetry is not None: + multiplicity = defaultdict(int) + for idx in numpy.ndindex(V_A.value_shape): + idx = symmetry.get(idx, idx) + multiplicity[idx] += 1 + symmetry = {i: multiplicity[idx] for i, idx in enumerate(multiplicity) + if multiplicity[idx] != 1} + if V_A.block_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements[0]) if V_B.block_size > 1: @@ -454,7 +471,7 @@ def likely(cell_A): lcols *= orig_block_size gcols *= orig_block_size size = ((lrows, grows), (lcols, gcols)) - context = BlockMatrix(mat, orig_block_size) + context = BlockMatrix(mat, orig_block_size, symmetry=symmetry) blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm) blockmat.setUp() return blockmat diff --git a/tests/firedrake/regression/test_interpolate_cross_mesh.py b/tests/firedrake/regression/test_interpolate_cross_mesh.py index f6c22e48e8..311b0e2568 100644 --- a/tests/firedrake/regression/test_interpolate_cross_mesh.py +++ b/tests/firedrake/regression/test_interpolate_cross_mesh.py @@ -396,11 +396,12 @@ def test_exact_refinement(): ) -def test_interpolate_unitsquare_tfs_shape(): +@pytest.mark.parametrize("shape,symmetry", [((1, 2, 3), None), ((3, 3), True)]) +def test_interpolate_unitsquare_tfs_shape(shape, symmetry): m_src = UnitSquareMesh(2, 3) m_dest = UnitSquareMesh(3, 5, quadrilateral=True) - V_src = TensorFunctionSpace(m_src, "CG", 3, shape=(1, 2, 3)) - V_dest = TensorFunctionSpace(m_dest, "CG", 4, shape=(1, 2, 3)) + V_src = TensorFunctionSpace(m_src, "CG", 3, shape=shape, symmetry=symmetry) + V_dest = TensorFunctionSpace(m_dest, "CG", 4, shape=shape, symmetry=symmetry) f_src = Function(V_src) assemble(interpolate(f_src, V_dest)) diff --git a/tests/firedrake/supermesh/test_galerkin_projection.py b/tests/firedrake/supermesh/test_galerkin_projection.py index 2770bfa7d2..cc20b8b3b9 100644 --- a/tests/firedrake/supermesh/test_galerkin_projection.py +++ b/tests/firedrake/supermesh/test_galerkin_projection.py @@ -2,6 +2,7 @@ from firedrake.petsc import DEFAULT_DIRECT_SOLVER_PARAMETERS from firedrake.supermeshing import * from itertools import product +from functools import partial import numpy import pytest @@ -14,7 +15,7 @@ def mesh(request): return UnitCubeMesh(3, 2, 1) -@pytest.fixture(params=["scalar", "vector", pytest.param("tensor", marks=pytest.mark.skip(reason="Prolongation fails for tensors"))]) +@pytest.fixture(params=["scalar", "vector", "tensor", "symmetric"]) def shapify(request): if request.param == "scalar": return lambda x: x @@ -22,6 +23,8 @@ def shapify(request): return VectorElement elif request.param == "tensor": return TensorElement + elif request.param == "symmetric": + return partial(TensorElement, symmetry=True) else: raise RuntimeError From 97419f291ce304583b3a0a34103dfa335ab7ed0c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 9 Dec 2025 13:46:04 +0000 Subject: [PATCH 4/8] fixup --- firedrake/supermeshing.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index a6f03e3596..aa12223a13 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -117,15 +117,19 @@ def likely(cell_A): assert V_A.block_size == V_B.block_size orig_block_size = V_A.block_size - symmetry = V_A.ufl_element().symmetry() - assert symmetry == V_B.ufl_element().symmetry() - if symmetry is not None: + if V_A.ufl_element().mapping() == "symmetries": + symmetry = V_A.ufl_element().symmetry() + assert V_B.ufl_element().mapping() == "symmetries" + assert V_B.ufl_element().symmetry() == symmetry + multiplicity = defaultdict(int) for idx in numpy.ndindex(V_A.value_shape): idx = symmetry.get(idx, idx) multiplicity[idx] += 1 symmetry = {i: multiplicity[idx] for i, idx in enumerate(multiplicity) if multiplicity[idx] != 1} + else: + symmetry = None if V_A.block_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements[0]) From 3fb7fb88c0c7055fdaede94bb16735be4f685ff8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 9 Dec 2025 14:41:16 +0000 Subject: [PATCH 5/8] docs --- firedrake/supermeshing.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index aa12223a13..a8757f806e 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -117,6 +117,7 @@ def likely(cell_A): assert V_A.block_size == V_B.block_size orig_block_size = V_A.block_size + # To deal with symmetry, each block of the mass matrix must be rescaled by the multiplicity if V_A.ufl_element().mapping() == "symmetries": symmetry = V_A.ufl_element().symmetry() assert V_B.ufl_element().mapping() == "symmetries" @@ -126,6 +127,8 @@ def likely(cell_A): for idx in numpy.ndindex(V_A.value_shape): idx = symmetry.get(idx, idx) multiplicity[idx] += 1 + + # dict mapping reference components to their multiplicity (if greater than 1) symmetry = {i: multiplicity[idx] for i, idx in enumerate(multiplicity) if multiplicity[idx] != 1} else: From 3d9e5509459ee3eeffdaa5754bb1754c4384a4a0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 15 Dec 2025 10:38:52 +0000 Subject: [PATCH 6/8] refactor for clarity --- firedrake/supermeshing.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index a8757f806e..6841f0543a 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -27,10 +27,10 @@ class BlockMatrix(object): - def __init__(self, mat, dimension, symmetry=None): + def __init__(self, mat, dimension, block_scale=None): self.mat = mat self.dimension = dimension - self.symmetry = symmetry or {} + self.block_scale = block_scale def mult(self, mat, x, y): sizes = self.mat.getSizes() @@ -43,8 +43,8 @@ def mult(self, mat, x, y): xi = PETSc.Vec().createWithArray(xa, size=sizes[1], comm=x.comm) yi = PETSc.Vec().createWithArray(ya, size=sizes[0], comm=y.comm) self.mat.mult(xi, yi) - if i in self.symmetry: - yi.array[:] *= self.symmetry[i] + if self.block_scale is not None: + yi.scale(self.block_scale[i]) y.array[start::stride] = yi.array_r def multTranspose(self, mat, x, y): @@ -58,8 +58,8 @@ def multTranspose(self, mat, x, y): xi = PETSc.Vec().createWithArray(xa, size=sizes[0], comm=x.comm) yi = PETSc.Vec().createWithArray(ya, size=sizes[1], comm=y.comm) self.mat.multTranspose(xi, yi) - if i in self.symmetry: - yi.array[:] *= self.symmetry[i] + if self.block_scale is not None: + yi.scale(self.block_scale[i]) y.array[start::stride] = yi.array_r @@ -128,11 +128,9 @@ def likely(cell_A): idx = symmetry.get(idx, idx) multiplicity[idx] += 1 - # dict mapping reference components to their multiplicity (if greater than 1) - symmetry = {i: multiplicity[idx] for i, idx in enumerate(multiplicity) - if multiplicity[idx] != 1} + block_scale = tuple(scale for idx, scale in multiplicity) else: - symmetry = None + block_scale = None if V_A.block_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements[0]) @@ -478,7 +476,7 @@ def likely(cell_A): lcols *= orig_block_size gcols *= orig_block_size size = ((lrows, grows), (lcols, gcols)) - context = BlockMatrix(mat, orig_block_size, symmetry=symmetry) + context = BlockMatrix(mat, orig_block_size, block_scale=block_scale) blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm) blockmat.setUp() return blockmat From 4412918d8f7778b54805d6bdf12f04b6e0beb19e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 15 Dec 2025 12:28:00 +0000 Subject: [PATCH 7/8] Apply suggestions from code review --- firedrake/supermeshing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index 6841f0543a..c1dc80dacb 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -26,6 +26,7 @@ __all__ = ["assemble_mixed_mass_matrix", "intersection_finder"] +# TODO replace with KAIJ (we require petsc4py wrappers) class BlockMatrix(object): def __init__(self, mat, dimension, block_scale=None): self.mat = mat @@ -128,7 +129,7 @@ def likely(cell_A): idx = symmetry.get(idx, idx) multiplicity[idx] += 1 - block_scale = tuple(scale for idx, scale in multiplicity) + block_scale = tuple(scale for idx, scale in multiplicity.items()) else: block_scale = None From f071c085b6fb3c70fef052d71cfd251774a22ec5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 16 Dec 2025 10:05:01 +0000 Subject: [PATCH 8/8] Update firedrake/interpolation.py --- firedrake/interpolation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index f94e2ad7eb..c4d0c8f6ae 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -473,6 +473,7 @@ def __init__( if bcs: raise NotImplementedError("bcs not implemented") + # TODO check V.finat_element.is_lagrange() once https://github.com/firedrakeproject/fiat/pull/200 is released target_element = V.ufl_element() if not ((isinstance(target_element, finat.ufl.MixedElement) and all(sub.mapping() == "identity" for sub in target_element.sub_elements))