# Extracting coefficients in sympy is only working for certain symbols

## Question:

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`?

You need to expand before calling coeff:

``````In : x, y, z = symbols('x, y, z')

In : expr = x*(y + z)

In : expr
Out: x⋅(y + z)

In : expr.coeff(x)
Out: y + z

In : expr.coeff(y) # y does not appear at top level
Out: 0

In : expr.expand() # Now y and z are at top level
Out: x⋅y + x⋅z

In : expr.expand().coeff(x)
Out: y + z

In : expr.expand().coeff(y)
Out: x

In : expr.expand().coeff(z)
Out: x
``````

It is also possible to find the term(s) of interest without having to do a full expansion by replacing products that are not of interest with 0:

``````>>> (2*x00+expr).replace(lambda x:x.is_Mul, lambda x: x if x.has_free(x00) else 0)
p_{x0}*p_{y0}*x_{00} + 2*x_{00}
>>> _.coeff(x00)
p_{x0}*p_{y0} + 2
``````
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.