address review comments

Author: pbrubeck

firedrake/assemble.py

@@ -466,13 +466,11 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                             petsc_mat.mult(v_vec, res_vec)
                     return res
                 elif isinstance(rhs, matrix.MatrixBase):
-                    petsc_mat = lhs.petscmat
-                    (row, col) = lhs.arguments()
-                    res = tensor.petscmat if tensor else PETSc.Mat()
-                    petsc_mat.matMult(rhs.petscmat, result=res)
-                    return tensor if tensor else matrix.AssembledMatrix(expr, self._bcs, res,
-                                                                        appctx=self._appctx,
-                                                                        options_prefix=self._options_prefix)
+                    result = tensor.petscmat if tensor else PETSc.Mat()
+                    lhs.petsc_mat.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)):
@@ -501,17 +499,18 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 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()
+                result = tensor.petscmat if tensor else PETSc.Mat()
                 for (op, w) in zip(args, expr.weights()):
-                    if res:
-                        res.axpy(w, op.petscmat)
+                    if result:
+                        # If result is not void, then accumulate on it
+                        result.axpy(w, op.petscmat)
                     else:
-                        # Make a copy to avoid in-place scaling
-                        res = op.petscmat.copy()
-                        res.scale(w)
-                return tensor if tensor else 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):
@@ -585,9 +584,9 @@ 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.")
@@ -600,6 +599,10 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
         else:
             raise TypeError(f"Unrecognised BaseForm instance: {expr}")
 
+    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={}):
         if expr in visited:

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,12 @@ 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):
+            return self.petscmat + other.petscmat
+        else:
+            return NotImplemented
+
     def assign(self, val):
         """Set matrix entries."""
         if isinstance(val, MatrixBase):
@@ -212,15 +220,11 @@ def __init__(self, a, bcs, petscmat, *args, **kwargs):
         super(AssembledMatrix, self).__init__(a, bcs, "assembled")
 
         self.petscmat = petscmat
+        options_prefix = kwargs.pop("options_prefix")
+        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, MatrixBase):
-            return self.petscmat + other.petscmat
-        else:
-            return NotImplemented