propagate functions with several arguments (#16)

.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

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

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

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)

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)