Polynomial expansion

Expanding series and sorting coefficients by hand can be tedious. The Python library sympy allows us to automate this process efficiently.

Polynomial expansion and coefficient extraction

from sympy import symbols, expand, collect

# define variables
a0, a1, a2, a3, x = symbols('a0 a1 a2 a3 x')

# expand the expression
expr = expand((a0 + a1*x + a2*x**2 + a3*x**3)**2)

print("Expanded expression:")
print(expr)

print("\nCollected by powers of x:")
print(collect(expr, x))

Output:

Expanded expression:
a0**2 + 2*a0*a1*x + 2*a0*a2*x**2 + 2*a0*a3*x**3
+ a1**2*x**2 + 2*a1*a2*x**3 + 2*a1*a3*x**4
+ a2**2*x**4 + 2*a2*a3*x**5 + a3**2*x**6

Collected by powers of x:
a0**2 + 2*a0*a1*x
+ (2*a0*a2 + a1**2)*x**2
+ (2*a0*a3 + 2*a1*a2)*x**3
+ (2*a1*a3 + a2**2)*x**4
+ 2*a2*a3*x**5
+ a3**2*x**6

We can also extract coefficients order by order:

coeffs = [expr.coeff(x, i) for i in range(7)]

for i, c in enumerate(coeffs):
    print(f"Coefficient of x^{i}:", c)  

Output:

Coefficient of x^0: a0**2
Coefficient of x^1: 2*a0*a1
Coefficient of x^2: 2*a0*a2 + a1**2
Coefficient of x^3: 2*a0*a3 + 2*a1*a2
Coefficient of x^4: 2*a1*a3 + a2**2
Coefficient of x^5: 2*a2*a3
Coefficient of x^6: a3**2

This approach allows us to systematically extract and match coefficients order by order, which is particularly useful in perturbation theory.