# Monte Carlo using Numba

I'm having trouble with the slow computation of my Monte Carlo simulation code. Based on the pycallgraph, the bottleneck seems to be the module named miepython.mie_S1_S2 (highlighted by pink), which takes around 0.5 seconds per call.

Below is an example of the implementation of miepython.mie_S1_S2.

m = 1.336-2.462e-09j
x = 8526.95
mu = np.array([-1., -0.7500396, 0.46037385, 0.5988121, 0.67384093, 0.72468684, 0.76421644, 0.79175856, 0.81723714, 0.83962897, 0.85924182, 0.87641596, 0.89383665, 0.90708978, 0.91931481, 0.93067567, 0.94073113, 0.94961222, 0.95689496, 0.96467123,  0.97138347, 0.97791831, 0.98339434, 0.98870543, 0.99414948, 0.9975728   0.9989995, 0.9989995, 0.9989995, 0.9989995, 0.9989995,0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899952,  0.99899952,
0.99899952,  0.99899952,  0.99899952,  0.99899952,  0.99899952, 0.99899952, 0.99899952,  1.        ])

result = mie_S1_S2(m, x, mu)


I somehow want to speed up this mie_S1_S2 function. The source code of this function is like this:

def mie_S1_S2(m, x, mu):
if np.isscalar(mu):
mu_array = np.array([mu], dtype=float)
s1, s2 = _mie_S1_S2(m, x, mu_array)
return s1[0], s2[0]

return _mie_S1_S2(m, x, mu)


and helper functions are like below:

import numpy as np
from numba import njit, int32, float64, complex128

__all__ = ('ez_mie',
'ez_intensities',
'generate_mie_costheta',
'i_par',
'i_per',
'i_unpolarized',
'mie',
'mie_S1_S2',
'mie_cdf',
'mie_mu_with_uniform_cdf',
)

@njit((complex128, int32), cache=True)
def _Lentz_Dn(z, N):
"""
Compute the logarithmic derivative of the Ricatti-Bessel function.

Args:
z: function argument
N: order of Ricatti-Bessel function

Returns:
This returns the Ricatti-Bessel function of order N with argument z
using the continued fraction technique of Lentz, Appl. Opt., 15,
668-671, (1976).
"""
zinv = 2.0 / z
alpha = (N + 0.5) * zinv
aj = -(N + 1.5) * zinv
alpha_j1 = aj + 1 / alpha
alpha_j2 = aj
ratio = alpha_j1 / alpha_j2
runratio = alpha * ratio

while np.abs(np.abs(ratio) - 1.0) > 1e-12:
aj = zinv - aj
alpha_j1 = 1.0 / alpha_j1 + aj
alpha_j2 = 1.0 / alpha_j2 + aj
ratio = alpha_j1 / alpha_j2
zinv *= -1
runratio = ratio * runratio

return -N / z + runratio

@njit((complex128, int32, complex128[:]), cache=True)
def _D_downwards(z, N, D):
"""
Compute the logarithmic derivative by downwards recurrence.

Args:
z: function argument
N: order of Ricatti-Bessel function
D: gets filled with the Ricatti-Bessel function values for orders
from 0 to N for an argument z using the downwards recurrence relations.
"""
last_D = _Lentz_Dn(z, N)
for n in range(N, 0, -1):
last_D = n / z - 1.0 / (last_D + n / z)
D[n - 1] = last_D

@njit((complex128, int32, complex128[:]), cache=True)
def _D_upwards(z, N, D):
"""
Compute the logarithmic derivative by upwards recurrence.

Args:
z: function argument
N: order of Ricatti-Bessel function
D: gets filled with the Ricatti-Bessel function values for orders
from 0 to N for an argument z using the upwards recurrence relations.
"""
exp = np.exp(-2j * z)
D[1] = -1 / z + (1 - exp) / ((1 - exp) / z - 1j * (1 + exp))
for n in range(2, N):
D[n] = 1 / (n / z - D[n - 1]) - n / z

@njit((complex128, float64, int32), cache=True)
def _D_calc(m, x, N):
"""
Compute the logarithmic derivative using best method.

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere
N: order of Ricatti-Bessel function

Returns:
The values of the Ricatti-Bessel function for orders from 0 to N.
"""
n = m.real
kappa = np.abs(m.imag)
D = np.zeros(N, dtype=np.complex128)

if n < 1 or n > 10 or kappa > 10 or x * kappa >= 3.9 - 10.8 * n + 13.78 * n**2:
_D_downwards(m * x, N, D)
else:
_D_upwards(m * x, N, D)
return D

@njit((complex128, float64, complex128[:], complex128[:]), cache=True)
def _mie_An_Bn(m, x, a, b):
"""
Compute arrays of Mie coefficients A and B for a sphere.

This estimates the size of the arrays based on Wiscombe's formula. The length
of the arrays is chosen so that the error when the series are summed is
around 1e-6.

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere

Returns:
An, Bn: arrays of Mie coefficents
"""
psi_nm1 = np.sin(x)                   # nm1 = n-1 = 0
psi_n = psi_nm1 / x - np.cos(x)       # n = 1
xi_nm1 = complex(psi_nm1, np.cos(x))
xi_n = complex(psi_n, np.cos(x) / x + np.sin(x))

nstop = len(a)
if m.real > 0.0:
D = _D_calc(m, x, nstop + 1)

for n in range(1, nstop):
temp = D[n] / m + n / x
a[n - 1] = (temp * psi_n - psi_nm1) / (temp * xi_n - xi_nm1)
temp = D[n] * m + n / x
b[n - 1] = (temp * psi_n - psi_nm1) / (temp * xi_n - xi_nm1)
xi = (2 * n + 1) * xi_n / x - xi_nm1
xi_nm1 = xi_n
xi_n = xi
psi_nm1 = psi_n
psi_n = xi_n.real

else:
for n in range(1, nstop):
a[n - 1] = (n * psi_n / x - psi_nm1) / (n * xi_n / x - xi_nm1)
b[n - 1] = psi_n / xi_n
xi = (2 * n + 1) * xi_n / x - xi_nm1
xi_nm1 = xi_n
xi_n = xi
psi_nm1 = psi_n
psi_n = xi_n.real

@njit((complex128, float64, float64[:]), cache=True)
def _mie_S1_S2(m, x, mu):
"""
Calculate the scattering amplitude functions for spheres.

The amplitude functions have been normalized so that when integrated
over all 4*pi solid angles, the integral will be qext*pi*x**2.

The units are weird, sr**(-0.5)

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere
mu: array of angles, cos(theta), to calculate scattering amplitudes

Returns:
S1, S2: the scattering amplitudes at each angle mu [sr**(-0.5)]
"""
nstop = int(x + 4.05 * x**0.33333 + 2.0) + 1
a = np.zeros(nstop - 1, dtype=np.complex128)
b = np.zeros(nstop - 1, dtype=np.complex128)
_mie_An_Bn(m, x, a, b)

nangles = len(mu)
S1 = np.zeros(nangles, dtype=np.complex128)
S2 = np.zeros(nangles, dtype=np.complex128)

nstop = len(a)
for k in range(nangles):
pi_nm2 = 0
pi_nm1 = 1
for n in range(1, nstop):
tau_nm1 = n * mu[k] * pi_nm1 - (n + 1) * pi_nm2

S1[k] += (2 * n + 1) * (pi_nm1 * a[n - 1]
+ tau_nm1 * b[n - 1]) / (n + 1) / n

S2[k] += (2 * n + 1) * (tau_nm1 * a[n - 1]
+ pi_nm1 * b[n - 1]) / (n + 1) / n

temp = pi_nm1
pi_nm1 = ((2 * n + 1) * mu[k] * pi_nm1 - (n + 1) * pi_nm2) / n
pi_nm2 = temp

# calculate norm = sqrt(pi * Qext * x**2)
n = np.arange(1, nstop + 1)
norm = np.sqrt(2 * np.pi * np.sum((2 * n + 1) * (a.real + b.real)))

S1 /= norm
S2 /= norm

return [S1, S2]


These functions are using Numba, which I believe is the fastest way to compute numpy arrays and for loops on Python. However, it currently takes 0.5 seconds per call even though I need to call it a million times. I'm wondering something may be hindering the performance of Numba. Is there anything that can be done to improve these functions?

• Please clarify your specific problem or provide additional details to highlight exactly what you need. As it's currently written, it's hard to tell exactly what you're asking.
– Community Bot
Sep 19, 2021 at 17:09
• At least explain what the code does and how it works, as per site rules. Sep 19, 2021 at 17:27
• Sorry about that. I added more explanations to my post. Let me know if there is anything missing. Sep 19, 2021 at 17:45
• You say it's a Monte Carlo simulation, but haven't said what you're simulating. What's the purpose of the program? When you have included that, please replace the title to state the purpose - see How to get the best value out of Code Review: Asking Questions for guidance on writing good question titles. Sep 21, 2021 at 7:49

By a very wide margin, the slowest part of your code is this loop:

for k in range(nangles):
pi_nm2 = 0
pi_nm1 = 1
for n in range(1, nstop):
tau_nm1 = n * mu[k] * pi_nm1 - (n + 1) * pi_nm2

S1[k] += (2 * n + 1) * (pi_nm1 * a[n - 1]
+ tau_nm1 * b[n - 1]) / (n + 1) / n

S2[k] += (2 * n + 1) * (tau_nm1 * a[n - 1]
+ pi_nm1 * b[n - 1]) / (n + 1) / n

temp = pi_nm1
pi_nm1 = ((2 * n + 1) * mu[k] * pi_nm1 - (n + 1) * pi_nm2) / n
pi_nm2 = temp


It has no vectorization at all. The trivial improvement is to

• Keep the initialization of a pi vector in a loop, but vectorize all other expressions
• Have vectors for pi, n, tau, an intermediate expression for $$\ \frac {2n + 1} {n(n + 1)}\$$, and the addends s1 and s2
• Sum over the addends and initialize the given element for S1 and S2

To go any further with vectorization, you need to solve the recurrence relation in pi which I don't know how to do. Despite this, the trivial improvement is much, much faster than the original code, and has not produced any numeric regressions in my testing. Also, add type hints. I've commented out your JIT decorators for debugging purposes.

from numbers import Real
from typing import Tuple

import numpy as np
from numba import njit, int32, float64, complex128

# @njit((complex128, int32), cache=True)
def _Lentz_Dn(z: complex, N: int) -> complex:
"""
Compute the logarithmic derivative of the Ricatti-Bessel function.

Args:
z: function argument
N: order of Ricatti-Bessel function

Returns:
This returns the Ricatti-Bessel function of order N with argument z
using the continued fraction technique of Lentz, Appl. Opt., 15,
668-671, (1976).
"""
zinv = 2.0 / z
alpha = (N + 0.5) * zinv
aj = -(N + 1.5) * zinv
alpha_j1 = aj + 1 / alpha
alpha_j2 = aj
ratio = alpha_j1 / alpha_j2
runratio = alpha * ratio

while np.abs(np.abs(ratio) - 1.0) > 1e-12:
aj = zinv - aj
alpha_j1 = 1.0 / alpha_j1 + aj
alpha_j2 = 1.0 / alpha_j2 + aj
ratio = alpha_j1 / alpha_j2
zinv *= -1
runratio = ratio * runratio

return -N / z + runratio

# @njit((complex128, int32, complex128[:]), cache=True)
def _D_downwards(z: complex, N: int, D: np.ndarray) -> None:
"""
Compute the logarithmic derivative by downwards recurrence.

Args:
z: function argument
N: order of Ricatti-Bessel function
D: gets filled with the Ricatti-Bessel function values for orders
from 0 to N for an argument z using the downwards recurrence relations.
"""
last_D = _Lentz_Dn(z, N)
for n in range(N, 0, -1):
last_D = n / z - 1.0 / (last_D + n / z)
D[n - 1] = last_D

# @njit((complex128, int32, complex128[:]), cache=True)
def _D_upwards(z: complex, N: int, D: np.ndarray) -> None:
"""
Compute the logarithmic derivative by upwards recurrence.

Args:
z: function argument
N: order of Ricatti-Bessel function
D: gets filled with the Ricatti-Bessel function values for orders
from 0 to N for an argument z using the upwards recurrence relations.
"""
exp = np.exp(-2j * z)
D[1] = -1 / z + (1 - exp) / ((1 - exp) / z - 1j * (1 + exp))
for n in range(2, N):
D[n] = 1 / (n / z - D[n - 1]) - n / z

# @njit((complex128, float64, int32), cache=True)
def _D_calc(m: complex, x: Real, N: int) -> np.ndarray:
"""
Compute the logarithmic derivative using best method.

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere
N: order of Ricatti-Bessel function

Returns:
The values of the Ricatti-Bessel function for orders from 0 to N.
"""
n = m.real
kappa = np.abs(m.imag)
D = np.zeros(N, dtype=np.complex128)

if n < 1 or n > 10 or kappa > 10 or x * kappa >= 3.9 - 10.8 * n + 13.78 * n**2:
_D_downwards(m * x, N, D)
else:
_D_upwards(m * x, N, D)
return D

# @njit((complex128, float64, complex128[:], complex128[:]), cache=True)
def _mie_An_Bn(m: complex, x: Real, a: np.ndarray, b: np.ndarray) -> None:
"""
Compute arrays of Mie coefficients A and B for a sphere.

This estimates the size of the arrays based on Wiscombe's formula. The length
of the arrays is chosen so that the error when the series are summed is
around 1e-6.

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere

Returns:
An, Bn: arrays of Mie coefficents
"""
psi_nm1 = np.sin(x)                   # nm1 = n-1 = 0
psi_n = psi_nm1 / x - np.cos(x)       # n = 1
xi_nm1 = complex(psi_nm1, np.cos(x))
xi_n = complex(psi_n, np.cos(x) / x + np.sin(x))

nstop = len(a)
if m.real > 0.0:
D = _D_calc(m, x, nstop + 1)

for n in range(1, nstop):
temp = D[n] / m + n / x
a[n - 1] = (temp * psi_n - psi_nm1) / (temp * xi_n - xi_nm1)
temp = D[n] * m + n / x
b[n - 1] = (temp * psi_n - psi_nm1) / (temp * xi_n - xi_nm1)
xi = (2 * n + 1) * xi_n / x - xi_nm1
xi_nm1 = xi_n
xi_n = xi
psi_nm1 = psi_n
psi_n = xi_n.real

else:
for n in range(1, nstop):
a[n - 1] = (n * psi_n / x - psi_nm1) / (n * xi_n / x - xi_nm1)
b[n - 1] = psi_n / xi_n
xi = (2 * n + 1) * xi_n / x - xi_nm1
xi_nm1 = xi_n
xi_n = xi
psi_nm1 = psi_n
psi_n = xi_n.real

# @njit((complex128, float64, float64[:]), cache=True)
def _mie_S1_S2(m: complex, x: Real, mu: np.ndarray) -> Tuple[
np.ndarray,  # S1
np.ndarray,  # S2
]:
"""
Calculate the scattering amplitude functions for spheres.

The amplitude functions have been normalized so that when integrated
over all 4*pi solid angles, the integral will be qext*pi*x**2.

The units are weird, sr**(-0.5)

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere
mu: array of angles, cos(theta), to calculate scattering amplitudes

Returns:
S1, S2: the scattering amplitudes at each angle mu [sr**(-0.5)]
"""
nstop = int(x + 4.05 * x**(1/3) + 2.0) + 1
a = np.zeros(nstop - 1, dtype=np.complex128)
b = np.zeros(nstop - 1, dtype=np.complex128)
_mie_An_Bn(m, x, a, b)

nangles = len(mu)
S1 = np.zeros(nangles, dtype=np.complex128)
S2 = np.zeros(nangles, dtype=np.complex128)

nstop = len(a)  # 8611
for k in range(nangles):  # 50
pi = np.empty(nstop)
pi[0] = 0
pi[1] = 1
for n in range(1, nstop-1):
pi[n+1] = ((2*n + 1) * mu[k] * pi[n] - (n + 1)*pi[n - 1]) / n

n = np.arange(1, nstop)
tau = n * mu[k] * pi[1:] - (n + 1) * pi[:-1]
pi = pi[1:]

nfac = (2*n + 1)/(n + 1)/n
s1 = nfac * ( pi * a[:-1] + tau * b[:-1])
s2 = nfac * (tau * a[:-1] +  pi * b[:-1])
S1[k] = np.sum(s1)
S2[k] = np.sum(s2)

# calculate norm = sqrt(pi * Qext * x**2)
n = np.arange(1, nstop + 1)
norm = np.sqrt(2 * np.pi * np.sum((2 * n + 1) * (a.real + b.real)))

S1 /= norm
S2 /= norm

return S1, S2

def mie_S1_S2(m: complex, x: Real, mu: np.ndarray) -> Tuple[
np.ndarray,  # S1
np.ndarray,  # S2
]:
if np.isscalar(mu):
mu_array = np.array([mu], dtype=float)
s1, s2 = _mie_S1_S2(m, x, mu_array)
return s1[0], s2[0]

return _mie_S1_S2(m, x, mu)

def test():
m = 1.336 - 2.462e-09j
x = 8526.95
mu = np.array([
-1., -0.75003960,
0.46037385, 0.59881210, 0.67384093, 0.72468684, 0.76421644, 0.79175856,
0.81723714, 0.83962897, 0.85924182, 0.87641596, 0.89383665, 0.90708978,
0.91931481, 0.93067567, 0.94073113, 0.94961222, 0.95689496, 0.96467123,
0.97138347, 0.97791831, 0.98339434, 0.98870543, 0.99414948, 0.99757280,
0.99899950, 0.99899950, 0.99899950, 0.99899950, 0.99899950, 0.99899951,
0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951,
0.99899951, 0.99899951, 0.99899952, 0.99899952, 0.99899952, 0.99899952,
0.99899952, 0.99899952, 0.99899952, 0.99899952, 0.99899952, 1.
])

s1, s2 = mie_S1_S2(m, x, mu)
assert s1.shape == (50,)
assert s2.shape == (50,)

assert np.all(np.isclose(
s1,
[
+1.59632905e-01 + 0.03502544j, +3.90504470e-01 + 0.20900939j,
-9.97914478e-02 + 0.02620006j, +9.21026962e-02 + 0.06207354j,
+2.10094434e-01 - 0.04797237j, +1.05940856e-01 + 0.13947792j,
-4.01215559e-01 - 0.16688863j, +2.00543669e-01 - 0.21176728j,
+2.35794245e-01 - 0.17652092j, +1.49608213e-01 + 0.35291572j,
+4.55588556e-01 - 0.13719540j, -3.31926579e-02 - 0.38015050j,
+5.76677965e-01 + 0.11300876j, +4.41880057e-01 + 0.02616034j,
-3.15436450e-01 + 0.23299934j, -2.64993116e-01 + 0.42528377j,
+3.91626422e-01 - 0.50254770j, +3.80587284e-01 - 0.35133698j,
-5.86213554e-01 + 0.11891230j, -6.92705719e-01 - 0.32413945j,
-4.40737481e-01 + 0.67130149j, +6.27519196e-01 - 0.56104087j,
+8.04106398e-01 - 0.23357289j, -3.80413473e-02 - 0.90094431j,
+4.05299405e-01 + 0.47019625j, -3.16373362e-02 + 0.98316482j,
+2.85522025e-01 + 0.52299093j, +2.85522025e-01 + 0.52299093j,
+2.85522025e-01 + 0.52299093j, +2.85522025e-01 + 0.52299093j,
+2.85522025e-01 + 0.52299093j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +1.70304446e+03 - 6.84700064j,
],
atol=0, rtol=1e-7,
))

assert np.all(np.isclose(
s2,
[
-1.59632905e-01 - 0.03502544j, +8.84797851e-02 + 0.05820884j,
-7.90527129e-02 + 0.08837562j, +1.71665708e-01 + 0.01900138j,
+2.33875478e-01 + 0.01646697j, +1.08591324e-01 + 0.22891712j,
-3.61692929e-01 - 0.11667708j, +2.60383195e-01 - 0.19467110j,
+2.63947149e-01 - 0.23064441j, +1.98888391e-01 + 0.34604412j,
+4.50369063e-01 - 0.08758544j, -9.92951922e-03 - 0.40840802j,
+5.41458693e-01 + 0.10689114j, +4.79442051e-01 + 0.07059381j,
-3.79039180e-01 + 0.24306640j, -3.19135376e-01 + 0.42468406j,
+3.54480077e-01 - 0.52377100j, +4.29151820e-01 - 0.33160320j,
-5.99542307e-01 + 0.07498510j, -6.62025318e-01 - 0.33146038j,
-4.15698108e-01 + 0.65755116j, +5.96137922e-01 - 0.55023812j,
+7.77715408e-01 - 0.24167316j, -2.78492552e-02 - 0.87653004j,
+3.94728230e-01 + 0.47730399j, -1.26273416e-02 + 0.96551473j,
+3.39065634e-01 + 0.53203896j, +3.39065634e-01 + 0.53203896j,
+3.39065634e-01 + 0.53203896j, +3.39065634e-01 + 0.53203896j,
+3.39065634e-01 + 0.53203896j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +1.70304446e+03 - 6.84700064j,
],
atol=0, rtol=1e-7,
))

if __name__ == '__main__':
test()


For me this executed in 450 ms. By contrast, the old method executes in 5195 ms.

## Edit

Vectorize more of the functions, and broaden the first vectorization to another dimension. That reduces this to 83 ms on my laptop:

from numbers import Real
from timeit import timeit
from typing import Tuple, Union

import numpy as np
from numba import njit, int32, float64, complex128

# @njit((complex128, int32), cache=True)
def _Lentz_Dn(z: complex, N: int) -> complex:
"""
Compute the logarithmic derivative of the Ricatti-Bessel function.

Args:
z: function argument
N: order of Ricatti-Bessel function

Returns:
This returns the Ricatti-Bessel function of order N with argument z
using the continued fraction technique of Lentz, Appl. Opt., 15,
668-671, (1976).
"""
zinv = 2.0 / z
alpha = (N + 0.5) * zinv
aj = -(N + 1.5) * zinv
alpha_j1 = aj + 1 / alpha
alpha_j2 = aj
ratio = alpha_j1 / alpha_j2
runratio = alpha * ratio

while np.abs(np.abs(ratio) - 1.0) > 1e-12:
aj = zinv - aj
alpha_j1 = 1.0 / alpha_j1 + aj
alpha_j2 = 1.0 / alpha_j2 + aj
ratio = alpha_j1 / alpha_j2
zinv *= -1
runratio = ratio * runratio

return -N / z + runratio

# @njit((complex128, int32, complex128[:]), cache=True)
def _D_downwards(z: complex, N: int, D: np.ndarray) -> None:
"""
Compute the logarithmic derivative by downwards recurrence.

Args:
z: function argument
N: order of Ricatti-Bessel function
D: gets filled with the Ricatti-Bessel function values for orders
from 0 to N for an argument z using the downwards recurrence relations.
"""
last_D = _Lentz_Dn(z, N)
for n in range(N, 0, -1):
last_D = n/z - 1.0/(last_D + n/z)
D[n - 1] = last_D

# @njit((complex128, int32, complex128[:]), cache=True)
def _D_upwards(z: complex, N: int, D: np.ndarray) -> None:
"""
Compute the logarithmic derivative by upwards recurrence.

Args:
z: function argument
N: order of Ricatti-Bessel function
D: gets filled with the Ricatti-Bessel function values for orders
from 0 to N for an argument z using the upwards recurrence relations.
"""
exp = np.exp(-2j * z)
D[0] = 0
D[1] = 1/(1/z - 1j*(1 + exp)/(1 - exp)) - 1/z
for n in range(2, N):
D[n] = 1/(n/z - D[n-1]) - n/z

# @njit((complex128, float64, int32), cache=True)
def _D_calc(m: complex, x: Real, N: int) -> np.ndarray:
"""
Compute the logarithmic derivative using best method.

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere
N: order of Ricatti-Bessel function

Returns:
The values of the Ricatti-Bessel function for orders from 0 to N.
"""
n = m.real
kappa = np.abs(m.imag)
D = np.empty(N, dtype=np.complex128)

if not (0 < n <= 10) or kappa > 10 or x * kappa >= 3.9 - 10.8*n + 13.78*n**2:
_D_downwards(m * x, N, D)
else:
_D_upwards(m * x, N, D)
return D

# @njit((complex128, float64, complex128[:], complex128[:]), cache=True)
def _mie_An_Bn(m: complex, x: Real, a: np.ndarray, b: np.ndarray) -> None:
"""
Compute arrays of Mie coefficients A and B for a sphere.

This estimates the size of the arrays based on Wiscombe's formula. The length
of the arrays is chosen so that the error when the series are summed is
around 1e-6.

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere

Returns:
An, Bn: arrays of Mie coefficents
"""
nstop = len(a)

psi = np.empty(nstop)
psi[0] = np.sin(x)                  # nm1 = n-1 = 0
psi[1] = psi[0]/x - np.cos(x)       # n = 1

xi = np.empty(nstop, dtype=np.complex128)
xi[0] = complex(psi[0], np.cos(x))
xi[1] = complex(psi[1], np.cos(x)/x + np.sin(x))

for n in range(1, nstop-1):
xi[n+1] = (2*n + 1) * xi[n]/x - xi[n-1]

psi[2:] = np.real(xi[2:])

if m.real > 0.0:
D = _D_calc(m, x, nstop)[1:]
n = np.arange(1, nstop)

def fill(Dm, dest):
temp = Dm + n/x
dest[:-1] = (temp*psi[1:] - psi[:-1]) / (temp*xi[1:] - xi[:-1])
dest[-1] = 0

fill(D/m, a)
fill(D*m, b)

else:
# for n in range(1, nstop):
#     a[n - 1] = (n * psi_n / x - psi_nm1) / (n * xi_n / x - xi_nm1)
#     b[n - 1] = psi_n / xi_n
#     xi = (2 * n + 1) * xi_n / x - xi_nm1
#     xi_nm1 = xi_n
#     xi_n = xi
#     psi_nm1 = psi_n
#     psi_n = xi_n.real
raise NotImplementedError(
'Not vectorized yet since this code path is currently not hit'
)

# @njit((complex128, float64, float64[:]), cache=True)
def _mie_S1_S2(m: complex, x: Real, mu: np.ndarray) -> Tuple[
np.ndarray,  # S1
np.ndarray,  # S2
]:
"""
Calculate the scattering amplitude functions for spheres.

The amplitude functions have been normalized so that when integrated
over all 4*pi solid angles, the integral will be qext*pi*x**2.

The units are weird, sr**(-0.5)

Args:
m: the complex index of refraction of the sphere
x: the size parameter of the sphere
mu: array of angles, cos(theta), to calculate scattering amplitudes

Returns:
S1, S2: the scattering amplitudes at each angle mu [sr**(-0.5)]
"""
nstop = int(x + 4.05 * x**(1/3) + 2.0) + 1
a = np.empty(nstop - 1, dtype=np.complex128)
b = np.empty_like(a)
_mie_An_Bn(m, x, a, b)

nangles = len(mu)
nstop = len(a)  # 8611

pi = np.empty((nangles, nstop))
pi[:, 0] = 0
pi[:, 1] = 1
for n in range(1, nstop-1):
pi[:, n+1] = ((2*n + 1) * mu * pi[:,n] - (n + 1)*pi[:,n-1]) / n

n = np.arange(1, nstop)
tau = n * pi[:,1:] * mu[:, np.newaxis] - (n + 1) * pi[:,:-1]
pi = pi[:,1:]

nfac = (2*n + 1)/(n + 1)/n
s1 = nfac * ( pi * a[:-1] + tau * b[:-1])
s2 = nfac * (tau * a[:-1] +  pi * b[:-1])
S1 = np.sum(s1, axis=1)
S2 = np.sum(s2, axis=1)

# calculate norm = sqrt(pi * Qext * x**2)
n = np.arange(1, nstop + 1)
norm = np.sqrt(2 * np.pi * np.sum((2*n + 1) * (a.real + b.real)))

S1 /= norm
S2 /= norm

return S1, S2

def mie_S1_S2(m: complex, x: Real, mu: Union[Real, np.ndarray]) -> Tuple[
np.ndarray,  # S1
np.ndarray,  # S2
]:
if np.isscalar(mu):
mu_array = np.array([mu], dtype=float)
s1, s2 = _mie_S1_S2(m, x, mu_array)
return s1[0], s2[0]

return _mie_S1_S2(m, x, mu)

def test():
m = 1.336 - 2.462e-09j
x = 8526.95
mu = np.array([
-1., -0.75003960,
0.46037385, 0.59881210, 0.67384093, 0.72468684, 0.76421644, 0.79175856,
0.81723714, 0.83962897, 0.85924182, 0.87641596, 0.89383665, 0.90708978,
0.91931481, 0.93067567, 0.94073113, 0.94961222, 0.95689496, 0.96467123,
0.97138347, 0.97791831, 0.98339434, 0.98870543, 0.99414948, 0.99757280,
0.99899950, 0.99899950, 0.99899950, 0.99899950, 0.99899950, 0.99899951,
0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951, 0.99899951,
0.99899951, 0.99899951, 0.99899952, 0.99899952, 0.99899952, 0.99899952,
0.99899952, 0.99899952, 0.99899952, 0.99899952, 0.99899952, 1.
])

s1, s2 = None, None
def f():
nonlocal s1, s2
s1, s2 = mie_S1_S2(m, x, mu)
print(timeit(f, number=1))

assert s1.shape == (50,)
assert s2.shape == (50,)

assert np.all(np.isclose(
s1,
[
+1.59632905e-01 + 0.03502544j, +3.90504470e-01 + 0.20900939j,
-9.97914478e-02 + 0.02620006j, +9.21026962e-02 + 0.06207354j,
+2.10094434e-01 - 0.04797237j, +1.05940856e-01 + 0.13947792j,
-4.01215559e-01 - 0.16688863j, +2.00543669e-01 - 0.21176728j,
+2.35794245e-01 - 0.17652092j, +1.49608213e-01 + 0.35291572j,
+4.55588556e-01 - 0.13719540j, -3.31926579e-02 - 0.38015050j,
+5.76677965e-01 + 0.11300876j, +4.41880057e-01 + 0.02616034j,
-3.15436450e-01 + 0.23299934j, -2.64993116e-01 + 0.42528377j,
+3.91626422e-01 - 0.50254770j, +3.80587284e-01 - 0.35133698j,
-5.86213554e-01 + 0.11891230j, -6.92705719e-01 - 0.32413945j,
-4.40737481e-01 + 0.67130149j, +6.27519196e-01 - 0.56104087j,
+8.04106398e-01 - 0.23357289j, -3.80413473e-02 - 0.90094431j,
+4.05299405e-01 + 0.47019625j, -3.16373362e-02 + 0.98316482j,
+2.85522025e-01 + 0.52299093j, +2.85522025e-01 + 0.52299093j,
+2.85522025e-01 + 0.52299093j, +2.85522025e-01 + 0.52299093j,
+2.85522025e-01 + 0.52299093j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.85900437e-01 + 0.52348502j, +2.85900437e-01 + 0.52348502j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +2.86280222e-01 + 0.52397854j,
+2.86280222e-01 + 0.52397854j, +1.70304446e+03 - 6.84700064j,
],
atol=0, rtol=1e-7,
))

assert np.all(np.isclose(
s2,
[
-1.59632905e-01 - 0.03502544j, +8.84797851e-02 + 0.05820884j,
-7.90527129e-02 + 0.08837562j, +1.71665708e-01 + 0.01900138j,
+2.33875478e-01 + 0.01646697j, +1.08591324e-01 + 0.22891712j,
-3.61692929e-01 - 0.11667708j, +2.60383195e-01 - 0.19467110j,
+2.63947149e-01 - 0.23064441j, +1.98888391e-01 + 0.34604412j,
+4.50369063e-01 - 0.08758544j, -9.92951922e-03 - 0.40840802j,
+5.41458693e-01 + 0.10689114j, +4.79442051e-01 + 0.07059381j,
-3.79039180e-01 + 0.24306640j, -3.19135376e-01 + 0.42468406j,
+3.54480077e-01 - 0.52377100j, +4.29151820e-01 - 0.33160320j,
-5.99542307e-01 + 0.07498510j, -6.62025318e-01 - 0.33146038j,
-4.15698108e-01 + 0.65755116j, +5.96137922e-01 - 0.55023812j,
+7.77715408e-01 - 0.24167316j, -2.78492552e-02 - 0.87653004j,
+3.94728230e-01 + 0.47730399j, -1.26273416e-02 + 0.96551473j,
+3.39065634e-01 + 0.53203896j, +3.39065634e-01 + 0.53203896j,
+3.39065634e-01 + 0.53203896j, +3.39065634e-01 + 0.53203896j,
+3.39065634e-01 + 0.53203896j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39499038e-01 + 0.53246355j, +3.39499038e-01 + 0.53246355j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +3.39933654e-01 + 0.53288754j,
+3.39933654e-01 + 0.53288754j, +1.70304446e+03 - 6.84700064j,
],
atol=0, rtol=1e-7,
))

if __name__ == '__main__':
test()

• Thank you, it helped me a lot! On my machine, the computation has become 8 times faster. I have a quick question about the function argument. If you specify the type of arguments like you do, does it contribute to the speedup? Or, is it just for the readability of the code? Sep 22, 2021 at 20:28
• Type hints won't change the execution speed one way or the other. They improve legibility, tell your IDE better information e.g. to facilitate more meaningful autocomplete, and allow for better static analysis using tools like mypy. Sep 22, 2021 at 21:43
• That makes sense. I have not used the type hints so far but I should start using it. It would be helpful especially when I share my code with my colleagues. Thank you! Sep 23, 2021 at 14:29