tinydiffeq 0.2.0__tar.gz → 0.3.0__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: tinydiffeq
-Version: 0.2.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,9 +30,9 @@ 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 integral or proportional-integral step-size control, 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
@@ -42,9 +43,9 @@ differentiates through a rollout. After a solve reaches its horizon, a
 
 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, full derivative-term 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
 
@@ -107,6 +108,38 @@ padding. Rejected attempts never appear in the returned trajectory.
 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

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

@@ -7,9 +7,9 @@
 [![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 integral or proportional-integral step-size control, 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
@@ -20,9 +20,9 @@ differentiates through a rollout. After a solve reaches its horizon, a
 
 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, full derivative-term 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
 
@@ -85,6 +85,38 @@ padding. Rejected attempts never appear in the returned trajectory.
 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

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

@@ -4,6 +4,8 @@
 
 ::: tinydiffeq.solve_ode
 
+::: tinydiffeq.solve_semi_explicit_dae
+
 ::: tinydiffeq.solve_sde
 
 ## Solvers
@@ -28,6 +30,10 @@
 
 ::: 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.2.0 → tinydiffeq-0.3.0}/docs/index.md

@@ -1,9 +1,9 @@
 # tinydiffeq
 
-`tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE integrators
-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
+`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
@@ -14,7 +14,7 @@ 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
+- 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
@@ -61,6 +61,9 @@ 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
 

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.2.0 → tinydiffeq-0.3.0}/mkdocs.yml

@@ -45,5 +45,6 @@ nav:
   - Home: index.md
   - Static Shapes: static_shapes.md
   - Adaptive Stepping and AD: adaptive_ad.md
+  - Semi-Explicit DAEs: dae.md
   - SDEs: sde.md
   - API Reference: api.md

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

@@ -4,8 +4,8 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tinydiffeq"
-version = "0.2.0"
-description = "Tiny differentiable ODE/SDE solvers for JAX — bounded-scan adaptive stepping, static shapes, jvp/vjp-safe"
+version = "0.3.0"
+description = "Tiny differentiable ODE/SDE/DAE solvers for JAX with static shapes and composable AD"
 readme = "README.md"
 license = "MIT"
 license-files = ["LICENSE"]
@@ -29,6 +29,7 @@ classifiers = [
 ]
 dependencies = [
     "jax>=0.7.0",
+    "nlls-gram>=1.6.0",
 ]
 
 [project.urls]

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

@@ -1,4 +1,4 @@
-"""Tiny differentiable ODE/SDE solvers for JAX.
+"""Tiny differentiable ODE/SDE/DAE solvers for JAX.
 
 solve_ode integrates dx/dt = f(x, t, args, p) with fixed-step (Euler, RK4)
 or adaptive (Tsit5 + IController/PIController) explicit Runge-Kutta methods
@@ -8,22 +8,26 @@ static and solves are differentiable in BOTH forward and reverse mode
 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, full derivative-term PID control, events, dense
-output, and adjoint methods are non-goals — use diffrax for those.
+diagonal noise. solve_semi_explicit_dae handles nonstiff index-1 systems with
+an implicitly differentiated algebraic root. States are arrays (scalar or
+vector); pytree states, stiff or fully implicit 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, PIController
+from tinydiffeq.dae import LMRootSolver, solve_semi_explicit_dae
 from tinydiffeq.interpolation import hermite_interpolate
 from tinydiffeq.ode import solve_ode
 from tinydiffeq.quadrature import cumulative_trapezoid
 from tinydiffeq.save_at import SaveAt
 from tinydiffeq.sde import solve_sde
-from tinydiffeq.solution import Solution
+from tinydiffeq.solution import DAESolution, Solution
 from tinydiffeq.solvers import RK4, Euler, EulerMaruyama, Tsit5
 
 __all__ = [
     "solve_ode",
+    "solve_semi_explicit_dae",
     "solve_sde",
     "Euler",
     "RK4",
@@ -34,6 +38,8 @@ __all__ = [
     "PIController",
     "SaveAt",
     "Solution",
+    "DAESolution",
+    "LMRootSolver",
     "hermite_interpolate",
     "cumulative_trapezoid",
 ]