[Python] Sympy: I cannot get lambdifying a summation with cse=True to work

I am trying to use Sympy, python's computer algebra system, to model a spherical harmonic gravity model, and it is either very slow, or not working. The slowness I understand, but when I use CSE with lambidfy, it produces a badly formed callback. CSE identifies and precomputes common subexpressions before they are used multiple times.

Here is a simplified example showing the problem:

#%%
import sympy as sy
from IPython.display import display

# normal symbols
x = sy.Symbol('x', real=True)
e = sy.Symbol('e', real=True, positive=True)

# indexed symbols
i = sy.Symbol('i', integer=True)
iMax = sy.Symbol("i_m", integer=True, positive=True)
xOfI = sy.Indexed('x',i)

# the expression
s = sy.Sum((xOfI+e*sy.sin(xOfI)+(e*sy.sin(xOfI))**2),(i,1,iMax))
display(s)

# lambdify normally, it works
callback = sy.lambdify([x, e, iMax], s)
answer = callback([5, 6, 7, 8, 9], 0.01, 4)
print(answer)

# lambdify with cse=True, and it fails
callbackWithCse = sy.lambdify([x, e, iMax], s, cse=True)
answer2 = callbackWithCse([5, 6, 7, 8, 9], 0.01, 4)
print(answer2)


The code starts by defining a simple summation (this is just one I made up, it mathematically means nothing) where sin(xOfI) is a repeated term. I then lambidfy it with default options (cse is False), and it evaluates as I expect it.

But the real summation is MUCH longer and I want to be able to use cse on it because many terms and elements are themselves complicated functions of the state, and repeated. But when I run it or this with cse=True I get the same problem (here's the relevant portion of my output when I run the script):

Sum(e**2*sin(x)**2 - e*sin(x) + x, (i, 1, i_m)) # first display
29.98237535925591 # first summation with cse=False
Traceback (most recent call last):
...
answer2 = callbackWithCse([5, 6, 7, 8, 9], 0.01, 4)
File "<lambdifygenerated-2>", line 2, in _lambdifygenerated
x0 = sin(x)
TypeError: list indices must be integers or slices, not Symbol


We get a better hint of whats wrong by looking at the error in a Jupyter notebook:

1 def _lambdifygenerated(x, e, i_m):
----> 2 x0 = sin(x)
3 return (builtins.sum(e**2*x0**2 + e*x0 + x for i in range(1, i_m+1)))


With cse=True, it extracted the indexed variables (or terms with indexed variables) out of the summation where the i, the loop variable, is undefined.

Am I doing something wrong? I've tried a bunch of variations but none of them have worked the way I would want them to. If this is an unfix'able problem, is there a way around it (besides calling .doit() on the Sum before lambdifying. That makes things REALLY slow for the real sum I'm working with).

Continue reading...