diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d7ba3ca..d787190 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -47,5 +47,12 @@ repos: rev: 'v0.971' hooks: - id: mypy - args: [--allow-redefinition, --ignore-missing-imports, src] + args: [src] pass_filenames: false + +# doc string checking +- repo: https://github.com/PyCQA/pydocstyle + rev: 6.1.1 + hooks: + - id: pydocstyle + files: src/jacobi/.*\.py diff --git a/pyproject.toml b/pyproject.toml index 6a2a1ea..c12514f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,3 +7,7 @@ build-backend = "setuptools.build_meta" [tool.setuptools_scm] write_to = "src/jacobi/_version.py" + +[tool.mypy] +ignore_missing_imports = true +allow_redefinition = true diff --git a/src/jacobi/__init__.py b/src/jacobi/__init__.py index ee09996..be87199 100644 --- a/src/jacobi/__init__.py +++ b/src/jacobi/__init__.py @@ -1,7 +1,4 @@ -"""Jacobi - -Fast numerical derivatives for real analytic functions with arbitrary round-off error. -""" +"""Fast numerical derivatives for analytic functions with arbitrary round-off error.""" from .core import jacobi, propagate # noqa from ._version import version as __version__ # noqa diff --git a/src/jacobi/core.py b/src/jacobi/core.py index 5e6ead8..1f4fe65 100644 --- a/src/jacobi/core.py +++ b/src/jacobi/core.py @@ -1,3 +1,5 @@ +"""Core functions of jacobi.""" + import numpy as np import typing as _tp @@ -74,15 +76,15 @@ def _first(method, f0, f, x, i, h, args): def jacobi( fn: _tp.Callable, - x: _tp.Union[int, float, _tp.Sequence], + x: _tp.Union[float, _Indexable[float]], *args, - method: _tp.Optional[int] = None, - mask: _tp.Optional[np.ndarray] = None, + method: int = None, + mask: np.ndarray = None, rtol: float = 0, maxiter: int = 10, maxgrad: int = 3, - step: _tp.Optional[_tp.Tuple[float, float]] = None, - diagnostic: _tp.Optional[dict] = None, + step: _tp.Tuple[float, float] = None, + diagnostic: dict = None, ): """ Return first derivative and its error estimate. @@ -240,6 +242,7 @@ def propagate( fn: _tp.Callable, x: _tp.Union[float, _Indexable[float]], cov: _tp.Union[float, _Indexable[float], _Indexable[_Indexable[float]]], + *args, **kwargs, ) -> _tp.Tuple[np.ndarray, np.ndarray]: """ @@ -252,14 +255,21 @@ def propagate( Parameters ---------- fn: callable - Function that computes y = fn(x). x and y are each allowed to be scalars or - one-dimensional arrays. + Function that computes r = fn(x, [y, ...]). The arguments of the function are + each allowed to be scalars or one-dimensional arrays. If the function accepts + several arguments, their uncertainties are treated as uncorrelated. + Functions that accept several correlated arguments must be wrapped, see examples. + The result of the function may be a scalar or a one-dimensional array with a + different lenth as the input. x: float or array-like with shape (N,) - Input vector. + Input vector. An array-like is converted before passing it to the callable. cov: float or array-like with shape (N,) or shape(N, N) Covariance matrix of input vector. If the array is one-dimensional, it is interpreted as the diagonal of a covariance matrix with zero off-diagonal elements. + *args: + If the function accepts several arguments that are mutually independent, these + is possible to pass those values and covariance matrices pairwise, see examples. **kwargs: Extra arguments are passed to :func:`jacobi`. @@ -270,14 +280,93 @@ def propagate( ycov is the propagated covariance matrix. If ycov is a matrix, unless y is a number. In that case, ycov is also reduced to a number. + + Examples + -------- + General error propagation maps input vectors to output vectors:: + + def fn(x): + return x ** 2 + 1 + + x = [1, 2] + xcov = [[3, 1], + [1, 4]] + + y, ycov = propagate(fn, x, xcov) + + If the function accepts several arguments, their uncertainties are treated as + uncorrelated:: + + def fn(x, y): + return x + y + + x = 1 + y = 2 + xcov = 2 + ycov = 3 + + z, zcov = propagate(fn, x, xcov, y, ycov) + + Functions that accept several correlated arguments must be wrapped:: + + def fn(x, y): + return x + y + + x = 1 + y = 2 + sigma_x = 3 + sigma_y = 4 + rho_xy = 0.5 + + r = [x, y] + cov_xy = rho_xy * sigma_x * sigma_y + rcov = [[sigma_x ** 2, cov_xy], [cov_xy, sigma_y ** 2]] + + def fn_wrapped(r): + return fn(r[0], r[1]) + + z, zcov = propagate(fn_wrapped, r, rcov) + + See Also + -------- + jacobi """ + if args: + if len(args) % 2 != 0: + raise ValueError("number of extra positional arguments must be even") + + x_parts: _tp.List[np.ndarray] = [ + np.atleast_1d(_) for _ in ([x] + [a for a in args[::2]]) + ] + cov_parts: _tp.List[np.ndarray] = [ + np.atleast_1d(_) for _ in ([cov] + [a for a in args[1::2]]) + ] + slices = [] + i = 0 + for xi in x_parts: + n = len(xi) + slices.append(slice(i, i + n)) + i += n + + r = np.concatenate(x_parts) + n = len(r) + rcov = np.zeros((n, n)) + for sl, covi in zip(slices, cov_parts): + rcov[sl, sl] = np.diag(covi) if covi.ndim == 1 else covi + + def wrapped(r): + args = [r[sl] for sl in slices] + return fn(*args) + + return propagate(wrapped, r, rcov) + x = np.array(x) y = fn(x) jac = jacobi(fn, x, **kwargs)[0] x_nd = np.ndim(x) y_nd = np.ndim(y) - x_len = len(x) if x_nd == 1 else 1 + x_len = len(x) if x_nd == 1 else 1 # type: ignore y_len = len(y) if y_nd == 1 else 1 jac_nd = np.ndim(jac) diff --git a/test/test_propagate.py b/test/test_propagate.py index 4975c5e..3f2e139 100644 --- a/test/test_propagate.py +++ b/test/test_propagate.py @@ -75,3 +75,66 @@ def fn(x): else: xcov2 = xcov assert_allclose(ycov, np.linalg.multi_dot([jac, xcov2, jac.T])) + + +@pytest.mark.parametrize("ndim", (1, 2)) +def test_cov_1d_2d(ndim): + def fn(x): + return x + + x = [1, 2] + xcov_1d = [3, 4] + xcov_2d = np.diag(xcov_1d) + + y, ycov = propagate(fn, x, xcov_1d if ndim == 1 else xcov_2d) + + assert np.ndim(ycov) == 2 + + assert_allclose(y, x) + assert_allclose(ycov, xcov_2d) + + +def test_two_arguments_1(): + def fn1(x, y): + return (x - y) / (x + y) + + x = 1 + xcov = 2 + y = 3 + ycov = 4 + + z1, zcov1 = propagate(fn1, x, xcov, y, ycov) + + def fn2(r): + return fn1(r[0], r[1]) + + r = [x, y] + rcov = np.diag([xcov, ycov]) + + z2, zcov2 = propagate(fn2, r, rcov) + + assert_allclose(z2, z1) + assert_allclose(zcov2, zcov1) + + +def test_two_arguments_2(): + def fn1(x, y): + return np.concatenate([x, np.atleast_1d(y)]) + + x = [1, 2] + xcov = [2, 3] + y = 3 + ycov = 4 + + z1, zcov1 = propagate(fn1, x, xcov, y, ycov) + + def fn2(r): + return fn1(r[:2], r[2]) + + r = [*x, y] + rcov = np.diag([*xcov, ycov]) + + z2, zcov2 = propagate(fn2, r, rcov) + + assert_allclose(z2, z1) + assert_allclose(zcov2, zcov1)