Simplify shape handling (#28)

src/jacobi/_jacobi.py

@@ -89,57 +89,61 @@ def jacobi(
     if method is not None and method not in (-1, 0, 1):
         raise ValueError("method must be -1, 0, 1")
 
-    squeeze = np.ndim(x) == 0
-    x = np.atleast_1d(x).astype(float)
-    assert x.ndim == 1
-
-    x_indices = np.arange(len(x))
+    x = np.asarray(x, dtype=float)
     if mask is not None:
-        x_indices = x_indices[mask]
-    nx = len(x_indices)
+        mask = np.asarray(mask)
+        if mask.dtype != bool:
+            raise ValueError("mask must be a boolean array")
+        if mask.shape != x.shape:
+            raise ValueError("mask shape must match x shape")
 
     if diagnostic is not None:
-        diagnostic["method"] = np.zeros(nx, dtype=np.int8)
-        diagnostic["iteration"] = np.zeros(len(x_indices), dtype=np.uint8)
-        diagnostic["residual"] = [[] for _ in range(nx)]
+        diagnostic["method"] = np.zeros(x.size, dtype=np.int8)
+        diagnostic["iteration"] = np.zeros(x.size, dtype=np.uint8)
+        diagnostic["residual"] = [[] for _ in range(x.size)]
 
     f0 = None
     jac = None
     err = None
-    for ik, k in enumerate(x_indices):
+    it = np.nditer(x, flags=["c_index", "multi_index"])
+    while not it.finished:
+        k = it.index
+        kx = it.multi_index if it.has_multi_index else ...
+        if mask is not None and not mask[kx]:
+            it.iternext()
+            continue
+        xk = it[0]
         # if step is None, use optimal step sizes for central derivatives
-        h = _steps(x[k], step or (0.25, 0.5), maxiter)
+        h = _steps(xk, step or (0.25, 0.5), maxiter)
         # if method is None, auto-detect for each x[k]
-        md, f0, r = _first(method, f0, fn, x, k, h[0], args)
+        md, f0, r = _first(method, f0, fn, x, kx, h[0], args)
+        # f0 is not guaranteed to be set here and can be still None
 
         if md != 0 and step is None:
-            # optimal step sizes for forward derivatives
-            h = _steps(x[k], (0.125, 0.125), maxiter)
+            # need different step sizes for forward derivatives to avoid overlap
+            h = _steps(xk, (0.25, 0.125), maxiter)
 
-        r_shape = np.shape(r)
-        r = np.reshape(r, -1)
-        nr = len(r)
-        re = np.full(nr, np.inf)
-        todo = np.ones(nr, dtype=bool)
-        fd = [r]
+        r = np.asarray(r, dtype=float)
+        re = np.full_like(r, np.inf)
+        todo = np.ones_like(r, dtype=bool)
+        fd = [np.reshape(r.copy(), -1)]
 
         if jac is None:  # first iteration
-            jac = np.empty(r_shape + (nx,), dtype=r.dtype)
-            err = np.empty(r_shape + (nx,), dtype=r.dtype)
+            jac = np.zeros(r.shape + x.shape, dtype=r.dtype)
+            err = np.zeros(r.shape + x.shape, dtype=r.dtype)
             if diagnostic is not None:
-                diagnostic["call"] = np.zeros((nr, nx), dtype=np.uint8)
+                diagnostic["call"] = np.zeros((r.size, x.size), dtype=np.uint8)
 
         if diagnostic is not None:
-            diagnostic["method"][ik] = md
-            diagnostic["call"][:, ik] = 2 if md == 0 else 3
+            diagnostic["method"][k] = md
+            diagnostic["call"][:, k] = 2 if md == 0 else 3
 
         for i in range(1, len(h)):
-            fdi = _derive(md, f0, fn, x, k, h[i], args)
-            fdi = np.reshape(fdi, -1)
-            fd.append(fdi if i == 1 else fdi[todo])
+            fdi = _derive(md, f0, fn, x, kx, h[i], args)
+            fd.append(np.reshape(fdi, -1) if i == 1 else fdi[todo])
             if diagnostic is not None:
-                diagnostic["call"][todo, ik] += 2
-                diagnostic["iteration"][ik] += 1
+                diagnostic["call"][todo, k] += 2
+                diagnostic["iteration"][k] += 1
 
             # polynomial fit with one extra degree of freedom;
             # use latest maxgrad + 1 data points
@@ -170,25 +174,25 @@ def jacobi(
             if diagnostic is not None:
                 re2 = re.copy()
                 re2[todo1] = rei
-                diagnostic["residual"][ik].append(re2)
+                diagnostic["residual"][k].append(re2)
 
             if np.sum(todo) == 0:
                 break
 
             # shrink previous vectors of estimates
             fd = [fdi[sub_todo] for fdi in fd]
 
-        jac[..., ik] = r.reshape(r_shape)
-        err[..., ik] = re.reshape(r_shape)
-
-    if diagnostic is not None:
-        diagnostic["call"].shape = r_shape + (nx,)
+        if jac.ndim == 0:
+            jac[...] = r
+            err[...] = re
+        elif jac.ndim == 1:
+            jac[kx] = r
+            err[kx] = re
+        else:
+            jac[(...,) + kx] = r
+            err[(...,) + kx] = re
 
-    if squeeze:
-        if diagnostic is not None:
-            diagnostic["call"] = np.squeeze(diagnostic["call"])
-        jac = np.squeeze(jac)
-        err = np.squeeze(err)
+        it.iternext()
 
     return jac, err
 

src/jacobi/_propagate.py

@@ -1,16 +1,19 @@
-import typing as _tp
-from ._typing import Indexable as _Indexable
+from typing import Callable, Union, Tuple, List
+from ._typing import Indexable
 from ._jacobi import jacobi
 import numpy as np
 
 
+__all__ = ["propagate"]
+
+
 def propagate(
-    fn: _tp.Callable,
-    x: _tp.Union[float, _Indexable[float]],
-    cov: _tp.Union[float, _Indexable[float], _Indexable[_Indexable[float]]],
+    fn: Callable,
+    x: Union[float, Indexable[float]],
+    cov: Union[float, Indexable[float], Indexable[Indexable[float]]],
     *args,
     **kwargs,
-) -> _tp.Tuple[np.ndarray, np.ndarray]:
+) -> Tuple[np.ndarray, np.ndarray]:
     """
     Numerically propagates the covariance of function inputs to function outputs.
 
@@ -127,36 +130,37 @@ def fn_wrapped(r):
     cov_a = np.asarray(cov)
     y_a = np.asarray(fn(x_a))
 
+    # TODO lift this limitation
     if x_a.ndim > 1:
         raise ValueError("x must have dimension 0 or 1")
 
-    if kwargs.get("diagonal", False):
-        return _propagate_diagonal(fn, y_a, x_a, cov_a, **kwargs)
+    # TODO lift this limitation
+    if y_a.ndim > 1:
+        raise ValueError("f(x) must have dimension 0 or 1")
 
+    # TODO lift this limitation
     if cov_a.ndim > 2:
         raise ValueError("cov must have dimension 0, 1, or 2")
 
-    return _propagate_full(fn, y_a, x_a, cov_a, **kwargs)
-
+    return _propagate(fn, y_a, x_a, cov_a, **kwargs)
 
-def _propagate_full(fn, y: np.ndarray, x: np.ndarray, xcov: np.ndarray, **kwargs):
-    x_a = np.atleast_1d(x)
 
-    _check_x_xcov_compatibility(x_a, xcov)
+def _propagate(fn: Callable, y: np.ndarray, x: np.ndarray, xcov: np.ndarray, **kwargs):
+    _check_x_xcov_compatibility(x, xcov)
 
-    jac = jacobi(fn, x_a, **kwargs)[0]
+    jac = jacobi(fn, x, **kwargs)[0]
 
-    y_len = len(y) if y.ndim == 1 else 1
+    diagonal = kwargs.get("diagonal", False)
 
-    if jac.ndim == 1:
-        jac = jac.reshape((y_len, len(x_a)))
-    assert np.ndim(jac) == 2
+    if jac.ndim == 2:
+        # Check if jacobian is diagonal, count NaN as zero.
+        # This is important to speed up the product below and
+        # to get the right answer for covariance matrices that
+        # contain NaN values.
+        jac = _try_reduce_jacobian(jac)
+    elif not diagonal:
+        jac.shape = (y.size, x.size)
 
-    # Check if jacobian is diagonal, count NaN as zero.
-    # This is important to speed up the product below and
-    # to get the right answer for covariance matrices that
-    # contain NaN values.
-    jac = _try_reduce_jacobian(jac)
     ycov = _jac_cov_product(jac, xcov)
 
     if y.ndim == 0:
@@ -165,52 +169,36 @@ def _propagate_full(fn, y: np.ndarray, x: np.ndarray, xcov: np.ndarray, **kwargs
     return y, ycov
 
 
-def _propagate_diagonal(fn, y: np.ndarray, x: np.ndarray, xcov: np.ndarray, **kwargs):
-    x_a = np.atleast_1d(x)
-
-    _check_x_xcov_compatibility(x_a, xcov)
-
-    jac = jacobi(fn, x_a, **kwargs)[0]
-    assert jac.ndim <= 1
-
-    ycov = _jac_cov_product(jac, xcov)
-
-    if y.ndim == 0:
-        assert ycov.ndim == 0
-
-    return y, ycov
-
-
 def _propagate_independent(
-    fn,
+    fn: Callable,
     y: np.ndarray,
-    x_parts: _tp.List[np.ndarray],
-    xcov_parts: _tp.List[np.ndarray],
+    x_parts: List[np.ndarray],
+    xcov_parts: List[np.ndarray],
     **kwargs,
 ):
-    ycov = 0
+    ycov: Union[float, np.ndarray] = 0
 
     for i, x in enumerate(x_parts):
         rest = x_parts[:i] + x_parts[i + 1 :]
 
-        def wrapped(x, *rest):
-            args = rest[:i] + (x,) + rest[i:]
+        def wrapped(x):
+            args = rest[:i] + [x] + rest[i:]
             return fn(*args)
 
-        x_a = np.atleast_1d(x)
         xcov = xcov_parts[i]
-        _check_x_xcov_compatibility(x_a, xcov)
-
-        jac = jacobi(wrapped, x_a, *rest, **kwargs)[0]
-        ycov += _jac_cov_product(jac, xcov)
 
-    if y.ndim == 0:
-        ycov = np.squeeze(ycov)
+        yc = _propagate(wrapped, y, x, xcov)[1]
+        if np.ndim(ycov) == 2 and yc.ndim == 1:
+            for i, yci in enumerate(yc):
+                ycov[i, i] += yci  # type:ignore
+        else:
+            ycov += yc
 
     return y, ycov
 
 
 def _jac_cov_product(jac: np.ndarray, xcov: np.ndarray):
+    # if jac or xcov are 1D, they represent diagonal matrices
     if xcov.ndim == 2:
         if jac.ndim == 2:
             return np.einsum("ij,kl,jl", jac, jac, xcov)

test/test_jacobi.py

@@ -53,8 +53,11 @@ def fd6(r):
     return r
 
 
+f7_a = np.array([[1, 2, 3], [4, 5, 6]])
+
+
 def f7(x):
-    return np.ones(3) * x**2
+    return f7_a * x**2
 
 
 @pytest.mark.parametrize(
@@ -70,22 +73,33 @@ def f7(x):
         (f6, fd6),
     ],
 )
-def test_jacobi(fn):
+def test_1d(fn):
     x = np.array([1, 2, 3], dtype=float)
     f, fd = fn
     y, ye = jacobi(f, x)
     assert_allclose(y, fd(x))
     assert_allclose(ye, np.zeros_like(y), atol=1e-10)
 
 
-def test_jacobi_0d():
+def test_0d():
     x = 2
     y, ye = jacobi(f7, x)
-    assert np.ndim(y) == 1
-    assert_allclose(y, f7(x))
+    assert np.ndim(y) == 2
+    assert_allclose(y, f7_a * 2 * x)
     assert_allclose(ye, np.zeros_like(y), atol=1e-10)
 
 
+def test_2d():
+    x = [[1, 2, 3], [3, 4, 5]]
+    fd, _ = jacobi(f7, x)
+    assert np.ndim(fd) == 4
+    fd_ref = np.zeros((2, 3, 2, 3))
+    for i in range(2):
+        for j in range(3):
+            fd_ref[i, j, i, j] = f7_a[i, j] * 2 * x[i][j]
+    assert_allclose(fd, fd_ref)
+
+
 def test_abs_at_zero():
     fp, fpe = jacobi(np.abs, 0)
     assert_equal(fp, 0)
@@ -115,9 +129,11 @@ def test_mask():
     mask = np.array([True, False, False, True])
     jac = jacobi(f1, x, 3, mask=mask)[0]
 
-    assert jac.shape == (4, 2)
+    assert jac.shape == (4, 4)
     jac_ref = np.diag(fd1(x, 3))
-    jac_ref = jac_ref[:, mask]
+    for i, mi in enumerate(mask):
+        if not mi:
+            jac_ref[i] = 0
 
     assert_allclose(jac, jac_ref)
 
@@ -170,12 +186,11 @@ def test_maxgrad(maxgrad):
 def test_rtol():
     x = [1, 2, 3]
 
-    jac, jace = jacobi(f1, x, 3)
-
-    jac2, jac2e = jacobi(f1, x, 3, rtol=0.1)
+    jac, jace = jacobi(f1, x, 3, rtol=0.1)
+    jac_ref, jace_ref = jacobi(f1, x, 3)
 
-    assert np.all(jac2e >= jace)
-    assert_allclose(jac, jac2, rtol=0.02)
+    assert np.all(jace >= jace_ref)
+    assert_allclose(jac, jac_ref, rtol=0.03)
 
 
 @pytest.mark.parametrize("maxiter", (0, -1))
@@ -208,7 +223,7 @@ def test_bad_method(method):
         jacobi(f1, 1, 3, method=method)
 
 
-def test_jacobi_on_nan():
+def test_on_nan():
     x = np.array([2.0, np.nan, 3.0])
 
     d = {}