From a4bfe934ac54fb35fe2b0feda909398e5d134070 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Tue, 20 Oct 2026 11:24:11 -0600 Subject: [PATCH 1/2] Make models declare their inputs, so that `least_squares()` and `evaluate()` can check `x`. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- README.md | 10 +++- docs/index.rst | 14 +++-- ndoptimize/_evaluate.py | 8 ++- ndoptimize/_least_squares.py | 8 ++- ndoptimize/_model.py | 78 ++++++++++++++++++------- ndoptimize/_tests/test_evaluate.py | 2 + ndoptimize/_tests/test_least_squares.py | 6 +- ndoptimize/_tests/test_models.py | 65 ++++++++++++++++----- ndoptimize/models.py | 24 ++++---- 9 files changed, 155 insertions(+), 60 deletions(-) diff --git a/README.md b/README.md index f31669e..f8be5cb 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,8 @@ standard deviation in the tests, but they differ in the following ways. Gauss-Newton step predicts replaces `ftol`, `xtol` and `gtol`. - **Model arguments.** A model takes only the points `x` and the parameters, so any other data it needs, such as the edges of each bin, goes in more - columns of `x`. SciPy passes extra data through `args` and `kwargs`. + columns of `x`, each one an input that the model names. SciPy passes extra + data through `args` and `kwargs`. - **Results.** `chi2` is the sum of the squares of the whitened residuals, without the priors, while SciPy's `cost` is half the sum of the squares of whatever residual it is given. `max_iterations` and `iterations` count @@ -133,7 +134,10 @@ value at the points `x`, an array of shape `(num_points, num_inputs)`, into `f`, and its derivatives with respect to the `parameters` into `jacobian`. The [`ndoptimize.model`](https://ndoptimize.readthedocs.io/en/latest/_autosummary/ndoptimize.model-decorator.html) -decorator compiles it and names its parameters. +decorator compiles it and names its parameters and its inputs, the columns of +`x`. +`least_squares()` and `evaluate()` check that `x` has one column for each +input, since the compiled model does not check its indices. The compiled model is cached on disk, so define it in a file or a notebook. A model defined in the interactive interpreter needs `@ndoptimize.model(..., cache=False)`. @@ -143,7 +147,7 @@ import math import numpy as np import ndoptimize -@ndoptimize.model(parameters=("amplitude", "scale")) +@ndoptimize.model(parameters=("amplitude", "scale"), inputs=("x",)) def exponential( x: np.ndarray, parameters: np.ndarray, diff --git a/docs/index.rst b/docs/index.rst index 1b94b6d..89a3e1e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -35,7 +35,7 @@ The package provides: - :func:`ndoptimize.evaluate`, which evaluates a model for many sets of parameters at once, such as the best-fit ones. - :func:`ndoptimize.model`, a decorator which compiles a function as a model, - an :class:`ndoptimize.Model`, and names its parameters. + an :class:`ndoptimize.Model`, and names its parameters and inputs. - :mod:`ndoptimize.models`, some common models: a Gaussian, a Gaussian averaged over each bin, a Gaussian plus a straight line, and two Gaussians plus a straight line. @@ -62,8 +62,8 @@ standard deviation in the tests, but they differ in the following ways. Gauss-Newton step predicts replaces ``ftol``, ``xtol`` and ``gtol``. - **Model arguments.** A model takes only the points ``x`` and the parameters, so any other data it needs, such as the edges of each bin, goes - in more columns of ``x``. SciPy passes extra data through ``args`` and - ``kwargs``. + in more columns of ``x``, each one an input that the model names. SciPy + passes extra data through ``args`` and ``kwargs``. - **Results.** :attr:`~ndoptimize.LeastSquaresResult.chi2` is the sum of the squares of the whitened residuals, without the priors, while SciPy's ``cost`` is half the sum of the squares of whatever residual it is given. @@ -154,7 +154,11 @@ A model is a function, ``model(x, parameters, f, jacobian)``, which writes its value at the points ``x``, an array of shape ``(num_points, num_inputs)``, into ``f``, and its derivatives with respect to the ``parameters`` into ``jacobian``. -The :func:`ndoptimize.model` decorator compiles it and names its parameters. +The :func:`ndoptimize.model` decorator compiles it and names its parameters +and its inputs, the columns of ``x``. +:func:`~ndoptimize.least_squares` and :func:`~ndoptimize.evaluate` check that +``x`` has one column for each input, since the compiled model does not check +its indices. The compiled model is cached on disk, so define it in a file or a notebook. A model defined in the interactive interpreter needs ``@ndoptimize.model(..., cache=False)``. @@ -166,7 +170,7 @@ keep the amplitude from going negative and the scale at least 0.1. import math - @ndoptimize.model(parameters=("amplitude", "scale")) + @ndoptimize.model(parameters=("amplitude", "scale"), inputs=("x",)) def exponential( x: np.ndarray, parameters: np.ndarray, diff --git a/ndoptimize/_evaluate.py b/ndoptimize/_evaluate.py index c3b130b..654f576 100644 --- a/ndoptimize/_evaluate.py +++ b/ndoptimize/_evaluate.py @@ -85,7 +85,8 @@ def evaluate( 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)``. + ``(..., num_points, num_inputs)``, where ``num_inputs`` is the number + of names in :attr:`ndoptimize.Model.inputs`. parameters The parameters of the model, an array of shape ``(..., num_parameters)``, where ``num_parameters`` is the number of @@ -146,6 +147,11 @@ def _evaluate_arrays( num_point, num_input = x.shape[-2:] num_parameter = len(model.parameters) + if num_input != len(model.inputs): + raise ValueError( + f"The last axis of `x`, {x.shape}, must match the number of " + f"inputs of the model, {model.inputs}." + ) if parameters.shape[-1:] != (num_parameter,): raise ValueError( f"The last axis of `parameters`, {parameters.shape}, must match " diff --git a/ndoptimize/_least_squares.py b/ndoptimize/_least_squares.py index 97b75d0..6ce6ba5 100644 --- a/ndoptimize/_least_squares.py +++ b/ndoptimize/_least_squares.py @@ -609,6 +609,11 @@ def _flatten( num_point, num_input = x.shape[-2:] num_parameter = len(model.parameters) + if num_input != len(model.inputs): + raise ValueError( + f"The last axis of `x`, {x.shape}, must match the number of " + f"inputs of the model, {model.inputs}." + ) if data.shape[-1:] != (num_point,): raise ValueError( f"The last axis of `data`, {data.shape}, must match the " @@ -721,7 +726,8 @@ def least_squares( The model to fit, made with :func:`ndoptimize.model`. x The independent variable, an array of shape - ``(..., num_points, num_inputs)``. + ``(..., num_points, num_inputs)``, where ``num_inputs`` is the number + of names in :attr:`ndoptimize.Model.inputs`. data The data to fit, an array of shape ``(..., num_points)``. Points which are NaN are left out of their fit. diff --git a/ndoptimize/_model.py b/ndoptimize/_model.py index 5403640..2baa638 100644 --- a/ndoptimize/_model.py +++ b/ndoptimize/_model.py @@ -36,7 +36,8 @@ 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`. + function and the names of its parameters and inputs, 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. @@ -54,11 +55,19 @@ class Model: and of the columns of the ``jacobian`` array. """ + inputs: tuple[str, ...] + """ + The names of the inputs, in the order of the columns of the ``x`` 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})" + return ( + f"{type(self).__name__}(parameters={self.parameters}, " + f"inputs={self.inputs})" + ) def __get__(self, instance: object, owner: None | type = None) -> "Model": """ @@ -92,7 +101,8 @@ def __call__( ---------- x The points at which to evaluate the model, - ``(num_points, num_inputs)``. + ``(num_points, num_inputs)``, where ``num_inputs`` is the number + of names in :attr:`inputs`. parameters The parameters of the model, ``(num_parameters,)``. f @@ -102,38 +112,57 @@ def __call__( to the parameters, ``(num_points, num_parameters)``. """ num_point = np.shape(x)[0] if np.ndim(x) == 2 else -1 + num_input = len(self.inputs) num_parameter = len(self.parameters) if ( np.ndim(x) != 2 + or np.shape(x)[1] != num_input 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"For a model with {num_input} inputs and {num_parameter} " + f"parameters, `x` must have shape (num_points, {num_input}), " + f"`parameters` ({num_parameter},), `f` (num_points,) and " + f"`jacobian` (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 _names(names: Sequence[str], argument: str) -> tuple[str, ...]: + """ + Check that `names`, the argument of :func:`model` named `argument`, is a + sequence of unique names, and return it as a tuple. + """ + if isinstance(names, str): + raise TypeError( + f"`{argument}` must be a sequence of names, got the string {names!r}." + ) + names = tuple(names) + if len(set(names)) != len(names): + raise ValueError(f"The names in `{argument}`, {names}, must be unique.") + return names + + def model( parameters: Sequence[str], + inputs: 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. + fit, with the given names for its parameters and inputs. 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)``. + ``(num_points, num_inputs)`` with a column for each input, 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 @@ -152,6 +181,13 @@ def model( :func:`ndoptimize.least_squares` checks the initial guess against their number, since the compiled function does not check its indices. + inputs + The names of the inputs, in the order of the columns of the ``x`` + array, such as ``("x",)`` for a function of one variable, or + ``("lower", "upper")`` for a function of the edges of a bin. + :func:`ndoptimize.least_squares` and :func:`ndoptimize.evaluate` + check the last axis of ``x`` against their number, for the same + reason. cache Whether to cache the compiled function on disk, so that later processes need not compile it again. @@ -170,7 +206,7 @@ def model( import ndoptimize - @ndoptimize.model(parameters=("intercept", "slope")) + @ndoptimize.model(parameters=("intercept", "slope"), inputs=("x",)) def line(x, parameters, f, jacobian): a, b = parameters[0], parameters[1] for i in range(x.shape[0]): @@ -180,17 +216,15 @@ def line(x, parameters, f, jacobian): 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.") + parameters = _names(parameters, "parameters") + inputs = _names(inputs, "inputs") def decorator(function: ModelFunction) -> Model: compiled = numba.njit(signature_model, cache=cache, error_model="numpy") - return Model(function=compiled(function), parameters=parameters) + return Model( + function=compiled(function), + parameters=parameters, + inputs=inputs, + ) return decorator diff --git a/ndoptimize/_tests/test_evaluate.py b/ndoptimize/_tests/test_evaluate.py index 9ab9a11..688c41f 100644 --- a/ndoptimize/_tests/test_evaluate.py +++ b/ndoptimize/_tests/test_evaluate.py @@ -86,6 +86,8 @@ def test_evaluate_least_squares() -> None: argnames="kwargs", argvalues=[ dict(x=np.linspace(-50, 50, 21)), + dict(x=np.zeros((21, 2))), + dict(model=ndoptimize.models.gaussian_binned), dict(parameters=np.array([1, 0.0])), dict(parameters=1.0), ], diff --git a/ndoptimize/_tests/test_least_squares.py b/ndoptimize/_tests/test_least_squares.py index 6f45de6..7d6522a 100644 --- a/ndoptimize/_tests/test_least_squares.py +++ b/ndoptimize/_tests/test_least_squares.py @@ -284,7 +284,7 @@ def test_least_squares_unconstrained() -> None: assert np.isfinite(result.covariance[0, 0]) -@ndoptimize.model(parameters=("intercept", "slope")) +@ndoptimize.model(parameters=("intercept", "slope"), inputs=("x",)) def _line( x: np.ndarray, parameters: np.ndarray, @@ -324,7 +324,7 @@ def test_least_squares_linear() -> None: assert result.iterations <= 3 -@ndoptimize.model(parameters=("scale",)) +@ndoptimize.model(parameters=("scale",), inputs=("x",)) def _scaled( x: np.ndarray, parameters: np.ndarray, @@ -455,6 +455,8 @@ def test_least_squares_max_iterations() -> None: dict(correlation=np.array([1, np.nan])), dict(data=np.zeros(5)), dict(x=_velocity), + dict(x=np.stack([_velocity, _velocity], axis=-1)), + dict(model=ndoptimize.models.gaussian_binned), dict(guess=np.array([1, 0.0])), dict(guess=1.0), dict(lower=np.array([0, 5, 0]), upper=np.array([10, 2, 30])), diff --git a/ndoptimize/_tests/test_models.py b/ndoptimize/_tests/test_models.py index c8ce849..49b4e98 100644 --- a/ndoptimize/_tests/test_models.py +++ b/ndoptimize/_tests/test_models.py @@ -60,30 +60,37 @@ def test_models( @pytest.mark.parametrize( - argnames="shape_parameters,shape_f,shape_jacobian", + argnames="shape_x,shape_parameters,shape_f,shape_jacobian", argvalues=[ - ((2,), (21,), (21, 2)), - ((3,), (20,), (21, 3)), - ((3,), (21,), (21, 5)), + ((21, 1), (2,), (21,), (21, 2)), + ((21, 1), (3,), (20,), (21, 3)), + ((21, 1), (3,), (21,), (21, 5)), + ((21, 2), (3,), (21,), (21, 3)), + ((21,), (3,), (21,), (21, 3)), ], ) def test_model_shapes( + shape_x: tuple[int, ...], 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.""" + """Calling a model with arrays of the wrong shapes, including points with + more columns than the model has inputs, raises an error rather than + reading or writing outside them.""" with pytest.raises(ValueError): ndoptimize.models.gaussian( - np.zeros((21, 1)), + np.zeros(shape_x), np.ones(shape_parameters), np.empty(shape_f), np.empty(shape_jacobian), ) -def test_model_parameters() -> None: +def test_model_names() -> None: + """A model keeps the names of its parameters and inputs, and the + docstring of its function.""" + def constant( x: np.ndarray, parameters: np.ndarray, @@ -95,15 +102,11 @@ def constant( 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) + model = ndoptimize.model(parameters=["offset"], inputs=["x"], cache=False)(constant) assert model.parameters == ("offset",) + assert model.inputs == ("x",) assert model.__doc__ == "A constant." - assert repr(model) == "Model(parameters=('offset',))" + assert repr(model) == "Model(parameters=('offset',), inputs=('x',))" f = np.empty(3) jacobian = np.empty((3, 1)) @@ -112,11 +115,41 @@ def constant( assert np.all(jacobian == 1) +@pytest.mark.parametrize( + argnames="kwargs,error", + argvalues=[ + (dict(parameters="offset", inputs=("x",)), TypeError), + (dict(parameters=("offset", "offset"), inputs=("x",)), ValueError), + (dict(parameters=("offset",), inputs="x"), TypeError), + (dict(parameters=("offset",), inputs=("x", "x")), ValueError), + ], +) +def test_model_invalid(kwargs: dict[str, Any], error: type[Exception]) -> None: + """The names of the parameters and of the inputs must each be a sequence + of unique names, not a single string.""" + with pytest.raises(error): + ndoptimize.model(**kwargs) + + +@pytest.mark.parametrize( + argnames="model,inputs", + argvalues=[ + (ndoptimize.models.gaussian, ("x",)), + (ndoptimize.models.gaussian_linear, ("x",)), + (ndoptimize.models.two_gaussians_linear, ("x",)), + (ndoptimize.models.gaussian_binned, ("lower", "upper")), + ], +) +def test_models_inputs(model: ndoptimize.Model, inputs: tuple[str, ...]) -> None: + """The common models name the columns of `x` that they read.""" + assert model.inputs == inputs + + 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" + "@ndoptimize.model(parameters=('slope',), inputs=('x',), 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" diff --git a/ndoptimize/models.py b/ndoptimize/models.py index dc6182a..86ca24c 100644 --- a/ndoptimize/models.py +++ b/ndoptimize/models.py @@ -16,6 +16,7 @@ @_model.model( parameters=("amplitude", "center", "width"), + inputs=("x",), ) def gaussian( x: np.ndarray, @@ -31,12 +32,12 @@ def gaussian( 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`. + width, of a single input, :math:`x`. Parameters ---------- x - The points at which to evaluate the model, ``(num_points, num_inputs)``. + The points at which to evaluate the model, ``(num_points, 1)``. parameters The amplitude, center, and width. f @@ -57,6 +58,7 @@ def gaussian( @_model.model( parameters=("amplitude", "center", "width"), + inputs=("lower", "upper"), ) def gaussian_binned( x: np.ndarray, @@ -77,8 +79,8 @@ def gaussian_binned( \right], with the same parameters as :func:`gaussian`, the amplitude, center and - width, where :math:`x_0` and :math:`x_1`, the first two components of - :math:`x`, are the edges of each bin. + width, where its two inputs, ``lower`` and ``upper``, are the edges of + each bin, :math:`x_0` and :math:`x_1`. The integral of the Gaussian is :math:`\sqrt{2 \pi} A \sigma`. Averaging matters when the bins are not much narrower than the line. @@ -90,8 +92,8 @@ def gaussian_binned( Parameters ---------- x - The edges of each bin, ``(num_points, num_inputs)``, with the lower - edge in the first column and the upper edge in the second. + The edges of each bin, ``(num_points, 2)``, with the lower edge in the + first column and the upper edge in the second. parameters The amplitude, center, and width. f @@ -171,6 +173,7 @@ def gaussian_binned( @_model.model( parameters=("amplitude", "center", "width", "slope", "intercept"), + inputs=("x",), ) def gaussian_linear( x: np.ndarray, @@ -186,13 +189,13 @@ def gaussian_linear( 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`. + width, slope and intercept, of a single input, :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)``. + The points at which to evaluate the model, ``(num_points, 1)``. parameters The amplitude, center, width, slope and intercept. f @@ -226,6 +229,7 @@ def gaussian_linear( "center_2", "width_2", ), + inputs=("x",), ) def two_gaussians_linear( x: np.ndarray, @@ -242,14 +246,14 @@ def two_gaussians_linear( + 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`. + a single input, :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)``. + The points at which to evaluate the model, ``(num_points, 1)``. parameters The amplitude, center, width, slope and intercept, then the amplitude, center and width of the second Gaussian. From 6f9ef4a975bf995faab567325047cc5a69d17f55 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Tue, 20 Oct 2026 17:25:36 -0600 Subject: [PATCH 2/2] Check the names of every `Model`, take `inputs` and `cache` by keyword, and check `x` and the parameters in one `Model` method. - The names of the parameters and inputs must be a sequence of strings, so a set, whose order changes from one process to the next, bytes or `None` raise a `TypeError`. `Model.__post_init__` checks them too, so a model made directly is checked like one made with `model()`. - `inputs` and `cache` are keyword-only, so the inputs cannot be mistaken for the parameters, and 0.1 code which passed `cache` by position gets a clear error. - `Model._check_shapes()` checks the last axes of `x` and of the parameters against the model for `Model.__call__`, `least_squares()` and `evaluate()`, instead of three copies. The error from calling a model names its inputs and parameters. - `test_model_shapes` calls `gaussian_binned` with one column, and `test_model_invalid` makes each invalid model both ways. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PZRXiFhGzzrF87vsFg73B9 --- ndoptimize/_evaluate.py | 17 +------ ndoptimize/_least_squares.py | 16 +----- ndoptimize/_model.py | 86 ++++++++++++++++++++++---------- ndoptimize/_tests/test_models.py | 54 +++++++++++++++----- 4 files changed, 105 insertions(+), 68 deletions(-) diff --git a/ndoptimize/_evaluate.py b/ndoptimize/_evaluate.py index 654f576..af998f4 100644 --- a/ndoptimize/_evaluate.py +++ b/ndoptimize/_evaluate.py @@ -139,25 +139,10 @@ def _evaluate_arrays( 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}." - ) + model._check_shapes(x, parameters, "parameters") num_point, num_input = x.shape[-2:] num_parameter = len(model.parameters) - if num_input != len(model.inputs): - raise ValueError( - f"The last axis of `x`, {x.shape}, must match the number of " - f"inputs of the model, {model.inputs}." - ) - 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) diff --git a/ndoptimize/_least_squares.py b/ndoptimize/_least_squares.py index 6ce6ba5..458aca6 100644 --- a/ndoptimize/_least_squares.py +++ b/ndoptimize/_least_squares.py @@ -601,29 +601,15 @@ def _flatten( 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}." - ) + model._check_shapes(x, guess, "guess") num_point, num_input = x.shape[-2:] num_parameter = len(model.parameters) - if num_input != len(model.inputs): - raise ValueError( - f"The last axis of `x`, {x.shape}, must match the number of " - f"inputs of the model, {model.inputs}." - ) 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): diff --git a/ndoptimize/_model.py b/ndoptimize/_model.py index 2baa638..15dbe69 100644 --- a/ndoptimize/_model.py +++ b/ndoptimize/_model.py @@ -3,6 +3,7 @@ """ from typing import Callable, Sequence +import collections.abc import dataclasses import functools import numpy as np @@ -32,6 +33,26 @@ """ +def _names(names: Sequence[str], argument: str) -> tuple[str, ...]: + """ + Check that `names`, the argument named `argument`, is a sequence of unique + names, in order, and return it as a tuple. + """ + if ( + isinstance(names, str) + or not isinstance(names, collections.abc.Sequence) + or not all(isinstance(name, str) for name in names) + ): + raise TypeError( + f"`{argument}` must be a sequence of names, such as a tuple of " + f"strings, got {names!r}." + ) + names = tuple(names) + if len(set(names)) != len(names): + raise ValueError(f"The names in `{argument}`, {names}, must be unique.") + return names + + @dataclasses.dataclass(eq=False) class Model: """ @@ -61,6 +82,8 @@ class Model: """ def __post_init__(self) -> None: + self.parameters = _names(self.parameters, "parameters") + self.inputs = _names(self.inputs, "inputs") functools.update_wrapper(self, self.function, updated=()) def __repr__(self) -> str: @@ -87,6 +110,34 @@ def __reduce__(self) -> str: """ return self.function.__qualname__ + def _check_shapes( + self, + x: np.ndarray, + parameters: np.ndarray, + name: str, + ) -> None: + """ + Check the last axis of the points `x`, ``(..., num_points, num_inputs)``, + against the inputs of this model, and the last axis of `parameters`, + ``(..., num_parameters)``, the argument named `name`, against its + parameters, since the compiled function does not check its indices. + """ + if np.ndim(x) < 2: + raise ValueError( + f"`x` must have at least two axes, (..., num_points, num_inputs), " + f"got shape {np.shape(x)}." + ) + if np.shape(x)[-1] != len(self.inputs): + raise ValueError( + f"The last axis of `x`, {np.shape(x)}, must match the number of " + f"inputs of the model, {self.inputs}." + ) + if np.shape(parameters)[-1:] != (len(self.parameters),): + raise ValueError( + f"The last axis of `{name}`, {np.shape(parameters)}, must match " + f"the number of parameters of the model, {self.parameters}." + ) + def __call__( self, x: np.ndarray, @@ -111,44 +162,29 @@ def __call__( 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_input = len(self.inputs) + self._check_shapes(x, parameters, "parameters") + num_point = np.shape(x)[0] num_parameter = len(self.parameters) if ( np.ndim(x) != 2 - or np.shape(x)[1] != num_input - or np.shape(parameters) != (num_parameter,) + or np.ndim(parameters) != 1 or np.shape(f) != (num_point,) or np.shape(jacobian) != (num_point, num_parameter) ): raise ValueError( - f"For a model with {num_input} inputs and {num_parameter} " - f"parameters, `x` must have shape (num_points, {num_input}), " - f"`parameters` ({num_parameter},), `f` (num_points,) and " - f"`jacobian` (num_points, {num_parameter}), got {np.shape(x)}, " - f"{np.shape(parameters)}, {np.shape(f)} and " - f"{np.shape(jacobian)}." + f"For a model with the inputs {self.inputs} and the parameters " + f"{self.parameters}, `x` must have shape " + f"(num_points, {len(self.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 {np.shape(jacobian)}." ) self.function(x, parameters, f, jacobian) -def _names(names: Sequence[str], argument: str) -> tuple[str, ...]: - """ - Check that `names`, the argument of :func:`model` named `argument`, is a - sequence of unique names, and return it as a tuple. - """ - if isinstance(names, str): - raise TypeError( - f"`{argument}` must be a sequence of names, got the string {names!r}." - ) - names = tuple(names) - if len(set(names)) != len(names): - raise ValueError(f"The names in `{argument}`, {names}, must be unique.") - return names - - def model( parameters: Sequence[str], + *, inputs: Sequence[str], cache: bool = True, ) -> Callable[[ModelFunction], Model]: diff --git a/ndoptimize/_tests/test_models.py b/ndoptimize/_tests/test_models.py index 49b4e98..b7248ed 100644 --- a/ndoptimize/_tests/test_models.py +++ b/ndoptimize/_tests/test_models.py @@ -1,5 +1,6 @@ from typing import Any, Callable import copy +import functools import math import pickle import numpy as np @@ -60,26 +61,30 @@ def test_models( @pytest.mark.parametrize( - argnames="shape_x,shape_parameters,shape_f,shape_jacobian", + argnames="model,shape_x,shape_parameters,shape_f,shape_jacobian", argvalues=[ - ((21, 1), (2,), (21,), (21, 2)), - ((21, 1), (3,), (20,), (21, 3)), - ((21, 1), (3,), (21,), (21, 5)), - ((21, 2), (3,), (21,), (21, 3)), - ((21,), (3,), (21,), (21, 3)), + (ndoptimize.models.gaussian, (21, 1), (2,), (21,), (21, 2)), + (ndoptimize.models.gaussian, (21, 1), (3,), (20,), (21, 3)), + (ndoptimize.models.gaussian, (21, 1), (3,), (21,), (21, 5)), + (ndoptimize.models.gaussian, (21, 2), (3,), (21,), (21, 3)), + (ndoptimize.models.gaussian, (21,), (3,), (21,), (21, 3)), + (ndoptimize.models.gaussian, (2, 21, 1), (3,), (2,), (2, 3)), + (ndoptimize.models.gaussian, (21, 1), (2, 3), (21,), (21, 3)), + (ndoptimize.models.gaussian_binned, (21, 1), (3,), (21,), (21, 3)), ], ) def test_model_shapes( + model: ndoptimize.Model, shape_x: tuple[int, ...], shape_parameters: tuple[int, ...], shape_f: tuple[int, ...], shape_jacobian: tuple[int, ...], ) -> None: """Calling a model with arrays of the wrong shapes, including points with - more columns than the model has inputs, raises an error rather than - reading or writing outside them.""" + more or fewer columns than the model has inputs, raises an error rather + than reading or writing outside them.""" with pytest.raises(ValueError): - ndoptimize.models.gaussian( + model( np.zeros(shape_x), np.ones(shape_parameters), np.empty(shape_f), @@ -120,15 +125,40 @@ def constant( argvalues=[ (dict(parameters="offset", inputs=("x",)), TypeError), (dict(parameters=("offset", "offset"), inputs=("x",)), ValueError), + (dict(parameters={"offset", "slope"}, inputs=("x",)), TypeError), (dict(parameters=("offset",), inputs="x"), TypeError), + (dict(parameters=("offset",), inputs=b"x"), TypeError), (dict(parameters=("offset",), inputs=("x", "x")), ValueError), + (dict(parameters=("offset",), inputs={"lower", "upper"}), TypeError), + (dict(parameters=("offset",), inputs=("x", None)), TypeError), + (dict(parameters=("offset",), inputs=None), TypeError), + ], +) +@pytest.mark.parametrize( + argnames="make", + argvalues=[ + ndoptimize.model, + functools.partial(ndoptimize.Model, ndoptimize.models.gaussian.function), ], + ids=["model", "Model"], ) -def test_model_invalid(kwargs: dict[str, Any], error: type[Exception]) -> None: +def test_model_invalid( + make: Callable[..., Any], + kwargs: dict[str, Any], + error: type[Exception], +) -> None: """The names of the parameters and of the inputs must each be a sequence - of unique names, not a single string.""" + of unique names, in order, not a single string or a set, whether the + model is made with the decorator or directly.""" with pytest.raises(error): - ndoptimize.model(**kwargs) + make(**kwargs) + + +def test_model_keywords() -> None: + """The inputs of a model are given by keyword, so that they cannot be + mistaken for its parameters.""" + with pytest.raises(TypeError): + ndoptimize.model(("offset",), ("x",)) # type: ignore[misc] @pytest.mark.parametrize(