diff --git a/firedrake/assemble.py b/firedrake/assemble.py index dcf57978ca..fa931c0434 100644 --- a/firedrake/assemble.py +++ b/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={}): diff --git a/firedrake/constant.py b/firedrake/constant.py index 42854145ad..74b7953905 100644 --- a/firedrake/constant.py +++ b/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") diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 07f97caa8f..533d0c698f 100644 --- a/firedrake/matrix.py +++ b/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))) diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 54cc0b183c..444c23cc8e 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/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,9 +180,7 @@ 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() @@ -150,9 +188,7 @@ def test_tensor_copy(a, 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): diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index b2d8f64463..5e36610a33 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/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