pythonsympycoefficients

Extracting coefficients in sympy is only working for certain symbols


If I define the following equation

import sympy as sp

x00, x01, x02 = sp.symbols('x_{00} x_{01} x_{02}')
x10, x11, x12 = sp.symbols('x_{10} x_{11} x_{12}')
x20, x21, x22 = sp.symbols('x_{20} x_{21} x_{22}')

px0, px1, px2 = sp.symbols('p_{x0} p_{x1} p_{x2}')
py0, py1, py2 = sp.symbols('p_{y0} p_{y1} p_{y2}')

Z = sp.Matrix([ [x00, x01, x02],
                [x10, x11, x12],
                [x20, x21, x22]])
px = sp.Matrix([px0, px1, px2])
py = sp.Matrix([py0, py1, py2])

expr = (py.T * Z * px)[0, 0]

print(expr.coeff(x00))
print(expr.coeff(px0))

I get the output

0
p_{y0}*x_{00} + p_{y1}*x_{10} + p_{y2}*x_{20}

Why is sympy returning the correct output for px0 but not for x00?


Solution

  • You need to expand before calling coeff:

    In [20]: x, y, z = symbols('x, y, z')
    
    In [12]: expr = x*(y + z)
    
    In [13]: expr
    Out[13]: x⋅(y + z)
    
    In [14]: expr.coeff(x)
    Out[14]: y + z
    
    In [15]: expr.coeff(y) # y does not appear at top level
    Out[15]: 0
    
    In [16]: expr.expand() # Now y and z are at top level
    Out[16]: x⋅y + x⋅z
    
    In [17]: expr.expand().coeff(x)
    Out[17]: y + z
    
    In [18]: expr.expand().coeff(y)
    Out[18]: x
    
    In [19]: expr.expand().coeff(z)
    Out[19]: x