diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..da2e352 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,52 @@ + +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 + - 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/.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..9f06046 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,73 @@ # 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) [![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) + +# 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, +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 2eda6fd..d3983c0 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 --------------------------------------------------- @@ -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 2b7a971..524fe86 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,11 +1,28 @@ 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, which also names their +parameters. +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 b36e432..f8ebabf 100644 --- a/ndoptimize/__init__.py +++ b/ndoptimize/__init__.py @@ -1,3 +1,17 @@ """ Numba-accelerated minimization of objective functions. """ + +from ._model import Model, model +from . import models +from ._least_squares import LeastSquaresResult, least_squares +from ._evaluate import evaluate + +__all__ = [ + "Model", + "model", + "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/_least_squares.py b/ndoptimize/_least_squares.py new file mode 100644 index 0000000..28e3605 --- /dev/null +++ b/ndoptimize/_least_squares.py @@ -0,0 +1,914 @@ +""" +Nonlinear least squares for many independent fits at once. +""" + +import dataclasses +import numpy as np +import numpy.typing as npt +import numba +from numba import types +from ._model import ModelFunction, 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 +""" +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_band(band: np.ndarray, m: int) -> bool: + """ + 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 the matrix is not positive definite. + """ + w = band.shape[1] - 1 + for j in range(m): + 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 + pivot = np.sqrt(s) + band[j, 0] = pivot + for i in range(j + 1, min(m, j + w + 1)): + s = band[i, i - j] + for k in range(max(0, i - w), j): + 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_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 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 -= band[i, i - k] * b[k] + b[i] = s / band[i, 0] + + +@numba.njit(cache=True, error_model="numpy") +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 the lower-triangular `factor` stored as a band by + :func:`_cholesky_band`. + """ + w = band.shape[1] - 1 + for i in range(m): + 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 -= 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") +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. + """ + work[:, :] = a + if not _cholesky(work, 0.0): + return False + _solve_factored(work, b, out) + return True + + +@numba.njit(cache=True, error_model="numpy") +def _whitened_residual( + factor: np.ndarray, + m: 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_band(factor, out, m) + 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, + 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_band_columns(factor, out, m) + + +@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 + 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 + 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") +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 + + +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], + types.float64[:, :, ::1], + types.float64[::1], + types.int64[::1], + types.int64[::1], + types.boolean[::1], + ) +)( + types.FunctionType(signature_model), + _input(3), + _input(2), + _input(2), + _input(1), + _input(2), + _input(2), + _input(2), + _input(2), + _input(2), + 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: ModelFunction, + 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 compiled model function, :attr:`ndoptimize.Model.function`. + 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, not negative. + 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) + + 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.empty((num_point, num_lag)) + f = np.empty(num_point) + 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)) + 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 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, d] = sigma[r] * sigma[r - d] * rho + if not _cholesky_band(factor, m): + continue + _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] + 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, 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 + while True: + _normal(jacobian, residual, m, p, mean, precision, a, g) + + # 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] = 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 + # 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 + 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): + 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] * (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, data_whitened, f, residual_new + ) + objective_new = chi2_new + _chi2_prior(p_new, mean, precision) + accepted = objective_new < objective + 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 + 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, jac, jacobian) + + # the covariance of the parameters is the inverse of the normal + # 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) + for k in range(num_parameter): + 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 + _solve_factored(work, unit, step) + for q in range(num_parameter): + covariance[i, q, k] = step[q] + for k in range(num_parameter): + if 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, eq=False) +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 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 + """ + 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( + model: Model, + x: np.ndarray, + data: np.ndarray, + variance: np.ndarray, + guess: np.ndarray, + lower: np.ndarray, + upper: np.ndarray, + prior_mean: np.ndarray, + prior_sigma: np.ndarray, +) -> tuple[tuple[int, ...], dict[str, np.ndarray]]: + """ + 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. + """ + 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 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(variance.shape, (num_point,)) + ) + lower, upper, prior_mean, prior_sigma = ( + np.broadcast_to(a, np.broadcast_shapes(a.shape, (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` 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 a.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: 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: + 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, made 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)``, 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)``, 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, 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. + 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, 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 + -------- + + 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(); + """ + if max_iterations < 0: + raise ValueError( + f"`max_iterations` must not be negative, got {max_iterations}." + ) + + shape, a = _flatten( + 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 = 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 {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] + 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`, {rho}, does not describe a positive " + f"definite noise covariance for {num_point} points." + ) + + result = _levenberg_marquardt( + model.function, + a["x"], + a["data"], + a["variance"], + rho, + 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..5403640 --- /dev/null +++ b/ndoptimize/_model.py @@ -0,0 +1,196 @@ +""" +The interface between a model function and the fitting kernels. +""" + +from typing import Callable, Sequence +import dataclasses +import functools +import numpy as np +import numba +from numba import types + +__all__ = [ + "ModelFunction", + "signature_model", + "Model", + "model", +] + +ModelFunction = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None] +""" +The type of a model function, ``function(x, parameters, f, jacobian)``. +""" + +signature_model = types.void( + types.float64[:, :], + types.float64[:], + types.float64[:], + types.float64[:, :], +) +""" +The Numba signature of a model function. +""" + + +@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 __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, + 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, with the given names for its parameters. + + 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 + 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 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`, + since they are the ones Numba compiles for every target. + + Parameters + ---------- + 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 + -------- + + A straight line, :math:`f(x) = a + b x`. + + .. jupyter-execute:: + + import ndoptimize + + @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 + """ + 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/__init__.py b/ndoptimize/_tests/__init__.py new file mode 100644 index 0000000..e69de29 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)) diff --git a/ndoptimize/_tests/test_least_squares.py b/ndoptimize/_tests/test_least_squares.py new file mode 100644 index 0000000..6f45de6 --- /dev/null +++ b/ndoptimize/_tests/test_least_squares.py @@ -0,0 +1,477 @@ +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, 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) + + # 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) + 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(parameters=("intercept", "slope")) +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 + + +@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.""" + 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(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: + 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..6cbbd7c --- /dev/null +++ b/ndoptimize/_tests/test_models.py @@ -0,0 +1,141 @@ +from typing import Callable +import copy +import pickle +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, + 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)) + + assert len(model.parameters) == 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) + + +@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) + + +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 diff --git a/ndoptimize/models.py b/ndoptimize/models.py new file mode 100644 index 0000000..817f22d --- /dev/null +++ b/ndoptimize/models.py @@ -0,0 +1,164 @@ +""" +Common models, made with :func:`ndoptimize.model`. +""" + +import math +import numpy as np +from . import _model + +__all__ = [ + "gaussian", + "gaussian_linear", + "two_gaussians_linear", +] + + +@_model.model( + parameters=("amplitude", "center", "width"), +) +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( + parameters=("amplitude", "center", "width", "slope", "intercept"), +) +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( + parameters=( + "amplitude", + "center", + "width", + "slope", + "intercept", + "amplitude_2", + "center_2", + "width_2", + ), +) +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"]