From 051280d5d439d531eaa589698e0f58a92ed08baa Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sun, 18 Oct 2026 21:41:46 -0600 Subject: [PATCH 1/5] Add `ndoptimize.least_squares()`, a batched Levenberg-Marquardt fitter. `ndoptimize.least_squares()` fits a model to many independent datasets at once, in parallel with Numba, each with its own independent variable, data, noise, initial guess, bounds and Gaussian priors. The noise of neighboring points may be correlated, making the fit generalized least squares through a banded Cholesky whitening. It uses Marquardt damping with Nielsen's gain-ratio update, projected bounds, and the Newton decrement as the convergence test, and returns the parameter covariance. Models are Numba functions compiled with `ndoptimize.model`, which the kernel takes as a typed function pointer so that one cached kernel serves every model. `ndoptimize.models` provides a Gaussian, a Gaussian plus a line, and two Gaussians plus a line. Also sets up the package: a BSD-3-Clause license, citation metadata, tests and typing workflows like `ndfilters`', and the documentation. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- .github/workflows/tests.yml | 43 ++ .github/workflows/typing.yml | 26 + .gitignore | 16 + CITATION.cff | 21 + LICENSE | 28 + README.md | 59 +- docs/conf.py | 2 +- docs/index.rst | 20 +- ndoptimize/__init__.py | 11 + ndoptimize/_least_squares.py | 837 ++++++++++++++++++++++++ ndoptimize/_model.py | 74 +++ ndoptimize/_tests/__init__.py | 0 ndoptimize/_tests/test_least_squares.py | 376 +++++++++++ ndoptimize/_tests/test_models.py | 53 ++ ndoptimize/models.py | 149 +++++ ndoptimize/py.typed | 0 pyproject.toml | 24 +- 17 files changed, 1731 insertions(+), 8 deletions(-) create mode 100644 .github/workflows/tests.yml create mode 100644 .github/workflows/typing.yml create mode 100644 .gitignore create mode 100644 CITATION.cff create mode 100644 LICENSE create mode 100644 ndoptimize/_least_squares.py create mode 100644 ndoptimize/_model.py create mode 100644 ndoptimize/_tests/__init__.py create mode 100644 ndoptimize/_tests/test_least_squares.py create mode 100644 ndoptimize/_tests/test_models.py create mode 100644 ndoptimize/models.py create mode 100644 ndoptimize/py.typed diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..2599471 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,43 @@ + +name: tests + +on: + push: + branches: + - main + pull_request: + + +jobs: + build: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ + ubuntu-latest, + windows-latest, + macOS-latest, + ] + python-version: ["3.10", "3.11", "3.12", "3.13"] + name: ${{ matrix.os }}, Python ${{ matrix.python-version }} lint and test + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install package + run: | + python -m pip install --upgrade pip + pip install setuptools wheel + pip install -e .[test] + - name: Test with pytest + run: | + pip install pytest pytest-cov + pytest --cov=ndoptimize --cov-report=xml + - name: Test with Numba disabled + env: + NUMBA_DISABLE_JIT: 1 + run: | + pytest --cov=ndoptimize --cov-report=xml diff --git a/.github/workflows/typing.yml b/.github/workflows/typing.yml new file mode 100644 index 0000000..c2b0e53 --- /dev/null +++ b/.github/workflows/typing.yml @@ -0,0 +1,26 @@ +name: typing + +on: + push: + branches: + - main + pull_request: + +jobs: + pyright: + name: Check type annotations with pyright + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install package + run: | + python -m pip install --upgrade pip + pip install -e .[test] + - name: Check types with pyright + uses: jakebailey/pyright-action@v2 + with: + version: 1.1.411 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..629d315 --- /dev/null +++ b/.gitignore @@ -0,0 +1,16 @@ +__pycache__/ +*.py[cod] + +/build/ +/dist/ +*.egg-info/ + +docs/_autosummary/ +docs/_build/ +docs/jupyter_execute/ + +/venv/ +.idea/ + +.coverage +coverage.xml diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..1174a71 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,21 @@ +cff-version: 1.2.0 +message: "If you use ndoptimize in your research, please cite it using these metadata." +type: software +title: ndoptimize +abstract: >- + Numba-accelerated minimization of objective functions for many independent + problems at once, including batched Levenberg-Marquardt least squares with + correlated noise, priors and bounds. +authors: + - family-names: Smart + given-names: Roy T. + orcid: "https://orcid.org/0000-0002-9997-5515" +repository-code: "https://github.com/sun-data/ndoptimize" +url: "https://ndoptimize.readthedocs.io" +license: BSD-3-Clause +keywords: + - nonlinear least squares + - Levenberg-Marquardt + - curve fitting + - spectral line fitting + - Numba diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d69756e --- /dev/null +++ b/LICENSE @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2024-2026, Roy T. Smart + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md index 2670a2a..8e6ddd1 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,64 @@ # ndoptimize +[![tests](https://github.com/sun-data/ndoptimize/actions/workflows/tests.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/tests.yml) +[![typing](https://github.com/sun-data/ndoptimize/actions/workflows/typing.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/typing.yml) [![Black](https://github.com/sun-data/ndoptimize/actions/workflows/black.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/black.yml) [![Ruff](https://github.com/sun-data/ndoptimize/actions/workflows/ruff.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/ruff.yml) [![Documentation Status](https://readthedocs.org/projects/ndoptimize/badge/?version=latest)](https://ndoptimize.readthedocs.io/en/latest/?badge=latest) +[![PyPI version](https://badge.fury.io/py/ndoptimize.svg)](https://badge.fury.io/py/ndoptimize) -Numba-accelerated minimization of objective functions. +`ndoptimize` minimizes objective functions for many independent problems at +once, similar to +[`scipy.optimize`](https://docs.scipy.org/doc/scipy/reference/optimize.html), +but compiled and parallelized using +[Numba](https://numba.readthedocs.io/en/stable/). + +Its main function, `ndoptimize.least_squares()`, fits a model to every +spectrum of a spectral cube, or any other batch of small datasets, using the +Levenberg-Marquardt method. Compared to calling +`scipy.optimize.least_squares()` in a loop, it offers: + +- **Speed.** The fits run in parallel on every CPU core, roughly half a + million fits of a five-parameter model to 39 points per second on a + 48-thread machine. +- **Correlated noise.** The noise of neighboring data points may be + correlated, as it is in resampled data, which makes the fit generalized + least squares, with parameter uncertainties to match. +- **Priors and bounds.** Each parameter of each fit may have a Gaussian prior + and bounds, to keep faint fits stable. +- **Missing data.** Data points which are NaN are left out of their fit. +- **Broadcasting.** The data, noise, initial guess, bounds and priors are + broadcast against each other, so each fit can have its own. + +## Installation + +```bash +pip install ndoptimize +``` + +## Example + +```python +import numpy as np +import ndoptimize + +velocity = np.linspace(-100, 100, 41) +data = ... # spectra, shape (..., 41) +variance = ... # their noise variance, the same shape + +result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian_linear, + x=velocity[:, np.newaxis], + data=data, + variance=variance, + guess=[5, 0, 15, 0, 1], # amplitude, center, width, slope, intercept +) + +result.parameters # shape (..., 5) +result.covariance # shape (..., 5, 5) +``` + +Models are short Numba functions which compute the model and its Jacobian, +compiled with the `ndoptimize.model` decorator. +See the [documentation](https://ndoptimize.readthedocs.io/en/latest) for +details. diff --git a/docs/conf.py b/docs/conf.py index 2eda6fd..c5d7ac3 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -20,7 +20,7 @@ # -- Project information ----------------------------------------------------- project = 'ndoptimize' -copyright = '2024, Roy T. Smart' +copyright = '2024-2026, Roy T. Smart' author = 'Roy T. Smart' # -- General configuration --------------------------------------------------- diff --git a/docs/index.rst b/docs/index.rst index 2b7a971..61a2f12 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,11 +1,25 @@ Introduction ============ -:mod:`ndoptimize` is an experimental library designed to parallelize and -accelerate objective function minimization, similar to the :mod:`scipy.optimize` -package, using the `Numba `_ +:mod:`ndoptimize` minimizes objective functions for many independent +problems at once, similar to the :mod:`scipy.optimize` package, but compiled +and parallelized using the `Numba `_ just-in-time compiler. +Its main function, :func:`ndoptimize.least_squares`, fits a model to every +spectrum of a spectral cube, or any other batch of small datasets, using the +Levenberg-Marquardt method. +Each fit has its own independent variable, data, noise, initial guess, +bounds and Gaussian priors, and the noise of neighboring data points may be +correlated, as it is in resampled data, which makes the fit generalized least +squares. +The fits run in parallel on every CPU core, roughly half a million small +fits per second on a 48-thread machine. + +Models are short Numba functions which compute the model and its Jacobian, +compiled with the :func:`ndoptimize.model` decorator. +Some common ones are provided in :mod:`ndoptimize.models`. + Installation ============ diff --git a/ndoptimize/__init__.py b/ndoptimize/__init__.py index b36e432..148cebc 100644 --- a/ndoptimize/__init__.py +++ b/ndoptimize/__init__.py @@ -1,3 +1,14 @@ """ Numba-accelerated minimization of objective functions. """ + +from ._model import model +from . import models +from ._least_squares import LeastSquaresResult, least_squares + +__all__ = [ + "model", + "models", + "LeastSquaresResult", + "least_squares", +] diff --git a/ndoptimize/_least_squares.py b/ndoptimize/_least_squares.py new file mode 100644 index 0000000..7fec6a6 --- /dev/null +++ b/ndoptimize/_least_squares.py @@ -0,0 +1,837 @@ +""" +Nonlinear least squares for many independent fits at once. +""" + +import dataclasses +import numpy as np +import numba +from numba import types +from ._model import Model, signature_model + +__all__ = [ + "LeastSquaresResult", + "least_squares", +] + +_damping_initial = 1e-3 +"""The starting value of the Levenberg-Marquardt damping parameter.""" + +_damping_min = 1e-7 +"""The smallest value of the damping parameter.""" + +_damping_max = 1e7 +"""The largest value of the damping parameter, at which a fit stalls.""" + +_ridge = 1e-12 +"""A small number added to the diagonal of the normal matrix, so that it can +be factored even where a parameter has no effect on the model.""" + + +@numba.njit(cache=True, error_model="numpy") +def _cholesky(a: np.ndarray, m: int, w: int) -> bool: + """ + Overwrite the lower triangle of the leading ``m x m`` block of the + symmetric matrix `a`, of bandwidth `w`, with its Cholesky factor. + Only the band of the lower triangle is read. + + Returns :obj:`False` if `a` is not positive definite. + """ + for j in range(m): + s = a[j, j] + for k in range(max(0, j - w), j): + s -= a[j, k] * a[j, k] + if not s > 0: + return False + d = np.sqrt(s) + a[j, j] = d + for i in range(j + 1, min(m, j + w + 1)): + s = a[i, j] + for k in range(max(0, i - w), j): + s -= a[i, k] * a[j, k] + a[i, j] = s / d + return True + + +@numba.njit(cache=True, error_model="numpy") +def _solve_lower(factor: np.ndarray, b: np.ndarray, m: int, w: int) -> None: + """ + Overwrite the first `m` elements of `b` with the solution of + ``factor @ z = b`` for a lower-triangular `factor` of bandwidth `w`. + """ + for i in range(m): + s = b[i] + for k in range(max(0, i - w), i): + s -= factor[i, k] * b[k] + b[i] = s / factor[i, i] + + +@numba.njit(cache=True, error_model="numpy") +def _solve_lower_columns(factor: np.ndarray, b: np.ndarray, m: int, w: int) -> None: + """ + Overwrite the first `m` rows of every column of `b` with the solution of + ``factor @ z = b`` for a lower-triangular `factor` of bandwidth `w`. + """ + for i in range(m): + d = factor[i, i] + for c in range(b.shape[1]): + s = b[i, c] + for k in range(max(0, i - w), i): + s -= factor[i, k] * b[k, c] + b[i, c] = s / d + + +@numba.njit(cache=True, error_model="numpy") +def _solve_positive_definite( + a: np.ndarray, + b: np.ndarray, + work: np.ndarray, + out: np.ndarray, +) -> bool: + """ + Solve ``a @ x = b`` into `out` for a small, symmetric, positive-definite + matrix `a`, using `work` for its Cholesky factor. + + Returns :obj:`False` if `a` is not positive definite. + """ + n = a.shape[0] + work[:, :] = a + if not _cholesky(work, n, n): + return False + out[:] = b + _solve_lower(work, out, n, n) + for i in range(n - 1, -1, -1): + s = out[i] + for k in range(i + 1, n): + s -= work[k, i] * out[k] + out[i] = s / work[i, i] + return True + + +@numba.njit(cache=True, error_model="numpy") +def _whitened_residual( + factor: np.ndarray, + m: int, + w: int, + data_whitened: np.ndarray, + f: np.ndarray, + out: np.ndarray, +) -> float: + """ + Compute the whitened residual, ``out = data_whitened - factor^-1 @ f``, + and return its sum of squares. + """ + for j in range(m): + out[j] = f[j] + _solve_lower(factor, out, m, w) + total = 0.0 + for j in range(m): + r = data_whitened[j] - out[j] + out[j] = r + total += r * r + return total + + +@numba.njit(cache=True, error_model="numpy") +def _whitened_jacobian( + factor: np.ndarray, + m: int, + w: int, + jacobian: np.ndarray, + out: np.ndarray, +) -> None: + """Compute the whitened Jacobian, ``out = factor^-1 @ jacobian``.""" + for j in range(m): + for k in range(jacobian.shape[1]): + out[j, k] = jacobian[j, k] + _solve_lower_columns(factor, out, m, w) + + +@numba.njit(cache=True, error_model="numpy") +def _normal( + jacobian_whitened: np.ndarray, + residual_whitened: np.ndarray, + m: int, + parameters: np.ndarray, + prior_mean: np.ndarray, + prior_precision: np.ndarray, + a: np.ndarray, + g: np.ndarray, +) -> None: + """ + Compute the normal matrix, ``a = J^T C^-1 J + P``, and the descent + direction, ``g = J^T C^-1 r - P (p - mean)``, from the whitened Jacobian + and residual, where ``P`` is the diagonal prior precision. + """ + n = a.shape[0] + for k in range(n): + s = 0.0 + for j in range(m): + s += jacobian_whitened[j, k] * residual_whitened[j] + g[k] = s - prior_precision[k] * (parameters[k] - prior_mean[k]) + for q in range(k + 1): + s = 0.0 + for j in range(m): + s += jacobian_whitened[j, k] * jacobian_whitened[j, q] + a[k, q] = s + a[q, k] = s + a[k, k] += prior_precision[k] + + +@numba.njit(cache=True, error_model="numpy") +def _chi2_prior( + parameters: np.ndarray, + prior_mean: np.ndarray, + prior_precision: np.ndarray, +) -> float: + """The chi-squared of the parameters against their priors.""" + s = 0.0 + for k in range(parameters.shape[0]): + if prior_precision[k] > 0: + d = parameters[k] - prior_mean[k] + s += prior_precision[k] * d * d + return s + + +_signature_kernel = types.Tuple( + ( + types.float64[:, ::1], + types.float64[:, :, ::1], + types.float64[::1], + types.int64[::1], + types.int64[::1], + types.boolean[::1], + ) +)( + types.FunctionType(signature_model), + types.float64[:, :, ::1], + types.float64[:, ::1], + types.float64[:, ::1], + types.float64[::1], + types.float64[:, ::1], + types.float64[:, ::1], + types.float64[:, ::1], + types.float64[:, ::1], + types.float64[:, ::1], + types.int64, + types.float64, + types.int64, +) +""" +The Numba signature of :func:`_levenberg_marquardt`. + +The model is a typed function pointer, rather than a dispatcher which would be +compiled into the kernel, because Numba cannot cache a function that takes a +dispatcher argument: such a function is compiled anew in every process. +""" + + +@numba.njit(_signature_kernel, parallel=True, cache=True, error_model="numpy") +def _levenberg_marquardt( + model: Model, + x: np.ndarray, + data: np.ndarray, + variance: np.ndarray, + correlation: np.ndarray, + guess: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + prior_mean: np.ndarray, + prior_sigma: np.ndarray, + max_iterations: int, + tolerance: float, + num_threads: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Fit a model to each row of `data` using the Levenberg-Marquardt method. + + The fits are split into `num_threads` interleaved chunks, one per thread, + so that an expensive region of the data is shared between the threads + rather than landing on one of them. + + Parameters + ---------- + model + The model function, compiled with :func:`ndoptimize.model`. + x + The independent variable, ``(num_fits, num_points, num_inputs)``. + data + The data, ``(num_fits, num_points)``, NaN where missing. + variance + The noise variance of each data point, ``(num_fits, num_points)``. + correlation + The correlation of the noise of two points by their separation, in + units of the sampling of the data, starting with 1. + guess + The initial guess for the parameters, ``(num_fits, num_parameters)``. + lower + The lower bound of each parameter, ``(num_fits, num_parameters)``. + upper + The upper bound of each parameter, ``(num_fits, num_parameters)``. + prior_mean + The mean of the Gaussian prior of each parameter, + ``(num_fits, num_parameters)``. + prior_sigma + The standard deviation of the Gaussian prior of each parameter, + ``(num_fits, num_parameters)``, infinite for no prior. + max_iterations + The maximum number of steps per fit. + tolerance + The Newton decrement below which a fit has converged. + num_threads + The number of chunks to split the fits into, the result of + :func:`numba.get_num_threads`, which disables caching if it is called + inside a compiled function. + """ + num_fit, num_point, num_input = x.shape + num_parameter = guess.shape[1] + num_lag = min(correlation.shape[0], num_point) + w = num_lag - 1 + + parameters = np.full((num_fit, num_parameter), np.nan) + covariance = np.full((num_fit, num_parameter, num_parameter), np.nan) + chi2 = np.full(num_fit, np.nan) + num_used = np.zeros(num_fit, dtype=np.int64) + iterations = np.zeros(num_fit, dtype=np.int64) + converged = np.zeros(num_fit, dtype=np.bool_) + + num_chunk = max(1, min(num_threads, num_fit)) + for chunk in numba.prange(num_chunk): # type: ignore[attr-defined] + xg = np.empty((num_point, num_input)) + data_whitened = np.empty(num_point) + sigma = np.empty(num_point) + index = np.empty(num_point, dtype=np.int64) + factor = np.zeros((num_point, num_point)) + f = np.empty(num_point) + jac = np.empty((num_point, num_parameter)) + residual = np.empty(num_point) + residual_new = np.empty(num_point) + jacobian = np.empty((num_point, num_parameter)) + a = np.empty((num_parameter, num_parameter)) + system = np.empty((num_parameter, num_parameter)) + work = np.empty((num_parameter, num_parameter)) + g = np.empty(num_parameter) + rhs = np.empty(num_parameter) + step = np.empty(num_parameter) + unit = np.empty(num_parameter) + p = np.empty(num_parameter) + p_new = np.empty(num_parameter) + mean = np.empty(num_parameter) + precision = np.empty(num_parameter) + free = np.empty(num_parameter, dtype=np.bool_) + + for i in range(chunk, num_fit, num_chunk): + + # gather the usable points + m = 0 + for j in range(num_point): + ok = np.isfinite(data[i, j]) + ok = ok and np.isfinite(variance[i, j]) and variance[i, j] > 0 + for q in range(num_input): + ok = ok and np.isfinite(x[i, j, q]) + if ok: + for q in range(num_input): + xg[m, q] = x[i, j, q] + data_whitened[m] = data[i, j] + sigma[m] = np.sqrt(variance[i, j]) + index[m] = j + m += 1 + num_used[i] = m + if m == 0: + continue + + # factor the banded noise covariance of the usable points and + # whiten the data with it + for r in range(m): + for s in range(max(0, r - w), r + 1): + lag = index[r] - index[s] + rho = correlation[lag] if lag < num_lag else 0.0 + factor[r, s] = sigma[r] * sigma[s] * rho + if not _cholesky(factor, m, w): + continue + _solve_lower(factor, data_whitened, m, w) + + for k in range(num_parameter): + precision[k] = 1 / prior_sigma[i, k] ** 2 + mean[k] = prior_mean[i, k] if np.isfinite(prior_mean[i, k]) else 0 + p[k] = min(max(guess[i, k], lower[i, k]), upper[i, k]) + + model(xg[:m], p, f[:m], jac[:m]) + chi2_data = _whitened_residual(factor, m, w, data_whitened, f, residual) + _whitened_jacobian(factor, m, w, jac, jacobian) + objective = chi2_data + _chi2_prior(p, mean, precision) + + damping = _damping_initial + nu = 2.0 + done = False + n_step = 0 + for _ in range(max_iterations + 1): + _normal(jacobian, residual, m, p, mean, precision, a, g) + + # hold the parameters which are at a bound and are pushed + # outward by the gradient + for k in range(num_parameter): + at_lower = p[k] <= lower[i, k] and g[k] < 0 + at_upper = p[k] >= upper[i, k] and g[k] > 0 + free[k] = not (at_lower or at_upper) + + # the Newton decrement of the free parameters, the decrease + # in chi-squared a Gauss-Newton step would give if the model + # were linear + for k in range(num_parameter): + rhs[k] = g[k] if free[k] else 0 + for q in range(num_parameter): + system[k, q] = a[k, q] if (free[k] and free[q]) else 0 + if free[k]: + system[k, k] += _ridge + else: + system[k, k] = 1 + if _solve_positive_definite(system, rhs, work, step): + decrement = 0.0 + for k in range(num_parameter): + decrement += rhs[k] * step[k] + if decrement < tolerance: + done = True + break + if n_step == max_iterations: + break + + # take a damped step, with more damping until the objective + # decreases + accepted = False + chi2_new = np.inf + objective_new = np.inf + while not accepted: + for k in range(num_parameter): + if free[k]: + system[k, k] = (a[k, k] + _ridge) * (1 + damping) + if _solve_positive_definite(system, rhs, work, step): + for k in range(num_parameter): + p_new[k] = p[k] + step[k] + p_new[k] = min(max(p_new[k], lower[i, k]), upper[i, k]) + model(xg[:m], p_new, f[:m], jac[:m]) + chi2_new = _whitened_residual( + factor, m, w, data_whitened, f, residual_new + ) + objective_new = chi2_new + _chi2_prior(p_new, mean, precision) + accepted = objective_new < objective + if np.isnan(objective) and not np.isnan(objective_new): + accepted = True + if accepted: + # Nielsen's rule: the better the quadratic model + # predicted the decrease, the less damping + predicted = 0.0 + for k in range(num_parameter): + dk = p_new[k] - p[k] + predicted += 2 * g[k] * dk + for q in range(num_parameter): + predicted -= dk * a[k, q] * (p_new[q] - p[q]) + if predicted > 0: + gain = (objective - objective_new) / predicted + else: + gain = -1.0 + factor_damping = max(1 / 3, 1 - (2 * gain - 1) ** 3) + damping = max(damping * factor_damping, _damping_min) + nu = 2.0 + elif damping >= _damping_max: + break + else: + damping = min(damping * nu, _damping_max) + nu *= 2 + n_step += 1 + if not accepted: + break + + p, p_new = p_new, p + residual, residual_new = residual_new, residual + chi2_data = chi2_new + objective = objective_new + _whitened_jacobian(factor, m, w, jac, jacobian) + + # the covariance of the parameters is the inverse of the normal + # matrix, with infinite variance for a parameter which nothing + # constrains, and a small ridge if the rest is still singular + _normal(jacobian, residual, m, p, mean, precision, a, g) + diagonal_max = 0.0 + for k in range(num_parameter): + diagonal_max = max(diagonal_max, a[k, k]) + for attempt in range(2): + for k in range(num_parameter): + for q in range(num_parameter): + constrained = a[k, k] > 0 and a[q, q] > 0 + system[k, q] = a[k, q] if constrained else 0 + if not a[k, k] > 0: + system[k, k] = 1 + elif attempt > 0: + system[k, k] += _ridge * diagonal_max + invertible = True + for k in range(num_parameter): + for q in range(num_parameter): + unit[q] = 1 if q == k else 0 + if not _solve_positive_definite(system, unit, work, step): + invertible = False + break + for q in range(num_parameter): + covariance[i, q, k] = step[q] + if invertible: + break + for k in range(num_parameter): + if not a[k, k] > 0: + for q in range(num_parameter): + covariance[i, k, q] = 0 + covariance[i, q, k] = 0 + covariance[i, k, k] = np.inf + + for k in range(num_parameter): + parameters[i, k] = p[k] + chi2[i] = chi2_data + iterations[i] = n_step + converged[i] = done + + return parameters, covariance, chi2, num_used, iterations, converged + + +@dataclasses.dataclass(frozen=True) +class LeastSquaresResult: + """ + The results of :func:`least_squares` for every fit of a batch. + + The leading axes of each array are the batch axes of the fit. + """ + + parameters: np.ndarray + """The best-fit parameters, ``(..., num_parameters)``.""" + + covariance: np.ndarray + """ + The covariance of the best-fit parameters, + ``(..., num_parameters, num_parameters)``. + + This is the inverse of the normal matrix, :math:`J^T C^{-1} J`, plus the + precision of the priors, at the solution. + A parameter which neither the data nor its prior constrains has an + infinite variance. + """ + + chi2: np.ndarray + """ + The chi-squared of the data at the solution, without the contribution of + the priors. + """ + + num_points: np.ndarray + """The number of data points used by each fit.""" + + iterations: np.ndarray + """The number of steps taken by each fit.""" + + converged: np.ndarray + """Whether each fit converged.""" + + +def _flatten( + x: np.ndarray, + data: np.ndarray, + variance: float | np.ndarray, + guess: np.ndarray, + lower: float | np.ndarray, + upper: float | np.ndarray, + prior_mean: float | np.ndarray, + prior_sigma: float | np.ndarray, +) -> tuple[tuple[int, ...], dict[str, np.ndarray]]: + """ + Broadcast the arguments of :func:`least_squares` against each other over + their batch axes and flatten those axes, into writable, contiguous + arrays of the types the kernel's signature requires. + + Returns the batch shape and the flattened arrays. + """ + x = np.asarray(x, dtype=float) + data = np.asarray(data, dtype=float) + guess = np.asarray(guess, dtype=float) + + if x.ndim < 2: + raise ValueError( + f"`x` must have at least two axes, (..., num_points, num_inputs), " + f"got shape {x.shape}." + ) + num_point, num_input = x.shape[-2:] + num_parameter = guess.shape[-1] + + if data.shape[-1:] != (num_point,): + raise ValueError( + f"The last axis of `data`, {data.shape}, must match the " + f"number of points in `x`, {num_point}." + ) + + variance = np.broadcast_to( + variance, np.broadcast_shapes(np.shape(variance), (num_point,)) + ) + lower, upper, prior_mean, prior_sigma = ( + np.broadcast_to(a, np.broadcast_shapes(np.shape(a), (num_parameter,))) + for a in (lower, upper, prior_mean, prior_sigma) + ) + + shape = np.broadcast_shapes( + x.shape[:-2], + data.shape[:-1], + variance.shape[:-1], + guess.shape[:-1], + lower.shape[:-1], + upper.shape[:-1], + prior_mean.shape[:-1], + prior_sigma.shape[:-1], + ) + num_fit = int(np.prod(shape)) + + def flat(a: np.ndarray, *shape_trailing: int) -> np.ndarray: + """A writable, contiguous copy of `a` with its batch axes flattened.""" + a = np.broadcast_to(a, shape + shape_trailing) + return np.array(a, dtype=float, order="C").reshape(num_fit, *shape_trailing) + + arrays = dict( + x=flat(x, num_point, num_input), + data=flat(data, num_point), + variance=flat(variance, num_point), + guess=flat(guess, num_parameter), + lower=flat(lower, num_parameter), + upper=flat(upper, num_parameter), + prior_mean=flat(prior_mean, num_parameter), + prior_sigma=flat(prior_sigma, num_parameter), + ) + + return shape, arrays + + +def least_squares( + model: Model, + x: np.ndarray, + data: np.ndarray, + variance: float | np.ndarray, + guess: np.ndarray, + lower: float | np.ndarray = -np.inf, + upper: float | np.ndarray = np.inf, + prior_mean: float | np.ndarray = 0, + prior_sigma: float | np.ndarray = np.inf, + correlation: None | np.ndarray = None, + max_iterations: int = 1000, + tolerance: float = 1e-8, +) -> LeastSquaresResult: + r""" + Fit a model to many independent sets of data at once, using the + Levenberg-Marquardt method. + + Every fit has the same model, but its own independent variable, data, + noise, initial guess, bounds and priors. + The fits run in parallel, each with its own damping and convergence test. + + Each fit minimizes + + .. math:: + + \chi^2 = (y - f)^T C^{-1} (y - f) + + \sum_k \left( \frac{p_k - \bar{p}_k}{\sigma_k} \right)^2, + + where :math:`y` is the data, :math:`f` is the model, + :math:`p_k` are the parameters, :math:`\bar{p}_k` and :math:`\sigma_k` + are the means and standard deviations of their Gaussian priors, + and :math:`C` is the covariance of the noise, + + .. math:: + + C_{ij} = \sigma_i \sigma_j \rho_{|i - j|}, + + with :math:`\sigma_i^2` the variance of data point :math:`i` + and :math:`\rho_k` the correlation of the noise of two points + :math:`k` samples apart. + Uncorrelated noise, :math:`\rho = (1,)`, is ordinary weighted least + squares. + Correlated noise, from resampling for instance, makes this generalized + least squares, which is computed by whitening the residual and the + Jacobian with the banded Cholesky factor of :math:`C`. + + Parameters + ---------- + model + The model to fit, a function compiled with :func:`ndoptimize.model`. + x + The independent variable, an array of shape + ``(..., num_points, num_inputs)``. + data + The data to fit, an array of shape ``(..., num_points)``. + Points which are NaN are left out of their fit. + variance + The variance of the noise of each data point, an array of shape + ``(..., num_points)``. + Points with a variance which is NaN or not positive are left out of + their fit. + guess + The initial guess for the parameters, an array of shape + ``(..., num_parameters)``. + lower + The lower bound of each parameter, an array of shape + ``(..., num_parameters)``. + upper + The upper bound of each parameter, an array of shape + ``(..., num_parameters)``. + prior_mean + The mean of the Gaussian prior of each parameter, an array of shape + ``(..., num_parameters)``. + prior_sigma + The standard deviation of the Gaussian prior of each parameter, an + array of shape ``(..., num_parameters)``. + An infinite standard deviation, the default, means no prior. + correlation + The correlation of the noise of two data points by their separation + along the last axis of `data`, in samples, starting with 1 for no + separation and zero beyond its length. + The separation counts the points left out of a fit. + If :obj:`None` (the default), the noise is uncorrelated. + max_iterations + The maximum number of steps for each fit. + tolerance + The convergence tolerance on the Newton decrement, + :math:`g^T A^{-1} g`, where :math:`g = J^T C^{-1} r` is the gradient + and :math:`A = J^T C^{-1} J` is the normal matrix, which is the + decrease in :math:`\chi^2` a Gauss-Newton step would give if the + model were linear. + The default, :math:`10^{-8}`, puts each parameter within about + :math:`10^{-4}` of its uncertainty of the minimum. + + Returns + ------- + The best-fit parameters, their covariance, chi-squared, and the + convergence state of each fit. + Fits with no usable data points are NaN. + + Notes + ----- + + The batch axes of the arguments, all but their last one (two for `x`), + are broadcast against each other. + + Each step solves :math:`(A + \lambda \, \text{diag} A) \delta = g`, + Marquardt's damping, and the damping parameter :math:`\lambda` is updated + by the gain-ratio rule of Madsen, Nielsen, and Tingleff (2004), "Methods + for Non-Linear Least Squares Problems", which converges in curved + valleys of :math:`\chi^2` (where the amplitude and width of a weak line + trade off) in about a third of the steps of the classic fixed factors. + + Steps are clipped to the bounds, and a parameter at a bound which the + gradient pushes outward is held fixed for the step and for the + convergence test. + + Examples + -------- + + Fit a Gaussian plus a straight line to a thousand noisy spectral lines. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import ndoptimize + + rng = np.random.default_rng(0) + + # the true parameters: amplitude, center, width, slope and intercept + num = 1000 + truth = np.stack([ + rng.uniform(2, 10, num), + rng.normal(0, 5, num), + rng.uniform(10, 20, num), + np.zeros(num), + rng.uniform(1, 2, num), + ], axis=-1) + + # the noisy spectra + velocity = np.linspace(-100, 100, 41) + x = velocity[:, np.newaxis] + f = np.empty((num, velocity.size)) + jacobian = np.empty((velocity.size, 5)) + for i in range(num): + ndoptimize.models.gaussian_linear(x, truth[i], f[i], jacobian) + variance = 0.5 + data = f + rng.normal(0, np.sqrt(variance), f.shape) + + # fit every spectrum + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian_linear, + x=x, + data=data, + variance=variance, + guess=[5, 0, 15, 0, 1], + ) + + # compare the fitted centers to the truth + sigma = np.sqrt(result.covariance[..., 1, 1]) + error = (result.parameters[..., 1] - truth[..., 1]) / sigma + + fig, ax = plt.subplots(constrained_layout=True) + ax.hist(error, bins=40, density=True, label="fits"); + z = np.linspace(-4, 4, 101) + ax.plot(z, np.exp(-z**2 / 2) / np.sqrt(2 * np.pi), label="unit normal"); + ax.set_xlabel("(fitted center - true center) / uncertainty"); + ax.legend(); + """ + shape, a = _flatten( + x=x, + data=data, + variance=variance, + guess=guess, + lower=lower, + upper=upper, + prior_mean=prior_mean, + prior_sigma=prior_sigma, + ) + + if correlation is None: + correlation = np.ones(1) + correlation = np.array(correlation, dtype=float).reshape(-1) + if correlation.size == 0 or correlation[0] != 1: + raise ValueError( + f"`correlation` must start with 1, the correlation of a point " + f"with itself, got {correlation}." + ) + + # the noise correlation of every fit is a principal submatrix of this + # one, so it is positive definite if this one is + num_point = a["x"].shape[1] + rho = np.zeros(num_point) + rho[: correlation.size] = correlation[:num_point] + index = np.arange(num_point) + try: + np.linalg.cholesky(rho[np.abs(index[:, np.newaxis] - index)]) + except np.linalg.LinAlgError: + raise ValueError( + f"`correlation`, {correlation}, does not describe a positive " + f"definite noise covariance for {num_point} points." + ) + + result = _levenberg_marquardt( + model, + a["x"], + a["data"], + a["variance"], + correlation, + a["guess"], + a["lower"], + a["upper"], + a["prior_mean"], + a["prior_sigma"], + max_iterations, + tolerance, + numba.get_num_threads(), + ) + + parameters, covariance, chi2, num_points, iterations, converged = ( + r.reshape(shape + r.shape[1:]) for r in result + ) + + return LeastSquaresResult( + parameters=parameters, + covariance=covariance, + chi2=chi2, + num_points=num_points, + iterations=iterations, + converged=converged, + ) diff --git a/ndoptimize/_model.py b/ndoptimize/_model.py new file mode 100644 index 0000000..c621ab5 --- /dev/null +++ b/ndoptimize/_model.py @@ -0,0 +1,74 @@ +""" +The interface between a model function and the fitting kernels. +""" + +from typing import Callable +import numpy as np +import numba +from numba import types + +__all__ = [ + "Model", + "signature_model", + "model", +] + +Model = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None] +""" +The type of a model function, ``model(x, parameters, f, jacobian)``. +""" + +signature_model = types.void( + types.float64[:, :], + types.float64[:], + types.float64[:], + types.float64[:, :], +) +""" +The Numba signature of a model function. +""" + + +def model(function: Model) -> Model: + r""" + Compile a function as a model which :func:`ndoptimize.least_squares` can fit. + + A model is a function ``model(x, parameters, f, jacobian)`` which writes + the model evaluated at the points ``x``, an array of shape + ``(num_points, num_inputs)``, for the given ``parameters``, an array of + shape ``(num_parameters,)``, into ``f``, an array of shape + ``(num_points,)``, and its derivatives with respect to the parameters + into ``jacobian``, an array of shape ``(num_points, num_parameters)``. + It returns nothing. + + The function is compiled with Numba in `nopython` mode, with this fixed + signature. + The fitting kernel takes the model as a typed function pointer, so a + single compiled kernel serves every model and is cached on disk. + + Scalar functions such as :func:`math.exp` should come from :mod:`math`, + since they are the ones Numba compiles for every target. + + Parameters + ---------- + function + The model function to compile. + + Examples + -------- + + A straight line, :math:`f(x) = a + b x`. + + .. jupyter-execute:: + + import ndoptimize + + @ndoptimize.model + def line(x, parameters, f, jacobian): + a, b = parameters[0], parameters[1] + for i in range(x.shape[0]): + f[i] = a + b * x[i, 0] + jacobian[i, 0] = 1 + jacobian[i, 1] = x[i, 0] + """ + return numba.njit(signature_model, cache=True, error_model="numpy")(function) diff --git a/ndoptimize/_tests/__init__.py b/ndoptimize/_tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ndoptimize/_tests/test_least_squares.py b/ndoptimize/_tests/test_least_squares.py new file mode 100644 index 0000000..8d24700 --- /dev/null +++ b/ndoptimize/_tests/test_least_squares.py @@ -0,0 +1,376 @@ +from typing import Any +import numpy as np +import pytest +import scipy.optimize +import ndoptimize + +_velocity = np.linspace(-112, 112, 39) +"""The independent variable of the test spectra.""" + +_correlation = np.array([1, 0.187, -0.048, 0.007, 0.002, 0.005]) +"""The correlation by lag of the noise of resampled spectra, from IRIS.""" + + +def _noise_covariance(sigma: np.ndarray, correlation: None | np.ndarray) -> np.ndarray: + """The banded covariance of the noise of one spectrum.""" + if correlation is None: + correlation = np.ones(1) + num = sigma.size + lag = np.abs(np.arange(num)[:, np.newaxis] - np.arange(num)) + rho = np.zeros(num) + rho[: correlation.size] = correlation[:num] + return sigma[:, np.newaxis] * rho[lag] * sigma + + +def _evaluate(model: ndoptimize._model.Model, parameters: np.ndarray) -> np.ndarray: + """The model at the test velocities for each row of `parameters`.""" + x = _velocity[:, np.newaxis] + result = np.empty((len(parameters), _velocity.size)) + jacobian = np.empty((_velocity.size, parameters.shape[-1])) + for i in range(len(parameters)): + model(x, parameters[i], result[i], jacobian) + return result + + +def _spectra( + num: int, + correlation: None | np.ndarray, + seed: int = 0, + missing: float = 0.05, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Bright Gaussian-plus-line spectra with noise of the given correlation and + some missing points: the truth, the data, the variance and a guess. + + The lines are bright so that every fit has a minimum without priors or + bounds; a faint line in correlated noise can fit better as the wing of + an ever broader Gaussian, which has no minimum at all. + """ + rng = np.random.default_rng(seed) + truth = np.stack( + [ + rng.uniform(4, 8, num), + rng.normal(0, 5, num), + rng.uniform(12, 25, num), + rng.normal(0, 0.001, num), + rng.uniform(0.5, 2, num), + ], + axis=-1, + ) + f = _evaluate(ndoptimize.models.gaussian_linear, truth) + variance = 0.05 + np.clip(f, 0, None) / 8 + factor = np.linalg.cholesky(_noise_covariance(np.ones(_velocity.size), correlation)) + noise = rng.standard_normal(f.shape) @ factor.T * np.sqrt(variance) + data = f + noise + data[rng.random(data.shape) < missing] = np.nan + guess = truth * np.array([1.2, 1, 1.1, 0, 1]) + np.array([0, 3, 0, 0, 0.1]) + return truth, data, variance, guess + + +def _fit_scipy( + data: np.ndarray, + variance: np.ndarray, + guess: np.ndarray, + correlation: None | np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + prior_mean: np.ndarray, + prior_sigma: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, float]: + """ + Fit one spectrum with :func:`scipy.optimize.least_squares`, whitening the + residual with the dense Cholesky factor of the noise covariance of the + usable points: the parameters, their covariance and the objective. + """ + good = np.isfinite(data) + covariance = _noise_covariance(np.sqrt(variance), correlation) + factor = np.linalg.cholesky(covariance[np.ix_(good, good)]) + has_prior = np.isfinite(prior_sigma) + x = _velocity[good, np.newaxis] + f = np.empty(good.sum()) + jacobian = np.empty((good.sum(), guess.size)) + + def residual(p: np.ndarray) -> np.ndarray: + """The whitened residual of the data and of the priors.""" + ndoptimize.models.gaussian_linear(x, p, f, jacobian) + r = np.linalg.solve(factor, data[good] - f) + return np.concatenate([r, ((p - prior_mean) / prior_sigma)[has_prior]]) + + bounded = np.isfinite(lower).any() or np.isfinite(upper).any() + result = scipy.optimize.least_squares( + fun=residual, + x0=np.clip(guess, lower, upper), + method="trf" if bounded else "lm", + bounds=(lower, upper), + ftol=1e-15, + xtol=1e-15, + gtol=1e-15, + x_scale="jac", + max_nfev=10000, + ) + jacobian = np.asarray(result.jac) + covariance = np.linalg.inv(jacobian.T @ jacobian) + return result.x, covariance, float(np.sum(result.fun**2)) + + +@pytest.mark.parametrize( + argnames="correlation", + argvalues=[None, _correlation], +) +@pytest.mark.parametrize( + argnames="prior", + argvalues=[False, True], +) +@pytest.mark.parametrize( + argnames="bounded", + argvalues=[False, True], +) +def test_least_squares( + correlation: None | np.ndarray, + prior: bool, + bounded: bool, +) -> None: + truth, data, variance, guess = _spectra(20, correlation) + + infinity = np.full(5, np.inf) + prior_mean = np.array([0, 0, 18, 0, 0]) + prior_sigma = np.array([np.inf, 10, 6, np.inf, np.inf]) if prior else infinity + lower = np.array([-np.inf, -90, 3, -np.inf, -np.inf]) if bounded else -infinity + upper = np.array([np.inf, 90, 16, np.inf, np.inf]) if bounded else infinity + + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + data=data, + variance=variance, + guess=guess, + lower=lower, + upper=upper, + prior_mean=prior_mean, + prior_sigma=prior_sigma, + correlation=correlation, + ) + + assert np.all(result.converged) + assert np.all(result.num_points == np.isfinite(data).sum(axis=-1)) + if bounded: + assert np.any(result.parameters[..., 2] == upper[2]) + + for i in range(len(data)): + parameters, covariance, objective = _fit_scipy( + data=data[i], + variance=variance[i], + guess=guess[i], + correlation=correlation, + lower=lower, + upper=upper, + prior_mean=prior_mean, + prior_sigma=prior_sigma, + ) + chi2_prior = np.where( + np.isfinite(prior_sigma), + ((result.parameters[i] - prior_mean) / prior_sigma) ** 2, + 0, + ) + assert result.chi2[i] + chi2_prior.sum() <= objective + 1e-6 + + sigma = np.sqrt(np.diag(result.covariance[i])) + assert np.all(np.abs(result.parameters[i] - parameters) < 1e-3 * sigma) + difference = (result.covariance[i] - covariance) / np.outer(sigma, sigma) + assert np.all(np.abs(difference) < 1e-3) + + +def test_least_squares_calibration() -> None: + """ + The covariance of the fits is the scatter of the parameters over many + realizations of the noise, if the noise correlation is given, and + underestimates it if it is not. + """ + num = 500 + rng = np.random.default_rng(1) + truth = np.array([4, 2, 16, 0, 1.0]) + f = _evaluate(ndoptimize.models.gaussian_linear, truth[np.newaxis])[0] + variance = 0.05 + f / 8 + covariance = _noise_covariance(np.sqrt(variance), _correlation) + data = f + rng.standard_normal((num, f.size)) @ np.linalg.cholesky(covariance).T + + def ratio(correlation: None | np.ndarray) -> np.ndarray: + """The scatter of the fits over their covariance, for each parameter.""" + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + data=data, + variance=variance, + guess=truth, + correlation=correlation, + ) + assert np.all(result.converged) + scatter = np.var(result.parameters, axis=0) + return scatter / np.mean( + np.diagonal(result.covariance, axis1=-2, axis2=-1), axis=0 + ) + + assert np.all(np.abs(ratio(_correlation) - 1) < 0.2) + assert np.all(ratio(None)[:3] > 1.15) + + +def test_least_squares_broadcasting() -> None: + truth, data, variance, guess = _spectra(6, None) + kwargs: dict[str, Any] = dict( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + variance=variance.mean(), + guess=guess.mean(axis=0), + ) + + result = ndoptimize.least_squares(data=data.reshape(2, 3, -1), **kwargs) + result_flat = ndoptimize.least_squares(data=data, **kwargs) + + assert result.parameters.shape == (2, 3, 5) + assert result.covariance.shape == (2, 3, 5, 5) + for name in ("chi2", "num_points", "iterations", "converged"): + assert getattr(result, name).shape == (2, 3) + assert np.array_equal(result.parameters.reshape(6, 5), result_flat.parameters) + + +def test_least_squares_missing() -> None: + truth, data, variance, guess = _spectra(3, None, missing=0) + data[0] = np.nan + data[1, 10:15] = np.nan + kwargs: dict[str, Any] = dict( + model=ndoptimize.models.gaussian_linear, + variance=variance[1], + guess=guess[1], + ) + + result = ndoptimize.least_squares(x=_velocity[:, np.newaxis], data=data, **kwargs) + + assert np.all(np.isnan(result.parameters[0])) + assert result.num_points[0] == 0 + assert not result.converged[0] + assert result.num_points[1] == _velocity.size - 5 + + good = np.isfinite(data[1]) + result_good = ndoptimize.least_squares( + model=kwargs["model"], + x=_velocity[good, np.newaxis], + data=data[1, good], + variance=variance[1, good], + guess=guess[1], + ) + assert np.allclose(result.parameters[1], result_good.parameters, rtol=1e-10) + + +def test_least_squares_unconstrained() -> None: + """A parameter which nothing constrains has an infinite variance.""" + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian, + x=_velocity[:, np.newaxis], + data=np.zeros(_velocity.size), + variance=1, + guess=np.array([0, 0, 15.0]), + lower=np.array([0, -np.inf, 0]), + upper=np.array([0, np.inf, np.inf]), + ) + assert result.converged + assert np.all(np.isinf(np.diagonal(result.covariance)[1:])) + assert np.isfinite(result.covariance[0, 0]) + + +@ndoptimize.model +def _line( + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, +) -> None: + """A straight line.""" + for i in range(x.shape[0]): + f[i] = parameters[0] + parameters[1] * x[i, 0] + jacobian[i, 0] = 1 + jacobian[i, 1] = x[i, 0] + + +def test_least_squares_linear() -> None: + """A user-defined linear model gives the weighted linear least squares + solution in a few steps, to well within the tolerance.""" + rng = np.random.default_rng(2) + x = np.linspace(0, 1, 11) + variance = rng.uniform(0.5, 2, x.size) + data = 1 + 2 * x + rng.normal(0, np.sqrt(variance)) + + result = ndoptimize.least_squares( + model=_line, + x=x[:, np.newaxis], + data=data, + variance=variance, + guess=np.zeros(2), + ) + + design = np.stack([np.ones_like(x), x], axis=-1) / np.sqrt(variance)[:, np.newaxis] + expected = np.linalg.lstsq(design, data / np.sqrt(variance), rcond=None)[0] + covariance = np.linalg.inv(design.T @ design) + sigma = np.sqrt(np.diag(covariance)) + assert np.all(np.abs(result.parameters - expected) < 1e-4 * sigma) + assert np.allclose(result.covariance, covariance, rtol=1e-10) + assert result.converged + assert result.iterations <= 3 + + +def test_least_squares_far_guess() -> None: + """Fits started far from the solution, which take damped and rejected + steps, reach the same minimum as fits started near it.""" + truth, data, variance, guess = _spectra(10, _correlation) + kwargs: dict[str, Any] = dict( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + data=data, + variance=variance, + lower=np.array([-np.inf, -np.inf, 3, -np.inf, -np.inf]), + correlation=_correlation, + ) + + result = ndoptimize.least_squares(guess=guess, **kwargs) + result_far = ndoptimize.least_squares(guess=np.array([2, 10, 30, 0, 1]), **kwargs) + + assert np.all(result_far.converged) + assert result_far.iterations.sum() > result.iterations.sum() + sigma = np.sqrt(np.diagonal(result.covariance, axis1=-2, axis2=-1)) + assert np.all(np.abs(result_far.parameters - result.parameters) < 1e-3 * sigma) + + +def test_least_squares_max_iterations() -> None: + truth, data, variance, guess = _spectra(5, None) + + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + data=data, + variance=variance, + guess=guess, + max_iterations=1, + ) + + assert not np.any(result.converged) + assert np.all(result.iterations == 1) + + +@pytest.mark.parametrize( + argnames="kwargs", + argvalues=[ + dict(correlation=np.array([0.5, 0.1])), + dict(correlation=np.array([1, 0.9])), + dict(data=np.zeros(5)), + dict(x=_velocity), + ], +) +def test_least_squares_invalid(kwargs: dict[str, Any]) -> None: + arguments: dict[str, Any] = dict( + model=ndoptimize.models.gaussian, + x=_velocity[:, np.newaxis], + data=np.zeros(_velocity.size), + variance=1, + guess=np.array([1, 0, 15.0]), + ) + with pytest.raises(ValueError): + ndoptimize.least_squares(**(arguments | kwargs)) diff --git a/ndoptimize/_tests/test_models.py b/ndoptimize/_tests/test_models.py new file mode 100644 index 0000000..e0aa024 --- /dev/null +++ b/ndoptimize/_tests/test_models.py @@ -0,0 +1,53 @@ +from typing import Callable +import numpy as np +import pytest +import ndoptimize + + +def _gaussian(v: np.ndarray, a: float, c: float, w: float) -> np.ndarray: + """A Gaussian evaluated with numpy.""" + return a * np.exp(-0.5 * ((v - c) / w) ** 2) + + +@pytest.mark.parametrize( + argnames="model,parameters,function", + argvalues=[ + ( + ndoptimize.models.gaussian, + np.array([2.0, 3.0, 15.0]), + lambda v, p: _gaussian(v, *p), + ), + ( + ndoptimize.models.gaussian_linear, + np.array([2.0, 3.0, 15.0, 0.01, 1.5]), + lambda v, p: _gaussian(v, *p[:3]) + p[3] * v + p[4], + ), + ( + ndoptimize.models.two_gaussians_linear, + np.array([2.0, 3.0, 15.0, 0.01, 1.5, 1.0, -20.0, 4.0]), + lambda v, p: _gaussian(v, *p[:3]) + p[3] * v + p[4] + _gaussian(v, *p[5:]), + ), + ], +) +def test_models( + model: ndoptimize._model.Model, + parameters: np.ndarray, + function: Callable[[np.ndarray, np.ndarray], np.ndarray], +) -> None: + v = np.linspace(-50, 50, 21) + x = v[:, np.newaxis] + f = np.empty(v.size) + jacobian = np.empty((v.size, parameters.size)) + + model(x, parameters, f, jacobian) + + assert np.allclose(f, function(v, parameters)) + + step = 1e-6 + for k in range(parameters.size): + dp = np.zeros(parameters.size) + dp[k] = step + derivative = (function(v, parameters + dp) - function(v, parameters - dp)) / ( + 2 * step + ) + assert np.allclose(jacobian[:, k], derivative, rtol=1e-5, atol=1e-8) diff --git a/ndoptimize/models.py b/ndoptimize/models.py new file mode 100644 index 0000000..43ab7c2 --- /dev/null +++ b/ndoptimize/models.py @@ -0,0 +1,149 @@ +""" +Common model functions, compiled with :func:`ndoptimize.model`. +""" + +import math +import numpy as np +from . import _model + +__all__ = [ + "gaussian", + "gaussian_linear", + "two_gaussians_linear", +] + + +@_model.model +def gaussian( + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, +) -> None: + r""" + A Gaussian, + + .. math:: + + f(x) = A \exp \left( -\frac{(x - \mu)^2}{2 \sigma^2} \right), + + with parameters :math:`(A, \mu, \sigma)`, the amplitude, center and + width, of the first component of :math:`x`. + + Parameters + ---------- + x + The points at which to evaluate the model, ``(num_points, num_inputs)``. + 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)``. + """ + amplitude, center, width = parameters[0], parameters[1], parameters[2] + for i in range(x.shape[0]): + z = (x[i, 0] - 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 + + +@_model.model +def gaussian_linear( + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, +) -> None: + r""" + A Gaussian plus a straight line, + + .. math:: + + f(x) = A \exp \left( -\frac{(x - \mu)^2}{2 \sigma^2} \right) + m x + b, + + with parameters :math:`(A, \mu, \sigma, m, b)`, the amplitude, center, + width, slope and intercept, of the first component of :math:`x`. + The line takes up a continuum or background under a spectral line. + + Parameters + ---------- + x + The points at which to evaluate the model, ``(num_points, num_inputs)``. + parameters + The amplitude, center, width, slope and intercept. + 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, 5)``. + """ + amplitude, center, width = parameters[0], parameters[1], parameters[2] + slope, intercept = parameters[3], parameters[4] + for i in range(x.shape[0]): + v = x[i, 0] + z = (v - center) / width + e = math.exp(-0.5 * z * z) + f[i] = amplitude * e + slope * v + intercept + jacobian[i, 0] = e + jacobian[i, 1] = amplitude * e * z / width + jacobian[i, 2] = amplitude * e * z * z / width + jacobian[i, 3] = v + jacobian[i, 4] = 1 + + +@_model.model +def two_gaussians_linear( + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, +) -> None: + r""" + Two Gaussians plus a straight line, + + .. math:: + + f(x) = A \exp \left( -\frac{(x - \mu)^2}{2 \sigma^2} \right) + m x + b + + A_2 \exp \left( -\frac{(x - \mu_2)^2}{2 \sigma_2^2} \right), + + with parameters :math:`(A, \mu, \sigma, m, b, A_2, \mu_2, \sigma_2)`, of + the first component of :math:`x`. + The first five parameters are those of :func:`gaussian_linear`, so its + fit is a natural starting point for this one. + + Parameters + ---------- + x + The points at which to evaluate the model, ``(num_points, num_inputs)``. + parameters + The amplitude, center, width, slope and intercept, then the + amplitude, center and width of the second Gaussian. + 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, 8)``. + """ + amplitude, center, width = parameters[0], parameters[1], parameters[2] + slope, intercept = parameters[3], parameters[4] + amplitude_2, center_2, width_2 = parameters[5], parameters[6], parameters[7] + for i in range(x.shape[0]): + v = x[i, 0] + z = (v - center) / width + e = math.exp(-0.5 * z * z) + z_2 = (v - center_2) / width_2 + e_2 = math.exp(-0.5 * z_2 * z_2) + f[i] = amplitude * e + slope * v + intercept + amplitude_2 * e_2 + jacobian[i, 0] = e + jacobian[i, 1] = amplitude * e * z / width + jacobian[i, 2] = amplitude * e * z * z / width + jacobian[i, 3] = v + jacobian[i, 4] = 1 + jacobian[i, 5] = e_2 + jacobian[i, 6] = amplitude_2 * e_2 * z_2 / width_2 + jacobian[i, 7] = amplitude_2 * e_2 * z_2 * z_2 / width_2 diff --git a/ndoptimize/py.typed b/ndoptimize/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml index 2681603..d5ebbad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools", "setuptools-scm"] +requires = ["setuptools>=77", "setuptools-scm"] build-backend = "setuptools.build_meta" [project] @@ -10,13 +10,20 @@ authors = [ description = "Numba-accelerated minimization of objective functions." readme = "README.md" requires-python = ">=3.10" +license = "BSD-3-Clause" +license-files = ["LICENSE"] classifiers = [ + "Intended Audience :: Science/Research", "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Astronomy", ] dependencies = [ - "numpy", + "numpy>2", "numba", - "astropy", ] dynamic = ["version"] @@ -27,6 +34,7 @@ test = [ ] doc = [ "pytest", + "scipy", "matplotlib", "graphviz", "sphinx-autodoc-typehints", @@ -40,7 +48,17 @@ doc = [ Homepage = "https://github.com/sun-data/ndoptimize" Documentation = "https://ndoptimize.readthedocs.io/en/latest" +[tool.pyright] +include = ["ndoptimize"] +typeCheckingMode = "standard" + [tool.setuptools] packages = ["ndoptimize"] +[tool.setuptools.package-data] +ndoptimize = ["py.typed"] + [tool.setuptools_scm] + +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F"] From d456eb068069a724cc66f4312243ce312861ff63 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Mon, 19 Oct 2026 09:24:51 -0600 Subject: [PATCH 2/5] Upload test coverage to Codecov. The repository now has a `CODECOV_TOKEN` secret, so the tests workflow uploads the coverage report like `ndfilters`' does, and the README gets a Codecov badge. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- .github/workflows/tests.yml | 9 +++++++++ README.md | 1 + 2 files changed, 10 insertions(+) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2599471..da2e352 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -41,3 +41,12 @@ jobs: NUMBA_DISABLE_JIT: 1 run: | pytest --cov=ndoptimize --cov-report=xml + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v7 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: coverage.xml + flags: unittests + env_vars: OS,PYTHON + name: codecov-umbrella + fail_ci_if_error: true diff --git a/README.md b/README.md index 8e6ddd1..9f27b97 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,7 @@ # ndoptimize [![tests](https://github.com/sun-data/ndoptimize/actions/workflows/tests.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/tests.yml) +[![codecov](https://codecov.io/gh/sun-data/ndoptimize/graph/badge.svg)](https://codecov.io/gh/sun-data/ndoptimize) [![typing](https://github.com/sun-data/ndoptimize/actions/workflows/typing.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/typing.yml) [![Black](https://github.com/sun-data/ndoptimize/actions/workflows/black.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/black.yml) [![Ruff](https://github.com/sun-data/ndoptimize/actions/workflows/ruff.yml/badge.svg)](https://github.com/sun-data/ndoptimize/actions/workflows/ruff.yml) From 4278d3d8fd1b3e2124664e1c8a72c6d723f07253 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Mon, 19 Oct 2026 10:33:22 -0600 Subject: [PATCH 3/5] Fix the issues found in code review. - Models name their parameters, `ndoptimize.model(parameters=...)`, and are `ndoptimize.Model` objects. `least_squares()` checks the guess against them, since a mismatch made the compiled model read and write outside its arrays, and calling a model directly checks the shapes of its arrays. - `ndoptimize.model` takes `cache=False`, for models not defined in a file or reading globals computed at run time. - The ridge in the convergence test is a fraction of each diagonal element of the normal matrix rather than 1e-12, which stopped fits of parameters with a very different scale than the data at their guess, reported as converged. Parameters with no effect on the model are held instead. - The covariance is NaN where the normal matrix is singular, as in a fit with fewer usable points than parameters, rather than finite numbers set by a ridge. - `least_squares()` rejects `lower > upper`, NaN bounds, a non-positive or NaN `prior_sigma`, a non-finite `prior_mean` with a prior, a non-finite `correlation`, a guess of the wrong shape and a negative `max_iterations`. A `prior_mean` without a prior is ignored. - The noise factor is stored as a band, not a dense matrix per thread, and the positive-definite check of `correlation` is a banded Cholesky too. - Arguments shared by every fit are passed to the kernel as views rather than copied. Together with the band, the benchmark of 51,400 fits takes 74 ms instead of 97 ms, with the same results. - `iterations` no longer counts a rejected step, an unreachable branch for a NaN objective is gone, `LeastSquaresResult` no longer generates a broken `__eq__`, and the inputs are annotated as `ArrayLike`. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- README.md | 3 +- docs/conf.py | 3 + docs/index.rst | 3 +- ndoptimize/__init__.py | 3 +- ndoptimize/_least_squares.py | 381 ++++++++++++++---------- ndoptimize/_model.py | 140 ++++++++- ndoptimize/_tests/test_least_squares.py | 105 ++++++- ndoptimize/_tests/test_models.py | 80 ++++- ndoptimize/models.py | 23 +- 9 files changed, 566 insertions(+), 175 deletions(-) diff --git a/README.md b/README.md index 9f27b97..7bb8ec1 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,7 @@ result.covariance # shape (..., 5, 5) ``` Models are short Numba functions which compute the model and its Jacobian, -compiled with the `ndoptimize.model` decorator. +compiled with the `ndoptimize.model` decorator, which also names their +parameters. See the [documentation](https://ndoptimize.readthedocs.io/en/latest) for details. diff --git a/docs/conf.py b/docs/conf.py index c5d7ac3..d3983c0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -39,6 +39,9 @@ ] autosummary_generate = True # Turn on sphinx.ext.autosummary autosummary_imported_members = True +# the stubs of `ndoptimize.Model` and `ndoptimize.model` would have the same +# name on a case-insensitive file system +autosummary_filename_map = {'ndoptimize.model': 'ndoptimize.model-decorator'} # autoclass_content = 'both' autodoc_typehints = "description" diff --git a/docs/index.rst b/docs/index.rst index 61a2f12..ede37d8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -17,7 +17,8 @@ The fits run in parallel on every CPU core, roughly half a million small fits per second on a 48-thread machine. Models are short Numba functions which compute the model and its Jacobian, -compiled with the :func:`ndoptimize.model` decorator. +compiled with the :func:`ndoptimize.model` decorator, which also names their +parameters. Some common ones are provided in :mod:`ndoptimize.models`. diff --git a/ndoptimize/__init__.py b/ndoptimize/__init__.py index 148cebc..b50c733 100644 --- a/ndoptimize/__init__.py +++ b/ndoptimize/__init__.py @@ -2,11 +2,12 @@ Numba-accelerated minimization of objective functions. """ -from ._model import model +from ._model import Model, model from . import models from ._least_squares import LeastSquaresResult, least_squares __all__ = [ + "Model", "model", "models", "LeastSquaresResult", diff --git a/ndoptimize/_least_squares.py b/ndoptimize/_least_squares.py index 7fec6a6..28e3605 100644 --- a/ndoptimize/_least_squares.py +++ b/ndoptimize/_least_squares.py @@ -4,9 +4,10 @@ import dataclasses import numpy as np +import numpy.typing as npt import numba from numba import types -from ._model import Model, signature_model +from ._model import ModelFunction, Model, signature_model __all__ = [ "LeastSquaresResult", @@ -23,61 +24,123 @@ """The largest value of the damping parameter, at which a fit stalls.""" _ridge = 1e-12 -"""A small number added to the diagonal of the normal matrix, so that it can -be factored even where a parameter has no effect on the model.""" +""" +The fraction of each diagonal element of the normal matrix which is added to +it for the convergence test, so that it can be factored where the parameters +are degenerate. +""" + +_degenerate = 1e-12 +""" +The smallest ratio of the precision of a parameter, given the parameters +before it, to its precision alone, below which the normal matrix is singular +and the covariance of the parameters is NaN. +""" @numba.njit(cache=True, error_model="numpy") -def _cholesky(a: np.ndarray, m: int, w: int) -> bool: +def _cholesky_band(band: np.ndarray, m: int) -> bool: """ - Overwrite the lower triangle of the leading ``m x m`` block of the - symmetric matrix `a`, of bandwidth `w`, with its Cholesky factor. - Only the band of the lower triangle is read. + Overwrite the first `m` rows of `band`, the lower band of a symmetric + matrix with ``band[i, d]`` its element ``(i, i - d)``, with the same band + of its Cholesky factor. - Returns :obj:`False` if `a` is not positive definite. + Returns :obj:`False` if the matrix is not positive definite. """ + w = band.shape[1] - 1 for j in range(m): - s = a[j, j] - for k in range(max(0, j - w), j): - s -= a[j, k] * a[j, k] + s = band[j, 0] + for d in range(1, min(j, w) + 1): + s -= band[j, d] * band[j, d] if not s > 0: return False - d = np.sqrt(s) - a[j, j] = d + pivot = np.sqrt(s) + band[j, 0] = pivot for i in range(j + 1, min(m, j + w + 1)): - s = a[i, j] + s = band[i, i - j] for k in range(max(0, i - w), j): - s -= a[i, k] * a[j, k] - a[i, j] = s / d + s -= band[i, i - k] * band[j, j - k] + band[i, i - j] = s / pivot return True @numba.njit(cache=True, error_model="numpy") -def _solve_lower(factor: np.ndarray, b: np.ndarray, m: int, w: int) -> None: +def _solve_lower_band(band: np.ndarray, b: np.ndarray, m: int) -> None: """ Overwrite the first `m` elements of `b` with the solution of - ``factor @ z = b`` for a lower-triangular `factor` of bandwidth `w`. + ``factor @ z = b``, for the lower-triangular `factor` stored as a band by + :func:`_cholesky_band`. """ + w = band.shape[1] - 1 for i in range(m): s = b[i] for k in range(max(0, i - w), i): - s -= factor[i, k] * b[k] - b[i] = s / factor[i, i] + s -= band[i, i - k] * b[k] + b[i] = s / band[i, 0] @numba.njit(cache=True, error_model="numpy") -def _solve_lower_columns(factor: np.ndarray, b: np.ndarray, m: int, w: int) -> None: +def _solve_lower_band_columns(band: np.ndarray, b: np.ndarray, m: int) -> None: """ Overwrite the first `m` rows of every column of `b` with the solution of - ``factor @ z = b`` for a lower-triangular `factor` of bandwidth `w`. + ``factor @ z = b``, for the lower-triangular `factor` stored as a band by + :func:`_cholesky_band`. """ + w = band.shape[1] - 1 for i in range(m): - d = factor[i, i] + pivot = band[i, 0] for c in range(b.shape[1]): s = b[i, c] for k in range(max(0, i - w), i): - s -= factor[i, k] * b[k, c] - b[i, c] = s / d + s -= band[i, i - k] * b[k, c] + b[i, c] = s / pivot + + +@numba.njit(cache=True, error_model="numpy") +def _cholesky(a: np.ndarray, degenerate: float) -> bool: + """ + Overwrite the lower triangle of the small symmetric matrix `a` with its + Cholesky factor. + + Returns :obj:`False` if `a` is not positive definite, or if a pivot is + no more than `degenerate` times the diagonal element it came from. + For a normal matrix, that ratio is the precision of a parameter, given + the parameters before it, over its precision alone. + """ + n = a.shape[0] + for j in range(n): + s = a[j, j] + for k in range(j): + s -= a[j, k] * a[j, k] + if not s > degenerate * a[j, j]: + return False + pivot = np.sqrt(s) + a[j, j] = pivot + for i in range(j + 1, n): + s = a[i, j] + for k in range(j): + s -= a[i, k] * a[j, k] + a[i, j] = s / pivot + return True + + +@numba.njit(cache=True, error_model="numpy") +def _solve_factored(factor: np.ndarray, b: np.ndarray, out: np.ndarray) -> None: + """ + Solve ``a @ x = b`` into `out`, given the Cholesky factor of `a` from + :func:`_cholesky`. + """ + n = factor.shape[0] + for i in range(n): + s = b[i] + for k in range(i): + s -= factor[i, k] * out[k] + out[i] = s / factor[i, i] + for i in range(n - 1, -1, -1): + s = out[i] + for k in range(i + 1, n): + s -= factor[k, i] * out[k] + out[i] = s / factor[i, i] @numba.njit(cache=True, error_model="numpy") @@ -93,17 +156,10 @@ def _solve_positive_definite( Returns :obj:`False` if `a` is not positive definite. """ - n = a.shape[0] work[:, :] = a - if not _cholesky(work, n, n): + if not _cholesky(work, 0.0): return False - out[:] = b - _solve_lower(work, out, n, n) - for i in range(n - 1, -1, -1): - s = out[i] - for k in range(i + 1, n): - s -= work[k, i] * out[k] - out[i] = s / work[i, i] + _solve_factored(work, b, out) return True @@ -111,7 +167,6 @@ def _solve_positive_definite( def _whitened_residual( factor: np.ndarray, m: int, - w: int, data_whitened: np.ndarray, f: np.ndarray, out: np.ndarray, @@ -122,7 +177,7 @@ def _whitened_residual( """ for j in range(m): out[j] = f[j] - _solve_lower(factor, out, m, w) + _solve_lower_band(factor, out, m) total = 0.0 for j in range(m): r = data_whitened[j] - out[j] @@ -135,7 +190,6 @@ def _whitened_residual( def _whitened_jacobian( factor: np.ndarray, m: int, - w: int, jacobian: np.ndarray, out: np.ndarray, ) -> None: @@ -143,7 +197,7 @@ def _whitened_jacobian( for j in range(m): for k in range(jacobian.shape[1]): out[j, k] = jacobian[j, k] - _solve_lower_columns(factor, out, m, w) + _solve_lower_band_columns(factor, out, m) @numba.njit(cache=True, error_model="numpy") @@ -167,14 +221,16 @@ def _normal( s = 0.0 for j in range(m): s += jacobian_whitened[j, k] * residual_whitened[j] - g[k] = s - prior_precision[k] * (parameters[k] - prior_mean[k]) + g[k] = s for q in range(k + 1): s = 0.0 for j in range(m): s += jacobian_whitened[j, k] * jacobian_whitened[j, q] a[k, q] = s a[q, k] = s - a[k, k] += prior_precision[k] + if prior_precision[k] > 0: + g[k] -= prior_precision[k] * (parameters[k] - prior_mean[k]) + a[k, k] += prior_precision[k] @numba.njit(cache=True, error_model="numpy") @@ -192,6 +248,16 @@ def _chi2_prior( return s +def _input(ndim: int) -> types.Array: + """ + The Numba type of an array input of :func:`_levenberg_marquardt`. + + It is read-only and of any layout, so that an argument broadcast over the + fits can be passed as a view rather than copied. + """ + return types.Array(types.float64, ndim, "A", readonly=True) + + _signature_kernel = types.Tuple( ( types.float64[:, ::1], @@ -203,15 +269,15 @@ def _chi2_prior( ) )( types.FunctionType(signature_model), - types.float64[:, :, ::1], - types.float64[:, ::1], - types.float64[:, ::1], - types.float64[::1], - types.float64[:, ::1], - types.float64[:, ::1], - types.float64[:, ::1], - types.float64[:, ::1], - types.float64[:, ::1], + _input(3), + _input(2), + _input(2), + _input(1), + _input(2), + _input(2), + _input(2), + _input(2), + _input(2), types.int64, types.float64, types.int64, @@ -227,7 +293,7 @@ def _chi2_prior( @numba.njit(_signature_kernel, parallel=True, cache=True, error_model="numpy") def _levenberg_marquardt( - model: Model, + model: ModelFunction, x: np.ndarray, data: np.ndarray, variance: np.ndarray, @@ -251,7 +317,7 @@ def _levenberg_marquardt( Parameters ---------- model - The model function, compiled with :func:`ndoptimize.model`. + The compiled model function, :attr:`ndoptimize.Model.function`. x The independent variable, ``(num_fits, num_points, num_inputs)``. data @@ -274,7 +340,7 @@ def _levenberg_marquardt( The standard deviation of the Gaussian prior of each parameter, ``(num_fits, num_parameters)``, infinite for no prior. max_iterations - The maximum number of steps per fit. + The maximum number of steps per fit, not negative. tolerance The Newton decrement below which a fit has converged. num_threads @@ -285,7 +351,6 @@ def _levenberg_marquardt( num_fit, num_point, num_input = x.shape num_parameter = guess.shape[1] num_lag = min(correlation.shape[0], num_point) - w = num_lag - 1 parameters = np.full((num_fit, num_parameter), np.nan) covariance = np.full((num_fit, num_parameter, num_parameter), np.nan) @@ -300,9 +365,9 @@ def _levenberg_marquardt( data_whitened = np.empty(num_point) sigma = np.empty(num_point) index = np.empty(num_point, dtype=np.int64) - factor = np.zeros((num_point, num_point)) + factor = np.empty((num_point, num_lag)) f = np.empty(num_point) - jac = np.empty((num_point, num_parameter)) + jac = np.zeros((num_point, num_parameter)) residual = np.empty(num_point) residual_new = np.empty(num_point) jacobian = np.empty((num_point, num_parameter)) @@ -342,37 +407,38 @@ def _levenberg_marquardt( # factor the banded noise covariance of the usable points and # whiten the data with it for r in range(m): - for s in range(max(0, r - w), r + 1): - lag = index[r] - index[s] + for d in range(min(r, num_lag - 1) + 1): + lag = index[r] - index[r - d] rho = correlation[lag] if lag < num_lag else 0.0 - factor[r, s] = sigma[r] * sigma[s] * rho - if not _cholesky(factor, m, w): + factor[r, d] = sigma[r] * sigma[r - d] * rho + if not _cholesky_band(factor, m): continue - _solve_lower(factor, data_whitened, m, w) + _solve_lower_band(factor, data_whitened, m) for k in range(num_parameter): precision[k] = 1 / prior_sigma[i, k] ** 2 - mean[k] = prior_mean[i, k] if np.isfinite(prior_mean[i, k]) else 0 + mean[k] = prior_mean[i, k] p[k] = min(max(guess[i, k], lower[i, k]), upper[i, k]) model(xg[:m], p, f[:m], jac[:m]) - chi2_data = _whitened_residual(factor, m, w, data_whitened, f, residual) - _whitened_jacobian(factor, m, w, jac, jacobian) + chi2_data = _whitened_residual(factor, m, data_whitened, f, residual) + _whitened_jacobian(factor, m, jac, jacobian) objective = chi2_data + _chi2_prior(p, mean, precision) damping = _damping_initial nu = 2.0 done = False n_step = 0 - for _ in range(max_iterations + 1): + while True: _normal(jacobian, residual, m, p, mean, precision, a, g) - # hold the parameters which are at a bound and are pushed - # outward by the gradient + # hold the parameters which have no effect on the model, and + # those which are at a bound and are pushed outward by the + # gradient for k in range(num_parameter): at_lower = p[k] <= lower[i, k] and g[k] < 0 at_upper = p[k] >= upper[i, k] and g[k] > 0 - free[k] = not (at_lower or at_upper) + free[k] = a[k, k] != 0 and not (at_lower or at_upper) # the Newton decrement of the free parameters, the decrease # in chi-squared a Gauss-Newton step would give if the model @@ -381,10 +447,7 @@ def _levenberg_marquardt( rhs[k] = g[k] if free[k] else 0 for q in range(num_parameter): system[k, q] = a[k, q] if (free[k] and free[q]) else 0 - if free[k]: - system[k, k] += _ridge - else: - system[k, k] = 1 + system[k, k] = a[k, k] * (1 + _ridge) if free[k] else 1 if _solve_positive_definite(system, rhs, work, step): decrement = 0.0 for k in range(num_parameter): @@ -403,19 +466,17 @@ def _levenberg_marquardt( while not accepted: for k in range(num_parameter): if free[k]: - system[k, k] = (a[k, k] + _ridge) * (1 + damping) + system[k, k] = a[k, k] * (1 + damping) if _solve_positive_definite(system, rhs, work, step): for k in range(num_parameter): p_new[k] = p[k] + step[k] p_new[k] = min(max(p_new[k], lower[i, k]), upper[i, k]) model(xg[:m], p_new, f[:m], jac[:m]) chi2_new = _whitened_residual( - factor, m, w, data_whitened, f, residual_new + factor, m, data_whitened, f, residual_new ) objective_new = chi2_new + _chi2_prior(p_new, mean, precision) accepted = objective_new < objective - if np.isnan(objective) and not np.isnan(objective_new): - accepted = True if accepted: # Nielsen's rule: the better the quadratic model # predicted the decrease, the less damping @@ -437,45 +498,35 @@ def _levenberg_marquardt( else: damping = min(damping * nu, _damping_max) nu *= 2 - n_step += 1 if not accepted: break + n_step += 1 p, p_new = p_new, p residual, residual_new = residual_new, residual chi2_data = chi2_new objective = objective_new - _whitened_jacobian(factor, m, w, jac, jacobian) + _whitened_jacobian(factor, m, jac, jacobian) # the covariance of the parameters is the inverse of the normal - # matrix, with infinite variance for a parameter which nothing - # constrains, and a small ridge if the rest is still singular + # matrix, with infinite variance for a parameter which has no + # effect on the model, and NaN if the rest of it is singular _normal(jacobian, residual, m, p, mean, precision, a, g) - diagonal_max = 0.0 for k in range(num_parameter): - diagonal_max = max(diagonal_max, a[k, k]) - for attempt in range(2): - for k in range(num_parameter): - for q in range(num_parameter): - constrained = a[k, k] > 0 and a[q, q] > 0 - system[k, q] = a[k, q] if constrained else 0 - if not a[k, k] > 0: - system[k, k] = 1 - elif attempt > 0: - system[k, k] += _ridge * diagonal_max - invertible = True + for q in range(num_parameter): + constrained = a[k, k] != 0 and a[q, q] != 0 + work[k, q] = a[k, q] if constrained else 0 + if a[k, k] == 0: + work[k, k] = 1 + if _cholesky(work, _degenerate): for k in range(num_parameter): for q in range(num_parameter): unit[q] = 1 if q == k else 0 - if not _solve_positive_definite(system, unit, work, step): - invertible = False - break + _solve_factored(work, unit, step) for q in range(num_parameter): covariance[i, q, k] = step[q] - if invertible: - break for k in range(num_parameter): - if not a[k, k] > 0: + if a[k, k] == 0: for q in range(num_parameter): covariance[i, k, q] = 0 covariance[i, q, k] = 0 @@ -490,7 +541,7 @@ def _levenberg_marquardt( return parameters, covariance, chi2, num_used, iterations, converged -@dataclasses.dataclass(frozen=True) +@dataclasses.dataclass(frozen=True, eq=False) class LeastSquaresResult: """ The results of :func:`least_squares` for every fit of a batch. @@ -508,8 +559,11 @@ class LeastSquaresResult: This is the inverse of the normal matrix, :math:`J^T C^{-1} J`, plus the precision of the priors, at the solution. - A parameter which neither the data nor its prior constrains has an + A parameter which has no effect on the model and no prior has an infinite variance. + If the rest of the normal matrix is singular, because the data and the + priors cannot constrain some combination of the parameters, as in a fit + with fewer usable points than parameters, the covariance is NaN. """ chi2: np.ndarray @@ -529,45 +583,53 @@ class LeastSquaresResult: def _flatten( + model: Model, x: np.ndarray, data: np.ndarray, - variance: float | np.ndarray, + variance: np.ndarray, guess: np.ndarray, - lower: float | np.ndarray, - upper: float | np.ndarray, - prior_mean: float | np.ndarray, - prior_sigma: float | np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + prior_mean: np.ndarray, + prior_sigma: np.ndarray, ) -> tuple[tuple[int, ...], dict[str, np.ndarray]]: """ - Broadcast the arguments of :func:`least_squares` against each other over - their batch axes and flatten those axes, into writable, contiguous - arrays of the types the kernel's signature requires. + Check the arguments of :func:`least_squares`, arrays of floats, broadcast + them against each other over their batch axes and flatten those axes, + into views of the arguments where possible. Returns the batch shape and the flattened arrays. """ - x = np.asarray(x, dtype=float) - data = np.asarray(data, dtype=float) - guess = np.asarray(guess, dtype=float) - if x.ndim < 2: raise ValueError( f"`x` must have at least two axes, (..., num_points, num_inputs), " f"got shape {x.shape}." ) num_point, num_input = x.shape[-2:] - num_parameter = guess.shape[-1] + num_parameter = len(model.parameters) if data.shape[-1:] != (num_point,): raise ValueError( f"The last axis of `data`, {data.shape}, must match the " f"number of points in `x`, {num_point}." ) + if guess.shape[-1:] != (num_parameter,): + raise ValueError( + f"The last axis of `guess`, {guess.shape}, must match the " + f"number of parameters of the model, {model.parameters}." + ) + if not np.all(lower <= upper): + raise ValueError("`lower` must not exceed `upper`, and neither may be NaN.") + if not np.all(prior_sigma > 0): + raise ValueError("`prior_sigma` must be positive.") + if np.any(np.isfinite(prior_sigma) & ~np.isfinite(prior_mean)): + raise ValueError("`prior_mean` must be finite where `prior_sigma` is finite.") variance = np.broadcast_to( - variance, np.broadcast_shapes(np.shape(variance), (num_point,)) + variance, np.broadcast_shapes(variance.shape, (num_point,)) ) lower, upper, prior_mean, prior_sigma = ( - np.broadcast_to(a, np.broadcast_shapes(np.shape(a), (num_parameter,))) + np.broadcast_to(a, np.broadcast_shapes(a.shape, (num_parameter,))) for a in (lower, upper, prior_mean, prior_sigma) ) @@ -584,9 +646,12 @@ def _flatten( num_fit = int(np.prod(shape)) def flat(a: np.ndarray, *shape_trailing: int) -> np.ndarray: - """A writable, contiguous copy of `a` with its batch axes flattened.""" + """ + `a` broadcast to the batch shape with its batch axes flattened, a + view of `a` unless the flattening needs a copy. + """ a = np.broadcast_to(a, shape + shape_trailing) - return np.array(a, dtype=float, order="C").reshape(num_fit, *shape_trailing) + return a.reshape(num_fit, *shape_trailing) arrays = dict( x=flat(x, num_point, num_input), @@ -604,15 +669,15 @@ def flat(a: np.ndarray, *shape_trailing: int) -> np.ndarray: def least_squares( model: Model, - x: np.ndarray, - data: np.ndarray, - variance: float | np.ndarray, - guess: np.ndarray, - lower: float | np.ndarray = -np.inf, - upper: float | np.ndarray = np.inf, - prior_mean: float | np.ndarray = 0, - prior_sigma: float | np.ndarray = np.inf, - correlation: None | np.ndarray = None, + x: npt.ArrayLike, + data: npt.ArrayLike, + variance: npt.ArrayLike, + guess: npt.ArrayLike, + lower: npt.ArrayLike = -np.inf, + upper: npt.ArrayLike = np.inf, + prior_mean: npt.ArrayLike = 0, + prior_sigma: npt.ArrayLike = np.inf, + correlation: None | npt.ArrayLike = None, max_iterations: int = 1000, tolerance: float = 1e-8, ) -> LeastSquaresResult: @@ -652,7 +717,7 @@ def least_squares( Parameters ---------- model - The model to fit, a function compiled with :func:`ndoptimize.model`. + The model to fit, made with :func:`ndoptimize.model`. x The independent variable, an array of shape ``(..., num_points, num_inputs)``. @@ -666,25 +731,29 @@ def least_squares( their fit. guess The initial guess for the parameters, an array of shape - ``(..., num_parameters)``. + ``(..., num_parameters)``, where ``num_parameters`` is the number of + names in :attr:`ndoptimize.Model.parameters`. + It is clipped to the bounds. lower The lower bound of each parameter, an array of shape ``(..., num_parameters)``. upper The upper bound of each parameter, an array of shape - ``(..., num_parameters)``. + ``(..., num_parameters)``, no less than `lower`. prior_mean The mean of the Gaussian prior of each parameter, an array of shape ``(..., num_parameters)``. + It is ignored where `prior_sigma` is infinite. prior_sigma - The standard deviation of the Gaussian prior of each parameter, an - array of shape ``(..., num_parameters)``. + The standard deviation of the Gaussian prior of each parameter, a + positive array of shape ``(..., num_parameters)``. An infinite standard deviation, the default, means no prior. correlation The correlation of the noise of two data points by their separation along the last axis of `data`, in samples, starting with 1 for no separation and zero beyond its length. The separation counts the points left out of a fit. + It must describe a positive-definite noise covariance. If :obj:`None` (the default), the noise is uncorrelated. max_iterations The maximum number of steps for each fit. @@ -718,7 +787,10 @@ def least_squares( Steps are clipped to the bounds, and a parameter at a bound which the gradient pushes outward is held fixed for the step and for the - convergence test. + convergence test, as is a parameter which has no effect on the model. + + A fit whose :math:`\chi^2` is NaN at the initial guess, where the model + is NaN for instance, takes no steps and does not converge. Examples -------- @@ -773,46 +845,51 @@ def least_squares( ax.set_xlabel("(fitted center - true center) / uncertainty"); ax.legend(); """ + if max_iterations < 0: + raise ValueError( + f"`max_iterations` must not be negative, got {max_iterations}." + ) + shape, a = _flatten( - x=x, - data=data, - variance=variance, - guess=guess, - lower=lower, - upper=upper, - prior_mean=prior_mean, - prior_sigma=prior_sigma, + model=model, + x=np.asarray(x, dtype=float), + data=np.asarray(data, dtype=float), + variance=np.asarray(variance, dtype=float), + guess=np.asarray(guess, dtype=float), + lower=np.asarray(lower, dtype=float), + upper=np.asarray(upper, dtype=float), + prior_mean=np.asarray(prior_mean, dtype=float), + prior_sigma=np.asarray(prior_sigma, dtype=float), ) if correlation is None: - correlation = np.ones(1) - correlation = np.array(correlation, dtype=float).reshape(-1) - if correlation.size == 0 or correlation[0] != 1: + correlation = 1 + rho = np.array(correlation, dtype=float).reshape(-1) + if rho.size == 0 or rho[0] != 1: raise ValueError( f"`correlation` must start with 1, the correlation of a point " - f"with itself, got {correlation}." + f"with itself, got {rho}." ) # the noise correlation of every fit is a principal submatrix of this # one, so it is positive definite if this one is num_point = a["x"].shape[1] - rho = np.zeros(num_point) - rho[: correlation.size] = correlation[:num_point] - index = np.arange(num_point) - try: - np.linalg.cholesky(rho[np.abs(index[:, np.newaxis] - index)]) - except np.linalg.LinAlgError: + num_lag = min(rho.size, num_point) + band = np.zeros((num_point, num_lag)) + for d in range(num_lag): + band[d:, d] = rho[d] + if not np.all(np.isfinite(rho)) or not _cholesky_band(band, num_point): raise ValueError( - f"`correlation`, {correlation}, does not describe a positive " + f"`correlation`, {rho}, does not describe a positive " f"definite noise covariance for {num_point} points." ) result = _levenberg_marquardt( - model, + model.function, a["x"], a["data"], a["variance"], - correlation, + rho, a["guess"], a["lower"], a["upper"], diff --git a/ndoptimize/_model.py b/ndoptimize/_model.py index c621ab5..cf2b6b3 100644 --- a/ndoptimize/_model.py +++ b/ndoptimize/_model.py @@ -2,20 +2,23 @@ The interface between a model function and the fitting kernels. """ -from typing import Callable +from typing import Callable, Sequence +import dataclasses +import functools import numpy as np import numba from numba import types __all__ = [ - "Model", + "ModelFunction", "signature_model", + "Model", "model", ] -Model = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None] +ModelFunction = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None] """ -The type of a model function, ``model(x, parameters, f, jacobian)``. +The type of a model function, ``function(x, parameters, f, jacobian)``. """ signature_model = types.void( @@ -29,12 +32,96 @@ """ -def model(function: Model) -> Model: +@dataclasses.dataclass(eq=False) +class Model: + """ + A model which :func:`ndoptimize.least_squares` can fit: a compiled model + function and the names of its parameters, made with :func:`model`. + + Calling a model evaluates it, after checking the shapes of the arrays, + since the compiled function does not check its indices. + """ + + function: ModelFunction + """ + The compiled model function, ``function(x, parameters, f, jacobian)``, + described in :func:`model`. + """ + + parameters: tuple[str, ...] + """ + The names of the parameters, in the order of the ``parameters`` array + and of the columns of the ``jacobian`` array. + """ + + def __post_init__(self) -> None: + functools.update_wrapper(self, self.function, updated=()) + + def __repr__(self) -> str: + return f"{type(self).__name__}(parameters={self.parameters})" + + def __get__(self, instance: object, owner: None | type = None) -> "Model": + """ + Return this model unchanged when it is an attribute of a class, like + a static method, rather than binding it to an instance. + + This also makes it a method descriptor, which documentation tools + treat as a function. + """ + return self + + def __call__( + self, + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, + ) -> None: + """ + Evaluate the model and its Jacobian. + + Parameters + ---------- + x + The points at which to evaluate the model, + ``(num_points, num_inputs)``. + parameters + The parameters of the model, ``(num_parameters,)``. + 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, num_parameters)``. + """ + num_point = np.shape(x)[0] if np.ndim(x) == 2 else -1 + num_parameter = len(self.parameters) + if ( + np.ndim(x) != 2 + or np.shape(parameters) != (num_parameter,) + or np.shape(f) != (num_point,) + or np.shape(jacobian) != (num_point, num_parameter) + ): + raise ValueError( + f"For a model with {num_parameter} parameters, `x` must have " + f"shape (num_points, num_inputs), `parameters` " + f"({num_parameter},), `f` (num_points,) and `jacobian` " + f"(num_points, {num_parameter}), got {np.shape(x)}, " + f"{np.shape(parameters)}, {np.shape(f)} and " + f"{np.shape(jacobian)}." + ) + self.function(x, parameters, f, jacobian) + + +def model( + parameters: Sequence[str], + cache: bool = True, +) -> Callable[[ModelFunction], Model]: r""" - Compile a function as a model which :func:`ndoptimize.least_squares` can fit. + Compile a function as a model which :func:`ndoptimize.least_squares` can + fit, with the given names for its parameters. - A model is a function ``model(x, parameters, f, jacobian)`` which writes - the model evaluated at the points ``x``, an array of shape + A model function, ``function(x, parameters, f, jacobian)``, writes the + model evaluated at the points ``x``, an array of shape ``(num_points, num_inputs)``, for the given ``parameters``, an array of shape ``(num_parameters,)``, into ``f``, an array of shape ``(num_points,)``, and its derivatives with respect to the parameters @@ -43,7 +130,7 @@ def model(function: Model) -> Model: The function is compiled with Numba in `nopython` mode, with this fixed signature. - The fitting kernel takes the model as a typed function pointer, so a + The fitting kernel takes the function as a typed function pointer, so a single compiled kernel serves every model and is cached on disk. Scalar functions such as :func:`math.exp` should come from :mod:`math`, @@ -51,8 +138,20 @@ def model(function: Model) -> Model: Parameters ---------- - function - The model function to compile. + parameters + The names of the parameters, in the order of the ``parameters`` + array and of the columns of the ``jacobian`` array. + :func:`ndoptimize.least_squares` checks the initial guess against + their number, since the compiled function does not check its + indices. + cache + Whether to cache the compiled function on disk, so that later + processes need not compile it again. + Numba compiles the global arrays a function reads into it as + constants, and the cache does not notice when they change, so a model + which reads globals computed at run time should not be cached. + Neither can a model which is not defined in a file, such as one + defined in the interactive interpreter. Examples -------- @@ -63,12 +162,27 @@ def model(function: Model) -> Model: import ndoptimize - @ndoptimize.model + @ndoptimize.model(parameters=("intercept", "slope")) def line(x, parameters, f, jacobian): a, b = parameters[0], parameters[1] for i in range(x.shape[0]): f[i] = a + b * x[i, 0] jacobian[i, 0] = 1 jacobian[i, 1] = x[i, 0] + + line """ - return numba.njit(signature_model, cache=True, error_model="numpy")(function) + if isinstance(parameters, str): + raise TypeError( + f"`parameters` must be a sequence of names, got the string " + f"{parameters!r}." + ) + parameters = tuple(parameters) + if len(set(parameters)) != len(parameters): + raise ValueError(f"The names in `parameters`, {parameters}, must be unique.") + + def decorator(function: ModelFunction) -> Model: + compiled = numba.njit(signature_model, cache=cache, error_model="numpy") + return Model(function=compiled(function), parameters=parameters) + + return decorator diff --git a/ndoptimize/_tests/test_least_squares.py b/ndoptimize/_tests/test_least_squares.py index 8d24700..6f45de6 100644 --- a/ndoptimize/_tests/test_least_squares.py +++ b/ndoptimize/_tests/test_least_squares.py @@ -22,7 +22,7 @@ def _noise_covariance(sigma: np.ndarray, correlation: None | np.ndarray) -> np.n return sigma[:, np.newaxis] * rho[lag] * sigma -def _evaluate(model: ndoptimize._model.Model, parameters: np.ndarray) -> np.ndarray: +def _evaluate(model: ndoptimize.Model, parameters: np.ndarray) -> np.ndarray: """The model at the test velocities for each row of `parameters`.""" x = _velocity[:, np.newaxis] result = np.empty((len(parameters), _velocity.size)) @@ -232,6 +232,13 @@ def test_least_squares_broadcasting() -> None: assert getattr(result, name).shape == (2, 3) assert np.array_equal(result.parameters.reshape(6, 5), result_flat.parameters) + # an argument broadcast over only some of the batch axes + kwargs["guess"] = guess[:3] + result_partial = ndoptimize.least_squares(data=data.reshape(2, 3, -1), **kwargs) + kwargs["guess"] = np.broadcast_to(guess[:3], (2, 3, 5)).copy() + result_full = ndoptimize.least_squares(data=data.reshape(2, 3, -1), **kwargs) + assert np.array_equal(result_partial.parameters, result_full.parameters) + def test_least_squares_missing() -> None: truth, data, variance, guess = _spectra(3, None, missing=0) @@ -277,7 +284,7 @@ def test_least_squares_unconstrained() -> None: assert np.isfinite(result.covariance[0, 0]) -@ndoptimize.model +@ndoptimize.model(parameters=("intercept", "slope")) def _line( x: np.ndarray, parameters: np.ndarray, @@ -317,6 +324,91 @@ def test_least_squares_linear() -> None: assert result.iterations <= 3 +@ndoptimize.model(parameters=("scale",)) +def _scaled( + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, +) -> None: + """A multiple of the input.""" + for i in range(x.shape[0]): + f[i] = parameters[0] * x[i, 0] + jacobian[i, 0] = x[i, 0] + + +def test_least_squares_scale() -> None: + """A parameter of a very different scale than the data, like an emission + measure times a temperature response, converges to the minimum.""" + rng = np.random.default_rng(3) + response = 1e-27 * np.exp(-0.5 * np.linspace(-3, 3, 39) ** 2) + variance = 0.01 + data = 3e27 * response + rng.normal(0, np.sqrt(variance), response.size) + + result = ndoptimize.least_squares( + model=_scaled, + x=response[:, np.newaxis], + data=data, + variance=variance, + guess=[2e27], + ) + + information = response @ response / variance + expected = response @ data / variance / information + sigma = 1 / np.sqrt(information) + assert result.converged + assert np.abs(result.parameters[0] - expected) < 1e-4 * sigma + assert np.allclose(result.covariance, 1 / information, rtol=1e-10) + + +def test_least_squares_degenerate() -> None: + """A fit with fewer usable points than parameters has a NaN covariance.""" + result = ndoptimize.least_squares( + model=_line, + x=np.array([[1.0], [2.0]]), + data=np.array([3.0, np.nan]), + variance=1, + guess=np.zeros(2), + ) + assert result.num_points == 1 + assert np.allclose(result.parameters @ [1, 1], 3) + assert np.all(np.isnan(result.covariance)) + + +def test_least_squares_nan() -> None: + """A fit whose model is NaN at the initial guess takes no steps.""" + truth, data, variance, guess = _spectra(1, None) + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + data=data, + variance=variance, + guess=np.array([np.nan, 0, 15, 0, 1]), + ) + assert not np.any(result.converged) + assert np.all(result.iterations == 0) + + +def test_least_squares_prior_mean_ignored() -> None: + """The prior mean of a parameter without a prior is ignored, even if it is + NaN.""" + truth, data, variance, guess = _spectra(3, None) + kwargs: dict[str, Any] = dict( + model=ndoptimize.models.gaussian_linear, + x=_velocity[:, np.newaxis], + data=data, + variance=variance, + guess=guess, + prior_sigma=np.array([np.inf, 10, np.inf, np.inf, np.inf]), + ) + result = ndoptimize.least_squares(**kwargs) + result_nan = ndoptimize.least_squares( + prior_mean=np.array([np.nan, 0, np.nan, np.inf, np.nan]), + **kwargs, + ) + assert np.array_equal(result.parameters, result_nan.parameters) + + def test_least_squares_far_guess() -> None: """Fits started far from the solution, which take damped and rejected steps, reach the same minimum as fits started near it.""" @@ -360,8 +452,17 @@ def test_least_squares_max_iterations() -> None: argvalues=[ dict(correlation=np.array([0.5, 0.1])), dict(correlation=np.array([1, 0.9])), + dict(correlation=np.array([1, np.nan])), dict(data=np.zeros(5)), dict(x=_velocity), + dict(guess=np.array([1, 0.0])), + dict(guess=1.0), + dict(lower=np.array([0, 5, 0]), upper=np.array([10, 2, 30])), + dict(lower=np.nan), + dict(prior_sigma=0), + dict(prior_sigma=np.nan), + dict(prior_mean=np.nan, prior_sigma=1), + dict(max_iterations=-1), ], ) def test_least_squares_invalid(kwargs: dict[str, Any]) -> None: diff --git a/ndoptimize/_tests/test_models.py b/ndoptimize/_tests/test_models.py index e0aa024..23bfed9 100644 --- a/ndoptimize/_tests/test_models.py +++ b/ndoptimize/_tests/test_models.py @@ -30,7 +30,7 @@ def _gaussian(v: np.ndarray, a: float, c: float, w: float) -> np.ndarray: ], ) def test_models( - model: ndoptimize._model.Model, + model: ndoptimize.Model, parameters: np.ndarray, function: Callable[[np.ndarray, np.ndarray], np.ndarray], ) -> None: @@ -39,6 +39,8 @@ def test_models( f = np.empty(v.size) jacobian = np.empty((v.size, parameters.size)) + assert len(model.parameters) == parameters.size + model(x, parameters, f, jacobian) assert np.allclose(f, function(v, parameters)) @@ -51,3 +53,79 @@ def test_models( 2 * step ) assert np.allclose(jacobian[:, k], derivative, rtol=1e-5, atol=1e-8) + + +@pytest.mark.parametrize( + argnames="shape_parameters,shape_f,shape_jacobian", + argvalues=[ + ((2,), (21,), (21, 2)), + ((3,), (20,), (21, 3)), + ((3,), (21,), (21, 5)), + ], +) +def test_model_shapes( + shape_parameters: tuple[int, ...], + shape_f: tuple[int, ...], + shape_jacobian: tuple[int, ...], +) -> None: + """Calling a model with arrays of the wrong shapes raises an error rather + than reading or writing outside them.""" + with pytest.raises(ValueError): + ndoptimize.models.gaussian( + np.zeros((21, 1)), + np.ones(shape_parameters), + np.empty(shape_f), + np.empty(shape_jacobian), + ) + + +def test_model_parameters() -> None: + def constant( + x: np.ndarray, + parameters: np.ndarray, + f: np.ndarray, + jacobian: np.ndarray, + ) -> None: + """A constant.""" + for i in range(x.shape[0]): + f[i] = parameters[0] + jacobian[i, 0] = 1 + + with pytest.raises(TypeError): + ndoptimize.model(parameters="offset") + with pytest.raises(ValueError): + ndoptimize.model(parameters=("offset", "offset")) + + model = ndoptimize.model(parameters=["offset"], cache=False)(constant) + assert model.parameters == ("offset",) + assert model.__doc__ == "A constant." + assert repr(model) == "Model(parameters=('offset',))" + + f = np.empty(3) + jacobian = np.empty((3, 1)) + model(np.zeros((3, 1)), np.array([2.0]), f, jacobian) + assert np.all(f == 2) + assert np.all(jacobian == 1) + + +def test_model_uncached() -> None: + """A model which is not defined in a file can be compiled without the + cache.""" + source = ( + "@ndoptimize.model(parameters=('slope',), cache=False)\n" + "def proportional(x, parameters, f, jacobian):\n" + " for i in range(x.shape[0]):\n" + " f[i] = parameters[0] * x[i, 0]\n" + " jacobian[i, 0] = x[i, 0]\n" + ) + namespace: dict = dict(ndoptimize=ndoptimize) + exec(source, namespace) + model = namespace["proportional"] + + x = np.arange(4.0)[:, np.newaxis] + f = np.empty(4) + jacobian = np.empty((4, 1)) + model(x, np.array([2.0]), f, jacobian) + + assert np.array_equal(f, 2 * x[:, 0]) + assert np.array_equal(jacobian, x) diff --git a/ndoptimize/models.py b/ndoptimize/models.py index 43ab7c2..817f22d 100644 --- a/ndoptimize/models.py +++ b/ndoptimize/models.py @@ -1,5 +1,5 @@ """ -Common model functions, compiled with :func:`ndoptimize.model`. +Common models, made with :func:`ndoptimize.model`. """ import math @@ -13,7 +13,9 @@ ] -@_model.model +@_model.model( + parameters=("amplitude", "center", "width"), +) def gaussian( x: np.ndarray, parameters: np.ndarray, @@ -52,7 +54,9 @@ def gaussian( jacobian[i, 2] = amplitude * e * z * z / width -@_model.model +@_model.model( + parameters=("amplitude", "center", "width", "slope", "intercept"), +) def gaussian_linear( x: np.ndarray, parameters: np.ndarray, @@ -96,7 +100,18 @@ def gaussian_linear( jacobian[i, 4] = 1 -@_model.model +@_model.model( + parameters=( + "amplitude", + "center", + "width", + "slope", + "intercept", + "amplitude_2", + "center_2", + "width_2", + ), +) def two_gaussians_linear( x: np.ndarray, parameters: np.ndarray, From efeb3711a831c8bfc95887d96a3cc46294b39947 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Mon, 19 Oct 2026 11:37:15 -0600 Subject: [PATCH 4/5] Add `ndoptimize.evaluate()`, which evaluates a model for many sets of parameters at once. It runs in parallel with Numba, and broadcasts the batch axes of the points and the parameters like `least_squares()` does, so the best-fit model of every fit is one call. The named-arrays wrapper of `least_squares()` needs it to give back the fitted model. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- README.md | 7 ++ docs/index.rst | 4 +- ndoptimize/__init__.py | 2 + ndoptimize/_evaluate.py | 168 +++++++++++++++++++++++++++++ ndoptimize/_tests/test_evaluate.py | 100 +++++++++++++++++ 5 files changed, 280 insertions(+), 1 deletion(-) create mode 100644 ndoptimize/_evaluate.py create mode 100644 ndoptimize/_tests/test_evaluate.py diff --git a/README.md b/README.md index 7bb8ec1..9f06046 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,13 @@ result = ndoptimize.least_squares( result.parameters # shape (..., 5) result.covariance # shape (..., 5, 5) + +# the best-fit model at each point of each spectrum, shape (..., 41) +fit = ndoptimize.evaluate( + model=ndoptimize.models.gaussian_linear, + x=velocity[:, np.newaxis], + parameters=result.parameters, +) ``` Models are short Numba functions which compute the model and its Jacobian, diff --git a/docs/index.rst b/docs/index.rst index ede37d8..524fe86 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -19,7 +19,9 @@ fits per second on a 48-thread machine. Models are short Numba functions which compute the model and its Jacobian, compiled with the :func:`ndoptimize.model` decorator, which also names their parameters. -Some common ones are provided in :mod:`ndoptimize.models`. +Some common ones are provided in :mod:`ndoptimize.models`, and +:func:`ndoptimize.evaluate` evaluates a model for many sets of parameters at +once, such as the best-fit ones. Installation diff --git a/ndoptimize/__init__.py b/ndoptimize/__init__.py index b50c733..f8ebabf 100644 --- a/ndoptimize/__init__.py +++ b/ndoptimize/__init__.py @@ -5,6 +5,7 @@ from ._model import Model, model from . import models from ._least_squares import LeastSquaresResult, least_squares +from ._evaluate import evaluate __all__ = [ "Model", @@ -12,4 +13,5 @@ "models", "LeastSquaresResult", "least_squares", + "evaluate", ] diff --git a/ndoptimize/_evaluate.py b/ndoptimize/_evaluate.py new file mode 100644 index 0000000..c3b130b --- /dev/null +++ b/ndoptimize/_evaluate.py @@ -0,0 +1,168 @@ +""" +Evaluation of a model for many sets of parameters at once. +""" + +import math +import numpy as np +import numpy.typing as npt +import numba +from numba import types +from ._model import ModelFunction, Model, signature_model +from ._least_squares import _input + +__all__ = [ + "evaluate", +] + +_signature_kernel = types.float64[:, ::1]( + types.FunctionType(signature_model), + _input(3), + _input(2), + types.int64, +) +"""The Numba signature of :func:`_evaluate`.""" + + +@numba.njit(_signature_kernel, parallel=True, cache=True, error_model="numpy") +def _evaluate( + model: ModelFunction, + x: np.ndarray, + parameters: np.ndarray, + num_threads: int, +) -> np.ndarray: + """ + Evaluate a model for each row of `parameters`. + + Parameters + ---------- + model + The compiled model function, :attr:`ndoptimize.Model.function`. + x + The independent variable, ``(num_sets, num_points, num_inputs)``. + parameters + The parameters, ``(num_sets, num_parameters)``. + num_threads + The number of chunks to split the sets of parameters into, the result + of :func:`numba.get_num_threads`, which disables caching if it is + called inside a compiled function. + """ + num_set, num_point, num_input = x.shape + num_parameter = parameters.shape[1] + + f = np.empty((num_set, num_point)) + + num_chunk = max(1, min(num_threads, num_set)) + for chunk in numba.prange(num_chunk): # type: ignore[attr-defined] + xi = np.empty((num_point, num_input)) + p = np.empty(num_parameter) + jacobian = np.empty((num_point, num_parameter)) + + for i in range(chunk, num_set, num_chunk): + for j in range(num_point): + for q in range(num_input): + xi[j, q] = x[i, j, q] + for k in range(num_parameter): + p[k] = parameters[i, k] + model(xi, p, f[i], jacobian) + + return f + + +def evaluate( + model: Model, + x: npt.ArrayLike, + parameters: npt.ArrayLike, +) -> np.ndarray: + """ + Evaluate a model for many sets of parameters at once, such as the results + of :func:`ndoptimize.least_squares`. + + The sets run in parallel. + + Parameters + ---------- + model + The model to evaluate, made with :func:`ndoptimize.model`. + x + The points at which to evaluate the model, an array of shape + ``(..., num_points, num_inputs)``. + parameters + The parameters of the model, an array of shape + ``(..., num_parameters)``, where ``num_parameters`` is the number of + names in :attr:`ndoptimize.Model.parameters`. + + Returns + ------- + The model evaluated at `x`, an array of shape ``(..., num_points)``, + where the batch axes of `x` and `parameters`, all but their last one (two + for `x`), are broadcast against each other. + + Examples + -------- + + Evaluate Gaussians of three different widths. + + .. jupyter-execute:: + + import numpy as np + import matplotlib.pyplot as plt + import ndoptimize + + x = np.linspace(-50, 50, 101) + parameters = np.array([ + [1, 0, 5], + [1, 0, 10], + [1, 0, 20], + ]) + + f = ndoptimize.evaluate( + model=ndoptimize.models.gaussian, + x=x[:, np.newaxis], + parameters=parameters, + ) + + fig, ax = plt.subplots(constrained_layout=True) + ax.plot(x, f.T, label=[f"width {w}" for w in parameters[:, 2]]); + ax.legend(); + """ + return _evaluate_arrays( + model=model, + x=np.asarray(x, dtype=float), + parameters=np.asarray(parameters, dtype=float), + ) + + +def _evaluate_arrays( + model: Model, + x: np.ndarray, + parameters: np.ndarray, +) -> np.ndarray: + """:func:`evaluate` for arguments which are already arrays of floats.""" + if x.ndim < 2: + raise ValueError( + f"`x` must have at least two axes, (..., num_points, num_inputs), " + f"got shape {x.shape}." + ) + num_point, num_input = x.shape[-2:] + num_parameter = len(model.parameters) + + if parameters.shape[-1:] != (num_parameter,): + raise ValueError( + f"The last axis of `parameters`, {parameters.shape}, must match " + f"the number of parameters of the model, {model.parameters}." + ) + + shape = np.broadcast_shapes(x.shape[:-2], parameters.shape[:-1]) + num_set = math.prod(shape) + + x = np.broadcast_to(x, shape + (num_point, num_input)) + parameters = np.broadcast_to(parameters, shape + (num_parameter,)) + + f = _evaluate( + model.function, + x.reshape(num_set, num_point, num_input), + parameters.reshape(num_set, num_parameter), + numba.get_num_threads(), + ) + + return f.reshape(shape + (num_point,)) diff --git a/ndoptimize/_tests/test_evaluate.py b/ndoptimize/_tests/test_evaluate.py new file mode 100644 index 0000000..9ab9a11 --- /dev/null +++ b/ndoptimize/_tests/test_evaluate.py @@ -0,0 +1,100 @@ +from typing import Any +import numpy as np +import pytest +import ndoptimize + + +def _evaluate( + model: ndoptimize.Model, + x: np.ndarray, + parameters: np.ndarray, +) -> np.ndarray: + """Evaluate a model for each row of `parameters`, one at a time.""" + f = np.empty((len(parameters), len(x))) + jacobian = np.empty((len(x), parameters.shape[-1])) + for i in range(len(parameters)): + model(x, parameters[i], f[i], jacobian) + return f + + +def test_evaluate() -> None: + """Evaluating many sets of parameters at once matches evaluating them + one at a time, with the points shared or not.""" + rng = np.random.default_rng(0) + x = np.linspace(-50, 50, 21)[:, np.newaxis] + parameters = np.stack( + [ + rng.uniform(1, 2, 10), + rng.normal(0, 5, 10), + rng.uniform(5, 20, 10), + rng.normal(0, 0.01, 10), + rng.uniform(0, 1, 10), + ], + axis=-1, + ) + model = ndoptimize.models.gaussian_linear + expected = _evaluate(model, x, parameters) + + result = ndoptimize.evaluate(model=model, x=x, parameters=parameters) + assert np.array_equal(result, expected) + + result_x = ndoptimize.evaluate( + model=model, + x=np.broadcast_to(x, (10,) + x.shape), + parameters=parameters[0], + ) + assert np.array_equal(result_x, np.broadcast_to(expected[0], (10, len(x)))) + + +def test_evaluate_broadcasting() -> None: + """The batch axes of the points and the parameters broadcast.""" + x = np.linspace(-50, 50, 21)[:, np.newaxis] + x = np.stack([x, 2 * x])[:, np.newaxis] + parameters = np.array([[1, 0, 10], [2, 5, 15], [3, -5, 20.0]]) + + result = ndoptimize.evaluate( + model=ndoptimize.models.gaussian, + x=x, + parameters=parameters, + ) + + assert result.shape == (2, 3, 21) + for i in range(2): + assert np.array_equal( + result[i], + _evaluate(ndoptimize.models.gaussian, x[i, 0], parameters), + ) + + +def test_evaluate_least_squares() -> None: + """The model evaluated at the best-fit parameters of noiseless data is + the data.""" + x = np.linspace(-50, 50, 21)[:, np.newaxis] + data = ndoptimize.evaluate(ndoptimize.models.gaussian, x, [2, 3, 12.0]) + result = ndoptimize.least_squares( + model=ndoptimize.models.gaussian, + x=x, + data=data, + variance=1, + guess=[1, 0, 10], + ) + f = ndoptimize.evaluate(ndoptimize.models.gaussian, x, result.parameters) + assert np.allclose(f, data, atol=1e-6) + + +@pytest.mark.parametrize( + argnames="kwargs", + argvalues=[ + dict(x=np.linspace(-50, 50, 21)), + dict(parameters=np.array([1, 0.0])), + dict(parameters=1.0), + ], +) +def test_evaluate_invalid(kwargs: dict[str, Any]) -> None: + arguments: dict[str, Any] = dict( + model=ndoptimize.models.gaussian, + x=np.linspace(-50, 50, 21)[:, np.newaxis], + parameters=np.array([1, 0, 15.0]), + ) + with pytest.raises(ValueError): + ndoptimize.evaluate(**(arguments | kwargs)) From ed0651f4171b76b19f8c5155a7d25e1b13f4ab11 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Mon, 19 Oct 2026 13:16:36 -0600 Subject: [PATCH 5/5] Copy and pickle models by reference, like functions. A copy of a model is the model itself, and a pickled model unpickles as the model of the same name in its module. Before, a deep copy or a pickle round trip made a new `Model`, which isn't equal to the original since models compare by identity. This broke `copy()` and pickling of named-arrays function arrays which keep a model as a field. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- ndoptimize/_model.py | 8 ++++++++ ndoptimize/_tests/test_models.py | 10 ++++++++++ 2 files changed, 18 insertions(+) diff --git a/ndoptimize/_model.py b/ndoptimize/_model.py index cf2b6b3..5403640 100644 --- a/ndoptimize/_model.py +++ b/ndoptimize/_model.py @@ -70,6 +70,14 @@ def __get__(self, instance: object, owner: None | type = None) -> "Model": """ return self + def __reduce__(self) -> str: + """ + Pickle and copy this model by reference, like a function: a copy is + this model itself, and a pickled model is the one of the same name in + its module when it is unpickled. + """ + return self.function.__qualname__ + def __call__( self, x: np.ndarray, diff --git a/ndoptimize/_tests/test_models.py b/ndoptimize/_tests/test_models.py index 23bfed9..6cbbd7c 100644 --- a/ndoptimize/_tests/test_models.py +++ b/ndoptimize/_tests/test_models.py @@ -1,4 +1,6 @@ from typing import Callable +import copy +import pickle import numpy as np import pytest import ndoptimize @@ -129,3 +131,11 @@ def test_model_uncached() -> None: assert np.array_equal(f, 2 * x[:, 0]) assert np.array_equal(jacobian, x) + + +def test_model_reference() -> None: + """A model is copied and pickled by reference, like a function.""" + model = ndoptimize.models.gaussian + assert copy.copy(model) is model + assert copy.deepcopy(model) is model + assert pickle.loads(pickle.dumps(model)) is model