def _expression_mathfunction(expr, ctx):
if expr.name.startswith('cyl_bessel_'):
# Bessel functions
if is_complex(ctx.scalar_type):
raise NotImplementedError("Bessel functions for complex numbers: "
"missing implementation")
nu, arg = expr.children
nu_ = expression(nu, ctx)
arg_ = expression(arg, ctx)
# Modified Bessel functions (C++ only)
#
# These mappings work for FEniCS only, and fail with Firedrake
# since no Boost available.
if expr.name in {'cyl_bessel_i', 'cyl_bessel_k'}:
name = 'boost::math::' + expr.name
return p.Variable(name)(nu_, arg_)
else:
# cyl_bessel_{jy} -> {jy}
name = expr.name[-1:]
if nu == gem.Zero():
return p.Variable(f"{name}0")(arg_)
elif nu == gem.one:
return p.Variable(f"{name}1")(arg_)
else:
return p.Variable(f"{name}n")(nu_, arg_)
else:
if expr.name == "ln":
name = "log"
else:
name = expr.name
# Not all mathfunctions apply to complex numbers, but this
# will be picked up in loopy. This way we allow erf(real(...))
# in complex mode (say).
return p.Variable(name)(*(expression(c, ctx) for c in expr.children))
def compile_form(form,
prefix="form",
parameters=None,
interface=None,
coffee=True,
diagonal=False):
"""Compiles a UFL form into a set of assembly kernels.
:arg form: UFL form
:arg prefix: kernel name will start with this string
:arg parameters: parameters object
:arg coffee: compile coffee kernel instead of loopy kernel
:arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor?
:returns: list of kernels
"""
cpu_time = time.time()
assert isinstance(form, Form)
# Determine whether in complex mode:
# complex nodes would break the refactoriser.
complex_mode = parameters and is_complex(parameters.get("scalar_type"))
if complex_mode:
logger.warning("Disabling whole expression optimisations"
" in GEM for supporting complex mode.")
parameters = parameters.copy()
parameters["mode"] = 'vanilla'
fd = ufl_utils.compute_form_data(form, complex_mode=complex_mode)
logger.info(GREEN % "compute_form_data finished in %g seconds.",
time.time() - cpu_time)
kernels = []
for integral_data in fd.integral_data:
start = time.time()
kernel = compile_integral(integral_data,
fd,
prefix,
parameters,
interface=interface,
coffee=coffee,
diagonal=diagonal)
if kernel is not None:
kernels.append(kernel)
logger.info(GREEN % "compile_integral finished in %g seconds.",
time.time() - start)
logger.info(GREEN % "TSFC finished in %g seconds.", time.time() - cpu_time)
return kernels
def _expression_mathfunction(expr, ctx):
from tsfc.coffee import math_table
math_table = math_table.copy()
math_table['abs'] = ('abs', 'cabs')
complex_mode = int(is_complex(ctx.scalar_type))
# Bessel functions
if expr.name.startswith('cyl_bessel_'):
if complex_mode:
msg = "Bessel functions for complex numbers: missing implementation"
raise NotImplementedError(msg)
nu, arg = expr.children
nu_thunk = lambda: expression(nu, ctx)
arg_loopy = expression(arg, ctx)
if expr.name == 'cyl_bessel_j':
if nu == gem.Zero():
return p.Variable("j0")(arg_loopy)
elif nu == gem.one:
return p.Variable("j1")(arg_loopy)
else:
return p.Variable("jn")(nu_thunk(), arg_loopy)
if expr.name == 'cyl_bessel_y':
if nu == gem.Zero():
return p.Variable("y0")(arg_loopy)
elif nu == gem.one:
return p.Variable("y1")(arg_loopy)
else:
return p.Variable("yn")(nu_thunk(), arg_loopy)
# Modified Bessel functions (C++ only)
#
# These mappings work for FEniCS only, and fail with Firedrake
# since no Boost available.
if expr.name in ['cyl_bessel_i', 'cyl_bessel_k']:
name = 'boost::math::' + expr.name
return p.Variable(name)(nu_thunk(), arg_loopy)
assert False, "Unknown Bessel function: {}".format(expr.name)
# Other math functions
name = math_table[expr.name][complex_mode]
if name is None:
raise RuntimeError("{} not supported in complex mode".format(expr.name))
return p.Variable(name)(*[expression(c, ctx) for c in expr.children])
def _expression_mathfunction(expr, parameters):
complex_mode = int(is_complex(parameters.scalar_type))
# Bessel functions
if expr.name.startswith('cyl_bessel_'):
if complex_mode:
msg = "Bessel functions for complex numbers: missing implementation"
raise NotImplementedError(msg)
nu, arg = expr.children
nu_thunk = lambda: expression(nu, parameters)
arg_coffee = expression(arg, parameters)
if expr.name == 'cyl_bessel_j':
if nu == gem.Zero():
return coffee.FunCall('j0', arg_coffee)
elif nu == gem.one:
return coffee.FunCall('j1', arg_coffee)
else:
return coffee.FunCall('jn', nu_thunk(), arg_coffee)
if expr.name == 'cyl_bessel_y':
if nu == gem.Zero():
return coffee.FunCall('y0', arg_coffee)
elif nu == gem.one:
return coffee.FunCall('y1', arg_coffee)
else:
return coffee.FunCall('yn', nu_thunk(), arg_coffee)
# Modified Bessel functions (C++ only)
#
# These mappings work for FEniCS only, and fail with Firedrake
# since no Boost available.
if expr.name in ['cyl_bessel_i', 'cyl_bessel_k']:
name = 'boost::math::' + expr.name
return coffee.FunCall(name, nu_thunk(), arg_coffee)
assert False, "Unknown Bessel function: {}".format(expr.name)
# Other math functions
name = math_table[expr.name][complex_mode]
if name is None:
raise RuntimeError("{} not supported in complex mode".format(
expr.name))
return coffee.FunCall(name,
*[expression(c, parameters) for c in expr.children])
def _expression_power(expr, parameters):
base, exponent = expr.children
complex_mode = int(is_complex(parameters.scalar_type))
return coffee.FunCall(math_table['power'][complex_mode],
expression(base, parameters),
expression(exponent, parameters))
def complex_mode(self):
return is_complex(self.scalar_type)
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'})
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!")
# 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
builder = firedrake_interface.ExpressionKernelBuilder(
parameters["scalar_type"])
# 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)
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(parameters["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"], parameters["scalar_type"])
# Handle cell orientations
if builder.needs_cell_orientations([ir]):
builder.require_cell_orientations()
# Handle cell sizes (physically mapped elements)
if builder.needs_cell_sizes([ir]):
builder.require_cell_sizes()
# Build kernel tuple
return builder.construct_kernel(return_arg, body)
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)
