Metadata-Version: 2.4
Name: tinydiffeq
Version: 2.4.0
Summary: Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX with static shapes and composable AD
Project-URL: Homepage, https://github.com/HighDimensionalEconLab/tinydiffeq
Project-URL: Repository, https://github.com/HighDimensionalEconLab/tinydiffeq
Project-URL: Documentation, https://highdimensionaleconlab.github.io/tinydiffeq/
Project-URL: Issues, https://github.com/HighDimensionalEconLab/tinydiffeq/issues
Author-email: Jesse Perla <jesseperla@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: autodiff,differential-equations,jax,ode,runge-kutta,sde
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.11
Requires-Dist: jax>=0.7.0
Requires-Dist: nlls-gram>=2.7.0
Description-Content-Type: text/markdown

# tinydiffeq

[![CI](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml)
[![Docs](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/tinydiffeq/)
[![PyPI](https://img.shields.io/pypi/v/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
[![Python versions](https://img.shields.io/pypi/pyversions/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4,
adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs,
and Euler–Maruyama for Itô SDEs and semi-explicit index-1 SDAEs. The package
also includes primal, vmap-friendly finite-state DTMC and CTMC simulators with
sequential and associative parallel-prefix execution. Deterministic probability
forecasts are differentiable in the initial mass and include DTMC matrix powers,
dense CTMC exponentials, and matrix-free Arnoldi/Krylov actions over probability
pytrees. The same dense and matrix-free backends are available directly through
`solve_linear_ode` for any fixed homogeneous linear array or pytree operator;
`jvp_linear_ode` and `vjp_linear_ode` apply the exact initial-state tangent and
adjoint exponential actions without differentiating Arnoldi orthogonalization.
Fixed stepping and the default adaptive path use bounded `lax.scan` loops with
exactly `max_steps` attempt slots. Shapes stay static as tolerances or curvature
change, and these solves support forward mode, reverse mode, and
reverse-over-forward. Adaptive ODE and DAE solves may instead select
`adaptive_loop="forward"`: a dynamic `lax.while_loop` that executes only actual
attempts and supports JVP and nested forward AD, but not reverse mode. A vmapped
forward loop runs until its slowest lane finishes.

This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX
adaptation of Steinebach's method and follows SciML's
[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock)
implementation. Use [diffrax](https://docs.kidger.site/diffrax/) or
[SciML](https://docs.sciml.ai/DiffEqDocs/stable/) if you need general mass
matrices, fully implicit or higher-index DAEs, events, continuous solution
objects, sparse/Krylov linear solvers for ODE/DAE stages, or specialized
adjoints. Initial DAE consistency and explicit DAE stages use `nlls-gram`;
the same nlls solve supplies the square root's implicit derivative, whose
default is a direct nonsymmetric `LU()` solve. `LMRootSolver` requires
residual-only stopping (`gtol=xtol=0`) and accepts only `CONVERGED` roots whose
Euclidean residual norm is below the root `atol`. Its
`max_steps_is_success` field remains for source compatibility but does not make
`MAX_STEPS` a valid DAE root.

The linear exponential-action API follows SciML
[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/).
It includes fixed and residual-controlled adaptive matrix-free time slicing;
the latter keeps the Krylov dimension static for predictable JAX compilation.
SciML's
[`ExponentialIntegrators.jl`](https://docs.sciml.ai/ExponentialIntegrators/stable/)
is the reference for the broader nonlinear exponential-integrator family.

## 2.4.0 migration note

- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit
  fixed-step ODEs. Every query must align with a knot; adaptive methods,
  Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode.
- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts,
  including rejections. DAE results additionally expose `num_root_solves` and
  `num_root_steps`; `num_accepted` retains its existing meaning.
- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an
  actual-work loop. It supports primal, JVP, and nested forward AD but not
  reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable
  default. Under `vmap`, the forward loop runs to the slowest lane.
- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for
  locally unique algebraic branches; `predictor="previous"` remains the
  default.

DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root
solve and implicit derivative; square implicit AD defaults to direct `LU()`.
Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so
`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source
compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid
root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they
relied on budget exhaustion, increase the root budget or adjust the residual
tolerance instead.

## Install

```bash
uv add tinydiffeq
```

For GPU use, install the JAX accelerator build that matches your hardware,
for example:

```bash
uv add tinydiffeq "jax[cuda13]"
```

## Minimal example

The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or
`(x, t, args, p)` — always in that order. `args` is pass-through data (not an
AD target by convention); `p` holds differentiable parameters (any pytree).
The state may also be any JAX pytree. It must contain at least one leaf, and
every leaf must be a nonempty real floating array with the same dtype; vector
fields and `project` preserve that structure. Output keeps the structure and
adds the saved-time axis to each
leaf.

```python
import jax
import jax.numpy as jnp
from tinydiffeq import solve_ode, Tsit5, IController, SaveAt

jax.config.update("jax_enable_x64", True)  # your call — the library never sets it


def f(x, t, args, p):
    return -p * x


sol = solve_ode(
    f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
    p=jnp.asarray(1.3),
    dt_0=0.1,
    controller=IController(rtol=1e-8, atol=1e-10),
    max_steps=512,
    save_at=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
)                                                  # however many steps adapt
print(sol.xs)   # states on the grid
print(sol.ok)   # reached t_1 with every requested output valid?
```

`IController()` and `PIController()` choose tolerances from `x_0.dtype`:
`rtol=1e-4, atol=1e-6` for float32 and `rtol=1e-7, atol=1e-9` for
float64. Pass explicit values when tolerances are part of your model's
scientific specification. The default `dt_min` is
`10 * finfo(dtype).eps * max(1, abs(t_1))`.

`max_steps` is the total internal **attempt budget**: accepted steps plus
rejections. It is not normally the number of returned times. Endpoint mode
returns one time/state, `SaveAt(ts=...)` returns the requested grid, and
`SaveAt(steps=True)` returns the initial state and accepted internal steps as
a contiguous prefix of `max_steps + 1` rows. The remaining rows repeat the
last accepted state by default; `sol.accepted` distinguishes data from
padding. Rejected attempts never appear in the returned trajectory.
`sol.num_steps` reports the number of attempts actually made, while
`sol.num_accepted` excludes rejections.

`SaveAt(ts=...)` also accepts a Python sequence. These are observation times:
the adaptive controller still chooses its own internal mesh. Explicit methods
use cubic Hermite interpolation; Rodas5P uses its published stiff-aware
fourth-order continuous extension. For an explicit fixed-step ODE,
`SaveAt(ts=..., exact=True)` instead requires every requested time to be an
internal knot and gathers the stored state directly. Exact mode does not apply
to adaptive ODEs, Rodas5P, DAEs, SDEs, or SDAEs.

## Semi-explicit DAEs

For a square index-1 system `dy/dt = f(y, z, t, args, p)` and
`0 = g(y, z, t, args, p)`:

```python
from tinydiffeq import IController, Rodas5P, Tsit5, solve_semi_explicit_dae


def dae_f(y, z, t, args, p):
    dy = p * z
    return dy, {"flow": dy}


def dae_g(y, z, t, args, p):
    return z - y


dae_sol = solve_semi_explicit_dae(
    dae_f, dae_g, Tsit5(), 0.0, 1.0,
    jnp.asarray(1.0), jnp.asarray(0.5),
    p=jnp.asarray(2.0), dt_0=0.1,
    controller=IController(), max_steps=128,
)
print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])

# One initial nonlinear consistency solve, then linear Rodas5P stages.
stiff_dae_sol = solve_semi_explicit_dae(
    dae_f, dae_g, Rodas5P(), 0.0, 1.0,
    jnp.asarray(1.0), jnp.asarray(0.5),
    p=jnp.asarray(2.0), dt_0=0.1,
    controller=IController(), max_steps=128,
)
```

`z_0` is a guess and is made consistent automatically. RK4 and Tsit5 restore
the algebraic root at every stage. Rodas5P performs no nonlinear solves after
initialization: it advances the corresponding block mass-matrix system using
one reused LU factorization per attempt. Differential fields may return a
floating saved-aux pytree stored at accepted nodes and interpolated on requested
deterministic grids. Algebraic equations may separately return internal context
passed to the dynamics. On the default bounded path, JVP, VJP, and
reverse-over-forward propagate through both implicit initialization and the
time integrator. See the
[DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
for root controls, `SaveAt`, and scope limits.

DAE solutions expose `num_steps`, `num_root_solves`, and `num_root_steps` as
logical per-trajectory work counters. Explicit methods default to reusing the
previous algebraic root as the next stage guess; `LMRootSolver(predictor="secant")`
is an opt-in continuation predictor for locally unique root branches.

Fixed-step semi-explicit Itô SDAEs use the corresponding
`solve_semi_explicit_sdae` interface with `EulerMaruyama`, a PRNG key, and
`n_steps`; see the [SDAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/).

## Gradients through the solve

```python
def endpoint(p):
    return solve_ode(
        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
        dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
        max_steps=512,
    ).xs

jax.grad(endpoint)(jnp.asarray(1.3))                      # reverse mode
jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])(
    jnp.asarray(1.3)
)                                                          # reverse-over-forward
```

The step-size controller is wrapped in `stop_gradient` (accept/reject is
non-differentiable either way, and the error-ratio power blows up at exactly
zero error); states differentiate through the solver stages on the realized,
frozen mesh. In particular, adaptive `SaveAt(steps=True)` does not include mesh
motion in its time or state derivatives. See the
[docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
semantics, and the package API.

## License

MIT
