tinydiffeq 0.3.0__tar.gz → 1.0.0__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: tinydiffeq
-Version: 0.3.0
-Summary: Tiny differentiable ODE/SDE/DAE solvers for JAX with static shapes and composable AD
+Version: 1.0.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/
@@ -10,7 +10,7 @@ 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 :: 4 - Beta
+Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -18,7 +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
+Requires-Dist: nlls-gram>=1.7.0
 Description-Content-Type: text/markdown
 
 # tinydiffeq
@@ -30,9 +30,10 @@ 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/DAE solvers for JAX: fixed-step Euler/RK4,
+Tiny differentiable ODE/SDE/DAE/SDAE 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.
+Euler–Maruyama for Itô SDEs, and nonstiff semi-explicit index-1 deterministic
+and stochastic 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,10 +43,10 @@ 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 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`.
+[diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need stiff
+or fully implicit solvers, higher-index DAEs, full
+derivative-term PID control, events, continuous interpolation objects, or
+checkpointed/backsolve adjoints. The DAE algebraic solve uses `nlls-gram`.
 
 ## Install
 
@@ -65,6 +66,11 @@ uv add tinydiffeq "jax[cuda13]"
 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
@@ -121,25 +127,31 @@ def dae_f(y, z, t, args, p):
     return p * z
 
 
-def dae_g(y, z):
-    return z - y
+def dae_g(y, z, t, args, p):
+    return z - y, {"flow": p * z}
 
 
 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,
+    controller=IController(), max_steps=128, has_aux=True,
 )
-print(dae_sol.ys, dae_sol.zs)
+print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])
 ```
 
-`z_0` is a guess and is made consistent automatically. Algebraic roots use
-an implicitly differentiated square LM solve, so JVP, VJP, and
+`z_0` is a guess and is made consistent automatically. Algebraic equations
+may return a floating aux pytree stored at every accepted node and
+Hermite-interpolated with `z` on requested deterministic grids. 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.
 
+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

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

@@ -7,9 +7,10 @@
 [![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 solvers for JAX: fixed-step Euler/RK4,
+Tiny differentiable ODE/SDE/DAE/SDAE 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.
+Euler–Maruyama for Itô SDEs, and nonstiff semi-explicit index-1 deterministic
+and stochastic 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
@@ -19,10 +20,10 @@ 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 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`.
+[diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need stiff
+or fully implicit solvers, higher-index DAEs, full
+derivative-term PID control, events, continuous interpolation objects, or
+checkpointed/backsolve adjoints. The DAE algebraic solve uses `nlls-gram`.
 
 ## Install
 
@@ -42,6 +43,11 @@ uv add tinydiffeq "jax[cuda13]"
 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
@@ -98,25 +104,31 @@ def dae_f(y, z, t, args, p):
     return p * z
 
 
-def dae_g(y, z):
-    return z - y
+def dae_g(y, z, t, args, p):
+    return z - y, {"flow": p * z}
 
 
 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,
+    controller=IController(), max_steps=128, has_aux=True,
 )
-print(dae_sol.ys, dae_sol.zs)
+print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])
 ```
 
-`z_0` is a guess and is made consistent automatically. Algebraic roots use
-an implicitly differentiated square LM solve, so JVP, VJP, and
+`z_0` is a guess and is made consistent automatically. Algebraic equations
+may return a floating aux pytree stored at every accepted node and
+Hermite-interpolated with `z` on requested deterministic grids. 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.
 
+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

tinydiffeq-1.0.0/benchmarks/README.md

@@ -0,0 +1,38 @@
+# Benchmarks
+
+Run the opt-in CPU suite with:
+
+```bash
+JAX_PLATFORMS=cpu uv run --group benchmark pytest benchmarks --benchmark-only
+```
+
+It measures post-compilation primal, JVP, and VJP execution for RK4, Tsit5,
+Euler-Maruyama, and the semi-explicit DAE/SDAE solvers. Each method is
+exercised on a scalar array, a length-16 array, and an equal-sized two-leaf
+pytree. The array cases are the performance-regression baseline; pytree
+timings show the cost of executing separate leaves. Compilation is
+deliberately performed outside the timed region.
+
+Measure cold compilation separately (optionally selecting a subset) with:
+
+```bash
+JAX_PLATFORMS=cpu uv run --group benchmark python -m benchmarks.compile_times
+JAX_PLATFORMS=cpu uv run --group benchmark python -m benchmarks.compile_times \
+  --methods rk4 tsit5 --states vector16 --repeat 5
+```
+
+The script clears JAX's compilation caches between `timeit` repetitions and
+reports the median. Keep compilation and execution results separate.
+
+`test_aux_dense_output.py` separately measures the intentionally changed DAE
+grid-output path across accepted-step and query-grid counts, both with and
+without aux, plus batched adaptive interpolation. The no-aux cases can run
+against an older checkout for a direct cross-version comparison. Those
+results are re-baselined rather than held to the unchanged endpoint/steps
+array-path gate.
+
+For changes to the state arithmetic, compare the JSON outputs before and
+after the change using `--benchmark-json=PATH`. Treat an array-path slowdown
+larger than `max(5%, 1 us)` as a regression. GPU timing is informative rather
+than a release gate because available devices vary; GPU correctness belongs
+in `tests/test_gpu.py`.

tinydiffeq-1.0.0/benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Opt-in tinydiffeq benchmarks."""

tinydiffeq-1.0.0/benchmarks/compile_times.py

@@ -0,0 +1,64 @@
+"""Cold JAX compilation timings for the pytree benchmark matrix."""
+
+import argparse
+import statistics
+import timeit
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from benchmarks.test_pytree_overhead import (
+    make_solve,
+    make_state,
+    tree_map,
+    tree_sum,
+)
+
+
+def transformed_solve(method, state_kind, transform):
+    initial = make_state(state_kind)
+    solve = make_solve(method, initial)
+    tangent = tree_map(jnp.ones_like, initial)
+    if transform == "primal":
+        run = jax.jit(solve)
+    elif transform == "jvp":
+        run = jax.jit(lambda x: jax.jvp(solve, (x,), (tangent,)))
+    else:
+        run = jax.jit(lambda x: jax.grad(lambda value: tree_sum(solve(value)))(x))
+    return run, initial
+
+
+def compile_once(run, initial):
+    jax.clear_caches()
+    run.lower(initial).compile()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--methods", nargs="+", default=["rk4", "tsit5", "em", "dae", "sdae"]
+    )
+    parser.add_argument("--states", nargs="+", default=["scalar", "vector16", "tree16"])
+    parser.add_argument("--transforms", nargs="+", default=["primal", "jvp", "vjp"])
+    parser.add_argument("--repeat", type=int, default=3)
+    args = parser.parse_args()
+
+    print("method,state,transform,median_compile_seconds")
+    for method in args.methods:
+        for state_kind in args.states:
+            for transform in args.transforms:
+                run, initial = transformed_solve(method, state_kind, transform)
+                samples = timeit.repeat(
+                    partial(compile_once, run, initial),
+                    repeat=args.repeat,
+                    number=1,
+                )
+                print(
+                    f"{method},{state_kind},{transform},"
+                    f"{statistics.median(samples):.6f}"
+                )
+
+
+if __name__ == "__main__":
+    main()

tinydiffeq-1.0.0/benchmarks/test_aux_dense_output.py

@@ -0,0 +1,159 @@
+"""Opt-in DAE aux/dense-output and batched interpolation benchmarks."""
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from tinydiffeq import (
+    RK4,
+    ConstantStepSize,
+    IController,
+    LMRootSolver,
+    SaveAt,
+    Tsit5,
+    solve_semi_explicit_dae,
+)
+
+
+def make_solve(n_steps, n_queries):
+    grid = jnp.linspace(0.0, 1.0, n_queries)
+
+    def run(rate):
+        return solve_semi_explicit_dae(
+            lambda y, z, t, args, p: p * z,
+            lambda y, z, t, args, p: (
+                z - y**2 - 0.1 * t,
+                {"flow": p * z, "level": y + z},
+            ),
+            RK4(),
+            0.0,
+            1.0,
+            jnp.asarray(0.8),
+            jnp.asarray(0.6),
+            p=rate,
+            dt_0=1.0 / n_steps,
+            controller=ConstantStepSize(),
+            root_solver=LMRootSolver(),
+            max_steps=n_steps,
+            save_at=SaveAt(ts=grid),
+            has_aux=True,
+        )
+
+    return run
+
+
+def make_adaptive_solve(n_queries):
+    grid = jnp.linspace(0.0, 1.0, n_queries)
+
+    def run(rate):
+        return solve_semi_explicit_dae(
+            lambda y, z, t, args, p: p * z,
+            lambda y, z, t, args, p: (
+                z - y**2 - 0.1 * t,
+                {"flow": p * z, "level": y + z},
+            ),
+            Tsit5(),
+            0.0,
+            1.0,
+            jnp.asarray(0.8),
+            jnp.asarray(0.6),
+            p=rate,
+            dt_0=0.1,
+            controller=IController(),
+            root_solver=LMRootSolver(),
+            max_steps=128,
+            save_at=SaveAt(ts=grid),
+            has_aux=True,
+        )
+
+    return run
+
+
+def make_no_aux_solve(n_steps, n_queries):
+    grid = jnp.linspace(0.0, 1.0, n_queries)
+
+    def run(rate):
+        return solve_semi_explicit_dae(
+            lambda y, z, t, args, p: p * z,
+            lambda y, z, t, args, p: z - y**2 - 0.1 * t,
+            RK4(),
+            0.0,
+            1.0,
+            jnp.asarray(0.8),
+            jnp.asarray(0.6),
+            p=rate,
+            dt_0=1.0 / n_steps,
+            controller=ConstantStepSize(),
+            root_solver=LMRootSolver(),
+            max_steps=n_steps,
+            save_at=SaveAt(ts=grid),
+        )
+
+    return run
+
+
+@pytest.mark.parametrize("n_steps", [8, 32, 128])
+@pytest.mark.parametrize("n_queries", [8, 32, 128])
+@pytest.mark.parametrize("transform", ["primal", "jvp", "vjp"])
+def test_aux_dense_output_crossover(benchmark, n_steps, n_queries, transform):
+    solve = make_solve(n_steps, n_queries)
+    rate = jnp.asarray(0.2)
+    assert bool(solve(rate).ok)
+
+    def scalar_output(value):
+        sol = solve(value)
+        return jnp.sum(sol.ys + sol.zs + sol.aux["flow"] + sol.aux["level"])
+
+    if transform == "primal":
+        run = jax.jit(solve)
+    elif transform == "jvp":
+        run = jax.jit(
+            lambda value: jax.jvp(scalar_output, (value,), (jnp.ones_like(value),))
+        )
+    else:
+        run = jax.jit(jax.grad(scalar_output))
+    compiled = run.lower(rate).compile()
+    jax.block_until_ready(compiled(rate))
+    benchmark.pedantic(
+        lambda: jax.block_until_ready(compiled(rate)), rounds=10, iterations=5
+    )
+
+
+@pytest.mark.parametrize("n_steps", [8, 128])
+@pytest.mark.parametrize("n_queries", [8, 128])
+@pytest.mark.parametrize("transform", ["primal", "jvp", "vjp"])
+def test_no_aux_grid_output_crossover(benchmark, n_steps, n_queries, transform):
+    """Cross-version baseline for the intentionally changed DAE grid path."""
+    solve = make_no_aux_solve(n_steps, n_queries)
+    rate = jnp.asarray(0.2)
+    assert bool(solve(rate).ok)
+
+    def scalar_output(value):
+        sol = solve(value)
+        return jnp.sum(sol.ys + sol.zs)
+
+    if transform == "primal":
+        run = jax.jit(solve)
+    elif transform == "jvp":
+        run = jax.jit(
+            lambda value: jax.jvp(scalar_output, (value,), (jnp.ones_like(value),))
+        )
+    else:
+        run = jax.jit(jax.grad(scalar_output))
+    compiled = run.lower(rate).compile()
+    jax.block_until_ready(compiled(rate))
+    benchmark.pedantic(
+        lambda: jax.block_until_ready(compiled(rate)), rounds=10, iterations=5
+    )
+
+
+def test_batched_adaptive_dense_output(benchmark):
+    solve = make_adaptive_solve(32)
+    rates = jnp.linspace(0.1, 1.0, 16)
+    assert bool(jnp.all(jax.vmap(solve)(rates).ok))
+    run = jax.jit(jax.vmap(lambda value: solve(value).aux["flow"]))
+    compiled = run.lower(rates).compile()
+    jax.block_until_ready(compiled(rates))
+    benchmark.pedantic(
+        lambda: jax.block_until_ready(compiled(rates)), rounds=10, iterations=5
+    )

tinydiffeq-1.0.0/benchmarks/test_pytree_overhead.py

@@ -0,0 +1,138 @@
+"""Opt-in steady-state array/pytree and AD benchmarks.
+
+Run with ``uv run --group benchmark pytest benchmarks --benchmark-only``.
+"""
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+from tinydiffeq import (
+    RK4,
+    EulerMaruyama,
+    IController,
+    Tsit5,
+    solve_ode,
+    solve_sde,
+    solve_semi_explicit_dae,
+    solve_semi_explicit_sdae,
+)
+
+
+def make_state(kind):
+    if kind == "scalar":
+        return jnp.asarray(1.0)
+    if kind == "vector16":
+        return jnp.linspace(0.5, 1.5, 16)
+    return {"left": jnp.linspace(0.5, 1.0, 8), "right": jnp.linspace(1.0, 1.5, 8)}
+
+
+def tree_map(fn, tree):
+    return jax.tree.map(fn, tree)
+
+
+def tree_sum(tree):
+    return sum(jnp.sum(leaf) for leaf in jax.tree.leaves(tree))
+
+
+def make_solve(method, initial):
+    if method == "rk4":
+        return lambda x: (
+            solve_ode(
+                lambda state: tree_map(lambda leaf: -0.2 * leaf, state),
+                RK4(),
+                0.0,
+                1.0,
+                x,
+                dt_0=1 / 64,
+                max_steps=64,
+            ).xs
+        )
+    if method == "tsit5":
+        return lambda x: (
+            solve_ode(
+                lambda state: tree_map(lambda leaf: -0.2 * leaf, state),
+                Tsit5(),
+                0.0,
+                1.0,
+                x,
+                dt_0=0.05,
+                controller=IController(),
+                max_steps=64,
+            ).xs
+        )
+    if method == "em":
+        return lambda x: (
+            solve_sde(
+                lambda state: tree_map(lambda leaf: -0.2 * leaf, state),
+                lambda state: tree_map(lambda leaf: 0.1 * jnp.ones_like(leaf), state),
+                EulerMaruyama(),
+                0.0,
+                1.0,
+                x,
+                key=jax.random.key(7),
+                n_steps=64,
+            ).xs
+        )
+
+    z_0 = tree_map(lambda leaf: 0.9 * leaf, initial)
+
+    def constraint(y, z):
+        y_flat = jnp.concatenate([jnp.ravel(leaf) for leaf in jax.tree.leaves(y)])
+        z_flat = jnp.concatenate([jnp.ravel(leaf) for leaf in jax.tree.leaves(z)])
+        return z_flat - y_flat
+
+    if method == "sdae":
+        return lambda y: (
+            solve_semi_explicit_sdae(
+                lambda state, z: tree_map(lambda leaf: -0.2 * leaf, z),
+                lambda state, z: tree_map(
+                    lambda leaf: 0.1 * jnp.ones_like(leaf), state
+                ),
+                constraint,
+                EulerMaruyama(),
+                0.0,
+                1.0,
+                y,
+                z_0,
+                key=jax.random.key(7),
+                n_steps=64,
+            ).ys
+        )
+
+    return lambda y: (
+        solve_semi_explicit_dae(
+            lambda state, z: tree_map(lambda leaf: -0.2 * leaf, z),
+            constraint,
+            Tsit5(),
+            0.0,
+            1.0,
+            y,
+            z_0,
+            dt_0=0.1,
+            controller=IController(),
+            max_steps=32,
+        ).ys
+    )
+
+
+@pytest.mark.parametrize("method", ["rk4", "tsit5", "em", "dae", "sdae"])
+@pytest.mark.parametrize("state_kind", ["scalar", "vector16", "tree16"])
+@pytest.mark.parametrize("transform", ["primal", "jvp", "vjp"])
+def test_steady_state(benchmark, method, state_kind, transform):
+    initial = make_state(state_kind)
+    solve = make_solve(method, initial)
+    tangent = tree_map(jnp.ones_like, initial)
+
+    if transform == "primal":
+        run = jax.jit(solve)
+    elif transform == "jvp":
+        run = jax.jit(lambda x: jax.jvp(solve, (x,), (tangent,)))
+    else:
+        run = jax.jit(lambda x: jax.grad(lambda value: tree_sum(solve(value)))(x))
+
+    compiled = run.lower(initial).compile()
+    jax.block_until_ready(compiled(initial))
+    benchmark.pedantic(
+        lambda: jax.block_until_ready(compiled(initial)), rounds=20, iterations=10
+    )

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

@@ -6,6 +6,8 @@
 
 ::: tinydiffeq.solve_semi_explicit_dae
 
+::: tinydiffeq.solve_semi_explicit_sdae
+
 ::: tinydiffeq.solve_sde
 
 ## Solvers