assemble: introduce get_assembler

allocate_matrix(...) -> get_assembler(...).allocate()
get_form_assembler(..., tensor=tensor) -> get_assembler(...).assemble(tensor=tensor)

firedrake/external_operators/abstract_external_operators.py

@@ -3,7 +3,7 @@
 from ufl.argument import BaseArgument
 
 import firedrake.ufl_expr as ufl_expr
-from firedrake.assemble import allocate_matrix
+from firedrake.assemble import get_assembler
 from firedrake.function import Function
 from firedrake.cofunction import Cofunction
 from firedrake.matrix import MatrixBase
@@ -202,7 +202,6 @@ def _matrix_builder(self, bcs, opts, integral_types):
 
         This helper function provides a way to allocate matrices that can then be populated
         in the assembly method(s) of the external operator subclass.
-        This function relies on the :func:`firedrake.assemble.allocate_matrix` function.
 
         Parameters
         ----------
@@ -222,7 +221,7 @@ def _matrix_builder(self, bcs, opts, integral_types):
         # Remove `diagonal` keyword argument
         opts.pop('diagonal', None)
         # Allocate the matrix associated with `self`
-        return allocate_matrix(self, bcs=bcs, integral_types=integral_types, **opts)
+        return get_assembler(self, bcs=bcs, allocation_integral_types=integral_types, **opts).allocate()
 
     def _ufl_expr_reconstruct_(self, *operands, function_space=None, derivatives=None,
                                argument_slots=None, operator_data=None, add_kwargs={}):

firedrake/linear_solver.py

@@ -112,12 +112,12 @@ def trial_space(self):
 
     @cached_property
     def _rhs(self):
-        from firedrake.assemble import OneFormAssembler
+        from firedrake.assemble import get_assembler
 
         u = function.Function(self.trial_space)
         b = cofunction.Cofunction(self.test_space.dual())
         expr = -action(self.A.a, u)
-        return u, OneFormAssembler(expr).assemble, b
+        return u, get_assembler(expr).assemble, b
 
     def _lifted(self, b):
         u, update, blift = self._rhs

firedrake/matrix_free/operators.py

@@ -87,7 +87,7 @@ class ImplicitMatrixContext(object):
     @PETSc.Log.EventDecorator()
     def __init__(self, a, row_bcs=[], col_bcs=[],
                  fc_params=None, appctx=None):
-        from firedrake.assemble import get_form_assembler
+        from firedrake.assemble import get_assembler
 
         self.a = a
         self.aT = adjoint(a)
@@ -144,10 +144,10 @@ def __init__(self, a, row_bcs=[], col_bcs=[],
             elif isinstance(bc, EquationBCSplit):
                 self.bcs_action.append(bc.reconstruct(action_x=self._x))
 
-        self._assemble_action = get_form_assembler(self.action, tensor=self._ystar,
-                                                   bcs=self.bcs_action,
-                                                   form_compiler_parameters=self.fc_params,
-                                                   zero_bc_nodes=True)
+        self._assemble_action = get_assembler(self.action,
+                                              bcs=self.bcs_action,
+                                              form_compiler_parameters=self.fc_params,
+                                              zero_bc_nodes=True).assemble
 
         # For assembling action(adjoint(f), self._y)
         # Sorted list of equation bcs
@@ -161,13 +161,12 @@ def __init__(self, a, row_bcs=[], col_bcs=[],
         for bc in self.bcs:
             for ebc in bc.sorted_equation_bcs():
                 self._assemble_actionT.append(
-                    get_form_assembler(action(adjoint(ebc.f), self._y), tensor=self._xbc,
-                                       form_compiler_parameters=self.fc_params))
+                    get_assembler(action(adjoint(ebc.f), self._y),
+                                  form_compiler_parameters=self.fc_params).assemble)
         # Domain last
         self._assemble_actionT.append(
-            get_form_assembler(self.actionT,
-                               tensor=self._xstar if len(self.bcs) == 0 else self._xbc,
-                               form_compiler_parameters=self.fc_params))
+            get_assembler(self.actionT,
+                          form_compiler_parameters=self.fc_params).assemble)
 
     @cached_property
     def _diagonal(self):
@@ -177,13 +176,13 @@ def _diagonal(self):
 
     @cached_property
     def _assemble_diagonal(self):
-        from firedrake.assemble import get_form_assembler
-        return get_form_assembler(self.a, tensor=self._diagonal,
-                                  form_compiler_parameters=self.fc_params,
-                                  diagonal=True)
+        from firedrake.assemble import get_assembler
+        return get_assembler(self.a,
+                             form_compiler_parameters=self.fc_params,
+                             diagonal=True).assemble
 
     def getDiagonal(self, mat, vec):
-        self._assemble_diagonal()
+        self._assemble_diagonal(tensor=self._diagonal)
         diagonal_func = self._diagonal.riesz_representation(riesz_map="l2")
         for bc in self.bcs:
             # Operator is identity on boundary nodes
@@ -212,7 +211,7 @@ def mult(self, mat, X, Y):
         # If we are not, then the matrix just has 0s in the rows and columns.
         for bc in self.col_bcs:
             bc.zero(self._x)
-        self._assemble_action()
+        self._assemble_action(tensor=self._ystar)
         # This sets the essential boundary condition values on the
         # result.
         if self.on_diag:
@@ -307,14 +306,14 @@ def multTranspose(self, mat, Y, X):
                 # zero columns associated with DirichletBCs/EquationBCs
                 for obc in obj.bcs:
                     obc.zero(self._y)
-                aT()
+                aT(tensor=self._xbc)
                 self._xstar += self._xbc
         else:
             # No DirichletBC/EquationBC
             # There is only a single element in the list (for the domain equation).
             # Save to self._x directly
             aT, = self._assemble_actionT
-            aT()
+            aT(tensor=self._xstar)
 
         if self.on_diag:
             if len(self.col_bcs) > 0:

firedrake/preconditioners/assembled.py

@@ -20,7 +20,7 @@ class AssembledPC(PCBase):
     _prefix = "assembled_"
 
     def initialize(self, pc):
-        from firedrake.assemble import allocate_matrix, TwoFormAssembler
+        from firedrake.assemble import get_assembler
         A, P = pc.getOperators()
 
         if pc.getType() != "python":
@@ -51,11 +51,9 @@ def initialize(self, pc):
 
         (a, bcs) = self.form(pc, test, trial)
 
-        self.P = allocate_matrix(a, bcs=bcs,
-                                 form_compiler_parameters=fcp,
-                                 mat_type=mat_type,
-                                 options_prefix=options_prefix)
-        self._assemble_P = TwoFormAssembler(a, bcs=bcs, form_compiler_parameters=fcp, mat_type=mat_type).assemble
+        form_assembler = get_assembler(a, bcs=bcs, form_compiler_parameters=fcp, mat_type=mat_type, options_prefix=options_prefix)
+        self.P = form_assembler.allocate()
+        self._assemble_P = form_assembler.assemble
         self._assemble_P(tensor=self.P)
 
         # Transfer nullspace over

firedrake/preconditioners/facet_split.py

@@ -40,7 +40,7 @@ def get_permutation(self, V, W):
     def initialize(self, pc):
         from finat.ufl import RestrictedElement, MixedElement, TensorElement, VectorElement
         from firedrake import FunctionSpace, TestFunctions, TrialFunctions
-        from firedrake.assemble import allocate_matrix, TwoFormAssembler
+        from firedrake.assemble import get_assembler
 
         _, P = pc.getOperators()
         appctx = self.get_appctx(pc)
@@ -88,12 +88,9 @@ def restrict(ele, restriction_domain):
             self.iperm = self.perm.invertPermutation()
 
         if mat_type != "submatrix":
-            self.mixed_op = allocate_matrix(mixed_operator,
-                                            bcs=mixed_bcs,
-                                            form_compiler_parameters=fcp,
-                                            mat_type=mat_type,
-                                            options_prefix=options_prefix)
-            self._assemble_mixed_op = TwoFormAssembler(mixed_operator, bcs=mixed_bcs, form_compiler_parameters=fcp, mat_type=mat_type).assemble
+            form_assembler = get_assembler(mixed_operator, bcs=mixed_bcs, form_compiler_parameters=fcp, mat_type=mat_type, options_prefix=options_prefix)
+            self.mixed_op = form_assembler.allocate()
+            self._assemble_mixed_op = form_assembler.assemble
             self._assemble_mixed_op(tensor=self.mixed_op)
             mixed_opmat = self.mixed_op.petscmat
 

firedrake/preconditioners/fdm.py

@@ -149,10 +149,10 @@ def initialize(self, pc):
             self.work_vec_x = Amat.createVecLeft()
             self.work_vec_y = Amat.createVecRight()
             if use_amat:
-                from firedrake.assemble import allocate_matrix, TwoFormAssembler
-                self.A = allocate_matrix(J_fdm, bcs=bcs_fdm, form_compiler_parameters=fcp,
-                                         mat_type=mat_type, options_prefix=options_prefix)
-                self._assemble_A = TwoFormAssembler(J_fdm, bcs=bcs_fdm, form_compiler_parameters=fcp, mat_type=mat_type).assemble
+                from firedrake.assemble import get_assembler
+                form_assembler = get_assembler(J_fdm, bcs=bcs_fdm, form_compiler_parameters=fcp, mat_type=mat_type, options_prefix=options_prefix)
+                self.A = form_assembler.allocate()
+                self._assemble_A = form_assembler.assemble
                 self._assemble_A(tensor=self.A)
                 Amat = self.A.petscmat
 
@@ -538,15 +538,15 @@ def assemble_coefficients(self, J, fcp, block_diagonal=False):
             W = MixedFunctionSpace([c.function_space() for c in bdiags])
             tensor = Function(W, val=op2.MixedDat([c.dat for c in bdiags]))
         else:
-            from firedrake.assemble import OneFormAssembler
+            from firedrake.assemble import get_assembler
             tensor = Function(Z.dual())
-            assembly_callables.append(partial(OneFormAssembler(mixed_form, diagonal=True, form_compiler_parameters=fcp).assemble, tensor=tensor))
+            assembly_callables.append(partial(get_assembler(mixed_form, form_compiler_parameters=fcp, diagonal=True).assemble, tensor=tensor))
         coefficients = {"cell": tensor}
         facet_integrals = [i for i in J.integrals() if "facet" in i.integral_type()]
         J_facet = expand_indices(expand_derivatives(ufl.Form(facet_integrals)))
         if len(J_facet.integrals()) > 0:
             gamma = coefficients.setdefault("facet", Function(V.dual()))
-            assembly_callables.append(partial(OneFormAssembler(J_facet, diagonal=True, form_compiler_parameters=fcp).assemble, tensor=gamma))
+            assembly_callables.append(partial(get_assembler(J_facet, form_compiler_parameters=fcp, tensor=gamma, diagonal=True).assemble, tensor=gamma))
         return coefficients, assembly_callables
 
     @PETSc.Log.EventDecorator("FDMRefTensor")
@@ -2068,7 +2068,7 @@ def condense(self, A, J, bcs, fcp):
 
     @PETSc.Log.EventDecorator("FDMCoefficients")
     def assemble_coefficients(self, J, fcp):
-        from firedrake.assemble import OneFormAssembler
+        from firedrake.assemble import get_assembler
         coefficients = {}
         assembly_callables = []
 
@@ -2109,7 +2109,7 @@ def assemble_coefficients(self, J, fcp):
         if not isinstance(alpha, ufl.constantvalue.Zero):
             Q = FunctionSpace(mesh, finat.ufl.TensorElement(DG, shape=alpha.ufl_shape))
             tensor = coefficients.setdefault("alpha", Function(Q.dual()))
-            assembly_callables.append(partial(OneFormAssembler(ufl.inner(TestFunction(Q), alpha)*dx, form_compiler_parameters=fcp).assemble, tensor=tensor))
+            assembly_callables.append(partial(get_assembler(ufl.inner(TestFunction(Q), alpha)*dx, form_compiler_parameters=fcp).assemble, tensor=tensor))
 
         # get zero-th order coefficent
         ref_val = [ufl.variable(t) for t in args_J]
@@ -2130,7 +2130,7 @@ def assemble_coefficients(self, J, fcp):
                 beta = ufl.diag_vector(beta)
             Q = FunctionSpace(mesh, finat.ufl.TensorElement(DG, shape=beta.ufl_shape) if beta.ufl_shape else DG)
             tensor = coefficients.setdefault("beta", Function(Q.dual()))
-            assembly_callables.append(partial(OneFormAssembler(ufl.inner(TestFunction(Q), beta)*dx, form_compiler_parameters=fcp).assemble, tensor=tensor))
+            assembly_callables.append(partial(get_assembler(ufl.inner(TestFunction(Q), beta)*dx, form_compiler_parameters=fcp).assemble, tensor=tensor))
 
         family = "CG" if tdim == 1 else "DGT"
         degree = 1 if tdim == 1 else 0
@@ -2152,11 +2152,11 @@ def assemble_coefficients(self, J, fcp):
 
             Q = FunctionSpace(mesh, finat.ufl.TensorElement(DGT, shape=G.ufl_shape))
             tensor = coefficients.setdefault("Gq_facet", Function(Q.dual()))
-            assembly_callables.append(partial(OneFormAssembler(ifacet_inner(TestFunction(Q), G), form_compiler_parameters=fcp).assemble, tensor=tensor))
+            assembly_callables.append(partial(get_assembler(ifacet_inner(TestFunction(Q), G), form_compiler_parameters=fcp).assemble, tensor=tensor))
             PT = Piola.T
             Q = FunctionSpace(mesh, finat.ufl.TensorElement(DGT, shape=PT.ufl_shape))
             tensor = coefficients.setdefault("PT_facet", Function(Q.dual()))
-            assembly_callables.append(partial(OneFormAssembler(ifacet_inner(TestFunction(Q), PT), form_compiler_parameters=fcp).assemble, tensor=tensor))
+            assembly_callables.append(partial(get_assembler(ifacet_inner(TestFunction(Q), PT), form_compiler_parameters=fcp).assemble, tensor=tensor))
 
         # make DGT functions with BC flags
         shape = V.ufl_element().reference_value_shape
@@ -2182,7 +2182,7 @@ def assemble_coefficients(self, J, fcp):
         if len(forms):
             form = sum(forms)
             if len(form.arguments()) == 1:
-                assembly_callables.append(partial(OneFormAssembler(form, form_compiler_parameters=fcp).assemble, tensor=tensor))
+                assembly_callables.append(partial(get_assembler(form, form_compiler_parameters=fcp).assemble, tensor=tensor))
         # set arbitrary non-zero coefficients for preallocation
         for coef in coefficients.values():
             with coef.dat.vec as cvec:

firedrake/preconditioners/gtmg.py

@@ -13,7 +13,7 @@ class GTMGPC(PCBase):
 
     def initialize(self, pc):
         from firedrake import TestFunction, parameters
-        from firedrake.assemble import allocate_matrix, TwoFormAssembler
+        from firedrake.assemble import get_assembler
         from firedrake.interpolation import Interpolator
         from firedrake.solving_utils import _SNESContext
         from firedrake.matrix_free.operators import ImplicitMatrixContext
@@ -50,12 +50,9 @@ def initialize(self, pc):
 
             fine_mat_type = opts.getString(options_prefix + "mat_type",
                                            parameters["default_matrix_type"])
-            self.fine_op = allocate_matrix(fine_operator,
-                                           bcs=fine_bcs,
-                                           form_compiler_parameters=fcp,
-                                           mat_type=fine_mat_type,
-                                           options_prefix=options_prefix)
-            self._assemble_fine_op = TwoFormAssembler(fine_operator, bcs=fine_bcs, form_compiler_parameters=fcp, mat_type=fine_mat_type).assemble
+            fine_form_assembler = get_assembler(fine_operator, bcs=fine_bcs, form_compiler_parameters=fcp, mat_type=fine_mat_type, options_prefix=options_prefix)
+            self.fine_op = fine_form_assembler.allocate()
+            self._assemble_fine_op = fine_form_assembler.assemble
             self._assemble_fine_op(tensor=self.fine_op)
             fine_petscmat = self.fine_op.petscmat
         else:
@@ -88,12 +85,9 @@ def initialize(self, pc):
         get_coarse_nullspace = appctx.get("get_coarse_op_nullspace", None)
         get_coarse_transpose_nullspace = appctx.get("get_coarse_op_transpose_nullspace", None)
 
-        self.coarse_op = allocate_matrix(coarse_operator,
-                                         bcs=coarse_space_bcs,
-                                         form_compiler_parameters=fcp,
-                                         mat_type=coarse_mat_type,
-                                         options_prefix=coarse_options_prefix)
-        self._assemble_coarse_op = TwoFormAssembler(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type).assemble
+        coarse_form_assembler = get_assembler(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type, options_prefix=coarse_options_prefix)
+        self.coarse_op = coarse_form_assembler.allocate()
+        self._assemble_coarse_op = coarse_form_assembler.assemble
         self._assemble_coarse_op(tensor=self.coarse_op)
         coarse_opmat = self.coarse_op.petscmat
 

firedrake/preconditioners/hiptmair.py

@@ -32,7 +32,7 @@ def coarsen(self, pc):
         raise NotImplementedError
 
     def initialize(self, pc):
-        from firedrake.assemble import allocate_matrix, TwoFormAssembler
+        from firedrake.assemble import get_assembler
         A, P = pc.getOperators()
         appctx = self.get_appctx(pc)
         fcp = appctx.get("form_compiler_parameters")
@@ -47,12 +47,9 @@ def initialize(self, pc):
         coarse_options_prefix = options_prefix + "mg_coarse_"
         coarse_mat_type = opts.getString(coarse_options_prefix + "mat_type",
                                          parameters["default_matrix_type"])
-        self.coarse_op = allocate_matrix(coarse_operator,
-                                         bcs=coarse_space_bcs,
-                                         form_compiler_parameters=fcp,
-                                         mat_type=coarse_mat_type,
-                                         options_prefix=coarse_options_prefix)
-        self._assemble_coarse_op = TwoFormAssembler(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type).assemble
+        form_assembler = get_assembler(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type, options_prefix=coarse_options_prefix)
+        self.coarse_op = form_assembler.allocate()
+        self._assemble_coarse_op = form_assembler.assemble
         self._assemble_coarse_op(tensor=self.coarse_op)
         coarse_opmat = self.coarse_op.petscmat
 

firedrake/preconditioners/pcd.py

@@ -42,7 +42,7 @@ class PCDPC(PCBase):
     def initialize(self, pc):
         from firedrake import (TrialFunction, TestFunction, dx, inner,
                                grad, split, Constant, parameters)
-        from firedrake.assemble import allocate_matrix, assemble, TwoFormAssembler
+        from firedrake.assemble import assemble, get_assembler
         if pc.getType() != "python":
             raise ValueError("Expecting PC type python")
         prefix = pc.getOptionsPrefix() + "pcd_"
@@ -111,10 +111,9 @@ def initialize(self, pc):
         fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx
 
         self.Re = Re
-        self.Fp = allocate_matrix(fp, form_compiler_parameters=context.fc_params,
-                                  mat_type=self.Fp_mat_type,
-                                  options_prefix=prefix + "Fp_")
-        self._assemble_Fp = TwoFormAssembler(fp, form_compiler_parameters=context.fc_params, mat_type=self.Fp_mat_type).assemble
+        form_assembler = get_assembler(fp, bcs=None, form_compiler_parameters=context.fc_params, mat_type=self.Fp_mat_type, options_prefix=prefix + "Fp_")
+        self.Fp = form_assembler.allocate()
+        self._assemble_Fp = form_assembler.assemble
         self._assemble_Fp(tensor=self.Fp)
         Fpmat = self.Fp.petscmat
         self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]

firedrake/projection.py

@@ -218,9 +218,9 @@ def rhs_form(self):
 
     @cached_property
     def assembler(self):
-        from firedrake.assemble import OneFormAssembler
-        return OneFormAssembler(self.rhs_form,
-                                form_compiler_parameters=self.form_compiler_parameters).assemble
+        from firedrake.assemble import get_assembler
+        return get_assembler(self.rhs_form,
+                             form_compiler_parameters=self.form_compiler_parameters).assemble
 
     @property
     def rhs(self):