def test_adjoint():
cell = triangle
V1 = FiniteElement("CG", cell, 1)
V2 = FiniteElement("CG", cell, 2)
u = TrialFunction(V1)
v = TestFunction(V2)
assert u.number() > v.number()
u2 = Argument(V1, 2)
v2 = Argument(V2, 3)
assert u2.number() < v2.number()
a = u * v * dx
a_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(a)]
assert a_arg_degrees == [2, 1]
b = adjoint(a)
b_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(b)]
assert b_arg_degrees == [1, 2]
c = adjoint(a, (u2, v2))
c_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(c)]
assert c_arg_degrees == [1, 2]
d = adjoint(b)
d_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(d)]
assert d_arg_degrees == [2, 1]
def derivative(form, u, du=None, coefficient_derivatives=None):
if du is None:
# Get existing arguments from form and position the new one with the next argument number
from ufl.algorithms import extract_arguments
form_arguments = extract_arguments(form)
number = max([-1] + [arg.number() for arg in form_arguments]) + 1
if any(arg.part() is not None for arg in form_arguments):
cpp.dolfin_error("formmanipulation.py",
"compute derivative of form",
"Cannot automatically create third argument using parts, please supply one")
part = None
if isinstance(u, Function):
V = u.function_space()
du = Argument(V, number, part)
elif isinstance(u, (list,tuple)) and all(isinstance(w, Function) for w in u):
V = MixedFunctionSpace([w.function_space() for w in u])
du = ufl.split(Argument(V, number, part))
else:
cpp.dolfin_error("formmanipulation.py",
"compute derivative of form",
"Cannot automatically create third argument, please supply one")
return ufl.derivative(form, u, du, coefficient_derivatives)
def derivative(form, u, du=None, coefficient_derivatives=None):
if du is None:
# Get existing arguments from form and position the new one with the next argument number
from ufl.algorithms import extract_arguments
form_arguments = extract_arguments(form)
number = max([-1] + [arg.number() for arg in form_arguments]) + 1
if any(arg.part() is not None for arg in form_arguments):
cpp.dolfin_error(
"formmanipulation.py", "compute derivative of form",
"Cannot automatically create third argument using parts, please supply one"
)
part = None
if isinstance(u, Function):
V = u.function_space()
du = Argument(V, number, part)
elif isinstance(u, (list, tuple)) and all(
isinstance(w, Function) for w in u):
V = MixedFunctionSpace([w.function_space() for w in u])
du = ufl.split(Argument(V, number, part))
else:
cpp.dolfin_error(
"formmanipulation.py", "compute derivative of form",
"Cannot automatically create third argument, please supply one"
)
return ufl.derivative(form, u, du, coefficient_derivatives)
def make_interpolator(expr, V, subset, access):
assert isinstance(expr, ufl.classes.Expr)
if isinstance(expr, firedrake.Expression):
arguments = ()
else:
arguments = extract_arguments(expr)
if len(arguments) == 0:
if isinstance(V, firedrake.Function):
f = V
V = f.function_space()
else:
f = firedrake.Function(V)
tensor = f.dat
elif len(arguments) == 1:
if isinstance(V, firedrake.Function):
raise ValueError(
"Cannot interpolate an expression with an argument into a Function"
)
argfs = arguments[0].function_space()
sparsity = op2.Sparsity((V.dof_dset, argfs.dof_dset),
((V.cell_node_map(), argfs.cell_node_map()), ),
name="%s_%s_sparsity" % (V.name, argfs.name),
nest=False,
block_sparse=True)
tensor = op2.Mat(sparsity)
f = tensor
else:
raise ValueError("Cannot interpolate an expression with %d arguments" %
len(arguments))
# Make sure we have an expression of the right length i.e. a value for
# each component in the value shape of each function space
dims = [numpy.prod(fs.ufl_element().value_shape(), dtype=int) for fs in V]
loops = []
if numpy.prod(expr.ufl_shape, dtype=int) != sum(dims):
raise RuntimeError('Expression of length %d required, got length %d' %
(sum(dims), numpy.prod(expr.ufl_shape, dtype=int)))
if not isinstance(expr, firedrake.Expression):
if len(V) > 1:
raise NotImplementedError(
"UFL expressions for mixed functions are not yet supported.")
loops.extend(_interpolator(V, tensor, expr, subset, arguments, access))
elif hasattr(expr, 'eval'):
if len(V) > 1:
raise NotImplementedError(
"Python expressions for mixed functions are not yet supported."
)
loops.extend(_interpolator(V, tensor, expr, subset, arguments, access))
else:
raise ValueError("Don't know how to interpolate a %r" % expr)
def callable(loops, f):
for l in loops:
l()
return f
return partial(callable, loops, f), arguments
def __call__(self,M,b):
# FIXME we currently lose the variable names of the forms (this can be
# looked up in the uflObjects)
# b as the rhs should always only have a single argument
# Extract its element
element = extract_arguments(b)[0].element()
coeff = Coefficient(element)
self._solves[coeff.count()] = (M,b)
return coeff
def test_adjoint():
cell = triangle
V1 = FiniteElement("CG", cell, 1)
V2 = FiniteElement("CG", cell, 2)
u = TrialFunction(V1)
v = TestFunction(V2)
assert u.number() > v.number()
u2 = Argument(V1, 2)
v2 = Argument(V2, 3)
assert u2.number() < v2.number()
a = u * v * dx
a_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(a)
]
assert a_arg_degrees == [2, 1]
b = adjoint(a)
b_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(b)
]
assert b_arg_degrees == [1, 2]
c = adjoint(a, (u2, v2))
c_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(c)
]
assert c_arg_degrees == [1, 2]
d = adjoint(b)
d_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(d)
]
assert d_arg_degrees == [2, 1]
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py", "compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces (NB: Order matters
# here!)
v_1 = Argument(arguments[1].function_space())
v_0 = Argument(arguments[0].function_space())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces (NB: Order matters
# here!)
v_1 = Argument(arguments[1].function_space())
v_0 = Argument(arguments[0].function_space())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form, reordered_arguments=None):
"""UFL form operator:
Given a combined bilinear form, compute the adjoint form by
changing the ordering (number) of the test and trial functions.
By default, new Argument objects will be created with
opposite ordering. However, if the adjoint form is to
be added to other forms later, their arguments must match.
In that case, the user must provide a tuple reordered_arguments=(u2,v2).
"""
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.function_space(),
number=v.number(),
part=v.part()),
Argument(v.function_space(),
number=u.number(),
part=u.part()))
return ufl.adjoint(form, reordered_arguments)
def adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a form.
:arg form: A UFL form, or a Slate tensor.
:arg reordered_arguments: arguments to use when creating the
adjoint. Ignored if form is a Slate tensor.
If the form is a slate tensor, this just returns its transpose.
Otherwise, given a bilinear form, compute the adjoint form by
changing the ordering (number) of the test and trial functions.
By default, new Argument objects will be created with opposite
ordering. However, if the adjoint form is to be added to other
forms later, their arguments must match. In that case, the user
must provide a tuple reordered_arguments=(u2,v2).
"""
if isinstance(form, firedrake.slate.TensorBase):
if reordered_arguments is not None:
firedrake.warning(
"Ignoring arguments for adjoint of Slate tensor.")
if form.rank != 2:
raise ValueError("Expecting rank-2 tensor")
return form.T
else:
if len(form.arguments()) != 2:
raise ValueError("Expecting bilinear form")
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.function_space(),
number=v.number(),
part=v.part()),
Argument(v.function_space(),
number=u.number(),
part=u.part()))
return ufl.adjoint(form, reordered_arguments)
def adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a form.
:arg form: A UFL form, or a Slate tensor.
:arg reordered_arguments: arguments to use when creating the
adjoint. Ignored if form is a Slate tensor.
If the form is a slate tensor, this just returns its transpose.
Otherwise, given a bilinear form, compute the adjoint form by
changing the ordering (number) of the test and trial functions.
By default, new Argument objects will be created with opposite
ordering. However, if the adjoint form is to be added to other
forms later, their arguments must match. In that case, the user
must provide a tuple reordered_arguments=(u2,v2).
"""
if isinstance(form, firedrake.slate.TensorBase):
if reordered_arguments is not None:
firedrake.warning("Ignoring arguments for adjoint of Slate tensor.")
if form.rank != 2:
raise ValueError("Expecting rank-2 tensor")
return form.T
else:
if len(form.arguments()) != 2:
raise ValueError("Expecting bilinear form")
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.function_space(),
number=v.number(),
part=v.part()),
Argument(v.function_space(),
number=u.number(),
part=u.part()))
return ufl.adjoint(form, reordered_arguments)
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if any(arg.part() != None for arg in arguments):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"parts not supported")
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces
# (NB: Order does not matter anymore here because number is absolute)
v_1 = Argument(arguments[1].function_space(), arguments[0].number(), arguments[0].part())
v_0 = Argument(arguments[0].function_space(), arguments[1].number(), arguments[1].part())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if any(arg.part() != None for arg in arguments):
cpp.dolfin_error("formmanipulation.py", "compute adjoint of form",
"parts not supported")
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py", "compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces
# (NB: Order does not matter anymore here because number is absolute)
v_1 = Argument(arguments[1].function_space(), arguments[0].number(),
arguments[0].part())
v_0 = Argument(arguments[0].function_space(), arguments[1].number(),
arguments[1].part())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def _handle_derivative_arguments(form, coefficient, argument):
# Wrap single coefficient in tuple for uniform treatment below
if isinstance(coefficient, (list, tuple, ListTensor)):
coefficients = tuple(coefficient)
else:
coefficients = (coefficient,)
if argument is None:
# Try to create argument if not provided
if not all(isinstance(c, Coefficient) for c in coefficients):
error("Can only create arguments automatically for non-indexed coefficients.")
# Get existing arguments from form and position the new one
# with the next argument number
if isinstance(form, Form):
form_arguments = form.arguments()
else:
# To handle derivative(expression), which is at least used
# in tests. Remove?
form_arguments = extract_arguments(form)
numbers = sorted(set(arg.number() for arg in form_arguments))
number = max(numbers + [-1]) + 1
# Don't know what to do with parts, let the user sort it out
# in that case
parts = set(arg.part() for arg in form_arguments)
if len(parts - {None}) != 0:
error("Not expecting parts here, provide your own arguments.")
part = None
# Create argument and split it if in a mixed space
function_spaces = [c.ufl_function_space() for c in coefficients]
domains = [fs.ufl_domain() for fs in function_spaces]
elements = [fs.ufl_element() for fs in function_spaces]
if len(function_spaces) == 1:
arguments = (Argument(function_spaces[0], number, part),)
else:
# Create in mixed space over assumed (for now) same domain
assert all(fs.ufl_domain() == domains[0] for fs in function_spaces)
elm = MixedElement(*elements)
fs = FunctionSpace(domains[0], elm)
arguments = split(Argument(fs, number, part))
else:
# Wrap single argument in tuple for uniform treatment below
if isinstance(argument, (list, tuple)):
arguments = tuple(argument)
else:
n = len(coefficients)
if n == 1:
arguments = (argument,)
else:
if argument.ufl_shape == (n,):
arguments = tuple(argument[i] for i in range(n))
else:
arguments = split(argument)
# Build mapping from coefficient to argument
m = {}
for (c, a) in zip(coefficients, arguments):
if c.ufl_shape != a.ufl_shape:
error("Coefficient and argument shapes do not match!")
if isinstance(c, Coefficient) or isinstance(c, SpatialCoordinate):
m[c] = a
else:
if not isinstance(c, Indexed):
error("Invalid coefficient type for %s" % ufl_err_str(c))
f, i = c.ufl_operands
if not isinstance(f, Coefficient):
error("Expecting an indexed coefficient, not %s" % ufl_err_str(f))
if not (isinstance(i, MultiIndex) and all(isinstance(j, FixedIndex) for j in i)):
error("Expecting one or more fixed indices, not %s" % ufl_err_str(i))
i = tuple(int(j) for j in i)
if f not in m:
m[f] = {}
m[f][i] = a
# Merge coefficient derivatives (arguments) based on indices
for c, p in m.items():
if isinstance(p, dict):
a = zero_lists(c.ufl_shape)
for i, g in p.items():
set_list_item(a, i, g)
m[c] = as_tensor(a)
# Wrap and return generic tuples
items = sorted(m.items(), key=lambda x: x[0].count())
coefficients = ExprList(*[item[0] for item in items])
arguments = ExprList(*[item[1] for item in items])
return coefficients, arguments
def _handle_derivative_arguments(form, coefficient, argument):
# Wrap single coefficient in tuple for uniform treatment below
if isinstance(coefficient, (list, tuple, ListTensor)):
coefficients = tuple(coefficient)
else:
coefficients = (coefficient,)
if argument is None:
# Try to create argument if not provided
if not all(isinstance(c, Coefficient) for c in coefficients):
error("Can only create arguments automatically for non-indexed coefficients.")
# Get existing arguments from form and position the new one
# with the next argument number
if isinstance(form, Form):
form_arguments = form.arguments()
else:
# To handle derivative(expression), which is at least used
# in tests. Remove?
form_arguments = extract_arguments(form)
numbers = sorted(set(arg.number() for arg in form_arguments))
number = max(numbers + [-1]) + 1
# Don't know what to do with parts, let the user sort it out
# in that case
parts = set(arg.part() for arg in form_arguments)
if len(parts - {None}) != 0:
error("Not expecting parts here, provide your own arguments.")
part = None
# Create argument and split it if in a mixed space
function_spaces = [c.ufl_function_space() for c in coefficients]
domains = [fs.ufl_domain() for fs in function_spaces]
elements = [fs.ufl_element() for fs in function_spaces]
if len(function_spaces) == 1:
arguments = (Argument(function_spaces[0], number, part),)
else:
# Create in mixed space over assumed (for now) same domain
assert all(fs.ufl_domain() == domains[0] for fs in function_spaces)
elm = MixedElement(*elements)
fs = FunctionSpace(domains[0], elm)
arguments = split(Argument(fs, number, part))
else:
# Wrap single argument in tuple for uniform treatment below
if isinstance(argument, (list, tuple)):
arguments = tuple(argument)
else:
n = len(coefficients)
if n == 1:
arguments = (argument,)
else:
if argument.ufl_shape == (n,):
arguments = tuple(argument[i] for i in range(n))
else:
arguments = split(argument)
# Build mapping from coefficient to argument
m = {}
for (c, a) in zip(coefficients, arguments):
if c.ufl_shape != a.ufl_shape:
error("Coefficient and argument shapes do not match!")
if isinstance(c, Coefficient) or isinstance(c, SpatialCoordinate):
m[c] = a
else:
if not isinstance(c, Indexed):
error("Invalid coefficient type for %s" % ufl_err_str(c))
f, i = c.ufl_operands
if not isinstance(f, Coefficient):
error("Expecting an indexed coefficient, not %s" % ufl_err_str(f))
if not (isinstance(i, MultiIndex) and all(isinstance(j, FixedIndex) for j in i)):
error("Expecting one or more fixed indices, not %s" % ufl_err_str(i))
i = tuple(int(j) for j in i)
if f not in m:
m[f] = {}
m[f][i] = a
# Merge coefficient derivatives (arguments) based on indices
for c, p in m.items():
if isinstance(p, dict):
a = zero_lists(c.ufl_shape)
for i, g in p.items():
set_list_item(a, i, g)
m[c] = as_tensor(a)
# Wrap and return generic tuples
items = sorted(m.items(), key=lambda x: x[0].count())
coefficients = ExprList(*[item[0] for item in items])
arguments = ExprList(*[item[1] for item in items])
return coefficients, arguments
def compile_expression_at_points(expression, points, coordinates, parameters=None):
"""Compiles a UFL expression to be evaluated at compile-time known
reference points. Useful for interpolating UFL expressions onto
function spaces with only point evaluation nodes.
:arg expression: UFL expression
:arg points: reference coordinates of the evaluation points
:arg coordinates: the coordinate function
:arg parameters: parameters object
"""
import coffee.base as ast
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = ufl_utils.preprocess_expression(expression)
# Initialise kernel builder
builder = firedrake_interface.ExpressionKernelBuilder()
# Replace coordinates (if any)
domain = expression.ufl_domain()
if domain:
assert coordinates.ufl_domain() == domain
builder.domain_coordinate[domain] = coordinates
# Collect required coefficients
coefficients = extract_coefficients(expression)
if has_type(expression, GeometricQuantity):
coefficients = [coordinates] + coefficients
builder.set_coefficients(coefficients)
# Split mixed coefficients
expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# Translate to GEM
point_set = PointSet(points)
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_set=point_set)
ir, = fem.compile_ufl(expression, point_sum=False, **config)
# Deal with non-scalar expressions
value_shape = ir.shape
tensor_indices = tuple(gem.Index() for s in value_shape)
if value_shape:
ir = gem.Indexed(ir, tensor_indices)
# Build kernel body
return_shape = (len(points),) + value_shape
return_indices = point_set.indices + tensor_indices
return_var = gem.Variable('A', return_shape)
return_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('A', rank=return_shape))
return_expr = gem.Indexed(return_var, return_indices)
ir, = impero_utils.preprocess_gem([ir])
impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices)
point_index, = point_set.indices
body = generate_coffee(impero_c, {point_index: 'p'}, parameters["precision"])
# Handle cell orientations
if builder.needs_cell_orientations([ir]):
builder.require_cell_orientations()
# Build kernel tuple
return builder.construct_kernel(return_arg, body)
def derivative(form, u, du=None, coefficient_derivatives=None):
"""Compute the derivative of a form.
Given a form, this computes its linearization with respect to the
provided :class:`.Function`. The resulting form has one
additional :class:`Argument` in the same finite element space as
the Function.
:arg form: a :class:`~ufl.classes.Form` to compute the derivative of.
:arg u: a :class:`.Function` to compute the derivative with
respect to.
:arg du: an optional :class:`Argument` to use as the replacement
in the new form (constructed automatically if not provided).
:arg coefficient_derivatives: an optional :class:`dict` to
provide the derivative of a coefficient function.
:raises ValueError: If any of the coefficients in ``form`` were
obtained from ``u.split()``. UFL doesn't notice that these
are related to ``u`` and so therefore the derivative is
wrong (instead one should have written ``split(u)``).
See also :func:`ufl.derivative`.
"""
if isinstance(form, firedrake.slate.TensorBase):
raise TypeError(
f"Cannot take the derivative of a {type(form).__name__}"
)
# TODO: What about Constant?
u_is_x = isinstance(u, ufl.SpatialCoordinate)
uc, = (u,) if u_is_x else extract_coefficients(u)
if not u_is_x and len(uc.split()) > 1 and set(extract_coefficients(form)) & set(uc.split()):
raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.split()."
"\nYou probably meant to write split(u) when defining your form.")
mesh = form.ufl_domain()
if not mesh:
raise ValueError("Expression to be differentiated has no ufl domain."
"\nDo you need to add a domain to your Constant?")
is_dX = u_is_x or u is mesh.coordinates
try:
args = form.arguments()
except AttributeError:
args = extract_arguments(form)
def argument(V):
if du is None:
n = max(a.number() for a in args) if args else -1
return Argument(V, n + 1)
else:
return du
if is_dX:
coords = mesh.coordinates
u = ufl.SpatialCoordinate(mesh)
V = coords.function_space()
du = argument(V)
cds = {coords: du}
if coefficient_derivatives is not None:
cds.update(coefficient_derivatives)
coefficient_derivatives = cds
elif isinstance(uc, firedrake.Function):
V = uc.function_space()
du = argument(V)
elif isinstance(uc, firedrake.Constant):
if uc.ufl_shape != ():
raise ValueError("Real function space of vector elements not supported")
V = firedrake.FunctionSpace(mesh, "Real", 0)
du = argument(V)
else:
raise RuntimeError("Can't compute derivative for form")
if u.ufl_shape != du.ufl_shape:
raise ValueError("Shapes of u and du do not match.\n"
"If you passed an indexed part of split(u) into "
"derivative, you need to provide an appropriate du as well.")
return ufl.derivative(form, u, du, coefficient_derivatives)
def compile_element(expression, coordinates, parameters=None):
"""Generates C code for point evaluations.
:arg expression: UFL expression
:arg coordinates: coordinate field
:arg parameters: form compiler parameters
:returns: C code as string
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = tsfc.ufl_utils.preprocess_expression(
expression, complex_mode=utils.complex_mode)
# Collect required coefficients
coefficient, = extract_coefficients(expression)
# Point evaluation of mixed coefficients not supported here
if type(coefficient.ufl_element()) == MixedElement:
raise NotImplementedError("Cannot point evaluate mixed elements yet!")
# Replace coordinates (if any)
domain = expression.ufl_domain()
assert coordinates.ufl_domain() == domain
# Initialise kernel builder
builder = firedrake_interface.KernelBuilderBase(utils.ScalarType_c)
builder.domain_coordinate[domain] = coordinates
x_arg = builder._coefficient(coordinates, "x")
f_arg = builder._coefficient(coefficient, "f")
# TODO: restore this for expression evaluation!
# expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
point_arg = ast.Decl(utils.ScalarType_c, ast.Symbol('X', rank=(dim, )))
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_indices=(),
point_expr=point,
complex_mode=utils.complex_mode)
# TODO: restore this for expression evaluation!
# config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression, utils.complex_mode)
# Translate UFL -> GEM
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
tensor_indices = ()
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape),
tensor_indices)
result_arg = ast.Decl(utils.ScalarType_c,
ast.Symbol('R', rank=expression.ufl_shape))
result = gem.Indexed(result, tensor_indices)
else:
return_variable = gem.Indexed(gem.Variable('R', (1, )), (0, ))
result_arg = ast.Decl(utils.ScalarType_c, ast.Symbol('R', rank=(1, )))
# Unroll
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
result, = gem.optimise.unroll_indexsum([result], predicate=predicate)
# Translate GEM -> COFFEE
result, = gem.impero_utils.preprocess_gem([result])
impero_c = gem.impero_utils.compile_gem([(return_variable, result)],
tensor_indices)
body = generate_coffee(impero_c, {}, parameters["precision"],
utils.ScalarType_c)
# Build kernel tuple
kernel_code = builder.construct_kernel(
"evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)
# Fill the code template
extruded = isinstance(cell, TensorProductCell)
code = {
"geometric_dimension": cell.geometric_dimension(),
"layers_arg": ", int const *__restrict__ layers" if extruded else "",
"layers": ", layers" if extruded else "",
"IntType": as_cstr(IntType),
"scalar_type": utils.ScalarType_c,
}
# if maps are the same, only need to pass one of them
if coordinates.cell_node_map() == coefficient.cell_node_map():
code[
"wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map" % code
code["map_args"] = "f->coords_map"
else:
code[
"wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map, %(IntType)s const *__restrict__ f_map" % code
code["map_args"] = "f->coords_map, f->f_map"
evaluate_template_c = """
static inline void wrap_evaluate(%(scalar_type)s* const result, %(scalar_type)s* const X, int const start, int const end%(layers_arg)s,
%(scalar_type)s const *__restrict__ coords, %(scalar_type)s const *__restrict__ f, %(wrapper_map_args)s);
int evaluate(struct Function *f, %(scalar_type)s *x, %(scalar_type)s *result)
{
struct ReferenceCoords reference_coords;
%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &reference_coords);
if (cell == -1) {
return -1;
}
if (!result) {
return 0;
}
int layers[2] = {0, 0};
if (f->extruded != 0) {
int nlayers = f->n_layers;
layers[1] = cell %% nlayers + 2;
cell = cell / nlayers;
}
wrap_evaluate(result, reference_coords.X, cell, cell+1%(layers)s, f->coords, f->f, %(map_args)s);
return 0;
}
"""
return (evaluate_template_c % code) + kernel_code.gencode()
def compile_expression_dual_evaluation(expression,
to_element,
*,
domain=None,
interface=None,
parameters=None,
coffee=False):
"""Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis.
Useful for interpolating UFL expressions into e.g. N1curl spaces.
:arg expression: UFL expression
:arg to_element: A FInAT element for the target space
:arg domain: optional UFL domain the expression is defined on (required when expression contains no domain).
:arg interface: backend module for the kernel interface
:arg parameters: parameters object
:arg coffee: compile coffee kernel instead of loopy kernel
"""
import coffee.base as ast
import loopy as lp
# Just convert FInAT element to FIAT for now.
# Dual evaluation in FInAT will bring a thorough revision.
to_element = to_element.fiat_equivalent
if any(len(dual.deriv_dict) != 0 for dual in to_element.dual_basis()):
raise NotImplementedError(
"Can only interpolate onto dual basis functionals without derivative evaluation, sorry!"
)
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# Determine whether in complex mode
complex_mode = is_complex(parameters["scalar_type"])
# Find out which mapping to apply
try:
mapping, = set(to_element.mapping())
except ValueError:
raise NotImplementedError(
"Don't know how to interpolate onto zany spaces, sorry")
expression = apply_mapping(expression, mapping, domain)
# Apply UFL preprocessing
expression = ufl_utils.preprocess_expression(expression,
complex_mode=complex_mode)
# Initialise kernel builder
if interface is None:
if coffee:
import tsfc.kernel_interface.firedrake as firedrake_interface_coffee
interface = firedrake_interface_coffee.ExpressionKernelBuilder
else:
# Delayed import, loopy is a runtime dependency
import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy
interface = firedrake_interface_loopy.ExpressionKernelBuilder
builder = interface(parameters["scalar_type"])
arguments = extract_arguments(expression)
argument_multiindices = tuple(
builder.create_element(arg.ufl_element()).get_indices()
for arg in arguments)
# Replace coordinates (if any) unless otherwise specified by kwarg
if domain is None:
domain = expression.ufl_domain()
assert domain is not None
# Collect required coefficients
first_coefficient_fake_coords = False
coefficients = extract_coefficients(expression)
if has_type(expression, GeometricQuantity) or any(
fem.needs_coordinate_mapping(c.ufl_element())
for c in coefficients):
# Create a fake coordinate coefficient for a domain.
coords_coefficient = ufl.Coefficient(
ufl.FunctionSpace(domain, domain.ufl_coordinate_element()))
builder.domain_coordinate[domain] = coords_coefficient
builder.set_cell_sizes(domain)
coefficients = [coords_coefficient] + coefficients
first_coefficient_fake_coords = True
builder.set_coefficients(coefficients)
# Split mixed coefficients
expression = ufl_utils.split_coefficients(expression,
builder.coefficient_split)
# Translate to GEM
kernel_cfg = dict(
interface=builder,
ufl_cell=domain.ufl_cell(),
# FIXME: change if we ever implement
# interpolation on facets.
integral_type="cell",
argument_multiindices=argument_multiindices,
index_cache={},
scalar_type=parameters["scalar_type"])
if all(
isinstance(dual, PointEvaluation)
for dual in to_element.dual_basis()):
# This is an optimisation for point-evaluation nodes which
# should go away once FInAT offers the interface properly
qpoints = []
# Everything is just a point evaluation.
for dual in to_element.dual_basis():
ptdict = dual.get_point_dict()
qpoint, = ptdict.keys()
(qweight, component), = ptdict[qpoint]
assert allclose(qweight, 1.0)
assert component == ()
qpoints.append(qpoint)
point_set = PointSet(qpoints)
config = kernel_cfg.copy()
config.update(point_set=point_set)
# Allow interpolation onto QuadratureElements to refer to the quadrature
# rule they represent
if isinstance(to_element, FIAT.QuadratureElement):
assert allclose(asarray(qpoints), asarray(to_element._points))
quad_rule = QuadratureRule(point_set, to_element._weights)
config["quadrature_rule"] = quad_rule
expr, = fem.compile_ufl(expression, **config, point_sum=False)
# In some cases point_set.indices may be dropped from expr, but nothing
# new should now appear
assert set(expr.free_indices) <= set(
chain(point_set.indices, *argument_multiindices))
shape_indices = tuple(gem.Index() for _ in expr.shape)
basis_indices = point_set.indices
ir = gem.Indexed(expr, shape_indices)
else:
# This is general code but is more unrolled than necssary.
dual_expressions = [] # one for each functional
broadcast_shape = len(expression.ufl_shape) - len(
to_element.value_shape())
shape_indices = tuple(gem.Index()
for _ in expression.ufl_shape[:broadcast_shape])
expr_cache = {} # Sharing of evaluation of the expression at points
for dual in to_element.dual_basis():
pts = tuple(sorted(dual.get_point_dict().keys()))
try:
expr, point_set = expr_cache[pts]
except KeyError:
point_set = PointSet(pts)
config = kernel_cfg.copy()
config.update(point_set=point_set)
expr, = fem.compile_ufl(expression, **config, point_sum=False)
# In some cases point_set.indices may be dropped from expr, but
# nothing new should now appear
assert set(expr.free_indices) <= set(
chain(point_set.indices, *argument_multiindices))
expr = gem.partial_indexed(expr, shape_indices)
expr_cache[pts] = expr, point_set
weights = collections.defaultdict(list)
for p in pts:
for (w, cmp) in dual.get_point_dict()[p]:
weights[cmp].append(w)
qexprs = gem.Zero()
for cmp in sorted(weights):
qweights = gem.Literal(weights[cmp])
qexpr = gem.Indexed(expr, cmp)
qexpr = gem.index_sum(
gem.Indexed(qweights, point_set.indices) * qexpr,
point_set.indices)
qexprs = gem.Sum(qexprs, qexpr)
assert qexprs.shape == ()
assert set(qexprs.free_indices) == set(
chain(shape_indices, *argument_multiindices))
dual_expressions.append(qexprs)
basis_indices = (gem.Index(), )
ir = gem.Indexed(gem.ListTensor(dual_expressions), basis_indices)
# Build kernel body
return_indices = basis_indices + shape_indices + tuple(
chain(*argument_multiindices))
return_shape = tuple(i.extent for i in return_indices)
return_var = gem.Variable('A', return_shape)
if coffee:
return_arg = ast.Decl(parameters["scalar_type"],
ast.Symbol('A', rank=return_shape))
else:
return_arg = lp.GlobalArg("A",
dtype=parameters["scalar_type"],
shape=return_shape)
return_expr = gem.Indexed(return_var, return_indices)
# TODO: one should apply some GEM optimisations as in assembly,
# but we don't for now.
ir, = impero_utils.preprocess_gem([ir])
impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices)
index_names = dict(
(idx, "p%d" % i) for (i, idx) in enumerate(basis_indices))
# Handle kernel interface requirements
builder.register_requirements([ir])
# Build kernel tuple
return builder.construct_kernel(return_arg, impero_c, index_names,
first_coefficient_fake_coords)
def compile_element(expression, coordinates, parameters=None):
"""Generates C code for point evaluations.
:arg expression: UFL expression
:arg coordinates: coordinate field
:arg parameters: form compiler parameters
:returns: C code as string
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = tsfc.ufl_utils.preprocess_expression(expression)
# Collect required coefficients
coefficient, = extract_coefficients(expression)
# Point evaluation of mixed coefficients not supported here
if type(coefficient.ufl_element()) == MixedElement:
raise NotImplementedError("Cannot point evaluate mixed elements yet!")
# Replace coordinates (if any)
domain = expression.ufl_domain()
assert coordinates.ufl_domain() == domain
expression = tsfc.ufl_utils.replace_coordinates(expression, coordinates)
# Initialise kernel builder
builder = firedrake_interface.KernelBuilderBase()
x_arg = builder._coefficient(coordinates, "x")
f_arg = builder._coefficient(coefficient, "f")
# TODO: restore this for expression evaluation!
# expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
point_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('X', rank=(dim, )))
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_indices=(),
point_expr=point)
# TODO: restore this for expression evaluation!
# config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
tensor_indices = ()
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape),
tensor_indices)
result_arg = ast.Decl(SCALAR_TYPE,
ast.Symbol('R', rank=expression.ufl_shape))
result = gem.Indexed(result, tensor_indices)
else:
return_variable = gem.Indexed(gem.Variable('R', (1, )), (0, ))
result_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('R', rank=(1, )))
# Unroll
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
result, = gem.optimise.unroll_indexsum([result], predicate=predicate)
# Translate GEM -> COFFEE
result, = gem.impero_utils.preprocess_gem([result])
impero_c = gem.impero_utils.compile_gem([(return_variable, result)],
tensor_indices)
body = generate_coffee(impero_c, {}, parameters["precision"])
# Build kernel tuple
kernel_code = builder.construct_kernel(
"evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)
# Fill the code template
extruded = isinstance(cell, TensorProductCell)
#code = {
#"geometric_dimension": cell.geometric_dimension(),
#"extruded_arg": ", %s nlayers" % as_cstr(IntType) if extruded else "",
#"nlayers": ", f->n_layers" if extruded else "",
#"IntType": as_cstr(IntType),
#}
#evaluate_template_c = """static inline void wrap_evaluate(double *result, double *X, double *coords, %(IntType)s *coords_map, double *f, %(IntType)s *f_map%(extruded_arg)s, %(IntType)s cell);
#int evaluate(struct Function *f, double *x, double *result)
#{
#struct ReferenceCoords reference_coords;
#%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &reference_coords);
#if (cell == -1) {
#return -1;
#}
#if (!result) {
#return 0;
#}
#wrap_evaluate(result, reference_coords.X, f->coords, f->coords_map, f->f, f->f_map%(nlayers)s, cell);
#return 0;
#}
#"""
# return (evaluate_template_c % code) + kernel_code.gencode()
return kernel_code.gencode()
def compute_expression_ir(expression, analysis, parameters, visualise):
"""Compute IR for expression.
Parameters
----------
expression
Triple of (UFL expression, array of evaluation points, original UFL expression).
Note
----
Original UFL expression is needed to compute original positions of coefficients.
"""
logger.info("Computing uflacs representation of expression")
original_expression = expression[2]
points = expression[1]
expression = expression[0]
num_points = points.shape[0]
weights = numpy.array([1.0] * num_points)
cell = expression.ufl_domain().ufl_cell()
ir = {}
# Prepare dimensions of all unique element in expression,
# including elements for arguments, coefficients and coordinate mappings
ir["element_dimensions"] = {
ufl_element: create_element(ufl_element).space_dimension()
for ufl_element in analysis.unique_elements
}
# Extract dimensions for elements of arguments only
arguments = extract_arguments(expression)
argument_elements = tuple(f.ufl_element() for f in arguments)
argument_dimensions = [
ir["element_dimensions"][ufl_element]
for ufl_element in argument_elements
]
tensor_shape = argument_dimensions
ir["tensor_shape"] = tensor_shape
ir["expression_shape"] = list(expression.ufl_shape)
coefficients = extract_coefficients(expression)
coefficient_numbering = {}
for i, coeff in enumerate(coefficients):
coefficient_numbering[coeff] = i
# Add coefficient numbering to IR
ir["coefficient_numbering"] = coefficient_numbering
original_coefficient_positions = []
original_coefficients = extract_coefficients(original_expression)
for coeff in coefficients:
original_coefficient_positions.append(
original_coefficients.index(coeff))
ir["original_coefficient_positions"] = original_coefficient_positions
coefficient_elements = tuple(f.ufl_element() for f in coefficients)
offsets = {}
_offset = 0
for i, el in enumerate(coefficient_elements):
offsets[coefficients[i]] = _offset
_offset += ir["element_dimensions"][el]
# Copy offsets also into IR
ir["coefficient_offsets"] = offsets
ir["integral_type"] = "expression"
ir["entitytype"] = "cell"
# Build offsets for Constants
original_constant_offsets = {}
_offset = 0
for constant in extract_constants(expression):
original_constant_offsets[constant] = _offset
_offset += numpy.product(constant.ufl_shape, dtype=numpy.int)
ir["original_constant_offsets"] = original_constant_offsets
ir["points"] = points
integrands = {num_points: expression}
quadrature_rules = {num_points: (points, weights)}
uflacs_ir = build_uflacs_ir(cell, ir["integral_type"], ir["entitytype"],
integrands, tensor_shape, quadrature_rules,
parameters, visualise)
ir.update(uflacs_ir)
return ir
def make_interpolator(expr, V, subset, access):
assert isinstance(expr, ufl.classes.Expr)
arguments = extract_arguments(expr)
if len(arguments) == 0:
if isinstance(V, firedrake.Function):
f = V
V = f.function_space()
else:
f = firedrake.Function(V)
if access in {firedrake.MIN, firedrake.MAX}:
finfo = numpy.finfo(f.dat.dtype)
if access == firedrake.MIN:
val = firedrake.Constant(finfo.max)
else:
val = firedrake.Constant(finfo.min)
f.assign(val)
tensor = f.dat
elif len(arguments) == 1:
if isinstance(V, firedrake.Function):
raise ValueError(
"Cannot interpolate an expression with an argument into a Function"
)
argfs = arguments[0].function_space()
target_mesh = V.ufl_domain()
source_mesh = argfs.mesh()
argfs_map = argfs.cell_node_map()
if target_mesh is not source_mesh:
if not isinstance(target_mesh.topology,
firedrake.mesh.VertexOnlyMeshTopology):
raise NotImplementedError(
"Can only interpolate onto a Vertex Only Mesh")
if target_mesh.geometric_dimension(
) != source_mesh.geometric_dimension():
raise ValueError(
"Cannot interpolate onto a mesh of a different geometric dimension"
)
if not hasattr(target_mesh, "_parent_mesh"
) or target_mesh._parent_mesh is not source_mesh:
raise ValueError(
"Can only interpolate across meshes where the source mesh is the parent of the target"
)
if argfs_map:
# Since the par_loop is over the target mesh cells we need to
# compose a map that takes us from target mesh cells to the
# function space nodes on the source mesh. NOTE: argfs_map is
# allowed to be None when interpolating from a Real space, even
# in the trans-mesh case.
argfs_map = compose_map_and_cache(
target_mesh.cell_parent_cell_map, argfs_map)
sparsity = op2.Sparsity((V.dof_dset, argfs.dof_dset),
((V.cell_node_map(), argfs_map), ),
name="%s_%s_sparsity" % (V.name, argfs.name),
nest=False,
block_sparse=True)
tensor = op2.Mat(sparsity)
f = tensor
else:
raise ValueError("Cannot interpolate an expression with %d arguments" %
len(arguments))
# Make sure we have an expression of the right length i.e. a value for
# each component in the value shape of each function space
dims = [numpy.prod(fs.ufl_element().value_shape(), dtype=int) for fs in V]
loops = []
if numpy.prod(expr.ufl_shape, dtype=int) != sum(dims):
raise RuntimeError('Expression of length %d required, got length %d' %
(sum(dims), numpy.prod(expr.ufl_shape, dtype=int)))
if len(V) > 1:
raise NotImplementedError(
"UFL expressions for mixed functions are not yet supported.")
loops.extend(_interpolator(V, tensor, expr, subset, arguments, access))
def callable(loops, f):
for l in loops:
l()
return f
return partial(callable, loops, f), arguments
def compile_element(expression, coordinates, parameters=None):
"""Generates C code for point evaluations.
:arg expression: UFL expression
:arg coordinates: coordinate field
:arg parameters: form compiler parameters
:returns: C code as string
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# No arguments, please!
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
# Apply UFL preprocessing
expression = tsfc.ufl_utils.preprocess_expression(expression)
# Collect required coefficients
coefficient, = extract_coefficients(expression)
# Point evaluation of mixed coefficients not supported here
if type(coefficient.ufl_element()) == MixedElement:
raise NotImplementedError("Cannot point evaluate mixed elements yet!")
# Replace coordinates (if any)
domain = expression.ufl_domain()
assert coordinates.ufl_domain() == domain
# Initialise kernel builder
builder = firedrake_interface.KernelBuilderBase()
builder.domain_coordinate[domain] = coordinates
x_arg = builder._coefficient(coordinates, "x")
f_arg = builder._coefficient(coefficient, "f")
# TODO: restore this for expression evaluation!
# expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim,))
point_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('X', rank=(dim,)))
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_indices=(),
point_expr=point)
# TODO: restore this for expression evaluation!
# config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
tensor_indices = ()
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
return_variable = gem.Indexed(gem.Variable('R', expression.ufl_shape), tensor_indices)
result_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('R', rank=expression.ufl_shape))
result = gem.Indexed(result, tensor_indices)
else:
return_variable = gem.Indexed(gem.Variable('R', (1,)), (0,))
result_arg = ast.Decl(SCALAR_TYPE, ast.Symbol('R', rank=(1,)))
# Unroll
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
result, = gem.optimise.unroll_indexsum([result], predicate=predicate)
# Translate GEM -> COFFEE
result, = gem.impero_utils.preprocess_gem([result])
impero_c = gem.impero_utils.compile_gem([(return_variable, result)], tensor_indices)
body = generate_coffee(impero_c, {}, parameters["precision"])
# Build kernel tuple
kernel_code = builder.construct_kernel("evaluate_kernel", [result_arg, point_arg, x_arg, f_arg], body)
# Fill the code template
extruded = isinstance(cell, TensorProductCell)
code = {
"geometric_dimension": cell.geometric_dimension(),
"layers_arg": ", int const *__restrict__ layers" if extruded else "",
"layers": ", layers" if extruded else "",
"IntType": as_cstr(IntType),
}
# if maps are the same, only need to pass one of them
if coordinates.cell_node_map() == coefficient.cell_node_map():
code["wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map" % code
code["map_args"] = "f->coords_map"
else:
code["wrapper_map_args"] = "%(IntType)s const *__restrict__ coords_map, %(IntType)s const *__restrict__ f_map" % code
code["map_args"] = "f->coords_map, f->f_map"
evaluate_template_c = """
static inline void wrap_evaluate(double* const result, double* const X, int const start, int const end%(layers_arg)s,
double const *__restrict__ coords, double const *__restrict__ f, %(wrapper_map_args)s);
int evaluate(struct Function *f, double *x, double *result)
{
struct ReferenceCoords reference_coords;
%(IntType)s cell = locate_cell(f, x, %(geometric_dimension)d, &to_reference_coords, &to_reference_coords_xtr, &reference_coords);
if (cell == -1) {
return -1;
}
if (!result) {
return 0;
}
int layers[2] = {0, 0};
if (f->extruded != 0) {
int nlayers = f->n_layers;
layers[1] = cell %% nlayers + 2;
cell = cell / nlayers;
}
wrap_evaluate(result, reference_coords.X, cell, cell+1%(layers)s, f->coords, f->f, %(map_args)s);
return 0;
}
"""
return (evaluate_template_c % code) + kernel_code.gencode()
def compile_ufl_kernel(expression, to_pts, to_element, fs):
import collections
from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks
from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
from ufl.algorithms.apply_derivatives import apply_derivatives
from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering
from ufl.algorithms import extract_arguments, extract_coefficients
from gem import gem, impero_utils
from tsfc import fem, ufl_utils
from tsfc.coffee import generate as generate_coffee
from tsfc.kernel_interface import (KernelBuilderBase,
needs_cell_orientations,
cell_orientations_coffee_arg)
# Imitate the compute_form_data processing pipeline
#
# Unfortunately, we cannot call compute_form_data here, since
# we only have an expression, not a form
expression = apply_algebra_lowering(expression)
expression = apply_derivatives(expression)
expression = apply_function_pullbacks(expression)
expression = apply_geometry_lowering(expression)
expression = apply_derivatives(expression)
expression = apply_geometry_lowering(expression)
expression = apply_derivatives(expression)
# Replace coordinates (if any)
if expression.ufl_domain():
assert fs.mesh() == expression.ufl_domain()
expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates)
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
builder = KernelBuilderBase()
args = []
coefficients = extract_coefficients(expression)
for i, coefficient in enumerate(coefficients):
args.append(builder.coefficient(coefficient, "w_%d" % i))
point_index = gem.Index(name='p')
ir = fem.compile_ufl(expression,
cell=fs.mesh().ufl_cell(),
points=to_pts,
point_index=point_index,
coefficient_mapper=builder.coefficient_mapper)
assert len(ir) == 1
# Deal with non-scalar expressions
tensor_indices = ()
if fs.shape:
tensor_indices = tuple(gem.Index() for s in fs.shape)
ir = [gem.Indexed(ir[0], tensor_indices)]
# Build kernel body
return_var = gem.Variable('A', (len(to_pts),) + fs.shape)
return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices)
impero_c = impero_utils.compile_gem([return_expr], ir, [point_index])
body = generate_coffee(impero_c, index_names={point_index: 'p'})
oriented = needs_cell_orientations(ir)
if oriented:
args.insert(0, cell_orientations_coffee_arg)
# Build kernel
args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape)))
kernel_code = builder.construct_kernel("expression_kernel", args, body)
return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
def compile_ufl_kernel(expression, to_pts, to_element, fs):
import collections
from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks
from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
from ufl.algorithms.apply_derivatives import apply_derivatives
from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering
from ufl.algorithms import extract_arguments, extract_coefficients
from gem import gem, impero_utils
from tsfc import fem, ufl_utils
from tsfc.coffee import generate as generate_coffee
from tsfc.kernel_interface import (KernelBuilderBase,
needs_cell_orientations,
cell_orientations_coffee_arg)
# Imitate the compute_form_data processing pipeline
#
# Unfortunately, we cannot call compute_form_data here, since
# we only have an expression, not a form
expression = apply_algebra_lowering(expression)
expression = apply_derivatives(expression)
expression = apply_function_pullbacks(expression)
expression = apply_geometry_lowering(expression)
expression = apply_derivatives(expression)
expression = apply_geometry_lowering(expression)
expression = apply_derivatives(expression)
# Replace coordinates (if any)
if expression.ufl_domain():
assert fs.mesh() == expression.ufl_domain()
expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates)
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
builder = KernelBuilderBase()
args = []
coefficients = extract_coefficients(expression)
for i, coefficient in enumerate(coefficients):
args.append(builder.coefficient(coefficient, "w_%d" % i))
point_index = gem.Index(name='p')
ir = fem.process('cell', fs.mesh().ufl_cell(), to_pts, None,
point_index, (), expression,
builder.coefficient_mapper,
collections.defaultdict(gem.Index))
assert len(ir) == 1
# Deal with non-scalar expressions
tensor_indices = ()
if fs.shape:
tensor_indices = tuple(gem.Index() for s in fs.shape)
ir = [gem.Indexed(ir[0], tensor_indices)]
# Build kernel body
return_var = gem.Variable('A', (len(to_pts),) + fs.shape)
return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices)
impero_c = impero_utils.compile_gem([return_expr], ir, [point_index])
body = generate_coffee(impero_c, index_names={point_index: 'p'})
oriented = needs_cell_orientations(ir)
if oriented:
args.insert(0, cell_orientations_coffee_arg)
# Build kernel
args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape)))
kernel_code = builder.construct_kernel("expression_kernel", args, body)
return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
def test_extract_arguments_vs_fixture(arguments, forms):
assert arguments == tuple(extract_arguments(forms[0]))
assert tuple(arguments[:1]) == tuple(extract_arguments(forms[1]))
def test_extract_arguments_vs_fixture(arguments, forms):
assert arguments == tuple(extract_arguments(forms[0]))
assert tuple(arguments[:1]) == tuple(extract_arguments(forms[1]))
def compile_expression_at_points(expression,
points,
coordinates,
interface=None,
parameters=None,
coffee=True):
"""Compiles a UFL expression to be evaluated at compile-time known
reference points. Useful for interpolating UFL expressions onto
function spaces with only point evaluation nodes.
:arg expression: UFL expression
:arg points: reference coordinates of the evaluation points
:arg coordinates: the coordinate function
:arg interface: backend module for the kernel interface
:arg parameters: parameters object
:arg coffee: compile coffee kernel instead of loopy kernel
"""
import coffee.base as ast
import loopy as lp
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
# Determine whether in complex mode
complex_mode = is_complex(parameters["scalar_type"])
# Apply UFL preprocessing
expression = ufl_utils.preprocess_expression(expression,
complex_mode=complex_mode)
# Initialise kernel builder
if interface is None:
if coffee:
import tsfc.kernel_interface.firedrake as firedrake_interface_coffee
interface = firedrake_interface_coffee.ExpressionKernelBuilder
else:
# Delayed import, loopy is a runtime dependency
import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy
interface = firedrake_interface_loopy.ExpressionKernelBuilder
builder = interface(parameters["scalar_type"])
arguments = extract_arguments(expression)
argument_multiindices = tuple(
builder.create_element(arg.ufl_element()).get_indices()
for arg in arguments)
# Replace coordinates (if any)
domain = expression.ufl_domain()
if domain:
assert coordinates.ufl_domain() == domain
builder.domain_coordinate[domain] = coordinates
builder.set_cell_sizes(domain)
# Collect required coefficients
coefficients = extract_coefficients(expression)
if has_type(expression, GeometricQuantity) or any(
fem.needs_coordinate_mapping(c.ufl_element())
for c in coefficients):
coefficients = [coordinates] + coefficients
builder.set_coefficients(coefficients)
# Split mixed coefficients
expression = ufl_utils.split_coefficients(expression,
builder.coefficient_split)
# Translate to GEM
point_set = PointSet(points)
config = dict(interface=builder,
ufl_cell=coordinates.ufl_domain().ufl_cell(),
precision=parameters["precision"],
point_set=point_set,
argument_multiindices=argument_multiindices)
ir, = fem.compile_ufl(expression, point_sum=False, **config)
# Deal with non-scalar expressions
value_shape = ir.shape
tensor_indices = tuple(gem.Index() for s in value_shape)
if value_shape:
ir = gem.Indexed(ir, tensor_indices)
# Build kernel body
return_indices = point_set.indices + tensor_indices + tuple(
chain(*argument_multiindices))
return_shape = tuple(i.extent for i in return_indices)
return_var = gem.Variable('A', return_shape)
if coffee:
return_arg = ast.Decl(parameters["scalar_type"],
ast.Symbol('A', rank=return_shape))
else:
return_arg = lp.GlobalArg("A",
dtype=parameters["scalar_type"],
shape=return_shape)
return_expr = gem.Indexed(return_var, return_indices)
ir, = impero_utils.preprocess_gem([ir])
impero_c = impero_utils.compile_gem([(return_expr, ir)], return_indices)
point_index, = point_set.indices
# Handle kernel interface requirements
builder.register_requirements([ir])
# Build kernel tuple
return builder.construct_kernel(return_arg, impero_c,
parameters["precision"],
{point_index: 'p'})
