pythonsympyequationalgebra

How to extract all coefficients for a particular variable


I have done some calculations in sympy, and the result is below: (I use [expr].coeff(u1) method)

X1*X2*(u1 - u2 + u3 - u4)/(4*a*b) + X1*(-u1 + u2 + u3 - u4)/(4*a) + X2*(-u1 - u2 + u3 + u4)/(4*b) + u1/4 + u2/4 + u3/4 + u4/4

I need to extract all coefficients for variable u1 (for example) and it's supposed to look like this:

X1*X2/(4*a*b) + -X1/(4*a) + -X2/(4*b) + 1/4

But sympy do this:

1/4

In smaller scales - I need coefficient for u1 in this

(-u1 + u2 + u3 - u4)/(4*a)

Supposed to look like this:

-1/(4*a)

But I get something wrong:

0

Example of code here:

import sympy as sp
import numpy as np

a0, a1, a2, a3, X1, X2, a, b = sp.symbols('a0 a1 a2 a3 X1 X2 a b')
u1, u2, u3, u4 = sp.symbols('u1 u2 u3 u4')
Disp = a0 + a1*X1 + a2*X2 + a3*X1*X2

Eq1 = Disp.subs({X1: -a, X2: -b})
Eq2 = Disp.subs({X1: a, X2: -b})
Eq3 = Disp.subs({X1: a, X2: b})
Eq4 = Disp.subs({X1: -a, X2: b})

sol = sp.solve([Eq1 - u1, Eq2 - u2, Eq3 - u3, Eq4 - u4]
               , [a0, a1, a2, a3], dict=True)
a0, a1, a2, a3 = sol[0][a0], sol[0][a1], sol[0][a2], sol[0][a3]
Disp1 = a0 + a1*X1 + a2*X2 + a3*X1*X2
# print(Disp1)
N1 = Disp1.as_expr()
N2 = Disp1.coeff(u2)
print(N2)

Solution

  • I simply use sympy.expand(expr).coeff(u1) to get the:

    X1*X2/(4*a*b) - X1/(4*a) - X2/(4*b) + 1/4
    

    from there:

    X1*X2*(u1 - u2 + u3 - u4)/(4*a*b) + X1*(-u1 + u2 + u3 - u4)/(4*a) + X2*(-u1 - u2 + u3 + u4)/(4*b) + u1/4 + u2/4 + u3/4 + u4/4