Mie Performance and Jitting

Scott Prahl

Feb 2025

Unfortunately, switching between jit and non-jit during runtime is too complicated when combined with numba caching.

[1]:

import os
import tempfile
import numpy as np
import matplotlib.pyplot as plt

os.environ["MIEPYTHON_USE_JIT"] = "1"  # Set to "0" to disable JIT
os.environ["NUMBA_CACHE_DIR"] = tempfile.gettempdir()

import miepython as mie

Size Parameters

We will use %timeit to see speeds for unjitted code, then jitted code

[2]:

ntests = 6

m = 1.5
N = np.logspace(0, 3, ntests, dtype=int)
result = np.zeros(ntests)
resultj = np.zeros(ntests)

for i in range(ntests):
    x = np.linspace(0.1, 20, N[i])
    a = %timeit -o qext, qsca, qback, g = mie.mie(m,x)
    result[i] = a.best

# mie.use_numba(True)  # Ensure the JIT backend is used
# for i in range(ntests):
#     x = np.linspace(0.1, 20, N[i])
#     a = %timeit -o qext, qsca, qback, g = mie.mie(m,x)
#     resultj[i] = a.best

# improvement = result / resultj
# plt.loglog(N, resultj, ":r")
# plt.loglog(N, result, ":b")
# plt.loglog(N, resultj, "or", label="jit")
# plt.loglog(N, result, "ob", label="no jit")
# plt.legend()
# plt.xlabel("Number of sphere sizes calculated")
# plt.ylabel("Execution Time")
# plt.title("Jit improvement is %d to %dX" % (np.min(improvement), np.max(improvement)))
# plt.show()

16 μs ± 259 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
47.2 μs ± 267 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
233 μs ± 763 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
985 μs ± 6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
3.87 ms ± 20 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
15.7 ms ± 302 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Embedded spheres

[ ]:

ntests = 6
mwater = 4 / 3  # rough approximation
m = 1.0
mm = m / mwater
r = 500  # nm

N = np.logspace(0, 3, ntests, dtype=int)
result = np.zeros(ntests)
resultj = np.zeros(ntests)

for i in range(ntests):
    lambda0 = np.linspace(300, 800, N[i])  # also in nm
    xx = 2 * np.pi * r * mwater / lambda0
    a = %timeit -o qext, qsca, qback, g = mie.mie(mm,xx)
    result[i] = a.best

# mie.use_numba(True)  # Ensure the JIT backend is used
# for i in range(ntests):
#     lambda0 = np.linspace(300, 800, N[i])  # also in nm
#     xx = 2 * np.pi * r * mwater / lambda0
#     a = %timeit -o qext, qsca, qback, g = mie.mie(mm,xx)
#     resultj[i] = a.best

# improvement = result / resultj
# plt.loglog(N, resultj, ":r")
# plt.loglog(N, result, ":b")
# plt.loglog(N, resultj, "or", label="jit")
# plt.loglog(N, result, "ob", label="no jit")
# plt.legend()
# plt.xlabel("Number of Wavelengths Calculated")
# plt.ylabel("Execution Time")
# plt.title("Jit improvement is %d to %dX" % (np.min(improvement), np.max(improvement)))
# plt.show()

Testing `ez_mie`

Another high level function that should be sped up by jitting.

[ ]:

ntests = 6
m_sphere = 1.0
n_water = 4 / 3
d = 1000  # nm
N = np.logspace(0, 3, ntests, dtype=int)
result = np.zeros(ntests)
resultj = np.zeros(ntests)

for i in range(ntests):
    lambda0 = np.linspace(300, 800, N[i])  # also in nm
    a = %timeit -o qext, qsca, qback, g = mie.ez_mie(m_sphere, d, lambda0, n_water)
    result[i] = a.best

# mie.use_numba(True)  # Ensure the JIT backend is used
# for i in range(ntests):
#     lambda0 = np.linspace(300, 800, N[i])  # also in nm
#     a = %timeit -o qext, qsca, qback, g = mie.ez_mie(m_sphere, d, lambda0, n_water)
#     resultj[i] = a.best

# improvement = result / resultj
# plt.loglog(N, resultj, ":r")
# plt.loglog(N, result, ":b")
# plt.loglog(N, resultj, "or", label="jit")
# plt.loglog(N, result, "ob", label="no jit")
# plt.legend()
# plt.xlabel("Number of Wavelengths Calculated")
# plt.ylabel("Execution Time")
# plt.title("Jit improvement is %d to %dX" % (np.min(improvement), np.max(improvement)))
# plt.show()

Scattering Phase Function

[ ]:

ntests = 6
m = 1.5
x = np.pi / 3

N = np.logspace(0, 3, ntests, dtype=int)
result = np.zeros(ntests)
resultj = np.zeros(ntests)

for i in range(ntests):
    theta = np.linspace(-180, 180, N[i])
    mu = np.cos(theta / 180 * np.pi)
    a = %timeit -o s1, s2 = mie.mie_S1_S2(m,x,mu)
    result[i] = a.best

# for i in range(ntests):
#     theta = np.linspace(-180, 180, N[i])
#     mu = np.cos(theta / 180 * np.pi)
#     a = %timeit -o s1, s2 = mie.mie_S1_S2(m,x,mu)
#     resultj[i] = a.best

# improvement = result / resultj
# plt.loglog(N, resultj, ":r")
# plt.loglog(N, result, ":b")
# plt.loglog(N, resultj, "or", label="jit")
# plt.loglog(N, result, "ob", label="no jit")
# plt.legend()
# plt.xlabel("Number of Angles Calculated")
# plt.ylabel("Execution Time")
# plt.title("Jit improvement is %d to %dX" % (np.min(improvement), np.max(improvement)))
# plt.show()

And finally, as function of sphere size

[ ]:

ntests = 6
m = 1.5 - 0.1j
x = np.logspace(0, 3, ntests)
result = np.zeros(ntests)
resultj = np.zeros(ntests)

theta = np.linspace(-180, 180)
mu = np.cos(theta / 180 * np.pi)

for i in range(ntests):
    a = %timeit -o s1, s2 = mie.mie_S1_S2(m,x[i],mu)
    result[i] = a.best

# mie.use_numba(True)  # Ensure the JIT backend is used
# for i in range(ntests):
#     a = %timeit -o s1, s2 = mie.mie_S1_S2(m,x[i],mu)
#     resultj[i] = a.best

# improvement = result / resultj
# plt.loglog(N, resultj, ":r")
# plt.loglog(N, result, ":b")
# plt.loglog(N, resultj, "or", label="jit")
# plt.loglog(N, result, "ob", label="no jit")
# plt.legend()
# plt.xlabel("Sphere Size Parameter")
# plt.ylabel("Execution Time")
# plt.title("Jit improvement is %d to %dX" % (np.min(improvement), np.max(improvement)))
# plt.show()

[ ]: