#!/usr/bin/env python3
"""
dust_grain.py -- how many records can a sunlit dust grain carry? (Chapter 12)

Riedel & Zurek, PRL 105, 020404 (2010) state that after 1 us of sunlight "a grain of dust 1 um across will
have its location imprinted about 100 x 10^6 times in the scattered photons".  They use "the classical cross
section of a dielectric sphere in the dipole approximation (lambda >> a)".

Why scattered photons bound the records: in that model each photon's two conditional states are pure,
|psi_x> = sqrt(1-P)|k> + sqrt(P)|s_x>, so their overlap differs from 1 by at most 2P, where P is the photon's
scattering probability.  A fragment's overlap is the product over its photons, and it must be small for the
fragment to read the position, so each record uses up O(1) expected scattering events.  The redundancy R is
therefore at most the number of scattered photons, up to a factor of order one.  This script counts
them with exact Mie scattering (Bohren-Huffman series) for a sphere of radius a = 0.5 um, integrated over a
5772 K blackbody spectrum normalised to the solar constant 1361 W/m^2, and compares with the dipole formula.
The refractive index of the grain is not given, so three values are shown; both readings of "1 um"
(diameter or radius) are shown.
numpy only.  No files are written.
"""
import os
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, "1")
import numpy as np

TRAP = getattr(np, "trapezoid", None) or np.trapz
H, C, KB = 6.62607015e-34, 2.99792458e8, 1.380649e-23
T_SUN, S0 = 5772.0, 1361.0            # K, W/m^2 at 1 AU
A_GRAIN, T_EXP = 0.5e-6, 1e-6         # radius (1 um across), exposure time
A_ALT = 1.0e-6                        # in case "1 um" is read as the radius
RZ_CLAIM = 1e8


def mie_qsca(x, m):
    """Scattering efficiency of a homogeneous sphere, size parameter x = 2 pi a / lambda, index m (BHMIE)."""
    nstop = int(x + 4.05 * x ** (1 / 3) + 2)
    y = m * x
    nmx = int(max(nstop, abs(y)) + 16)
    D = np.zeros(nmx + 1, complex)                      # logarithmic derivative, downward recurrence
    for n in range(nmx, 0, -1):
        D[n - 1] = n / y - 1 / (D[n] + n / y)
    psi0, psi1 = np.cos(x), np.sin(x)
    chi0, chi1 = -np.sin(x), np.cos(x)
    xi1 = psi1 - 1j * chi1
    q = 0.0
    for n in range(1, nstop + 1):
        psi = (2 * n - 1) * psi1 / x - psi0
        chi = (2 * n - 1) * chi1 / x - chi0
        xi = psi - 1j * chi
        an = ((D[n] / m + n / x) * psi - psi1) / ((D[n] / m + n / x) * xi - xi1)
        bn = ((D[n] * m + n / x) * psi - psi1) / ((D[n] * m + n / x) * xi - xi1)
        q += (2 * n + 1) * (abs(an) ** 2 + abs(bn) ** 2)
        psi0, psi1, chi0, chi1, xi1 = psi1, psi, chi1, chi, xi
    return 2 * q / x ** 2


def rayleigh_qsca(x, m):
    return (8 / 3) * x ** 4 * abs((m * m - 1) / (m * m + 2)) ** 2


def solar_photon_flux(lams):
    """Photons m^-2 s^-1 per metre of wavelength, 5772 K blackbody scaled to S0."""
    B = 2 * H * C ** 2 / lams ** 5 / np.expm1(H * C / (lams * KB * T_SUN))    # spectral radiance shape
    E = TRAP(B, lams)
    return (S0 / E) * B / (H * C / lams)


def main():
    lams = np.geomspace(0.15e-6, 20e-6, 6000)
    phi = solar_photon_flux(lams)
    total_flux = TRAP(phi, lams)
    geo = np.pi * A_GRAIN ** 2
    x = 2 * np.pi * A_GRAIN / lams
    print("dust_grain.py -- photons scattered by a sunlit dust grain (radius 0.5 um, 1 us of sunlight at 1 AU)")
    print(f"solar photon flux {total_flux:.3e} m^-2 s^-1;  photons crossing the geometric cross-section in 1 us: {total_flux * geo * T_EXP:.3e}")
    vis = (lams > 0.4e-6) & (lams < 0.7e-6)
    print(f"size parameter ka over the visible band: {x[vis].min():.1f} to {x[vis].max():.1f};  ka = 3.4 at lambda = {2 * np.pi * A_GRAIN / 3.4 * 1e6:.2f} um")
    print(f"{'index m':>12} {'N_scat (Mie)':>13} {'N_scat (dipole)':>16} {'dipole / Mie':>13} {'dipole/Mie at ka>=3.4':>22} {'RZ 1e8 / N_scat':>16}")
    ok = True
    for m in (1.33 + 0j, 1.5 + 0j, 1.7 + 0.01j):
        qm = np.array([mie_qsca(xi, m) for xi in x])
        qr = rayleigh_qsca(x, m)
        n_mie = TRAP(phi * qm * geo, lams) * T_EXP
        n_ray = TRAP(phi * qr * geo, lams) * T_EXP
        big = x >= 3.4
        ratio_big = qr[big] / qm[big]
        print(f"{str(m):>12} {n_mie:>13.3e} {n_ray:>16.3e} {n_ray / n_mie:>13.1f} {ratio_big.min():>10.0f} to {ratio_big.max():<9.0f} {RZ_CLAIM / n_mie:>16.0f}")
        ok &= n_mie < 1e5
    # the other reading of "1 um": radius a = 1 um
    x2 = 2 * np.pi * A_ALT / lams
    geo2 = np.pi * A_ALT ** 2
    print("if '1 um' is the radius (a = 1 um):")
    for m in (1.5 + 0j, 1.7 + 0.01j):
        qm = np.array([mie_qsca(xi, m) for xi in x2])
        n_mie = TRAP(phi * qm * geo2, lams) * T_EXP
        n_ray = TRAP(phi * rayleigh_qsca(x2, m) * geo2, lams) * T_EXP
        print(f"{str(m):>12} {n_mie:>13.3e} {n_ray:>16.3e} {n_ray / n_mie:>13.1f} {'':>22} {RZ_CLAIM / n_mie:>16.0f}")
        ok &= n_mie < 1e5
    # sanity: Mie efficiency approaches 2 for large spheres (extinction paradox, non-absorbing) and Rayleigh for small ones
    print(f"checks: Q_sca(x=0.05, m=1.5) Mie/Rayleigh = {mie_qsca(0.05, 1.5) / rayleigh_qsca(0.05, 1.5):.6f} (-> 1);  Q_sca(x=400, m=1.5) = {mie_qsca(400.0, 1.5):.3f} (-> 2)")
    print("about 1e4 (1 um diameter) to 5e4 (1 um radius) photons scatter, so the grain carries at most that many records, not ~1e8" if ok else "CHECK FAILED")


if __name__ == "__main__":
    main()