I ended up using sympy.collect
. If the equation doesn't have too many variables, it's possible to simply brute force all combinations, and to recurse down into the "collected" terms.
Here's the code I came up with. There's probably room for a lot of improvement:
def collect_best(expr, measure=sympy.count_ops):
# This method performs sympy.collect over all permutations of the free variables, and returns the best collection
best = expr
best_score = measure(expr)
perms = itertools.permutations(expr.free_symbols)
permlen = np.math.factorial(len(expr.free_symbols))
print(permlen)
for i, perm in enumerate(perms):
if (permlen > 1000) and not (i%int(permlen/100)):
print(i)
collected = sympy.collect(expr, perm)
if measure(collected) < best_score:
best_score = measure(collected)
best = collected
return best
def product(args):
arg = next(args)
try:
return arg*product(args)
except:
return arg
def rcollect_best(expr, measure=sympy.count_ops):
# This method performs collect_best recursively on the collected terms
best = collect_best(expr, measure)
best_score = measure(best)
if expr == best:
return best
if isinstance(best, sympy.Mul):
return product(map(rcollect_best, best.args))
if isinstance(best, sympy.Add):
return sum(map(rcollect_best, best.args))
rcollect_best
turns this (count_ops = 136):
4*a**3*d*e - 6*a**2*b*d*e - 6*a**2*c*d*e + 16*a**2*e**3 + 6*a**2*e*f**2 + 6*a**2*e*g**2 + 2*a*b**2*d*e + 8*a*b*c*d*e - 14*a*b*e**3 - 2*a*b*e*f**2 - 8*a*b*e*g**2 + 2*a*c**2*d*e - 14*a*c*e**3 - 8*a*c*e*f**2 - 2*a*c*e*g**2 - 2*b**2*c*d*e + 2*b**2*e**3 + 2*b**2*e*g**2 - 2*b*c**2*d*e + 8*b*c*e**3 + 2*b*c*e*f**2 + 2*b*c*e*g**2 + 2*c**2*e**3 + 2*c**2*e*f**2
Into this (count_ops = 68):
2*e*(d*(2*a**3 - 3*a**2*b + a*b**2 + c**2*(a - b) + c*(-3*a**2 + 4*a*b - b**2)) + e**2*(8*a**2 - 7*a*b + b**2 + c**2 + c*(-7*a + 4*b)) + f**2*(3*a**2 - a*b + c**2 + c*(-4*a + b)) + g**2*(3*a**2 - 4*a*b + b**2 + c*(-a + b)))
Which is a 5th degree polynomial in 7 variables. Run-time is about 10 or 15 minutes, and will increase super-exponentially, so I don't recommend this for anything much more demanding than this. I'm sure there are some basic improvements that could fix the super-exponential growth, but this has solved my problem, so I'm cashing out now. :)