def abs(self, o, expr):
if o.ufl_shape:
indices = tuple(Index() for i in range(len(o.ufl_shape)))
return ComponentTensor(MathFunction('abs', Indexed(expr, indices)),
indices)
else:
return MathFunction('abs', expr)
def sum(self, o, *ops):
if o.ufl_shape:
indices = tuple(Index() for i in range(len(o.ufl_shape)))
return ComponentTensor(Sum(*[Indexed(op, indices) for op in ops]),
indices)
else:
return Sum(*ops)
def conditional(self, o, condition, then, else_):
assert condition.shape == ()
if o.ufl_shape:
indices = tuple(Index() for i in range(len(o.ufl_shape)))
return ComponentTensor(
Conditional(condition, Indexed(then, indices),
Indexed(else_, indices)), indices)
else:
return Conditional(condition, then, else_)
def index_sum(self, o, summand, indices):
# ufl.IndexSum technically has a MultiIndex, but it must have
# exactly one index in it.
index, = indices
if o.ufl_shape:
indices = tuple(Index() for i in range(len(o.ufl_shape)))
return ComponentTensor(IndexSum(Indexed(summand, indices), (index,)), indices)
else:
return IndexSum(summand, (index,))
def _slate2gem_diagonal(expr, self):
if not self.matfree:
A, = map(self, expr.children)
assert A.shape[0] == A.shape[1]
i, j = (Index(extent=s) for s in A.shape)
return ComponentTensor(Product(Indexed(A, (i, i)), Delta(i, j)),
(i, j))
else:
raise NotImplementedError("Diagonals on Slate expressions are \
not implemented in a matrix-free manner yet."
)
def _slate2gem_inverse(expr, self):
tensor, = expr.children
if expr.diagonal:
# optimise inverse on diagonal tensor by translating to
# matrix which contains the reciprocal values of the diagonal tensor
A, = map(self, expr.children)
i, j = (Index(extent=s) for s in A.shape)
return ComponentTensor(
Product(Division(Literal(1), Indexed(A, (i, i))), Delta(i, j)),
(i, j))
else:
return Inverse(self(tensor))
