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 sum(*args):
"""Sum of multiple :py:class:`MonomialSum`s"""
result = MonomialSum()
for arg in args:
assert isinstance(arg, MonomialSum)
# Optimised implementation: no need to decompose and
# reconstruct key.
for key, rest in arg.monomials.items():
result.monomials[key] = Sum(result.monomials[key], rest)
for key, value in arg.ordering.items():
result.ordering.setdefault(key, value)
return result
def add(self, sum_indices, atomics, rest):
"""Updates the :py:class:`MonomialSum` adding a new monomial."""
sum_indices = tuple(sum_indices)
sum_indices_set = frozenset(sum_indices)
# Sum indices cannot have duplicates
assert len(sum_indices) == len(sum_indices_set)
atomics = tuple(atomics)
atomics_set = frozenset(Counter(atomics).items())
assert isinstance(rest, Node)
key = (sum_indices_set, atomics_set)
self.monomials[key] = Sum(self.monomials[key], rest)
self.ordering.setdefault(key, (sum_indices, atomics))
def test_loop_optimise():
I = 20
J = K = 10
i = Index('i', I)
j = Index('j', J)
k = Index('k', K)
A1 = Variable('a1', (I,))
A2 = Variable('a2', (I,))
A3 = Variable('a3', (I,))
A1i = Indexed(A1, (i,))
A2i = Indexed(A2, (i,))
A3i = Indexed(A3, (i,))
B = Variable('b', (J,))
C = Variable('c', (J,))
Bj = Indexed(B, (j,))
Cj = Indexed(C, (j,))
E = Variable('e', (K,))
F = Variable('f', (K,))
G = Variable('g', (K,))
Ek = Indexed(E, (k,))
Fk = Indexed(F, (k,))
Gk = Indexed(G, (k,))
Z = Variable('z', ())
# Bj*Ek + Bj*Fk => (Ek + Fk)*Bj
expr = Sum(Product(Bj, Ek), Product(Bj, Fk))
result, = optimise_expressions([expr], (j, k))
expected = Product(Sum(Ek, Fk), Bj)
assert result == expected
# Bj*Ek + Bj*Fk + Bj*Gk + Cj*Ek + Cj*Fk =>
# (Ek + Fk + Gk)*Bj + (Ek+Fk)*Cj
expr = Sum(Sum(Sum(Sum(Product(Bj, Ek), Product(Bj, Fk)), Product(Bj, Gk)),
Product(Cj, Ek)), Product(Cj, Fk))
result, = optimise_expressions([expr], (j, k))
expected = Sum(Product(Sum(Sum(Ek, Fk), Gk), Bj), Product(Sum(Ek, Fk), Cj))
assert result == expected
# Z*A1i*Bj*Ek + Z*A2i*Bj*Ek + A3i*Bj*Ek + Z*A1i*Bj*Fk =>
# Bj*(Ek*(Z*A1i + Z*A2i) + A3i) + Z*A1i*Fk)
expr = Sum(Sum(Sum(Product(Z, Product(A1i, Product(Bj, Ek))),
Product(Z, Product(A2i, Product(Bj, Ek)))),
Product(A3i, Product(Bj, Ek))),
Product(Z, Product(A1i, Product(Bj, Fk))))
result, = optimise_expressions([expr], (j, k))
expected = Product(Sum(Product(Ek, Sum(Sum(Product(Z, A1i), Product(Z, A2i)), A3i)),
Product(Fk, Product(Z, A1i))), Bj)
assert result == expected