Monte Carlo using Numba

Question

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, Commented Sep 19, 2021 at 17:09
At least explain what the code does and how it works, as per site rules. — Cris Luengo, Commented Sep 19, 2021 at 17:27
Sorry about that. I added more explanations to my post. Let me know if there is anything missing. — Ogiad, Commented 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. — Toby Speight, Commented Sep 21, 2021 at 7:49

Reinderien · Accepted Answer · 2021-09-24 00:37:21Z

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? — Ogiad, Commented 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. — Reinderien, Commented 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! — Ogiad, Commented Sep 23, 2021 at 14:29

Stack Exchange Network

Monte Carlo using Numba

1 Answer 1

Edit

Your Answer

Not the answer you're looking for? Browse other questions tagged
python
numba
or ask your own question.

Hot Network Questions

Monte Carlo using Numba

1 Answer 1

Edit

Your Answer

Sign up or log in

Post as a guest

Not the answer you're looking for? Browse other questions tagged pythonnumba or ask your own question.

Related

Hot Network Questions

Not the answer you're looking for? Browse other questions tagged
python
numba
or ask your own question.