Reusable Interpolator

Author: pbrubeck

firedrake/assemble.py

@@ -538,6 +538,8 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             result = expr.assemble(assembly_opts=opts)
             return tensor.assign(result) if tensor else result
         elif isinstance(expr, ufl.Interpolate):
+            if not isinstance(expr, firedrake.Interpolate):
+                expr = firedrake.Interpolate(*reversed(expr.dual_args()))
             orig_expr = expr
             # Replace assembled children
             _, operand = expr.argument_slots()

firedrake/interpolation.py

@@ -1006,8 +1006,13 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     parameters = {}
     parameters['scalar_type'] = utils.ScalarType
 
+    callables = ()
+    if access == op2.INC:
+        callables += (tensor.zero,)
+
     needs_weight = isinstance(dual_arg, ufl.Cofunction) and not to_element.is_dg()
     if needs_weight:
+        # Compute the reciprocal of the DOF multiplicity
         W = dual_arg.function_space()
         shapes = (W.finat_element.space_dimension(), W.block_size)
         domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
@@ -1018,14 +1023,14 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
         """
         weight = firedrake.Function(W)
         firedrake.par_loop((domain, instructions), ufl.dx, {"w": (weight, op2.INC)})
+        with weight.dat.vec as w:
+            w.reciprocal()
 
-        # Create a copy and apply the weight
-        # TODO include this in the callables
-        v = firedrake.Function(dual_arg)
-        with v.dat.vec as x, weight.dat.vec as w:
-            x.pointwiseDivide(x, w)
-
+        # Create a buffer for the weighted Cofunction and a callable to apply the weight
+        v = firedrake.Function(W)
         expr = expr._ufl_expr_reconstruct_(operand, v=v)
+        with weight.dat.vec_ro as w, dual_arg.dat.vec_ro as x, v.dat.vec_wo as y:
+            callables += (partial(y.pointwiseMult, x, w),)
 
     # We need to pass both the ufl element and the finat element
     # because the finat elements might not have the right mapping
@@ -1043,6 +1048,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     name = kernel.name
     kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True,
                         flop_count=kernel.flop_count, events=(kernel.event,))
+
     parloop_args = [kernel, cell_set]
 
     coefficients = tsfc_interface.extract_numbered_coefficients(expr, coefficient_numbers)
@@ -1158,7 +1164,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
     if isinstance(tensor, op2.Mat):
         return parloop_compute_callable, tensor.assemble
     else:
-        return copyin + (parloop_compute_callable, ) + copyout
+        return copyin + callables + (parloop_compute_callable, ) + copyout
 
 
 try:

firedrake/variational_solver.py

@@ -7,7 +7,7 @@
 from firedrake import dmhooks, slate, solving, solving_utils, ufl_expr, utils
 from firedrake.petsc import PETSc, DEFAULT_KSP_PARAMETERS, DEFAULT_SNES_PARAMETERS
 from firedrake.function import Function
-from firedrake.interpolation import Interpolate
+from firedrake.interpolation import interpolate
 from firedrake.matrix import MatrixBase
 from firedrake.ufl_expr import TrialFunction, TestFunction
 from firedrake.bcs import DirichletBC, EquationBC, extract_subdomain_ids, restricted_function_space
@@ -98,7 +98,7 @@ def __init__(self, F, u, bcs=None, J=None,
                 F_arg, = F.arguments()
                 self.F = replace(F, {F_arg: v_res, self.u: self.u_restrict})
             else:
-                self.F = Interpolate(v_res, replace(F, {self.u: self.u_restrict}))
+                self.F = interpolate(v_res, replace(F, {self.u: self.u_restrict}))
 
             v_arg, u_arg = self.J.arguments()
             self.J = replace(self.J, {v_arg: v_res, u_arg: u_res, self.u: self.u_restrict})

tests/firedrake/regression/test_interpolate.py

@@ -329,7 +329,7 @@ def test_trace():
 
 @pytest.mark.parametrize("rank", (0, 1))
 @pytest.mark.parametrize("mat_type", ("matfree", "aij"))
-@pytest.mark.parametrize("degree", range(1, 4))
+@pytest.mark.parametrize("degree", (1, 3))
 @pytest.mark.parametrize("cell", ["triangle", "quadrilateral"])
 @pytest.mark.parametrize("shape", ("scalar", "vector", "tensor"))
 def test_adjoint_Pk(rank, mat_type, degree, cell, shape):
@@ -350,14 +350,15 @@ def test_adjoint_Pk(rank, mat_type, degree, cell, shape):
         operand = TestFunction(Pk)
 
     if mat_type == "matfree":
-        result = assemble(interpolate(operand, v))
+        interp = interpolate(operand, v)
     else:
         adj_interp = assemble(interpolate(operand, TrialFunction(Pkp1.dual())))
         if rank == 0:
-            result = assemble(action(v, adj_interp))
+            interp = action(v, adj_interp)
         else:
-            result = assemble(action(adj_interp, v))
+            interp = action(adj_interp, v)
 
+    result = assemble(interp)
     expect = assemble(inner(expr, operand) * dx)
     if rank == 0:
         assert np.allclose(result, expect)