Normalize scattering CDFs before inverse angle sampling

CHANGELOG.rst

@@ -3,6 +3,28 @@ Changelog
 
 unreleased
 -------------------
+*   fix ``monte_carlo.mu_with_uniform_cdf`` returning wrong angles for fine tables.
+    ``cdf`` began at I(-1)/num rather than zero, so a target below that value left
+    the inversion loop at index zero and interpolated against the last entry.
+    ``cdf`` now integrates the phase function with the trapezoidal rule, starting at
+    exactly 0 and ending at exactly 1 (it used to end slightly above 1), so tables
+    differ slightly from 3.3.0.  ``cdf`` and ``mu_with_uniform_cdf`` also raise
+    ``ValueError`` when num is less than 2 (thanks @antonsoo)
+*   replace the inversion loop in ``monte_carlo.mu_with_uniform_cdf`` with
+    ``np.interp``
+*   warn when ``monte_carlo.cdf`` uses too few angles to resolve the phase function.
+    The raw integral should be one, so a ``RuntimeWarning`` is issued when it is off
+    by more than 1%; normalizing used to hide the error silently.  A large sphere
+    needs roughly x**2 angles, because its forward peak is only about 8/x**2 wide in
+    cos(theta)
+*   size the table inside ``monte_carlo.mu_with_uniform_cdf`` from x, as
+    max(2000, 5*x**2) angles instead of a fixed 2000.  At x = 100 the sampled mean
+    cosine was 0.930 against g = 0.868; it is now 0.869.  Large spheres are slower
+    to tabulate: seconds at x = 100, about a minute at x = 300
+*   raise ``ValueError`` from ``monte_carlo.cdf`` and ``mu_with_uniform_cdf`` when the
+    phase function is not finite
+*   rewrite the ``monte_carlo`` docstrings, with tested examples of drawing random
+    scattering angles
 *   name the encoding on every text read and write.  ``open(path)`` and
     ``Path.read_text()`` decode with the locale's preferred encoding, which is UTF-8
     on Linux and macOS but cp1252 on Windows, so ``tests/test_makefile.py`` passed

miepython/monte_carlo.py

@@ -1,37 +1,64 @@
 """
-Sample scattering angles from the Mie phase function.
-
-Monte Carlo photon transport needs a random exit direction distributed according to
-the phase function.  This module builds the inverse-transform table that provides
-one: ``mu_with_uniform_cdf`` returns angles spaced so that equal steps along the
-cumulative distribution are equally likely, and ``generate_mie_costheta`` draws from
-that table.  ``docs/06_random_deviates.ipynb`` works through the whole procedure.
-
-Importing
----------
-
-This module is *not* imported with the package, because it imports ``miepython``
-itself and doing so from ``miepython/__init__.py`` would be circular.  Ask for it by
-name::
-
-    import numpy as np
-    import miepython.monte_carlo as mc
-
-    m, x = 1.33, 2.0                              # a water droplet
-    mu, cdf = mc.mu_with_uniform_cdf(m, x, 500)   # 500-entry lookup table
-
-    rng = np.random.default_rng(12345)
-    cosines = [mc.generate_mie_costheta(mu, rng=rng) for _ in range(10000)]
-
-Build the table once and reuse it: it costs a full Mie calculation, while each draw
-off it costs one random number.  The number of table entries sets how finely the
-phase function is resolved, so strongly forward-scattering spheres need more of
-them.
+Sample random scattering angles from the Mie phase function.
+
+A Monte Carlo simulation of light transport needs, at every scattering event, a
+random direction in which the photon leaves the sphere. For unpolarized light
+only the polar angle theta depends on the sphere, and it must be drawn so that
+its cosine mu follows the phase function. This module does that by inverse
+transform sampling: build the cumulative distribution of mu once, then map each
+uniform random number u to the cosine at which the distribution reaches u.
+
+``cdf``
+    The cumulative distribution at evenly spaced cosines.
+``mu_with_uniform_cdf``
+    The cosines at which the cumulative distribution reaches evenly spaced
+    probabilities, a lookup table that makes each draw cheap.
+``generate_mie_costheta``
+    One random cosine drawn from that table.
+
+``docs/06_random_deviates.ipynb`` works through the whole procedure.
+
+This module is *not* imported with the package, because it imports
+``miepython`` itself and doing so from ``miepython/__init__.py`` would be
+circular. Import it by name, as the examples below do.
+
+Examples:
+    Invert the distribution with ``np.interp`` to draw many cosines at once.
+    Their average is the asymmetry parameter g:
+
+    >>> import numpy as np
+    >>> import miepython as mie
+    >>> import miepython.monte_carlo as mc
+    >>> m, x = 1.33, 2.0                          # a water droplet
+    >>> mu, cdf = mc.cdf(m, x, 1000)
+    >>> rng = np.random.default_rng(42)
+    >>> cosines = np.interp(rng.random(100_000), cdf, mu)
+    >>> print(f"{np.mean(cosines):.2f}")
+    0.67
+    >>> qext, qsca, qback, g = mie.efficiencies_mx(m, x)
+    >>> print(f"{g:.2f}")
+    0.67
+
+    Or build a lookup table once and draw one cosine at a time, as a photon
+    transport loop does:
+
+    >>> mu_table, _ = mc.mu_with_uniform_cdf(m, x, 1000)
+    >>> cosines = [mc.generate_mie_costheta(mu_table, rng=rng) for _ in range(100_000)]
+    >>> print(f"{np.mean(cosines):.2f}")
+    0.67
+
+Building either table costs a full Mie calculation, while each draw costs only a
+random number or two, so build it once and reuse it. Evenly spaced cosines
+resolve the narrow forward peak of a large sphere poorly; ``cdf`` explains how
+many are needed and warns when there are too few, while ``mu_with_uniform_cdf``
+chooses enough by itself.
 
 Passing ``rng`` keeps the draws reproducible without disturbing the process-wide
 ``numpy.random``; omitting it falls back to that global generator.
 """
 
+import warnings
+
 import numpy as np
 import miepython as mie
 
@@ -41,87 +68,102 @@
     "generate_mie_costheta",
 )
 
+# Largest acceptable relative error in the integral of the phase function
+CDF_TOLERANCE = 0.01
+
 
 def cdf(m, x, num):
     """
-    Create a CDF for unpolarized scattering uniformly spaced in cos(theta).
+    Cumulative distribution of the unpolarized phase function in cos(theta).
 
-    The CDF covers scattered (exit) angles ranging from 180 to 0 degrees.
-    (The cosines are uniformly distributed over -1 to 1.) Because the angles
-    are uniformly distributed in cos(theta), the scattering function is not
-    sampled uniformly and therefore huge array sizes are needed to adequately
-    sample highly anisotropic phase functions.
+    The cosines mu are evenly spaced from -1 (backscattering) to 1 (forward
+    scattering), and cdf[i] is the probability that a scattered photon leaves
+    with a cosine no greater than mu[i]. It is found by integrating the phase
+    function with the trapezoidal rule, so cdf[0] is exactly 0 and the values
+    never decrease.
 
-    Since this is a cumulative distribution function, the maximum value
-    should be 1.
+    A large sphere scatters much of its light into a diffraction peak about
+    4/x radians wide around the forward direction. Near theta = 0, mu = cos(theta)
+    is close to 1 - theta**2/2, so that peak spans only about 8/x**2 in mu,
+    while the cosines are 2/(num - 1) apart. To put several cosines inside the
+    peak, num must therefore grow roughly as x**2; x = 30 needs about 5000.
+
+    The phase function integrates to exactly one, so the raw trapezoidal sum
+    shows how well num resolves it. If that sum is off by more than
+    ``CDF_TOLERANCE``, a ``RuntimeWarning`` is issued; increase num until it
+    stops. The result is then divided by the sum so that cdf[-1] is exactly 1.
 
     Args:
-        m: the complex index of refraction of the sphere
-        x: the size parameter of the sphere
-        num: length of desired CDF array
+        m: complex index of refraction of the sphere
+        x: size parameter of the sphere
+        num: number of cosines, at least two
 
     Returns:
-        mu: array of cosines of angles
-        cdf: array of cumulative distribution function values
+        mu: evenly spaced cosines of the scattering angle, from -1 to 1
+        cdf: cumulative probability at each cosine, from 0 to 1
+
+    Raises:
+        ValueError: if num is less than two or the phase function is not finite
     """
+    if num < 2:
+        raise ValueError("num must be at least 2")
+    return _trapezoid_cdf(m, x, num, stacklevel=3)
+
+
+def _trapezoid_cdf(m, x, num, stacklevel):
+    """Integrate the phase function at num cosines, warning at the caller's caller."""
     mu = np.linspace(-1, 1, num)
-    intensity_per_mu = mie.i_unpolarized(m, x, mu, norm="4pi") / num
-    cdf_values = np.cumsum(intensity_per_mu)
+    intensity_per_mu = mie.i_unpolarized(m, x, mu, norm="4pi") / 2
+    cdf_values = np.zeros(num)
+    cdf_values[1:] = np.cumsum((intensity_per_mu[1:] + intensity_per_mu[:-1]) * np.diff(mu) / 2)
+    total = cdf_values[-1]
+    if not np.isfinite(total):
+        raise ValueError(f"phase function is not finite for m={m}, x={x}")
+    if abs(total - 1) > CDF_TOLERANCE:
+        warnings.warn(
+            f"phase function under-resolved for x={x}: integral over {num} angles is {total:.4g}, "
+            "not 1, so sampled angles will be inaccurate; use more angles",
+            RuntimeWarning,
+            stacklevel=stacklevel,
+        )
+    cdf_values /= total
     return mu, cdf_values
 
 
 def mu_with_uniform_cdf(m, x, num):
     """
-    Create a CDF for unpolarized scattering for uniform CDF.
-
-    The CDF covers scattered (exit) angles ranging from 180 to 0 degrees.
-    (The cosines are uniformly distributed over -1 to 1.) These angles mu
-    correspond to uniform spacing of the cumulative distribution function
-    for unpolarized Mie scattering where cdf[i] = i/(num-1).
+    Inverse-transform table for sampling the unpolarized phase function.
 
-    This is a brute force implementation that solves the problem by
-    calculating the CDF at many points and then scanning to find the
-    specific angles that correspond to uniform interval of the CDF.
+    The table holds num cosines mu running from -1 (backscattering) to 1
+    (forward scattering), spaced so that a scattered photon is equally likely
+    to fall between any two neighbours: cdf[i] is exactly i/(num - 1). The
+    cosines are therefore bunched where scattering is strong and sparse where
+    it is weak. ``generate_mie_costheta`` draws random cosines from this table.
 
-    Since this is a cumulative distribution function, the maximum value
-    should be 1.
+    The table is found by calculating ``cdf`` at max(2000, 5*x**2) evenly
+    spaced cosines, enough to resolve the forward peak whatever num is (see
+    ``cdf``), and interpolating to find where it reaches each step. Because
+    that count grows as x**2, building the table takes seconds at x = 100 and
+    over a minute at x = 300.
 
     Args:
-        m: the complex index of refraction of the sphere
-        x: the size parameter of the sphere
-        num: length of desired CDF array
+        m: complex index of refraction of the sphere
+        x: size parameter of the sphere
+        num: number of entries in the table, at least two
 
     Returns:
-        mu: array of cosines of angles (irregularly spaced)
-        cdf: array of cumulative distribution function values
-    """
-    big_num = 2000  # large to work with x up to 10
-    big_mu, big_cdf = cdf(m, x, big_num)
-    mu = np.empty(num)
-    cdf_values = np.empty(num)
-
-    mu[0] = -1  # cos[180 degrees] is -1
-    cdf_values[0] = 0  # initial cdf is zero
-
-    big_k = 0  # index into big_cdf
-    for k in range(1, num - 1):
+        mu: unevenly spaced cosines of the scattering angle, from -1 to 1
+        cdf: evenly spaced cumulative probabilities, from 0 to 1
 
-        target = k / (num - 1)
-        while big_cdf[big_k] < target:
-            big_k += 1
-
-        delta = big_cdf[big_k] - target
-        delta_cdf = big_cdf[big_k] - big_cdf[big_k - 1]
-        delta_mu = big_mu[big_k] - big_mu[big_k - 1]
-
-        mu[k] = big_mu[big_k] - delta / delta_cdf * delta_mu  # interpolate
-        cdf_values[k] = target
-
-    #       print(' mu[',k,']=% .5f'%mu[k],' cdf[',k,']=% .5f'%cdf[k],
-    #       'cdf=',big_cdf[big_k], fraction)
-
-    mu[num - 1] = 1  # cos[0 degrees] is 1
-    cdf_values[num - 1] = 1  # last cdf is one
+    Raises:
+        ValueError: if num is less than two or the phase function is not finite
+    """
+    if num < 2:
+        raise ValueError("num must be at least 2")
+    big_num = max(2000, int(5 * x**2))  # resolve the forward peak, see cdf
+    big_mu, big_cdf = _trapezoid_cdf(m, x, big_num, stacklevel=3)
+    cdf_values = np.arange(num) / (num - 1)
+    mu = np.interp(cdf_values, big_cdf, big_mu)
 
     return [mu, cdf_values]