def translate_coefficient(terminal, mt, ctx):
vec = ctx.coefficient(terminal, mt.restriction)
if terminal.ufl_element().family() == 'Real':
assert mt.local_derivatives == 0
return vec
element = ctx.create_element(terminal.ufl_element())
# Collect FInAT tabulation for all entities
per_derivative = collections.defaultdict(list)
for entity_id in ctx.entity_ids:
finat_dict = ctx.basis_evaluation(element, mt.local_derivatives, entity_id)
for alpha, table in finat_dict.items():
# Filter out irrelevant derivatives
if sum(alpha) == mt.local_derivatives:
# A numerical hack that FFC used to apply on FIAT
# tables still lives on after ditching FFC and
# switching to FInAT.
table = ffc_rounding(table, ctx.epsilon)
per_derivative[alpha].append(table)
# Merge entity tabulations for each derivative
if len(ctx.entity_ids) == 1:
def take_singleton(xs):
x, = xs # asserts singleton
return x
per_derivative = {alpha: take_singleton(tables)
for alpha, tables in per_derivative.items()}
else:
f = ctx.entity_number(mt.restriction)
per_derivative = {alpha: gem.select_expression(tables, f)
for alpha, tables in per_derivative.items()}
# Coefficient evaluation
ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices())
beta = ctx.index_cache[terminal.ufl_element()]
zeta = element.get_value_indices()
vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)])
value_dict = {}
for alpha, table in per_derivative.items():
table_qi = gem.Indexed(table, beta + zeta)
summands = []
for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache):
value = gem.IndexSum(gem.Product(expr, var), var.index_ordering())
summands.append(gem.optimise.contraction(value))
optimised_value = gem.optimise.make_sum(summands)
value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta)
# Change from FIAT to UFL arrangement
result = fiat_to_ufl(value_dict, mt.local_derivatives)
assert result.shape == mt.expr.ufl_shape
assert set(result.free_indices) <= set(ctx.point_indices)
# Detect Jacobian of affine cells
if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2
for node in traversal((result,))
if isinstance(node, gem.Literal)):
result = gem.optimise.aggressive_unroll(result)
return result
def entity_selector(self, callback, restriction):
"""Selects code for the correct entity at run-time. Callback
generates code for a specified entity.
This function passes ``callback`` the entity number.
:arg callback: A function to be called with an entity number
that generates code for that entity.
:arg restriction: Restriction of the modified terminal, used
for entity selection.
"""
if len(self.entity_ids) == 1:
return callback(self.entity_ids[0])
else:
f = self.entity_number(restriction)
return gem.select_expression(list(map(callback, self.entity_ids)), f)
def entity_selector(self, callback, restriction):
"""Selects code for the correct entity at run-time. Callback
generates code for a specified entity.
This function passes ``callback`` the entity number.
:arg callback: A function to be called with an entity number
that generates code for that entity.
:arg restriction: Restriction of the modified terminal, used
for entity selection.
"""
if len(self.entity_ids) == 1:
return callback(self.entity_ids[0])
else:
f = self.entity_number(restriction)
return gem.select_expression(list(map(callback, self.entity_ids)), f)