fix for shape inconsistency (#10)

HDembinski · Mar 10, 2022 · 7e625e5 · 7e625e5
1 parent 97caa09
commit 7e625e5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/jacobi/core.py b/src/jacobi/core.py
@@ -266,8 +266,10 @@ def propagate(
     Returns
     -------
     y, ycov
-        y is the result of fn(x)
+        y is the result of fn(x).
         ycov is the propagated covariance matrix.
+        If ycov is a matrix, unless y is a number. In that case, ycov is also
+        reduced to a number.
     """
     x = np.array(x)
     y = fn(x)
@@ -291,4 +293,8 @@ def propagate(
         raise ValueError("x and cov have incompatible shapes")
 
     ycov = np.einsum("il,kl,l" if xcov_nd == 1 else "ij,kl,jl", jac, jac, xcov)
+
+    if np.ndim(y) == 0:
+        ycov = ycov[0, 0]
+
     return y, ycov
diff --git a/test/test_propagate.py b/test/test_propagate.py
@@ -44,6 +44,8 @@ def fn(x):
     assert_allclose(y, fn(x))
     jac = np.ones((1, len(x)))
     assert_allclose(ycov, np.linalg.multi_dot([jac, xcov, jac.T]))
+    assert np.ndim(y) == 0
+    assert np.ndim(ycov) == 0
 
 
 def test_11():