Fix for vom-to-vom permutation matrix on Vector-valued function spaces (#4510)

* fix readonly=True

* remove _get_sizes; save nleaves as attribute

* add tensor-valued test

---------

Co-authored-by: Pablo Brubeck <[email protected]>

Author: leo-collins

firedrake/interpolation.py

@@ -1521,10 +1521,13 @@ def __init__(self, sf, forward_reduce, V, source_vom, expr, arguments):
         self.arguments = arguments
         # Calculate correct local and global sizes for the matrix
         nroots, leaves, _ = sf.getGraph()
-        nleaves = len(leaves)
+        self.nleaves = len(leaves)
         self._local_sizes = V.comm.allgather(nroots)
-        self.source_size = (nroots, sum(self._local_sizes))
-        self.target_size = (nleaves, self.V.comm.allreduce(nleaves, op=MPI.SUM))
+        self.source_size = (self.V.block_size * nroots, self.V.block_size * sum(self._local_sizes))
+        self.target_size = (
+            self.V.block_size * self.nleaves,
+            self.V.block_size * V.comm.allreduce(self.nleaves, op=MPI.SUM),
+        )
 
     @property
     def mpi_type(self):
@@ -1565,7 +1568,7 @@ def expr_as_coeff(self, source_vec=None):
                     raise ValueError("Need to provide a source dat for the argument!")
                 arg = self.arguments[0]
                 arg_coeff = firedrake.Function(arg.function_space())
-                arg_coeff.dat.data_wo[:] = source_vec.getArray().reshape(
+                arg_coeff.dat.data_wo[:] = source_vec.getArray(readonly=True).reshape(
                     arg_coeff.dat.data_wo.shape
                 )
                 coeff_expr = ufl.replace(self.expr, {arg: arg_coeff})
@@ -1643,14 +1646,6 @@ def multTranspose(self, mat, source_vec, target_vec):
             target_vec.zeroEntries()
             self.reduce(source_vec, target_vec)
 
-    def _get_sizes(self):
-        nroots, leaves, _ = self.sf.getGraph()
-        nleaves = len(leaves)
-        local_sizes = self.V.comm.allgather(nroots)
-        source_size = (nroots, sum(local_sizes))
-        target_size = (nleaves, self.V.comm.allreduce(nleaves, op=MPI.SUM))
-        return source_size, target_size
-
     def _create_permutation_mat(self):
         """Creates the PETSc matrix that represents the interpolation operator from a vertex-only mesh to
         its input ordering vertex-only mesh"""
@@ -1659,25 +1654,23 @@ def _create_permutation_mat(self):
         start = sum(self._local_sizes[:self.V.comm.rank])
         end = start + self.source_size[0]
         contiguous_indices = numpy.arange(start, end, dtype=utils.IntType)
-        perm = numpy.zeros(self.target_size[0], dtype=utils.IntType)
+        perm = numpy.zeros(self.nleaves, dtype=utils.IntType)
         self.sf.bcastBegin(MPI.INT, contiguous_indices, perm, MPI.REPLACE)
         self.sf.bcastEnd(MPI.INT, contiguous_indices, perm, MPI.REPLACE)
         rows = numpy.arange(self.target_size[0] + 1, dtype=utils.IntType)
-        mat.setValuesCSR(rows, perm, numpy.ones_like(perm, dtype=utils.IntType))
+        cols = (self.V.block_size * perm[:, None] + numpy.arange(self.V.block_size, dtype=utils.IntType)[None, :]).reshape(-1)
+        mat.setValuesCSR(rows, cols, numpy.ones_like(cols, dtype=utils.IntType))
         mat.assemble()
         if self.forward_reduce:
             mat.transpose()
         return mat
 
     def _wrap_dummy_mat(self):
         mat = PETSc.Mat().create(comm=self.V.comm)
-        dim = self.V.value_size
-        source_size = tuple(dim * i for i in self.source_size)
-        target_size = tuple(dim * i for i in self.target_size)
         if self.forward_reduce:
-            mat_size = (source_size, target_size)
+            mat_size = (self.source_size, self.target_size)
         else:
-            mat_size = (target_size, source_size)
+            mat_size = (self.target_size, self.source_size)
         mat.setSizes(mat_size)
         mat.setType(mat.Type.PYTHON)
         mat.setPythonContext(self)

tests/firedrake/vertexonly/test_vertex_only_fs.py

@@ -120,7 +120,10 @@ def functionspace_tests(vm, petsc_raises):
     idxs_to_include = input_ordering_parent_cell_nums != -1
     assert np.allclose(h.dat.data_ro_with_halos[idxs_to_include], np.prod(vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include].reshape(-1, vm.input_ordering.geometric_dimension()), axis=1))
     assert np.all(h.dat.data_ro_with_halos[~idxs_to_include] == -1)
-
+    # Using permutation matrix
+    perm_mat = assemble(interpolate(TestFunction(V), W, matfree=False))
+    h2 = assemble(perm_mat @ g)
+    assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], h.dat.data_ro_with_halos[idxs_to_include])
     # check other interpolation APIs work identically
     h2 = assemble(interpolate(g, W))
     assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], h.dat.data_ro_with_halos[idxs_to_include])
@@ -225,6 +228,10 @@ def vectorfunctionspace_tests(vm, petsc_raises):
     idxs_to_include = input_ordering_parent_cell_nums != -1
     assert np.allclose(h.dat.data_ro[idxs_to_include], 2*vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include])
     assert np.all(h.dat.data_ro_with_halos[~idxs_to_include] == -1)
+    # Using permutation matrix
+    perm_mat = assemble(interpolate(TestFunction(V), W, matfree=False))
+    h2 = assemble(perm_mat @ g)
+    assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], h.dat.data_ro_with_halos[idxs_to_include])
     # check other interpolation APIs work identically
     h2 = assemble(interpolate(g, W))
     assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], h.dat.data_ro_with_halos[idxs_to_include])
@@ -353,3 +360,48 @@ def test_input_ordering_missing_point():
         assert not len(data_input_ordering.dat.data_ro)
         # Accessing data_ro [*here] is collective, hence this redundant call
         _ = len(data_input_ordering.dat.data_ro)
+
+
+@pytest.fixture(
+    params=[
+        ((2, 2), None),
+        (None, True),
+        ((), None),
+        ((2, 3), None),
+    ]
+)
+def tensorfs_and_expr(request):
+    shape, symmetry = request.param
+    np.random.seed(0)
+    mesh = UnitSquareMesh(2, 2)
+    coords = np.random.random_sample(size=(10, 2))
+    vom = VertexOnlyMesh(mesh, coords)
+
+    V = TensorFunctionSpace(vom, "DG", 0, shape=shape, symmetry=symmetry)
+    W = TensorFunctionSpace(vom.input_ordering, "DG", 0, shape=shape, symmetry=symmetry)
+
+    x = SpatialCoordinate(vom)
+    if shape == ():
+        expr = inner(x, x)
+    elif shape is None or shape == (2, 2):
+        expr = outer(x, x) + Identity(2)
+    elif shape == (2, 3):
+        a = as_vector([x[0], x[1]])
+        b = as_vector([x[0], x[1], Constant(1.0)])
+        expr = outer(a, b)
+
+    return V, W, expr
+
+
+@pytest.mark.parallel([1, 3])
+def test_tensorfs_permutation(tensorfs_and_expr):
+    V, W, expr = tensorfs_and_expr
+    f = Function(V)
+    f.interpolate(expr)
+    f_in_W = assemble(interpolate(f, W))
+    python_mat = assemble(interpolate(TestFunction(V), W, matfree=False))
+    f_in_W_2 = assemble(python_mat @ f)
+    assert np.allclose(f_in_W.dat.data_ro, f_in_W_2.dat.data_ro)
+    petsc_mat = assemble(interpolate(TestFunction(V), W, matfree=True))
+    f_in_W_petsc = assemble(petsc_mat @ f)
+    assert np.allclose(f_in_W.dat.data_ro, f_in_W_petsc.dat.data_ro)