def to_reference_coordinates(ufl_coordinate_element):
# Set up UFL form
cell = ufl_coordinate_element.cell()
domain = ufl.Mesh(ufl_coordinate_element)
K = ufl.JacobianInverse(domain)
x = ufl.SpatialCoordinate(domain)
x0_element = ufl.VectorElement("Real", cell, 0)
x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element))
expr = ufl.dot(K, x - x0)
# Translation to GEM
C = ufl_utils.coordinate_coefficient(domain)
expr = ufl_utils.preprocess_expression(expr)
expr = ufl_utils.replace_coordinates(expr, C)
expr = ufl_utils.simplify_abs(expr)
builder = firedrake_interface.KernelBuilderBase()
builder._coefficient(C, "C")
builder._coefficient(x0, "x0")
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
context = tsfc.fem.GemPointContext(
interface=builder,
ufl_cell=cell,
precision=parameters["precision"],
point_indices=(),
point_expr=point,
)
translator = tsfc.fem.Translator(context)
ir = map_expr_dag(translator, expr)
# Unroll result
ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)]
# Unroll IndexSums
max_extent = parameters["unroll_indexsum"]
if max_extent:
def predicate(index):
return index.extent <= max_extent
ir = gem.optimise.unroll_indexsum(ir, predicate=predicate)
# Translate to COFFEE
ir = impero_utils.preprocess_gem(ir)
return_variable = gem.Variable('dX', (dim, ))
assignments = [(gem.Indexed(return_variable, (i, )), e)
for i, e in enumerate(ir)]
impero_c = impero_utils.compile_gem(assignments, ())
body = tsfc.coffee.generate(impero_c, {}, parameters["precision"])
body.open_scope = False
return body
def compile_element(expression,
dual_space=None,
parameters=None,
name="evaluate"):
"""Generate code for point evaluations.
:arg expression: A UFL expression (may contain up to one coefficient, or one argument)
:arg dual_space: if the expression has an argument, should we also distribute residual data?
:returns: Some coffee AST
"""
if parameters is None:
parameters = default_parameters()
else:
_ = default_parameters()
_.update(parameters)
parameters = _
expression = tsfc.ufl_utils.preprocess_expression(expression)
# # Collect required coefficients
try:
arg, = extract_coefficients(expression)
argument_multiindices = ()
coefficient = True
if expression.ufl_shape:
tensor_indices = tuple(gem.Index() for s in expression.ufl_shape)
else:
tensor_indices = ()
except ValueError:
arg, = extract_arguments(expression)
finat_elem = create_element(arg.ufl_element())
argument_multiindices = (finat_elem.get_indices(), )
argument_multiindex, = argument_multiindices
value_shape = finat_elem.value_shape
if value_shape:
tensor_indices = argument_multiindex[-len(value_shape):]
else:
tensor_indices = ()
coefficient = False
# Replace coordinates (if any)
builder = firedrake_interface.KernelBuilderBase(scalar_type=ScalarType_c)
domain = expression.ufl_domain()
# Translate to GEM
cell = domain.ufl_cell()
dim = cell.topological_dimension()
point = gem.Variable('X', (dim, ))
point_arg = ast.Decl(ScalarType_c, ast.Symbol('X', rank=(dim, )))
config = dict(interface=builder,
ufl_cell=cell,
precision=parameters["precision"],
point_indices=(),
point_expr=point,
argument_multiindices=argument_multiindices)
context = tsfc.fem.GemPointContext(**config)
# Abs-simplification
expression = tsfc.ufl_utils.simplify_abs(expression)
# Translate UFL -> GEM
if coefficient:
assert dual_space is None
f_arg = [builder._coefficient(arg, "f")]
else:
f_arg = []
translator = tsfc.fem.Translator(context)
result, = map_expr_dags(translator, [expression])
b_arg = []
if coefficient:
if expression.ufl_shape:
return_variable = gem.Indexed(
gem.Variable('R', expression.ufl_shape), tensor_indices)
result_arg = ast.Decl(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(ScalarType_c, ast.Symbol('R', rank=(1, )))
else:
return_variable = gem.Indexed(
gem.Variable('R', finat_elem.index_shape), argument_multiindex)
result = gem.Indexed(result, tensor_indices)
if dual_space:
elem = create_element(dual_space.ufl_element())
if elem.value_shape:
var = gem.Indexed(gem.Variable("b", elem.value_shape),
tensor_indices)
b_arg = [
ast.Decl(ScalarType_c,
ast.Symbol("b", rank=elem.value_shape))
]
else:
var = gem.Indexed(gem.Variable("b", (1, )), (0, ))
b_arg = [ast.Decl(ScalarType_c, ast.Symbol("b", rank=(1, )))]
result = gem.Product(result, var)
result_arg = ast.Decl(ScalarType_c,
ast.Symbol('R', rank=finat_elem.index_shape))
# 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"], ScalarType_c)
# Build kernel tuple
kernel_code = builder.construct_kernel(
"pyop2_kernel_" + name, [result_arg] + b_arg + f_arg + [point_arg],
body)
return kernel_code
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_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()
