accept functions which return sequences and not arrays (#42)

Make jacobi and propagate work with functions that return sequences
instead of numpy arrays. Produce better error messages if the return
values are incorrect.

src/jacobi/_jacobi.py

@@ -22,8 +22,11 @@ def jacobi(
     Parameters
     ----------
     fn : Callable
-        Function with the signature `fn(x, *args)`, where `x` is a number or an
-        array of numbers and `*args` are optional auxiliary arguments.
+        Function with the signature `fn(x, *args)`, where `x` is a number or a sequence
+        of numbers and `*args` are optional auxiliary arguments. The function must
+        return a number or a regular shape of numbers (ideally as a numpy array). The
+        length of `x` can differ from the output sequence. Derivatives are only
+        computed with respect to `x`, the auxiliary arguments are ignored.
     x : Number or array of numbers
         The derivative is computed with respect to `x`. If `x` is an array, the Jacobi
         matrix is computed with respect to each element of `x`.
@@ -120,7 +123,7 @@ def jacobi(
         # if step is None, use optimal step sizes for central derivatives
         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, kx, h[0], args)
+        fn, 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:
@@ -224,7 +227,10 @@ def _derive(mode, f0, f, x, i, h, args):
     return (-3 * f0 + 4 * f1 - f2) * (0.5 / h)
 
 
-def _first(method, f0, f, x, i, h, args):
+def _first(method, f0, fn, x, i, h, args):
+    # This is the first derivative that we calculate.
+    # This function is special because we collect a lot of diagnostic
+    # information about the function for the remainder of the iterations.
     norm = 0.5 / h
     f1 = None
     f2 = None
@@ -233,8 +239,9 @@ def _first(method, f0, f, x, i, h, args):
         x2 = x.copy()
         x1[i] -= h
         x2[i] += h
-        f1 = f(x1, *args)
-        f2 = f(x2, *args)
+        f1 = fn(x1, *args)
+        fn, f1 = _wrap_function_if_needed(fn, f1)
+        f2 = fn(x2, *args)
         if method is None:
             if np.any(np.isnan(f1)):  # forward method
                 method = 1
@@ -243,19 +250,33 @@ def _first(method, f0, f, x, i, h, args):
             else:
                 method = 0
     if method == 0:
-        return method, None, (f1 - f2) * norm
+        return fn, method, None, (f1 - f2) * norm
     if f0 is None:
-        f0 = f(x, *args)
+        f0 = fn(x, *args)
+        fn, f0 = _wrap_function_if_needed(fn, f0)
     if method == -1:
         h = -h
         norm = -norm
     if f1 is None:
         x1 = x.copy()
         x1[i] += h
-        f1 = f(x1, *args)
+        f1 = fn(x1, *args)
     elif method == 1:
         f1 = f2
     x2 = x.copy()
     x2[i] += 2 * h
-    f2 = f(x2, *args)
-    return method, f0, (-3 * f0 + 4 * f1 - f2) * norm
+    f2 = fn(x2, *args)
+    return fn, method, f0, (-3 * f0 + 4 * f1 - f2) * norm
+
+
+def _wrap_function_if_needed(fn, fval):
+    if not isinstance(fval, float):
+        try:
+            fval = np.asarray(fval, dtype=float)
+        except ValueError as e:
+            raise ValueError(
+                "function return value cannot be converted into "
+                "1D numpy array of floats"
+            ) from e
+        return lambda *args: np.asarray(fn(*args)), fval
+    return fn, fval

src/jacobi/_propagate.py

@@ -24,12 +24,11 @@ def propagate(
     Parameters
     ----------
     fn: callable
-        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.
+        Function with the signature `fn(x, *args)`, where `x` is a number or a sequence
+        of numbers and `*args` are optional auxiliary arguments. The function must
+        return a number or a sequence of numbers (ideally as a numpy array). The
+        length of `x` can differ from the output sequence. Error propagation is only
+        performed with respect to `x`, the auxiliary arguments are ignored.
     x: float or array-like with shape (N,)
         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)
@@ -128,15 +127,20 @@ def fn_wrapped(r):
 
     x_a = np.asarray(x)
     cov_a = np.asarray(cov)
-    y_a = np.asarray(fn(x_a))
+    try:
+        y_a = np.asarray(fn(x_a))
+    except ValueError as e:
+        raise ValueError(
+            "function return value cannot be converted into numpy array"
+        ) from e
 
     # TODO lift this limitation
     if x_a.ndim > 1:
         raise ValueError("x must have dimension 0 or 1")
 
     # TODO lift this limitation
     if y_a.ndim > 1:
-        raise ValueError("f(x) must have dimension 0 or 1")
+        raise ValueError("function return value must have dimension 0 or 1")
 
     # TODO lift this limitation
     if cov_a.ndim > 2:

test/test_jacobi.py

@@ -238,3 +238,22 @@ def test_on_nan():
         [[0.0, np.inf, 0.0], [np.inf, np.inf, np.inf], [0.0, np.inf, 0.0]],
         atol=1e-15,
     )
+
+
+@pytest.mark.parametrize("method", (None, -1, 0, 1))
+@pytest.mark.parametrize("fn", (lambda x: (x[0], 2 * x[1], x[1]), lambda x: x[0]))
+def test_non_array_arguments_and_return_value(method, fn):
+    def fn(x):
+        return [x[0], 2 * x[1], x[1] ** 2]
+
+    j, _ = jacobi(fn, (1, 2), method=method)
+    assert_allclose(j, [(1, 0), (0, 2), (0, 4)])
+
+
+@pytest.mark.parametrize("method", (None, -1, 0, 1))
+@pytest.mark.parametrize(
+    "fn", (lambda x: [1, [1, 2]], lambda x: "s", lambda x: ("a", "b"))
+)
+def test_bad_return_value_2(method, fn):
+    with pytest.raises(ValueError, match="function return value cannot be converted"):
+        jacobi(fn, (1, 2), method=method)

test/test_propagate.py

@@ -307,3 +307,29 @@ def f(a, b):
     c2, c2_var = propagate(f, a, a_var, b, b_var)
 
     assert np.sum(np.diag(c2_var) > np.diag(c1_var)) > 0
+
+
+@pytest.mark.parametrize("method", (None, -1, 0, 1))
+@pytest.mark.parametrize("fn", (lambda x: (x[0], 2 * x[1], x[1]), lambda x: x[0]))
+def test_non_array_arguments_and_return_value(method, fn):
+    def fn(x):
+        return [x[0], 2 * x[1], x[1] ** 3]
+
+    x = (1, 2)
+    xcov = ((1, 0), (0, 2))
+    y, ycov = propagate(fn, x, xcov, method=method)
+
+    j = np.array([[1, 0], [0, 2], [0, 3 * x[1] ** 2]])
+    ycov_ref = j @ xcov @ j.T
+
+    assert_allclose(y, [1, 4, 8])
+    assert_allclose(ycov, ycov_ref)
+
+
+@pytest.mark.parametrize("method", (None, -1, 0, 1))
+@pytest.mark.parametrize(
+    "fn", (lambda x: [1, [1, 2]], lambda x: "s", lambda x: ("a", "b"))
+)
+def test_bad_return_value_2(method, fn):
+    with pytest.raises(ValueError, match="function return value cannot be converted"):
+        propagate(fn, (1, 2), ((1, 0), (0, 1)), method=method)