Numerical code generation from symbolic expressions in Python

Fabio Rossi

email: fabio.rossi@inaf.it

Github: https://github.com/FabioRossiArcetri

This work was developed in the context of the simulations for the MAVIS/VLT instrument

MAVIS simulations team @Arcetri:

• AO expertise: Cedric Plantet, Guido Agapito - Thanks for the support!
• Python implementation: Fabio Rossi

“Astrobigné” presentation, Arcetri 06/10/2020

Introduction

• Numpy/SciPy are widely used in numerical codes, at different scales
• NumPy can be seen as an implementation on the CPU of the “Array Programming” paradigm:
  • https://www.nature.com/articles/s41586-020-2649-2#article-info
• Other implementations of the same paradigm have been developed (similar APIs, different backend) :
  • CuPy, GPU, based on NVIDIA CUDA
  • ClPy, GPU/CPU based on openCL
  • Tensorflow (GPU/CPU mainly used for Neural Networks development)

Motivation

• What happens when you have to develop a numerical code?
  • Somewhere the math of the problem is stated, a bunch of expressions, functions, integrals, matrices etc
  • During the implementation phase, the needed math is embedded in the code in more or less clean/modular fashion
• Are the algorithms well apart from the mathematical definitions?
• How reusable is your code?
• Can you easily use the same expression in different algorithms?
• Could you run your code (or just its critical parts) on a GPU? How easily could you port it to use some library that allows it?

Motivation

• Usually numerical codes are developed independently from symbolic functions definition
• SymPy is the Python library implementing symbolic computation
  • it is a CAS, Computer Algebra System like Mathematica, MathCad, Maple, but developed in Python and OpenSource
• Can we benefit by introducing symbolic expression definition and manipulation Numerical Code development cycle?
  • One initial step defining the symbolic functions
  • Symbolic functions can be manipulated in ways numerical functions can’t
  • Sympy Symbolic functions can be translated automatically to numerical functions: "lambdify" mechanism
  • Easily import/export your math from/to other CAS softwares or languages (including LaTex)
• See appendix A for a 5 minutes introduction to SymPy

Approach outline

Software Stack

SEEING: Sympy Expressions Evaluation Implemented oN the GPU

• Github Repo: https://github.com/FabioRossiArcetri/SEEING
  • feel free to checkout, experiment and give feedbacks
• Goals:
  • Factor out common tasks when using SymPy expressions as a starting point for numerical code development
  • Provide helper functions to handle groups of related SymPy expressions
    • evaluate, plot, substitute parameters, share variables
  • Provide a simple way to generate NumPy of CuPy code, so that you quickly evaluate expressions on the GPU
    • Implement functions that are not provided by CuPy or NumPy but do exist in SymPy
  • Develop Backend agnostic Numerical Algorithms
    • use of xp.something calls (where xp can be numpy or cupy)
    • currently numpy and cupy are supported, in the future clpy or other numerical backends might be added
    • for now just a few Numerical Integration methods, over n-dimensional domains

SEEING Overview

• Define a sympy function (or just load some formulas someone else saved to disk)
• Specialize some parameters
• Evaluate a function or an integral

from seeing import *
propMethodsCartesian = Formulary.loadFromFile('Propagation100.frm')
approximations = ["Rayleigh-Sommerfeld", "Approximate Rayleigh-Sommerfeld", "Near Fresnel"]
propMethodsCartesian.display(approximations)

Rayleigh-Sommerfeld

$\displaystyle \operatorname{E_{1}}{\left(x_{1},y_{1},z_{1} \right)} = \int\limits_{- a}^{a}\int\limits_{- a}^{a} - \frac{i E_{0} z_{1} \left(i \lambda \frac{1}{2 \pi \sqrt{z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}}} + 1\right) e^{\frac{2 i \pi \sqrt{z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}}}{\lambda}}}{\lambda \left(z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}\right)}\, dx_{0}\, dy_{0}$

Approximate Rayleigh-Sommerfeld

$\displaystyle \operatorname{E_{1}}{\left(x_{1},y_{1},z_{1} \right)} = \int\limits_{- a}^{a}\int\limits_{- a}^{a} - \frac{i E_{0} z_{1} e^{\frac{2 i \pi \sqrt{z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}}}{\lambda}}}{\lambda \left(z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}\right)}\, dx_{0}\, dy_{0}$

Near Fresnel

$\displaystyle \operatorname{E_{1}}{\left(x_{1},y_{1},z_{1} \right)} = \int\limits_{- a}^{a}\int\limits_{- a}^{a} - \frac{i E_{0} e^{\frac{2 i \pi z_{1}}{\lambda}} e^{\frac{i \pi \left(\left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}\right)}{\lambda z_{1}}}}{\lambda z_{1}}\, dx_{0}\, dy_{0}$

propMethodsCartesian.display(["Rayleigh-Sommerfeld Arg"])

Rayleigh-Sommerfeld Arg

$\displaystyle - \frac{i E_{0} z_{1} \left(i \lambda \frac{1}{2 \pi \sqrt{z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}}} + 1\right) e^{\frac{2 i \pi \sqrt{z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}}}{\lambda}}}{\lambda \left(z_{1}^{2} + \left(- x_{0} + x_{1}\right)^{2} + \left(- y_{0} + y_{1}\right)^{2}\right)}$

waveLength = 10e-6
apertureRadius =  waveLength*10
mCalc = Calculator(cp, cp.float64, 'intensity')
subdiv_points = 512
subsDictC = {'E_0': 1, 'z_1': 2*apertureRadius, 'x_0': 0, 'y_0': apertureRadius, 'lambda': waveLength, 'a':apertureRadius}
paramAndRanges = [( 'x_1', -apertureRadius, apertureRadius, subdiv_points, 'linear' ),
                  ( 'y_1', -apertureRadius, apertureRadius, subdiv_points, 'linear' )]
propMethodsCartesian.plotFormula("Rayleigh-Sommerfeld Arg", subsDictC, paramAndRanges, mCalc)

subdiv_points = 256
ez0_c =  subsParamsByName( propMethodsCartesian.getFormula('xyCircle'), {'a':apertureRadius} )
subsDictC = {'E_0': ez0_c, 'x_1': 0, 'y_1': 0, 'lambda': waveLength, 'a':apertureRadius}
paramAndRange = ( 'z_1', 0.01*apertureRadius, apertureRadius*100, subdiv_points, 'geometric' )
fig, ax = plt.subplots(figsize=(10,5))
plt.xscale('log')
plt.yscale('linear')
for appr in reversed(approximations):
    eeq = propMethodsCartesian.getFormula(appr)
    xplot2, zplot2 = mCalc.IntegralEvalE(subsParamsByName(eeq, subsDictC), [paramAndRange], None, 'trap')
    plt.plot(xplot2[0]/apertureRadius, np.clip(zplot2, -4.5, 4.5), label=appr)
plt.legend(loc=(1.0, 0.78))
plt.show()

approximationsF = ["Rayleigh-Sommerfeld", "Far Fresnel"]
fig, ax = plt.subplots(figsize=(10,5))
plt.xscale('log')
plt.yscale('linear')
paramAndRange = ( 'z_1', apertureRadius, apertureRadius*10, 500, 'geometric' )
subsDictC = {'E_0': ez0_c, 'x_1': 0, 'y_1': 0, 'lambda': waveLength, 'a':apertureRadius}
for appr in reversed(approximationsF):
    eeq = propMethodsCartesian.getFormula(appr)
    xplot2, zplot2 = mCalc.IntegralEvalE(subsParamsByName(eeq, subsDictC), [paramAndRange], [(subdiv_points, 'linear'), (subdiv_points, 'linear')], 'rect')
    plt.plot(xplot2[0]/apertureRadius, np.clip(zplot2, -4.5, 4.5), label=appr)
plt.legend(loc=(1.0, 0.78))
plt.show()

subdiv_points = 256
subsDict = {'E_0': ez0_c, 'y_1': 0, 'lambda': waveLength, 'a':apertureRadius}
paramsAndRanges = [ ('x_1', -1.5*apertureRadius, 1.5*apertureRadius, subdiv_points, 'linear'), ('z_1',  waveLength, 15*apertureRadius, subdiv_points, 'linear')]
eeq = propMethodsCartesian.getFormula("Rayleigh-Sommerfeld")
xplot, fplot1 = mCalc.IntegralEvalE(subsParamsByName(eeq, subsDict), paramsAndRanges) 
fig, ax = plt.subplots(figsize=(10,10))
ax.imshow( fplot1, cmap='hot' )

<matplotlib.image.AxesImage at 0x7f200bdd1e50>

subsDict = {'E_0': ez0_c, 'z_1': 2*apertureRadius, 'lambda': waveLength, 'a':apertureRadius}
paramsAndRanges = [ ('x_1', -2*apertureRadius, 2*apertureRadius, subdiv_points, 'linear'), ('y_1', -2*apertureRadius, 2*apertureRadius, subdiv_points, 'linear')]
eeq = propMethodsCartesian.getFormula("Rayleigh-Sommerfeld")
xplot, fplot1 = mCalc.IntegralEvalE(subsParamsByName(eeq, subsDict), paramsAndRanges) 
fig, ax = plt.subplots(figsize=(10,10))
ax.imshow( np.log(fplot1+1), cmap='hot' )

<matplotlib.image.AxesImage at 0x7f200be0ffd0>

Using SEEING for MAVIS Simulations

• We are using the approach described so far and the SEEING library to develop simulations related to the development of the MAVIS instrument for the VLT telescope

  • Development of a tool for Natural Guide Star asterism selection
  • Synthetic PSF estimation (both for design purposes and observation planning/evaluation)

• Quite a lot, quite huge expressions!

• Example: Covariance of two Zernike Modes

from seeing import *
from zernike import *
mf = Formulary.loadFromFile('Mavis.frm')
mIt = Integrator(cp, cp.float64, 'none')
covIntegral = mf['ZernikeCovarianceI']
display(covIntegral)

$\displaystyle \operatorname{W_{\phi}}{\left(\rho \right)} = \int\limits_{f_{min}}^{f_{max}} \frac{0.0229 \left(-1\right)^{m_{k}} 2^{\left(-0.5\right) \delta_{0 m_{j}} - 0.5 \delta_{0 m_{k}} + 1} i^{n_{j} + n_{k}} \sqrt{\left(n_{j} + 1\right) \left(n_{k} + 1\right)} \left(i^{3 m_{j} + 3 m_{k}} \cos{\left(\theta \left(m_{j} + m_{k}\right) + \frac{\pi \left(\left(1 - \delta_{0 m_{j}}\right) \left(\left(-1\right)^{j} - 1\right) + \left(1 - \delta_{0 m_{k}}\right) \left(\left(-1\right)^{k} - 1\right)\right)}{4} \right)} J_{m_{j} + m_{k}}\left(2 \pi f h \rho\right) + i^{3 \left|{m_{j} - m_{k}}\right|} \cos{\left(\theta \left(m_{j} - m_{k}\right) + \frac{\pi \left(\left(1 - \delta_{0 m_{j}}\right) \left(\left(-1\right)^{j} - 1\right) - \left(1 - \delta_{0 m_{k}}\right) \left(\left(-1\right)^{k} - 1\right)\right)}{4} \right)} J_{\left|{m_{j} - m_{k}}\right|}\left(2 \pi f h \rho\right)\right) J_{n_{j} + 1}\left(2 \pi R_{1} f\right) J_{n_{j} + 1}\left(2 \pi R_{2} f\right)}{\pi R_{1} R_{2} f r_{0}^{\frac{5}{3}} \left(\left(\frac{1}{L_{0}}\right)^{2} + f^{2}\right)^{\frac{11}{6}}}\, df$

r0_Value = 0.15 # [m]
L0_Value = 25.0 # [m]
TelescopeDiameter = 8.0 #[m]
DM_height = 10000 # [m]
rho_max = 120.0/206265.0
paramsDict = {'f_min':0.001, 'f_max':10, 'theta': np.pi/4.0, 'L_0': L0_Value, 
              'r_0': r0_Value, 'R_1': TelescopeDiameter/2.0, 'R_2': TelescopeDiameter/2.0, 'h': DM_height }
_integral1 = subsParamsByName( cov_expr_jk(covIntegral, 2, 3), paramsDict)
display(_integral1)

$\displaystyle \operatorname{W_{\phi}}{\left(\rho \right)} = \int\limits_{0.001}^{10} - \frac{0.135193784634899 J^{2}_{2}\left(8.0 \pi f\right) J_{2}\left(20000 \pi f \rho\right)}{\pi f \left(f^{2} + 0.0016\right)^{\frac{11}{6}}}\, df$

xplot1, zplot1 = mIt.IntegralEvalE(_integral1, [('rho',  0.0, rho_max, 1000, 'linear')], [(1024, 'linear')], method='rect')
fig, ax = plt.subplots(figsize=(10,5))
ax.plot(xplot1[0], zplot1)
plt.show()

Sympy Lambdify

• Sympy provides the lambdify method to transform a sympy expression to a lambda which is also a composition of ufuncs, for different backends ("math", "mpmath", "numpy", "numexpr", "scipy" etc)
• What about a CuPy?
• We can pass to lambdify a dictionary which maps sympy functions to cupy functions
• The dictionary will include many obvious translations like: gpulib = { ... 'sin': cp.sin …

import inspect
from seeing import *

x = sp.symbols('x')
fsympy = sp.sin(x) * sp.exp(x) / (1+sp.sqrt(sp.Abs(x)) ) 
display(fsympy)
fnumpy = sp.lambdify(x, fsympy, "numpy" )
print(inspect.getsource(fnumpy))
print(fnumpy.__globals__['sin'])
fcupy = sp.lambdify(x, fsympy, gpulib )
print(inspect.getsource(fcupy))
print(fcupy.__globals__['sin'])

$\displaystyle \frac{e^{x} \sin{\left(x \right)}}{\sqrt{\left|{x}\right|} + 1}$

def _lambdifygenerated(x):
    return (exp(x)*sin(x)/(sqrt(abs(x)) + 1))

<ufunc 'sin'>
def _lambdifygenerated(x):
    return (exp(x)*sin(x)/(sqrt(abs(x)) + 1))

<ufunc 'cupy_sin'>

Problem

• What about SymPy functions which do not have a CuPy analogue?
  • i.e. CuPy only provides the Bessel Functions of the first order of index 0 and 1... but we can implent and use in the lambdify process our own version for general n

# somewhere we defined:
def besselj__n(n, z):
    if n==0:
        return cupyx.scipy.special.j0(z)
    elif n==1:
        return cupyx.scipy.special.j1(z)
    elif n>=2:
        return 2*(n-1)*besselj__n(n-1, z)/z - besselj__n(int(n)-2, z)

# and in gpulib dictionary we have: ... 'besselj: 'besselj__n' ... 

from seeing import *

n = sp.symbols('n', integer=True)
fsympy = sp.besselj(n, x)
display(fsympy)
fnumpy = sp.lambdify((x, n), fsympy, "scipy" )
print(inspect.getsource(fnumpy))
print(fnumpy.__globals__['jv'])
fcupy = sp.lambdify((x,n), fsympy, gpulib )
print(inspect.getsource(fcupy))
print(fcupy.__globals__['besselj'])

$\displaystyle J_{n}\left(x\right)$

def _lambdifygenerated(x, n):
    return (jv(n, x))

<ufunc 'jv'>
def _lambdifygenerated(x, n):
    return (besselj(n, x))

<function besselj__n at 0x7f202eeefa70>

This is interesting, now what?

• Read the reminder of this presentation, including the Appendixes
• Refresh your NumPy skills (ufuncs, array programming)
  • remember: loops are bad !!!
• Start using SymPy: don’t need to become an expert, just need to able to define your expressions/functions
• Understand basic concepts of GPU programming by experimenting with CuPy: if you know NumPy it’s going to be almost straightforward
• “SEEING” library could help you to easily manipulate/lambdify/evaluate your math on CPU/GPU # Be relieved! You don’t need to learn C/C++ and CUDA if you don't want to!

More details about NumPy, Cupy, Sympy and the development of SEEING

Array programming

• Do as much as possible using Numpy built-in functions and avoid loops and element-wise access
  • understand what ufuncs are
  • See Appeindix C for a glimpse of performance gains
• Performances: Python-like vs C-like !
• Simplyfied transition to different backends (i.e. CuPy)
• Crucial for GPU programming

What is a ufunc

• A (basic) ufunc in NumPy (or CuPy) is a function that can be applied to vectors, working an element-wise basis.
  • for example: np.sin, np.cos, np.power
  • np.meshgrid() is not!
  • np.sum() is a reduction ufunc...
  • more details here: https://numpy.org/doc/stable/reference/ufuncs.html

Array programming on the GPU with CuPy

• CuPy API tries to be as similar as possible to NumPy's, https://docs.cupy.dev/en/stable/
• Comparison table of the two libraries: https://docs.cupy.dev/en/stable/reference/comparison.html
• Important issue to take care of: CPU<->GPU data transfer
  • if the data set you are performing your computation on is quite small, depending on how complex is your computations, it might get less efficient to transfer it to the GPU and get the result back than simply perform your computation on the CPU
• For a 5 minutes introduction to CuPy see Appendix B

What is a lambda?

• A lambda is just a (usually small) anonymous function:

f = lambda a,b,c : a+b*c
print(f(5, 6, 2))

17

# also useful when you want to do have a function that returns a function:
def myfunc(n):
    return lambda a:a*n

mydoubler = myfunc(2)
print(mydoubler(11))

22

Adding some functionalities to SymPy

• When manipulating an expression imported from another module, we would like an easy way to access its symbols, symbols' names

# lets say we have this function returning a SymPy expression defined somewhere
def niceExpr():
    x0, y0, ss = sp.symbols('x_0 y_0 sigma')
    return sp.exp(-x0/ss) + sp.log(1+sp.Abs(y0)*ss)

# then in out main program we could have:
anExpr=niceExpr()
display(anExpr)

$\displaystyle \log{\left(\sigma \left|{y_{0}}\right| + 1 \right)} + e^{- \frac{x_{0}}{\sigma}}$

• just looking at the expression, we can tell there is a parameter named sigma
• we would like to substitute it with some specific real value, how to do it?
• we might not have an easy way to access to the 'ss' variable!!!

print(anExpr.free_symbols)
# cannot do this: ss = anExpr.free_symbols[1]

# we can do this:
ss = sp.symbols('sigma')
anExprSpec = anExpr.subs(ss, 0.5)
display(anExprSpec)

# ugly, what if sigma was complex?

{y_0, sigma, x_0}

$\displaystyle \log{\left(0.5 \left|{y_{0}}\right| + 1 \right)} + e^{- 2.0 x_{0}}$

• We want to have group of formulas sharing the same symbols, this is not obviosly done in SymPy
• We want to implement numerical algorithms that can have in input SymPy expressions and run relying on NumPy or CuPy (same algorithm AND code, different backend)
  • Think about a Monte Carlo integration method that can work using calls like:
    • sp.lambdify(integrand_functions, xp), xp.sum(), xp.random.uniform() etc etc, where xp is a variable point to either np or cp

Appendix A: 5 minutes intro to Sympy

• Define some symbols
• Define a function
• .. or an integral
• Specialize
• Evaluate

import sympy as sp

#Define some symbols
x, y, a, b = sp.symbols("x y a b")
# Define a function
f1 = sp.exp(-a*x**2-b*y**2)
# see how nicely it is displayed
display(f1)

$\displaystyle e^{- a x^{2} - b y^{2}}$

# Or define an integral
i1 = sp.Integral(f1, (x,-sp.oo,+sp.oo), (y,-sp.oo,+sp.oo))
display(i1)

$\displaystyle \int\limits_{-\infty}^{\infty}\int\limits_{-\infty}^{\infty} e^{- a x^{2} - b y^{2}}\, dx\, dy$

# Specialize
i2 = i1.subs([(a, 1), (b,2)])
# Evaluate
display(i2.doit())

$\displaystyle \frac{\sqrt{2} \pi}{2}$

#You could also do real symbolic computations… : 
fdev = sp.diff(f1, y) 
display(fdev)

$\displaystyle - 2 b y e^{- a x^{2} - b y^{2}}$

Appendix B: 5 minutes intro to CuPy

How to:

• Allocate an array, compose and use some ufuncs
• Move data from the CPU to the GPU
• Move data from the GPU to the CPU

import numpy as np
import cupy as cp

npoints = int(1e6)
fcpu = lambda x : np.sum( np.sin(x) * np.exp(x) / (1+np.sqrt(np.abs(x)) ) ) / npoints
x_cpu = np.linspace(-1.0, 1.0, npoints)
result = fcpu(x_cpu) 
print(type(result), result)

x_gpu = cp.linspace(-1.0, 1.0, npoints)
fgpu = lambda x : cp.sum( cp.sin(x) * cp.exp(x) / (1+cp.sqrt(cp.abs(x)) ) ) / npoints
result = fgpu(x_gpu)
print(type(result), result)

<class 'numpy.float64'> 0.17955110645016475
<class 'cupy.core.core.ndarray'> 0.1795511064501647

x_cpu = np.array([1,2,3])
print(type(x_cpu), x_cpu)
x_gpu = cp.asarray(x_cpu)  # move the data to the current device.
print(type(x_gpu), x_gpu)

<class 'numpy.ndarray'> [1 2 3]
<class 'cupy.core.core.ndarray'> [1 2 3]

x_gpu = cp.array([1, 2, 3])  # create an array in the current device
print(type(x_gpu), x_gpu)
x_cpu = cp.asnumpy(x_gpu)  # move the array to the host.
print(type(x_cpu), x_cpu)

<class 'cupy.core.core.ndarray'> [1 2 3]
<class 'numpy.ndarray'> [1 2 3]

Appendix C: Performance comparison: Python vs NumPy vs Cupy

import numpy as np
import cupy as cp
import math

def mySquaredSum(xx):
    r = 0.0
    for i in range(xx.shape[0]):
        r += xx[i]*xx[i]*math.exp(x[i])*math.sin(x[i])**3
    return r

npoints = int(1e5)

x_gpu = cp.linspace(0, 1, npoints)
x = np.linspace(0, 1, npoints)
print(mySquaredSum(x))

26717.919072457928

f = lambda x: np.sum(np.square(x)*np.exp(x)*np.sin(x)**3)
print(f(x))

26717.919072458004

fgpu = lambda x: cp.sum(cp.square(x)*cp.exp(x)*cp.sin(x)**3)
print(fgpu(x_gpu))

26717.919072458

%%timeit
y = mySquaredSum(x)

143 ms ± 3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
y = f(x)

9.28 ms ± 21.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
y = fgpu(x_gpu)

189 µs ± 109 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

paramAndRanges = [( 'x_1', -apertureRadius, apertureRadius, subdiv_points, 'linear' )]
subsDictC = {'E_0': 1, 'z_1': 2*apertureRadius, 'y_1':0.0, 'x_0': 0, 'y_0': apertureRadius, 'lambda': waveLength, 'a':apertureRadius}
propMethodsCartesian.plotFormula("Rayleigh-Sommerfeld Arg", subsDictC, paramAndRanges, mCalc)