Sympy codegen optimisation by temporary variables

[The Pauli Principle]

Here I am computing the distance d between two points A and B and the derivatives of d towards the coordinates of A and B

import sympy as sp
sp.init_printing()


X_A=sp.symbols('X_A')
Y_A=sp.symbols('Y_A')
Z_A=sp.symbols('Z_A')
X_B=sp.symbols('X_B')
Y_B=sp.symbols('Y_B')
Z_B=sp.symbols('Z_B')

d=sp.sqrt((X_A-X_B)**2+(Y_A-Y_B)**2+(Z_A-Z_B)**2)
Q=sp.Matrix([d])
R=sp.Matrix([X_A,Y_A,Z_A,X_B,Y_B,Z_B])
dQ_dR=Q.jacobian(R)
J=sp.MatrixSymbol('J',1,6)
print sp.fcode(dQ_dR,J)

This yields the (valid) code:

      J(1, 1) = (X_A - X_B)/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)
      J(1, 2) = (Y_A - Y_B)/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)
      J(1, 3) = (Z_A - Z_B)/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)
      J(1, 4) = (-X_A + X_B)/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)
      J(1, 5) = (-Y_A + Y_B)/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)
      J(1, 6) = (-Z_A + Z_B)/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)

unfortunately, this could have been written much more efficiently:

      overd=1.0/sqrt((X_A - X_B)**2 + (Y_A - Y_B)**2 + (Z_A
     @ - Z_B)**2)
      J(1, 1) = (X_A - X_B)*overd
      J(1, 2) = (Y_A - Y_B)*overd
      J(1, 3) = (Z_A - Z_B)*overd
      J(1, 4) = (-X_A + X_B)*overd
      J(1, 5) = (-Y_A + Y_B)*overd
      J(1, 6) = (-Z_A + Z_B)*overd

or even

      J(1, 4) = -J(1,1)
      J(1, 5) = -J(1,2)
      J(1, 6) = -J(1,3)

Is there any way to have sympy recognise the possible use of such "temporary variables", or if not, at least enforce when you specify them?

Would you expect compiler optimisation to catch this?

[Leonid Kovalev]

I would indeed expect the detection of common subexpressions to be compiler's job. SymPy should not do this by default. But it can do it if you explicitly ask: 

cs = sp.cse(dQ_dR)
print "\n".join(sp.fcode(sub[1], sub[0]) for sub in cs[0]) + "\n" + sp.fcode(cs[1][0], J)

prints

x0 = X_A - X_B
x1 = Y_A - Y_B
x2 = Z_A - Z_B
x3 = 1/sqrt(x0**2 + x1**2 + x2**2)
J(1, 1) = x0*x3
J(1, 2) = x1*x3
J(1, 3) = x2*x3
J(1, 4) = x3*(-X_A + X_B)
J(1, 5) = x3*(-Y_A + Y_B)
J(1, 6) = x3*(-Z_A + Z_B)

[The Pauli Principle]

Thank you, this was almost exactly what I was looking for!

It did not manage to catch

J(1, 4) = x3*(-x0)

Is there any more aggressive 'optimisation' possible?

or would I have to manually try to see if you can substitute

x0 = X_A - X_B

in


x3*(-X_A + X_B)

by trying all possible substitutions (intuition here telling me to loop backwards in this case)?

Not asking for someone else to do it for me, trying to find the right direction, which is always hard when you are not yet familiar with the terminology and don't know what to look for specifically.



[Leonid Kovalev]

The more aggressive (and slower) form is

cs = sp.cse(dQ_dR, optimizations='basic')

which results in  

x0 = X_A - X_B
x1 = Y_A - Y_B
x2 = Z_A - Z_B
x3 = 1/sqrt(x0**2 + x1**2 + x2**2)
x4 = x0*x3
x5 = x1*x3
x6 = x2*x3
J(1, 1) = x4
J(1, 2) = x5
J(1, 3) = x6
J(1, 4) = -x4
J(1, 5) = -x5
J(1, 6) = -x6

Reference: http://docs.sympy.org/dev/modules/rewriting.html#module-sympy.simplify.cse_main