Utilities¶

def odeint(
    drift: Callable[..., PyTree[Array]],
    y0: PyTree[Array],
    ts: Array,
    *args: Any,
    method: str = "rk4",
    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``.
        method: Integration method name.
        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,
        ...             method="dopri5")
    """
    drift = _wrap_if_plain_callable(drift)
    return _odeint_custom(
        y0,
        drift,
        ts,
        tuple(args),
        logdet_rng,
        method=method,
        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,
    method: str = "euler_maruyama",
    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.
        method: Integration method name. Available:

            - ``"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)
            - ``"srk"``: Stochastic Runge-Kutta methods
        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 (resamples noise on
            rejection — cheaper) or
            :class:`~probjax.utils.sdeutil.adaptive.StrongStepSizeAdaptor`
            for per-path consistency (refines a single Brownian path via
            a virtual tree). Currently restricted to **diagonal noise**
            and ignores ``method`` (always uses Euler-Maruyama with
            step-doubling error estimation). Defaults to ``None``
            (fixed-step integration via ``method``).

    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)
    """
    drift = _wrap_if_plain_callable(drift)
    diffusion = _wrap_if_plain_callable(diffusion)
    result = _sdeint(
        rng,
        drift,
        diffusion,
        y0,
        ts,
        tuple(args),
        method=method,
        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 isinstance(diag_hits, jax.core.Tracer):
        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 bracket_x(a, b, p):
    """Return (lo, hi) such that betainc(a,b,lo) <= p <= betainc(a,b,hi)."""
    beta_ab = jax.scipy.special.beta(a, b)

    lower_bound = (p * a * beta_ab) ** (1.0 / a)
    lower_bound = jnp.where(b >= 1.0, lower_bound, 0.0)

    upper_bound1 = 1.0 - ((1.0 - p) * b * beta_ab) ** (1.0 / b)
    upper_bound2 = 1.0
    upper_bound = jnp.where(a < 1.0, upper_bound2, upper_bound1)

    lower_bound = jnp.clip(lower_bound, 1e-10, 1.0 - 1e-7)
    upper_bound = jnp.clip(upper_bound, 1e-10, 1.0 - 1e-7)
    return lower_bound, upper_bound

def betaincinv(a, b, p, *, max_halley_steps=6, max_bisection_steps=15):
    """
    Inverse of the regularized incomplete beta function:
    returns x in [0,1] s.t. betainc(a,b,x) = p.

    a, b > 0
    p in [0,1]

    max_halley_steps, max_bisection_steps control solver refinement;
    they are treated as static "tuning knobs", not differentiable inputs.
    """
    core = _make_betaincinv_core(max_halley_steps, max_bisection_steps)
    # We can jit the core. The closure constants (step counts) are static,
    # and the JIT only sees (a,b,p) so there's no weird kwarg plumbing.
    core_jit = jax.jit(core)
    return core_jit(a, b, p)

@jax.jit
def gammaincinv(a, p):
    """
    Inverse of the *regularized* lower incomplete gamma function.
    Solves for x >= 0 such that gammainc(a, x) = p.

    Args:
        a (jnp.ndarray): Shape parameter (a > 0).
        p (jnp.ndarray): Probability in [0, 1].

    Returns:
        jnp.ndarray: x in [0, ∞) satisfying gammainc(a, x) = p.
    """

    # Clip p to [0,1]
    p = jnp.clip(p, 0.0, 1.0)

    # Handle trivial cases
    trivial_low = p == 0.0
    trivial_high = p == 1.0
    # By convention: gammainc(a, x=0) = 0 => inverse at p=0 => x=0
    #                gammainc(a, x->∞) = 1 => inverse at p=1 => x=∞
    x_trivial = jnp.where(trivial_low, 0.0, jnp.inf)
    trivial = trivial_low | trivial_high

    # Initial guess
    x_init = _compute_initial_guess_gammaincinv(a, p)
    x_init = jnp.where(trivial, x_trivial, x_init)

    # Refine with safe solver
    x_sol = _safe_gammaincinv_solve(a, p, x_init)

    return jnp.where(trivial, x_trivial, x_sol)

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)