Make models declare their inputs, so that least_squares() and evaluate() can check x.

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,

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,

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
@@ -138,20 +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 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)
 

ndoptimize/_least_squares.py

@@ -601,11 +601,7 @@ 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)
 
@@ -614,11 +610,6 @@ def _flatten(
             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):
@@ -721,7 +712,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.

ndoptimize/_model.py

@@ -3,6 +3,7 @@
 """
 
 from typing import Callable, Sequence
+import collections.abc
 import dataclasses
 import functools
 import numpy as np
@@ -32,11 +33,32 @@
 """
 
 
+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:
     """
     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 +76,21 @@ 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:
+        self.parameters = _names(self.parameters, "parameters")
+        self.inputs = _names(self.inputs, "inputs")
         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":
         """
@@ -78,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,
@@ -92,7 +152,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
@@ -101,39 +162,43 @@ 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
+        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(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_parameter} parameters, `x` must have "
-                f"shape (num_points, num_inputs), `parameters` "
+                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 "
-                f"{np.shape(jacobian)}."
+                f"{np.shape(parameters)}, {np.shape(f)} and {np.shape(jacobian)}."
             )
         self.function(x, parameters, f, jacobian)
 
 
 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 +217,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 +242,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 +252,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

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),
     ],

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])),