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,18 +460,19 @@ 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)
-                    return matrix.AssembledMatrix(expr, self._bcs, res,
-                                                  appctx=self._appctx,
-                                                  options_prefix=self._options_prefix)
+                    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)
                 else:
                     raise TypeError("Incompatible RHS for Action.")
             elif isinstance(lhs, (firedrake.Cofunction, firedrake.Function)):
@@ -493,30 +490,28 @@ 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
-                return matrix.AssembledMatrix(expr, self._bcs, res,
-                                              appctx=self._appctx,
-                                              options_prefix=self._options_prefix)
+                        # 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)
             else:
                 raise TypeError("Mismatching FormSum shapes")
         elif isinstance(expr, ufl.ExternalOperator):
@@ -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/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/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 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.
@@ -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

@@ -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_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)