fix propagate for nan input (#25)

setup.cfg

@@ -34,7 +34,6 @@ install_requires = numpy >= 1.10
 test =
     pytest
     pytest-benchmark
-    pre-commit
     types-setuptools
 doc =
     sphinx

src/jacobi/_jacobi.py

@@ -85,6 +85,8 @@ def jacobi(
             raise ValueError("step[0] must be between 0 and 0.5")
         if not (0 < step[1] < 1):
             raise ValueError("step[1] must be between 0 and 1")
+    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)
@@ -98,9 +100,7 @@ def jacobi(
     if diagnostic is not None:
         diagnostic["method"] = np.zeros(nx, dtype=np.int8)
         diagnostic["iteration"] = np.zeros(len(x_indices), dtype=np.uint8)
-
-    if method is not None and method not in (-1, 0, 1):
-        raise ValueError("method must be -1, 0, 1")
+        diagnostic["residual"] = [[] for _ in range(nx)]
 
     f0 = None
     jac = None
@@ -111,9 +111,6 @@ def jacobi(
         # if method is None, auto-detect for each x[k]
         md, f0, r = _first(method, f0, fn, x, k, h[0], args)
 
-        if diagnostic is not None:
-            diagnostic["method"][ik] = md
-
         if md != 0 and step is None:
             # optimal step sizes for forward derivatives
             h = _steps(x[k], (0.125, 0.125), maxiter)
@@ -125,13 +122,14 @@ def jacobi(
         todo = np.ones(nr, dtype=bool)
         fd = [r]
 
-        if jac is None:
+        if jac is None:  # first iteration
             jac = np.empty(r_shape + (nx,), dtype=r.dtype)
             err = np.empty(r_shape + (nx,), dtype=r.dtype)
             if diagnostic is not None:
                 diagnostic["call"] = np.zeros((nr, nx), dtype=np.uint8)
 
         if diagnostic is not None:
+            diagnostic["method"][ik] = md
             diagnostic["call"][:, ik] = 2 if md == 0 else 3
 
         for i in range(1, len(h)):
@@ -168,6 +166,11 @@ def jacobi(
                 sub_todo &= rei > rtol * np.abs(ri)
                 todo[todo1] = sub_todo
 
+            if diagnostic is not None:
+                re2 = re.copy()
+                re2[todo1] = rei
+                diagnostic["residual"][ik].append(re2)
+
             if np.sum(todo) == 0:
                 break
 
@@ -192,7 +195,7 @@ def jacobi(
 def _steps(p, step, maxiter):
     h0, factor = step
     h = p * h0
-    if h == 0:
+    if not h != 0:  # also works if p is NaN
         h = h0
     return h * factor ** np.arange(maxiter)
 

test/test_jacobi.py

@@ -206,3 +206,20 @@ def test_bad_step_1(step):
 def test_bad_method(method):
     with pytest.raises(ValueError, match="method"):
         jacobi(f1, 1, method=method)
+
+
+def test_jacobi_on_nan():
+    x = np.array([2.0, np.nan, 3.0])
+
+    d = {}
+    yd, yde = jacobi(f1, x, diagnostic=d)
+
+    assert d["iteration"][1] == 1
+    assert_allclose(
+        yd, [[4.0, np.nan, 0.0], [np.nan, np.nan, np.nan], [0.0, np.nan, 6.0]]
+    )
+    assert_allclose(
+        yde,
+        [[0.0, np.inf, 0.0], [np.inf, np.inf, np.inf], [0.0, np.inf, 0.0]],
+        atol=1e-15,
+    )

test/test_propagate.py

@@ -1,6 +1,6 @@
 import numpy as np
 from numpy.testing import assert_allclose
-from jacobi import propagate
+from jacobi import propagate, jacobi
 import pytest
 
 
@@ -213,3 +213,29 @@ def fn(x):
 
     assert_allclose(y, y_ref)
     assert_allclose(ycov, ycov_ref)
+
+
+def test_on_nan():
+    def fn(x):
+        return x**2 + 1
+
+    x = [1.0, np.nan, 2.0]
+    xcov = [[3.0, np.nan, 1.0], [np.nan, np.nan, np.nan], [1.0, np.nan, 5.0]]
+
+    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)
+
+    jac = jacobi(fn, x, diagonal=True)[0]
+    jac_ref = [2.0, np.nan, 4.0]
+    assert_allclose(jac, jac_ref)
+
+    # 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)