Utilities¶

def odeint(
    drift: Callable[..., PyTree[Array]],
    y0: PyTree[Array],
    ts: Array,
    *args: Any,
    solver: Union[str, type, ConfiguredSolver] = "rk4",
    method: Optional[str] = None,
    dtype: Optional[jnp.dtype] = jnp.float32,
    filter_state: Optional[Callable[[PyTree[Array]], Optional[PyTree[Array]]]] = None,
    collect_trace: bool = True,
    check_points: Optional[Sequence[int]] = None,
    step_size_adaptor: Optional[StepSizeAdaptor] = None,
    logdet_rng: Optional[Array] = None,
    trace_estimator: TraceEstimator = "exact",
    num_samples: int = 1,
    sample_dist: SampleDist = "rademacher",
) -> Optional[PyTree[Array]]:
    """Solve an ordinary differential equation ``dy/dt = drift(t, y, *args)``.

    ``drift`` may be any of:

    - a plain Python callable ``drift(t, y, *args)`` — it is automatically
      wrapped in :class:`probjax.utils.functions.generic_drift` so it rides
      as a pytree through ``jax.jit`` and the ``custom_inverse`` primitive;
    - a registered JAX pytree node (``eqx.Module``, ``flax.struct.PyTreeNode``,
      :class:`~probjax.utils.functions.Drift` subclass, ...) — its array
      leaves participate in transformations natively, non-array fields ride
      as aux.

    Drift keyword arguments are no longer supported. Pass parameters
    positionally via ``*args``, or bind them up-front with ``functools.partial``
    (or ``drift.bind_args(...)`` on a :class:`~probjax.utils.functions.Drift`).

    Args:
        drift: The drift function ``f(t, y, *args)``.
        y0: Initial state. Single array or pytree of arrays.
        ts: Time points at which to evaluate the solution.
        *args: Positional arguments forwarded to ``drift``.
        solver: Integration method: a registered name (``"rk4"``,
            ``"dopri5"``, ``"bdf2"``, ...), a solver class, or a configured
            solver such as ``bdf2(newton_iters=3, jacobian="matrix_free")``.
        method: Deprecated alias of ``solver``.
        dtype: Computation dtype (default ``float32``).
        filter_state: Optional function to filter the state during integration.
        collect_trace: If ``True`` (default) return the filtered state at
            every time point; otherwise return only the filtered terminal state.
        check_points: Optional index sequence for checkpointed grid integration.
        step_size_adaptor: Step-size controller for **adaptive** solver
            methods (``"dopri5"``, ``"dopri8"``, ``"bosh3"``, …). Pass an
            instance of :class:`~probjax.utils.odeutil.adaptive.StepSizeAdaptor`
            (or a subclass) to tune ``rtol`` / ``atol`` / clip bounds /
            controller behavior. Ignored entirely by fixed-step methods
            (``"rk4"``, ``"euler"``, ``"midpoint"``, …). Defaults to
            :class:`StepSizeAdaptor()` with the standard Hairer–Wanner
            controller.
        logdet_rng: RNG key for stochastic log-determinant estimators.
            Required when ``trace_estimator="hutchinson"``. Ignored on the
            forward path; only consumed by
            :func:`probjax.core.inverse_and_logabsdet`.
        trace_estimator: Log-det trace estimator used when the function is
            inverted via ``inverse_and_logabsdet``. One of:

            - ``"exact"`` (default): full Jacobian per step (O(d²) cost).
            - ``"hutchinson"``: FFJORD-style stochastic estimator
              ``tr(J) ≈ mean_k vᵀ_k J v_k`` via one JVP per probe vector;
              probe vectors are fixed across the trajectory so
              ``∫tr(J)dt`` stays unbiased. Requires ``logdet_rng``.
            - a callable ``trace_fn(drift_flat, t, x_flat, args) -> scalar``
              for custom structured-Jacobian strategies.
        num_samples: Hutchinson probe-vector count per trajectory. Higher
            values reduce variance linearly in cost.
        sample_dist: Hutchinson probe distribution — ``"rademacher"``
            (default, minimum-variance for general matrices) or
            ``"normal"``.

    Returns:
        Pytree containing either the time-series trace (when
        ``collect_trace=True``) or the filtered terminal state.

    Example:
        >>> import jax.numpy as jnp
        >>> from probjax.utils.odeint import odeint
        >>>
        >>> def lotka_volterra(t, y, alpha, beta, delta, gamma):
        ...     prey, predator = y
        ...     dprey = alpha * prey - beta * prey * predator
        ...     dpredator = delta * prey * predator - gamma * predator
        ...     return jnp.array([dprey, dpredator])
        >>>
        >>> y0 = jnp.array([40.0, 9.0])
        >>> ts = jnp.linspace(0, 10, 100)
        >>> ys = odeint(lotka_volterra, y0, ts, 1.0, 0.1, 0.075, 0.5, solver="dopri5")
    """
    solver = solver_or_legacy_method(solver, method, "odeint", default="rk4")
    drift = wrap_if_plain_callable(drift)
    return _odeint_custom(
        y0,
        drift,
        ts,
        tuple(args),
        logdet_rng,
        method=solver,
        dtype=dtype,
        filter_state=filter_state,
        collect_trace=collect_trace,
        check_points=check_points,
        step_size_adaptor=step_size_adaptor,
        trace_estimator=trace_estimator,
        num_samples=num_samples,
        sample_dist=sample_dist,
    )

def sdeint(
    rng: Key,
    drift: Callable[..., PyTree[Array]],
    diffusion: Callable[..., PyTree[Array]],
    y0: PyTree[Array],
    ts: Array,
    *args: Any,
    solver: Union[str, type, ConfiguredSolver] = "euler_maruyama",
    method: Optional[str] = None,
    dtype: Optional[jnp.dtype] = jnp.float32,
    sde_type: str = "ito",
    return_brownian: bool = False,
    return_state: bool = False,
    filter_state: Optional[Callable[[PyTree[Array]], Optional[PyTree[Array]]]] = None,
    collect_trace: bool = True,
    check_points: Optional[Sequence[int]] = None,
    step_size_adaptor: Optional[SDEStepSizeAdaptor] = None,
) -> Union[
    Optional[PyTree[Array]],
    Tuple[Any, Optional[PyTree[Array]]],
    Tuple[
        Optional[PyTree[Array]],
        Optional[PyTree[Array]],
    ],
    Tuple[
        Any,
        Tuple[Optional[PyTree[Array]], Optional[PyTree[Array]]],
    ],
]:
    """Solve a stochastic differential equation
    ``dy = drift(t, y, *args) dt + diffusion(t, y, *args) dW``.

    ``drift`` and ``diffusion`` may be plain Python callables (automatically
    wrapped as :class:`~probjax.utils.functions.generic_drift`) or any
    pytree-registered callable (marker
    :class:`~probjax.utils.functions.Drift` subclass, ``eqx.Module``, etc.).

    Keyword arguments to drift/diffusion are no longer supported. Pass
    parameters positionally via ``*args``, or bind them with
    ``functools.partial`` / ``drift.bind_args(...)``.

    Args:
        rng: Random number generator key.
        drift: Drift function ``f(t, y, *args)`` — deterministic part of
            ``dy = f(t, y, *args) dt + g(t, y, *args) dW_t``.
        diffusion: Diffusion function ``g(t, y, *args)`` — stochastic part.
        y0: Initial state. Single array or pytree of arrays.
        ts: Strictly increasing 1D array of time points.
        *args: Positional arguments forwarded to both drift and diffusion.
        solver: Integration method: a registered name, a solver class or a
            configured solver (``milstein(...)``). Registered names include:

            - ``"euler_maruyama"``: Euler-Maruyama (order 0.5)
            - ``"exp_euler_maruyama"``: Exponential Euler-Maruyama
              (requires :class:`~probjax.utils.functions.split_drift`)
            - ``"milstein"``: Milstein (order 1.0)
            - ``"sri1"``, ``"sri2"``: Stochastic Runge-Kutta methods
              (strong order 1.0); ``"srk"`` is an alias of ``"sri1"``
            - ``"linear_exact_sde"``: exact transition sampling for linear SDEs
              (requires :class:`~probjax.utils.functions.linear_drift` and
              :class:`~probjax.utils.functions.const_diffusion`)
        method: Deprecated alias of ``solver``.
        dtype: Computation dtype (default ``float32``).
        sde_type: ``"ito"`` (default) or ``"stratonovich"``. Noise layout
            is inferred from the diffusion output shape — scalar/vector
            outputs are diagonal, matrix outputs are full (including
            rectangular).
        return_brownian: Whether to return Brownian paths (requires
            ``collect_trace=True``).
        return_state: Whether to return solver state.
        filter_state: Optional state filter; returning ``None`` disables
            tracing entirely.
        collect_trace: Record the filtered quantity at every time step
            (``True``, default) or return only the filtered terminal state
            (``False``). Must be ``True`` when returning Brownian paths.
        check_points: Optional index sequence for checkpointed grid
            integration.
        step_size_adaptor: When provided, switches integration to an
            adaptive step-doubling Euler-Maruyama scheme driven by the
            given controller. Pass
            :class:`~probjax.utils.sdeutil.adaptive.WeakStepSizeAdaptor`
            for distributional / weak quantities (fresh noise from the key)
            or :class:`~probjax.utils.sdeutil.adaptive.StrongStepSizeAdaptor`
            for per-path consistency (increments from a single virtual
            Brownian tree, so runs at different tolerances are comparable
            per path). Steps are never rejected on their own noise; the
            local error only sets the next step size. Currently restricted to
            **diagonal noise** and ignores ``solver`` (always uses Euler-Maruyama with
            step-doubling error estimation). Defaults to ``None``
            (fixed-step integration via ``solver``).

    Returns:
        When ``return_brownian=False``: the filtered trajectory (if
        ``collect_trace=True``) or the filtered terminal state. When
        ``return_brownian=True``: ``(state_trace, brownian_trace)``. If
        ``return_state=True``, the solver state is prepended.

    Example:
        >>> import jax.numpy as jnp
        >>> from probjax.utils.sdeint import sdeint
        >>> from jax import random
        >>>
        >>> def drift(t, state, mu, sigma):
        ...     return {"price": mu * state["price"]}
        >>> def diffusion(t, state, mu, sigma):
        ...     return {"price": sigma * state["price"]}
        >>>
        >>> rng = random.PRNGKey(0)
        >>> y0 = {"price": jnp.array([1.0])}
        >>> ts = jnp.linspace(0, 1, 100)
        >>> ys = sdeint(rng, drift, diffusion, y0, ts, 0.1, 0.2)
    """
    solver = solver_or_legacy_method(solver, method, "sdeint", default="euler_maruyama")
    drift = wrap_if_plain_callable(drift)
    diffusion = wrap_if_plain_callable(diffusion)
    result = _sdeint(
        rng,
        drift,
        diffusion,
        y0,
        ts,
        tuple(args),
        method=solver,
        dtype=dtype,
        sde_type=sde_type,
        return_brownian=return_brownian,
        return_state=return_state,
        filter_state=filter_state,
        collect_trace=collect_trace,
        check_points=check_points,
        step_size_adaptor=step_size_adaptor,
    )
    # ``_sdeint`` appends an adaptive-diagnostic hit count as the last
    # element of its return so we can warn host-side (no per-vmap-element
    # callback overhead). Strip it before returning to the user; warn iff
    # the controller ran out of budget on a meaningful fraction of the
    # output segments.
    if return_state:
        state_obj, payload, diag_hits = result
        out = (state_obj, payload)
    else:
        payload, diag_hits = result
        out = payload
    # Host-side boundary warning. Skip when a tracer flows through (under
    # ``jax.vmap`` / ``jax.jit`` of ``sdeint`` itself); the warning will
    # surface from the outermost concrete invocation. ``jax.core.Tracer``
    # check costs nothing under vmap and avoids the per-element callback
    # dispatch we'd pay if the warn lived inside the JIT graph.
    if step_size_adaptor is not None and not is_tracing(diag_hits):
        n_segments = int(jnp.atleast_1d(jnp.asarray(ts)).shape[0]) - 1
        warn_boundary_hits(step_size_adaptor, diag_hits, n_segments)
    return out

def betaincinv(a, b, p, *, max_halley_steps=6, max_bisection_steps=15):
    """Invert regularized lower beta for finite positive shapes.

    Probabilities are clipped to [0, 1]; invalid shapes return NaN. The total
    safeguarded Halley/bisection budget is 2 plus both step settings.
    Reverse AD uses the inverse density and central differences for shapes;
    forward AD and higher shape derivatives are not supported.
    """
    return _inverse(a, b, p, max_halley_steps, max_bisection_steps, False)

def betainccinv(a, b, q, *, max_halley_steps=6, max_bisection_steps=15):
    """Invert regularized upper beta directly, preserving tiny q.

    q=0 returns one; q=1 returns zero. Domain, step settings and differentiation
    have the same semantics as betaincinv.
    """
    return _inverse(a, b, q, max_halley_steps, max_bisection_steps, True)

@jax.jit
def gammaincinv(a, p):
    """Invert regularized lower gamma; finite a > 0, p clipped to [0, 1]."""
    return _inverse(a, p, False)

@jax.jit
def gammainccinv(a, q):
    """Invert regularized upper gamma directly, preserving tiny q.

    Finite a > 0 is required. Probabilities are clipped to [0, 1]; q=0
    returns infinity and q=1 returns zero. Invalid shapes return NaN.
    Supports implicit forward and reverse differentiation.
    """
    return _inverse(a, q, True)

def digammainv(
    y: Union[float, jnp.ndarray], maxiter: int = 5, tol: float = 1e-14
) -> Union[float, jnp.ndarray]:
    """
    Inverse of the digamma function using Newton's method with asymptotic approximations.

    Implementation follows the approach described in the literature:
    For Ψ(x) = y, use Newton's method with the update:
    x_new = x_old - (Ψ(x) - y) / Ψ'(x)

    NOTE: Digamm is only invertible for x > 0. and this function assumes that y is in the domain of invertibility.

    Args:
        y: The value to find the inverse digamma for
        maxiter: Maximum number of Newton iterations (5 iterations typically sufficient for 14 digits)
        tol: Tolerance for convergence

    Returns:
        x such that digamma(x) = y
    """

    # Initial guess based on asymptotic formulas (eq. 149 in the paper)
    def initial_guess(y):
        gamma = jnp.euler_gamma  # Euler-Mascheroni constant

        # Use asymptotic approximations
        # For y ≥ -2.22: x ≈ exp(y) + 1/2
        # For y < -2.22: x ≈ -1/(y + γ)
        x_init = jnp.where(y >= -2.22, jnp.exp(y) + 0.5, -1.0 / (y + gamma))

        return x_init

    # Define Newton step function (eq. 146 in the paper)
    def newton_step(x_old):
        # Calculate Ψ(x) - y and Ψ'(x)
        psi_x = digamma(x_old)
        psi_prime_x = polygamma(1, x_old)  # First derivative of digamma

        # Newton update
        x_new = x_old - (psi_x - y) / psi_prime_x

        return x_new

    # Initialize with the asymptotic approximation
    x = initial_guess(y)

    # Run fixed number of Newton iterations using lax.fori_loop
    def body_fun(i, x):
        return newton_step(x)

    # The paper states 5 iterations are sufficient for 14 digits of precision
    x = lax.fori_loop(0, maxiter, body_fun, x)

    return jnp.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)