diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 7423eb548e..c4d0c8f6ae 100644 --- a/firedrake/interpolation.py +++ b/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 diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index ec037aa283..c1dc80dacb 100644 --- a/firedrake/supermeshing.py +++ b/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 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