def dist(prod):
if not isinstance(prod, Product):
return prod
leading = []
for i in prod.children:
if isinstance(i, Sum):
break
else:
leading.append(i)
if len(leading) == len(prod.children):
# no more sums found
result = pymbolic.flattened_product(prod.children)
return result
else:
sum = prod.children[len(leading)]
assert isinstance(sum, Sum)
rest = prod.children[len(leading)+1:]
if rest:
rest = dist(Product(rest))
else:
rest = 1
result = self.collect(pymbolic.flattened_sum(
pymbolic.flattened_product(leading) * dist(sumchild*rest)
for sumchild in sum.children
))
return result
def map_call(self, expr, *args):
return pymbolic.flattened_sum(
self.function_map(i,
expr.function,
self.rec_undiff(expr.parameters, *args),
allowed_nonsmoothness=self.allowed_nonsmoothness)
* self.rec(par, *args) for i, par in enumerate(expr.parameters))
def map_product(self, expr, *args):
return pymbolic.flattened_sum(
pymbolic.flattened_product(
[self.rec_undiff(ch, *args)
for ch in expr.children[0:i]] + [self.rec(child, *args)] +
[self.rec_undiff(ch, *args) for ch in expr.children[i + 1:]])
for i, child in enumerate(expr.children))
def map_product(self, expr):
return pymbolic.flattened_sum(
pymbolic.flattened_product(
expr.children[0:i] +
(self.rec(child),) +
expr.children[i+1:])
for i, child in enumerate(expr.children))
def map_call(self, expr, *args):
return pymbolic.flattened_sum(
self.function_map(
i, expr.function, self.rec_undiff(expr.parameters, *args))
* self.rec(par, *args)
for i, par in enumerate(expr.parameters)
)
def expand(prod):
if not isinstance(prod, Product):
return prod
leading = []
for i in prod.children:
if isinstance(i, Sum):
break
else:
leading.append(i)
if len(leading) == len(prod.children):
# no more sums found
result = pymbolic.flattened_product(prod.children)
return result
else:
sum = prod.children[len(leading)]
assert isinstance(sum, Sum)
rest = prod.children[len(leading)+1:]
if rest:
rest = expand(Product(rest))
else:
rest = 1
result = self.collector(pymbolic.flattened_sum(
pymbolic.flattened_product(leading) * expand(sumchild*rest)
for sumchild in sum.children
))
return result
def map_product(self, expr, *args):
return pymbolic.flattened_sum(
pymbolic.flattened_product(
[self.rec_undiff(ch, *args) for ch in expr.children[0:i]] +
[self.rec(child, *args)] +
[self.rec_undiff(ch, *args) for ch in expr.children[i+1:]]
)
for i, child in enumerate(expr.children))
def map_sum(self, mysum):
term2coeff = {}
for child in mysum.children:
term, coeff = self.split_term(child)
term2coeff[term] = term2coeff.get(term, 0) + coeff
def rep2term(rep):
return pymbolic.flattened_product(base ** exp for base, exp in rep)
result = pymbolic.flattened_sum(coeff * rep2term(termrep) for termrep, coeff in term2coeff.iteritems())
return result
def map_sum(self, mysum):
term2coeff = {}
for child in mysum.children:
term, coeff = self.split_term(child)
term2coeff[term] = term2coeff.get(term, 0) + coeff
def rep2term(rep):
return pymbolic.flattened_product(base**exp for base, exp in rep)
result = pymbolic.flattened_sum(
coeff * rep2term(termrep)
for termrep, coeff in six.iteritems(term2coeff))
return result
def map_product(self, expr):
return pymbolic.flattened_sum(
pymbolic.flattened_product(expr.children[0:i] +
(self.rec(child), ) +
expr.children[i + 1:])
for i, child in enumerate(expr.children))
def map_sum(self, expr):
return pymbolic.flattened_sum(
self.rec(child) for child in expr.children)
def map_call(self, expr):
return pymbolic.flattened_sum(
self.function_map(i, expr.function, expr.parameters) *
self.rec(par) for i, par in enumerate(expr.parameters))
def map_sum(self, expr, *args):
return pymbolic.flattened_sum(
self.rec(child, *args) for child in expr.children)
def map_call(self, expr):
return pymbolic.flattened_sum(
self.function_map(i, expr.function, expr.parameters)
* self.rec(par)
for i, par in enumerate(expr.parameters)
)
