Skip to content

Commit

Permalink
Rewrite assign
Browse files Browse the repository at this point in the history
This commit:
- Restricts `assign` to only work for weighted sums of coefficients
  (plus addition of constants).
- Expunges codegen in favour of directly manipulating numpy arrays.
- Introduces an `Assigner` class to speed up repeated `assign` calls.
  • Loading branch information
connorjward committed Oct 13, 2022
1 parent 07ead98 commit 9c80185
Show file tree
Hide file tree
Showing 8 changed files with 62 additions and 801 deletions.
21 changes: 19 additions & 2 deletions firedrake/assemble.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from tsfc import kernel_args
from tsfc.finatinterface import create_element
import ufl
from firedrake import (assemble_expressions, extrusion_utils as eutils, matrix, parameters, solving,
from firedrake import (extrusion_utils as eutils, matrix, parameters, solving,
tsfc_interface, utils)
from firedrake.adjoint import annotate_assemble
from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit
Expand Down Expand Up @@ -100,7 +100,7 @@ def assemble(expr, *args, **kwargs):
if isinstance(expr, (ufl.form.Form, slate.TensorBase)):
return _assemble_form(expr, *args, **kwargs)
elif isinstance(expr, ufl.core.expr.Expr):
return assemble_expressions.assemble_expression(expr)
return _assemble_expr(expr)
else:
raise TypeError(f"Unable to assemble: {expr}")

Expand Down Expand Up @@ -290,6 +290,23 @@ def _assemble_form(form, tensor=None, bcs=None, *,
return assembler.assemble()


def _assemble_expr(expr):
"""Assemble a pointwise expression.
:arg expr: The :class:`ufl.core.expr.Expr` to be evaluated.
:returns: A :class:`firedrake.Function` containing the result of this evaluation.
"""
from firedrake.assign import Assigner
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")
result = firedrake.Function(V)
Assigner(result, expr).assign()
return result


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).
Expand Down
Loading

0 comments on commit 9c80185

Please sign in to comment.