def aggregate_coeffs(expr, mapper, nn_derivs=None):
nn_derivs = nn_derivs or mapper.get(expr)
args = [aggregate_coeffs(a, mapper, nn_derivs) for a in expr.args]
expr = reuse_if_untouched(expr, args, evaluate=True)
return expr
def _(expr, mapper, nn_derivs=None):
nn_derivs = nn_derivs or mapper.get(expr)
args = [aggregate_coeffs(a, mapper, nn_derivs) for a in expr.args]
expr = reuse_if_untouched(expr, args)
# Separate arguments containing derivatives from those which do not
hope_coeffs = []
with_derivs = []
for a in args:
if isinstance(a, sympy.Derivative):
with_derivs.append((a, [a], []))
else:
derivs, others = split(a.args, lambda i: isinstance(i, sympy.Derivative))
if a.is_Add and derivs:
with_derivs.append((a, derivs, others))
else:
hope_coeffs.append(a)
# E.g., non-linear term, expansion won't help (in fact, it would only
# cause an increase in operation count), so we skip
if len(with_derivs) > 1:
return expr
try:
with_deriv, derivs, others = with_derivs.pop(0)
except IndexError:
# No derivatives found, give up
return expr
# Aggregating the potential coefficient won't help if, in the current scope
# at least one derivative type does not appear more than once. In fact, aggregation
# might even have a detrimental effect due to increasing the operation count by
# expanding Muls), so we rather give if that's the case
if not any(nn_derivs[i._metadata] > 1 for i in derivs):
return expr
# Is the potential coefficient really a coefficient?
csymbols = set().union(*[i.free_symbols for i in hope_coeffs])
cdims = [i._defines for i in csymbols if i.is_Dimension]
ddims = [set(i.dims) for i in derivs]
if any(i & j for i, j in product(cdims, ddims)):
return expr
# Redundancies unlikely to pop up along the time dimension
if any(d.is_Time for d in flatten(ddims)):
return expr
if len(derivs) == 1 and with_deriv is derivs[0]:
expr = with_deriv._new_from_self(expr=expr.func(*hope_coeffs, with_deriv.expr))
else:
others = [expr.func(*hope_coeffs, a) for a in others]
derivs = [a._new_from_self(expr=expr.func(*hope_coeffs, a.expr)) for a in derivs]
expr = with_deriv.func(*(derivs + others))
return expr
def _deindexify(expr):
args = []
mapper = defaultdict(list)
for a in expr.args:
arg, m = _deindexify(a)
args.append(arg)
for k, v in m.items():
mapper[k].extend(v)
rexpr = reuse_if_untouched(expr, args)
if rexpr is not expr:
mapper[rexpr] = [expr]
return rexpr, mapper
def _(expr):
args = [factorize_derivatives(a) for a in expr.args]
derivs, others = split(args, lambda a: isinstance(a, sympy.Derivative))
if not derivs:
return reuse_if_untouched(expr, args)
# Map by type of derivative
# Note: `D0(a) + D1(b) == D(a + b)` <=> `D0` and `D1`'s metadata match,
# i.e. they are the same type of derivative
mapper = as_mapper(derivs, lambda i: i._metadata)
if len(mapper) == len(derivs):
return expr
args = list(others)
for v in mapper.values():
c = v[0]
if len(v) == 1:
args.append(c)
else:
args.append(c._new_from_self(expr=expr.func(*[i.expr for i in v])))
expr = expr.func(*args)
return expr
def _(expr, mapper, nn_derivs=None):
# Opens up a new derivative scope, so do not propagate `nn_derivs`
args = [aggregate_coeffs(a, mapper) for a in expr.args]
expr = reuse_if_untouched(expr, args)
return expr
def factorize_derivatives(expr):
args = [factorize_derivatives(a) for a in expr.args]
expr = reuse_if_untouched(expr, args)
return expr