perform correct error propagation if jacobian is diagonal (#26)

src/jacobi/_jacobi.py

@@ -154,7 +154,7 @@ def jacobi(
             # more like student's t
             rei = c[-1, -1] ** 0.5
 
-            # update estimates that have significantly smaller error
+            # update estimates that have smaller estimated error
             sub_todo = rei < re[todo]
             todo1 = todo.copy()
             todo[todo1] = sub_todo

src/jacobi/_propagate.py

@@ -47,6 +47,15 @@ def propagate(
         If ycov is a matrix, unless y is a number. In that case, ycov is also
         reduced to a number.
 
+    Notes
+    -----
+    For callables `fn` which perform only element-wise computation, the jacobian is
+    a diagonal matrix. This special case is detected and the computation optimised,
+    although can further speed up the computation by passing the argumet `diagonal=True`.
+
+    In this special case, error propagation works correctly even if the output of `fn`
+    is NaN for some inputs.
+
     Examples
     --------
     General error propagation maps input vectors to output vectors::
@@ -135,14 +144,19 @@ def _propagate_full(fn, y: np.ndarray, x: np.ndarray, xcov: np.ndarray, **kwargs
 
     _check_x_xcov_compatibility(x_a, xcov)
 
-    jac = np.asarray(jacobi(fn, x_a, **kwargs)[0])
+    jac = jacobi(fn, x_a, **kwargs)[0]
 
     y_len = len(y) if y.ndim == 1 else 1
 
     if jac.ndim == 1:
         jac = jac.reshape((y_len, len(x_a)))
     assert np.ndim(jac) == 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)
     ycov = _jac_cov_product(jac, xcov)
 
     if y.ndim == 0:
@@ -156,7 +170,7 @@ def _propagate_diagonal(fn, y: np.ndarray, x: np.ndarray, xcov: np.ndarray, **kw
 
     _check_x_xcov_compatibility(x_a, xcov)
 
-    jac = np.asarray(jacobi(fn, x_a, **kwargs)[0])
+    jac = jacobi(fn, x_a, **kwargs)[0]
     assert jac.ndim <= 1
 
     ycov = _jac_cov_product(jac, xcov)
@@ -187,7 +201,7 @@ def wrapped(x, *rest):
         xcov = xcov_parts[i]
         _check_x_xcov_compatibility(x_a, xcov)
 
-        jac = np.asarray(jacobi(wrapped, x_a, *rest, **kwargs)[0])
+        jac = jacobi(wrapped, x_a, *rest, **kwargs)[0]
         ycov += _jac_cov_product(jac, xcov)
 
     if y.ndim == 0:
@@ -198,12 +212,16 @@ def wrapped(x, *rest):
 
 def _jac_cov_product(jac: np.ndarray, xcov: np.ndarray):
     if xcov.ndim == 2:
-        return np.einsum(
-            "i,j,ij -> ij" if jac.ndim == 1 else "ij,kl,jl", jac, jac, xcov
-        )
-    elif jac.ndim == 2:
+        if jac.ndim == 2:
+            return np.einsum("ij,kl,jl", jac, jac, xcov)
+        if jac.ndim == 1:
+            return np.einsum("i,j,ij -> ij", jac, jac, xcov)
+        return jac**2 * xcov
+    assert xcov.ndim < 2
+    if jac.ndim == 2:
         if xcov.ndim == 1:
             return np.einsum("ij,kj,j", jac, jac, xcov)
+        assert xcov.ndim == 0  # xcov.ndim == 2 is already covered above
         return np.einsum("ij,kj", jac, jac) * xcov
     assert jac.ndim < 2 and xcov.ndim < 2
     return xcov * jac**2
@@ -213,3 +231,23 @@ def _check_x_xcov_compatibility(x: np.ndarray, xcov: np.ndarray):
     if xcov.ndim > 0 and len(xcov) != (len(x) if x.ndim == 1 else 1):
         # this works for 1D and 2D xcov
         raise ValueError("x and cov have incompatible shapes")
+
+
+def _try_reduce_jacobian(jac: np.ndarray):
+    if jac.ndim != 2 or jac.shape[0] != jac.shape[1]:
+        return jac
+    # if jacobian contains only off-diagonal elements
+    # that are zero or NaN, we reduce it to diagonal form
+    ndv = _nodiag_view(jac)
+    m = np.isnan(ndv)
+    ndv[m] = 0
+    if np.count_nonzero(ndv) == 0:
+        return np.diag(jac)
+    return jac
+
+
+def _nodiag_view(a: np.ndarray):
+    # https://stackoverflow.com/a/43761941/ @Divakar
+    m = a.shape[0]
+    p, q = a.strides
+    return np.lib.stride_tricks.as_strided(a[:, 1:], (m - 1, m), (p + q, q))

test/test_propagate.py

@@ -78,20 +78,23 @@ def fn(x):
 
 
 @pytest.mark.parametrize("ndim", (1, 2))
-def test_cov_1d_2d(ndim):
+@pytest.mark.parametrize("diagonal", (False, True))
+@pytest.mark.parametrize("len", (1, 2))
+def test_cov_1d_2d(ndim, diagonal, len):
     def fn(x):
-        return x
+        return 2 * x
 
-    x = [1, 2]
-    xcov_1d = [3, 4]
-    xcov_2d = np.diag(xcov_1d)
+    x = [1, 2][:len]
+    xcov = [3, 4][:len]
+    if ndim == 2:
+        xcov = np.diag(xcov)
 
-    y, ycov = propagate(fn, x, xcov_1d if ndim == 1 else xcov_2d)
+    y, ycov = propagate(fn, x, xcov, diagonal=diagonal)
 
-    assert np.ndim(ycov) == 2
+    assert np.ndim(ycov) == ndim
 
-    assert_allclose(y, x)
-    assert_allclose(ycov, xcov_2d)
+    assert_allclose(y, np.multiply(x, 2))
+    assert_allclose(ycov, np.multiply(xcov, 4))
 
 
 def test_two_arguments_1():
@@ -224,11 +227,6 @@ def fn(x):
 
     y, ycov = propagate(fn, x, xcov, diagonal=True)
 
-    # Beware: this produces a matrix with all NaNs
-    #   y_ref, ycov_ref = propagate(fn, x, xcov)
-    # The derivative cannot figure out that the off-diagonal elements
-    # of the jacobian are zero.
-
     y_ref = [2, np.nan, 5]
     assert_allclose(y, y_ref)
 
@@ -239,3 +237,10 @@ def fn(x):
     # ycov_ref = jac @ np.array(xcov) @ jac.T
     ycov_ref = [[12, np.nan, 8], [np.nan, np.nan, np.nan], [8, np.nan, 80]]
     assert_allclose(ycov, ycov_ref)
+
+    # propagate now detects the special case where jac is effectively diagonal
+    # and does the equivalent of propagate(fn, x, xcov, diagonal=True), which
+    # is nevertheless faster
+    y2, ycov2 = propagate(fn, x, xcov)
+    assert_allclose(y2, y_ref)
+    assert_allclose(ycov2, ycov_ref)