def simplify_abs(expression, complex_mode):
"""Simplify absolute values in a UFL expression. Its primary
purpose is to "neutralise" CellOrientation nodes that are
surrounded by absolute values and thus not at all necessary."""
mapper = MemoizerArg(_simplify_abs)
mapper.complex_mode = complex_mode
return mapper(expression, False)
def delta_elimination(sum_indices, factors):
"""IndexSum-Delta cancellation.
:arg sum_indices: free indices for contractions
:arg factors: product factors
:returns: optimised (sum_indices, factors)
"""
sum_indices = list(sum_indices) # copy for modification
delta_queue = [(f, index)
for f in factors if isinstance(f, Delta)
for index in (f.i, f.j) if index in sum_indices]
while delta_queue:
delta, from_ = delta_queue[0]
to_, = list({delta.i, delta.j} - {from_})
sum_indices.remove(from_)
mapper = MemoizerArg(filtered_replace_indices)
factors = [mapper(e, ((from_, to_),)) for e in factors]
delta_queue = [(f, index)
for f in factors if isinstance(f, Delta)
for index in (f.i, f.j) if index in sum_indices]
# Drop ones
return sum_indices, [e for e in factors if e != one]
def substitute(expression, from_, to_):
if from_ not in expression.free_indices:
return expression
elif isinstance(expression, Delta):
mapper = MemoizerArg(filtered_replace_indices)
return mapper(expression, ((from_, to_), ))
else:
return Indexed(ComponentTensor(expression, (from_, )), (to_, ))
def push_mul(tensor, options):
"""Executes a Slate compiler optimisation pass.
The optimisation is achieved by pushing coefficients from
the outside to the inside of an expression.
The optimisation pass essentially changes the order of operations
in the expressions so that only matrix-vector products are executed.
:arg tensor: A (potentially unoptimised) Slate expression.
:arg options: Optimisation pass options,
e.g. if the multiplication should be replaced by an action.
Returns: An optimised Slate expression,
where only matrix-vector products are executed whereever possible.
"""
from gem.node import MemoizerArg
mapper = MemoizerArg(_push_mul)
mapper.action = options["replace_mul"]
a = mapper(tensor, ActionBag(None, 1))
return a
def push_diag(expression):
"""Executes a Slate compiler optimisation pass.
The optimisation is achieved by pushing DiagonalTensor from the outside to the inside of an expression.
:arg expression: A (potentially unoptimised) Slate expression.
Returns: An optimised Slate expression, where DiagonalTensors are sitting
on terminal tensors whereever possible.
"""
mapper = MemoizerArg(_push_diag)
return mapper(expression, False)
def __init__(self, lvalue, rvalue):
"""
:arg lvalue: The coefficient to assign into.
:arg rvalue: The pointwise expression.
"""
if not isinstance(lvalue, ufl.Coefficient):
raise ValueError("lvalue for pointwise assignment must be a coefficient")
self.lvalue = lvalue
self.rvalue = ufl.as_ufl(rvalue)
n = len(self.lvalue.function_space())
if n > 1:
self.splitter = MemoizerArg(_split)
self.splitter.n = n
def push_block(expression):
"""Executes a Slate compiler optimisation pass.
The optimisation is achieved by pushing blocks from the outside to the inside of an expression.
Without the optimisation the local TSFC kernels are assembled first
and then the result of the assembly kernel gets indexed in the Slate kernel
(and further linear algebra operations maybe done on it).
The optimisation pass essentially changes the order of assembly and indexing.
:arg expression: A (potentially unoptimised) Slate expression.
Returns: An optimised Slate expression, where Blocks are terminal whereever possible.
"""
mapper = MemoizerArg(_push_block)
ret = mapper(expression, ())
return ret
def remove_componenttensors(expressions):
"""Removes all ComponentTensors in multi-root expression DAG."""
mapper = MemoizerArg(filtered_replace_indices)
return [mapper(expression, ()) for expression in expressions]
def simplify_abs(expression):
"""Simplify absolute values in a UFL expression. Its primary
purpose is to "neutralise" CellOrientation nodes that are
surrounded by absolute values and thus not at all necessary."""
return MemoizerArg(_simplify_abs)(expression, False)