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 test_conditional_zero_folding():
b = Variable("B", ())
a = Variable("A", (3, ))
i = Index()
expr = Conditional(LogicalAnd(b, b), Product(Indexed(a, (i, )), Zero()),
Zero())
assert expr == Zero()
def test_replace_div():
i = Index()
A = Variable('A', ())
B = Variable('B', (6,))
Bi = Indexed(B, (i,))
d = Division(Bi, A)
result, = replace_division([d])
expected = Product(Bi, Division(Literal(1.0), A))
assert result == expected
def product(*args, **kwargs):
"""Product of multiple :py:class:`MonomialSum`s"""
rename_map = kwargs.pop('rename_map', None)
if rename_map is None:
rename_map = make_rename_map()
if kwargs:
raise ValueError("Unrecognised keyword argument: " + kwargs.pop())
result = MonomialSum()
for monomials in product(*args):
renamer = make_renamer(rename_map)
sum_indices = []
atomics = []
rest = one
for s, a, r in monomials:
s_, applier = renamer(s)
sum_indices.extend(s_)
atomics.extend(map(applier, a))
rest = Product(applier(r), rest)
result.add(sum_indices, atomics, rest)
return result
def make_expression(i, j):
A = Variable('A', (6, ))
s = IndexSum(Indexed(A, (j, )), (j, ))
return Product(Indexed(A, (i, )), s)
def _replace_division_division(node, self):
a, b = node.children
return Product(self(a), Division(one, self(b)))
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