Interpolate physically-mapped elements into VertexOnlyMesh (#4213)

---------

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

Author: pbrubeck

firedrake/assemble.py

@@ -1114,7 +1114,7 @@ class ZeroFormAssembler(ParloopFormAssembler):
 
     Parameters
     ----------
-    form : ufl.Form or slate.TensorBasehe
+    form : ufl.Form or slate.TensorBase
         0-form.
 
     Notes

firedrake/interpolation.py

@@ -1194,7 +1194,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
         co = target_mesh.cell_orientations()
         parloop_args.append(co.dat(op2.READ, co.cell_node_map()))
     if needs_cell_sizes:
-        cs = target_mesh.cell_sizes
+        cs = source_mesh.cell_sizes
         parloop_args.append(cs.dat(op2.READ, cs.cell_node_map()))
     for coefficient in coefficients:
         if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):

tests/firedrake/regression/test_interpolate_zany.py

@@ -51,7 +51,7 @@ def expect(V, which):
         return a + b
 
 
-def test_interpolate(V, mesh, which, expect, tolerance):
+def test_interpolate_zany_into_cg(V, mesh, which, expect, tolerance):
     degree = V.ufl_element().degree()
     Vcg = FunctionSpace(mesh, "P", degree)
 
@@ -71,3 +71,64 @@ def test_interpolate(V, mesh, which, expect, tolerance):
         g.interpolate(a + b)
 
     assert numpy.allclose(norm(g - expect), 0, atol=tolerance)
+
+
+@pytest.fixture
+def vom(mesh):
+    return VertexOnlyMesh(mesh, [(0.5, 0.5), (0.31, 0.72)])
+
+
+@pytest.fixture
+def expr_at_vom(V, which, vom):
+    mesh = V.mesh()
+    degree = V.ufl_element().degree()
+    x, y = SpatialCoordinate(mesh)
+    expr = (x + y)**degree
+
+    if which == "coefficient":
+        P0 = FunctionSpace(vom, "DG", 0)
+    elif which == "grad":
+        expr = ufl.algorithms.expand_derivatives(grad(expr))
+        P0 = VectorFunctionSpace(vom, "DG", 0)
+
+    fvom = Function(P0)
+    point = Constant([0] * mesh.geometric_dimension())
+    expr_at_pt = ufl.replace(expr, {SpatialCoordinate(mesh): point})
+    for i, pt in enumerate(vom.coordinates.dat.data_ro):
+        point.assign(pt)
+        fvom.dat.data[i] = numpy.asarray(expr_at_pt, dtype=float)
+    return fvom
+
+
+def test_interpolate_zany_into_vom(V, mesh, which, expr_at_vom):
+    degree = V.ufl_element().degree()
+    x, y = SpatialCoordinate(mesh)
+    expr = (x + y)**degree
+
+    f = Function(V)
+    f.project(expr, solver_parameters={"ksp_type": "preonly",
+                                       "pc_type": "lu"})
+    fexpr = f
+    vexpr = TestFunction(V)
+    if which == "grad":
+        fexpr = grad(fexpr)
+        vexpr = grad(vexpr)
+
+    P0 = expr_at_vom.function_space()
+
+    # Interpolate a Function into P0(vom)
+    f_at_vom = assemble(Interpolate(fexpr, P0))
+    assert numpy.allclose(f_at_vom.dat.data_ro, expr_at_vom.dat.data_ro)
+
+    # Construct a Cofunction on P0(vom)*
+    Fvom = Cofunction(P0.dual()).assign(1)
+    expected_action = assemble(action(Fvom, expr_at_vom))
+
+    # Interpolate a Function into Fvom
+    f_at_vom = assemble(Interpolate(fexpr, Fvom))
+    assert numpy.allclose(f_at_vom, expected_action)
+
+    # Interpolate a TestFunction into Fvom
+    expr_vom = assemble(Interpolate(vexpr, Fvom))
+    f_at_vom = assemble(action(expr_vom, f))
+    assert numpy.allclose(f_at_vom, expected_action)

tsfc/driver.py

@@ -208,11 +208,13 @@ def compile_expression_dual_evaluation(expression, to_element, ufl_element, *,
     # Collect required coefficients and determine numbering
     coefficients = extract_coefficients(expression)
     orig_coefficients = extract_coefficients(orig_expression)
-    coefficient_numbers = tuple(orig_coefficients.index(c) for c in coefficients)
+    coefficient_numbers = tuple(map(orig_coefficients.index, coefficients))
     builder.set_coefficient_numbers(coefficient_numbers)
 
+    elements = [f.ufl_element() for f in (*coefficients, *arguments)]
+
     needs_external_coords = False
-    if has_type(expression, GeometricQuantity) or any(fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients):
+    if has_type(expression, GeometricQuantity) or any(map(fem.needs_coordinate_mapping, elements)):
         # Create a fake coordinate coefficient for a domain.
         coords_coefficient = ufl.Coefficient(ufl.FunctionSpace(domain, domain.ufl_coordinate_element()))
         builder.domain_coordinate[domain] = coords_coefficient

tsfc/fem.py

@@ -336,7 +336,8 @@ class GemPointContext(ContextBase):
     def basis_evaluation(self, finat_element, mt, entity_id):
         return finat_element.point_evaluation(mt.local_derivatives,
                                               self.point_expr,
-                                              (self.integration_dim, entity_id))
+                                              (self.integration_dim, entity_id),
+                                              CoordinateMapping(mt, self))
 
 
 class Translator(MultiFunction, ModifiedTerminalMixin, ufl2gem.Mixin):

tsfc/kernel_interface/firedrake_loopy.py

@@ -104,6 +104,16 @@ def _coefficient(self, coefficient, name):
         self.coefficient_map[coefficient] = expr
         return expr
 
+    def set_coordinates(self, domain):
+        """Prepare the coordinate field.
+
+        :arg domain: :class:`ufl.Domain`
+        """
+        # Create a fake coordinate coefficient for a domain.
+        f = Coefficient(FunctionSpace(domain, domain.ufl_coordinate_element()))
+        self.domain_coordinate[domain] = f
+        self._coefficient(f, "coords")
+
     def set_cell_sizes(self, domain):
         """Setup a fake coefficient for "cell sizes".
 
@@ -304,16 +314,6 @@ def set_arguments(self, arguments):
         self.return_variables = return_variables
         self.argument_multiindices = argument_multiindices
 
-    def set_coordinates(self, domain):
-        """Prepare the coordinate field.
-
-        :arg domain: :class:`ufl.Domain`
-        """
-        # Create a fake coordinate coefficient for a domain.
-        f = Coefficient(FunctionSpace(domain, domain.ufl_coordinate_element()))
-        self.domain_coordinate[domain] = f
-        self._coefficient(f, "coords")
-
     def set_coefficients(self, integral_data, form_data):
         """Prepare the coefficients of the form.