Expand along a specific symbol in sympy

Question

When I do symbolic integration using Sympy, I would like to expand an expression along a specific symbol so that later I can collect and integrate w.r.t. that symbol. Is there a way to do such "expand" without altering the forms of other (irrelevant) symbols?

For example, consider this Python code and run results:

>>> from sympy import *
>>> var('x,a,b,c')
(x, a, b, c)
>>> f = (a + 1) * a + x * (b + exp(-c*x))
>>> collect(f, (exp(-c*x), x))
a*(a + 1) + x*(b + exp(-c*x))
>>> collect(expand(f), (exp(-c*x), x))
a**2 + a + b*x + x*exp(-c*x)

The outputs are all expected. Without "expand" first, "collect" just gives me back the original form. If we use "expand" first, I get what I want. Imagine if we have a sum of many above 'f' and symbols b and c are complex expressions, the integration will take a long time if we use the original form of f. Indeed, I have an integration taking seconds to complete if "expand" is applied first, but could not be finished after nearly one hour of run.

The problem of "expand" is that it can be expensive to "fully expand" the expression. In the above example, we see a*(a+1) was also expanded (and took computing time).

I have expressions (to be integrated later), each expanded to about 40 thousand terms. The result of expansion, as function of x, is similar to the form in the above example -- I knew it must be of that form. The most time (if it would ever finish) was spent on expanding those terms having nothing to do with 'x'. Can I avoid those unnecessary expansions?

Answer 1

You could try masking off the non-x-containing terms like this:

>>> def ex(e,x):
...  m = [i for i in e.atoms(Mul) if not i.has(x)]
...  reps = dict(zip(m,[Dummy() for i in m]))
...  return expand(e.xreplace(reps)).subs([(v,k) for k,v in reps.items()])
...
>>> f = (a + 1) * a + x * (b + exp(-c*x))
>>> ex(f,x)
a*(a + 1) + b*x + x*exp(-c*x)

Notice that only the x-containing term was expanded.

Also, instead of using expand, you might only want to use expand_mul; expand does several types of expansion (see docstring).