tinydiffeq 0.1.0__tar.gz → 0.3.0__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: tinydiffeq
-Version: 0.1.0
-Summary: Tiny differentiable ODE/SDE solvers for JAX — bounded-scan adaptive stepping, static shapes, jvp/vjp-safe
+Version: 0.3.0
+Summary: Tiny differentiable ODE/SDE/DAE 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/
@@ -18,6 +18,7 @@ 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>=1.6.0
 Description-Content-Type: text/markdown
 
 # tinydiffeq
@@ -29,20 +30,22 @@ Description-Content-Type: text/markdown
 [![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 solvers for JAX: fixed-step Euler/RK4, adaptive
-Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+Tiny differentiable ODE/SDE/DAE solvers for JAX: fixed-step Euler/RK4,
+adaptive Tsit5 with integral or proportional-integral step-size control,
+Euler–Maruyama for Itô SDEs, and nonstiff semi-explicit index-1 DAEs.
 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 or fully implicit solvers, higher-index DAEs, full
+derivative-term PID control, events, dense output, or checkpointed/backsolve
+adjoints. The DAE algebraic solve uses `nlls-gram`.
 
 ## Install
 
@@ -78,22 +81,72 @@ 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.
+
+## 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, Tsit5, solve_semi_explicit_dae
+
+
+def dae_f(y, z, t, args, p):
+    return p * z
+
+
+def dae_g(y, z):
+    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)
+```
+
+`z_0` is a guess and is made consistent automatically. Algebraic roots use
+an implicitly differentiated square LM solve, so JVP, VJP, and
+reverse-over-forward propagate with respect to `y`, `t`, and `p`. See the
+[DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
+for root controls, `SaveAt`, and scope limits.
+
 ## 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 +162,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.3.0}/README.md

@@ -7,20 +7,22 @@
 [![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 solvers for JAX: fixed-step Euler/RK4, adaptive
-Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+Tiny differentiable ODE/SDE/DAE solvers for JAX: fixed-step Euler/RK4,
+adaptive Tsit5 with integral or proportional-integral step-size control,
+Euler–Maruyama for Itô SDEs, and nonstiff semi-explicit index-1 DAEs.
 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 or fully implicit solvers, higher-index DAEs, full
+derivative-term PID control, events, dense output, or checkpointed/backsolve
+adjoints. The DAE algebraic solve uses `nlls-gram`.
 
 ## Install
 
@@ -56,22 +58,72 @@ 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.
+
+## 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, Tsit5, solve_semi_explicit_dae
+
+
+def dae_f(y, z, t, args, p):
+    return p * z
+
+
+def dae_g(y, z):
+    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)
+```
+
+`z_0` is a guess and is made consistent automatically. Algebraic roots use
+an implicitly differentiated square LM solve, so JVP, VJP, and
+reverse-over-forward propagate with respect to `y`, `t`, and `p`. See the
+[DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
+for root controls, `SaveAt`, and scope limits.
+
 ## 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 +139,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.3.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.3.0}/docs/api.md

@@ -4,6 +4,8 @@
 
 ::: tinydiffeq.solve_ode
 
+::: tinydiffeq.solve_semi_explicit_dae
+
 ::: tinydiffeq.solve_sde
 
 ## Solvers
@@ -22,10 +24,16 @@
 
 ::: tinydiffeq.IController
 
+::: tinydiffeq.PIController
+
 ## Output selection and results
 
 ::: tinydiffeq.SaveAt
 
+::: tinydiffeq.LMRootSolver
+
+::: tinydiffeq.DAESolution
+
 ::: tinydiffeq.Solution
 
 ## Utilities

tinydiffeq-0.3.0/docs/dae.md

@@ -0,0 +1,120 @@
+# Semi-Explicit Index-1 DAEs
+
+`solve_semi_explicit_dae` integrates nonstiff systems of the form
+
+$$
+\dot y = f(y, z, t, \mathrm{args}, p), \qquad
+0 = g(y, z, t, \mathrm{args}, p),
+$$
+
+where the algebraic equation is square and $g_z$ is nonsingular along the
+solution. `y` and `z` are array states. The implementation supports fixed-step
+RK4 and Tsit5 with fixed or adaptive control.
+
+The algebraic solve uses
+[`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/square_systems/)'s
+solve-only `SquareLevenbergMarquardt`: an augmented-QR damped step for the
+primal root and a direct implicit Jacobian solve for derivatives.
+
+## Minimal examples
+
+Consider
+
+$$
+\dot y = pz, \qquad 0=z-y,
+$$
+
+whose reduced solution is $y(t)=z(t)=y_0e^{pt}$.
+
+```python
+import jax.numpy as jnp
+
+from tinydiffeq import (
+    IController,
+    RK4,
+    Tsit5,
+    solve_semi_explicit_dae,
+)
+
+
+def f(y, z, t, args, p):
+    return p * z
+
+
+def g(y, z):
+    return z - y
+
+
+fixed = solve_semi_explicit_dae(
+    f, g, RK4(), 0.0, 1.0,
+    jnp.asarray(1.0), jnp.asarray(0.5),
+    p=jnp.asarray(2.0), dt_0=0.01, max_steps=100,
+)
+
+adaptive = solve_semi_explicit_dae(
+    f, 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,
+)
+```
+
+`z_0` is a root-finding guess, not an assumed-consistent initial value. Both
+calls first solve `g(y_0, z, t_0, args, p) = 0`, so the `0.5` guess becomes
+the consistent value `1.0`.
+
+## Nonlinear-solve and AD contract
+
+The default root configuration is:
+
+```python
+from tinydiffeq import LMRootSolver
+
+root_solver = LMRootSolver(
+    max_steps=8,
+    atol=None,          # 1e-6 float32, 1e-10 float64
+    init_damping=1e-3,
+)
+```
+
+The outer `max_steps` counts attempted time steps, including adaptive
+rejections. `root_solver.max_steps` separately bounds one algebraic root.
+
+Every root passes `(y, t, p)` through nlls-gram's differentiated parameter
+pytree. Thus it differentiates the defining constraint,
+
+$$
+\dot z = -g_z^{-1}
+  (g_y\dot y + g_t\dot t + g_p\dot p),
+$$
+
+rather than differentiating the LM iterations. The warm-start guess has zero
+derivative by design. `args` is fixed data; put every differentiated model
+quantity in `p`. JVP, VJP, `vmap`, and reverse-over-forward compose through
+the complete DAE solve.
+
+An adaptive stage-root failure rejects the time-step attempt and asks the
+controller for a smaller step. A fixed-step root failure terminates the solve.
+In either case `sol.ok` is false if the endpoint is not reached with valid
+algebraic states.
+
+## Saving output
+
+All `SaveAt` modes are supported:
+
+- `SaveAt(t_1=True)` returns the consistent endpoint.
+- `SaveAt(steps=True)` returns the initial point and accepted internal steps
+  as a padded `max_steps + 1` buffer with the usual `accepted` mask.
+- `SaveAt(ts=grid)` Hermite-interpolates `y` onto the requested grid and then
+  solves `g=0` at each requested time. It never interpolates `z` independently,
+  so every returned algebraic state satisfies the constraint.
+
+The result is `DAESolution(ts, ys, zs, ok, num_accepted, accepted)`.
+
+## Deliberate limits
+
+This is an explicit, nonstiff, semi-explicit index-1 integrator. It does not
+support mass-matrix or fully implicit DAEs, stiff implicit time stepping,
+higher-index constraints, or automatic index reduction. It is an initial-value
+solver: it does not determine unknown initial costates or solve boundary-value
+or saddle-path conditions.

{tinydiffeq-0.1.0 → tinydiffeq-0.3.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
+`tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE/DAE
+integrators for JAX: fixed-step Euler and RK4, adaptive Tsit5 with integral or
+proportional-integral step-size control, fixed-step Euler–Maruyama for Itô
+SDEs, and nonstiff semi-explicit index-1 DAEs. Time integration 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
@@ -13,8 +14,8 @@ It is a jvp/vjp-friendly subset of
 need any of:**
 
 - pytree states (tinydiffeq states are single arrays, scalar or vector)
-- stiff or implicit solvers
-- PID step-size control
+- stiff or fully implicit solvers, or higher-index DAEs
+- 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,12 +55,15 @@ 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
 no special autonomous code path: an unused `t` is dead-code-eliminated.
 `drift` and `diffusion` in [`solve_sde`](sde.md) follow the same convention.
+Semi-explicit DAE fields use `(y, z)`, `(y, z, t)`,
+`(y, z, t, args)`, or `(y, z, t, args, p)`; see
+[Semi-Explicit DAEs](dae.md).
 
 ## Minimal example
 
@@ -78,22 +82,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 +116,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.3.0/docs/llms.txt

@@ -0,0 +1,12 @@
+# tinydiffeq
+
+> Tiny differentiable ODE/SDE/DAE 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. Nonstiff semi-explicit index-1 DAEs use RK4 or Tsit5 plus nlls-gram's square LM algebraic solve, whose custom JVP implicitly differentiates each root with respect to the differential state, time, and parameters. 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. For DAE observation grids, only the differential state is interpolated and the algebraic constraint is re-solved. 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 or fully implicit solvers, higher-index DAEs, 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/)
+- [Semi-explicit DAEs — index-1 contract, algebraic LM roots, implicit AD, failure behavior, and SaveAt](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
+- [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 — ODE/SDE/DAE solves, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/)

{tinydiffeq-0.1.0 → tinydiffeq-0.3.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