CODEWITHSUNDEEP
sympysymbolic-mathalgebra
Jarrod Haas
Jarrod Haas

Reputation: 65

SymPy: How to collect multi-variable terms?

I have a long expression in sympy that looks something like this:

2𝑤^2𝑥𝑦𝑧 + 3𝑤^2 𝑥𝑧^2 − 4𝑤𝑥𝑦^2 − 5𝑤𝑥𝑦𝑧 + 2𝑤𝑥𝑧^2 − 3𝑥𝑦^2 − 4𝑥𝑦𝑧 − 5𝑥𝑧^2 ...

where the variables are combinations of w,x,y,z and can also have powers of 2.

I would like to collect terms [x,y,z] such that I end up with a quadratic polynomial in w as the coefficient to each term, for example:

xyz(−0.285𝑤^2 - 1.09⋅10−5𝑤, 1.60⋅10−10) + xy^2(...w^2, ...w, ...) + x^2y(...) + xz^2(...) + ...

When I use collect(exp, [x,y,z]) I'm not able to achieve this; sympy seems to want to nest some terms:

𝑥(𝑦2(−0.00341083824360158𝑤−2.60668077412341⋅10−8) + 
𝑦𝑧(−0.28569359766975𝑤2−1.09161331685904⋅10−5𝑤−1.60378772636814⋅10−10) + 𝑧2(0.000708269071473656𝑤2+1.72432957139821⋅10−10𝑤−2.29362549750881⋅10−13))

I've also tried several combinations of collecting terms one at a time, simplifying, expanding etc. but to no avail. Any help is greatly appreciated!

Thanks!

Upvotes: 3

Views: 2644

Answers (3)

Patosai
Patosai

Reputation: 760

I faced the same problem and wrote this recursive function.

def collect_with_respect_to_vars(eq, vars):
    assert isinstance(vars, list)
    eq = eq.expand()
    if len(vars) == 0:
        return {1: eq}

    var_map = eq.collect(vars[0], evaluate=False)
    final_var_map = {}
    for var_power in var_map:
        sub_expression = var_map[var_power]
        sub_var_map = collect_with_respect_to_vars(sub_expression, vars[1:])
        for sub_var_power in sub_var_map:
            final_var_map[var_power*sub_var_power] = sub_var_map[sub_var_power]
    return final_var_map

It returns a map with all the combined powers of the given vars as keys.

collect_with_respect_to_vars(x*2 + y**2*z*5 + x**3*y*4, [x, y, z])
=> {x: 2, x**3*y: 4, y**2*z: 5}

collect_with_respect_to_vars(x*2 + x*4 + z**3 + x*z*2, [x, y])
=> {x: 2*z + 6, 1: z**3}

I suppose in your case, you can use it like this.

collect_with_respect_to_vars(your_expr, [x, y, z])

Upvotes: 3

smichr
smichr

Reputation: 19077

As a reiteration of Oscar's point about getting SymPy to do what you want, here is another approach using cse which collects terms of more than 1 operation which are repeated in an expression. Since it will also collect squares that are repeated we will replace w**2 with u so it won't be collected.

>>> r, e = cse(expr.subs(w**2,u))

Here's the output of cse

>>> r
[(x0, x*y*z), (x1, x*y**2), (x2, x*z**2)]
>>> e
[2*u*x0 + 3*u*x2 - 5*w*x0 - 4*w*x1 + 2*w*x2 - 4*x0 - 3*x1 - 5*x2]

Now collect on the patterns:

>>> collect(e, [v for v, _ in r])

And restore the originals

>>> _.subs(list(reversed(r))).subs(u, w**2)
x*y**2*(-4*w - 3) + x*y*z*(2*w**2 - 5*w - 4) + x*z**2*(3*w**2 + 2*w - 5)

This collection scheme works because there are multi-factor coefficients of w-related factors in all cases. If there were some linear factors on w that would need some additional wrangling (for this and the other approach):

>>> expr += x*w + x*w**2 + y + y*w**2
>>> terms = set(e.as_coeff_Mul()[1] for c in expr.as_poly(w).all_coeffs() for e in Add.make_args(c))
>>> new_expr = sum(t * expr.coeff(t) for t in terms)
>>> new_expr.equals(expr)
False

Upvotes: 2

Oscar Benjamin
Oscar Benjamin

Reputation: 14480

I don't know that there is a ready function for this sort of thing in SymPy. What I like about SymPy though is that once you familiarise yourself with the internals it gives you a lot of power to define elementary operations like this.

This might work or hopefully at least gives a starting point:

In [124]: x, y, z, w = symbols('x, y, z, w')                                                                                      

In [125]: expr = 2*w**2*x*y*z + 3*w**2*x*z**2 - 4*w*x*y**2 - 5*w*x*y*z + 2*w*x*z**2 - 3*x*y**2 - 4*x*y*z - 5*x*z**2               

In [126]: expr                                                                                                                    
Out[126]: 
   2            2    2          2                      2        2                  2
2⋅w ⋅x⋅y⋅z + 3⋅w ⋅x⋅z  - 4⋅w⋅x⋅y  - 5⋅w⋅x⋅y⋅z + 2⋅w⋅x⋅z  - 3⋅x⋅y  - 4⋅x⋅y⋅z - 5⋅x⋅z 

In [127]: terms = set(e.as_coeff_Mul()[1] for c in expr.as_poly(w).all_coeffs() for e in Add.make_args(c))                        

In [128]: new_expr = sum(t * expr.coeff(t) for t in terms)                                                                        

In [129]: new_expr                                                                                                                
Out[129]: 
   2                    ⎛   2          ⎞      2 ⎛   2          ⎞
x⋅y ⋅(-4⋅w - 3) + x⋅y⋅z⋅⎝2⋅w  - 5⋅w - 4⎠ + x⋅z ⋅⎝3⋅w  + 2⋅w - 5⎠

Upvotes: 2

Related Questions