feat: solve_bvp collocation solver with implicit differentiation

AGENTS.md

@@ -0,0 +1,28 @@
+# tinydiffeq — agent notes
+
+## Recompilation hygiene: a manual audit, not a test
+
+Do not write tests against `fn._cache_size()`. It reads JAX's globally shared
+C++ executable cache (capacity 8192, LRU, shared by every jitted function and
+every `jnp` op in the process), so absolute entry-count assertions turn flaky
+as the suite grows: entries get evicted before the assert and the count reads
+zero.
+
+Instead, after a big refactor or a substantial new feature, audit
+recompilation with the environment-variable protocol from the `jax-project`
+skill: two solves in one process with changed data leaves, a flushed marker
+before each, and zero trace/compile events after the second marker.
+
+```bash
+JAX_EXPLAIN_CACHE_MISSES=1 JAX_LOG_COMPILES=1 \
+  uv run python -m benchmarks.bvp_scaling --cache-audit > /tmp/audit.log 2>&1
+grep -c "Compiling" /tmp/audit.log   # after the "=== SOLVE 1" marker: zero
+```
+
+A per-call `eval_shape` validation trace is expected and compiles nothing.
+Leaf-value changes (tolerances, meshes, `p`, `args`, initial states) must
+not recompile; only static configuration (attempt budgets, `max_nodes`,
+pytree structure, function identity, jacobian modes, solver objects) may. A
+changed initial mesh *length* for `solve_bvp` reuses the solve executable
+but compiles trivial eager padding ops (`concatenate`, `broadcast_in_dim`)
+once per new length — those events are expected in the audit log.

CLAUDE.md

@@ -0,0 +1 @@
+@AGENTS.md

README.md

@@ -7,48 +7,45 @@
 [![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
-fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and semi-explicit
-index-1 SDAEs. Solves run in bounded `lax.scan` loops with static shapes and
-compose with `jit`, `vmap`, forward mode, reverse mode, and
-reverse-over-forward. Finite-state Markov simulation, probability forecasts,
-and general fixed homogeneous linear solves (dense or matrix-free Krylov
-exponential actions, after SciML's
-[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/))
-round out the package.
-
-This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX
-adaptation of Steinebach's method following SciML's
-[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock)
-implementation, and DAE algebraic roots delegate both the primal solve and
-the implicit derivative to
-[`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/). Use
-[diffrax](https://docs.kidger.site/diffrax/) or
+tinydiffeq is an unsupported research repo of vibe-coded ports of
+well-established ODE, DAE, and SDE algorithms to JAX. Heavily AI-generated —
+but the algorithms are well established, with
+[SciML](https://docs.sciml.ai/DiffEqDocs/stable/),
+[scipy](https://docs.scipy.org/doc/scipy/reference/integrate.html), and
+[diffrax](https://docs.kidger.site/diffrax/) as the reference
+implementations — so correctness and performance are often reasonable. The
+method set is intentionally minimal, though the package is no longer
+especially tiny.
+
+Fixed-step Euler/RK4, adaptive Tsit5, and linearly implicit Rodas5P for
+stiff ODEs and index-1 DAEs; Euler–Maruyama, Milstein, and SRA1 for Itô SDEs
+and SDAEs; a port of scipy's collocation solver for two-point BVPs with
+unknown parameters; finite-state Markov chains and dense or Krylov linear
+exponential actions. Every solve runs in bounded `lax` loops with static
+shapes and composes with `jit`, `vmap`, forward mode, reverse mode, and
+reverse-over-forward; iterative solves (BVP, DAE roots) differentiate
+implicitly at the solution, never through the iterations.
+
+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, adaptive SDE stepping, events,
-continuous solution objects, sparse/Krylov ODE/DAE stages, or specialized
-adjoints.
+matrices, fully implicit or higher-index DAEs, adaptive SDE stepping,
+events, continuous solution objects, or specialized adjoints.
 
 ## Install
 
 ```bash
 uv add tinydiffeq
 ```
 
-For GPU use, install the JAX accelerator build that matches your hardware,
-for example:
+For GPU use, add the JAX accelerator build matching your hardware, for
+example `uv add tinydiffeq "jax[cuda13]"`.
 
-```bash
-uv add tinydiffeq "jax[cuda13]"
-```
-
-## Minimal example
+## 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, and the state
-may be any pytree of same-dtype real floating arrays.
+`(x, t, args, p)` — always in that order. `args` is pass-through data (not
+an AD target by convention); `p` holds differentiable parameters, and the
+state may be any pytree of same-dtype real floating arrays.
 
 ```python
 import jax
@@ -68,125 +65,57 @@ sol = solve_ode(
     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?
-```
-
-`max_steps` is the internal attempt budget (accepted plus rejected steps),
-not the number of returned times: `SaveAt` picks the endpoint, a fixed
-interpolation grid, or the padded accepted-step prefix, so output shapes
-never depend on how many steps the controller took. Omitted controller
-tolerances follow the state dtype (`1e-4`/`1e-6` in float32,
-`1e-7`/`1e-9` in float64).
-
-## SDEs with first-class noise
-
-`solve_sde` integrates diagonal-noise Itô SDEs with `EulerMaruyama` (strong
-order 0.5), `Milstein` (1.0, commutative diagonal noise), or `SRA1` (1.5,
-additive noise). An Ornstein–Uhlenbeck process under SRA1:
-
-```python
-from tinydiffeq import solve_sde, SRA1
-
-theta, sigma, n = 1.0, 0.5, 256
-
-
-def ou_drift(x):
-    return -theta * x
-
-
-def ou_diffusion(x):
-    return sigma * jnp.ones_like(x)
-
-
-sol = solve_sde(
-    ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, jnp.asarray(1.0),
-    key=jax.random.key(0), n_steps=n,
+    save_at=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),
 )
+sol.xs  # states on the grid, however many internal steps were taken
+sol.ok  # False if integration or a requested output failed
 ```
 
-The noise realization can also be passed explicitly — the same pytree
-`sample_noise` would draw, now inspectable, storable data that is
-differentiable like any other input:
+`max_steps` bounds attempted internal steps (accepted plus rejected);
+`SaveAt` fixes the output shape regardless of how many steps the controller
+takes. Gradients go straight through the solve:
 
 ```python
-x_0 = jnp.asarray(1.0)
-noise = SRA1().sample_noise(x_0, jax.random.key(0), n, jnp.asarray(1.0 / n), x_0.dtype)
-same_sol = solve_sde(
-    ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n
-)  # bit-identical to the key= call
-d_endpoint_d_noise = jax.grad(
-    lambda noise: solve_sde(
-        ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n
+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
-)(noise)
-```
-
-A fixed key (or fixed noise) pins the whole path, so gradients with respect
-to `x_0`, `p`, and `noise` are pathwise derivatives under common random
-numbers — the setup simulation-based estimators want. `vmap` over
-trajectories with per-trajectory `x_0` and noise composes with `jit` and
-`grad`.
 
-## Semi-explicit DAEs
+jax.grad(endpoint)(jnp.asarray(1.3))                         # reverse mode
+jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
+```
 
-For a square index-1 system `dy/dt = f(y, z, t, args, p)` and
-`0 = g(y, z, t, args, p)`:
+An SDE ensemble — per-key noise, `vmap` over trajectories:
 
 ```python
-from tinydiffeq import solve_semi_explicit_dae
-
-
-def dae_f(y, z, t, args, p):
-    dy = p * z
-    return dy, {"flow": dy}
+from tinydiffeq import solve_sde, SRA1
 
 
-def dae_g(y, z, t, args, p):
-    return z - y
+def ou_drift(x):
+    return -x
 
 
-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"])
-```
-
-`z_0` is a guess and is made consistent automatically. RK4 and Tsit5 restore
-the algebraic root at every stage through `nlls-gram`, which also supplies
-the root's implicit derivative; `Rodas5P()` instead performs one initial
-consistency solve and then advances the block mass-matrix system with one
-reused LU factorization per attempt — the stiff path. Stochastic
-semi-explicit systems use `solve_semi_explicit_sdae` with `EulerMaruyama` or
-`SRA1`. See the
-[DAE](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) and
-[SDAE](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) docs.
+def ou_diffusion(x):
+    return 0.5 * jnp.ones_like(x)
 
-## 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,
+def ou_path(key):
+    return solve_sde(
+        ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, jnp.asarray(1.0),
+        key=key, n_steps=256, save_at=SaveAt(steps=True),
     ).xs
 
-jax.grad(endpoint)(jnp.asarray(1.3))                         # reverse mode
-jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
+
+paths = jax.vmap(ou_path)(jax.random.split(jax.random.key(0), 1000))  # (1000, 257)
 ```
 
-The step-size controller is wrapped in `stop_gradient` (accept/reject is
-non-differentiable either way); states differentiate through the solver
-stages on the realized, frozen mesh. See the
-[docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
-contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE
-noise semantics, and the package API.
+SDEs with first-class differentiable noise, semi-explicit DAEs and SDAEs,
+two-point BVPs, Markov chains, linear exponential solves, and the design
+contracts (static shapes and `SaveAt`, AD through adaptive stepping,
+failure-as-data) are in the
+[docs](https://highdimensionaleconlab.github.io/tinydiffeq/).
 
 ## License