def _expand_conditional_conditional(node, self):
if self.predicate(node):
condition, then, else_ = map(self, node.children)
return Sum(Product(Conditional(condition, one, Zero()), then),
Product(Conditional(condition, Zero(), one), else_))
else:
return reuse_if_untouched(node, self)
def _(node, self):
if node.index.extent <= self.max_extent:
# Unrolling
summand = self(node.children[0])
return reduce(
Sum, (Indexed(ComponentTensor(summand, (node.index,)), (i,)) for i in range(node.index.extent)), Zero()
)
else:
return reuse_if_untouched(node, self)
def _replace_node(node, self):
"""Replace subexpressions using a given mapping.
:param node: root of expression
:param self: function for recursive calls
"""
assert isinstance(node, Node)
if self.cut(node):
return node
try:
return self.mapping[node]
except KeyError:
return reuse_if_untouched(node, self)
def _replace_node(node, self):
"""Replace subexpressions using a given mapping.
:param node: root of expression
:param self: function for recursive calls
"""
assert isinstance(node, Node)
if self.cut(node):
return node
try:
return self.mapping[node]
except KeyError:
return reuse_if_untouched(node, self)
def _(node, self):
unroll = tuple(filter(self.predicate, node.multiindex))
if unroll:
# Unrolling
summand = self(node.children[0])
shape = tuple(index.extent for index in unroll)
unrolled = reduce(Sum,
(Indexed(ComponentTensor(summand, unroll), alpha)
for alpha in numpy.ndindex(shape)), Zero())
return IndexSum(
unrolled,
tuple(index for index in node.multiindex if index not in unroll))
else:
return reuse_if_untouched(node, self)