Numerical Jacobian¶

The numerical Jacobian utility estimates gradients and Jacobians with finite differences. It supports first-order two-point differences, second-order three-point differences, and one-sided schemes when variable bounds prevent a central difference.

Use approx_derivative when an analytical or automatic Jacobian is unavailable or when an independent derivative check is useful.

API Reference¶

soromox.utils.numerical_jacobian ¶

Routines for numerical differentiation.

approx_derivative ¶

approx_derivative(fun: Callable[..., Array | float], x0: Array | float, method: str = '3-point', rel_step: Array | float | None = None, abs_step: Array | float | None = None, f0: Array | float | None = None, bounds: tuple[Array | float, Array | float] = (-inf, inf), args: tuple = (), kwargs: dict | None = None) -> Array

Compute a finite-difference approximation of a function's Jacobian.

If a function maps from R^n to R^m, the Jacobian has shape (m, n) and contains the partial derivative of output i with respect to input j at entry (i, j).

Parameters:

Name	Type	Description	Default
`fun`	`Callable[..., Array \| float]`	Function whose derivatives should be estimated. It receives a one-dimensional array of shape `(n,)` and returns a scalar or a one-dimensional array of shape `(m,)`.	required
`x0`	`Array \| float`	Point at which to estimate the derivatives. Scalars are converted to a one-dimensional array.	required
`method`	`str`	Finite-difference method. `"2-point"` uses a first-order forward or backward difference. `"3-point"` uses a central difference in the interior and a second-order one-sided difference near bounds.	`'3-point'`
`rel_step`	`Array \| float \| None`	Relative step size. If `None`, an appropriate step is selected from the input and output dtypes.	`None`
`abs_step`	`Array \| float \| None`	Absolute step size. If provided, it takes precedence over `rel_step` and is adjusted to fit within `bounds`.	`None`
`f0`	`Array \| float \| None`	Optional value of `fun(x0)`. Providing it avoids evaluating the function at `x0`.	`None`
`bounds`	`tuple[Array \| float, Array \| float]`	Lower and upper bounds on the independent variables. Each bound may be a scalar or an array matching `x0`. Defaults to no bounds.	`(-inf, inf)`
`args`	`tuple`	Additional positional arguments passed to `fun`.	`()`
`kwargs`	`dict \| None`	Additional keyword arguments passed to `fun`.	`None`

Returns:

Type	Description
`Array`	A dense Jacobian. For a scalar-valued function the result has shape
`Array`	`(n,)`; otherwise it has shape `(m, n)`.

Raises:

Type	Description
`ValueError`	If `method` is unsupported, `x0` is not one dimensional, the bounds have incompatible shapes, or `x0` lies outside the bounds.

Notes

If a requested step is indistinguishable from x0 in the working dtype, an automatic step is substituted for that entry. For the "3-point" method, central differences are used when both sides fit within the bounds; otherwise a second-order one-sided difference is used.

Scalar-valued functions return a one-dimensional gradient for consistency with common autodiff APIs. Use jnp.atleast_2d when a two-dimensional Jacobian is required in all cases.

References

W. H. Press et al., Numerical Recipes: The Art of Scientific Computing, 3^rd edition, section 5.7.
B. Fornberg, "Generation of Finite Difference Formulas on Arbitrarily Spaced Grids," Mathematics of Computation, 51, 1988.

Examples:

import jax.numpy as jnp

from soromox.utils.numerical_jacobian import approx_derivative

def f(x, c1, c2):
    return jnp.array(
        [
            x[0] * jnp.sin(c1 * x[1]),
            x[0] * jnp.cos(c2 * x[1]),
        ]
    )

x0 = jnp.array([1.0, 0.5 * jnp.pi])
jacobian = approx_derivative(f, x0, args=(1, 2))
# jacobian is approximately [[1, 0], [-1, 0]]

Bounds can limit the region of function evaluation:

def g(x):
    return x**2 if x >= 1 else x

x0 = 1.0
left_derivative = approx_derivative(g, x0, bounds=(-jnp.inf, 1.0))
right_derivative = approx_derivative(g, x0, bounds=(1.0, jnp.inf))
# left_derivative is [1] and right_derivative is [2]