# tinydiffeq

> Tiny differentiable ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Fixed stepping and the default adaptive path use bounded lax.scan loops with static output shapes and forward/reverse AD, including reverse-over-forward; adaptive Tsit5 and Rodas5P ODE/DAE solves can instead select adaptive_loop="forward", an actual-work lax.while_loop for primal, JVP, and nested forward AD only. max_steps is an attempt budget, Solution.num_steps records actual attempts, and fixed-step times are arithmetic and budget-invariant. SaveAt picks the endpoint, a fixed interpolation grid (exact=True gathers aligned knots for explicit constant-step ODEs), or the padded accepted-step prefix. Adaptive AD freezes the internal mesh (stop-gradiented controllers) and omits mesh motion. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, and Steinebach's stiff-aware fourth-order dense output, following SciML's OrdinaryDiffEqRosenbrock. RK4/Tsit5 DAEs delegate stage roots and implicit derivatives to nlls-gram; LMRootSolver requires residual-only stopping (gtol=xtol=0) and accepts only CONVERGED roots with Euclidean residual norm below the root atol, with previous-root or opt-in secant stage predictors. solve_sde integrates diagonal-noise Ito SDEs with EulerMaruyama (strong order 0.5), Milstein (1.0, commutative diagonal), or SRA1 (1.5, additive); each solver declares its per-step noise via sample_noise, drawn from a key or passed explicitly through noise= as validated, differentiable data. solve_semi_explicit_sdae applies EulerMaruyama or SRA1 to the reduced index-1 stochastic system with root-restored consistency; SaveAt(ts=...) raises for SDEs/SDAEs because interpolating rough paths is wrong. solve_bvp is a faithful port of scipy's collocation solve_bvp (4th-order Lobatto IIIA, damped Newton, residual-controlled insert-1/insert-2 mesh refinement, optional singular term S y/(t-a)) with pointwise fun(t, y, z, args, p) and bc(ya, yb, z, args, p): z are scipy's unknown parameters solved from a z_0 guess, p the differentiable inputs, outputs padded to a static max_nodes (default 128) with num_nodes marking the active prefix, scipy's 0/1/2/3 statuses as data, AD-computed local Jacobians replacing fun_jac/bc_jac and finite differences, one custom_jvp implicit-function-theorem rule at the solution (reverse mode by transposition; failed solves and all guesses get exact-zero tangents), and dense output via hermite_interpolate/hermite_derivative which evaluate exactly scipy's returned cubic spline. Fields and stochastic drifts may return (value, saved_aux); DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary pytrees of same-dtype real floating arrays, and the library never sets jax_enable_x64. Finite-state DTMC/CTMC sampling is primal-only; deterministic distribution forecasts are differentiable in their initial mass through matrix powers, dense exponentials, or matrix-free Arnoldi/Krylov actions. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints.

## Docs

- [Home — positioning, install, vector-field signature convention f(x, t, args, p), minimal example, design contracts](https://highdimensionaleconlab.github.io/tinydiffeq/)
- [ODEs — solver table, fixed vs adaptive control, static shapes and SaveAt, AD through adaptive stepping, aux outputs, stiff Rodas5P with SciML credit](https://highdimensionaleconlab.github.io/tinydiffeq/ode/)
- [SDEs — EulerMaruyama/Milstein/SRA1 orders and noise contracts, explicit noise= as differentiable data, fixed-key semantics, shared-path convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/)
- [Semi-explicit DAEs — index-1 contract, nlls-gram LM roots and predictors, implicit AD, aux contracts, Rodas5P mass-matrix path, dense output, failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
- [Semi-explicit SDAEs — reduced-SDE EulerMaruyama and SRA1, algebraic roots, aux, convergence assumptions, pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/)
- [Boundary value problems — scipy solve_bvp port, z unknowns vs differentiable p, static max_nodes padding, statuses, implicit AD contract, dense output via hermite_interpolate, credit and deviations](https://highdimensionaleconlab.github.io/tinydiffeq/bvp/)
- [Finite-state Markov chains — sampling, deterministic PMF forecasts, matrix-free Krylov CTMC actions, pytrees, vmap, and AD scope](https://highdimensionaleconlab.github.io/tinydiffeq/markov_chains/)
- [Linear exponential solves — dense expm and matrix-free Arnoldi/Krylov actions for fixed linear array or pytree operators, with traced and hand-coded initial-state JVP/VJP](https://highdimensionaleconlab.github.io/tinydiffeq/exponential/)
- [API reference — solve functions, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/)
