Assemble: use tensor kwarg in FormSum maxpy

firedrake/assemble.py

@@ -409,11 +409,7 @@ def visitor(e, *operands):
             for bc in self._bcs:
                 OneFormAssembler._apply_bc(self, result, bc, u=current_state)
 
-        if tensor:
-            BaseFormAssembler.update_tensor(result, tensor)
-            return tensor
-        else:
-            return result
+        return result
 
     def base_form_assembly_visitor(self, expr, tensor, *args):
         r"""Assemble a :class:`~ufl.classes.BaseForm` object given its assembled operands.
@@ -453,7 +449,6 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             petsc_mat.hermitianTranspose(out=res)
             (row, col) = mat.arguments()
             return matrix.AssembledMatrix((col, row), self._bcs, res,
-                                          appctx=self._appctx,
                                           options_prefix=self._options_prefix)
         elif isinstance(expr, ufl.Action):
             if len(args) != 2:
@@ -464,18 +459,16 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                     petsc_mat = lhs.petscmat
                     (row, col) = lhs.arguments()
                     # The matrix-vector product lives in the dual of the test space.
-                    res = firedrake.Function(row.function_space().dual())
-                    with rhs.dat.vec_ro as v_vec:
-                        with res.dat.vec as res_vec:
-                            petsc_mat.mult(v_vec, res_vec)
+                    res = tensor if tensor else firedrake.Function(row.function_space().dual())
+                    with rhs.dat.vec_ro as v_vec, res.dat.vec as res_vec:
+                        petsc_mat.mult(v_vec, res_vec)
                     return res
                 elif isinstance(rhs, matrix.MatrixBase):
-                    petsc_mat = lhs.petscmat
-                    (row, col) = lhs.arguments()
-                    res = petsc_mat.matMult(rhs.petscmat)
-                    return matrix.AssembledMatrix(expr, self._bcs, res,
-                                                  appctx=self._appctx,
-                                                  options_prefix=self._options_prefix)
+                    result = tensor.petscmat if tensor else PETSc.Mat()
+                    lhs.petscmat.matMult(rhs.petscmat, result=result)
+                    if tensor is None:
+                        tensor = self.assembled_matrix(expr, result)
+                    return tensor
                 else:
                     raise TypeError("Incompatible RHS for Action.")
             elif isinstance(lhs, (firedrake.Cofunction, firedrake.Function)):
@@ -493,30 +486,29 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 raise TypeError("Mismatching weights and operands in FormSum")
             if len(args) == 0:
                 raise TypeError("Empty FormSum")
+            if tensor:
+                tensor.zero()
             if all(isinstance(op, numbers.Complex) for op in args):
-                return sum(weight * arg for weight, arg in zip(expr.weights(), args))
+                result = sum(weight * arg for weight, arg in zip(expr.weights(), args))
+                return tensor.assign(result) if tensor else result
             elif all(isinstance(op, firedrake.Cofunction) for op in args):
                 V, = set(a.function_space() for a in args)
-                result = firedrake.Cofunction(V)
+                result = tensor if tensor else firedrake.Cofunction(V)
                 result.dat.maxpy(expr.weights(), [a.dat for a in args])
                 return result
             elif all(isinstance(op, ufl.Matrix) for op in args):
-                res = tensor.petscmat if tensor else PETSc.Mat()
-                is_set = False
+                result = tensor.petscmat if tensor else PETSc.Mat()
                 for (op, w) in zip(args, expr.weights()):
-                    # Make a copy to avoid in-place scaling
-                    petsc_mat = op.petscmat.copy()
-                    petsc_mat.scale(w)
-                    if is_set:
-                        # Modify output tensor in-place
-                        res += petsc_mat
+                    if result:
+                        # If result is not void, then accumulate on it
+                        result.axpy(w, op.petscmat)
                     else:
-                        # Copy to output tensor
-                        petsc_mat.copy(result=res)
-                        is_set = True
-                return matrix.AssembledMatrix(expr, self._bcs, res,
-                                              appctx=self._appctx,
-                                              options_prefix=self._options_prefix)
+                        # If result is void, then allocate it with first term
+                        op.petscmat.copy(result=result)
+                        result.scale(w)
+                if tensor is None:
+                    tensor = self.assembled_matrix(expr, result)
+                return tensor
             else:
                 raise TypeError("Mismatching FormSum shapes")
         elif isinstance(expr, ufl.ExternalOperator):
@@ -537,7 +529,8 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 # It is also convenient when we have a Form in that slot since Forms don't play well with `ufl.replace`
                 expr = expr._ufl_expr_reconstruct_(*expr.ufl_operands, argument_slots=(v,) + expr.argument_slots()[1:])
             # Call the external operator assembly
-            return expr.assemble(assembly_opts=opts)
+            result = expr.assemble(assembly_opts=opts)
+            return tensor.assign(result) if tensor else result
         elif isinstance(expr, ufl.Interpolate):
             # Replace assembled children
             _, expression = expr.argument_slots()
@@ -589,33 +582,24 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 else:
                     # Copy the interpolation matrix into the output tensor
                     petsc_mat.copy(result=res)
-                return matrix.AssembledMatrix(expr.arguments(), self._bcs, res,
-                                              appctx=self._appctx,
-                                              options_prefix=self._options_prefix)
+                if tensor is None:
+                    tensor = self.assembled_matrix(expr, res)
+                return tensor
             else:
                 # The case rank == 0 is handled via the DAG restructuring
                 raise ValueError("Incompatible number of arguments.")
-        elif isinstance(expr, (ufl.Cofunction, ufl.Coargument, ufl.Argument, ufl.Matrix, ufl.ZeroBaseForm)):
-            return expr
-        elif isinstance(expr, ufl.Coefficient):
+        elif tensor and isinstance(expr, (firedrake.Function, firedrake.Cofunction, firedrake.MatrixBase)):
+            return tensor.assign(expr)
+        elif tensor and isinstance(expr, ufl.ZeroBaseForm):
+            return tensor.zero()
+        elif isinstance(expr, (ufl.Coefficient, ufl.Cofunction, ufl.Matrix, ufl.Argument, ufl.Coargument, ufl.ZeroBaseForm)):
             return expr
         else:
             raise TypeError(f"Unrecognised BaseForm instance: {expr}")
 
-    @staticmethod
-    def update_tensor(assembled_base_form, tensor):
-        if isinstance(tensor, (firedrake.Function, firedrake.Cofunction)):
-            if isinstance(assembled_base_form, ufl.ZeroBaseForm):
-                tensor.dat.zero()
-            else:
-                assembled_base_form.dat.copy(tensor.dat)
-        elif isinstance(tensor, matrix.MatrixBase):
-            if isinstance(assembled_base_form, ufl.ZeroBaseForm):
-                tensor.petscmat.zeroEntries()
-            else:
-                assembled_base_form.petscmat.copy(tensor.petscmat)
-        else:
-            raise NotImplementedError("Cannot update tensor of type %s" % type(tensor))
+    def assembled_matrix(self, expr, petscmat):
+        return matrix.AssembledMatrix(expr.arguments(), self._bcs, petscmat,
+                                      options_prefix=self._options_prefix)
 
     @staticmethod
     def base_form_postorder_traversal(expr, visitor, visited={}):

firedrake/constant.py

@@ -187,6 +187,10 @@ def assign(self, value):
         except (DataTypeError, DataValueError) as e:
             raise ValueError(e)
 
+    def zero(self):
+        """Set the value of this constant to zero."""
+        return self.assign(0)
+
     def __iadd__(self, o):
         raise NotImplementedError("Augmented assignment to Constant not implemented")
 

firedrake/matrix.py

@@ -42,6 +42,8 @@ def __init__(self, a, bcs, mat_type):
         self._analyze_form_arguments()
         self._arguments = arguments
 
+        if bcs is None:
+            bcs = ()
         self.bcs = bcs
         self.comm = test.function_space().comm
         self._comm = internal_comm(self.comm, self)
@@ -97,6 +99,33 @@ def __str__(self):
         return "assembled %s(a=%s, bcs=%s)" % (type(self).__name__,
                                                self.a, self.bcs)
 
+    def __add__(self, other):
+        if isinstance(other, MatrixBase):
+            mat = self.petscmat + other.petscmat
+            return AssembledMatrix(self.arguments(), (), mat)
+        else:
+            return NotImplemented
+
+    def __sub__(self, other):
+        if isinstance(other, MatrixBase):
+            mat = self.petscmat - other.petscmat
+            return AssembledMatrix(self.arguments(), (), mat)
+        else:
+            return NotImplemented
+
+    def assign(self, val):
+        """Set matrix entries."""
+        if isinstance(val, MatrixBase):
+            val.petscmat.copy(self.petscmat)
+        else:
+            raise TypeError(f"Cannot assign a {type(val).__name__} to a {type(self).__name__}.")
+        return self
+
+    def zero(self):
+        """Set all matrix entries to zero."""
+        self.petscmat.zeroEntries()
+        return self
+
 
 class Matrix(MatrixBase):
     """A representation of an assembled bilinear form.
@@ -123,10 +152,11 @@ class Matrix(MatrixBase):
     def __init__(self, a, bcs, mat_type, *args, **kwargs):
         # sets self._a, self._bcs, and self._mat_type
         MatrixBase.__init__(self, a, bcs, mat_type)
-        options_prefix = kwargs.pop("options_prefix")
+        options_prefix = kwargs.pop("options_prefix", None)
         self.M = op2.Mat(*args, mat_type=mat_type, **kwargs)
         self.petscmat = self.M.handle
-        self.petscmat.setOptionsPrefix(options_prefix)
+        if options_prefix is not None:
+            self.petscmat.setOptionsPrefix(options_prefix)
         self.mat_type = mat_type
 
     def assemble(self):
@@ -196,19 +226,15 @@ class AssembledMatrix(MatrixBase):
     :arg petscmat: the already constructed petsc matrix this object represents.
     """
     def __init__(self, a, bcs, petscmat, *args, **kwargs):
+        options_prefix = kwargs.pop("options_prefix", None)
         super(AssembledMatrix, self).__init__(a, bcs, "assembled")
 
         self.petscmat = petscmat
+        if options_prefix is not None:
+            self.petscmat.setOptionsPrefix(options_prefix)
 
         # this allows call to self.M.handle without a new class
         self.M = SimpleNamespace(handle=self.mat())
 
     def mat(self):
         return self.petscmat
-
-    def __add__(self, other):
-        if isinstance(other, AssembledMatrix):
-            return self.petscmat + other.petscmat
-        else:
-            raise TypeError("Unable to add %s to AssembledMatrix"
-                            % (type(other)))

tests/firedrake/regression/test_assemble_baseform.py

@@ -90,7 +90,9 @@ def test_assemble_adjoint(M):
 def test_assemble_action(M, f):
     res = assemble(action(M, f))
     assembledM = assemble(M)
-    res2 = assemble(action(assembledM, f))
+    expr = action(assembledM, f)
+
+    res2 = assemble(expr)
     assert isinstance(res2, Cofunction)
     assert isinstance(res, Cofunction)
     for f, f2 in zip(res.subfunctions, res2.subfunctions):
@@ -100,6 +102,36 @@ def test_assemble_action(M, f):
         else:
             assert abs(f.dat.data.sum() - 0.5*f.function_space().value_size) < 1.0e-12
 
+    out = assemble(expr, tensor=res2)
+    assert out is res2
+    for f, f2 in zip(res.subfunctions, res2.subfunctions):
+        assert abs(f.dat.data.sum() - f2.dat.data.sum()) < 1.0e-12
+        if f.function_space().rank == 2:
+            assert abs(f.dat.data.sum() - 0.5*sum(f.function_space().shape)) < 1.0e-12
+        else:
+            assert abs(f.dat.data.sum() - 0.5*f.function_space().value_size) < 1.0e-12
+
+
+def test_scalar_formsum(f):
+    c = Cofunction(f.function_space().dual())
+    c.assign(1)
+
+    q = action(c, f)
+    expected = 2 * assemble(q)
+
+    formsum = 1.5 * q + 0.5 * q
+    assert isinstance(formsum, ufl.form.FormSum)
+    res2 = assemble(formsum)
+    assert res2 == expected
+
+    mesh = f.function_space().mesh()
+    R = FunctionSpace(mesh, "R", 0)
+    tensor = Cofunction(R.dual())
+
+    out = assemble(formsum, tensor=tensor)
+    assert out is tensor
+    assert tensor.dat.data[0] == expected
+
 
 def test_vector_formsum(a):
     res = assemble(a)
@@ -111,7 +143,12 @@ def test_vector_formsum(a):
     assert isinstance(res2, Cofunction)
     assert isinstance(preassemble, Cofunction)
     for f, f2 in zip(preassemble.subfunctions, res2.subfunctions):
-        assert abs(f.dat.data.sum() - f2.dat.data.sum()) < 1.0e-12
+        assert np.allclose(f.dat.data, f2.dat.data, atol=1e-12)
+
+    out = assemble(formsum, tensor=res2)
+    assert out is res2
+    for f, f2 in zip(preassemble.subfunctions, res2.subfunctions):
+        assert np.allclose(f.dat.data, f2.dat.data, atol=1e-12)
 
 
 def test_matrix_formsum(M):
@@ -121,8 +158,11 @@ def test_matrix_formsum(M):
     assert isinstance(formsum, ufl.form.FormSum)
     res2 = assemble(formsum)
     assert isinstance(res2, ufl.Matrix)
-    assert np.allclose(sumfirst.petscmat[:, :],
-                       res2.petscmat[:, :], rtol=1e-14)
+    assert np.allclose(sumfirst.petscmat[:, :], res2.petscmat[:, :], rtol=1E-14)
+
+    out = assemble(formsum, tensor=res2)
+    assert out is res2
+    assert np.allclose(sumfirst.petscmat[:, :], res2.petscmat[:, :], rtol=1E-14)
 
 
 def test_zero_form(M, f, one):
@@ -131,7 +171,7 @@ def test_zero_form(M, f, one):
     assert abs(zero_form - 0.5 * np.prod(f.ufl_shape)) < 1.0e-12
 
 
-def test_tensor_copy(a, M):
+def test_tensor_output(a, M):
 
     # 1-form tensor
     V = a.arguments()[0].function_space()
@@ -140,19 +180,15 @@ def test_tensor_copy(a, M):
     res = assemble(formsum, tensor=tensor)
 
     assert isinstance(formsum, ufl.form.FormSum)
-    assert isinstance(res, Cofunction)
-    for f, f2 in zip(res.subfunctions, tensor.subfunctions):
-        assert abs(f.dat.data.sum() - f2.dat.data.sum()) < 1.0e-12
+    assert res is tensor
 
     # 2-form tensor
     tensor = get_assembler(M).allocate()
     formsum = assemble(M) + M
     res = assemble(formsum, tensor=tensor)
 
     assert isinstance(formsum, ufl.form.FormSum)
-    assert isinstance(res, ufl.Matrix)
-    assert np.allclose(res.petscmat[:, :],
-                       tensor.petscmat[:, :], rtol=1e-14)
+    assert res is tensor
 
 
 def test_cofunction_assign(a, M, f):

tests/firedrake/regression/test_interp_dual.py

@@ -181,7 +181,7 @@ def test_assemble_base_form_operator_expressions(mesh):
     res = assemble(Iv1 + Iv2)
     mat_Iv1 = assemble(Iv1)
     mat_Iv2 = assemble(Iv2)
-    assert np.allclose((mat_Iv1 + mat_Iv2)[:, :], res.petscmat[:, :], rtol=1e-14)
+    assert np.allclose(mat_Iv1.petscmat[:, :] + mat_Iv2.petscmat[:, :], res.petscmat[:, :], rtol=1e-14)
 
     # Linear combination of BaseFormOperator (1-form)
     alpha = 0.5