diff --git a/demos/higher_order_mass_lumping/higher_order_mass_lumping.py.rst b/demos/higher_order_mass_lumping/higher_order_mass_lumping.py.rst index bb217a8e4f..f33f9f6348 100644 --- a/demos/higher_order_mass_lumping/higher_order_mass_lumping.py.rst +++ b/demos/higher_order_mass_lumping/higher_order_mass_lumping.py.rst @@ -146,9 +146,9 @@ The source is injected at the center of the unit square:: ricker = Constant(0.0) ricker.assign(RickerWavelet(t, freq)) -We also create a function `R` to save the assembled RHS vector:: +We also create a cofunction `R` to save the assembled RHS vector:: - R = Function(V) + R = Cofunction(V.dual()) Finally, we define the whole variational form :math:`F`, assemble it, and then create a cached PETSc `LinearSolver` object to efficiently timestep with:: diff --git a/firedrake/adjoint_utils/assembly.py b/firedrake/adjoint_utils/assembly.py index f7bdb287df..cac94728ae 100644 --- a/firedrake/adjoint_utils/assembly.py +++ b/firedrake/adjoint_utils/assembly.py @@ -16,7 +16,7 @@ def wrapper(form, *args, **kwargs): ad_block_tag = kwargs.pop("ad_block_tag", None) annotate = annotate_tape(kwargs) with stop_annotating(): - from firedrake.assemble import preprocess_base_form + from firedrake.assemble import BaseFormAssembler from firedrake.slate import slate if not isinstance(form, slate.TensorBase): # Preprocess the form at the annotation stage so that the `AssembleBlock` @@ -25,7 +25,7 @@ def wrapper(form, *args, **kwargs): # -> `interp = Action(Interpolate(v1, v0), f)` with `v1` and `v0` being respectively `Argument` # and `Coargument`. Differentiating `interp` is not currently supported as the action's left slot # is a 2-form. However, after preprocessing, we obtain `Interpolate(f, v0)`, which can be differentiated. - form = preprocess_base_form(form) + form = BaseFormAssembler.preprocess_base_form(form) kwargs['is_base_form_preprocessed'] = True output = assemble(form, *args, **kwargs) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index cc95b0a4ca..fb03494065 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1,5 +1,6 @@ import abc -from collections import OrderedDict +from collections import defaultdict +from collections.abc import Sequence # noqa: F401 import functools import itertools from itertools import product @@ -41,1018 +42,991 @@ @PETSc.Log.EventDecorator() @annotate_assemble def assemble(expr, *args, **kwargs): - r"""Evaluate expr. - - :arg expr: a :class:`~ufl.classes.BaseForm`, :class:`~ufl.classes.Expr` or - a :class:`~.slate.TensorBase` expression. - :arg tensor: Existing tensor object to place the result in. - :arg bcs: Iterable of boundary conditions to apply. - :kwarg diagonal: If assembling a matrix is it diagonal? - :kwarg form_compiler_parameters: Dictionary of parameters to pass to - the form compiler. Ignored if not assembling a :class:`~ufl.classes.Form`. + """Assemble. + + Parameters + ---------- + expr : ufl.classes.Expr or ufl.classes.BaseForm or slate.TensorBase + Object to assemble. + tensor : firedrake.function.Function or firedrake.cofunction.Cofunction or matrix.MatrixBase or None + Existing tensor object to place the result in. + bcs : Sequence + Iterable of boundary conditions to apply. + form_compiler_parameters : dict + Dictionary of parameters to pass to + the form compiler. Ignored if not assembling a `ufl.classes.Form`. Any parameters provided here will be overridden by parameters set on the - :class:`~ufl.classes.Measure` in the form. For example, if a + `ufl.classes.Measure` in the form. For example, if a ``quadrature_degree`` of 4 is specified in this argument, but a degree of 3 is requested in the measure, the latter will be used. - :kwarg mat_type: String indicating how a 2-form (matrix) should be + mat_type : str + String indicating how a 2-form (matrix) should be assembled -- either as a monolithic matrix (``"aij"`` or ``"baij"``), - a block matrix (``"nest"``), or left as a :class:`.ImplicitMatrix` giving + a block matrix (``"nest"``), or left as a `matrix.ImplicitMatrix` giving matrix-free actions (``'matfree'``). If not supplied, the default value in ``parameters["default_matrix_type"]`` is used. BAIJ differs from AIJ in that only the block sparsity rather than the dof sparsity is constructed. This can result in some memory savings, but does not work with all PETSc preconditioners. BAIJ matrices only make sense for non-mixed matrices. - :kwarg sub_mat_type: String indicating the matrix type to + sub_mat_type : str + String indicating the matrix type to use *inside* a nested block matrix. Only makes sense if ``mat_type`` is ``nest``. May be one of ``"aij"`` or ``"baij"``. If not supplied, defaults to ``parameters["default_sub_matrix_type"]``. - :kwarg appctx: Additional information to hang on the assembled + options_prefix : str + PETSc options prefix to apply to matrices. + appctx : dict + Additional information to hang on the assembled matrix if an implicit matrix is requested (mat_type ``"matfree"``). - :kwarg options_prefix: PETSc options prefix to apply to matrices. - :kwarg zero_bc_nodes: If ``True``, set the boundary condition nodes in the + zero_bc_nodes : bool + If `True`, set the boundary condition nodes in the output tensor to zero rather than to the values prescribed by the - boundary condition. Default is ``False``. - :kwarg weight: weight of the boundary condition, i.e. the scalar in front of the + boundary condition. Default is `False`. + diagonal : bool + If assembling a matrix is it diagonal? + weight : float + Weight of the boundary condition, i.e. the scalar in front of the identity matrix corresponding to the boundary nodes. To discretise eigenvalue problems set the weight equal to 0.0. - - :returns: See below. - - If expr is a :class:`~ufl.classes.Form` or Slate tensor expression then - this evaluates the corresponding integral(s) and returns a :class:`float` - for 0-forms, a :class:`.Function` for 1-forms and a :class:`.Matrix` or - :class:`.ImplicitMatrix` for 2-forms. In the case of 2-forms the rows - correspond to the test functions and the columns to the trial functions. + allocation_integral_types : Sequence + `Sequence` of integral types to be used when allocating the output + `matrix.Matrix`. + is_base_form_preprocessed : bool + If `True`, skip preprocessing of the form. + + Returns + ------- + float or firedrake.function.Function or firedrake.cofunction.Cofunction or matrix.MatrixBase + Result of assembly. + + Notes + ----- + Input arguments are all optional, except ``expr``. + + If expr is a `ufl.classes.Form` or `slate.TensorBase` then this evaluates + the corresponding integral(s) and returns a `float` for 0-forms, + a `firedrake.function.Function` or a `firedrake.cofunction.Cofunction` + for 1-forms and a `matrix.Matrix` or a `matrix.ImplicitMatrix` for 2-forms. + In the case of 2-forms the rows correspond to the test functions and the + columns to the trial functions. If expr is an expression other than a form, it will be evaluated - pointwise on the :class:`.Function`\s in the expression. This will + pointwise on the `firedrake.function.Function`s in the expression. This will only succeed if all the Functions are on the same - :class:`.FunctionSpace`. + `firedrake.functionspace.FunctionSpace`. If ``tensor`` is supplied, the assembled result will be placed there, otherwise a new object of the appropriate type will be returned. If ``bcs`` is supplied and ``expr`` is a 2-form, the rows and columns - of the resulting :class:`.Matrix` corresponding to boundary nodes + of the resulting `matrix.Matrix` corresponding to boundary nodes will be set to 0 and the diagonal entries to 1. If ``expr`` is a 1-form, the vector entries at boundary nodes are set to the boundary condition values. - """ - if isinstance(expr, (ufl.form.BaseForm, slate.TensorBase)): - return assemble_base_form(expr, *args, **kwargs) - elif isinstance(expr, ufl.core.expr.Expr): - return _assemble_expr(expr) - else: - raise TypeError(f"Unable to assemble: {expr}") - - -def base_form_postorder_traversal(expr, visitor, visited={}): - if expr in visited: - return visited[expr] - - stack = [expr] - while stack: - e = stack.pop() - unvisited_children = [] - operands = base_form_operands(e) - for arg in operands: - if arg not in visited: - unvisited_children.append(arg) - - if unvisited_children: - stack.append(e) - stack.extend(unvisited_children) - else: - visited[e] = visitor(e, *(visited[arg] for arg in operands)) - - return visited[expr] - - -def base_form_preorder_traversal(expr, visitor, visited={}): - if expr in visited: - return visited[expr] - - stack = [expr] - while stack: - e = stack.pop() - unvisited_children = [] - operands = base_form_operands(e) - for arg in operands: - if arg not in visited: - unvisited_children.append(arg) - - if unvisited_children: - stack.extend(unvisited_children) - visited[e] = visitor(e) - - return visited[expr] - - -def reconstruct_node_from_operands(expr, operands): - if isinstance(expr, (ufl.Adjoint, ufl.Action)): - return expr._ufl_expr_reconstruct_(*operands) - elif isinstance(expr, ufl.FormSum): - return ufl.FormSum(*[(op, w) for op, w in zip(operands, expr.weights())]) - return expr - - -def base_form_operands(expr): - if isinstance(expr, (ufl.FormSum, ufl.Adjoint, ufl.Action)): - return expr.ufl_operands - if isinstance(expr, ufl.Form): - # Use reversed to treat base form operators - # in the order in which they have been made. - return list(reversed(expr.base_form_operators())) - if isinstance(expr, ufl.core.base_form_operator.BaseFormOperator): - # Conserve order - children = dict.fromkeys(e for e in (expr.argument_slots() + expr.ufl_operands) - if isinstance(e, ufl.form.BaseForm)) - return list(children) - return [] - - -def restructure_base_form_postorder(expression, visited=None): - visited = visited or {} + """ + if args: + raise RuntimeError(f"Got unexpected args: {args}") + tensor = kwargs.pop("tensor", None) + return get_assembler(expr, *args, **kwargs).assemble(tensor=tensor) - def visitor(expr, *operands): - # Need to reconstruct the expression with its visited operands! - expr = reconstruct_node_from_operands(expr, operands) - # Perform the DAG restructuring when needed - return restructure_base_form(expr, visited) - return base_form_postorder_traversal(expression, visitor, visited) +def get_assembler(form, *args, **kwargs): + """Create an assembler. + Notes + ----- + See `assemble` for descriptions of the parameters. ``tensor`` should not be passed to this function. -def restructure_base_form_preorder(expression, visited=None): - visited = visited or {} + """ + is_base_form_preprocessed = kwargs.pop('is_base_form_preprocessed', False) + bcs = kwargs.get('bcs', None) + fc_params = kwargs.get('form_compiler_parameters', None) + if isinstance(form, ufl.form.BaseForm) and not is_base_form_preprocessed: + mat_type = kwargs.get('mat_type', None) + # Preprocess the DAG and restructure the DAG + # Only pre-process `form` once beforehand to avoid pre-processing for each assembly call + form = BaseFormAssembler.preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params) + if isinstance(form, (ufl.form.Form, slate.TensorBase)) and not BaseFormAssembler.base_form_operands(form): + diagonal = kwargs.pop('diagonal', False) + if len(form.arguments()) == 0: + return ZeroFormAssembler(form, form_compiler_parameters=fc_params) + elif len(form.arguments()) == 1 or diagonal: + return OneFormAssembler(form, *args, bcs=bcs, form_compiler_parameters=fc_params, needs_zeroing=kwargs.get('needs_zeroing', True), + zero_bc_nodes=kwargs.get('zero_bc_nodes', False), diagonal=diagonal) + elif len(form.arguments()) == 2: + return TwoFormAssembler(form, *args, **kwargs) + else: + raise ValueError('Expecting a 0-, 1-, or 2-form: got %s' % (form)) + elif isinstance(form, ufl.core.expr.Expr) and not isinstance(form, ufl.core.base_form_operator.BaseFormOperator): + # BaseForm preprocessing can turn BaseForm into an Expr (cf. case (6) in `restructure_base_form`) + return ExprAssembler(form) + elif isinstance(form, ufl.form.BaseForm): + return BaseFormAssembler(form, *args, **kwargs) + else: + raise ValueError(f'Expecting a BaseForm, slate.TensorBase, or Expr object: got {form}') - def visitor(expr): - # Perform the DAG restructuring when needed - return restructure_base_form(expr, visited) - expression = base_form_preorder_traversal(expression, visitor, visited) - # Need to reconstruct the expression at the end when all its operands have been visited! - operands = [visited.get(args, args) for args in base_form_operands(expression)] - return reconstruct_node_from_operands(expression, operands) +class ExprAssembler(object): + """Expression assembler. + Parameters + ---------- + expr : ufl.core.expr.Expr + Expression. -def restructure_base_form(expr, visited=None): - r"""Perform a preorder traversal to simplify and optimize the DAG. - Example: Let's consider F(u, N(u; v*); v) with N(u; v*) a base form operator. + """ - We have: dFdu = \frac{\partial F}{\partial u} + Action(dFdN, dNdu) - Now taking the action on a rank-1 object w (e.g. Coefficient/Cofunction) results in: + def __init__(self, expr): + self._expr = expr - (1) Action(Action(dFdN, dNdu), w) + def assemble(self, tensor=None): + """Assemble the pointwise expression. - Action Action - / \ / \ - Action w -----> dFdN Action - / \ / \ - dFdN dNdu dNdu w + Parameters + ---------- + tensor : firedrake.function.Function or firedrake.cofunction.Cofunction or matrix.MatrixBase + Output tensor. - This situations does not only arise for BaseFormOperator but also when we have a 2-form instead of dNdu! + Returns + ------- + float or firedrake.cofunction.Cofunction or firedrake.function.Function or matrix.MatrixBase + Result of evaluation: `float` for 0-forms, `firedrake.cofunction.Cofunction` or `firedrake.function.Function` for 1-forms, and `matrix.MatrixBase` for 2-forms. - (2) Action(dNdu, w) + """ + from ufl.algorithms.analysis import extract_base_form_operators + from ufl.checks import is_scalar_constant_expression - Action - / \ - / w -----> dNdu(u; w, v*) - / - dNdu(u; uhat, v*) + assert tensor is None + expr = self._expr + # Get BaseFormOperators (e.g. `Interpolate` or `ExternalOperator`) + base_form_operators = extract_base_form_operators(expr) - (3) Action(F, N) + # -- Linear combination involving 2-form BaseFormOperators -- # + # Example: a * dNdu1(u1, u2; v1, v*) + b * dNdu2(u1, u2; v2, v*) + # with u1, u2 Functions, v1, v2, v* BaseArguments, dNdu1, dNdu2 BaseFormOperators, and a, b scalars. + if len(base_form_operators) and any(len(e.arguments()) > 1 for e in base_form_operators): + if isinstance(expr, ufl.algebra.Sum): + a, b = [assemble(e) for e in expr.ufl_operands] + # Only Expr resulting in a Matrix if assembled are BaseFormOperator + if not all(isinstance(op, matrix.AssembledMatrix) for op in (a, b)): + raise TypeError('Mismatching Sum shapes') + return get_assembler(ufl.FormSum((a, 1), (b, 1))).assemble() + elif isinstance(expr, ufl.algebra.Product): + a, b = expr.ufl_operands + scalar = [e for e in expr.ufl_operands if is_scalar_constant_expression(e)] + if scalar: + base_form = a if a is scalar else b + assembled_mat = assemble(base_form) + return get_assembler(ufl.FormSum((assembled_mat, scalar[0]))).assemble() + a, b = [assemble(e) for e in (a, b)] + return get_assembler(ufl.action(a, b)).assemble() + # -- Linear combination of Functions and 1-form BaseFormOperators -- # + # Example: a * u1 + b * u2 + c * N(u1; v*) + d * N(u2; v*) + # with u1, u2 Functions, N a BaseFormOperator, and a, b, c, d scalars or 0-form BaseFormOperators. + else: + base_form_operators = extract_base_form_operators(expr) + assembled_bfops = [firedrake.assemble(e) for e in base_form_operators] + # Substitute base form operators with their output before examining the expression + # which avoids conflict when determining function space, for example: + # extract_coefficients(Interpolate(u, V2)) with u \in V1 will result in an output function space V1 + # instead of V2. + if base_form_operators: + expr = ufl.replace(expr, dict(zip(base_form_operators, assembled_bfops))) + try: + coefficients = ufl.algorithms.extract_coefficients(expr) + V, = set(c.function_space() for c in coefficients) - {None} + except ValueError: + raise ValueError("Cannot deduce correct target space from pointwise expression") + return firedrake.Function(V).assign(expr) - Action F - / \ -----> F(..., N)[v] = | - F[v] N N - (4) Adjoint(dNdu) +class AbstractFormAssembler(abc.ABC): + """Abstract assembler class for forms. - Adjoint - | -----> dNdu(u; v*, uhat) - dNdu(u; uhat, v*) + Parameters + ---------- + form : ufl.form.BaseForm or slate.TensorBase + ``form`` to assemble. + bcs : DirichletBC or EquationBCSplit or Sequence + Boundary conditions. + form_compiler_parameters : dict + ``form_compiler_parameters`` to use. - (5) N(u; w) (scalar valued) + """ + def __init__(self, form, bcs=None, form_compiler_parameters=None): + self._form = form + self._bcs = solving._extract_bcs(bcs) + if any(isinstance(bc, EquationBC) for bc in self._bcs): + raise TypeError("EquationBC objects not expected here. " + "Preprocess by extracting the appropriate form with bc.extract_form('Jp') or bc.extract_form('J')") + self._form_compiler_params = form_compiler_parameters or {} - Action - N(u; w) ----> / \ = Action(N, w) - N(u; v*) w + @abc.abstractmethod + def allocate(self): + """Allocate memory for the output tensor.""" - So from Action(Action(dFdN, dNdu(u; v*)), w) we get: + @abc.abstractmethod + def assemble(self, tensor=None): + """Assemble the form. - Action Action Action - / \ (1) / \ (2) / \ (4) dFdN - Action w ----> dFdN Action ----> dFdN dNdu(u; w, v*) ----> dFdN(..., dNdu(u; w, v*)) = | - / \ / \ dNdu(u; w, v*) - dFdN dNdu dNdu w + Parameters + ---------- + tensor : firedrake.cofunction.Cofunction or firedrake.function.Function or matrix.MatrixBase + Output tensor to contain the result of assembly; if `None`, a tensor of appropriate type is created. - (6) ufl.FormSum(dN1du(u; w, v*), dN2du(u; w, v*)) -> ufl.Sum(dN1du(u; w, v*), dN2du(u; w, v*)) + Returns + ------- + float or firedrake.cofunction.Cofunction or firedrake.function.Function or matrix.MatrixBase + Result of assembly: `float` for 0-forms, `firedrake.cofunction.Cofunction` or `firedrake.function.Function` for 1-forms, and `matrix.MatrixBase` for 2-forms. - Let's consider `Action(dN1du, w) + Action(dN2du, w)`, we have: + """ - FormSum (2) FormSum (6) Sum - / \ ----> / \ ----> / \ - / \ / \ / \ - Action(dN1du, w) Action(dN2du, w) dN1du(u; w, v*) dN2du(u; w, v*) dN1du(u; w, v*) dN2du(u; w, v*) - This case arises as a consequence of (2) which turns sum of `Action`s (i.e. ufl.FormSum since Action is a BaseForm) - into sum of `BaseFormOperator`s (i.e. ufl.Sum since BaseFormOperator is an Expr as well). +class BaseFormAssembler(AbstractFormAssembler): + """Base form assembler. - (7) Action(w*, dNdu) + Parameters + ---------- + form : ufl.form.BaseForm + `ufl.form.BaseForm` to assemble. - Action - / \ - w* \ -----> dNdu(u; v0, w*) - \ - dNdu(u; v1, v0*) + Notes + ----- + See `AbstractFormAssembler` and `assemble` for descriptions of the other parameters. - It uses a recursive approach to reconstruct the DAG as we traverse it, enabling to take into account - various dag rotations/manipulations in expr. """ - if isinstance(expr, ufl.Action): - left, right = expr.ufl_operands - is_rank_1 = lambda x: isinstance(x, (firedrake.Cofunction, firedrake.Function, firedrake.Argument)) or len(x.arguments()) == 1 - is_rank_2 = lambda x: len(x.arguments()) == 2 - - # -- Case (1) -- # - # If left is Action and has a rank 2, then it is an action of a 2-form on a 2-form - if isinstance(left, ufl.Action) and is_rank_2(left): - return ufl.action(left.left(), ufl.action(left.right(), right)) - # -- Case (2) (except if left has only 1 argument, i.e. we have done case (5)) -- # - if isinstance(left, ufl.core.base_form_operator.BaseFormOperator) and is_rank_1(right) and len(left.arguments()) != 1: - # Retrieve the highest numbered argument - arg = max(left.arguments(), key=lambda v: v.number()) - return ufl.replace(left, {arg: right}) - # -- Case (3) -- # - if isinstance(left, ufl.Form) and is_rank_1(right): - # 1) Replace the highest-numbered argument of left by right when needed - # -> e.g. if right is a BaseFormOperator with 1 argument. - # Or - # 2) Let expr as it is by returning `ufl.Action(left, right)`. - return ufl.action(left, right) - # -- Case (7) -- # - if is_rank_1(left) and isinstance(right, ufl.core.base_form_operator.BaseFormOperator) and len(right.arguments()) != 1: - # Action(w*, dNdu(u; v1, v*)) -> dNdu(u; v0, w*) - # Get lowest numbered argument - arg = min(right.arguments(), key=lambda v: v.number()) - # Need to replace lowest numbered argument of right by left - replace_map = {arg: left} - # Decrease number for all the other arguments since the lowest numbered argument will be replaced. - other_args = [a for a in right.arguments() if a is not arg] - new_args = [firedrake.Argument(a.function_space(), number=a.number()-1, part=a.part()) for a in other_args] - replace_map.update(dict(zip(other_args, new_args))) - # Replace arguments - return ufl.replace(right, replace_map) - - # -- Case (4) -- # - if isinstance(expr, ufl.Adjoint) and isinstance(expr.form(), ufl.core.base_form_operator.BaseFormOperator): - B = expr.form() - u, v = B.arguments() - # Let V1 and V2 be primal spaces, B: V1 -> V2 and B*: V2* -> V1*: - # Adjoint(B(Argument(V1, 1), Argument(V2.dual(), 0))) = B(Argument(V1, 0), Argument(V2.dual(), 1)) - reordered_arguments = (firedrake.Argument(u.function_space(), number=v.number(), part=v.part()), - firedrake.Argument(v.function_space(), number=u.number(), part=u.part())) - # Replace arguments in argument slots - return ufl.replace(B, dict(zip((u, v), reordered_arguments))) - - # -- Case (5) -- # - if isinstance(expr, ufl.core.base_form_operator.BaseFormOperator) and not expr.arguments(): - # We are assembling a BaseFormOperator of rank 0 (no arguments). - # B(f, u*) be a BaseFormOperator with u* a Cofunction and f a Coefficient, then: - # B(f, u*) <=> Action(B(f, v*), f) where v* is a Coargument - ustar, *_ = expr.argument_slots() - vstar = firedrake.Argument(ustar.function_space(), 0) - expr = ufl.replace(expr, {ustar: vstar}) - return ufl.action(expr, ustar) - - # -- Case (6) -- # - if isinstance(expr, ufl.FormSum) and all(isinstance(c, ufl.core.base_form_operator.BaseFormOperator) for c in expr.components()): - # Return ufl.Sum - return sum([c for c in expr.components()]) - return expr - - -def assemble_base_form(expression, tensor=None, bcs=None, - diagonal=False, - mat_type=None, - sub_mat_type=None, - form_compiler_parameters=None, - appctx=None, - options_prefix=None, - zero_bc_nodes=False, - is_base_form_preprocessed=False, - weight=1.0, - visited=None): - r"""Evaluate expression. - - :arg expression: a :class:`~ufl.classes.BaseForm` - :kwarg tensor: Existing tensor object to place the result in. - :kwarg bcs: Iterable of boundary conditions to apply. - :kwarg diagonal: If assembling a matrix is it diagonal? - :kwarg mat_type: String indicating how a 2-form (matrix) should be - assembled -- either as a monolithic matrix (``"aij"`` or ``"baij"``), - a block matrix (``"nest"``), or left as a :class:`.ImplicitMatrix` giving - matrix-free actions (``'matfree'``). If not supplied, the default value in - ``parameters["default_matrix_type"]`` is used. BAIJ differs - from AIJ in that only the block sparsity rather than the dof - sparsity is constructed. This can result in some memory - savings, but does not work with all PETSc preconditioners. - BAIJ matrices only make sense for non-mixed matrices. - :kwarg sub_mat_type: String indicating the matrix type to - use *inside* a nested block matrix. Only makes sense if - ``mat_type`` is ``nest``. May be one of ``"aij"`` or ``"baij"``. If - not supplied, defaults to ``parameters["default_sub_matrix_type"]``. - :kwarg form_compiler_parameters: Dictionary of parameters to pass to - the form compiler. Ignored if not assembling a :class:`~ufl.classes.Form`. - Any parameters provided here will be overridden by parameters set on the - :class:`~ufl.classes.Measure` in the form. For example, if a - ``quadrature_degree`` of 4 is specified in this argument, but a degree of - 3 is requested in the measure, the latter will be used. - :kwarg appctx: Additional information to hang on the assembled - matrix if an implicit matrix is requested (mat_type ``"matfree"``). - :kwarg options_prefix: PETSc options prefix to apply to matrices. - :kwarg zero_bc_nodes: If ``True``, set the boundary condition nodes in the - output tensor to zero rather than to the values prescribed by the - boundary condition. Default is ``False``. - :kwarg is_base_form_preprocessed: If ``True``, skip preprocessing of the form. - :kwarg weight: weight of the boundary condition, i.e. the scalar in front of the - identity matrix corresponding to the boundary nodes. - To discretise eigenvalue problems set the weight equal to 0.0. - :returns: a :class:`float` for 0-forms, a :class:`.Cofunction` or a :class:`.Function` for 1-forms, - and a :class:`.MatrixBase` for 2-forms. - - This function assembles a :class:`~ufl.classes.BaseForm` object by traversing the corresponding DAG - in a post-order fashion and evaluating the nodes on the fly. - """ + def __init__(self, + form, + bcs=None, + form_compiler_parameters=None, + mat_type=None, + sub_mat_type=None, + options_prefix=None, + appctx=None, + zero_bc_nodes=False, + diagonal=False, + weight=1.0, + allocation_integral_types=None): + super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters) + self._mat_type = mat_type + self._sub_mat_type = sub_mat_type + self._options_prefix = options_prefix + self._appctx = appctx + self._zero_bc_nodes = zero_bc_nodes + self._diagonal = diagonal + self._weight = weight + self._allocation_integral_types = allocation_integral_types + + def allocate(self): + rank = len(self._form.arguments()) + if rank == 2 and not self._diagonal: + if self._mat_type == "matfree": + return MatrixFreeAssembler(self._form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, + options_prefix=self._options_prefix, + appctx=self._appctx).allocate() + else: + test, trial = self._form.arguments() + sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, self._mat_type, self._sub_mat_type, self.maps_and_regions) + return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, options_prefix=self._options_prefix) + else: + raise NotImplementedError("Only implemented for rank = 2 and diagonal = False") - # Preprocess the DAG and restructure the DAG - if not is_base_form_preprocessed and not isinstance(expression, slate.TensorBase): - expr = preprocess_base_form(expression, mat_type, form_compiler_parameters) - # BaseForm preprocessing can turn BaseForm into an Expr (cf. case (6) in `restructure_base_form`) - if isinstance(expr, ufl.core.expr.Expr) and not isinstance(expr, ufl.core.base_form_operator.BaseFormOperator): - return _assemble_expr(expr) - else: - expr = expression + @cached_property + def maps_and_regions(self): + # The sparsity could be made tighter by inspecting the form DAG. + test, trial = self._form.arguments() + return ExplicitMatrixAssembler._make_maps_and_regions_default(test, trial, self.allocation_integral_types) - # Define assembly DAG visitor - assembly_visitor = functools.partial(base_form_assembly_visitor, bcs=bcs, diagonal=diagonal, - form_compiler_parameters=form_compiler_parameters, - mat_type=mat_type, sub_mat_type=sub_mat_type, - appctx=appctx, options_prefix=options_prefix, - zero_bc_nodes=zero_bc_nodes, weight=weight) + @cached_property + def allocation_integral_types(self): + if self._allocation_integral_types is None: + # Use the most conservative integration types. + test, _ = self._form.arguments() + if test.function_space().mesh().extruded: + return ("interior_facet_vert", "interior_facet_horiz") + else: + return ("interior_facet", ) + else: + return self._allocation_integral_types - def visitor(e, *operands): - t = tensor if e is expr else None - return assembly_visitor(e, t, *operands) + def assemble(self, tensor=None): + """Assemble the form. - # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly. - visited = visited or {} - result = base_form_postorder_traversal(expr, visitor, visited) + Parameters + ---------- + tensor : firedrake.cofunction.Cofunction or firedrake.function.Function or matrix.MatrixBase + Output tensor to contain the result of assembly. - if tensor: - update_tensor(result, tensor) - return tensor - else: - return result - - -def preprocess_base_form(expr, mat_type=None, form_compiler_parameters=None): - """Preprocess ufl.BaseForm objects""" - original_expr = expr - if mat_type != "matfree": - # For "matfree", Form evaluation is delayed - expr = expand_derivatives_form(expr, form_compiler_parameters) - if not isinstance(expr, (ufl.form.Form, slate.TensorBase)): - # => No restructuring needed for Form and slate.TensorBase - expr = restructure_base_form_preorder(expr) - expr = restructure_base_form_postorder(expr) - # Preprocessing the form makes a new object -> current form caching mechanism - # will populate `expr`'s cache which is now different than `original_expr`'s cache so we need - # to transmit the cache. All of this only holds when both are `ufl.Form` objects. - if isinstance(original_expr, ufl.form.Form) and isinstance(expr, ufl.form.Form): - expr._cache = original_expr._cache - return expr - - -def update_tensor(assembled_base_form, tensor): - if isinstance(tensor, (firedrake.Function, firedrake.Cofunction)): - assembled_base_form.dat.copy(tensor.dat) - elif isinstance(tensor, matrix.MatrixBase): - # Uses the PETSc copy method. - assembled_base_form.petscmat.copy(tensor.petscmat) - else: - raise NotImplementedError("Cannot update tensor of type %s" % type(tensor)) + Returns + ------- + float or firedrake.cofunction.Cofunction or firedrake.function.Function or matrix.MatrixBase + Result of assembly: `float` for 0-forms, `firedrake.cofunction.Cofunction` or `firedrake.function.Function` for 1-forms, and `matrix.MatrixBase` for 2-forms. + Notes + ----- + This function assembles a `ufl.form.BaseForm` object by traversing the corresponding DAG + in a post-order fashion and evaluating the nodes on the fly. -def expand_derivatives_form(form, fc_params): - """Expand derivatives of ufl.BaseForm objects - :arg form: a :class:`~ufl.classes.BaseForm` - :arg fc_params:: Dictionary of parameters to pass to the form compiler. + """ + def visitor(e, *operands): + t = tensor if e is self._form else None + return self.base_form_assembly_visitor(e, t, *operands) - :returns: The resulting preprocessed :class:`~ufl.classes.BaseForm`. - This function preprocess the form, mainly by expanding the derivatives, in order to determine - if we are dealing with a :class:`~ufl.classes.Form` or another :class:`~ufl.classes.BaseForm` object. - This function is called in :func:`base_form_assembly_visitor`. Depending on the type of the resulting tensor, - we may call :func:`assemble_form` or traverse the sub-DAG via :func:`assemble_base_form`. - """ - if isinstance(form, ufl.form.Form): - from firedrake.parameters import parameters as default_parameters - from tsfc.parameters import is_complex + # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly. + visited = {} + result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited) - if fc_params is None: - fc_params = default_parameters["form_compiler"].copy() + if tensor: + BaseFormAssembler.update_tensor(result, tensor) + return tensor else: - # Override defaults with user-specified values - _ = fc_params - fc_params = default_parameters["form_compiler"].copy() - fc_params.update(_) - - complex_mode = fc_params and is_complex(fc_params.get("scalar_type")) - - return ufl.algorithms.preprocess_form(form, complex_mode) - # We also need to expand derivatives for `ufl.BaseForm` objects that are not `ufl.Form` - # Example: `Action(A, derivative(B, f))`, where `A` is a `ufl.BaseForm` and `B` can - # be `ufl.BaseForm`, or even an appropriate `ufl.Expr`, since assembly of expressions - # containing derivatives is not supported anymore but might be needed if the expression - # in question is within a `ufl.BaseForm` object. - return ufl.algorithms.ad.expand_derivatives(form) - - -def base_form_assembly_visitor(expr, tensor, *args, bcs, diagonal, - form_compiler_parameters, - mat_type, sub_mat_type, - appctx, options_prefix, - zero_bc_nodes, weight): - r"""Assemble a :class:`~ufl.classes.BaseForm` object given its assembled operands. - - This functions contains the assembly handlers corresponding to the different nodes that - can arise in a `~ufl.classes.BaseForm` object. It is called by :func:`assemble_base_form` - in a post-order fashion. - """ + return result - if isinstance(expr, (ufl.form.Form, slate.TensorBase)): - - if args and mat_type != "matfree": - # Retrieve the Form's children - base_form_operators = base_form_operands(expr) - # Substitute the base form operators by their output - expr = ufl.replace(expr, dict(zip(base_form_operators, args))) - - return _assemble_form(expr, tensor=tensor, bcs=bcs, - diagonal=diagonal, - mat_type=mat_type, - sub_mat_type=sub_mat_type, - appctx=appctx, - options_prefix=options_prefix, - form_compiler_parameters=form_compiler_parameters, - zero_bc_nodes=zero_bc_nodes, weight=weight) - - elif isinstance(expr, ufl.Adjoint): - if len(args) != 1: - raise TypeError("Not enough operands for Adjoint") - mat, = args - res = tensor.petscmat if tensor else PETSc.Mat() - petsc_mat = mat.petscmat - # Out-of-place Hermitian transpose - petsc_mat.hermitianTranspose(out=res) - (row, col) = mat.arguments() - return matrix.AssembledMatrix((col, row), bcs, res, - appctx=appctx, - options_prefix=options_prefix) - elif isinstance(expr, ufl.Action): - if (len(args) != 2): - raise TypeError("Not enough operands for Action") - lhs, rhs = args - if isinstance(lhs, matrix.MatrixBase): - if isinstance(rhs, (firedrake.Cofunction, firedrake.Function)): - 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) - 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, bcs, res, - appctx=appctx, - options_prefix=options_prefix) - else: - raise TypeError("Incompatible RHS for Action.") - elif isinstance(lhs, (firedrake.Cofunction, firedrake.Function)): - if isinstance(rhs, (firedrake.Cofunction, firedrake.Function)): - # Return scalar value - with lhs.dat.vec_ro as x, rhs.dat.vec_ro as y: - res = x.dot(y) - return res + def base_form_assembly_visitor(self, expr, tensor, *args): + r"""Assemble a :class:`~ufl.classes.BaseForm` object given its assembled operands. + + This functions contains the assembly handlers corresponding to the different nodes that + can arise in a `~ufl.classes.BaseForm` object. It is called by :func:`assemble_base_form` + in a post-order fashion. + """ + if isinstance(expr, (ufl.form.Form, slate.TensorBase)): + if args and self._mat_type != "matfree": + # Retrieve the Form's children + base_form_operators = BaseFormAssembler.base_form_operands(expr) + # Substitute the base form operators by their output + expr = ufl.replace(expr, dict(zip(base_form_operators, args))) + form = expr + rank = len(form.arguments()) + if rank == 0: + assembler = ZeroFormAssembler(form, form_compiler_parameters=self._form_compiler_params) + elif rank == 1 or (rank == 2 and self._diagonal): + assembler = OneFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, + zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal) + elif rank == 2: + assembler = TwoFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, + mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, + options_prefix=self._options_prefix, appctx=self._appctx, weight=self._weight, + allocation_integral_types=self.allocation_integral_types) else: - raise TypeError("Incompatible RHS for Action.") - else: - raise TypeError("Incompatible LHS for Action.") - elif isinstance(expr, ufl.FormSum): - if len(args) != len(expr.weights()): - raise TypeError("Mismatching weights and operands in FormSum") - if len(args) == 0: - raise TypeError("Empty FormSum") - if all(isinstance(op, float) for op in args): - return sum(args) - elif all(isinstance(op, firedrake.Cofunction) for op in args): - V, = set(a.function_space() for a in args) - res = sum([w*op.dat for (op, w) in zip(args, expr.weights())]) - return firedrake.Cofunction(V, res) - elif all(isinstance(op, ufl.Matrix) for op in args): + raise AssertionError + return assembler.assemble(tensor=tensor) + elif isinstance(expr, ufl.Adjoint): + if len(args) != 1: + raise TypeError("Not enough operands for Adjoint") + mat, = 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 + petsc_mat = mat.petscmat + # Out-of-place Hermitian transpose + 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: + raise TypeError("Not enough operands for Action") + lhs, rhs = args + if isinstance(lhs, matrix.MatrixBase): + if isinstance(rhs, (firedrake.Cofunction, firedrake.Function)): + 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) + 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) else: - # Copy to output tensor + raise TypeError("Incompatible RHS for Action.") + elif isinstance(lhs, (firedrake.Cofunction, firedrake.Function)): + if isinstance(rhs, (firedrake.Cofunction, firedrake.Function)): + # Return scalar value + with lhs.dat.vec_ro as x, rhs.dat.vec_ro as y: + res = x.dot(y) + return res + else: + raise TypeError("Incompatible RHS for Action.") + else: + raise TypeError("Incompatible LHS for Action.") + elif isinstance(expr, ufl.FormSum): + if len(args) != len(expr.weights()): + raise TypeError("Mismatching weights and operands in FormSum") + if len(args) == 0: + raise TypeError("Empty FormSum") + if all(isinstance(op, float) for op in args): + return sum(args) + elif all(isinstance(op, firedrake.Cofunction) for op in args): + V, = set(a.function_space() for a in args) + res = sum([w*op.dat for (op, w) in zip(args, expr.weights())]) + return firedrake.Cofunction(V, res) + 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 + 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) + else: + raise TypeError("Mismatching FormSum shapes") + elif isinstance(expr, ufl.ExternalOperator): + opts = {'form_compiler_parameters': self._form_compiler_params, + 'mat_type': self._mat_type, 'sub_mat_type': self._sub_mat_type, + 'appctx': self._appctx, 'options_prefix': self._options_prefix, + 'diagonal': self._diagonal} + # External operators might not have any children that needs to be assembled + # -> e.g. N(u; v0, w) with v0 a ufl.Argument and w a ufl.Coefficient + if args: + # Replace base forms in the operands and argument slots of the external operator by their result + v, *assembled_children = args + if assembled_children: + _, *children = BaseFormAssembler.base_form_operands(expr) + # Replace assembled children by their results + expr = ufl.replace(expr, dict(zip(children, assembled_children))) + # Always reconstruct the dual argument (0-slot argument) since it is a BaseForm + # 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) + elif isinstance(expr, ufl.Interpolate): + # Replace assembled children + _, expression = expr.argument_slots() + v, *assembled_expression = args + if assembled_expression: + # Occur in situations such as Interpolate composition + expression = assembled_expression[0] + expr = expr._ufl_expr_reconstruct_(expression, v) + + # Different assembly procedures: + # 1) Interpolate(Argument(V1, 1), Argument(V2.dual(), 0)) -> Jacobian (Interpolate matrix) + # 2) Interpolate(Coefficient(...), Argument(V2.dual(), 0)) -> Operator (or Jacobian action) + # 3) Interpolate(Argument(V1, 0), Argument(V2.dual(), 1)) -> Jacobian adjoint + # 4) Interpolate(Argument(V1, 0), Cofunction(...)) -> Action of the Jacobian adjoint + # This can be generalized to the case where the first slot is an arbitray expression. + rank = len(expr.arguments()) + # If argument numbers have been swapped => Adjoint. + arg_expression = ufl.algorithms.extract_arguments(expression) + is_adjoint = (arg_expression and arg_expression[0].number() == 0) + # Workaround: Renumber argument when needed since Interpolator assumes it takes a zero-numbered argument. + if not is_adjoint and rank != 1: + _, v1 = expr.arguments() + expression = ufl.replace(expression, {v1: firedrake.Argument(v1.function_space(), number=0, part=v1.part())}) + # Get the interpolator + interp_data = expr.interp_data + default_missing_val = interp_data.pop('default_missing_val', None) + interpolator = firedrake.Interpolator(expression, expr.function_space(), **interp_data) + # Assembly + if rank == 1: + # Assembling the action of the Jacobian adjoint. + if is_adjoint: + output = tensor or firedrake.Cofunction(arg_expression[0].function_space().dual()) + return interpolator._interpolate(v, output=output, transpose=True, default_missing_val=default_missing_val) + # Assembling the Jacobian action. + if interpolator.nargs: + return interpolator._interpolate(expression, output=tensor, default_missing_val=default_missing_val) + # Assembling the operator + if tensor is None: + return interpolator._interpolate(default_missing_val=default_missing_val) + return firedrake.Interpolator(expression, tensor, **interp_data)._interpolate(default_missing_val=default_missing_val) + elif rank == 2: + res = tensor.petscmat if tensor else PETSc.Mat() + # Get the interpolation matrix + op2_mat = interpolator.callable() + petsc_mat = op2_mat.handle + if is_adjoint: + # Out-of-place Hermitian transpose + petsc_mat.hermitianTranspose(out=res) + else: + # Copy the interpolation matrix into the output tensor petsc_mat.copy(result=res) - is_set = True - return matrix.AssembledMatrix(expr, bcs, res, - appctx=appctx, - options_prefix=options_prefix) - else: - raise TypeError("Mismatching FormSum shapes") - elif isinstance(expr, ufl.ExternalOperator): - opts = {'form_compiler_parameters': form_compiler_parameters, - 'mat_type': mat_type, 'sub_mat_type': sub_mat_type, - 'appctx': appctx, 'options_prefix': options_prefix, - 'diagonal': diagonal} - # External operators might not have any children that needs to be assembled - # -> e.g. N(u; v0, w) with v0 a ufl.Argument and w a ufl.Coefficient - if args: - # Replace base forms in the operands and argument slots of the external operator by their result - v, *assembled_children = args - if assembled_children: - _, *children = base_form_operands(expr) - # Replace assembled children by their results - expr = ufl.replace(expr, dict(zip(children, assembled_children))) - # Always reconstruct the dual argument (0-slot argument) since it is a BaseForm - # 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) - elif isinstance(expr, ufl.Interpolate): - # Replace assembled children - _, expression = expr.argument_slots() - v, *assembled_expression = args - if assembled_expression: - # Occur in situations such as Interpolate composition - expression = assembled_expression[0] - expr = expr._ufl_expr_reconstruct_(expression, v) - - # Different assembly procedures: - # 1) Interpolate(Argument(V1, 1), Argument(V2.dual(), 0)) -> Jacobian (Interpolate matrix) - # 2) Interpolate(Coefficient(...), Argument(V2.dual(), 0)) -> Operator (or Jacobian action) - # 3) Interpolate(Argument(V1, 0), Argument(V2.dual(), 1)) -> Jacobian adjoint - # 4) Interpolate(Argument(V1, 0), Cofunction(...)) -> Action of the Jacobian adjoint - # This can be generalized to the case where the first slot is an arbitray expression. - rank = len(expr.arguments()) - # If argument numbers have been swapped => Adjoint. - arg_expression = ufl.algorithms.extract_arguments(expression) - is_adjoint = (arg_expression and arg_expression[0].number() == 0) - # Workaround: Renumber argument when needed since Interpolator assumes it takes a zero-numbered argument. - if not is_adjoint and rank != 1: - _, v1 = expr.arguments() - expression = ufl.replace(expression, {v1: firedrake.Argument(v1.function_space(), number=0, part=v1.part())}) - # Get the interpolator - interp_data = expr.interp_data - default_missing_val = interp_data.pop('default_missing_val', None) - interpolator = firedrake.Interpolator(expression, expr.function_space(), **interp_data) - # Assembly - if rank == 1: - # Assembling the action of the Jacobian adjoint. - if is_adjoint: - output = tensor or firedrake.Cofunction(arg_expression[0].function_space().dual()) - return interpolator._interpolate(v, output=output, transpose=True, default_missing_val=default_missing_val) - # Assembling the Jacobian action. - if interpolator.nargs: - return interpolator._interpolate(expression, output=tensor, default_missing_val=default_missing_val) - # Assembling the operator - if tensor is None: - return interpolator._interpolate(default_missing_val=default_missing_val) - return firedrake.Interpolator(expression, tensor, **interp_data)._interpolate(default_missing_val=default_missing_val) - elif rank == 2: - res = tensor.petscmat if tensor else PETSc.Mat() - # Get the interpolation matrix - op2_mat = interpolator.callable() - petsc_mat = op2_mat.handle - if is_adjoint: - # Out-of-place Hermitian transpose - petsc_mat.hermitianTranspose(out=res) + return matrix.AssembledMatrix(expr.arguments(), self._bcs, res, + appctx=self._appctx, + options_prefix=self._options_prefix) else: - # Copy the interpolation matrix into the output tensor - petsc_mat.copy(result=res) - return matrix.AssembledMatrix(expr.arguments(), bcs, res, - appctx=appctx, - options_prefix=options_prefix) + # 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): + return expr 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): - return expr - else: - raise TypeError(f"Unrecognised BaseForm instance: {expr}") + raise TypeError(f"Unrecognised BaseForm instance: {expr}") + @staticmethod + def update_tensor(assembled_base_form, tensor): + if isinstance(tensor, (firedrake.Function, firedrake.Cofunction)): + assembled_base_form.dat.copy(tensor.dat) + elif isinstance(tensor, matrix.MatrixBase): + # Uses the PETSc copy method. + assembled_base_form.petscmat.copy(tensor.petscmat) + else: + raise NotImplementedError("Cannot update tensor of type %s" % type(tensor)) -@PETSc.Log.EventDecorator() -def allocate_matrix(expr, bcs=None, *, mat_type=None, sub_mat_type=None, - appctx=None, form_compiler_parameters=None, - integral_types=None, options_prefix=None): - r"""Allocate a matrix given an expression. + @staticmethod + def base_form_postorder_traversal(expr, visitor, visited={}): + if expr in visited: + return visited[expr] + + stack = [expr] + while stack: + e = stack.pop() + unvisited_children = [] + operands = BaseFormAssembler.base_form_operands(e) + for arg in operands: + if arg not in visited: + unvisited_children.append(arg) + + if unvisited_children: + stack.append(e) + stack.extend(unvisited_children) + else: + visited[e] = visitor(e, *(visited[arg] for arg in operands)) - .. warning:: + return visited[expr] - Do not use this function unless you know what you're doing. - """ - bcs = bcs or () - appctx = appctx or {} + @staticmethod + def base_form_preorder_traversal(expr, visitor, visited={}): + if expr in visited: + return visited[expr] - matfree = mat_type == "matfree" - arguments = expr.arguments() - if bcs is None: - bcs = () - else: - if any(isinstance(bc, EquationBC) for bc in bcs): - raise TypeError("EquationBC objects not expected here. " - "Preprocess by extracting the appropriate form with bc.extract_form('Jp') or bc.extract_form('J')") - if matfree: - return matrix.ImplicitMatrix(expr, bcs, - appctx=appctx, - fc_params=form_compiler_parameters, - options_prefix=options_prefix) - - integral_types = integral_types or set(i.integral_type() for i in expr.integrals()) - for bc in bcs: - integral_types.update(integral.integral_type() - for integral in bc.integrals()) - nest = mat_type == "nest" - if nest: - baij = sub_mat_type == "baij" - else: - baij = mat_type == "baij" - - if any(len(a.function_space()) > 1 for a in arguments) and mat_type == "baij": - raise ValueError("BAIJ matrix type makes no sense for mixed spaces, use 'aij'") - - get_cell_map = operator.methodcaller("cell_node_map") - get_extf_map = operator.methodcaller("exterior_facet_node_map") - get_intf_map = operator.methodcaller("interior_facet_node_map") - domains = OrderedDict((k, set()) for k in (get_cell_map, - get_extf_map, - get_intf_map)) - mapping = {"cell": (get_cell_map, op2.ALL), - "exterior_facet_bottom": (get_cell_map, op2.ON_BOTTOM), - "exterior_facet_top": (get_cell_map, op2.ON_TOP), - "interior_facet_horiz": (get_cell_map, op2.ON_INTERIOR_FACETS), - "exterior_facet": (get_extf_map, op2.ALL), - "exterior_facet_vert": (get_extf_map, op2.ALL), - "interior_facet": (get_intf_map, op2.ALL), - "interior_facet_vert": (get_intf_map, op2.ALL)} - for integral_type in integral_types: - try: - get_map, region = mapping[integral_type] - except KeyError: - raise ValueError(f"Unknown integral type '{integral_type}'") - domains[get_map].add(region) - - test, trial = arguments - map_pairs, iteration_regions = zip(*(((get_map(test), get_map(trial)), - tuple(sorted(regions))) - for get_map, regions in domains.items() - if regions)) - try: - sparsity = op2.Sparsity((test.function_space().dof_dset, - trial.function_space().dof_dset), - tuple(map_pairs), - iteration_regions=tuple(iteration_regions), - nest=nest, - block_sparse=baij) - except SparsityFormatError: - raise ValueError("Monolithic matrix assembly not supported for systems " - "with R-space blocks") - - return matrix.Matrix(expr, bcs, mat_type, sparsity, ScalarType, - options_prefix=options_prefix) + stack = [expr] + while stack: + e = stack.pop() + unvisited_children = [] + operands = BaseFormAssembler.base_form_operands(e) + for arg in operands: + if arg not in visited: + unvisited_children.append(arg) + if unvisited_children: + stack.extend(unvisited_children) -@PETSc.Log.EventDecorator() -def create_assembly_callable(expr, tensor=None, bcs=None, form_compiler_parameters=None, - mat_type=None, sub_mat_type=None, diagonal=False): - r"""Create a callable object than be used to assemble expr into a tensor. + visited[e] = visitor(e) - This is really only designed to be used inside residual and - jacobian callbacks, since it always assembles back into the - initially provided tensor. See also :func:`allocate_matrix`. + return visited[expr] - .. warning:: + @staticmethod + def reconstruct_node_from_operands(expr, operands): + if isinstance(expr, (ufl.Adjoint, ufl.Action)): + return expr._ufl_expr_reconstruct_(*operands) + elif isinstance(expr, ufl.FormSum): + return ufl.FormSum(*[(op, w) for op, w in zip(operands, expr.weights())]) + return expr - This function is now deprecated. + @staticmethod + def base_form_operands(expr): + if isinstance(expr, (ufl.FormSum, ufl.Adjoint, ufl.Action)): + return expr.ufl_operands + if isinstance(expr, ufl.Form): + # Use reversed to treat base form operators + # in the order in which they have been made. + return list(reversed(expr.base_form_operators())) + if isinstance(expr, ufl.core.base_form_operator.BaseFormOperator): + # Conserve order + children = dict.fromkeys(e for e in (expr.argument_slots() + expr.ufl_operands) + if isinstance(e, ufl.form.BaseForm)) + return list(children) + return [] - .. warning:: + @staticmethod + def restructure_base_form_postorder(expression, visited=None): + visited = visited or {} - Really do not use this function unless you know what you're doing. - """ - import warnings - with warnings.catch_warnings(): - warnings.simplefilter("once", DeprecationWarning) - warnings.warn("create_assembly_callable is now deprecated. Please use assemble or FormAssembler instead.", - DeprecationWarning) + def visitor(expr, *operands): + # Need to reconstruct the expression with its visited operands! + expr = BaseFormAssembler.reconstruct_node_from_operands(expr, operands) + # Perform the DAG restructuring when needed + return BaseFormAssembler.restructure_base_form(expr, visited) - if tensor is None: - raise ValueError("Have to provide tensor to write to") + return BaseFormAssembler.base_form_postorder_traversal(expression, visitor, visited) - rank = len(expr.arguments()) - if rank == 0: - return ZeroFormAssembler(expr, tensor, form_compiler_parameters).assemble - elif rank == 1 or (rank == 2 and diagonal): - return OneFormAssembler(expr, tensor, bcs, diagonal=diagonal, - form_compiler_parameters=form_compiler_parameters).assemble - elif rank == 2: - return TwoFormAssembler(expr, tensor, bcs, form_compiler_parameters).assemble - else: - raise AssertionError + @staticmethod + def restructure_base_form_preorder(expression, visited=None): + visited = visited or {} + def visitor(expr): + # Perform the DAG restructuring when needed + return BaseFormAssembler.restructure_base_form(expr, visited) -def _assemble_form(form, tensor=None, bcs=None, *, - diagonal=False, - mat_type=None, - sub_mat_type=None, - appctx=None, - options_prefix=None, - form_compiler_parameters=None, - zero_bc_nodes=False, - weight=1.0): - """Assemble a form. + expression = BaseFormAssembler.base_form_preorder_traversal(expression, visitor, visited) + # Need to reconstruct the expression at the end when all its operands have been visited! + operands = [visited.get(args, args) for args in BaseFormAssembler.base_form_operands(expression)] + return BaseFormAssembler.reconstruct_node_from_operands(expression, operands) - See :func:`assemble` for a description of the arguments to this function. - """ - bcs = solving._extract_bcs(bcs) + @staticmethod + def restructure_base_form(expr, visited=None): + r"""Perform a preorder traversal to simplify and optimize the DAG. + Example: Let's consider F(u, N(u; v*); v) with N(u; v*) a base form operator. - _check_inputs(form, tensor, bcs, diagonal) + We have: dFdu = \frac{\partial F}{\partial u} + Action(dFdN, dNdu) + Now taking the action on a rank-1 object w (e.g. Coefficient/Cofunction) results in: - if tensor is not None: - _zero_tensor(tensor, form, diagonal) - else: - tensor = _make_tensor(form, bcs, diagonal=diagonal, mat_type=mat_type, - sub_mat_type=sub_mat_type, appctx=appctx, - form_compiler_parameters=form_compiler_parameters, - options_prefix=options_prefix) - - # It is expensive to construct new assemblers because extracting the data - # from the form is slow. Since all of the data structures in the assembler - # are persistent apart from the output tensor, we stash the assembler on the - # form and swap out the tensor if needed. - # The cache key only needs to contain the boundary conditions, diagonal and - # form compiler parameters since all other assemble kwargs are only used for - # creating the tensor which is handled above and has no bearing on the assembler. - # Note: This technically creates a memory leak since bcs are 'heavy' and so - # repeated assembly of the same form but with different boundary conditions - # will lead to old bcs getting stored along with old tensors. - - # FIXME This only works for 1-forms at the moment - is_cacheable = len(form.arguments()) == 1 - if is_cacheable: - try: - key = tuple(bcs), diagonal, tuplify(form_compiler_parameters), zero_bc_nodes - assembler = form._cache[_FORM_CACHE_KEY][key] - assembler.replace_tensor(tensor) - return assembler.assemble() - except KeyError: - pass + (1) Action(Action(dFdN, dNdu), w) - rank = len(form.arguments()) - if rank == 0: - assembler = ZeroFormAssembler(form, tensor, form_compiler_parameters) - elif rank == 1 or (rank == 2 and diagonal): - assembler = OneFormAssembler(form, tensor, bcs, diagonal=diagonal, - form_compiler_parameters=form_compiler_parameters, - needs_zeroing=False, zero_bc_nodes=zero_bc_nodes) - elif rank == 2: - assembler = TwoFormAssembler(form, tensor, bcs, form_compiler_parameters, - needs_zeroing=False, weight=weight) - else: - raise AssertionError + Action Action + / \ / \ + Action w -----> dFdN Action + / \ / \ + dFdN dNdu dNdu w - if is_cacheable: - if _FORM_CACHE_KEY not in form._cache: - form._cache[_FORM_CACHE_KEY] = {} - form._cache[_FORM_CACHE_KEY][key] = assembler + This situations does not only arise for BaseFormOperator but also when we have a 2-form instead of dNdu! - return assembler.assemble() + (2) Action(dNdu, w) + Action + / \ + / w -----> dNdu(u; w, v*) + / + dNdu(u; uhat, v*) -def _assemble_expr(expr): - """Assemble a pointwise expression. + (3) Action(F, N) - :arg expr: The :class:`ufl.core.expr.Expr` to be evaluated. - :returns: A :class:`firedrake.Function` containing the result of this evaluation. - """ - from ufl.algorithms.analysis import extract_base_form_operators - from ufl.checks import is_scalar_constant_expression - - # Get BaseFormOperators (e.g. `Interpolate` or `ExternalOperator`) - base_form_operators = extract_base_form_operators(expr) - - # -- Linear combination involving 2-form BaseFormOperators -- # - # Example: a * dNdu1(u1, u2; v1, v*) + b * dNdu2(u1, u2; v2, v*) - # with u1, u2 Functions, v1, v2, v* BaseArguments, dNdu1, dNdu2 BaseFormOperators, and a, b scalars. - if len(base_form_operators) and any(len(e.arguments()) > 1 for e in base_form_operators): - if isinstance(expr, ufl.algebra.Sum): - a, b = [assemble(e) for e in expr.ufl_operands] - # Only Expr resulting in a Matrix if assembled are BaseFormOperator - if not all(isinstance(op, matrix.AssembledMatrix) for op in (a, b)): - raise TypeError('Mismatching Sum shapes') - return assemble_base_form(ufl.FormSum((a, 1), (b, 1))) - elif isinstance(expr, ufl.algebra.Product): - a, b = expr.ufl_operands - scalar = [e for e in expr.ufl_operands if is_scalar_constant_expression(e)] - if scalar: - base_form = a if a is scalar else b - assembled_mat = assemble(base_form) - return assemble_base_form(ufl.FormSum((assembled_mat, scalar[0]))) - a, b = [assemble(e) for e in (a, b)] - return assemble_base_form(ufl.action(a, b)) - # -- Linear combination of Functions and 1-form BaseFormOperators -- # - # Example: a * u1 + b * u2 + c * N(u1; v*) + d * N(u2; v*) - # with u1, u2 Functions, N a BaseFormOperator, and a, b, c, d scalars or 0-form BaseFormOperators. - else: - base_form_operators = extract_base_form_operators(expr) - assembled_bfops = [firedrake.assemble(e) for e in base_form_operators] - # Substitute base form operators with their output before examining the expression - # which avoids conflict when determining function space, for example: - # extract_coefficients(Interpolate(u, V2)) with u \in V1 will result in an output function space V1 - # instead of V2. - if base_form_operators: - expr = ufl.replace(expr, dict(zip(base_form_operators, assembled_bfops))) - try: - coefficients = ufl.algorithms.extract_coefficients(expr) - V, = set(c.function_space() for c in coefficients) - {None} - except ValueError: - raise ValueError("Cannot deduce correct target space from pointwise expression") - return firedrake.Function(V).assign(expr) + Action F + / \ -----> F(..., N)[v] = | + F[v] N N + (4) Adjoint(dNdu) -def _check_inputs(form, tensor, bcs, diagonal): - # Ensure mesh is 'initialised' as we could have got here without building a - # function space (e.g. if integrating a constant). - for mesh in form.ufl_domains(): - mesh.init() + Adjoint + | -----> dNdu(u; v*, uhat) + dNdu(u; uhat, v*) - if diagonal and any(isinstance(bc, EquationBCSplit) for bc in bcs): - raise NotImplementedError("Diagonal assembly and EquationBC not supported") + (5) N(u; w) (scalar valued) - rank = len(form.arguments()) - if rank == 0: - assert tensor is None - assert not bcs - elif rank == 1: - test, = form.arguments() - - if tensor is not None and test.function_space() != tensor.function_space(): - raise ValueError("Form's argument does not match provided result tensor") - elif rank == 2 and diagonal: - test, trial = form.arguments() - if test.function_space() != trial.function_space(): - raise ValueError("Can only assemble the diagonal of 2-form if the function spaces match") - elif rank == 2: - if tensor is not None and tensor.a.arguments() != form.arguments(): - raise ValueError("Form's arguments do not match provided result tensor") - else: - raise AssertionError + Action + N(u; w) ----> / \ = Action(N, w) + N(u; v*) w - if any(c.dat.dtype != ScalarType for c in form.coefficients()): - raise ValueError("Cannot assemble a form containing coefficients where the " - "dtype is not the PETSc scalar type.") + So from Action(Action(dFdN, dNdu(u; v*)), w) we get: + Action Action Action + / \ (1) / \ (2) / \ (4) dFdN + Action w ----> dFdN Action ----> dFdN dNdu(u; w, v*) ----> dFdN(..., dNdu(u; w, v*)) = | + / \ / \ dNdu(u; w, v*) + dFdN dNdu dNdu w -def _zero_tensor(tensor, form, diagonal): - rank = len(form.arguments()) - assert rank != 0 - if rank == 1 or (rank == 2 and diagonal): - tensor.dat.zero() - elif rank == 2: - if not isinstance(tensor, matrix.ImplicitMatrix): - tensor.M.zero() - else: - raise AssertionError + (6) ufl.FormSum(dN1du(u; w, v*), dN2du(u; w, v*)) -> ufl.Sum(dN1du(u; w, v*), dN2du(u; w, v*)) + Let's consider `Action(dN1du, w) + Action(dN2du, w)`, we have: -def _make_tensor(form, bcs, *, diagonal, mat_type, sub_mat_type, appctx, - form_compiler_parameters, options_prefix): - rank = len(form.arguments()) - if rank == 0: - # Getting the comm attribute of a form isn't straightforward - # form.ufl_domains()[0]._comm seems the most robust method - # revisit in a refactor - return op2.Global( - 1, - [0.0], - dtype=utils.ScalarType, - comm=form.ufl_domains()[0]._comm - ) - elif rank == 1: - test, = form.arguments() - return firedrake.Cofunction(test.function_space().dual()) - elif rank == 2 and diagonal: - test, _ = form.arguments() - return firedrake.Cofunction(test.function_space().dual()) - elif rank == 2: - mat_type, sub_mat_type = _get_mat_type(mat_type, sub_mat_type, form.arguments()) - return allocate_matrix(form, bcs, mat_type=mat_type, sub_mat_type=sub_mat_type, - appctx=appctx, form_compiler_parameters=form_compiler_parameters, - options_prefix=options_prefix) - else: - raise AssertionError + FormSum (2) FormSum (6) Sum + / \ ----> / \ ----> / \ + / \ / \ / \ + Action(dN1du, w) Action(dN2du, w) dN1du(u; w, v*) dN2du(u; w, v*) dN1du(u; w, v*) dN2du(u; w, v*) + + This case arises as a consequence of (2) which turns sum of `Action`s (i.e. ufl.FormSum since Action is a BaseForm) + into sum of `BaseFormOperator`s (i.e. ufl.Sum since BaseFormOperator is an Expr as well). + + (7) Action(w*, dNdu) + + Action + / \ + w* \ -----> dNdu(u; v0, w*) + \ + dNdu(u; v1, v0*) + + It uses a recursive approach to reconstruct the DAG as we traverse it, enabling to take into account + various dag rotations/manipulations in expr. + """ + if isinstance(expr, ufl.Action): + left, right = expr.ufl_operands + is_rank_1 = lambda x: isinstance(x, (firedrake.Cofunction, firedrake.Function, firedrake.Argument)) or len(x.arguments()) == 1 + is_rank_2 = lambda x: len(x.arguments()) == 2 + + # -- Case (1) -- # + # If left is Action and has a rank 2, then it is an action of a 2-form on a 2-form + if isinstance(left, ufl.Action) and is_rank_2(left): + return ufl.action(left.left(), ufl.action(left.right(), right)) + # -- Case (2) (except if left has only 1 argument, i.e. we have done case (5)) -- # + if isinstance(left, ufl.core.base_form_operator.BaseFormOperator) and is_rank_1(right) and len(left.arguments()) != 1: + # Retrieve the highest numbered argument + arg = max(left.arguments(), key=lambda v: v.number()) + return ufl.replace(left, {arg: right}) + # -- Case (3) -- # + if isinstance(left, ufl.Form) and is_rank_1(right): + # 1) Replace the highest-numbered argument of left by right when needed + # -> e.g. if right is a BaseFormOperator with 1 argument. + # Or + # 2) Let expr as it is by returning `ufl.Action(left, right)`. + return ufl.action(left, right) + # -- Case (7) -- # + if is_rank_1(left) and isinstance(right, ufl.core.base_form_operator.BaseFormOperator) and len(right.arguments()) != 1: + # Action(w*, dNdu(u; v1, v*)) -> dNdu(u; v0, w*) + # Get lowest numbered argument + arg = min(right.arguments(), key=lambda v: v.number()) + # Need to replace lowest numbered argument of right by left + replace_map = {arg: left} + # Decrease number for all the other arguments since the lowest numbered argument will be replaced. + other_args = [a for a in right.arguments() if a is not arg] + new_args = [firedrake.Argument(a.function_space(), number=a.number()-1, part=a.part()) for a in other_args] + replace_map.update(dict(zip(other_args, new_args))) + # Replace arguments + return ufl.replace(right, replace_map) + + # -- Case (4) -- # + if isinstance(expr, ufl.Adjoint) and isinstance(expr.form(), ufl.core.base_form_operator.BaseFormOperator): + B = expr.form() + u, v = B.arguments() + # Let V1 and V2 be primal spaces, B: V1 -> V2 and B*: V2* -> V1*: + # Adjoint(B(Argument(V1, 1), Argument(V2.dual(), 0))) = B(Argument(V1, 0), Argument(V2.dual(), 1)) + reordered_arguments = (firedrake.Argument(u.function_space(), number=v.number(), part=v.part()), + firedrake.Argument(v.function_space(), number=u.number(), part=u.part())) + # Replace arguments in argument slots + return ufl.replace(B, dict(zip((u, v), reordered_arguments))) + + # -- Case (5) -- # + if isinstance(expr, ufl.core.base_form_operator.BaseFormOperator) and not expr.arguments(): + # We are assembling a BaseFormOperator of rank 0 (no arguments). + # B(f, u*) be a BaseFormOperator with u* a Cofunction and f a Coefficient, then: + # B(f, u*) <=> Action(B(f, v*), f) where v* is a Coargument + ustar, *_ = expr.argument_slots() + vstar = firedrake.Argument(ustar.function_space(), 0) + expr = ufl.replace(expr, {ustar: vstar}) + return ufl.action(expr, ustar) + + # -- Case (6) -- # + if isinstance(expr, ufl.FormSum) and all(isinstance(c, ufl.core.base_form_operator.BaseFormOperator) for c in expr.components()): + # Return ufl.Sum + return sum([c for c in expr.components()]) + return expr + + @staticmethod + def preprocess_base_form(expr, mat_type=None, form_compiler_parameters=None): + """Preprocess ufl.BaseForm objects""" + original_expr = expr + if mat_type != "matfree": + # Don't expand derivatives if `mat_type` is 'matfree' + # For "matfree", Form evaluation is delayed + expr = BaseFormAssembler.expand_derivatives_form(expr, form_compiler_parameters) + if not isinstance(expr, (ufl.form.Form, slate.TensorBase)): + # => No restructuring needed for Form and slate.TensorBase + expr = BaseFormAssembler.restructure_base_form_preorder(expr) + expr = BaseFormAssembler.restructure_base_form_postorder(expr) + # Preprocessing the form makes a new object -> current form caching mechanism + # will populate `expr`'s cache which is now different than `original_expr`'s cache so we need + # to transmit the cache. All of this only holds when both are `ufl.Form` objects. + if isinstance(original_expr, ufl.form.Form) and isinstance(expr, ufl.form.Form): + expr._cache = original_expr._cache + return expr + @staticmethod + def expand_derivatives_form(form, fc_params): + """Expand derivatives of ufl.BaseForm objects + :arg form: a :class:`~ufl.classes.BaseForm` + :arg fc_params:: Dictionary of parameters to pass to the form compiler. + + :returns: The resulting preprocessed :class:`~ufl.classes.BaseForm`. + This function preprocess the form, mainly by expanding the derivatives, in order to determine + if we are dealing with a :class:`~ufl.classes.Form` or another :class:`~ufl.classes.BaseForm` object. + This function is called in :func:`base_form_assembly_visitor`. Depending on the type of the resulting tensor, + we may call :func:`assemble_form` or traverse the sub-DAG via :func:`assemble_base_form`. + """ + if isinstance(form, ufl.form.Form): + from firedrake.parameters import parameters as default_parameters + from tsfc.parameters import is_complex -class FormAssembler(abc.ABC): - """Abstract base class for assembling forms. + if fc_params is None: + fc_params = default_parameters["form_compiler"].copy() + else: + # Override defaults with user-specified values + _ = fc_params + fc_params = default_parameters["form_compiler"].copy() + fc_params.update(_) + + complex_mode = fc_params and is_complex(fc_params.get("scalar_type")) + + return ufl.algorithms.preprocess_form(form, complex_mode) + # We also need to expand derivatives for `ufl.BaseForm` objects that are not `ufl.Form` + # Example: `Action(A, derivative(B, f))`, where `A` is a `ufl.BaseForm` and `B` can + # be `ufl.BaseForm`, or even an appropriate `ufl.Expr`, since assembly of expressions + # containing derivatives is not supported anymore but might be needed if the expression + # in question is within a `ufl.BaseForm` object. + return ufl.algorithms.ad.expand_derivatives(form) + + +class FormAssembler(AbstractFormAssembler): + """Form assembler. + + Parameters + ---------- + form : ufl.Form or slate.TensorBase + Variational form to be assembled. + bcs : Sequence + Iterable of boundary conditions to apply. + form_compiler_parameters : dict + Optional parameters to pass to the TSFC and/or Slate compilers. - :param form: The variational form to be assembled. - :param tensor: The output tensor to store the result. - :param bcs: Iterable of boundary conditions to apply. - :param form_compiler_parameters: Optional parameters to pass to the - TSFC and/or Slate compilers. - :param needs_zeroing: Should ``tensor`` be zeroed before assembling? """ - def __init__(self, form, tensor, bcs=(), form_compiler_parameters=None, needs_zeroing=True, weight=1.0): - assert tensor is not None + def __new__(cls, *args, **kwargs): + form = args[0] + if not isinstance(form, (ufl.Form, slate.TensorBase)): + raise TypeError(f"The first positional argument must be of ufl.Form or slate.TensorBase: got {type(form)} ({form})") + # It is expensive to construct new assemblers because extracting the data + # from the form is slow. Since all of the data structures in the assembler + # are persistent apart from the output tensor, we stash the assembler on the + # form and swap out the tensor if needed. + # The cache key only needs to contain the boundary conditions, diagonal and + # form compiler parameters since all other assemble kwargs are only used for + # creating the tensor which is handled above and has no bearing on the assembler. + # Note: This technically creates a memory leak since bcs are 'heavy' and so + # repeated assembly of the same form but with different boundary conditions + # will lead to old bcs getting stored along with old tensors. + # FIXME This only works for 1-forms at the moment + key = cls._cache_key(*args, **kwargs) + if key is not None: + try: + return form._cache[_FORM_CACHE_KEY][key] + except KeyError: + pass + self = super().__new__(cls) + self._initialised = False + self.__init__(*args, **kwargs) + if _FORM_CACHE_KEY not in form._cache: + form._cache[_FORM_CACHE_KEY] = {} + form._cache[_FORM_CACHE_KEY][key] = self + return self - bcs = solving._extract_bcs(bcs) + @classmethod + @abc.abstractmethod + def _cache_key(cls, *args, **kwargs): + """Return cache key.""" - self._form = form - self._tensor = tensor - self._bcs = bcs - self._form_compiler_params = form_compiler_parameters or {} + @staticmethod + def _skip_if_initialised(init): + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + if not self._initialised: + init(self, *args, **kwargs) + self._initialised = True + return wrapper + + def __init__(self, form, bcs=None, form_compiler_parameters=None): + super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters) + # Ensure mesh is 'initialised' as we could have got here without building a + # function space (e.g. if integrating a constant). + for mesh in form.ufl_domains(): + mesh.init() + if any(c.dat.dtype != ScalarType for c in form.coefficients()): + raise ValueError("Cannot assemble a form containing coefficients where the " + "dtype is not the PETSc scalar type.") + + +class ParloopFormAssembler(FormAssembler): + """FormAssembler that uses Parloops. + + Parameters + ---------- + form : ufl.Form or slate.TensorBase + Variational form to be assembled. + bcs : Sequence + Iterable of boundary conditions to apply. + form_compiler_parameters : dict + Optional parameters to pass to the TSFC and/or Slate compilers. + needs_zeroing : bool + Should ``tensor`` be zeroed before assembling? + + """ + def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True): + super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters) self._needs_zeroing = needs_zeroing - self.weight = weight - @property - @abc.abstractmethod - def result(self): - """The result of the assembly operation.""" + def assemble(self, tensor=None): + """Assemble the form. - @property - @abc.abstractmethod - def diagonal(self): - """Are we assembling the diagonal of a 2-form?""" + Parameters + ---------- + tensor : firedrake.cofunction.Cofunction or matrix.MatrixBase + Output tensor to contain the result of assembly; if `None`, a tensor of appropriate type is created. - def assemble(self): - """Perform the assembly. + Returns + ------- + float or firedrake.cofunction.Cofunction or firedrake.function.Function or matrix.MatrixBase + Result of assembly: `float` for 0-forms, `firedrake.cofunction.Cofunction` or `firedrake.function.Function` for 1-forms, and `matrix.MatrixBase` for 2-forms. - :returns: The assembled object. """ + self._check_tensor(tensor) + if tensor is None: + tensor = self.allocate() + needs_zeroing = False + else: + needs_zeroing = self._needs_zeroing if annotate_tape(): raise NotImplementedError( "Taping with explicit FormAssembler objects is not supported yet. " "Use assemble instead." ) - - if self._needs_zeroing: - self._as_pyop2_type(self._tensor).zero() - - self.execute_parloops() - + if needs_zeroing: + type(self)._as_pyop2_type(tensor).zero() + self.execute_parloops(tensor) for bc in self._bcs: - if isinstance(bc, EquationBC): # can this be lifted? - bc = bc.extract_form("F") - self._apply_bc(bc) + self._apply_bc(tensor, bc) + return self.result(tensor) - return self.result + @abc.abstractmethod + def _apply_bc(self, tensor, bc): + """Apply boundary condition.""" - def replace_tensor(self, tensor): - if tensor is self._tensor: - return + @abc.abstractmethod + def _check_tensor(self, tensor): + """Check input tensor.""" - # TODO We should have some proper checks here - for (lknl, _), parloop in zip(self.local_kernels, self.parloops): - data = _FormHandler.index_tensor(tensor, self._form, lknl.indices, self.diagonal) - parloop.arguments[0].data = data - self._tensor = tensor + @staticmethod + def _as_pyop2_type(tensor): + """Return tensor as pyop2 type.""" + raise NotImplementedError - def execute_parloops(self): - for parloop in self.parloops: + def execute_parloops(self, tensor): + for parloop in self.parloops(tensor): parloop() + def parloops(self, tensor): + if hasattr(self, "_parloops"): + for (lknl, _), parloop in zip(self.local_kernels, self._parloops): + data = _FormHandler.index_tensor(tensor, self._form, lknl.indices, self.diagonal) + parloop.arguments[0].data = data + else: + # Make parloops for one concrete output tensor and cache them. + # TODO: Make parloops only with some symbolic information of the output tensor. + self._parloops = tuple(parloop_builder.build(tensor) for parloop_builder in self.parloop_builders) + return self._parloops + + @cached_property + def parloop_builders(self): + out = [] + for local_kernel, subdomain_id in self.local_kernels: + out.append( + ParloopBuilder( + self._form, + self._bcs, + local_kernel, + subdomain_id, + self.all_integer_subdomain_ids, + diagonal=self.diagonal, + ) + ) + return tuple(out) + @cached_property def local_kernels(self): """Return local kernels and their subdomain IDs. @@ -1062,6 +1036,7 @@ def local_kernels(self): tuple Collection of ``(local_kernel, subdomain_id)`` 2-tuples, one for each possible combination. + """ try: topology, = set(d.topology for d in self._form.ufl_domains()) @@ -1089,219 +1064,350 @@ def local_kernels(self): (k, subdomain_id) for k in kernels for subdomain_id in k.kinfo.subdomain_id ) + @property + @abc.abstractmethod + def diagonal(self): + """Are we assembling the diagonal of a 2-form?""" + @cached_property def all_integer_subdomain_ids(self): return tsfc_interface.gather_integer_subdomain_ids( {k for k, _ in self.local_kernels} ) - @cached_property - def global_kernels(self): - return tuple( - _make_global_kernel( - self._form, tsfc_knl, subdomain_id, self.all_integer_subdomain_ids, - diagonal=self.diagonal, unroll=self.needs_unrolling(tsfc_knl, self._bcs) - ) - for tsfc_knl, subdomain_id in self.local_kernels - ) + @abc.abstractmethod + def result(self, tensor): + """The result of the assembly operation.""" - @cached_property - def parloops(self): - loops = [] - for (local_kernel, subdomain_id), global_kernel in zip( - self.local_kernels, self.global_kernels - ): - loops.append( - ParloopBuilder( - self._form, - local_kernel, - global_kernel, - self._tensor, - subdomain_id, - self.all_integer_subdomain_ids, - diagonal=self.diagonal, - lgmaps=self.collect_lgmaps(local_kernel, self._bcs) - ).build() - ) - return tuple(loops) - def needs_unrolling(self, local_knl, bcs): - """Do we need to address matrix elements directly rather than in - a blocked fashion? +class ZeroFormAssembler(ParloopFormAssembler): + """Class for assembling a 0-form. - This is slower but required for the application of some boundary conditions - to 2-forms. + Parameters + ---------- + form : ufl.Form or slate.TensorBasehe + 0-form. - :param local_knl: A :class:`tsfc_interface.SplitKernel`. - :param bcs: Iterable of boundary conditions. - """ - return False + Notes + ----- + See `FormAssembler` and `assemble` for descriptions of the other parameters. - def collect_lgmaps(self, local_knl, bcs): - """Return any local-to-global maps that need to be swapped out. + """ - This is only needed when applying boundary conditions to 2-forms. + diagonal = False + """Diagonal assembly not possible for zero forms.""" - :param local_knl: A :class:`tsfc_interface.SplitKernel`. - :param bcs: Iterable of boundary conditions. - """ - return None + @classmethod + def _cache_key(cls, *args, **kwargs): + return - @staticmethod - def _as_pyop2_type(tensor): - if isinstance(tensor, op2.Global): - return tensor - elif isinstance(tensor, firedrake.Cofunction): - return tensor.dat - elif isinstance(tensor, matrix.Matrix): - return tensor.M - else: - raise AssertionError + @FormAssembler._skip_if_initialised + def __init__(self, form, form_compiler_parameters=None): + super().__init__(form, bcs=None, form_compiler_parameters=form_compiler_parameters) + def allocate(self): + # Getting the comm attribute of a form isn't straightforward + # form.ufl_domains()[0]._comm seems the most robust method + # revisit in a refactor + return op2.Global( + 1, + [0.0], + dtype=utils.ScalarType, + comm=self._form.ufl_domains()[0]._comm + ) -class ZeroFormAssembler(FormAssembler): - """Class for assembling a 0-form.""" + def _apply_bc(self, tensor, bc): + pass - diagonal = False - """Diagonal assembly not possible for zero forms.""" + def _check_tensor(self, tensor): + assert tensor is None - def __init__(self, form, tensor, form_compiler_parameters=None): - super().__init__(form, tensor, (), form_compiler_parameters) + @staticmethod + def _as_pyop2_type(tensor): + return tensor - @property - def result(self): - return self._tensor.data[0] + def result(self, tensor): + return tensor.data[0] -class OneFormAssembler(FormAssembler): +class OneFormAssembler(ParloopFormAssembler): """Class for assembling a 1-form. - :param diagonal: Are we actually assembling the diagonal of a 2-form? - :param zero_bc_nodes: If ``True``, set the boundary condition nodes in the - output tensor to zero rather than to the values prescribed by the - boundary condition. + Parameters + ---------- + form : ufl.Form or slate.TensorBasehe + 1-form. + + Notes + ----- + See `FormAssembler` and `assemble` for descriptions of the other parameters. - For all other arguments see :class:`FormAssembler` for more information. """ - def __init__(self, form, tensor, bcs=(), diagonal=False, zero_bc_nodes=False, - form_compiler_parameters=None, needs_zeroing=True): - super().__init__(form, tensor, bcs, form_compiler_parameters, needs_zeroing) + @classmethod + def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, + zero_bc_nodes=False, diagonal=False): + bcs = solving._extract_bcs(bcs) + return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal + + @FormAssembler._skip_if_initialised + def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, + zero_bc_nodes=False, diagonal=False): + super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) self._diagonal = diagonal self._zero_bc_nodes = zero_bc_nodes + if self._diagonal and any(isinstance(bc, EquationBCSplit) for bc in self._bcs): + raise NotImplementedError("Diagonal assembly and EquationBC not supported") + rank = len(self._form.arguments()) + if rank == 2 and self._diagonal: + test, trial = self._form.arguments() + if test.function_space() != trial.function_space(): + raise ValueError("Can only assemble the diagonal of 2-form if the function spaces match") + + def allocate(self): + rank = len(self._form.arguments()) + if rank == 1: + test, = self._form.arguments() + return firedrake.Cofunction(test.function_space().dual()) + elif rank == 2 and self._diagonal: + test, _ = self._form.arguments() + return firedrake.Cofunction(test.function_space().dual()) + else: + raise RuntimeError(f"Not expected: found rank = {rank} and diagonal = {self._diagonal}") - @property - def diagonal(self): - return self._diagonal - - @property - def result(self): - return self._tensor - - def execute_parloops(self): - # We are repeatedly incrementing into the same Dat so intermediate halo exchanges - # can be skipped. - with self._tensor.dat.frozen_halo(op2.INC): - for parloop in self.parloops: - parloop() - - def _apply_bc(self, bc): + def _apply_bc(self, tensor, bc): # TODO Maybe this could be a singledispatchmethod? if isinstance(bc, DirichletBC): - self._apply_dirichlet_bc(bc) + self._apply_dirichlet_bc(tensor, bc) elif isinstance(bc, EquationBCSplit): - bc.zero(self._tensor) - - type(self)(bc.f, self._tensor, bc.bcs, self._diagonal, self._zero_bc_nodes, - self._form_compiler_params, needs_zeroing=False).assemble() + bc.zero(tensor) + type(self)(bc.f, bcs=bc.bcs, form_compiler_parameters=self._form_compiler_params, needs_zeroing=False, + zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal).assemble(tensor=tensor) else: raise AssertionError - def _apply_dirichlet_bc(self, bc): + def _apply_dirichlet_bc(self, tensor, bc): if not self._zero_bc_nodes: - tensor_func = self._tensor.riesz_representation(riesz_map="l2") + tensor_func = tensor.riesz_representation(riesz_map="l2") if self._diagonal: bc.set(tensor_func, 1) else: bc.apply(tensor_func) - self._tensor.assign(tensor_func.riesz_representation(riesz_map="l2")) + tensor.assign(tensor_func.riesz_representation(riesz_map="l2")) else: - bc.zero(self._tensor) + bc.zero(tensor) + def _check_tensor(self, tensor): + rank = len(self._form.arguments()) + if rank == 1: + test, = self._form.arguments() + if tensor is not None and test.function_space() != tensor.function_space(): + raise ValueError("Form's argument does not match provided result tensor") -def TwoFormAssembler(form, tensor, *args, **kwargs): - if isinstance(tensor, matrix.ImplicitMatrix): - return MatrixFreeAssembler(tensor) + @staticmethod + def _as_pyop2_type(tensor): + return tensor.dat + + def execute_parloops(self, tensor): + # We are repeatedly incrementing into the same Dat so intermediate halo exchanges + # can be skipped. + with tensor.dat.frozen_halo(op2.INC): + for parloop in self.parloops(tensor): + parloop() + + @property + def diagonal(self): + return self._diagonal + + def result(self, tensor): + return tensor + + +def TwoFormAssembler(form, *args, **kwargs): + assert isinstance(form, (ufl.form.Form, slate.TensorBase)) + mat_type = kwargs.pop('mat_type', None) + sub_mat_type = kwargs.pop('sub_mat_type', None) + mat_type, sub_mat_type = _get_mat_type(mat_type, sub_mat_type, form.arguments()) + if mat_type == "matfree": + kwargs.pop('needs_zeroing', None) + kwargs.pop('weight', None) + kwargs.pop('allocation_integral_types', None) + return MatrixFreeAssembler(form, *args, **kwargs) else: - return ExplicitMatrixAssembler(form, tensor, *args, **kwargs) + return ExplicitMatrixAssembler(form, *args, mat_type=mat_type, sub_mat_type=sub_mat_type, **kwargs) -class ExplicitMatrixAssembler(FormAssembler): - """Class for assembling a 2-form.""" +def _get_mat_type(mat_type, sub_mat_type, arguments): + """Validate the matrix types provided by the user and set any that are + undefined to default values. - diagonal = False - """Diagonal assembly not possible for two forms.""" + Parameters + ---------- + mat_type : str + PETSc matrix type for the assembled matrix. + sub_mat_type : str + PETSc matrix type for blocks if ``mat_type`` is ``"nest"``. + arguments : Sequence + The test and trial functions of the expression being assembled. - @property - def test_function_space(self): - test, _ = self._form.arguments() - return test.function_space() + Returns + ------- + tuple + Tuple of validated/default ``mat_type`` and ``sub_mat_type``. - @property - def trial_function_space(self): - _, trial = self._form.arguments() - return trial.function_space() + """ + if mat_type is None: + mat_type = parameters.parameters["default_matrix_type"] + if any(V.ufl_element().family() == "Real" + for arg in arguments + for V in arg.function_space()): + mat_type = "nest" + if mat_type not in {"matfree", "aij", "baij", "nest", "dense"}: + raise ValueError(f"Unrecognised matrix type, '{mat_type}'") + if sub_mat_type is None: + sub_mat_type = parameters.parameters["default_sub_matrix_type"] + if sub_mat_type not in {"aij", "baij"}: + raise ValueError(f"Invalid submatrix type, '{sub_mat_type}' (not 'aij' or 'baij')") + return mat_type, sub_mat_type - def get_indicess(self, knl): - if all(i is None for i in knl.indices): - return numpy.ndindex(self._tensor.block_shape) - else: - assert all(i is not None for i in knl.indices) - return knl.indices, - @property - def result(self): - self._tensor.M.assemble() - return self._tensor - - def needs_unrolling(self, knl, bcs): - for i, j in self.get_indicess(knl): - for bc in itertools.chain(*self._filter_bcs(bcs, i, j)): - if bc.function_space().component is not None: - return True - return False +class ExplicitMatrixAssembler(ParloopFormAssembler): + """Class for assembling a matrix. - def collect_lgmaps(self, knl, bcs): - if not bcs: - return None + Parameters + ---------- + form : ufl.Form or slate.TensorBasehe + 2-form. - lgmaps = [] - for i, j in self.get_indicess(knl): - row_bcs, col_bcs = self._filter_bcs(bcs, i, j) - rlgmap, clgmap = self._tensor.M[i, j].local_to_global_maps - rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap) - clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap) - lgmaps.append((rlgmap, clgmap)) - return tuple(lgmaps) + Notes + ----- + See `FormAssembler` and `assemble` for descriptions of the other parameters. - def _filter_bcs(self, bcs, row, col): - if len(self.test_function_space) > 1: - bcrow = tuple(bc for bc in bcs - if bc.function_space_index() == row) - else: - bcrow = bcs + """ - if len(self.trial_function_space) > 1: - bccol = tuple(bc for bc in bcs - if bc.function_space_index() == col - and isinstance(bc, DirichletBC)) + diagonal = False + """Diagonal assembly not possible for two forms.""" + + @classmethod + def _cache_key(cls, *args, **kwargs): + return + + @FormAssembler._skip_if_initialised + def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, + mat_type=None, sub_mat_type=None, options_prefix=None, appctx=None, weight=1.0, + allocation_integral_types=None): + super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) + self._mat_type = mat_type + self._sub_mat_type = sub_mat_type + self._options_prefix = options_prefix + self._appctx = appctx + self.weight = weight + self._allocation_integral_types = allocation_integral_types + + def allocate(self): + test, trial = self._form.arguments() + sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, + self._mat_type, + self._sub_mat_type, + self._make_maps_and_regions()) + return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, + options_prefix=self._options_prefix) + + @staticmethod + def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions): + assert mat_type != "matfree" + nest = mat_type == "nest" + if nest: + baij = sub_mat_type == "baij" else: - bccol = tuple(bc for bc in bcs if isinstance(bc, DirichletBC)) - return bcrow, bccol + baij = mat_type == "baij" + if any(len(a.function_space()) > 1 for a in [test, trial]) and mat_type == "baij": + raise ValueError("BAIJ matrix type makes no sense for mixed spaces, use 'aij'") + try: + sparsity = op2.Sparsity((test.function_space().dof_dset, + trial.function_space().dof_dset), + maps_and_regions, + nest=nest, + block_sparse=baij) + except SparsityFormatError: + raise ValueError("Monolithic matrix assembly not supported for systems " + "with R-space blocks") + return sparsity + + def _make_maps_and_regions(self): + test, trial = self._form.arguments() + if self._allocation_integral_types is not None: + return ExplicitMatrixAssembler._make_maps_and_regions_default(test, trial, self._allocation_integral_types) + elif any(local_kernel.indices == (None, None) for local_kernel in self._all_local_kernels): + # Handle special cases: slate or split=False + assert all(local_kernel.indices == (None, None) for local_kernel in self._all_local_kernels) + allocation_integral_types = set(local_kernel.kinfo.integral_type + for local_kernel in self._all_local_kernels) + return ExplicitMatrixAssembler._make_maps_and_regions_default(test, trial, allocation_integral_types) + else: + maps_and_regions = defaultdict(lambda: defaultdict(set)) + for local_kernel in self._all_local_kernels: + i, j = local_kernel.indices + # Make Sparsity independent of _iterset, which can be a Subset, for better reusability. + get_map, region = ExplicitMatrixAssembler.integral_type_op2_map[local_kernel.kinfo.integral_type] + rmap_ = get_map(test).split[i] if get_map(test) is not None else None + cmap_ = get_map(trial).split[j] if get_map(trial) is not None else None + maps_and_regions[(i, j)][(rmap_, cmap_)].add(region) + return {block_indices: [map_pair + (tuple(region_set), ) for map_pair, region_set in map_pair_to_region_set.items()] + for block_indices, map_pair_to_region_set in maps_and_regions.items()} - def _apply_bc(self, bc): - op2tensor = self._tensor.M - spaces = tuple(a.function_space() for a in self._tensor.a.arguments()) + @staticmethod + def _make_maps_and_regions_default(test, trial, allocation_integral_types): + # Make maps using outer-product of the component maps + # using the given allocation_integral_types. + if allocation_integral_types is None: + raise ValueError("allocation_integral_types can not be None") + maps_and_regions = defaultdict(lambda: defaultdict(set)) + # Use outer product of component maps. + for integral_type in allocation_integral_types: + get_map, region = ExplicitMatrixAssembler.integral_type_op2_map[integral_type] + for i, rmap_ in enumerate(get_map(test)): + for j, cmap_ in enumerate(get_map(trial)): + maps_and_regions[(i, j)][(rmap_, cmap_)].add(region) + return {block_indices: [map_pair + (tuple(region_set), ) for map_pair, region_set in map_pair_to_region_set.items()] + for block_indices, map_pair_to_region_set in maps_and_regions.items()} + + @classmethod + @property + def integral_type_op2_map(cls): + try: + return cls._integral_type_op2_map + except AttributeError: + get_cell_map = operator.methodcaller("cell_node_map") + get_extf_map = operator.methodcaller("exterior_facet_node_map") + get_intf_map = operator.methodcaller("interior_facet_node_map") + cls._integral_type_op2_map = {"cell": (get_cell_map, op2.ALL), + "exterior_facet_bottom": (get_cell_map, op2.ON_BOTTOM), + "exterior_facet_top": (get_cell_map, op2.ON_TOP), + "interior_facet_horiz": (get_cell_map, op2.ON_INTERIOR_FACETS), + "exterior_facet": (get_extf_map, op2.ALL), + "exterior_facet_vert": (get_extf_map, op2.ALL), + "interior_facet": (get_intf_map, op2.ALL), + "interior_facet_vert": (get_intf_map, op2.ALL)} + return cls._integral_type_op2_map + + @cached_property + def _all_local_kernels(self): + """Collection of parloop_builders used for sparsity construction. + + When constructing sparsity, we use all parloop_builders + that are to be used in the actual assembly. + """ + all_local_kernels = tuple(local_kernel for local_kernel, _ in self.local_kernels) + for bc in self._bcs: + if isinstance(bc, EquationBCSplit): + _assembler = type(self)(bc.f, bcs=bc.bcs, form_compiler_parameters=self._form_compiler_params, needs_zeroing=False) + all_local_kernels += _assembler._all_local_kernels + return all_local_kernels + + def _apply_bc(self, tensor, bc): + op2tensor = tensor.M + spaces = tuple(a.function_space() for a in tensor.a.arguments()) V = bc.function_space() component = V.component if component is not None: @@ -1334,8 +1440,7 @@ def _apply_bc(self, bc): for j, s in enumerate(spaces[1]): if s.ufl_element().family() == "Real": self._apply_bcs_mat_real_block(op2tensor, index, j, component, bc.node_set) - type(self)(bc.f, self._tensor, bc.bcs, self._form_compiler_params, - needs_zeroing=False).assemble() + type(self)(bc.f, bcs=bc.bcs, form_compiler_parameters=self._form_compiler_params, needs_zeroing=False).assemble(tensor=tensor) else: raise AssertionError @@ -1346,41 +1451,60 @@ def _apply_bcs_mat_real_block(op2tensor, i, j, component, node_set): dat = op2.DatView(dat, component) dat.zero(subset=node_set) + def _check_tensor(self, tensor): + if tensor is not None and tensor.a.arguments() != self._form.arguments(): + raise ValueError("Form's arguments do not match provided result tensor") -class MatrixFreeAssembler: - """Stub class wrapping matrix-free assembly.""" + @staticmethod + def _as_pyop2_type(tensor): + return tensor.M - def __init__(self, tensor): - self._tensor = tensor + def result(self, tensor): + tensor.M.assemble() + return tensor - def assemble(self): - self._tensor.assemble() - return self._tensor +class MatrixFreeAssembler(FormAssembler): + """Stub class wrapping matrix-free assembly. -def get_form_assembler(form, tensor, *args, **kwargs): - """Provide the assemble method for `form`""" + Parameters + ---------- + form : ufl.Form or slate.TensorBasehe + 2-form. - # Don't expand derivatives if `mat_type` is 'matfree' - mat_type = kwargs.pop('mat_type', None) - if not isinstance(form, slate.TensorBase): - fc_params = kwargs.get('form_compiler_parameters') - # Only pre-process `form` once beforehand to avoid pre-processing for each assembly call - form = preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params) + Notes + ----- + See `FormAssembler` and `assemble` for descriptions of the other parameters. - if isinstance(form, (ufl.form.Form, slate.TensorBase)) and not base_form_operands(form): - diagonal = kwargs.pop('diagonal', False) - if len(form.arguments()) == 1 or diagonal: - return OneFormAssembler(form, tensor, *args, diagonal=diagonal, **kwargs).assemble - elif len(form.arguments()) == 2: - return TwoFormAssembler(form, tensor, *args, **kwargs).assemble - else: - raise ValueError('Expecting a 1-form or 2-form and not %s' % (form)) - elif isinstance(form, ufl.form.BaseForm): - return functools.partial(assemble_base_form, form, *args, tensor=tensor, mat_type=mat_type, - is_base_form_preprocessed=True, **kwargs) - else: - raise ValueError('Expecting a BaseForm or a slate.TensorBase object and not %s' % form) + """ + + @classmethod + def _cache_key(cls, *args, **kwargs): + return + + @FormAssembler._skip_if_initialised + def __init__(self, form, bcs=None, form_compiler_parameters=None, + options_prefix=None, appctx=None): + super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters) + self._options_prefix = options_prefix + self._appctx = appctx + + def allocate(self): + return matrix.ImplicitMatrix(self._form, self._bcs, + fc_params=self._form_compiler_params, + options_prefix=self._options_prefix, + appctx=self._appctx or {}) + + def assemble(self, tensor=None): + self._check_tensor(tensor) + if tensor is None: + tensor = self.allocate() + tensor.assemble() + return tensor + + def _check_tensor(self, tensor): + if tensor is not None and tensor.a.arguments() != self._form.arguments(): + raise ValueError("Form's arguments do not match provided result tensor") def _global_kernel_cache_key(form, local_knl, subdomain_id, all_integer_subdomain_ids, **kwargs): @@ -1702,42 +1826,137 @@ def _as_global_kernel_arg_layer_count(_, self): class ParloopBuilder: """Class that builds a :class:`op2.Parloop`. - :param form: The variational form. - :param local_knl: :class:`tsfc_interface.SplitKernel` compiled by either - TSFC or Slate. - :param global_knl: A :class:`pyop2.GlobalKernel` instance. - :param tensor: The output tensor to write to (cannot be ``None``). - :param subdomain_id: The subdomain of the mesh to iterate over. - :param all_integer_subdomain_ids: See :func:`tsfc_interface.gather_integer_subdomain_ids`. - :param diagonal: Are we assembling the diagonal of a 2-form? - :param lgmaps: Optional iterable of local-to-global maps needed for applying - boundary conditions to 2-forms. - """ + Parameters + ---------- + form : ufl.Form or slate.TensorBase + Variational form. + bcs : Sequence + Boundary conditions. + local_knl : tsfc_interface.SplitKernel + Kernel compiled by either TSFC or Slate. + subdomain_id : int + Subdomain of the mesh to iterate over. + all_integer_subdomain_ids : Sequence + See `tsfc_interface.gather_integer_subdomain_ids`. + diagonal : bool + Are we assembling the diagonal of a 2-form? - def __init__(self, form, local_knl, global_knl, tensor, subdomain_id, - all_integer_subdomain_ids, diagonal=False, lgmaps=None): + """ + def __init__(self, form, bcs, local_knl, subdomain_id, + all_integer_subdomain_ids, diagonal): self._form = form self._local_knl = local_knl - self._global_knl = global_knl self._subdomain_id = subdomain_id self._all_integer_subdomain_ids = all_integer_subdomain_ids - self._tensor = tensor self._diagonal = diagonal - self._lgmaps = lgmaps + self._bcs = bcs self._active_coefficients = _FormHandler.iter_active_coefficients(form, local_knl.kinfo) self._constants = _FormHandler.iter_constants(form, local_knl.kinfo) - def build(self): - """Construct the parloop.""" + def build(self, tensor): + """Construct the parloop. + + Parameters + ---------- + tensor : op2.Global or firedrake.cofunction.Cofunction or matrix.MatrixBase + The output tensor. + + """ + self._tensor = tensor parloop_args = [self._as_parloop_arg(tsfc_arg) for tsfc_arg in self._kinfo.arguments] + _global_knl = _make_global_kernel( + self._form, + self._local_knl, + self._subdomain_id, + self._all_integer_subdomain_ids, + diagonal=self._diagonal, + unroll=self.needs_unrolling() + ) try: - return op2.Parloop(self._global_knl, self._iterset, parloop_args) + return op2.Parloop(_global_knl, self._iterset, parloop_args) except MapValueError: raise RuntimeError("Integral measure does not match measure of all " "coefficients/arguments") + @property + def test_function_space(self): + assert len(self._form.arguments()) == 2 and not self._diagonal + test, _ = self._form.arguments() + return test.function_space() + + @property + def trial_function_space(self): + assert len(self._form.arguments()) == 2 and not self._diagonal + _, trial = self._form.arguments() + return trial.function_space() + + def get_indicess(self): + assert len(self._form.arguments()) == 2 and not self._diagonal + if all(i is None for i in self._local_knl.indices): + test, trial = self._form.arguments() + return numpy.ndindex((len(test.function_space()), + len(trial.function_space()))) + else: + assert all(i is not None for i in self._local_knl.indices) + return self._local_knl.indices, + + def _filter_bcs(self, row, col): + assert len(self._form.arguments()) == 2 and not self._diagonal + if len(self.test_function_space) > 1: + bcrow = tuple(bc for bc in self._bcs + if bc.function_space_index() == row) + else: + bcrow = self._bcs + + if len(self.trial_function_space) > 1: + bccol = tuple(bc for bc in self._bcs + if bc.function_space_index() == col + and isinstance(bc, DirichletBC)) + else: + bccol = tuple(bc for bc in self._bcs if isinstance(bc, DirichletBC)) + return bcrow, bccol + + def needs_unrolling(self): + """Do we need to address matrix elements directly rather than in + a blocked fashion? + + This is slower but required for the application of some boundary conditions + to 2-forms. + + :param local_knl: A :class:`tsfc_interface.SplitKernel`. + :param bcs: Iterable of boundary conditions. + """ + if len(self._form.arguments()) == 2 and not self._diagonal: + for i, j in self.get_indicess(): + for bc in itertools.chain(*self._filter_bcs(i, j)): + if bc.function_space().component is not None: + return True + return False + + def collect_lgmaps(self): + """Return any local-to-global maps that need to be swapped out. + + This is only needed when applying boundary conditions to 2-forms. + + :param local_knl: A :class:`tsfc_interface.SplitKernel`. + :param bcs: Iterable of boundary conditions. + """ + if len(self._form.arguments()) == 2 and not self._diagonal: + if not self._bcs: + return None + lgmaps = [] + for i, j in self.get_indicess(): + row_bcs, col_bcs = self._filter_bcs(i, j) + rlgmap, clgmap = self._tensor.M[i, j].local_to_global_maps + rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap) + clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap) + lgmaps.append((rlgmap, clgmap)) + return tuple(lgmaps) + else: + return None + @property def _indices(self): return self._local_knl.indices @@ -1835,7 +2054,7 @@ def _as_parloop_arg_output(_, self): m = rmap or cmap return op2.DatParloopArg(tensor.handle.getPythonContext().dat, m) else: - return op2.MatParloopArg(tensor, (rmap, cmap), lgmaps=self._lgmaps) + return op2.MatParloopArg(tensor, (rmap, cmap), lgmaps=self.collect_lgmaps()) else: raise AssertionError @@ -1955,29 +2174,3 @@ def index_tensor(tensor, form, indices, diagonal): return tensor.M[i, j] if is_indexed else tensor.M else: raise AssertionError - - -def _get_mat_type(mat_type, sub_mat_type, arguments): - """Validate the matrix types provided by the user and set any that are - undefined to default values. - - :arg mat_type: (:class:`str`) PETSc matrix type for the assembled matrix. - :arg sub_mat_type: (:class:`str`) PETSc matrix type for blocks if - ``mat_type`` is ``"nest"``. - :arg arguments: The test and trial functions of the expression being assembled. - :raises ValueError: On bad arguments. - :returns: 2-:class:`tuple` of validated/default ``mat_type`` and ``sub_mat_type``. - """ - if mat_type is None: - mat_type = parameters.parameters["default_matrix_type"] - if any(V.ufl_element().family() == "Real" - for arg in arguments - for V in arg.function_space()): - mat_type = "nest" - if mat_type not in {"matfree", "aij", "baij", "nest", "dense"}: - raise ValueError(f"Unrecognised matrix type, '{mat_type}'") - if sub_mat_type is None: - sub_mat_type = parameters.parameters["default_sub_matrix_type"] - if sub_mat_type not in {"aij", "baij"}: - raise ValueError(f"Invalid submatrix type, '{sub_mat_type}' (not 'aij' or 'baij')") - return mat_type, sub_mat_type diff --git a/firedrake/external_operators/abstract_external_operators.py b/firedrake/external_operators/abstract_external_operators.py index 40cd8650bc..7f8cc17b21 100644 --- a/firedrake/external_operators/abstract_external_operators.py +++ b/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={}): diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index cbd458d3fe..7b9a1890a0 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -938,7 +938,7 @@ def make_interpolator(expr, V, subset, access, bcs=None): tensor = None else: sparsity = op2.Sparsity((V.dof_dset, argfs.dof_dset), - ((V.cell_node_map(), argfs_map),), + [(V.cell_node_map(), argfs_map, None)], # non-mixed name="%s_%s_sparsity" % (V.name, argfs.name), nest=False, block_sparse=True) diff --git a/firedrake/linear_solver.py b/firedrake/linear_solver.py index d423b36bbb..c1dfbcc07e 100644 --- a/firedrake/linear_solver.py +++ b/firedrake/linear_solver.py @@ -112,19 +112,19 @@ 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, tensor=b).assemble, b + return u, get_assembler(expr).assemble, b def _lifted(self, b): u, update, blift = self._rhs u.dat.zero() for bc in self.A.bcs: bc.apply(u) - update() + update(tensor=blift) # blift contains -A u_bc blift += b if isinstance(blift, cofunction.Cofunction): diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 8aea90d99f..3ee448730e 100644 --- a/firedrake/matrix_free/operators.py +++ b/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: diff --git a/firedrake/preconditioners/assembled.py b/firedrake/preconditioners/assembled.py index 9d2d502633..4b2e9f52af 100644 --- a/firedrake/preconditioners/assembled.py +++ b/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,13 +51,10 @@ 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, tensor=self.P, bcs=bcs, - form_compiler_parameters=fcp).assemble - self._assemble_P() + 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 Pmat = self.P.petscmat @@ -87,7 +84,7 @@ def initialize(self, pc): pc.setFromOptions() def update(self, pc): - self._assemble_P() + self._assemble_P(tensor=self.P) def form(self, pc, test, trial): _, P = pc.getOperators() diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index a15d2af662..53ce1e6979 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/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,15 +88,10 @@ 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, tensor=self.mixed_op, - form_compiler_parameters=fcp, - bcs=mixed_bcs).assemble - self._assemble_mixed_op() + 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 def _permute_nullspace(nsp): @@ -147,7 +142,7 @@ def _permute_nullspace(nsp): def update(self, pc): if hasattr(self, "mixed_op"): - self._assemble_mixed_op() + self._assemble_mixed_op(tensor=self.mixed_op) elif hasattr(self, "_permute_op"): for mat in self.pc.getOperators(): mat.destroy() diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 40d2ae7b15..893920d49b 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -149,12 +149,11 @@ 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, tensor=self.A, bcs=bcs_fdm, - form_compiler_parameters=fcp).assemble - self._assemble_A() + 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 if len(bcs) > 0: @@ -333,7 +332,7 @@ def _assemble_P(self): @PETSc.Log.EventDecorator("FDMUpdate") def update(self, pc): if hasattr(self, "A"): - self._assemble_A() + self._assemble_A(tensor=self.A) self._assemble_P() def apply(self, pc, x, y): @@ -539,17 +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(OneFormAssembler(mixed_form, tensor=tensor, diagonal=True, - form_compiler_parameters=fcp).assemble) + 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(OneFormAssembler(J_facet, tensor=gamma, diagonal=True, - form_compiler_parameters=fcp).assemble) + 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") @@ -2071,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 = [] @@ -2112,8 +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(OneFormAssembler(ufl.inner(TestFunction(Q), alpha)*dx, tensor=tensor, - form_compiler_parameters=fcp).assemble) + 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] @@ -2134,8 +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(OneFormAssembler(ufl.inner(TestFunction(Q), beta)*dx, tensor=tensor, - form_compiler_parameters=fcp).assemble) + 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 @@ -2157,13 +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(OneFormAssembler(ifacet_inner(TestFunction(Q), G), tensor=tensor, - form_compiler_parameters=fcp).assemble) + 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(OneFormAssembler(ifacet_inner(TestFunction(Q), PT), tensor=tensor, - form_compiler_parameters=fcp).assemble) + 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 @@ -2189,8 +2182,7 @@ def assemble_coefficients(self, J, fcp): if len(forms): form = sum(forms) if len(form.arguments()) == 1: - assembly_callables.append(OneFormAssembler(form, tensor=tensor, - form_compiler_parameters=fcp).assemble) + 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: diff --git a/firedrake/preconditioners/gtmg.py b/firedrake/preconditioners/gtmg.py index d11cc548c1..a1014f0baf 100644 --- a/firedrake/preconditioners/gtmg.py +++ b/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,15 +50,10 @@ 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, tensor=self.fine_op, - form_compiler_parameters=fcp, - bcs=fine_bcs).assemble - self._assemble_fine_op() + 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: fine_petscmat = P @@ -90,15 +85,10 @@ 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, tensor=self.coarse_op, - form_compiler_parameters=fcp, - bcs=coarse_space_bcs).assemble - self._assemble_coarse_op() + 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 # Set nullspace if provided @@ -158,9 +148,9 @@ def initialize(self, pc): def update(self, pc): if hasattr(self, "fine_op"): - self._assemble_fine_op() + self._assemble_fine_op(tensor=self.fine_op) - self._assemble_coarse_op() + self._assemble_coarse_op(tensor=self.coarse_op) self.pc.setUp() def apply(self, pc, X, Y): diff --git a/firedrake/preconditioners/hiptmair.py b/firedrake/preconditioners/hiptmair.py index 9981076a68..262f666b51 100644 --- a/firedrake/preconditioners/hiptmair.py +++ b/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,15 +47,10 @@ 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, tensor=self.coarse_op, - form_compiler_parameters=fcp, - bcs=coarse_space_bcs).assemble - self._assemble_coarse_op() + 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 # We set up a PCMG object that uses the constructed interpolation @@ -103,7 +98,7 @@ def initialize(self, pc): coarse_solver.setFromOptions() def update(self, pc): - self._assemble_coarse_op() + self._assemble_coarse_op(tensor=self.coarse_op) self.pc.setUp() def apply(self, pc, X, Y): diff --git a/firedrake/preconditioners/pcd.py b/firedrake/preconditioners/pcd.py index 978000e930..12e60ed51e 100644 --- a/firedrake/preconditioners/pcd.py +++ b/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,17 +111,15 @@ 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, tensor=self.Fp, - form_compiler_parameters=context.fc_params).assemble - self._assemble_Fp() + 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)] def update(self, pc): - self._assemble_Fp() + self._assemble_Fp(tensor=self.Fp) def apply(self, pc, x, y): a, b = self.workspace diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index 82b538aea6..0d808eed6c 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -1535,8 +1535,10 @@ def prolongation_matrix_aij(P1, Pk, P1_bcs=[], Pk_bcs=[]): Pk = Pk.function_space() sp = op2.Sparsity((Pk.dof_dset, P1.dof_dset), - (Pk.cell_node_map(), - P1.cell_node_map())) + {(i, j): [(rmap, cmap, None)] + for i, rmap in enumerate(Pk.cell_node_map()) + for j, cmap in enumerate(P1.cell_node_map()) + if i == j}) mat = op2.Mat(sp, PETSc.ScalarType) mesh = Pk.mesh() diff --git a/firedrake/projection.py b/firedrake/projection.py index ccb2147d41..9d51482d10 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -218,13 +218,13 @@ def rhs_form(self): @cached_property def assembler(self): - from firedrake.assemble import OneFormAssembler - return OneFormAssembler(self.rhs_form, tensor=self.residual, - 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): - self.assembler() + self.assembler(tensor=self.residual) return self.residual diff --git a/firedrake/slate/static_condensation/hybridization.py b/firedrake/slate/static_condensation/hybridization.py index db5328b1c8..5dd592c44e 100644 --- a/firedrake/slate/static_condensation/hybridization.py +++ b/firedrake/slate/static_condensation/hybridization.py @@ -1,3 +1,4 @@ +import functools import numbers import numpy as np @@ -41,7 +42,7 @@ def initialize(self, pc): from firedrake import (FunctionSpace, Cofunction, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC) - from firedrake.assemble import allocate_matrix, OneFormAssembler, TwoFormAssembler + from firedrake.assemble import get_assembler from ufl.algorithms.replace import replace # Extract the problem context @@ -205,20 +206,15 @@ def initialize(self, pc): # Assemble the Schur complement operator and right-hand side self.schur_rhs = Cofunction(TraceSpace.dual()) - self._assemble_Srhs = OneFormAssembler(schur_rhs, tensor=self.schur_rhs, - form_compiler_parameters=self.ctx.fc_params).assemble + self._assemble_Srhs = get_assembler(schur_rhs, form_compiler_parameters=self.ctx.fc_params).assemble mat_type = PETSc.Options().getString(prefix + "mat_type", "aij") - self.S = allocate_matrix(schur_comp, bcs=trace_bcs, - form_compiler_parameters=self.ctx.fc_params, - mat_type=mat_type, - options_prefix=prefix, - appctx=self.get_appctx(pc)) - self._assemble_S = TwoFormAssembler(schur_comp, tensor=self.S, bcs=trace_bcs, - form_compiler_parameters=self.ctx.fc_params).assemble + form_assembler = get_assembler(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) + self.S = form_assembler.allocate() + self._assemble_S = form_assembler.assemble with PETSc.Log.Event("HybridOperatorAssembly"): - self._assemble_S() + self._assemble_S(tensor=self.S) Smat = self.S.petscmat @@ -268,7 +264,7 @@ def _reconstruction_calls(self): """This generates the reconstruction calls for the unknowns using the Lagrange multipliers. """ - from firedrake.assemble import OneFormAssembler + from firedrake.assemble import get_assembler # We always eliminate the velocity block first id0, id1 = (self.vidx, self.pidx) @@ -303,20 +299,18 @@ def _reconstruction_calls(self): rhs = Shat * rhs u_rec = S.solve(rhs, decomposition="PartialPivLU") - self._sub_unknown = OneFormAssembler(u_rec, tensor=u, - form_compiler_parameters=self.ctx.fc_params).assemble + self._sub_unknown = functools.partial(get_assembler(u_rec, form_compiler_parameters=self.ctx.fc_params).assemble, tensor=u) sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar, decomposition="PartialPivLU") - self._elim_unknown = OneFormAssembler(sigma_rec, tensor=sigma, - form_compiler_parameters=self.ctx.fc_params).assemble + self._elim_unknown = functools.partial(get_assembler(sigma_rec, form_compiler_parameters=self.ctx.fc_params).assemble, tensor=sigma) @PETSc.Log.EventDecorator("HybridUpdate") def update(self, pc): """Update by assembling into the operator. No need to reconstruct symbolic objects. """ - self._assemble_S() + self._assemble_S(tensor=self.S) def forward_elimination(self, pc, x): """Perform the forward elimination of fields and @@ -354,7 +348,7 @@ def forward_elimination(self, pc, x): with PETSc.Log.Event("HybridRHS"): # Compute the rhs for the multiplier system - self._assemble_Srhs() + self._assemble_Srhs(tensor=self.schur_rhs) def sc_solve(self, pc): """Solve the condensed linear system for the diff --git a/firedrake/slate/static_condensation/scpc.py b/firedrake/slate/static_condensation/scpc.py index 011a802627..35fa4742eb 100644 --- a/firedrake/slate/static_condensation/scpc.py +++ b/firedrake/slate/static_condensation/scpc.py @@ -1,3 +1,4 @@ +import functools import firedrake.dmhooks as dmhooks import numpy as np from firedrake.slate.static_condensation.sc_base import SCBase @@ -27,7 +28,7 @@ def initialize(self, pc): variables are recovered via back-substitution. """ - from firedrake.assemble import allocate_matrix, OneFormAssembler, TwoFormAssembler + from firedrake.assemble import get_assembler from firedrake.bcs import DirichletBC from firedrake.function import Function from firedrake.cofunction import Cofunction @@ -101,22 +102,14 @@ def initialize(self, pc): r_expr = reduced_sys.rhs # Construct the condensed right-hand side - self._assemble_Srhs = OneFormAssembler(r_expr, tensor=self.condensed_rhs, - bcs=bcs, zero_bc_nodes=True, - form_compiler_parameters=self.cxt.fc_params).assemble + self._assemble_Srhs = get_assembler(r_expr, bcs=bcs, zero_bc_nodes=True, form_compiler_parameters=self.cxt.fc_params).assemble # Allocate and set the condensed operator - self.S = allocate_matrix(S_expr, - bcs=bcs, - form_compiler_parameters=self.cxt.fc_params, - mat_type=mat_type, - options_prefix=prefix, - appctx=self.get_appctx(pc)) + form_assembler = get_assembler(S_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) + self.S = form_assembler.allocate() + self._assemble_S = form_assembler.assemble - self._assemble_S = TwoFormAssembler(S_expr, tensor=self.S, bcs=bcs, - form_compiler_parameters=self.cxt.fc_params).assemble - - self._assemble_S() + self._assemble_S(tensor=self.S) Smat = self.S.petscmat # If a different matrix is used for preconditioning, @@ -131,17 +124,11 @@ def initialize(self, pc): self.S_pc_expr = S_pc_expr # Allocate and set the condensed operator - self.S_pc = allocate_matrix(S_expr, - bcs=bcs, - form_compiler_parameters=self.cxt.fc_params, - mat_type=mat_type, - options_prefix=prefix, - appctx=self.get_appctx(pc)) - - self._assemble_S_pc = TwoFormAssembler(S_pc_expr, tensor=self.S_pc, bcs=bcs, - form_compiler_parameters=self.cxt.fc_params).assemble + form_assembler = get_assembler(S_pc_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) + self.S_pc = form_assembler.allocate() + self._assemble_S_pc = form_assembler.assemble - self._assemble_S_pc() + self._assemble_S_pc(tensor=self.S_pc) Smat_pc = self.S_pc.petscmat else: @@ -216,7 +203,7 @@ def local_solver_calls(self, A, rhs, x, elim_fields, schur_builder): to recover. :arg schur_builder: a `SchurComplementBuilder`. """ - from firedrake.assemble import OneFormAssembler + from firedrake.assemble import get_assembler from firedrake.slate.static_condensation.la_utils import backward_solve fields = x.subfunctions @@ -228,8 +215,7 @@ def local_solver_calls(self, A, rhs, x, elim_fields, schur_builder): be = local_system.rhs i, = local_system.field_idx local_solve = Aeinv * be - solve_call = OneFormAssembler(local_solve, tensor=fields[i], - form_compiler_parameters=self.cxt.fc_params).assemble + solve_call = functools.partial(get_assembler(local_solve, form_compiler_parameters=self.cxt.fc_params).assemble, tensor=fields[i]) local_solvers.append(solve_call) return local_solvers @@ -240,12 +226,12 @@ def update(self, pc): need to reconstruct symbolic objects. """ - self._assemble_S() + self._assemble_S(tensor=self.S) # Only reassemble if a preconditioning operator # is provided for the condensed system if hasattr(self, "S_pc"): - self._assemble_S_pc() + self._assemble_S_pc(tensor=self.S_pc) def forward_elimination(self, pc, x): """Perform the forward elimination of fields and @@ -264,7 +250,7 @@ def forward_elimination(self, pc, x): vc.pointwiseMult(vc, wc) # Now assemble residual for the reduced problem - self._assemble_Srhs() + self._assemble_Srhs(tensor=self.condensed_rhs) def sc_solve(self, pc): """Solve the condensed linear system for the diff --git a/firedrake/slope_limiter/vertex_based_limiter.py b/firedrake/slope_limiter/vertex_based_limiter.py index d4f8943c1f..e06682ee26 100644 --- a/firedrake/slope_limiter/vertex_based_limiter.py +++ b/firedrake/slope_limiter/vertex_based_limiter.py @@ -1,5 +1,6 @@ from firedrake import dx, assemble, LinearSolver from firedrake.function import Function +from firedrake.cofunction import Cofunction from firedrake.functionspace import FunctionSpace from firedrake.parloops import par_loop, READ, RW, MIN, MAX from firedrake.ufl_expr import TrialFunction, TestFunction @@ -35,7 +36,7 @@ def __init__(self, space): # Storage containers for cell means, max and mins self.centroids = Function(self.P0) - self.centroids_rhs = Function(self.P0) + self.centroids_rhs = Cofunction(self.P0.dual()) self.max_field = Function(self.P1CG) self.min_field = Function(self.P1CG) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index f307b9a382..0d62c9be13 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -178,7 +178,7 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None, post_jacobian_callback=None, post_function_callback=None, options_prefix=None, transfer_manager=None): - from firedrake.assemble import get_form_assembler + from firedrake.assemble import get_assembler if pmat_type is None: pmat_type = mat_type @@ -233,9 +233,9 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None, self.bcs_J = tuple(bc.extract_form('J') for bc in problem.bcs) self.bcs_Jp = tuple(bc.extract_form('Jp') for bc in problem.bcs) - self._assemble_residual = get_form_assembler(self.F, self._F, bcs=self.bcs_F, - form_compiler_parameters=self.fcp, - zero_bc_nodes=True) + self._assemble_residual = get_assembler(self.F, bcs=self.bcs_F, + form_compiler_parameters=self.fcp, + zero_bc_nodes=True).assemble self._jacobian_assembled = False self._splits = {} @@ -412,7 +412,7 @@ def form_function(snes, X, F): if ctx._pre_function_callback is not None: ctx._pre_function_callback(X) - ctx._assemble_residual() + ctx._assemble_residual(tensor=ctx._F) if ctx._post_function_callback is not None: with ctx._F.dat.vec as F_: @@ -450,15 +450,14 @@ def form_jacobian(snes, X, J, P): if ctx._pre_jacobian_callback is not None: ctx._pre_jacobian_callback(X) - - ctx._assemble_jac() + ctx._assemble_jac(ctx._jac) if ctx._post_jacobian_callback is not None: ctx._post_jacobian_callback(X, J) if ctx.Jp is not None: assert P.handle == ctx._pjac.petscmat.handle - ctx._assemble_pjac() + ctx._assemble_pjac(ctx._pjac) ises = problem.J.arguments()[0].function_space()._ises ctx.set_nullspace(ctx._nullspace, ises, transpose=False, near=False) @@ -493,50 +492,46 @@ def compute_operators(ksp, J, P): if isinstance(bc, DirichletBC): bc.apply(ctx._x) - ctx._assemble_jac() + ctx._assemble_jac(ctx._jac) if ctx.Jp is not None: assert P.handle == ctx._pjac.petscmat.handle - ctx._assemble_pjac() + ctx._assemble_pjac(ctx._pjac) + + @cached_property + def _assembler_jac(self): + from firedrake.assemble import get_assembler + return get_assembler(self.J, bcs=self.bcs_J, form_compiler_parameters=self.fcp, mat_type=self.mat_type, options_prefix=self.options_prefix, appctx=self.appctx) @cached_property def _jac(self): - from firedrake.assemble import allocate_matrix - return allocate_matrix(self.J, - bcs=self.bcs_J, - form_compiler_parameters=self.fcp, - mat_type=self.mat_type, - appctx=self.appctx, - options_prefix=self.options_prefix) + return self._assembler_jac.allocate() @cached_property def _assemble_jac(self): - from firedrake.assemble import get_form_assembler - return get_form_assembler(self.J, self._jac, bcs=self.bcs_J, - mat_type=self.mat_type, - form_compiler_parameters=self.fcp) + return self._assembler_jac.assemble @cached_property def is_mixed(self): return self._jac.block_shape != (1, 1) + @cached_property + def _assembler_pjac(self): + from firedrake.assemble import get_assembler + if self.mat_type != self.pmat_type or self._problem.Jp is not None: + return get_assembler(self.Jp, bcs=self.bcs_Jp, form_compiler_parameters=self.fcp, mat_type=self.pmat_type, options_prefix=self.options_prefix, appctx=self.appctx) + else: + return self._assembler_jac + @cached_property def _pjac(self): if self.mat_type != self.pmat_type or self._problem.Jp is not None: - from firedrake.assemble import allocate_matrix - return allocate_matrix(self.Jp, - bcs=self.bcs_Jp, - form_compiler_parameters=self.fcp, - mat_type=self.pmat_type, - appctx=self.appctx, - options_prefix=self.options_prefix) + return self._assembler_pjac.allocate() else: return self._jac @cached_property def _assemble_pjac(self): - from firedrake.assemble import get_form_assembler - return get_form_assembler(self.Jp, self._pjac, bcs=self.bcs_Jp, - form_compiler_parameters=self.fcp) + return self._assembler_pjac.assemble @cached_property def _F(self): diff --git a/tests/equation_bcs/test_equation_bcs_assemble.py b/tests/equation_bcs/test_equation_bcs_assemble.py index 21d67415c9..402148edfd 100644 --- a/tests/equation_bcs/test_equation_bcs_assemble.py +++ b/tests/equation_bcs/test_equation_bcs_assemble.py @@ -15,8 +15,6 @@ def test_equation_bcs_direct_assemble_one_form(): g = assemble(F, bcs=bc.extract_form('F')) assert np.allclose(g.dat.data, [0.5, 0.5, 0, 0]) - g = assemble(F, bcs=bc) - assert np.allclose(g.dat.data, [0.5, 0.5, 0, 0]) def test_equation_bcs_direct_assemble_two_form(): diff --git a/tests/regression/test_assemble.py b/tests/regression/test_assemble.py index 7751f95355..a80b46d5f0 100644 --- a/tests/regression/test_assemble.py +++ b/tests/regression/test_assemble.py @@ -1,7 +1,7 @@ import pytest import numpy as np from firedrake import * -from firedrake.utils import ScalarType +from firedrake.utils import ScalarType, IntType @pytest.fixture(scope='module') @@ -104,11 +104,11 @@ def test_assemble_with_tensor(mesh): V = FunctionSpace(mesh, "CG", 1) v = TestFunction(V) L = conj(v) * dx - f = Function(V) + f = Cofunction(V.dual()) # Assemble a form into f - f = assemble(L, f) + f = assemble(L, tensor=f) # Assemble a different form into f - f = assemble(Constant(2)*L, f) + f = assemble(Constant(2)*L, tensor=f) # Make sure we get the result of the last assembly assert np.allclose(f.dat.data, 2*assemble(L).dat.data, rtol=1e-14) @@ -120,7 +120,7 @@ def test_assemble_mat_with_tensor(mesh): a = inner(u, v) * dx M = assemble(a) # Assemble a different form into M - M = assemble(Constant(2)*a, M) + M = assemble(Constant(2)*a, tensor=M) # Make sure we get the result of the last assembly assert np.allclose(M.M.values, 2*assemble(a).M.values, rtol=1e-14) @@ -193,7 +193,7 @@ def test_one_form_assembler_cache(mesh): assert len(L._cache[_FORM_CACHE_KEY]) == 1 # changing tensor should not increase the cache size - tensor = Function(V) + tensor = Cofunction(V.dual()) assemble(L, tensor=tensor) assert len(L._cache[_FORM_CACHE_KEY]) == 1 @@ -291,3 +291,28 @@ def test_assemble_vector_rspace_one_form(mesh): U = inner(u, u)*dx L = derivative(U, u) assemble(L) + + +def test_assemble_sparsity_no_redundant_entries(): + mesh = UnitSquareMesh(2, 2, quadrilateral=True) + V = FunctionSpace(mesh, "CG", 1) + W = V * V * V + u = TrialFunction(W) + v = TestFunction(W) + A = assemble(inner(u, v) * dx, mat_type="nest") + for i in range(len(W)): + for j in range(len(W)): + if i != j: + assert np.all(A.M.sparsity[i][j].nnz == np.zeros(9, dtype=IntType)) + + +def test_assemble_sparsity_diagonal_entries_for_bc(): + mesh = UnitSquareMesh(1, 1, quadrilateral=True) + V = FunctionSpace(mesh, "CG", 1) + W = V * V + u = TrialFunction(W) + v = TestFunction(W) + bc = DirichletBC(W.sub(1), 0, "on_boundary") + A = assemble(inner(u[1], v[0]) * dx, bcs=[bc], mat_type="nest") + # Make sure that diagonals are allocated. + assert np.all(A.M.sparsity[1][1].nnz == np.ones(4, dtype=IntType)) diff --git a/tests/regression/test_assemble_baseform.py b/tests/regression/test_assemble_baseform.py index 41a2965944..db2f7278dc 100644 --- a/tests/regression/test_assemble_baseform.py +++ b/tests/regression/test_assemble_baseform.py @@ -1,7 +1,7 @@ import pytest import numpy as np from firedrake import * -from firedrake.assemble import preprocess_base_form, allocate_matrix +from firedrake.assemble import BaseFormAssembler, get_assembler from firedrake.utils import ScalarType import ufl @@ -159,7 +159,7 @@ def test_preprocess_form(M, a, f): from ufl.algorithms import expand_indices, expand_derivatives expr = action(action(M, M), f) - A = preprocess_base_form(expr) + A = BaseFormAssembler.preprocess_base_form(expr) B = action(expand_derivatives(M), action(M, f)) assert isinstance(A, ufl.Action) @@ -186,7 +186,7 @@ def test_tensor_copy(a, M): assert abs(f.dat.data.sum() - f2.dat.data.sum()) < 1.0e-12 # 2-form tensor - tensor = allocate_matrix(M) + tensor = get_assembler(M).allocate() formsum = assemble(M) + M res = assemble(formsum, tensor=tensor) diff --git a/tests/regression/test_periodic_interval_advection.py b/tests/regression/test_periodic_interval_advection.py index faa6796ff8..aa72582a11 100644 --- a/tests/regression/test_periodic_interval_advection.py +++ b/tests/regression/test_periodic_interval_advection.py @@ -46,7 +46,7 @@ def run_test(degree): nstep = 200 dt = Constant(5e-5) arhs = action(a_mass - dt * (a_int + a_flux), D1) - rhs = Function(V) + rhs = Cofunction(V.dual()) # Since DG mass-matrix is diagonal, just assemble the # diagonal and then "solve" is an entry-wise division. diff --git a/tests/regression/test_periodic_rectangle_advection.py b/tests/regression/test_periodic_rectangle_advection.py index b78a09c328..d818194136 100644 --- a/tests/regression/test_periodic_rectangle_advection.py +++ b/tests/regression/test_periodic_rectangle_advection.py @@ -72,7 +72,7 @@ def test_periodic_rectangle_advection(degree, threshold, nstep = 200 dt = Constant(5e-5) arhs = action(a_mass - dt * (a_int + a_flux), D1) - rhs = Function(V) + rhs = Cofunction(V.dual()) # Since DG mass-matrix is diagonal, just assemble the # diagonal and then "solve" is an entry-wise division. diff --git a/tests/slate/test_assemble_tensors.py b/tests/slate/test_assemble_tensors.py index ac6b4663b6..5aff159b9b 100644 --- a/tests/slate/test_assemble_tensors.py +++ b/tests/slate/test_assemble_tensors.py @@ -133,11 +133,11 @@ def test_assemble_matrix(rank_two_tensor): def test_assemble_vector_into_tensor(mesh): V = FunctionSpace(mesh, "DG", 1) v = TestFunction(V) - f = Function(V) + f = Cofunction(V.dual()) # Assemble a SLATE tensor into f - f = assemble(Tensor(v * dx), f) + f = assemble(Tensor(v * dx), tensor=f) # Assemble a different tensor into f - f = assemble(Tensor(Constant(2) * v * dx), f) + f = assemble(Tensor(Constant(2) * v * dx), tensor=f) assert np.allclose(f.dat.data, 2*assemble(Tensor(v * dx)).dat.data, rtol=1e-14) @@ -147,7 +147,7 @@ def test_assemble_matrix_into_tensor(mesh): v = TrialFunction(V) M = assemble(Tensor(u * v * dx)) # Assemble a different SLATE tensor into M - M = assemble(Tensor(Constant(2) * u * v * dx), M) + M = assemble(Tensor(Constant(2) * u * v * dx), tensor=M) assert np.allclose(M.M.values, 2*assemble(Tensor(u * v * dx)).M.values, rtol=1e-14)