How to factor powers of an matrix in a equation in Sympy?

I have an expression that comes from

from sympy.physics.quantum import Commutator as cmm

    x, t, A, V, W, D = sp.symbols('x t A V W D', commutative = False)
    Q = sp.Function('Q', commutative = False)
    F = (sp.diff(Q(x,t),x)+ cmm(W,Q(x,t)).doit() - sp.I*cmm(A,Q(x,t)+ cmm(W,Q(x,t)).doit()).doit())*(sp.diff(Q(x,t),x)+ cmm(W,Q(x,t)).doit() - sp.I*cmm(A,Q(x,t)+ cmm(W,Q(x,t)).doit()).doit())
    F.expand()

This gives me an expression with elements with zero order in W, with first order in W and second order. I want to get only the first order W's. I tried the factorization procedures but it appears that, due to the fact of no commutation, it doesn't recognize powers of W. It always gives me 0. Is there any simple way of doing this ? Of course I can do it by hand but this is not my goal.

Thanks

Solution

You can get all the terms that are first order in W by collecting them as you traverse the arguments of F:

>>> from sympy import Add
>>> first_order_terms = []
>>> for i in Add.make_args(F.expand()):
...   if i == W or i.is_Mul and i.has(W) and i.subs(W,y).as_independent(y)[1] == y:
...     first_order_terms.append(i)
... 
>>> Add(*first_order_terms)
-A*W*Q(x, t)*A*Q(x, t) - I*A*W*Q(x, t)*Derivative(Q(x, t), x) + A*W*Q(x, t)**2*A - 
A*Q(x, t)*A*W*Q(x, t) + A*Q(x, t)*A*Q(x, t)*W + A*Q(x, t)*W*A*Q(x, t) - I*A*Q(x, 
t)*W*Q(x, t) + I*A*Q(x, t)*W*Derivative(Q(x, t), x) + I*A*Q(x, t)**2*W - A*Q(x, 
t)**2*W*A - I*W*Q(x, t)*A*Q(x, t) - W*Q(x, t)*A*Q(x, t)*A + I*W*Q(x, 
t)*A*Derivative(Q(x, t), x) + W*Q(x, t)*A**2*Q(x, t) + W*Q(x, t)*Derivative(Q(x, t), 
x) + I*W*Q(x, t)**2*A + I*Q(x, t)*A*W*Q(x, t) - Q(x, t)*A*W*Q(x, t)*A - I*Q(x, 
t)*A*Q(x, t)*W + Q(x, t)*A*Q(x, t)*W*A + Q(x, t)*A**2*W*Q(x, t) - Q(x, t)*A**2*Q(x, 
t)*W + I*Q(x, t)*W*A*Q(x, t) + Q(x, t)*W*A*Q(x, t)*A - I*Q(x, t)*W*A*Derivative(Q(x, 
t), x) - Q(x, t)*W*A**2*Q(x, t) - I*Q(x, t)*W*Q(x, t)*A - Q(x, t)*W*Derivative(Q(x, 
t), x) - I*Derivative(Q(x, t), x)*A*W*Q(x, t) + I*Derivative(Q(x, t), x)*A*Q(x, t)*W + 
Derivative(Q(x, t), x)*W*Q(x, t) + I*Derivative(Q(x, t), x)*W*Q(x, t)*A - 
Derivative(Q(x, t), x)*Q(x, t)*W - I*Derivative(Q(x, t), x)*Q(x, t)*W*A

A more compact way to do this, after expanding F, is to do as follows:

>>> Add(*[i for i in F.expand().atoms(Mul) if i.xreplace({W:z}).coeff(z) != 0])

coeff will only return a non-zero value if z (in this case) appears as a linear factor. Of course, this does not work if you just replace W with z in F because terms which didn't cancel because of their distince non-commutative factors will do so after the substitution so you will not get the same result.