Polynomial calculation in SymPy

Question

Thie programs is written in Sympy 1.9 and it should find a polynomial g of degree dim for a given polynomial f such that f o g = g o f as described here, where I already posted program written in Maxima. This is my first program in SymPy so I am interested in your feedback on this program except on the algorithm itself.

from sympy import symbols, pprint, expand, Poly, IndexedBase, Idx,Eq,solve
from sympy.abc import x,f,g,i
fg,gf,d=symbols('fg gf d')
f=x**3+3*x # a polynomial g exists
#f=x**3+4*x # a polynomial g does not exist
dim=5
a=IndexedBase('a')
i=Idx('i',dim)
# create a polynomial of degree dim 
# with unknown coefficients,
# but the coefficient of the highest power is 1
# an the lowest power is 1
g=a[0]+x**dim
for i in range(1,dim):
 g+=x**i*a[i]
# calculate fg = f o g 
# and gf = g o f
fg=f
gf=g
fg=expand(f.subs(x,g))
gf=expand(g.subs(x,f))
# calculate the difference d=(f o g) - (g o f)
# and check how to choose the coefficients of g
# such that it will become 0
difference=(fg-gf).as_poly(x)
candidates=solve(difference.all_coeffs()[:dim],[a[i] for i in range(dim)])
replacements=[(a[i],candidates[0][i]) for i in range(dim)]
if (difference.subs(replacements)==0):
 pprint(g.subs(replacements))
else:
 print("no solution exists")

Answer 1

Turning into a Function

I would advise turning your program into a function. This allows people to use your function in other Python scripts. That is:

def poly_commute(f, dim):
 a=IndexedBase('a')
 i=Idx('i',dim)
 g=a[0]+x**dim
 for i in range(1,dim):
 g+=x**i*a[i]
 # calculate fg = f o g 
 # and gf = g o f
 fg=f
 gf=g
 fg=expand(f.subs(x,g))
 gf=expand(g.subs(x,f))
 # calculate the difference d=(f o g) - (g o f)
 # and check how to choose the coefficients of g
 # such that it will become 0
 difference=(fg-gf).as_poly(x)
 candidates=solve(difference.all_coeffs()[:dim],[a[i] for i in range(dim)])
 replacements=[(a[i],candidates[0][i]) for i in range(dim)]
 if (difference.subs(replacements)==0):
 return g.subs(replacements)
 else:
 return None
 

Also, you should add some documentation so others can understand what your function does.

Related: Removing Global Variables

Managing state can be a pain. Keeping track of what can be in which state can be a headache. It isn't of course unmanageable in this case, but it would be useful to get rid of the state. To do so, I would add a main function.

def main():
 f = x**3 + 3*x
 dim = 5
 g = poly_commute(f, dim)
 if g:
 pprint(g)
 else:
 print("No solution exists.")
if __name__ == '__main__':
 main()

PEP 8:

There is some spacing of equations that seem to be off. This is explained in "Other recommendations" of PEP 8:

If operators with different priorities are used, consider adding whitespace around the operators with the lowest priority(ies). Use your own judgment; however, never use more than one space, and always have the same amount of whitespace on both sides of a binary operator:

# Correct:
i = i + 1
submitted += 1
x = x*2 - 1
hypot2 = x*x + y*y
c = (a+b) * (a-b)
# Wrong:
i=i+1
submitted +=1
x = x * 2 - 1
hypot2 = x * x + y * y
c = (a + b) * (a - b)