How to extract all coefficients in sympy

Question

You can get a coefficient of a specific term by using coeff();

x, a = symbols("x, a")
expr = 3 + x + x**2 + a*x*2
expr.coeff(x)
# 2*a + 1

Here I want to extract all the coefficients of x, x**2 (and so on), like;

# for example
expr.coefficients(x)
# want {1: 3, x: (2*a + 1), x**2: 1}

There is a method as_coefficients_dict(), but it seems this doesn't work in the way I want;

expr.as_coefficients_dict()
# {1: 3, x: 1, x**2: 1, a*x: 2}
expr.collect(x).as_coefficients_dict()
# {1: 3, x**2: 1, x*(2*a + 1): 1}

Mylan

Are there limits on your expression expr, i.e. is it always a polynomial? Is its degree limited? – Ladybird

@Ladybird It would be great if there's an "unlimited" way, and in fact I think there must be especially because collect() is (probably) unlimited, but at least a method that works on a polynomial is needed. – Mylan

Answer 1

The easiest way is to use Poly

>>> a = Poly(expr, x)
>>> a.coeffs()
[1, 2*a + 1, 3]

Noelianoell

Note that this hides zeros. If you do Poly(x**2 + 1).coeffs() you get [1,1] and not [1,0,1]. – Pardo

Answer 2

all_coeffs() can be sometime better than using coeffs() for a Poly. The difference lies in output of these both. coeffs() returns a list containing all coefficients which has a value and ignores those whose coefficient is 0 whereas all_coeffs() returns all coefficients including those whose coefficient is zero.

>>> a = Poly(x**3 + a*x**2 - b, x)
>>> a.coeffs()
[1, a, -b]

>>> a.all_coeffs()
[1, a, 0, -b]

Fess

Thanks. I was caught with my pants down by coeffs() ignoring zeroed coefficients. It's also worth noting that numpy uses the reverse ordering so when using numpy polynomial its worthwhile doing a.reverse() – Kai

Can I just ask, what's going on when you write a = Poly(x**3 + a*x**2 - b, x)? It looks like you're entering an equation with a on both sides, but I guess that's not what you're doing since SymPy doesn't know that you're assigning the polynomial to a variable that also has the name "a". But did you choose that name for some specific reason, instead of calling it for example "p", for polynomial? Otherwise, I find it a bit confusing. – Lapidary

Answer 3

Collection of coefficients can be handled with Poly and then separation of the monomials into dependent and independent parts can be handled with Expr.as_independent:

def codict(expr, *x):
  collected = Poly(expr, *x).as_expr()
  i, d = collected.as_independent(*x, as_Add=True)
  rv = dict(i.as_independent(*x, as_Mul=True)[::-1] for i in Add.make_args(d))
  if i:
      assert 1 not in rv
      rv.update({S.One: i})
  return rv

>>> var('a x z y')
(a, x, z, y)
>>> expr = 3 + x + x**2 + a*x*2
>>> codict(expr, x)
{x**2: 1, x: 2*a + 1, 1: 3}
>>> codict(expr+y+z, x)
{x**2: 1, x: 2*a + 1, 1: y + z + 3}
>>> codict(expr+y+z, x,y)
{y: 1, x**2: 1, x: 2*a + 1, 1: z + 3}
>>> codict(expr+y+z, x,y,z)
{y: 1, z: 1, x**2: 1, x: 2*a + 1, 1: 3}

Hesperidium

This does not work. Sometimes there are more than one entry in the dict for 1, in which case only the last is stored. You need to sum over all entries with same key in the dict. – Wadding

Thanks. The code and example has been updated. – Hesperidium

Simply outstanding. – Sector

Answer 4

One thing you can do is use a dictionary comprehension like so:

dict = {x**p: expr.collect(x).coeff(x**p) for p in range(1,n)}

where n is the highest power+1. In this case n=3. So you would have the list [1,2]

This would give

dict = {x: (2*a+1), x**2: 1}

Then you can add in the single term with

dict[1] = 3

So

 dict = {1:3,x:(2*a+1),x**2:1}

You may also try:

a = list(reversed(expr.collect(x).as_ordered_terms()))
dict = {x**p: a[p],coeff(x**p) for p in range(1,n)}
dict[1] = a[0] # Would only apply if there is single term such as the 3 in the example

where n is the highest power + 1.

Izzard

As an exception, this works on series() while Poly() doesn't. ■ Side note, you can also write .coeff(x, p) which would allow p=0 (constant term). – Incogitant

Answer 5

Based on RDizzl3's answer, I modified like this:

d_c = collect(my_poly, x)
# replace n by highest power + 1 
dict = {x**p: d_c.coeff(x, p) for p in range(0,n)}

This does not skip the constant.

Zooplasty