Assemble: use tensor kwarg in FormSum maxpy

demos/netgen/netgen_mesh.py.rst

@@ -156,7 +156,8 @@ Then a SLEPc Eigenvalue Problem Solver (`EPS`) is initialised and set up to use
         bc = DirichletBC(V, 0, labels)
         A = assemble(a, bcs=bc)
         M = assemble(m, bcs=bc, weight=0.)
-        Asc, Msc = A.M.handle, M.M.handle
+        Asc = A.petscmat
+        Msc = M.petscmat
         E = SLEPc.EPS().create()
         E.setType(SLEPc.EPS.Type.ARNOLDI)
         E.setProblemType(SLEPc.EPS.ProblemType.GHEP)

docs/source/petsc-interface.rst

@@ -51,7 +51,7 @@ different.  For a bilinear form, the matrix is obtained with:
 
 .. code-block:: python3
 
-   petsc_mat = assemble(bilinear_form).M.handle
+   petsc_mat = assemble(bilinear_form).petscmat
 
 For a linear form, we need to use a context manager.  There are two
 options available here, depending on whether we want read-only or
@@ -136,7 +136,7 @@ newly defined class to compute the matrix action:
 
    # Assemble the bilinear form that defines A and get the concrete
    # PETSc matrix
-   A = assemble(bilinear_form).M.handle
+   A = assemble(bilinear_form).petscmat
 
    # Now do the same for the linear forms for u and v, making a copy
 

docs/source/solving-interface.rst

@@ -932,7 +932,7 @@ with the following:
 .. code-block:: python3
 
    A = assemble(a)  # use bcs keyword if there are boundary conditions
-   print A.M.handle.isSymmetric(tol=1e-13)
+   print A.petscmat.isSymmetric(tol=1e-13)
 
 If the problem is not symmetric, try using a method such as GMRES
 instead.  PETSc uses restarted GMRES with a default restart of 30, for

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.
@@ -464,15 +460,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())
+                    res = tensor if tensor else 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)
                     return res
                 elif isinstance(rhs, matrix.MatrixBase):
                     petsc_mat = lhs.petscmat
                     (row, col) = lhs.arguments()
-                    res = petsc_mat.matMult(rhs.petscmat)
+                    res = tensor.petscmat if tensor else PETSc.Mat()
+                    petsc_mat.matMult(rhs.petscmat, result=res)
                     return matrix.AssembledMatrix(expr, self._bcs, res,
                                                   appctx=self._appctx,
                                                   options_prefix=self._options_prefix)
@@ -493,27 +490,25 @@ 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
                 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 res:
+                        res.axpy(w, op.petscmat)
                     else:
-                        # Copy to output tensor
-                        petsc_mat.copy(result=res)
-                        is_set = True
+                        # Make a copy to avoid in-place scaling
+                        res = op.petscmat.copy()
+                        res.scale(w)
                 return matrix.AssembledMatrix(expr, self._bcs, res,
                                               appctx=self._appctx,
                                               options_prefix=self._options_prefix)
@@ -537,7 +532,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()
@@ -595,28 +591,15 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             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))
-
     @staticmethod
     def base_form_postorder_traversal(expr, visitor, visited={}):
         if expr in visited:

firedrake/eigensolver.py

@@ -186,11 +186,11 @@ def solve(self):
         int
             The number of Eigenvalues found.
         """
-        self.A_mat = assemble(self._problem.A, bcs=self._problem.bcs).M.handle
+        self.A_mat = assemble(self._problem.A, bcs=self._problem.bcs).petscmat
         self.M_mat = assemble(
             self._problem.M, bcs=self._problem.bcs,
             weight=self._problem.bc_shift and 1./self._problem.bc_shift
-        ).M.handle
+        ).petscmat
 
         self.es.setDimensions(nev=self.n_evals, ncv=self.ncv, mpd=self.mpd)
         self.es.setOperators(self.A_mat, self.M_mat)

firedrake/matrix.py

@@ -97,6 +97,19 @@ def __str__(self):
         return "assembled %s(a=%s, bcs=%s)" % (type(self).__name__,
                                                self.a, self.bcs)
 
+    def assign(self, val):
+        """Set matrix entries."""
+        if isinstance(val, MatrixBase):
+            val.petscmat.copy(self.petscmat)
+        else:
+            raise ValueError(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.
@@ -207,8 +220,7 @@ def mat(self):
         return self.petscmat
 
     def __add__(self, other):
-        if isinstance(other, AssembledMatrix):
+        if isinstance(other, MatrixBase):
             return self.petscmat + other.petscmat
         else:
-            raise TypeError("Unable to add %s to AssembledMatrix"
-                            % (type(other)))
+            return NotImplemented

firedrake/supermeshing.py

@@ -193,7 +193,7 @@ def likely(cell_A):
     V_S_A = FunctionSpace(reference_mesh, V_A.ufl_element())
     V_S_B = FunctionSpace(reference_mesh, V_B.ufl_element())
     M_SS = assemble(inner(TrialFunction(V_S_A), TestFunction(V_S_B)) * dx)
-    M_SS = M_SS.M.handle[:, :]
+    M_SS = M_SS.petscmat[:, :]
     node_locations_A = utils.physical_node_locations(V_S_A).dat.data_ro_with_halos
     node_locations_B = utils.physical_node_locations(V_S_B).dat.data_ro_with_halos
     num_nodes_A = node_locations_A.shape[0]

tests/firedrake/regression/test_assemble_baseform.py

@@ -111,7 +111,13 @@ 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)
+
+    res3 = Cofunction(res.function_space())
+    out = assemble(formsum, tensor=res3)
+    assert out is res3
+    for f, f3 in zip(preassemble.subfunctions, res3.subfunctions):
+        assert np.allclose(f.dat.data, f3.dat.data, atol=1e-12)
 
 
 def test_matrix_formsum(M):

tests/firedrake/regression/test_custom_pc_python_pmat.py

@@ -32,7 +32,7 @@ def test_python_pc_valueerror():
     u = TrialFunction(V)
     v = TestFunction(V)
 
-    A = assemble(inner(u, v) * dx).M.handle
+    A = assemble(inner(u, v) * dx).petscmat
     pc = PETSc.PC().create(comm=mesh.comm)
     pc.setType(pc.Type.PYTHON)
     pc.setOperators(A, A)

tests/firedrake/regression/test_matrix.py

@@ -22,7 +22,7 @@ def trial(V):
 
 @pytest.fixture
 def a(test, trial):
-    return inner(trial, test)*dx
+    return inner(trial, test) * dx
 
 
 @pytest.fixture(params=["nest", "aij", "matfree"])
@@ -36,17 +36,13 @@ def test_assemble_returns_matrix(a):
     assert isinstance(A, matrix.Matrix)
 
 
-def test_solve_with_assembled_matrix():
-    mesh = UnitIntervalMesh(3)
-    V = FunctionSpace(mesh, "CG", 1)
-
-    u = TrialFunction(V)
-    v = TestFunction(V)
-    x, = SpatialCoordinate(mesh)
+def test_solve_with_assembled_matrix(a):
+    (v, u) = a.arguments()
+    V = v.function_space()
+    x, = SpatialCoordinate(V.mesh())
     f = Function(V).interpolate(x)
 
-    a = inner(u, v) * dx
-    A = AssembledMatrix((v, u), bcs=(), petscmat=assemble(a).M.handle)
+    A = AssembledMatrix((v, u), bcs=(), petscmat=assemble(a).petscmat)
     L = inner(f, v) * dx
 
     solution = Function(V)

tests/firedrake/regression/test_projection_zany.py

@@ -88,8 +88,7 @@ def test_mass_conditioning(element, degree, hierarchy):
         u = TrialFunction(V)
         v = TestFunction(V)
         a = inner(u, v)*dx
-        B = assemble(a, mat_type="aij").M.handle
-        A = B.convert("dense").getDenseArray()
+        A = assemble(a, mat_type="aij").petscmat[:, :]
         kappa = np.linalg.cond(A)
 
         mass_cond.append(kappa)

tests/firedrake/regression/test_stress_elements.py

@@ -152,8 +152,7 @@ def test_mass_conditioning(stress_element, mesh_hierarchy):
         tau = TestFunction(Sig)
         mass = inner(sigh, tau)*dx
         a = derivative(mass, sigh)
-        B = assemble(a, mat_type="aij").M.handle
-        A = B.convert("dense").getDenseArray()
+        A = assemble(a, mat_type="aij").petscmat[:, :]
         kappa = np.linalg.cond(A)
 
         mass_cond.append(kappa)

tests/firedrake/regression/test_vfs_component_bcs.py

@@ -214,15 +214,12 @@ def test_component_full_bcs(V):
     v = TestFunction(V)
     a = inner(grad(u), grad(v))*dx
 
-    def asarray(A):
-        return A.M.handle[:, :]
+    A_full = assemble(a, bcs=bcs_full, mat_type="aij")
+    A_cmp = assemble(a, bcs=bcs_cmp, mat_type="aij")
+    A_mixed = assemble(a, bcs=bcs_mixed, mat_type="aij")
 
-    A_full = asarray(assemble(a, bcs=bcs_full, mat_type="aij"))
-    A_cmp = asarray(assemble(a, bcs=bcs_cmp, mat_type="aij"))
-    A_mixed = asarray(assemble(a, bcs=bcs_mixed, mat_type="aij"))
-
-    assert np.allclose(A_full, A_cmp)
-    assert np.allclose(A_mixed, A_full)
+    assert A_full.petscmat.equal(A_cmp.petscmat)
+    assert A_mixed.equal(A_full)
 
 
 def test_component_full_bcs_overlap(V):
@@ -240,10 +237,7 @@ def test_component_full_bcs_overlap(V):
 
     a = inner(grad(u), grad(v)) * dx
 
-    def asarray(A):
-        return A.M.handle[:, :]
-
-    A_1 = asarray(assemble(a, bcs=bcs_1, mat_type="aij"))
-    A_2 = asarray(assemble(a, bcs=bcs_2, mat_type="aij"))
+    A_1 = assemble(a, bcs=bcs_1, mat_type="aij")
+    A_2 = assemble(a, bcs=bcs_2, mat_type="aij")
 
-    assert np.allclose(A_1, A_2)
+    assert A_1.petscmat.equal(A_2.petscmat)

tests/firedrake/supermesh/test_assemble_mixed_mass_matrix.py

@@ -49,7 +49,7 @@ def test_assemble_mixed_mass_matrix(mesh, family_A, family_B, degree_A, degree_B
     M = assemble_mixed_mass_matrix(V_A, V_B)
 
     M_ex = assemble(inner(TrialFunction(V_A), TestFunction(V_B)) * dx)
-    M_ex = M_ex.M.handle
+    M_ex = M_ex.petscmat
 
     M_ex.axpy(-1.0, M)
     nrm = M_ex.norm(PETSc.NormType.NORM_INFINITY)