Question for speed-up sp.expand() on specific function (or: How to use hints?)

[MarioKrenn6240]

I use sympy for some physics simulations, for which i before used mathematica. It works great and easy. Now, I need to apply sp.expand() on functions like this:

fct=(vHa_*(trans_0*trans_3*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*phase_2)*a[H] - 1.0*sqrt(1 - trans_0**2)*sqrt(1 - trans_3**2)*a[H]) + sqrt(1 - vHa_**2)*(trans_0*trans_3*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*phase_2)*a[V] - 1.0*sqrt(1 - trans_0**2)*sqrt(1 - trans_3**2)*a[V])*exp(2.0*I*pi*vPa_))*(vHc_*(1.0*I*trans_0*sqrt(1 - trans_3**2)*a[H] + 1.0*I*trans_3*sqrt(1 - trans_0**2)*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*phase_2)*a[H]) + sqrt(1 - vHc_**2)*(1.0*I*trans_0*sqrt(1 - trans_3**2)*a[V] + 1.0*I*trans_3*sqrt(1 - trans_0**2)*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*phase_2)*a[V])*exp(2.0*I*pi*vPc_))*(vHb_*exp(2.0*I*pi*phase_4)*b[H] + sqrt(1 - vHb_**2)*exp(2.0*I*pi*phase_4)*exp(2.0*I*pi*vPb_)*b[V])

All variables are symbols. I need sp.expand(), because i need the coefficients of the monomials build from [a[H], a[V], b[H], b[V]] (e.g. b[H]*b[V]**2, b[V]**3, etc).

The function has several special properties, e.g.

• all symbols are between 0 and 1
• it is a homogeneous function of degree 3

Can i exploit this property in the expansion via hints or other arguments that might speed up the expansion?

Thank you very much!

[oscarbenjamin]

In general expansion is faster if you use Poly rather than Expr e.g. fct.as_poly(). You can convert back to Expr with fct.as_poly().as_expr() but if you want to do lots of arithmetic and have everything always expanded it is better to just keep everything as Poly while you do the arithmetic:

In [1]: e = '(vHa_*(trans_0*trans_3*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*phase_2)*a[H] - 1.0*sqrt(1 - t
   ...: rans_0**2)*sqrt(1 - trans_3**2)*a[H]) + sqrt(1 - vHa_**2)*(trans_0*trans_3*exp(2.0*I*pi*phase
   ...: _1)*exp(2.0*I*pi*phase_2)*a[V] - 1.0*sqrt(1 - trans_0**2)*sqrt(1 - trans_3**2)*a[V])*exp(2.0*
   ...: I*pi*vPa_))*(vHc_*(1.0*I*trans_0*sqrt(1 - trans_3**2)*a[H] + 1.0*I*trans_3*sqrt(1 - trans_0**
   ...: 2)*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*phase_2)*a[H]) + sqrt(1 - vHc_**2)*(1.0*I*trans_0*sqrt(
   ...: 1 - trans_3**2)*a[V] + 1.0*I*trans_3*sqrt(1 - trans_0**2)*exp(2.0*I*pi*phase_1)*exp(2.0*I*pi*
   ...: phase_2)*a[V])*exp(2.0*I*pi*vPc_))*(vHb_*exp(2.0*I*pi*phase_4)*b[H] + sqrt(1 - vHb_**2)*exp(2
   ...: .0*I*pi*phase_4)*exp(2.0*I*pi*vPb_)*b[V])'

In [2]: a, b = symbols('a, b', cls=IndexedBase)

In [4]: rep = {'a':a,'b':b}

In [5]: fct = parse_expr(e, local_dict=rep)

In [6]: %time ok = fct.expand()
CPU times: user 371 ms, sys: 0 ns, total: 371 ms
Wall time: 370 ms

In [7]: %time fct_p = fct.as_poly()
CPU times: user 70.3 ms, sys: 0 ns, total: 70.3 ms
Wall time: 67.9 ms

In [8]: %time ok = (fct**2).expand()
CPU times: user 13.2 s, sys: 2.69 ms, total: 13.2 s
Wall time: 13.2 s

In [9]: %time ok = fct_p ** 2
CPU times: user 1.27 s, sys: 3.8 ms, total: 1.28 s
Wall time: 1.28 s

Since these are multivariate polynomials the sparse representation is more efficient:

In [14]: dom = fct_p.domain[fct_p.gens]

In [15]: dom
Out[15]: CC[a[H],a[V],b[H],b[V],exp(2.0*I*pi*phase_1),exp(2.0*I*pi*phase_2),exp(2.0*I*pi*phase_4),exp(2.0*I*pi*vPa_),exp(2.0*I*pi*vPb_),exp(2.0*I*pi*vPc_),exp(4.0*I*pi*phase_1),exp(4.0*I*pi*phase_2),sqrt(1 - trans_0**2),sqrt(1 - trans_3**2),sqrt(1 - vHa_**2),sqrt(1 - vHb_**2),sqrt(1 - vHc_**2),trans_0,trans_3,vHa_,vHb_,vHc_]

In [16]: fct_p_s = dom(fct)

In [17]: %time ok = fct_p_s ** 2
CPU times: user 24.4 ms, sys: 7 µs, total: 24.4 ms
Wall time: 23.5 ms

Again fct_p_s.as_expr() can convert back to Expr.