Collecting sub-expressions of a multivariate polynomial function in Sympy - sympy

Question

I have a degree 6 multivariate equation in x and y written in Sympy, e.g.
eqn = a*x**6 + b*x**5*y + c*x**4*y + d*x**3*y + e*x**3*y**2 + ...
Is there a way to collect (x**2+y**2) and rearrange them into the following format?
eqn2 = A*(x**2+y**2)**3 + B*(x**2+y**2)**2 + C*(x**2+y**2) + D
A, B, C, D can be in x, y.
So far I have only tried collect(eqn, x**2 + y**2) and it returned the original equation.
Thank you!

Answer 1

Consider using a temporary symbol z = x**2 + y**2 and replace x**2 with z - y**2, then expand and restore:
>>> ex
A*x**6 + 3*A*x**4*y**2 + 3*A*x**2*y**4 + A*y**6 + B*x**4 + 2*B*x**2*y**2 +
B*y**4 + C*x**2 + C*y**2 + D
>>> ex.subs(x**2, z - y**2).expand().subs(z, x**2 + y**2)
A*(x**2 + y**2)**3 + B*(x**2 + y**2)**2 + C*(x**2 + y**2) + D
Although that works, perhaps a more direct thing to do is separate the expression by coefficients A-D and then factor those collections of terms:
def separatevars_additively(expr, symbols=[]):
free = set(symbols) or expr.free_symbols
d = {}
while free:
f = free.pop()
expr, dep = expr.as_independent(f, as_Add=True)
if dep.has(*free):
return None
d[f] = dep
if expr:
d[0] = expr
return d
>>> coeff = var("A:D")
>>> separatevars_additively(ex, coeff)
{B: B*x**4 + 2*B*x**2*y**2 + B*y**4, A: A*x**6 + 3*A*x**4*y**2 + 3*A*x**2*y**4 + A*y**6, D: D, C: C*x**2 + C*y**2}
>>> Add(*[factor(i) for i in _.values()])
A*(x**2 + y**2)**3 + B*(x**2 + y**2)**2 + C*(x**2 + y**2) + D