Collecting like term of an expression in Sympy

Question

I am currently dealing with functions of more than one variable and need to collect like terms in an attempt to simplify an expression.

Say the expression is written as follows:

x = sympy.Symbol('x')
y = sympy.Symbol('y')
k = sympy.Symbol('k')
a = sympy.Symbol('a')

z = k*(y**2*(a + x) + (a + x)**3/3) - k((2*k*y*(a + x)*(n - 1)*(-k*(y**2*(-a + x) + (-a + x)**3/3) + k*(y**2*(a + x) + (a + x)**3/3)) + y)**2*(-a + k*(n - 1)*(y**2 + (a + x)**2)*(-k*(y**2*(-a + x)))))
zEx = z.expand()
print type(z)
print type(zEx)

EDIT: Formatting to add clarity and changed the expression z to make the problem easier to understand.

Say z contains so many terms, that sifting through them by eye. and selecting the appropriate terms, would take an unsatisfactory amount of time.

I want to collect all of the terms which are ONLY a multiple of a**1. I do not care for quadratic or higher powers of a, and I do not care for terms which do not contain a.

The type of z and zEx return the following:

print type(z)
print type(zEx)
>>>
<class 'sympy.core.add.Add'>
<class 'sympy.core.mul.Mul'>

Does anyone know how I can collect the terms which are a multiple of a , not a^0 or a^2?

tl'dr

Where z(x,y) with constants a and k described by z and zEx and their type(): How can one remove all non-a terms from z AND remove all quadratic or higher terms of a from the expression? Such that what is left is only the terms which contain a unity power of a.

Grandfather

Answer 1

In addition to the other answers given, you can also use collect as a dictionary.

print(collect(zEx,a,evaluate=False)[a])

yields the expression

k*x**2 + k*y**2

Banwell

Answer 2

In the end it is just an one-liner. @asmeurer brought me on the right track (check the comments below this post). Here is the code; explanations can be found below:

from sympy import *
from sympy.parsing.sympy_parser import parse_expr
import sys

x, y, k, a = symbols('x y k a')

# modified string: I added a few terms
z = x*(k*a**9) + (k**1)*x**2 - k*a**8 + y*x*(k**2) + y*(x**2)*k**3 + x*(k*a**1) - k*a**3 + y*a**5

zmod = Add(*[argi for argi in z.args if argi.has(a)])

Then zmod is

a**9*k*x - a**8*k + a**5*y - a**3*k + a*k*x

So let's look at this more carefully:

z.args

is just a collection of all individual terms in your expression (please note, that also the sign is parsed which makes things easier):

(k*x**2, a**5*y, -a**3*k, -a**8*k, a*k*x, a**9*k*x, k**2*x*y, k**3*x**2*y)

In the list comprehension you then select all the terms that contain an a using the function has. All these terms can then be glued back together using Add which gives you the desired output.

EDIT

The above returns all all the expressions that contain an a. If you only want to filter out the expressions that contain a with unity power, you can use collect and Mul:

from sympy import *
from sympy.parsing.sympy_parser import parse_expr
import sys

x, y, k, a = symbols('x y k a')

z2 = x**2*(k*a**1) + (k**1)*x**2 - k*a**8 + y*x*(k**2) + y*(x**2)*k**3 + x*k*a - k*a**3 + y*a**1

zc = collect(z2, a, evaluate=False)
zmod2 = Mul(zc[a], a)

then zmod2 is

a*(k*x**2 + k*x + y)

and zmod2.expand()

a*k*x**2 + a*k*x + a*y

which is correct.

With the updated z you provide I run:

z3 =  k*(y**2*(a + x) + (a + x)**3/3) - k((2*k*y*(a + x)*(n - 1)*(-k*(y**2*(-a + x) + (-a + x)**3/3) + k*(y**2*(a + x) + (a + x)**3/3)) + y)**2*(-a + k*(n - 1)*(y**2 + (a + x)**2)*(-k*(y**2*(-a + x)))))
zc3 = collect(z3.expand(), a, evaluate=False)
zmod3 = Mul(zc3[a], a)

and then obtain for zmod3.expand():

a*k*x**2 + a*k*y**2

Is this the result you were looking for?

PS: Thanks to @asmeurer for all these helpful comments!

Brosy

If you're using strings to manipulate sympy expressions you're doing it wrong. – Wiburg

@asmeurer: As far as I understood the question, the OP wants to filter all the terms in z that contain a parameter a. Converting it to a string first was the only way I could get this done since iterating over z.args does not allow(?) to check whether a is in the element. The collect seems not suitable for that either but maybe I use it in a wrong manner. Maybe you could extend your answer, using z like I defined it and try to get the same result as I did without converting it to a string first!? I am always glad to broaden my knowledge :) – Brosy

Thanks for sharing! Splitting it into a string seems to work well when the expression is all positive. But what if there are negative terms in the expression? I don't think it wouldsplit the terms correctly. How could I go about doing this is everyother term was negative? – Grandfather

If ALL of them are negative you have to do zstr.split('-') and ' - '.join(res). If there are both + and -, it becomes a bit trickier. I could try to update the post later on. But according to @Wiburg and at least one other person that is not the way to go - it got even downvoted although it does the job. – Brosy

@user2970710: I updated the code. It now also works with mixed signs. Let me know if that answers your question. – Brosy

@Brosy you can check if an expression contains a subexpression using has, like term.has(a). – Wiburg

@asmeurer: Ok, thanks, I did not know that. But how would you then glue these terms back together using the appropriate sign for each term? – Brosy

In SymPy, subtraction is represented as multiplication by -1. So x - y is the same as Add(x, Mul(-1, y)). Hence if you split it by args the sign is already on the terms, so to put them back together you just need to do Add(*terms). – Wiburg

@asmeurer: Thanks so much for those comments! I updated my code; hope it is correct now... Could learn a lot which is worth your downvote ;) – Brosy

@Brosy I removed the downvote. I would just remove the old answer, though. It will still be visible in the edit history. – Wiburg

I have also learned a lot about how to ask a question concisely on this site ^^ Your method workd well @Cleb. However it seems to collect all terms which contain a, as opposed to ONLY terms with a unity power of a. I only wish to "glue together" the terms which contain a1. Whereas your method provides me with a9, a8, a5 etc. Apologies for not making this stunningly clear in the first place! I really appreciate your efforts and I think you're pretty close to cracking it :D Hope this helps. – Grandfather

@user2970710: I updated my post once again; please check the EDIT. Hope that this is the result you were looking for?! :) – Brosy

Answer 3

To iterate over the terms of an expression use expr.args.

I'm unclear what a is supposed to be, but the collect function may do what you want.

Wiburg

Sorry. a should have been defined as another variable: a = sympy.Symbol('a') – Grandfather