Repository: scipy/scipy
Pull request: DOC: differentiate.jacobian: correct/improve documentation about callable interface
#### Reference issue
Closes gh-22128

#### What does this implement/fix?
The Notes section of [`scipy.differentiate.jacobian`](https://scipy.github.io/devdocs/reference/generated/scipy.differentiate.jacobian.html) included statements like

> argument `f` must be vectorized to accept an array of shape `(m, p)`

but it was not clear that `p` could refer to either a single integer `p` or a tuple of integers `p1, p2`. This PR addresses this by replacing `p` with `...` and documenting that `...` may represent an arbitrary tuple of integers. I'd rather not be more specific than that because I would like to give us some flexibility in the implementation.

To help users satisfy this interface if their function is not vectorized or is vectorized for only one batch dimension, the Notes now provide simple wrappers that the user can copy-paste to adapt their function to the required interface. A documentation improvement is the best we can do before the 1.15.0 release. (And personally, I'd rather let the user take the small steps necessary to translate between their preferred interface and the one required by the function rather than requiring us to accommodate all possibilities efficiently + exhaustively test them.)

#### Additional information
It looks like `hessian` already uses the `...` notation, so I didn't make any changes there. We could copy-paste the wrappers over to those notes if desired, or maybe just refer users of `hessian` to the notes of `jacobian` for the wrappers.

Thanks @andyfaff for noting this before release!

Problem statement:
BUG/DOC: it's not clear how to use `differentiate.jacobian` for a scalar valued function
### Describe your issue.

Could be a bug, but probably documentation.

I'm familiarising myself with `differentiate.jacobian`, and my first step is to differentiate a scalar valued function, i.e. $R^m \rightarrow R^1$. The following example shows my process from trying to follow the notes faithfully.

You can see from the stack trace that fvec ends up being sent a 3-D array, which is incongruent with with what the notes say should happen for evaluating the jacobian at a single point.

I assume that it can probably be done, but if so the documentation definitely needs expanding in this area.

### Reproducing Code Example

```python
import numpy as np
from scipy.differentiate import jacobian

# a scalar valued function R^3 --> R^1
def f(x):
    assert np.size(x) == 3
    return np.sin(x[0]) + 5 * x[1] + np.log10(x[2])
 
v = f(np.array([2, 3, 4.]))
assert (np.isscalar(v))

# Notes say that "argument f must be vectorized to accept an array of shape (m, p)"

def fvec(x):
    # the notes section states "f must be vectorized to accept an array of shape (m, p).
    # The first axis represents the m inputs of f, the second is for evaluating the function
    #  at multiple points in a single call."
    assert x.shape[0] == 3     # m = x.shape[0] = 3
    assert x.ndim == 2         # x.shape hould be (m, p)
    p = x.shape[1]

    # let's call our scalar function. x.T should have length `p`.
    assert len(x.T) == p

    output = np.array([f(l) for l in x.T])
    # at this point output should have length `p`.
    assert len(output) == 3
 
    # documentation states that "argument f must return an array of shape (n, p)"
    # because our function is scalar the array it should return needs to be (1, 3)
    # this means we need to insert an axis at the start
    output = output[None, :]
    assert output.shape == (1, p)
    return output

jacobian(fvec, np.array([2, 3., 4.]))
```


### Error message

```shell
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[88], line 23
     20     assert output.shape == (1, p)
     21     return output
---> 23 jacobian(fvec, np.array([2, 3., 4.]))

File ~/Documents/Andy/programming/scipy/build-install/lib/python3.12/site-packages/scipy/differentiate/_differentiate.py:911, in jacobian(f, x, tolerances, maxiter, order, initial_step, step_factor, step_direction)
    908     xph[i, i] = x
    909     return f(xph)
--> 911 res = derivative(wrapped, x, tolerances=tolerances,
    912                  maxiter=maxiter, order=order, initial_step=initial_step,
    913                  step_factor=step_factor, preserve_shape=True,
    914                  step_direction=step_direction)
    916 del res.x  # the user knows `x`, and the way it gets broadcasted is meaningless here
    917 return res

File ~/Documents/Andy/programming/scipy/build-install/lib/python3.12/site-packages/scipy/differentiate/_differentiate.py:584, in derivative(f, x, args, tolerances, maxiter, order, initial_step, step_factor, step_direction, preserve_shape, callback)
    581 def customize_result(res, shape):
    582     return shape
--> 584 return eim._loop(work, callback, shape, maxiter, func, args, dtype,
    585                  pre_func_eval, post_func_eval, check_termination,
    586                  post_termination_check, customize_result, res_work_pairs,
    587                  xp, preserve_shape)

File ~/Documents/Andy/programming/scipy/build-install/lib/python3.12/site-packages/scipy/_lib/_elementwise_iterative_method.py:243, in _loop(work, callback, shape, maxiter, func, args, dtype, pre_func_eval, post_func_eval, check_termination, post_termination_check, customize_result, res_work_pairs, xp, preserve_shape)
    241 if preserve_shape:
    242     x = xp.reshape(x, (shape + (-1,)))
--> 243 f = func(x, *work.args)
    244 f = xp.asarray(f, dtype=dtype)
    245 if preserve_shape:

File ~/Documents/Andy/programming/scipy/build-install/lib/python3.12/site-packages/scipy/_lib/_elementwise_iterative_method.py:98, in _initialize.<locals>.func(x, shape, func, *args, **kwargs)
     96 def func(x, *args, shape=shape, func=func,  **kwargs):
     97     i = (0,)*(len(fshape) - len(shape))
---> 98     return func(x[i], *args, **kwargs)

File ~/Documents/Andy/programming/scipy/build-install/lib/python3.12/site-packages/scipy/differentiate/_differentiate.py:909, in jacobian.<locals>.wrapped(x)
    907 xph = xp_copy(xp.broadcast_to(xph, new_shape), xp=xp)
    908 xph[i, i] = x
--> 909 return f(xph)

Cell In[88], line 14, in fvec(x)
     12 def fvec(x):
     13     assert x.shape[0] == 3
---> 14     assert x.ndim == 2   # documentation states that x is a 2-D array
     15     p = x.shape[1]
     16     assert len(x.T) == p

AssertionError:
```


### SciPy/NumPy/Python version and system information

```shell
HEAD
```

Hints from issue discussion:
Probably a documentation issue? I called `jacobian` on your original function `f` (without the assertion) and it works. Does that help to pinpoint what needs to change about the documentation? 

```python3
import numpy as np
from scipy.differentiate import jacobian
def f(x):
    return np.sin(x[0]) + 5 * x[1] + np.log10(x[2])
res = jacobian(f, np.array([2, 3., 4.]))
res.df
# array([-0.41614684,  5.        ,  0.10857362])
```
Probably a documentation issue? I called `jacobian` on your original function `f` (without the assertion) and it works. Does that help to pinpoint what needs to change about the documentation?

```python3
import numpy as np
from scipy.differentiate import jacobian
def f(x):
    return np.sin(x[0]) + 5 * x[1] + np.log10(x[2])
res = jacobian(f, np.array([2, 3., 4.]))
res.df
# array([-0.41614684,  5.        ,  0.10857362])
```
I second guessed myself and hid my comment, but I posted it again. It seems to be working, but maybe it's not correct about exactly what gets passed to the callable.

Maybe you already knew it was working, but the documentation about the shapes is incorrect?
Unfortunately I don't think that's sufficient. Let's assume that `f` is a black box, and can *only* accept an array of shape `(m,)`, returning `(1,)`. In that case one will need to use the `fvec` approach. The documentation specifically states that the `x` sent to `fvec` will be of shape `(m, p)`. In contrast `fvec` is being sent a mixture of 2-D and 3-D arrays. So I'm guessing that `fvec` actually needs to be vectorised/iterate over the ellipsis in `(m, ...)`?
Yeah, maybe `p` should have been "`p` *or* `p, q`" to reflect what is currently happening. Seems like we could provide a wrapper in the docs (or the function itself) to make it easier: we just get `shape = x.shape`, pass the user callable `reshape(x, (shape[0], -1))`, and then reshape the output to whatever `jacobian` was expecting?
The following works:

```python
def f(x):
    return np.sin(x[0]) + 5 * x[1] + np.log10(x[2])
 
v = f(np.array([2,3,4]))
assert (np.isscalar(v))

# argument f must be vectorized to accept an array of shape (m, p)

def fvec(x):
    shp = x.shape
    xtemp = np.reshape(x, (shp[0], -1))
        
    output = np.array([f(l) for l in xtemp.T])
    output = np.reshape(output, (1,) + shp[1:])
    return output

jacobian(fvec, np.array([2, 3., 4.]))
```

It would be nice if `jacobian` took care of all the reshaping internally, I don't know how easy that is.

I think the point of confusion arises here because the final paragraph gives the impression that `(m, k, p)` only happens if the original `x` has ndim > 1.
It is easy enough to make it work in `jacobian`. In `wrapped`, we can replace:
```python3
        return f(xph)
```
with
```python3
        # reshape address gh-22128
        batch = xph.shape[1:]  # shape is (m, ...)
        xph = xp.reshape(xph, (m, -1))  # pass shape (m, p) to function
        res = f(xph)  # allow function to return shape (n, p) or just (p,)
        n = (res.shape[0],) if res.ndim > 1 else ()
        return xp.reshape(res, n + batch)  # return shape (n, ...) or just (...)
```
I think this would ensure that the user's callable is always evaluated with an argument of shape `(m, p)` (always 2D) and it is welcome to return an array of shape `p` or `(n, p)`, depending on whether the function inherently has a scalar or vector output.

However, this causes problems in `hessian` because of the admittedly peculiar way we have to calculate function evaluations (because we calculate the Hessian as the Jacobian of the Jacobian).

I don't think the function evaluation count is an easy fix... it might be, but it could also be a time sink (like when I got it working originally). At this point, we might have to take the number of function evaluations out of the result object of the Hessian for this release. Is that a good trade?

Actually, I think it's easy to fix, just kind of inelegant, and it makes the Hessian function counts hard to test.

It might be better to just do a doc fix, because I think the actual shapes make sense and can be described, too. We can provide a complete wrapper that will work for non -vectorized functions, if that would help.
Working on a big project, sure as eggs are eggs we'll get an issue about not having nfev. 

Once one knows what will get sent to `f` it's not that hard to deal with. So I'm 35:65 about it being documentation or a code change. I'll leave the design choice to you.

Is `f` evaluated at `x0` at any point? It might be nice to put those in the rich result as well.
> Is f evaluated at x0 at any point?

Not for central difference (`direction=0`), only for one-sided differences.
Are they recorded in `RichResult` for one-sided differences?

Commit message: DOC: differentiate.jacobian: correct/improve documentation about callable interface

[docs only]

Patch to review:
diff --git a/scipy/differentiate/_differentiate.py b/scipy/differentiate/_differentiate.py
index 2b50b1732dca..a8616635e555 100644
--- a/scipy/differentiate/_differentiate.py
+++ b/scipy/differentiate/_differentiate.py
@@ -802,27 +802,47 @@ def jacobian(f, x, *, tolerances=None, maxiter=10, order=8, initial_step=0.5,
     Notes
     -----
     Suppose we wish to evaluate the Jacobian of a function
-    :math:`f: \mathbf{R}^m \rightarrow \mathbf{R}^n`, and assign to variables
+    :math:`f: \mathbf{R}^m \rightarrow \mathbf{R}^n`. Assign to variables
     ``m`` and ``n`` the positive integer values of :math:`m` and :math:`n`,
-    respectively. If we wish to evaluate the Jacobian at a single point,
-    then:
+    respectively, and let ``...`` represent an arbitrary tuple of integers.
+    If we wish to evaluate the Jacobian at a single point, then:
 
     - argument `x` must be an array of shape ``(m,)``
-    - argument `f` must be vectorized to accept an array of shape ``(m, p)``.
-      The first axis represents the :math:`m` inputs of :math:`f`; the second
-      is for evaluating the function at multiple points in a single call.
-    - argument `f` must return an array of shape ``(n, p)``. The first
-      axis represents the :math:`n` outputs of :math:`f`; the second
-      is for the result of evaluating the function at multiple points.
+    - argument `f` must be vectorized to accept an array of shape ``(m, ...)``.
+      The first axis represents the :math:`m` inputs of :math:`f`; the remainder
+      are for evaluating the function at multiple points in a single call.
+    - argument `f` must return an array of shape ``(n, ...)``. The first
+      axis represents the :math:`n` outputs of :math:`f`; the remainder
+      are for the result of evaluating the function at multiple points.
     - attribute ``df`` of the result object will be an array of shape ``(n, m)``,
       the Jacobian.
 
     This function is also vectorized in the sense that the Jacobian can be
     evaluated at ``k`` points in a single call. In this case, `x` would be an
     array of shape ``(m, k)``, `f` would accept an array of shape
-    ``(m, k, p)`` and return an array of shape ``(n, k, p)``, and the ``df``
+    ``(m, k, ...)`` and return an array of shape ``(n, k, ...)``, and the ``df``
     attribute of the result would have shape ``(n, m, k)``.
 
+`   Suppose the desired callable ``f_not_vectorized`` is not vectorized; it can
+    only accept an array of shape ``(m,)``. A simple solution to satisfy the required
+    interface is to wrap ``f_not_vectorized`` as follows::
+
+        def f(x):
+            return np.apply_along_axis(f_not_vectorized, axis=0, arr=x)
+
+    Alternatively, suppose the desired callable ``f_vec_q`` is vectorized, but
+    only for 2-D arrays of shape ``(m, q)``. To satisfy the required interface,
+    consider::
+
+        def f(x):
+            m, batch = x.shape[0], x.shape[1:]  # x.shape is (m, ...)
+            x = np.reshape(x, (m, -1))  # `-1` is short for q = prod(batch)
+            res = f_vec_q(x)  # pass shape (m, q) to function
+            n = res.shape[0]
+            return np.reshape(res, (n,) + batch)  # return shape (n, ...)
+
+    Then pass the wrapped callable ``f`` as the first argument of `jacobian`.
+
     References
     ----------
     .. [1] Jacobian matrix and determinant, *Wikipedia*,


Code review comments:
diff_hunk: @@ -802,27 +802,47 @@ def jacobian(f, x, *, tolerances=None, maxiter=10, order=8, initial_step=0.5,
     Notes
     -----
     Suppose we wish to evaluate the Jacobian of a function
-    :math:`f: \mathbf{R}^m \rightarrow \mathbf{R}^n`, and assign to variables
+    :math:`f: \mathbf{R}^m \rightarrow \mathbf{R}^n`. Assign to variables
     ``m`` and ``n`` the positive integer values of :math:`m` and :math:`n`,
-    respectively. If we wish to evaluate the Jacobian at a single point,
-    then:
+    respectively, and let ``...`` represent an arbitrary tuple of integers.
+    If we wish to evaluate the Jacobian at a single point, then:
 
     - argument `x` must be an array of shape ``(m,)``
-    - argument `f` must be vectorized to accept an array of shape ``(m, p)``.
-      The first axis represents the :math:`m` inputs of :math:`f`; the second
-      is for evaluating the function at multiple points in a single call.
-    - argument `f` must return an array of shape ``(n, p)``. The first
-      axis represents the :math:`n` outputs of :math:`f`; the second
-      is for the result of evaluating the function at multiple points.
+    - argument `f` must be vectorized to accept an array of shape ``(m, ...)``.
+      The first axis represents the :math:`m` inputs of :math:`f`; the remainder
+      are for evaluating the function at multiple points in a single call.
+    - argument `f` must return an array of shape ``(n, ...)``. The first
+      axis represents the :math:`n` outputs of :math:`f`; the remainder
+      are for the result of evaluating the function at multiple points.
     - attribute ``df`` of the result object will be an array of shape ``(n, m)``,
       the Jacobian.
 
     This function is also vectorized in the sense that the Jacobian can be
     evaluated at ``k`` points in a single call. In this case, `x` would be an
     array of shape ``(m, k)``, `f` would accept an array of shape
-    ``(m, k, p)`` and return an array of shape ``(n, k, p)``, and the ``df``
+    ``(m, k, ...)`` and return an array of shape ``(n, k, ...)``, and the ``df``
     attribute of the result would have shape ``(n, m, k)``.
 
+`   Suppose the desired callable ``f_not_vectorized`` is not vectorized; it can
path: scipy/differentiate/_differentiate.py
text: @author:
```suggestion
    Suppose the desired callable ``f_not_vectorized`` is not vectorized; it can
```