def _slate2gem_mul(expr, self):
A, B = map(self, expr.children)
*i, k = tuple(make_indices(len(A.shape)))
_, *j = tuple(make_indices(len(B.shape)))
ABikj = Product(Indexed(A, tuple(i + [k])),
Indexed(B, tuple([k] + j)))
return ComponentTensor(IndexSum(ABikj, (k, )), tuple(i + j))
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 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 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))
def indexed(self, o, aggregate, index):
return Indexed(aggregate, index)
def _slate2gem_add(expr, self):
A, B = map(self, expr.children)
indices = tuple(make_indices(len(A.shape)))
return ComponentTensor(Sum(Indexed(A, indices), Indexed(B, indices)),
indices)
def _slate2gem_negative(expr, self):
child, = map(self, expr.children)
indices = tuple(make_indices(len(child.shape)))
return ComponentTensor(Product(Literal(-1), Indexed(child, indices)),
indices)
def _slate2gem_transpose(expr, self):
child, = map(self, expr.children)
indices = tuple(make_indices(len(child.shape)))
return ComponentTensor(Indexed(child, indices), tuple(indices[::-1]))
def _slate2gem_reciprocal(expr, self):
child, = map(self, expr.children)
indices = tuple(make_indices(len(child.shape)))
return ComponentTensor(Division(Literal(1.), Indexed(child, indices)),
indices)