diff --git a/ndoptimize/_tests/test_models.py b/ndoptimize/_tests/test_models.py index 6cbbd7c..c8ce849 100644 --- a/ndoptimize/_tests/test_models.py +++ b/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] diff --git a/ndoptimize/models.py b/ndoptimize/models.py index 817f22d..dc6182a 100644 --- a/ndoptimize/models.py +++ b/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"), )