tinydiffeq 0.1.0__tar.gz → 0.2.0__tar.gz

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tinydiffeq
-Version: 0.1.0
+Version: 0.2.0
 Summary: Tiny differentiable ODE/SDE solvers for JAX — bounded-scan adaptive stepping, static shapes, jvp/vjp-safe
 Project-URL: Homepage, https://github.com/HighDimensionalEconLab/tinydiffeq
 Project-URL: Repository, https://github.com/HighDimensionalEconLab/tinydiffeq
@@ -30,19 +30,21 @@ Description-Content-Type: text/markdown
 [![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 solvers for JAX: fixed-step Euler/RK4, adaptive
-Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+Tsit5 with integral or proportional-integral step-size control, and
+Euler–Maruyama for Itô SDEs.
 One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
 adaptive stepping, so shapes are static, nothing recompiles as tolerances or
 curvature change, and every solve is differentiable in **both** forward and
 reverse mode — including reverse-over-forward, the pattern a
 Levenberg–Marquardt optimizer with geodesic acceleration needs when it
-differentiates through a rollout.
+differentiates through a rollout. After a solve reaches its horizon, a
+`lax.cond` skips solver and controller work during the padded scan tail.
 
 This is a deliberately small, jvp/vjp-friendly subset of
 [diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
-states, stiff/implicit solvers, PID control, events, dense output, or
-checkpointed/backsolve adjoints. tinydiffeq's single runtime dependency is
-`jax`.
+states, stiff/implicit solvers, full derivative-term PID control, events,
+dense output, or checkpointed/backsolve adjoints. tinydiffeq's single runtime
+dependency is `jax`.
 
 ## Install
 
@@ -78,22 +80,40 @@ def f(x, t, args, p):
 sol = solve_ode(
     f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
     p=jnp.asarray(1.3),
-    dt0=0.1,
+    dt_0=0.1,
     controller=IController(rtol=1e-8, atol=1e-10),
     max_steps=512,
-    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
+    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 t1 within the max_steps budget?
+print(sol.ok)   # reached t_1 within the max_steps budget?
 ```
 
+`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.
+
+`SaveAt(ts=...)` also accepts a Python sequence. These are observation times:
+the adaptive controller still chooses its own internal mesh, and cubic
+Hermite interpolation evaluates the solution at every requested point.
+
 ## Gradients through the solve
 
 ```python
 def endpoint(p):
     return solve_ode(
         f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
-        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
         max_steps=512,
     ).xs
 
@@ -109,7 +129,7 @@ non-differentiable either way, and the error-ratio power blows up at exactly
 zero error); the states differentiate fully through the RK stages. See the
 [docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
 contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
-semantics, and migration recipes from hand-rolled RK4/Tsit5 loops.
+semantics, and the package API.
 
 ## License
 

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/README.md

@@ -8,19 +8,21 @@
 [![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 solvers for JAX: fixed-step Euler/RK4, adaptive
-Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+Tsit5 with integral or proportional-integral step-size control, and
+Euler–Maruyama for Itô SDEs.
 One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
 adaptive stepping, so shapes are static, nothing recompiles as tolerances or
 curvature change, and every solve is differentiable in **both** forward and
 reverse mode — including reverse-over-forward, the pattern a
 Levenberg–Marquardt optimizer with geodesic acceleration needs when it
-differentiates through a rollout.
+differentiates through a rollout. After a solve reaches its horizon, a
+`lax.cond` skips solver and controller work during the padded scan tail.
 
 This is a deliberately small, jvp/vjp-friendly subset of
 [diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
-states, stiff/implicit solvers, PID control, events, dense output, or
-checkpointed/backsolve adjoints. tinydiffeq's single runtime dependency is
-`jax`.
+states, stiff/implicit solvers, full derivative-term PID control, events,
+dense output, or checkpointed/backsolve adjoints. tinydiffeq's single runtime
+dependency is `jax`.
 
 ## Install
 
@@ -56,22 +58,40 @@ def f(x, t, args, p):
 sol = solve_ode(
     f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
     p=jnp.asarray(1.3),
-    dt0=0.1,
+    dt_0=0.1,
     controller=IController(rtol=1e-8, atol=1e-10),
     max_steps=512,
-    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
+    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 t1 within the max_steps budget?
+print(sol.ok)   # reached t_1 within the max_steps budget?
 ```
 
+`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.
+
+`SaveAt(ts=...)` also accepts a Python sequence. These are observation times:
+the adaptive controller still chooses its own internal mesh, and cubic
+Hermite interpolation evaluates the solution at every requested point.
+
 ## Gradients through the solve
 
 ```python
 def endpoint(p):
     return solve_ode(
         f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
-        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
         max_steps=512,
     ).xs
 
@@ -87,7 +107,7 @@ non-differentiable either way, and the error-ratio power blows up at exactly
 zero error); the states differentiate fully through the RK stages. See the
 [docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
 contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
-semantics, and migration recipes from hand-rolled RK4/Tsit5 loops.
+semantics, and the package API.
 
 ## License
 

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/docs/adaptive_ad.md

@@ -5,10 +5,31 @@ differentiable in forward mode, reverse mode, and reverse-over-forward. That
 holds because of three deliberate choices about **what is not
 differentiated**.
 
+## Default tolerances follow the state precision
+
+`IController()` and `PIController()` resolve omitted tolerances from the
+state dtype at trace time:
+
+| State/time dtype | `rtol` | `atol` | unit-scale `dt_min` |
+|---|---:|---:|---:|
+| float32 | `1e-4` | `1e-6` | `10 * eps` ≈ `1.19e-6` |
+| float64 | `1e-7` | `1e-9` | `10 * eps` ≈ `2.22e-15` |
+
+The policy never changes `jax_enable_x64`: a float32 `x_0` stays float32 even
+when x64 is enabled globally. Explicit tolerances override the defaults and
+are cast to the state dtype. Use explicit values whenever tolerances are part
+of a reproducibility or accuracy contract.
+
+Automatic `dt_min` is `10 * eps * max(1, abs(t_1))`; set it explicitly when
+the relevant time scale differs materially from the absolute horizon. The
+controller also floors exact-zero error ratios at machine epsilon before
+applying a negative power; the growth-factor clip then selects maximal step
+growth without introducing infinities.
+
 ## The controller is stop-gradiented
 
-The step-size controller (`IController`) computes its scaled error norm and
-next-step factor inside `stop_gradient`:
+The adaptive step-size controllers (`IController` and `PIController`) compute
+their scaled error norms and next-step factors inside `stop_gradient`:
 
 > accept/reject is a non-differentiable branch either way, the gradient of
 > `E**(-1/order)` blows up at the exact-zero error of a flat-start policy,
@@ -21,13 +42,27 @@ pattern actually taken*. This is the same convention diffrax uses.
 
 The `E**(-1/5)` blow-up is not hypothetical: a policy initialized flat gives
 an exactly-zero error estimate on the first step, and without the
-`stop_gradient` (plus the `max(E, 1e-12)` floor) the backward pass is NaN
+`stop_gradient` (plus the machine-epsilon error floor) the backward pass is NaN
 from iteration one. `tests/test_ad.py::test_grad_finite_on_flat_field` pins
 this.
 
+`PIController` additionally carries the previous accepted error ratio. Its
+step-size factor is
+
+```text
+safety * E_n**(-(p_coeff + i_coeff) / order)
+       * E_prev**(p_coeff / order)
+```
+
+and `E_prev` changes only after acceptance. The whole recurrence is
+controller-internal and stop-gradiented. Setting `p_coeff=0, i_coeff=1`
+reproduces `IController` bit for bit; the default `p_coeff=0.4, i_coeff=0.3`
+is less sensitive to oscillatory error estimates and typically rejects fewer
+attempts on harder problems.
+
 ## The horizon clip is the growth guard
 
-Every attempt is clipped so it cannot step past `t1`, and — deliberately
+Every attempt is clipped so it cannot step past `t_1`, and — deliberately
 unlike diffrax — the controller's next-step proposal is computed from the
 **clipped** step:
 
@@ -36,15 +71,15 @@ unlike diffrax — the controller's next-step proposal is computed from the
 > whose Gauss–Newton linearization stalls a trust-region optimizer
 > differentiating through the rollout.
 
-With `factormax = 5`, a flat field would otherwise reach `dt ≈ t1/4` within
+With `factor_max = 5`, a flat field would otherwise reach `dt ≈ t_1/4` within
 a few accepted steps; residuals sampled from three or four giant steps make
 the optimizer's linear model useless. Clipping first means the proposal can
-never exceed `factormax × remaining horizon`.
+never exceed `factor_max × remaining horizon`.
 
 One refinement over a bare `min(dt, remaining)`: when the remaining horizon
 is within `max_steps × eps` of the desired step, the step is stretched to
-land on `t1` exactly. Summing `n` rounded steps of `(t1 - t0)/n` can leave
-`t` one accumulated ulp short of `t1`, and without the stretch that sliver
+land on `t_1` exactly. Summing `n` rounded steps of `(t_1 - t_0)/n` can leave
+`t` one accumulated ulp short of `t_1`, and without the stretch that sliver
 would cost an extra iteration a `max_steps = n` budget doesn't have.
 
 ## Interpolation knots are non-differentiable
@@ -74,15 +109,15 @@ on the output is not enough — reverse mode differentiates both branches, and
   Levenberg–Marquardt geodesic-acceleration pattern) all work through
   adaptive solves and interpolated output, verified against closed forms in
   `tests/test_ad.py`.
-- `jax.vmap` over `x0` or `p` gives genuinely per-lane adaptivity: each lane
+- `jax.vmap` over `x_0` or `p` gives genuinely per-lane adaptivity: each lane
   accepts/rejects independently through the masked scan, and batched results
   equal the individual solves exactly.
 
 ## What to watch
 
-- Reductions over `SaveAt(steps=True)` rows are **discontinuous** in the
-  inputs: the number of duplicate rows changes when an accept flips to a
-  reject. Consume steps mode with residuals that vanish at every state (so
-  duplicates are harmless), or reduce to `xs[-1]`.
+- Reductions over the valid prefix of `SaveAt(steps=True)` are
+  **discontinuous** in the inputs: its length changes when an accept flips to
+  a reject. Use `sol.accepted` when padding must not contribute. Because the
+  default tail repeats the endpoint, `xs[-1]` remains the reached final state.
 - Finite-difference checks of adaptive solves are noisy for the same reason;
   compare AD against closed forms or use fixed-step solvers for FD tests.

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/docs/api.md

@@ -22,6 +22,8 @@
 
 ::: tinydiffeq.IController
 
+::: tinydiffeq.PIController
+
 ## Output selection and results
 
 ::: tinydiffeq.SaveAt

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/docs/index.md

@@ -1,8 +1,9 @@
 # tinydiffeq
 
 `tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE integrators
-for JAX: fixed-step Euler and RK4, adaptive Tsit5 with an integral step-size
-controller, and fixed-step Euler–Maruyama for Itô SDEs. Everything runs
+for JAX: fixed-step Euler and RK4, adaptive Tsit5 with integral or
+proportional-integral step-size control, and fixed-step Euler–Maruyama for
+Itô SDEs. Everything runs
 inside one bounded `lax.scan` with static shapes, and every solve is
 differentiable in **both** forward and reverse mode — including
 reverse-over-forward, the pattern a Levenberg–Marquardt optimizer with
@@ -14,7 +15,7 @@ need any of:**
 
 - pytree states (tinydiffeq states are single arrays, scalar or vector)
 - stiff or implicit solvers
-- PID step-size control
+- full derivative-term PID step-size control
 - events, root-finding, or backward-time integration
 - dense output / continuous interpolation objects
 - checkpointed or backsolve adjoints for long horizons
@@ -54,7 +55,7 @@ f(x, t, args, p)
   nothing stops you differentiating with respect to it, but the library's
   contracts and tests treat it as constants.
 - `p` holds differentiable parameters — any pytree, e.g. neural-network
-  weights. jvp/vjp with respect to `p` and `x0` are first-class and tested.
+  weights. jvp/vjp with respect to `p` and `x_0` are first-class and tested.
 
 The arity is inspected once and the function is wrapped into the canonical
 four-argument form, so the compiled code is identical for all four. There is
@@ -78,22 +79,28 @@ def f(x, t, args, p):
 sol = solve_ode(
     f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
     p=jnp.asarray(1.3),
-    dt0=0.1,
+    dt_0=0.1,
     controller=IController(rtol=1e-8, atol=1e-10),
     max_steps=512,
-    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),
+    save_at=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),
 )
 sol.xs  # (21,) states on the grid, however many internal steps were taken
-sol.ok  # False if max_steps ran out before t1
+sol.ok  # False if max_steps ran out before t_1
 ```
 
+With no arguments, `IController()` and `PIController()` use precision-aware
+tolerances: `rtol=1e-4, atol=1e-6` for float32 states and
+`rtol=1e-7, atol=1e-9` for float64 states. Explicit values override the
+policy and are cast to the state dtype. Automatic `dt_min` is
+`10 * eps * max(1, abs(t_1))` in the time dtype.
+
 Gradients go straight through the solve:
 
 ```python
 def endpoint(p):
     return solve_ode(
         f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
-        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
         max_steps=512,
     ).xs
 
@@ -106,20 +113,24 @@ jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])(
 
 ## Design contracts at a glance
 
-- **`dt0` is required.** There is no initial-step heuristic.
-- **Forward time only**: `t1 > t0`.
-- **Never poisons.** `sol.ok` reports whether `t1` was reached; callers that
+- **`dt_0` is required.** There is no initial-step heuristic.
+- **`max_steps` counts attempted internal steps**, including rejections. It
+  controls the bounded scan and only becomes an output-row count in
+  `SaveAt(steps=True)`, which returns `max_steps + 1` padded rows including
+  the initial state. Accepted steps form a contiguous prefix; rejected
+  attempts are omitted and the tail repeats the last accepted state.
+- **Forward time only**: `t_1 > t_0`.
+- **Never poisons.** `sol.ok` reports whether `t_1` was reached; callers that
   want diverging residuals do `jnp.where(sol.ok, sol.xs, jnp.inf)`.
 - **`project`** (an idempotent clamp, e.g. positivity) is applied at every
   point where the vector field is evaluated and to every accepted state.
 - **Never sets `jax_enable_x64`.** The time dtype follows
-  `jnp.result_type(x0, float)`; float32 problems stay float32 even under
+  `jnp.result_type(x_0, float)`; float32 problems stay float32 even under
   x64.
 - Solvers, controllers, `SaveAt`, and `Solution` are frozen dataclasses
-  registered as pytrees: numeric fields (tolerances, grids, `dt0`, `x0`) are
+  registered as pytrees: numeric fields (tolerances, grids, `dt_0`, `x_0`) are
   data leaves, so changing them never recompiles.
 
 Read next: [Static Shapes](static_shapes.md) for the bounded-scan design and
 `SaveAt`, [Adaptive Stepping and AD](adaptive_ad.md) for what is and is not
-differentiated, [SDEs](sde.md), and
-[Migration](migration.md) if you are replacing hand-rolled RK4/Tsit5 loops.
+differentiated, [SDEs](sde.md), and the [API Reference](api.md).

tinydiffeq-0.2.0/docs/llms.txt

@@ -0,0 +1,11 @@
+# tinydiffeq
+
+> Tiny differentiable ODE/SDE solvers for JAX. Fixed-step (Euler, RK4) and adaptive (Tsit5 + integral or proportional-integral controller) explicit Runge-Kutta inside one bounded lax.scan of exactly max_steps attempted steps: static shapes, one compilation across explicit controller values/curvature/initial conditions, and solves differentiable in both forward and reverse mode (including reverse-over-forward) with O(max_steps) memory. Adaptive-controller defaults are precision-aware: rtol=1e-4, atol=1e-6 for float32 and rtol=1e-7, atol=1e-9 for float64; dt_min is 10*eps*max(1, abs(t_1)) in the time dtype. Expensive solver work is skipped after the horizon via lax.cond while the scan emits its padded tail. Array states only (scalar or vector). SaveAt selects endpoint, cubic-Hermite interpolation onto a requested grid, or a compact accepted-step prefix padded to max_steps + 1. solve_sde is fixed-step Euler-Maruyama with presampled diagonal noise keyed on a PRNG key. A deliberately small subset of diffrax: use diffrax for pytree states, stiff/implicit solvers, full derivative-term PID control, events, dense output, and adjoint methods.
+
+## Docs
+
+- [Home — positioning, vector-field signature convention f(x, t, args, p), minimal examples](https://highdimensionaleconlab.github.io/tinydiffeq/)
+- [Static shapes — the bounded-scan design, SaveAt observation grids, compact accepted steps and padding, why nothing recompiles](https://highdimensionaleconlab.github.io/tinydiffeq/static_shapes/)
+- [Adaptive stepping and AD — stop-gradiented controller rationale, horizon-clip growth guard, non-differentiable interpolation knots, double-where NaN safety](https://highdimensionaleconlab.github.io/tinydiffeq/adaptive_ad/)
+- [SDEs — Euler-Maruyama orders, fixed-noise key semantics, shared-path strong-convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/)
+- [API reference — solve_ode, solve_sde, solvers, controllers, SaveAt, Solution, hermite_interpolate, cumulative_trapezoid](https://highdimensionaleconlab.github.io/tinydiffeq/api/)

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/docs/sde.md

@@ -14,11 +14,11 @@ increments.
 from tinydiffeq import solve_sde, EulerMaruyama, SaveAt
 
 sol = solve_sde(
-    drift, diffusion, EulerMaruyama(), 0.0, 1.0, x0,
+    drift, diffusion, EulerMaruyama(), 0.0, 1.0, x_0,
     key=jax.random.PRNGKey(0),
     n_steps=256,
     p=(mu, sigma),
-    saveat=SaveAt(steps=True),
+    save_at=SaveAt(steps=True),
 )
 ```
 
@@ -38,13 +38,13 @@ for now; a roadmap issue sketches what it would take here.
 The Brownian increments are presampled once,
 
 ```python
-dW = jnp.sqrt(dt) * jax.random.normal(key, (n_steps,) + x0.shape)
+d_w = jnp.sqrt(dt) * jax.random.normal(key, (n_steps,) + x_0.shape)
 ```
 
 so a fixed `key` pins the entire noise path:
 
 - **Reproducible** — the same key gives the same path, every call.
-- **Differentiable with respect to `x0` and `p`** (not `key`): with the path
+- **Differentiable with respect to `x_0` and `p`** (not `key`): with the path
   held fixed, the solution map is smooth, and jvp/vjp against finite
   differences are tested in `tests/test_sde.py`. This is exactly the "common
   random numbers" setup simulation-based estimators want.
@@ -61,7 +61,7 @@ X_T = X_0 \exp\left((\mu - \tfrac{\sigma^2}{2})T + \sigma W_T\right),
 \qquad W_T = \textstyle\sum_k \Delta W_k .
 $$
 
-Because the increments `dW` are exactly reproducible from the key, the test
+Because the increments `d_w` are exactly reproducible from the key, the test
 regenerates them, computes the exact endpoint on the same path, and checks
 the mean absolute error slope across `dt` levels lands in `[0.35, 0.65]`.
 That shared-path construction is the right way to measure strong error —
@@ -69,7 +69,7 @@ comparing against an independently sampled exact solution measures nothing.
 
 ## SaveAt for SDEs
 
-`SaveAt(t1=True)` (default) and `SaveAt(steps=True)` (`n_steps + 1` rows,
+`SaveAt(t_1=True)` (default) and `SaveAt(steps=True)` (`n_steps + 1` rows,
 all accepted) are supported. `SaveAt(ts=...)` **raises**: cubic Hermite
 interpolation assumes smooth trajectories and is simply wrong between the
 points of a rough path. Land your grid on the step boundaries instead by

tinydiffeq-0.2.0/docs/static_shapes.md

@@ -0,0 +1,124 @@
+# Static Shapes
+
+JAX jits fixed-shape programs. An adaptive integrator is naturally
+dynamic — the number of steps depends on the data — so something must give.
+tinydiffeq's answer is a **bounded scan**: `solve_ode` always runs one
+`lax.scan` of exactly `max_steps` iterations, whatever the controller does.
+
+Each iteration attempts one step:
+
+- an **accepted** attempt advances `(t, x)` and (for FSAL solvers) reuses the
+  last stage as the next first stage;
+- a **rejected** attempt leaves the state in place and retries with a smaller
+  step;
+- once `t_1` is reached, the remaining iterations **freeze**.
+
+The raw internal scan buffer contains repeated rows for rejected and frozen
+iterations. That buffer is an implementation detail used by interpolation;
+step output compacts it into accepted rows plus tail padding.
+
+The scan loop itself always has `max_steps` iterations, preserving static
+shapes and reverse-mode AD. A scalar `lax.cond` nevertheless keeps the
+expensive vector-field, solver-stage, and controller computations out of the
+frozen tail, so a scalar solve's compute cost follows its attempted steps plus
+cheap loop overhead. Under `vmap`, JAX may turn each lane's conditional into
+selection; batched work can therefore continue until the slowest lane
+finishes.
+
+Fixed-step and adaptive integration share this single code path:
+`ConstantStepSize` accepts every attempt, so `dt_0 = (t_1 - t_0)/n` with
+`max_steps = n` reproduces a fixed grid exactly. The conditional has modest
+overhead for very cheap solves whose budget is already exact; it pays off
+when the budget has a padded tail or the vector field is expensive.
+
+If the budget runs out before `t_1`, `sol.ok` is `False` and the outputs hold
+the reached prefix. The package never poisons values; the caller decides:
+
+```python
+xs = jnp.where(sol.ok, sol.xs, jnp.inf)  # kernels-style rejection
+```
+
+## SaveAt is the shape contract
+
+Exactly one of three modes:
+
+### `SaveAt(t_1=True)` — endpoint only (default)
+
+`sol.ts` is the reached time (equals `t_1` when `ok`), `sol.xs` the final
+state.
+
+### `SaveAt(ts=grid)` — interpolation onto a fixed grid
+
+This is the answer to "adaptive steps vs static shapes". Internal steps
+adapt freely; the output is cubic-Hermite interpolation onto your fixed
+query grid, so `sol.xs.shape == (len(grid),) + x_0.shape` **regardless of how
+many steps the controller took**. Changing tolerances, initial conditions,
+or curvature changes the internal knots but never the output shape — no
+recompilation. A one-dimensional JAX/NumPy array or Python sequence is
+accepted; times must be nondecreasing and within `[t_0, t_1]`, while repeated
+times and omitted endpoints are allowed. Changing values without changing
+the grid length does not recompile.
+
+These are observation times, not mandatory internal stops. The adaptive
+controller chooses exactly the same mesh regardless of the requested grid,
+then the solver evaluates every requested point through dense interpolation.
+Forcing exact internal landing times is a distinct feature and is not part of
+this API.
+
+The interpolation runs directly over the raw padded rows: duplicate knots
+from rejections or the frozen tail form zero-width brackets, and the
+bracketing `searchsorted` lands on the last duplicate at-or-before each
+query, so no compaction pass is needed. Queries outside the knot span clamp
+to the boundary values — in particular, when `ok` is `False`, queries beyond
+the reached time return the last state (flat extrapolation) rather than
+evaluating a cubic outside its bracket.
+
+The interpolant is 4th-order accurate between 5th-order-accurate knots:
+expect grid values slightly less accurate than the knots themselves, which
+is the standard dense-output trade-off.
+
+### `SaveAt(steps=True)` — accepted steps with padding
+
+`max_steps + 1` rows including the initial state. Accepted internal steps are
+gathered chronologically into a contiguous prefix; rejected attempts are not
+returned. `sol.accepted` is the validity mask (`accepted[0]` is always
+`True`, so `accepted.sum() == num_accepted + 1`). On a successful solve, the
+endpoint is at index `num_accepted`.
+
+- `fill="last"` (default) repeats the last valid time and state through the
+  padded tail.
+- `fill="inf"` fills only the invalid tail with `inf`.
+
+If the attempt budget is exhausted, the same contract holds for the reached
+accepted prefix, the tail repeats or masks its last state, and `sol.ok` is
+`False`. No fake endpoint is inserted. Compaction is a stable fixed-size
+gather; it performs no sorting.
+
+## Why one compilation, precisely
+
+- The scan length `max_steps` is static; nothing else about the loop depends
+  on data shapes.
+- Tolerances and PI coefficients (`IController(...)` / `PIController(...)`),
+  `dt_0`, `t_0`, `t_1`, `x_0`,
+  `args`, `p`, and `SaveAt.ts` are pytree **data leaves**. Only genuine
+  structure — the solver type, `SaveAt` mode, `fill`, `max_steps`, the
+  functions themselves — is static.
+
+An omitted tolerance or `dt_min` is represented by `None`, so switching a
+jitted call between automatic and explicit values changes the controller
+pytree structure and compiles once for each policy. Changing the numeric
+values of already-explicit controller fields does not recompile.
+
+So this compiles once:
+
+```python
+@jax.jit
+def run(x_0, dt_0, controller, args):
+    return solve_ode(f, Tsit5(), 0.0, 1.0, x_0, args=args, dt_0=dt_0,
+                     controller=controller, max_steps=128,
+                     save_at=SaveAt(steps=True))
+```
+
+across different curvatures (different accepted counts), tolerances, initial
+steps, and initial conditions — pinned by `tests/test_recompile.py` with
+`_cache_size() == 1` assertions.

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/mkdocs.yml

@@ -46,5 +46,4 @@ nav:
   - Static Shapes: static_shapes.md
   - Adaptive Stepping and AD: adaptive_ad.md
   - SDEs: sde.md
-  - Migration from Hand-Rolled Integrators: migration.md
   - API Reference: api.md

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tinydiffeq"
-version = "0.1.0"
+version = "0.2.0"
 description = "Tiny differentiable ODE/SDE solvers for JAX — bounded-scan adaptive stepping, static shapes, jvp/vjp-safe"
 readme = "README.md"
 license = "MIT"

{tinydiffeq-0.1.0 → tinydiffeq-0.2.0}/src/tinydiffeq/__init__.py

@@ -1,22 +1,23 @@
 """Tiny differentiable ODE/SDE solvers for JAX.
 
 solve_ode integrates dx/dt = f(x, t, args, p) with fixed-step (Euler, RK4)
-or adaptive (Tsit5 + IController) explicit Runge-Kutta methods inside one
-bounded lax.scan of exactly max_steps iterations, so shapes are static and
-solves are differentiable in BOTH forward and reverse mode (including
-reverse-over-forward) with O(max_steps) memory. SaveAt picks the output:
-the endpoint, cubic-Hermite interpolation onto a fixed grid, or the raw
-padded step rows. solve_sde is fixed-step Euler-Maruyama with presampled
+or adaptive (Tsit5 + IController/PIController) explicit Runge-Kutta methods
+inside one bounded lax.scan of exactly max_steps iterations, so shapes are
+static and solves are differentiable in BOTH forward and reverse mode
+(including reverse-over-forward) with O(max_steps) memory. SaveAt picks the
+output: the endpoint, cubic-Hermite interpolation onto a fixed grid, or
+accepted internal steps with fixed-shape padding. solve_sde is fixed-step
+Euler-Maruyama with presampled
 diagonal noise. States are arrays (scalar or vector); pytree states,
-implicit/stiff solvers, PID control, events, dense output, and adjoint
-methods are non-goals — use diffrax for those.
+implicit/stiff solvers, full derivative-term PID control, events, dense
+output, and adjoint methods are non-goals — use diffrax for those.
 """
 
-from tinydiffeq.controllers import ConstantStepSize, IController
+from tinydiffeq.controllers import ConstantStepSize, IController, PIController
 from tinydiffeq.interpolation import hermite_interpolate
 from tinydiffeq.ode import solve_ode
 from tinydiffeq.quadrature import cumulative_trapezoid
-from tinydiffeq.saveat import SaveAt
+from tinydiffeq.save_at import SaveAt
 from tinydiffeq.sde import solve_sde
 from tinydiffeq.solution import Solution
 from tinydiffeq.solvers import RK4, Euler, EulerMaruyama, Tsit5
@@ -30,6 +31,7 @@ __all__ = [
     "EulerMaruyama",
     "ConstantStepSize",
     "IController",
+    "PIController",
     "SaveAt",
     "Solution",
     "hermite_interpolate",