miepython 3.2.0__py3-none-any.whl → 3.3.0__py3-none-any.whl

miepython/mie_jit.py

@@ -82,10 +82,79 @@ def _D_upwards(z, N, D):
         D[n] = 1 / (n / z - D[n - 1]) - n / z
 
 
+@njit((complex128, int64), cache=True)
+def _psi_downwards_nb(z, nstop):
+    """
+    Compute the Riccati-Bessel functions psi_0(z)..psi_nstop(z).
+
+    The three-term upwards recurrence
+    ``psi_{n+1} = (2n+1)/z psi_n - psi_{n-1}`` is contaminated by the growing
+    chi_n solution once n exceeds ``|z|``, which is the normal case for psi_n(mx)
+    whenever the relative index is below one.  Miller's algorithm is used
+    instead: seed the recurrence well above ``|z|`` where psi_n is the decaying
+    solution, run it downwards, then fix the arbitrary scale against whichever
+    of psi_0 or psi_1 is larger.  sin(z) and cos(z) cannot both be small, so one
+    of those two seeds is always well conditioned.
+
+    Args:
+        z: function argument (the refractive index times the size parameter)
+        nstop: highest order needed
+
+    Returns:
+        Array of psi_k(z) for k=0 to nstop
+    """
+    abs_z = abs(z)
+    n_start = int(max(float(nstop), abs_z)) + 25 + int(1.5 * np.sqrt(abs_z))
+
+    psi = np.zeros(n_start + 2, dtype=np.complex128)
+    psi[n_start] = 1e-50
+    for n in range(n_start, 0, -1):
+        psi[n - 1] = (2 * n + 1) / z * psi[n] - psi[n + 1]
+        if abs(psi[n - 1]) > 1e200:  # keep the growing tail in range
+            for k in range(n - 1, n_start + 2):
+                psi[k] /= 1e200
+
+    sin_z = np.sin(z)
+    psi_0 = sin_z
+    psi_1 = sin_z / z - np.cos(z)
+    if abs(psi_1) > abs(psi_0):
+        scale = psi_1 / psi[1]
+    else:
+        scale = psi_0 / psi[0]
+    for k in range(n_start + 2):
+        psi[k] *= scale
+
+    return psi[: nstop + 1]
+
+
+@njit((complex128, int64), cache=True)
+def _D_calc_down_nb(z, N):
+    """
+    Compute D_1(z)..D_N(z) using only the downwards recurrence.
+
+    ``_D_calc_nb`` picks between the upwards and downwards recurrences with
+    Wiscombe's criterion, which considers the refractive index but not the number
+    of terms.  The upwards recurrence loses accuracy once N exceeds ``|z|`` (about
+    4 digits for a small sphere), and the resulting noise differs between the
+    numba and pure-python backends.  The internal-field coefficients are not on
+    the hot path, so they always take the stable route.
+
+    Args:
+        z: function argument (already multiplied by the refractive index)
+        N: highest order needed
+
+    Returns:
+        Array of logarithmic derivatives D_k(z) for k=1 to N
+    """
+    D = np.zeros(N + 1, dtype=np.complex128)
+    _D_downwards(z, N, D)
+    return D[1:]
+
+
 @njit((complex128, float64, int64), cache=True)
 def _D_calc_nb(m, x, N):
     """
-    Compute the logarithmic derivative of ψ_n(z) using the best method.
+    Compute the logarithmic derivative of ψ_n(z) using the downwards recurrence.
 
     D_n(z) = d[log ψ_n(z)] = ψ_n'(z)/ψ_n(z)
 
@@ -94,6 +163,14 @@ def _D_calc_nb(m, x, N):
 
     The zero-based array, D[:], is shifted so that D[0] = D₁(z) = ψ₁'(z)/ψ₁(z)
 
+    Wiscombe's criterion used to pick between the upwards and downwards
+    recurrences from the refractive index alone.  The upwards recurrence is
+    contaminated once N passes ``|mx|``, which is the case for every small sphere,
+    and it left the Mie coefficients wrong by as much as 1% near m=1.  The
+    downwards recurrence is stable everywhere, and the criterion already chose it
+    for most large spheres, so always taking it costs little.  ``_D_upwards`` is
+    kept for comparison and testing.
+
     Args:
         m: the np.complex128 index of refraction of the sphere
         x: the size parameter of the sphere
@@ -102,20 +179,15 @@ def _D_calc_nb(m, x, N):
     Returns:
         Array of logarithmic derivatives D_k(z) for k=1 to N-1.
     """
-    n = m.real
-    kappa = np.abs(m.imag)
-    D = np.zeros(N + 1, dtype=np.complex128)
-    mx = np.complex128(m * x)  # ensure complex
-
-    if n < 1 or n > 10 or kappa > 10 or x * kappa >= 3.9 - 10.8 * n + 13.78 * n**2:
-        _D_downwards(mx, N, D)
-    else:
-        _D_upwards(mx, N, D)
-    return D[1:]
+    return _D_calc_down_nb(np.complex128(m * x), N)
 
 
-@njit((complex128, float64, int64), cache=True, fastmath=True)
-def _an_bn_nb(m, x, n_pole):
+# Lazily compiled, unlike the kernels around it, so that n_pole can carry the same
+# default as ``_an_bn_py``: numba rejects an omitted argument against an eager
+# signature, which used to make ``an_bn(m, x)`` work with the JIT off and raise
+# TypeError with it on.
+@njit(cache=True)
+def _an_bn_nb(m, x, n_pole=0):
     """
     Compute arrays of Mie coefficients a_n and b_n for a sphere.
 
@@ -123,8 +195,8 @@ def _an_bn_nb(m, x, n_pole):
     formula. The length of the arrays is chosen so that the error when the series
     is summed is around 1e-6.
 
-    If n_pole>0, then the array sizes will be n_pole+1. This is useful when
-    trying to isolate the behavior of a particular multipole.
+    If n_pole>0, then the arrays hold exactly n_pole terms, orders 1 to n_pole.
+    This is useful when trying to isolate the behavior of a particular multipole.
 
     To support resonance calculations, one can specify the number of terms
     to be calculated.  In general, using too few or too many terms increases the
@@ -137,7 +209,8 @@ def _an_bn_nb(m, x, n_pole):
         n_pole: the number of An and Bn terms (0 does autosizing)
 
     Returns:
-        a, b: arrays of Mie coefficents An and Bn
+        a, b: arrays of Mie coefficents An and Bn, one entry per order,
+        with no padding.  ``_cn_dn`` returns the same number of terms.
     """
     # make sure sign of imaginary part is negative
     m_re = m.real
@@ -146,55 +219,70 @@ def _an_bn_nb(m, x, n_pole):
         m_im = -m_im
     m = complex(m_re, m_im)
 
-    # make sure array size is large enough
+    # Wiscombe's truncation: the series is summed through order n_terms
     if n_pole == 0:
-        nstop = int(x + 4.05 * x**0.33333 + 2.0) + 1
+        n_terms = int(x + 4.05 * x**0.33333 + 2.0)
     else:
-        nstop = n_pole + 1
+        n_terms = n_pole
 
-    a = np.zeros(nstop, dtype=np.complex128)
-    b = np.zeros(nstop, dtype=np.complex128)
+    a = np.zeros(n_terms, dtype=np.complex128)
+    b = np.zeros(n_terms, dtype=np.complex128)
     if x <= 0:
         return a, b
 
-    psi_nm1 = np.sin(x)  # nm1 = n-1 = 0
-    psi_n = psi_nm1 / x - np.cos(x)
-    xi_nm1 = np.complex128(psi_nm1 + 1j * np.cos(x))
-    xi_n = np.complex128(psi_n + 1j * (np.cos(x) / x + np.sin(x)))
+    # group the arithmetic exactly as _an_bn_py does: multiplying by a precomputed
+    # 1/x here and dividing by x there differed by one ulp, which the recurrences
+    # amplified into 2.6e-7 between the two backends.
+    inv_x = 1.0 / x
+    sin_x = np.sin(x)
+    cos_x = np.cos(x)
+
+    # psi comes from Miller's downwards recurrence.  The upwards one is fine while
+    # n stays below x, which is all Wiscombe's truncation needs, but its relative
+    # accuracy is already down to about 1e-6 at n = n_terms and it collapses
+    # completely past that, so a caller asking for extra orders used to get noise.
+    psi = _psi_downwards_nb(np.complex128(x), n_terms).real
+
+    # chi is the growing solution, so its upwards recurrence is the stable one.
+    # Building xi as psi + i*chi keeps xi exactly consistent with the psi above,
+    # which the lossless identity Re(a_n) == |a_n|**2 depends on.
+    chi_nm1 = cos_x
+    chi_n = cos_x * inv_x + sin_x
 
     if m.real > 0.0:
-        D = _D_calc_nb(m, x, nstop + 1)
-
-        for n in range(1, nstop):
-            temp = D[n - 1] / m + n / x
-            a[n - 1] = (temp * psi_n - psi_nm1) / (temp * xi_n - xi_nm1)
-            temp = D[n - 1] * m + n / x
-            b[n - 1] = (temp * psi_n - psi_nm1) / (temp * xi_n - xi_nm1)
-            psi = (2 * n + 1) * psi_n / x - psi_nm1
-            xi = (2 * n + 1) * xi_n / x - xi_nm1
-            xi_nm1 = xi_n
-            xi_n = xi
-            psi_nm1 = psi_n
-            psi_n = psi
+        D = _D_calc_nb(m, x, n_terms + 2)
+
+        for n in range(1, n_terms + 1):
+            n_over_x = n * inv_x
+            xi_nm1 = psi[n - 1] + 1j * chi_nm1
+            xi_n = psi[n] + 1j * chi_n
+            temp = D[n - 1] / m + n_over_x
+            a[n - 1] = (temp * psi[n] - psi[n - 1]) / (temp * xi_n - xi_nm1)
+            temp = D[n - 1] * m + n_over_x
+            b[n - 1] = (temp * psi[n] - psi[n - 1]) / (temp * xi_n - xi_nm1)
+            chi = (2 * n + 1) * inv_x * chi_n - chi_nm1
+            chi_nm1 = chi_n
+            chi_n = chi
 
     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
-
-    if n_pole != 0:
-        a = a[:-1]
-        b = b[:-1]
+        for n in range(1, n_terms + 1):
+            n_over_x = n * inv_x
+            xi_nm1 = psi[n - 1] + 1j * chi_nm1
+            xi_n = psi[n] + 1j * chi_n
+            a[n - 1] = (n_over_x * psi[n] - psi[n - 1]) / (n_over_x * xi_n - xi_nm1)
+            b[n - 1] = psi[n] / xi_n
+            chi = (2 * n + 1) * inv_x * chi_n - chi_nm1
+            chi_nm1 = chi_n
+            chi_n = chi
 
     return np.conjugate(a), np.conjugate(b)
 
 
-@njit((complex128, float64, int64), fastmath=True)
+# No kernel in this module uses fastmath.  It implies LLVM's nnan and ninf, which
+# folded the np.isinf guard below to False and let an infinite index divide by
+# zero, and its reassociation also broke the exact Re(a_n) == |a_n|**2 identity
+# that a lossless sphere has to satisfy.  Do not add it back.
+@njit((complex128, float64, int64))
 def _cn_dn_nb(m, x, n_pole):
     """
     Calculate Mie coefficients c_n and d_n for the internal field of a sphere.
@@ -215,51 +303,48 @@ def _cn_dn_nb(m, x, n_pole):
     m = complex(m_re, m_im)
     mx = m * x
 
+    # same truncation as _an_bn_nb so the internal and external series match
     if n_pole == 0:
-        nstop = int(x + 4.05 * x**0.33333 + 2.0) + 1
+        n_terms = int(x + 4.05 * x**0.33333 + 2.0)
     else:
-        nstop = n_pole + 1
+        n_terms = n_pole
 
-    c = np.zeros(nstop, dtype=np.complex128)
-    d = np.zeros(nstop, dtype=np.complex128)
+    c = np.zeros(n_terms, dtype=np.complex128)
+    d = np.zeros(n_terms, dtype=np.complex128)
     if x <= 0:
         return c, d
 
     # no need to calculate anything when sphere is perfectly conducting
-    if m.real > 0.0 and not np.isinf(m.real) or not np.isinf(m.imag):
-        psi_nm1 = np.sin(x)  # nm1 = n-1 = 0
-        psi_n = psi_nm1 / x - np.cos(x)
-
-        psi_nm1_mx = np.sin(mx)  # nm1 = n-1 = 0
-        psi_n_mx = psi_nm1_mx / mx - np.cos(mx)
-
-        xi_nm1 = np.complex128(psi_nm1 + 1j * np.cos(x))
-        xi_n = np.complex128(psi_n + 1j * (np.cos(x) / x + np.sin(x)))
-
-        Dmx = _D_calc_nb(np.complex128(m), x, nstop + 1)
-        Dx = _D_calc_nb(np.complex128(1), x, nstop + 1)
-
-        for n in range(1, nstop + 1):
-            common = (psi_n / psi_n_mx) * ((Dx[n - 1] + n / x) * xi_n - xi_nm1)
-
-            c[n - 1] = m * common / ((m * Dmx[n - 1] + n / x) * xi_n - xi_nm1)
-            d[n - 1] = common / ((Dmx[n - 1] / m + n / x) * xi_n - xi_nm1)
-
-            psi = (2 * n + 1) * psi_n / x - psi_nm1
-            psi_nm1 = psi_n
-            psi_n = psi
-
-            psi_mx = (2 * n + 1) * psi_n_mx / mx - psi_nm1_mx
-            psi_nm1_mx = psi_n_mx
-            psi_n_mx = psi_mx
-
-            xi = (2 * n + 1) * xi_n / x - xi_nm1
+    # (m.real <= 0 or an infinite index): there is no internal field, so c and d
+    # stay zero.  The `and` chain must not be mixed with `or` here, otherwise the
+    # test is true for every finite index and m=0 divides by zero below.
+    if m.real > 0.0 and not np.isinf(m.real) and not np.isinf(m.imag):
+        # same arithmetic grouping as _cn_dn_py, so the two backends agree to the
+        # last bits rather than only to about 1e-13
+        inv_x = 1.0 / x
+        sin_x = np.sin(x)
+        cos_x = np.cos(x)
+
+        # xi is the growing solution, so seeding it and running upwards is stable
+        xi_nm1 = np.complex128(sin_x + 1j * cos_x)
+        xi_n = np.complex128((sin_x * inv_x - cos_x) + 1j * (cos_x * inv_x + sin_x))
+
+        psi_x = _psi_downwards_nb(np.complex128(x), n_terms + 1)
+        psi_mx = _psi_downwards_nb(np.complex128(mx), n_terms + 1)
+        Dmx = _D_calc_down_nb(np.complex128(mx), n_terms + 2)
+        Dx = _D_calc_down_nb(np.complex128(x), n_terms + 2)
+
+        for n in range(1, n_terms + 1):
+            n_over_x = n * inv_x
+            common = (psi_x[n] / psi_mx[n]) * ((Dx[n - 1] + n_over_x) * xi_n - xi_nm1)
+
+            c[n - 1] = m * common / ((m * Dmx[n - 1] + n_over_x) * xi_n - xi_nm1)
+            d[n - 1] = common / ((Dmx[n - 1] / m + n_over_x) * xi_n - xi_nm1)
+
+            xi = (2 * n + 1) * inv_x * xi_n - xi_nm1
             xi_nm1 = xi_n
             xi_n = xi
 
-    if n_pole != 0:
-        c = c[:-1]
-        d = d[:-1]
     return np.conjugate(c), np.conjugate(d)
 
 
@@ -297,6 +382,9 @@ def _pi_tau_nb(mu, pi, tau):
         pi[n] = ((2 * n + 1) * mu * temp - (n + 1) * pi_nm2) / n
         pi_nm2 = temp
 
+    # each pass sets tau one order behind pi, so the highest order is still due
+    tau[n_terms - 1] = n_terms * mu * pi[n_terms - 1] - (n_terms + 1) * pi_nm2
+
 
 @njit((complex128, float64, float64[:], int64), cache=True)
 def _S1_S2_nb(m, x, mu, n_pole):
@@ -319,6 +407,9 @@ def _S1_S2_nb(m, x, mu, n_pole):
     """
     a, b = _an_bn_nb(m, x, 0)
     N = len(a)
+    if n_pole < 0 or n_pole > N:
+        raise ValueError("n_pole must be 0 (all terms) or a multipole order in 1.." + str(N))
+
     pi = np.zeros(N)
     tau = np.zeros(N)
     n = np.arange(1, N + 1, dtype=np.float64)
@@ -328,6 +419,9 @@ def _S1_S2_nb(m, x, mu, n_pole):
     S1 = np.zeros(nangles, dtype=np.complex128)
     S2 = np.zeros(nangles, dtype=np.complex128)
 
+    # zero-based arrays: multipole order n lives at index n-1
+    j = n_pole - 1
+
     for k in range(nangles):
         _pi_tau_nb(mu[k], pi, tau)
         if n_pole == 0:
@@ -340,13 +434,13 @@ def _S1_S2_nb(m, x, mu, n_pole):
             S1[k] = s1
             S2[k] = s2
         else:
-            S1[k] = scale[n_pole] * (pi[n_pole] * a[n_pole] + tau[n_pole] * b[n_pole])
-            S2[k] = scale[n_pole] * (tau[n_pole] * a[n_pole] + pi[n_pole] * b[n_pole])
+            S1[k] = scale[j] * (pi[j] * a[j] + tau[j] * b[j])
+            S2[k] = scale[j] * (tau[j] * a[j] + pi[j] * b[j])
 
     return np.conjugate(S1), np.conjugate(S2)
 
 
-@njit((complex128, float64), cache=True, fastmath=True)
+@njit((complex128, float64), cache=True)
 def _small_conducting_sphere_nb(_m, x):
     """
     Calculate the efficiencies for a small conducting spheres.
@@ -384,7 +478,7 @@ def _small_conducting_sphere_nb(_m, x):
     return qext, qsca, qback, g
 
 
-@njit((complex128, float64), cache=True, fastmath=True)
+@njit((complex128, float64), cache=True)
 def _small_sphere_nb(m, x):
     """
     Calculate the efficiencies for a small sphere.
@@ -432,7 +526,7 @@ def _small_sphere_nb(m, x):
     return qext, qsca, qback, g
 
 
-@njit((complex128, float64, int64, boolean), cache=True, fastmath=True)
+@njit((complex128, float64, int64, boolean), cache=True)
 def _single_sphere_nb(m, x, n_pole, e_field):
     """
     Calculate the efficiencies for a single sphere (both m and x are scalars).
@@ -441,7 +535,9 @@ def _single_sphere_nb(m, x, n_pole, e_field):
         m: the complex index of refraction of the sphere
         x: the size parameter of the sphere
         n_pole: a non-zero value returns the contribution by the n_pole multipole
-        e_field: Electric (True) or Magnetic Field otherwise
+        e_field: selects which multipole of order n_pole contributes, the
+            electric one a_n (True) or the magnetic one b_n (False).
+            Ignored when n_pole == 0.
 
     Returns:
         qext: the total extinction efficiency
@@ -449,7 +545,10 @@ def _single_sphere_nb(m, x, n_pole, e_field):
         qback: the backscatter efficiency
         g: the average cosine of the scattering phase function
     """
-    _ = e_field  # unused
+    # a sphere of zero size scatters nothing; every efficiency tends to zero,
+    # including qback, whose small-sphere form is 0/0 at x=0
+    if x <= 0:
+        return 0.0, 0.0, 0.0, 0.0
 
     # case when sphere matches its environment
     if abs(m.real - 1) <= 1e-8 and abs(m.imag) < 1e-8:
@@ -515,12 +614,18 @@ def _single_sphere_nb(m, x, n_pole, e_field):
         g = 4.0 * g_acc / qsca / x2
 
     else:
+        # isolate one multipole: the electric term a_n or the magnetic term b_n
         cn = 2.0 * n_pole + 1
-        qext = 2.0 * cn * (a[-1].real + b[-1].real) / x2
-        qback = np.abs((-1) ** n_pole * cn * (a[-1] - b[-1])) ** 2 / x2
+        coeff = a[-1] if e_field else b[-1]
+        qext = 2.0 * cn * coeff.real / x2
+        # the (-1)**n_pole and the sign of b_n drop out of the modulus
+        qback = np.abs(cn * coeff) ** 2 / x2
         qsca = qext
         if m.imag < 0:
-            qsca = 2.0 * cn * (np.abs(a[-1]) ** 2 + np.abs(b[-1]) ** 2) / x2
-        g = 0
+            qsca = 2.0 * cn * np.abs(coeff) ** 2 / x2
+
+        # pi_n**2 and tau_n**2 are both even in mu, so an isolated multipole of
+        # one parity scatters symmetrically about 90 degrees and <cos(theta)>=0
+        g = 0.0
 
     return qext, qsca, qback, g