BUG: Fix symmetry in supermeshing (#4764)

* BUG: Fix symmetry in supermeshing

* Fix cross mesh interpolation for symmetry=True

Author: pbrubeck

firedrake/interpolation.py

@@ -472,7 +472,12 @@ def __init__(
             raise NotImplementedError("freeze_expr not implemented")
         if bcs:
             raise NotImplementedError("bcs not implemented")
-        if V.ufl_element().mapping() != "identity":
+
+        # 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))
+                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 +556,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

firedrake/supermeshing.py

@@ -20,15 +20,18 @@
 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"]
 
 
+# TODO replace with KAIJ (we require petsc4py wrappers)
 class BlockMatrix(object):
-    def __init__(self, mat, dimension):
+    def __init__(self, mat, dimension, block_scale=None):
         self.mat = mat
         self.dimension = dimension
+        self.block_scale = block_scale
 
     def mult(self, mat, x, y):
         sizes = self.mat.getSizes()
@@ -41,6 +44,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 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):
@@ -54,6 +59,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 self.block_scale is not None:
+                yi.scale(self.block_scale[i])
             y.array[start::stride] = yi.array_r
 
 
@@ -68,14 +75,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 +115,39 @@ 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
+
+    # 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"
+        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
+
+        block_scale = tuple(scale for idx, scale in multiplicity.items())
+    else:
+        block_scale = None
+
+    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 +178,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 +468,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, block_scale=block_scale)
         blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm)
         blockmat.setUp()
         return blockmat

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))
 

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,14 +15,16 @@ 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
     elif request.param == "vector":
         return VectorElement
     elif request.param == "tensor":
         return TensorElement
+    elif request.param == "symmetric":
+        return partial(TensorElement, symmetry=True)
     else:
         raise RuntimeError