Add ndoptimize.models.gaussian_binned, a Gaussian averaged over each bin.

ndoptimize/_tests/test_models.py

@@ -1,8 +1,10 @@
-from typing import Callable
+from typing import Any, Callable
 import copy
+import math
 import pickle
 import numpy as np
 import pytest
+import scipy.integrate
 import ndoptimize
 
 
@@ -139,3 +141,113 @@ def test_model_reference() -> None:
     assert copy.copy(model) is model
     assert copy.deepcopy(model) is model
     assert pickle.loads(pickle.dumps(model)) is model
+
+
+_edges = np.array(
+    [
+        [-160.0, -150.0],
+        [-50.0, -40.0],
+        [-12.0, -2.0],
+        [-3.0, 8.0],
+        [2.0, 3.0],
+        [20.0, 35.0],
+        [150.0, 160.0],
+    ]
+)
+"""
+Bins in both tails, far enough out that the difference of the error functions
+cancels, on either side of the center, and across it.
+"""
+
+
+def _gaussian_scalar(v: float, a: float, c: float, w: float) -> float:
+    """A Gaussian at a single point, for numerical integration."""
+    return a * math.exp(-0.5 * ((v - c) / w) ** 2)
+
+
+def test_gaussian_binned() -> None:
+    """
+    The binned Gaussian is the average of the Gaussian over each bin, to full
+    precision even far in the tails, and its Jacobian is its derivative.
+    """
+    model = ndoptimize.models.gaussian_binned
+    parameters = np.array([2.0, 3.0, 10.0])
+    f = np.empty(len(_edges))
+    jacobian = np.empty((len(_edges), 3))
+
+    model(_edges, parameters, f, jacobian)
+
+    expected = np.array(
+        [
+            scipy.integrate.quad(
+                _gaussian_scalar,
+                lower,
+                upper,
+                args=tuple(parameters),
+                epsabs=0,
+                epsrel=1e-13,
+            )[0]
+            / (upper - lower)
+            for lower, upper in _edges
+        ]
+    )
+    assert np.all(expected[[0, -1]] < 1e-40)
+    assert np.allclose(f, expected, rtol=1e-11, atol=0)
+
+    step = 1e-6
+    for k in range(parameters.size):
+        dp = np.zeros(parameters.size)
+        dp[k] = step
+        f_plus, f_minus = np.empty(len(_edges)), np.empty(len(_edges))
+        model(_edges, parameters + dp, f_plus, np.empty((len(_edges), 3)))
+        model(_edges, parameters - dp, f_minus, np.empty((len(_edges), 3)))
+        derivative = (f_plus - f_minus) / (2 * step)
+        assert np.allclose(jacobian[:, k], derivative, rtol=1e-5, atol=1e-10)
+
+
+def test_gaussian_binned_narrow() -> None:
+    """A bin of zero width gives the Gaussian at that point, and narrow bins
+    approach it."""
+    parameters = np.array([2.0, 3.0, 10.0])
+    v = np.array([-7.0, 3.0, 15.0])
+
+    f = np.empty(v.size)
+    jacobian = np.empty((v.size, 3))
+    ndoptimize.models.gaussian(v[:, np.newaxis], parameters, f, jacobian)
+
+    for size in (0, 1e-6):
+        x = np.stack([v - size / 2, v + size / 2], axis=-1)
+        f_binned = np.empty(v.size)
+        jacobian_binned = np.empty((v.size, 3))
+        ndoptimize.models.gaussian_binned(x, parameters, f_binned, jacobian_binned)
+        assert np.allclose(f_binned, f, rtol=1e-9)
+        assert np.allclose(jacobian_binned, jacobian, rtol=1e-6, atol=1e-12)
+
+
+def test_gaussian_binned_fit() -> None:
+    """
+    Fitting the binned Gaussian to the averages of a line over bins as wide as
+    the line recovers the line, while fitting the Gaussian at the centers of
+    the bins overestimates its width.
+    """
+    truth = np.array([5.0, 2.0, 10.0])
+    edges = np.linspace(-60, 60, 13)
+    x = np.stack([edges[:-1], edges[1:]], axis=-1)
+    data = np.empty(len(x))
+    ndoptimize.models.gaussian_binned(x, truth, data, np.empty((len(x), 3)))
+
+    kwargs: dict[str, Any] = dict(data=data, variance=1e-4, guess=[4, 0, 8])
+    result = ndoptimize.least_squares(
+        model=ndoptimize.models.gaussian_binned,
+        x=x,
+        **kwargs,
+    )
+    assert result.converged
+    assert np.allclose(result.parameters, truth, rtol=1e-6)
+
+    result_sampled = ndoptimize.least_squares(
+        model=ndoptimize.models.gaussian,
+        x=x.mean(axis=-1, keepdims=True),
+        **kwargs,
+    )
+    assert result_sampled.parameters[2] > 1.02 * truth[2]

ndoptimize/models.py

@@ -8,6 +8,7 @@
 
 __all__ = [
     "gaussian",
+    "gaussian_binned",
     "gaussian_linear",
     "two_gaussians_linear",
 ]
@@ -54,6 +55,120 @@ def gaussian(
         jacobian[i, 2] = amplitude * e * z * z / width
 
 
+@_model.model(
+    parameters=("amplitude", "center", "width"),
+)
+def gaussian_binned(
+    x: np.ndarray,
+    parameters: np.ndarray,
+    f: np.ndarray,
+    jacobian: np.ndarray,
+) -> None:
+    r"""
+    A Gaussian averaged over each bin,
+
+    .. math::
+
+        f(x_0, x_1) = \frac{1}{x_1 - x_0} \int_{x_0}^{x_1}
+            A \exp \left( -\frac{(x - \mu)^2}{2 \sigma^2} \right) dx
+            = \sqrt{\frac{\pi}{2}} \frac{A \sigma}{x_1 - x_0} \left[
+                \text{erf} \left( \frac{x_1 - \mu}{\sqrt{2} \sigma} \right)
+                - \text{erf} \left( \frac{x_0 - \mu}{\sqrt{2} \sigma} \right)
+            \right],
+
+    with the same parameters as :func:`gaussian`, the amplitude, center and
+    width, where :math:`x_0` and :math:`x_1`, the first two components of
+    :math:`x`, are the edges of each bin.
+    The integral of the Gaussian is :math:`\sqrt{2 \pi} A \sigma`.
+
+    Averaging matters when the bins are not much narrower than the line.
+    Averaging over bins of width :math:`\Delta` broadens a Gaussian by about
+    :math:`\Delta^2 / 12` in variance, which a fit of :func:`gaussian` at the
+    centers of the bins would take as part of the width of the line.
+    A bin of zero width gives the Gaussian at that point.
+
+    Parameters
+    ----------
+    x
+        The edges of each bin, ``(num_points, num_inputs)``, with the lower
+        edge in the first column and the upper edge in the second.
+    parameters
+        The amplitude, center, and width.
+    f
+        The output array for the model, ``(num_points,)``.
+    jacobian
+        The output array for the derivatives of the model with respect to
+        the parameters, ``(num_points, 3)``.
+
+    Examples
+    --------
+
+    A Gaussian averaged over bins as wide as the line, compared to the
+    Gaussian itself.
+
+    .. jupyter-execute::
+
+        import numpy as np
+        import matplotlib.pyplot as plt
+        import ndoptimize
+
+        parameters = np.array([1.0, 3.0, 10.0])  # amplitude, center and width
+
+        # the Gaussian
+        v = np.linspace(-50, 50, 201)
+        f = np.empty(v.size)
+        ndoptimize.models.gaussian(v[:, np.newaxis], parameters, f, np.empty((v.size, 3)))
+
+        # its average over each bin
+        edges = np.linspace(-50, 50, 11)
+        x = np.stack([edges[:-1], edges[1:]], axis=-1)
+        f_binned = np.empty(len(x))
+        ndoptimize.models.gaussian_binned(x, parameters, f_binned, np.empty((len(x), 3)))
+
+        fig, ax = plt.subplots(constrained_layout=True)
+        ax.plot(v, f, label="gaussian");
+        ax.stairs(f_binned, edges, label="gaussian_binned");
+        ax.legend();
+    """
+    amplitude, center, width = parameters[0], parameters[1], parameters[2]
+    scale = math.sqrt(2) * width
+    for i in range(x.shape[0]):
+        lower = x[i, 0]
+        upper = x[i, 1]
+        size = upper - lower
+        if size == 0:
+            # the limit of a narrow bin, the Gaussian at that point
+            z = (lower - center) / width
+            e = math.exp(-0.5 * z * z)
+            f[i] = amplitude * e
+            jacobian[i, 0] = e
+            jacobian[i, 1] = amplitude * e * z / width
+            jacobian[i, 2] = amplitude * e * z * z / width
+            continue
+        a = (lower - center) / scale
+        b = (upper - center) / scale
+        # in a tail, where both error functions are close to 1 or to -1,
+        # their difference comes from the complementary error functions
+        # without cancellation
+        if a > 0 and b > 0:
+            difference = math.erfc(a) - math.erfc(b)
+        elif a < 0 and b < 0:
+            difference = math.erfc(-b) - math.erfc(-a)
+        else:
+            difference = math.erf(b) - math.erf(a)
+        e_a = math.exp(-a * a)
+        e_b = math.exp(-b * b)
+        factor = math.sqrt(math.pi / 2) * width / size
+        f[i] = amplitude * factor * difference
+        jacobian[i, 0] = factor * difference
+        jacobian[i, 1] = amplitude * (e_a - e_b) / size
+        jacobian[i, 2] = (
+            amplitude
+            * (math.sqrt(math.pi / 2) * difference + math.sqrt(2) * (a * e_a - b * e_b))
+            / size
+        )
+
+
 @_model.model(
     parameters=("amplitude", "center", "width", "slope", "intercept"),
 )