diff --git a/src/jacobi/_jacobi.py b/src/jacobi/_jacobi.py index 2658bc0..16e26e4 100644 --- a/src/jacobi/_jacobi.py +++ b/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 diff --git a/src/jacobi/_propagate.py b/src/jacobi/_propagate.py index e72cc71..fa8f5fd 100644 --- a/src/jacobi/_propagate.py +++ b/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,7 +127,12 @@ 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: @@ -136,7 +140,7 @@ def fn_wrapped(r): # 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: diff --git a/test/test_jacobi.py b/test/test_jacobi.py index a13e87c..52d44b9 100644 --- a/test/test_jacobi.py +++ b/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) diff --git a/test/test_propagate.py b/test/test_propagate.py index 620014e..79336a6 100644 --- a/test/test_propagate.py +++ b/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)