tinydiffeq 0.2.0__tar.gz → 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (56) hide show
  1. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/PKG-INFO +55 -10
  2. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/README.md +51 -7
  3. tinydiffeq-1.0.0/benchmarks/README.md +38 -0
  4. tinydiffeq-1.0.0/benchmarks/__init__.py +1 -0
  5. tinydiffeq-1.0.0/benchmarks/compile_times.py +64 -0
  6. tinydiffeq-1.0.0/benchmarks/test_aux_dense_output.py +159 -0
  7. tinydiffeq-1.0.0/benchmarks/test_pytree_overhead.py +138 -0
  8. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/api.md +8 -0
  9. tinydiffeq-1.0.0/docs/dae.md +184 -0
  10. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/index.md +20 -12
  11. tinydiffeq-1.0.0/docs/llms.txt +13 -0
  12. tinydiffeq-1.0.0/docs/sdae.md +99 -0
  13. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/sde.md +10 -4
  14. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/static_shapes.md +17 -2
  15. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/mkdocs.yml +3 -1
  16. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/pyproject.toml +4 -3
  17. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/__init__.py +15 -5
  18. tinydiffeq-1.0.0/src/tinydiffeq/_tree.py +122 -0
  19. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/controllers.py +12 -12
  20. tinydiffeq-1.0.0/src/tinydiffeq/dae.py +872 -0
  21. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/interpolation.py +22 -12
  22. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/ode.py +38 -18
  23. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/save_at.py +1 -1
  24. tinydiffeq-1.0.0/src/tinydiffeq/sdae.py +248 -0
  25. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/sde.py +50 -13
  26. tinydiffeq-1.0.0/src/tinydiffeq/solution.py +52 -0
  27. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/solvers.py +52 -22
  28. tinydiffeq-1.0.0/tests/test_aux.py +319 -0
  29. tinydiffeq-1.0.0/tests/test_dae.py +524 -0
  30. tinydiffeq-1.0.0/tests/test_float64_subprocess.py +488 -0
  31. tinydiffeq-1.0.0/tests/test_gpu.py +211 -0
  32. tinydiffeq-1.0.0/tests/test_pytree_states.py +269 -0
  33. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_save_at.py +1 -1
  34. tinydiffeq-1.0.0/tests/test_sdae.py +286 -0
  35. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/uv.lock +18 -2
  36. tinydiffeq-0.2.0/docs/llms.txt +0 -11
  37. tinydiffeq-0.2.0/src/tinydiffeq/solution.py +0 -28
  38. tinydiffeq-0.2.0/tests/test_float64_subprocess.py +0 -244
  39. tinydiffeq-0.2.0/tests/test_gpu.py +0 -75
  40. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.githooks/pre-commit +0 -0
  41. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.github/workflows/ci.yml +0 -0
  42. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.github/workflows/docs.yml +0 -0
  43. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.github/workflows/publish.yml +0 -0
  44. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.gitignore +0 -0
  45. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.python-version +0 -0
  46. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/LICENSE +0 -0
  47. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/adaptive_ad.md +0 -0
  48. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/quadrature.py +0 -0
  49. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/conftest.py +0 -0
  50. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/kernels_reference.py +0 -0
  51. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_ad.py +0 -0
  52. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_adaptive.py +0 -0
  53. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_quadrature.py +0 -0
  54. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_recompile.py +0 -0
  55. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_sde.py +0 -0
  56. {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_solvers_fixed.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tinydiffeq
3
- Version: 0.2.0
4
- Summary: Tiny differentiable ODE/SDE solvers for JAX bounded-scan adaptive stepping, static shapes, jvp/vjp-safe
3
+ Version: 1.0.0
4
+ Summary: Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX with static shapes and composable AD
5
5
  Project-URL: Homepage, https://github.com/HighDimensionalEconLab/tinydiffeq
6
6
  Project-URL: Repository, https://github.com/HighDimensionalEconLab/tinydiffeq
7
7
  Project-URL: Documentation, https://highdimensionaleconlab.github.io/tinydiffeq/
@@ -10,7 +10,7 @@ Author-email: Jesse Perla <jesseperla@gmail.com>
10
10
  License-Expression: MIT
11
11
  License-File: LICENSE
12
12
  Keywords: autodiff,differential-equations,jax,ode,runge-kutta,sde
13
- Classifier: Development Status :: 4 - Beta
13
+ Classifier: Development Status :: 5 - Production/Stable
14
14
  Classifier: Intended Audience :: Science/Research
15
15
  Classifier: Programming Language :: Python :: 3.11
16
16
  Classifier: Programming Language :: Python :: 3.12
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.13
18
18
  Classifier: Topic :: Scientific/Engineering :: Mathematics
19
19
  Requires-Python: >=3.11
20
20
  Requires-Dist: jax>=0.7.0
21
+ Requires-Dist: nlls-gram>=1.7.0
21
22
  Description-Content-Type: text/markdown
22
23
 
23
24
  # tinydiffeq
@@ -29,9 +30,10 @@ Description-Content-Type: text/markdown
29
30
  [![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
30
31
  [![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)
31
32
 
32
- Tiny differentiable ODE/SDE solvers for JAX: fixed-step Euler/RK4, adaptive
33
- Tsit5 with integral or proportional-integral step-size control, and
34
- Euler–Maruyama for Itô SDEs.
33
+ Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4,
34
+ adaptive Tsit5 with integral or proportional-integral step-size control,
35
+ Euler–Maruyama for Itô SDEs, and nonstiff semi-explicit index-1 deterministic
36
+ and stochastic DAEs.
35
37
  One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
36
38
  adaptive stepping, so shapes are static, nothing recompiles as tolerances or
37
39
  curvature change, and every solve is differentiable in **both** forward and
@@ -41,10 +43,10 @@ differentiates through a rollout. After a solve reaches its horizon, a
41
43
  `lax.cond` skips solver and controller work during the padded scan tail.
42
44
 
43
45
  This is a deliberately small, jvp/vjp-friendly subset of
44
- [diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
45
- states, stiff/implicit solvers, full derivative-term PID control, events,
46
- dense output, or checkpointed/backsolve adjoints. tinydiffeq's single runtime
47
- dependency is `jax`.
46
+ [diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need stiff
47
+ or fully implicit solvers, higher-index DAEs, full
48
+ derivative-term PID control, events, continuous interpolation objects, or
49
+ checkpointed/backsolve adjoints. The DAE algebraic solve uses `nlls-gram`.
48
50
 
49
51
  ## Install
50
52
 
@@ -64,6 +66,11 @@ uv add tinydiffeq "jax[cuda13]"
64
66
  The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or
65
67
  `(x, t, args, p)` — always in that order. `args` is pass-through data (not an
66
68
  AD target by convention); `p` holds differentiable parameters (any pytree).
69
+ The state may also be any JAX pytree. It must contain at least one leaf, and
70
+ every leaf must be a nonempty real floating array with the same dtype; vector
71
+ fields and `project` preserve that structure. Output keeps the structure and
72
+ adds the saved-time axis to each
73
+ leaf.
67
74
 
68
75
  ```python
69
76
  import jax
@@ -107,6 +114,44 @@ padding. Rejected attempts never appear in the returned trajectory.
107
114
  the adaptive controller still chooses its own internal mesh, and cubic
108
115
  Hermite interpolation evaluates the solution at every requested point.
109
116
 
117
+ ## Semi-explicit DAEs
118
+
119
+ For a square index-1 system `dy/dt = f(y, z, t, args, p)` and
120
+ `0 = g(y, z, t, args, p)`:
121
+
122
+ ```python
123
+ from tinydiffeq import IController, Tsit5, solve_semi_explicit_dae
124
+
125
+
126
+ def dae_f(y, z, t, args, p):
127
+ return p * z
128
+
129
+
130
+ def dae_g(y, z, t, args, p):
131
+ return z - y, {"flow": p * z}
132
+
133
+
134
+ dae_sol = solve_semi_explicit_dae(
135
+ dae_f, dae_g, Tsit5(), 0.0, 1.0,
136
+ jnp.asarray(1.0), jnp.asarray(0.5),
137
+ p=jnp.asarray(2.0), dt_0=0.1,
138
+ controller=IController(), max_steps=128, has_aux=True,
139
+ )
140
+ print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])
141
+ ```
142
+
143
+ `z_0` is a guess and is made consistent automatically. Algebraic equations
144
+ may return a floating aux pytree stored at every accepted node and
145
+ Hermite-interpolated with `z` on requested deterministic grids. Algebraic
146
+ roots use an implicitly differentiated square LM solve, so JVP, VJP, and
147
+ reverse-over-forward propagate with respect to `y`, `t`, and `p`. See the
148
+ [DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
149
+ for root controls, `SaveAt`, and scope limits.
150
+
151
+ Fixed-step semi-explicit Itô SDAEs use the corresponding
152
+ `solve_semi_explicit_sdae` interface with `EulerMaruyama`, a PRNG key, and
153
+ `n_steps`; see the [SDAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/).
154
+
110
155
  ## Gradients through the solve
111
156
 
112
157
  ```python
@@ -7,9 +7,10 @@
7
7
  [![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
8
8
  [![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)
9
9
 
10
- Tiny differentiable ODE/SDE solvers for JAX: fixed-step Euler/RK4, adaptive
11
- Tsit5 with integral or proportional-integral step-size control, and
12
- Euler–Maruyama for Itô SDEs.
10
+ Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4,
11
+ adaptive Tsit5 with integral or proportional-integral step-size control,
12
+ Euler–Maruyama for Itô SDEs, and nonstiff semi-explicit index-1 deterministic
13
+ and stochastic DAEs.
13
14
  One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
14
15
  adaptive stepping, so shapes are static, nothing recompiles as tolerances or
15
16
  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
19
20
  `lax.cond` skips solver and controller work during the padded scan tail.
20
21
 
21
22
  This is a deliberately small, jvp/vjp-friendly subset of
22
- [diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
23
- states, stiff/implicit solvers, full derivative-term PID control, events,
24
- dense output, or checkpointed/backsolve adjoints. tinydiffeq's single runtime
25
- dependency is `jax`.
23
+ [diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need stiff
24
+ or fully implicit solvers, higher-index DAEs, full
25
+ derivative-term PID control, events, continuous interpolation objects, or
26
+ checkpointed/backsolve adjoints. The DAE algebraic solve uses `nlls-gram`.
26
27
 
27
28
  ## Install
28
29
 
@@ -42,6 +43,11 @@ uv add tinydiffeq "jax[cuda13]"
42
43
  The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or
43
44
  `(x, t, args, p)` — always in that order. `args` is pass-through data (not an
44
45
  AD target by convention); `p` holds differentiable parameters (any pytree).
46
+ The state may also be any JAX pytree. It must contain at least one leaf, and
47
+ every leaf must be a nonempty real floating array with the same dtype; vector
48
+ fields and `project` preserve that structure. Output keeps the structure and
49
+ adds the saved-time axis to each
50
+ leaf.
45
51
 
46
52
  ```python
47
53
  import jax
@@ -85,6 +91,44 @@ padding. Rejected attempts never appear in the returned trajectory.
85
91
  the adaptive controller still chooses its own internal mesh, and cubic
86
92
  Hermite interpolation evaluates the solution at every requested point.
87
93
 
94
+ ## Semi-explicit DAEs
95
+
96
+ For a square index-1 system `dy/dt = f(y, z, t, args, p)` and
97
+ `0 = g(y, z, t, args, p)`:
98
+
99
+ ```python
100
+ from tinydiffeq import IController, Tsit5, solve_semi_explicit_dae
101
+
102
+
103
+ def dae_f(y, z, t, args, p):
104
+ return p * z
105
+
106
+
107
+ def dae_g(y, z, t, args, p):
108
+ return z - y, {"flow": p * z}
109
+
110
+
111
+ dae_sol = solve_semi_explicit_dae(
112
+ dae_f, dae_g, Tsit5(), 0.0, 1.0,
113
+ jnp.asarray(1.0), jnp.asarray(0.5),
114
+ p=jnp.asarray(2.0), dt_0=0.1,
115
+ controller=IController(), max_steps=128, has_aux=True,
116
+ )
117
+ print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])
118
+ ```
119
+
120
+ `z_0` is a guess and is made consistent automatically. Algebraic equations
121
+ may return a floating aux pytree stored at every accepted node and
122
+ Hermite-interpolated with `z` on requested deterministic grids. Algebraic
123
+ roots use an implicitly differentiated square LM solve, so JVP, VJP, and
124
+ reverse-over-forward propagate with respect to `y`, `t`, and `p`. See the
125
+ [DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
126
+ for root controls, `SaveAt`, and scope limits.
127
+
128
+ Fixed-step semi-explicit Itô SDAEs use the corresponding
129
+ `solve_semi_explicit_sdae` interface with `EulerMaruyama`, a PRNG key, and
130
+ `n_steps`; see the [SDAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/).
131
+
88
132
  ## Gradients through the solve
89
133
 
90
134
  ```python
@@ -0,0 +1,38 @@
1
+ # Benchmarks
2
+
3
+ Run the opt-in CPU suite with:
4
+
5
+ ```bash
6
+ JAX_PLATFORMS=cpu uv run --group benchmark pytest benchmarks --benchmark-only
7
+ ```
8
+
9
+ It measures post-compilation primal, JVP, and VJP execution for RK4, Tsit5,
10
+ Euler-Maruyama, and the semi-explicit DAE/SDAE solvers. Each method is
11
+ exercised on a scalar array, a length-16 array, and an equal-sized two-leaf
12
+ pytree. The array cases are the performance-regression baseline; pytree
13
+ timings show the cost of executing separate leaves. Compilation is
14
+ deliberately performed outside the timed region.
15
+
16
+ Measure cold compilation separately (optionally selecting a subset) with:
17
+
18
+ ```bash
19
+ JAX_PLATFORMS=cpu uv run --group benchmark python -m benchmarks.compile_times
20
+ JAX_PLATFORMS=cpu uv run --group benchmark python -m benchmarks.compile_times \
21
+ --methods rk4 tsit5 --states vector16 --repeat 5
22
+ ```
23
+
24
+ The script clears JAX's compilation caches between `timeit` repetitions and
25
+ reports the median. Keep compilation and execution results separate.
26
+
27
+ `test_aux_dense_output.py` separately measures the intentionally changed DAE
28
+ grid-output path across accepted-step and query-grid counts, both with and
29
+ without aux, plus batched adaptive interpolation. The no-aux cases can run
30
+ against an older checkout for a direct cross-version comparison. Those
31
+ results are re-baselined rather than held to the unchanged endpoint/steps
32
+ array-path gate.
33
+
34
+ For changes to the state arithmetic, compare the JSON outputs before and
35
+ after the change using `--benchmark-json=PATH`. Treat an array-path slowdown
36
+ larger than `max(5%, 1 us)` as a regression. GPU timing is informative rather
37
+ than a release gate because available devices vary; GPU correctness belongs
38
+ in `tests/test_gpu.py`.
@@ -0,0 +1 @@
1
+ """Opt-in tinydiffeq benchmarks."""
@@ -0,0 +1,64 @@
1
+ """Cold JAX compilation timings for the pytree benchmark matrix."""
2
+
3
+ import argparse
4
+ import statistics
5
+ import timeit
6
+ from functools import partial
7
+
8
+ import jax
9
+ import jax.numpy as jnp
10
+
11
+ from benchmarks.test_pytree_overhead import (
12
+ make_solve,
13
+ make_state,
14
+ tree_map,
15
+ tree_sum,
16
+ )
17
+
18
+
19
+ def transformed_solve(method, state_kind, transform):
20
+ initial = make_state(state_kind)
21
+ solve = make_solve(method, initial)
22
+ tangent = tree_map(jnp.ones_like, initial)
23
+ if transform == "primal":
24
+ run = jax.jit(solve)
25
+ elif transform == "jvp":
26
+ run = jax.jit(lambda x: jax.jvp(solve, (x,), (tangent,)))
27
+ else:
28
+ run = jax.jit(lambda x: jax.grad(lambda value: tree_sum(solve(value)))(x))
29
+ return run, initial
30
+
31
+
32
+ def compile_once(run, initial):
33
+ jax.clear_caches()
34
+ run.lower(initial).compile()
35
+
36
+
37
+ def main():
38
+ parser = argparse.ArgumentParser()
39
+ parser.add_argument(
40
+ "--methods", nargs="+", default=["rk4", "tsit5", "em", "dae", "sdae"]
41
+ )
42
+ parser.add_argument("--states", nargs="+", default=["scalar", "vector16", "tree16"])
43
+ parser.add_argument("--transforms", nargs="+", default=["primal", "jvp", "vjp"])
44
+ parser.add_argument("--repeat", type=int, default=3)
45
+ args = parser.parse_args()
46
+
47
+ print("method,state,transform,median_compile_seconds")
48
+ for method in args.methods:
49
+ for state_kind in args.states:
50
+ for transform in args.transforms:
51
+ run, initial = transformed_solve(method, state_kind, transform)
52
+ samples = timeit.repeat(
53
+ partial(compile_once, run, initial),
54
+ repeat=args.repeat,
55
+ number=1,
56
+ )
57
+ print(
58
+ f"{method},{state_kind},{transform},"
59
+ f"{statistics.median(samples):.6f}"
60
+ )
61
+
62
+
63
+ if __name__ == "__main__":
64
+ main()
@@ -0,0 +1,159 @@
1
+ """Opt-in DAE aux/dense-output and batched interpolation benchmarks."""
2
+
3
+ import jax
4
+ import jax.numpy as jnp
5
+ import pytest
6
+
7
+ from tinydiffeq import (
8
+ RK4,
9
+ ConstantStepSize,
10
+ IController,
11
+ LMRootSolver,
12
+ SaveAt,
13
+ Tsit5,
14
+ solve_semi_explicit_dae,
15
+ )
16
+
17
+
18
+ def make_solve(n_steps, n_queries):
19
+ grid = jnp.linspace(0.0, 1.0, n_queries)
20
+
21
+ def run(rate):
22
+ return solve_semi_explicit_dae(
23
+ lambda y, z, t, args, p: p * z,
24
+ lambda y, z, t, args, p: (
25
+ z - y**2 - 0.1 * t,
26
+ {"flow": p * z, "level": y + z},
27
+ ),
28
+ RK4(),
29
+ 0.0,
30
+ 1.0,
31
+ jnp.asarray(0.8),
32
+ jnp.asarray(0.6),
33
+ p=rate,
34
+ dt_0=1.0 / n_steps,
35
+ controller=ConstantStepSize(),
36
+ root_solver=LMRootSolver(),
37
+ max_steps=n_steps,
38
+ save_at=SaveAt(ts=grid),
39
+ has_aux=True,
40
+ )
41
+
42
+ return run
43
+
44
+
45
+ def make_adaptive_solve(n_queries):
46
+ grid = jnp.linspace(0.0, 1.0, n_queries)
47
+
48
+ def run(rate):
49
+ return solve_semi_explicit_dae(
50
+ lambda y, z, t, args, p: p * z,
51
+ lambda y, z, t, args, p: (
52
+ z - y**2 - 0.1 * t,
53
+ {"flow": p * z, "level": y + z},
54
+ ),
55
+ Tsit5(),
56
+ 0.0,
57
+ 1.0,
58
+ jnp.asarray(0.8),
59
+ jnp.asarray(0.6),
60
+ p=rate,
61
+ dt_0=0.1,
62
+ controller=IController(),
63
+ root_solver=LMRootSolver(),
64
+ max_steps=128,
65
+ save_at=SaveAt(ts=grid),
66
+ has_aux=True,
67
+ )
68
+
69
+ return run
70
+
71
+
72
+ def make_no_aux_solve(n_steps, n_queries):
73
+ grid = jnp.linspace(0.0, 1.0, n_queries)
74
+
75
+ def run(rate):
76
+ return solve_semi_explicit_dae(
77
+ lambda y, z, t, args, p: p * z,
78
+ lambda y, z, t, args, p: z - y**2 - 0.1 * t,
79
+ RK4(),
80
+ 0.0,
81
+ 1.0,
82
+ jnp.asarray(0.8),
83
+ jnp.asarray(0.6),
84
+ p=rate,
85
+ dt_0=1.0 / n_steps,
86
+ controller=ConstantStepSize(),
87
+ root_solver=LMRootSolver(),
88
+ max_steps=n_steps,
89
+ save_at=SaveAt(ts=grid),
90
+ )
91
+
92
+ return run
93
+
94
+
95
+ @pytest.mark.parametrize("n_steps", [8, 32, 128])
96
+ @pytest.mark.parametrize("n_queries", [8, 32, 128])
97
+ @pytest.mark.parametrize("transform", ["primal", "jvp", "vjp"])
98
+ def test_aux_dense_output_crossover(benchmark, n_steps, n_queries, transform):
99
+ solve = make_solve(n_steps, n_queries)
100
+ rate = jnp.asarray(0.2)
101
+ assert bool(solve(rate).ok)
102
+
103
+ def scalar_output(value):
104
+ sol = solve(value)
105
+ return jnp.sum(sol.ys + sol.zs + sol.aux["flow"] + sol.aux["level"])
106
+
107
+ if transform == "primal":
108
+ run = jax.jit(solve)
109
+ elif transform == "jvp":
110
+ run = jax.jit(
111
+ lambda value: jax.jvp(scalar_output, (value,), (jnp.ones_like(value),))
112
+ )
113
+ else:
114
+ run = jax.jit(jax.grad(scalar_output))
115
+ compiled = run.lower(rate).compile()
116
+ jax.block_until_ready(compiled(rate))
117
+ benchmark.pedantic(
118
+ lambda: jax.block_until_ready(compiled(rate)), rounds=10, iterations=5
119
+ )
120
+
121
+
122
+ @pytest.mark.parametrize("n_steps", [8, 128])
123
+ @pytest.mark.parametrize("n_queries", [8, 128])
124
+ @pytest.mark.parametrize("transform", ["primal", "jvp", "vjp"])
125
+ def test_no_aux_grid_output_crossover(benchmark, n_steps, n_queries, transform):
126
+ """Cross-version baseline for the intentionally changed DAE grid path."""
127
+ solve = make_no_aux_solve(n_steps, n_queries)
128
+ rate = jnp.asarray(0.2)
129
+ assert bool(solve(rate).ok)
130
+
131
+ def scalar_output(value):
132
+ sol = solve(value)
133
+ return jnp.sum(sol.ys + sol.zs)
134
+
135
+ if transform == "primal":
136
+ run = jax.jit(solve)
137
+ elif transform == "jvp":
138
+ run = jax.jit(
139
+ lambda value: jax.jvp(scalar_output, (value,), (jnp.ones_like(value),))
140
+ )
141
+ else:
142
+ run = jax.jit(jax.grad(scalar_output))
143
+ compiled = run.lower(rate).compile()
144
+ jax.block_until_ready(compiled(rate))
145
+ benchmark.pedantic(
146
+ lambda: jax.block_until_ready(compiled(rate)), rounds=10, iterations=5
147
+ )
148
+
149
+
150
+ def test_batched_adaptive_dense_output(benchmark):
151
+ solve = make_adaptive_solve(32)
152
+ rates = jnp.linspace(0.1, 1.0, 16)
153
+ assert bool(jnp.all(jax.vmap(solve)(rates).ok))
154
+ run = jax.jit(jax.vmap(lambda value: solve(value).aux["flow"]))
155
+ compiled = run.lower(rates).compile()
156
+ jax.block_until_ready(compiled(rates))
157
+ benchmark.pedantic(
158
+ lambda: jax.block_until_ready(compiled(rates)), rounds=10, iterations=5
159
+ )
@@ -0,0 +1,138 @@
1
+ """Opt-in steady-state array/pytree and AD benchmarks.
2
+
3
+ Run with ``uv run --group benchmark pytest benchmarks --benchmark-only``.
4
+ """
5
+
6
+ import jax
7
+ import jax.numpy as jnp
8
+ import pytest
9
+
10
+ from tinydiffeq import (
11
+ RK4,
12
+ EulerMaruyama,
13
+ IController,
14
+ Tsit5,
15
+ solve_ode,
16
+ solve_sde,
17
+ solve_semi_explicit_dae,
18
+ solve_semi_explicit_sdae,
19
+ )
20
+
21
+
22
+ def make_state(kind):
23
+ if kind == "scalar":
24
+ return jnp.asarray(1.0)
25
+ if kind == "vector16":
26
+ return jnp.linspace(0.5, 1.5, 16)
27
+ return {"left": jnp.linspace(0.5, 1.0, 8), "right": jnp.linspace(1.0, 1.5, 8)}
28
+
29
+
30
+ def tree_map(fn, tree):
31
+ return jax.tree.map(fn, tree)
32
+
33
+
34
+ def tree_sum(tree):
35
+ return sum(jnp.sum(leaf) for leaf in jax.tree.leaves(tree))
36
+
37
+
38
+ def make_solve(method, initial):
39
+ if method == "rk4":
40
+ return lambda x: (
41
+ solve_ode(
42
+ lambda state: tree_map(lambda leaf: -0.2 * leaf, state),
43
+ RK4(),
44
+ 0.0,
45
+ 1.0,
46
+ x,
47
+ dt_0=1 / 64,
48
+ max_steps=64,
49
+ ).xs
50
+ )
51
+ if method == "tsit5":
52
+ return lambda x: (
53
+ solve_ode(
54
+ lambda state: tree_map(lambda leaf: -0.2 * leaf, state),
55
+ Tsit5(),
56
+ 0.0,
57
+ 1.0,
58
+ x,
59
+ dt_0=0.05,
60
+ controller=IController(),
61
+ max_steps=64,
62
+ ).xs
63
+ )
64
+ if method == "em":
65
+ return lambda x: (
66
+ solve_sde(
67
+ lambda state: tree_map(lambda leaf: -0.2 * leaf, state),
68
+ lambda state: tree_map(lambda leaf: 0.1 * jnp.ones_like(leaf), state),
69
+ EulerMaruyama(),
70
+ 0.0,
71
+ 1.0,
72
+ x,
73
+ key=jax.random.key(7),
74
+ n_steps=64,
75
+ ).xs
76
+ )
77
+
78
+ z_0 = tree_map(lambda leaf: 0.9 * leaf, initial)
79
+
80
+ def constraint(y, z):
81
+ y_flat = jnp.concatenate([jnp.ravel(leaf) for leaf in jax.tree.leaves(y)])
82
+ z_flat = jnp.concatenate([jnp.ravel(leaf) for leaf in jax.tree.leaves(z)])
83
+ return z_flat - y_flat
84
+
85
+ if method == "sdae":
86
+ return lambda y: (
87
+ solve_semi_explicit_sdae(
88
+ lambda state, z: tree_map(lambda leaf: -0.2 * leaf, z),
89
+ lambda state, z: tree_map(
90
+ lambda leaf: 0.1 * jnp.ones_like(leaf), state
91
+ ),
92
+ constraint,
93
+ EulerMaruyama(),
94
+ 0.0,
95
+ 1.0,
96
+ y,
97
+ z_0,
98
+ key=jax.random.key(7),
99
+ n_steps=64,
100
+ ).ys
101
+ )
102
+
103
+ return lambda y: (
104
+ solve_semi_explicit_dae(
105
+ lambda state, z: tree_map(lambda leaf: -0.2 * leaf, z),
106
+ constraint,
107
+ Tsit5(),
108
+ 0.0,
109
+ 1.0,
110
+ y,
111
+ z_0,
112
+ dt_0=0.1,
113
+ controller=IController(),
114
+ max_steps=32,
115
+ ).ys
116
+ )
117
+
118
+
119
+ @pytest.mark.parametrize("method", ["rk4", "tsit5", "em", "dae", "sdae"])
120
+ @pytest.mark.parametrize("state_kind", ["scalar", "vector16", "tree16"])
121
+ @pytest.mark.parametrize("transform", ["primal", "jvp", "vjp"])
122
+ def test_steady_state(benchmark, method, state_kind, transform):
123
+ initial = make_state(state_kind)
124
+ solve = make_solve(method, initial)
125
+ tangent = tree_map(jnp.ones_like, initial)
126
+
127
+ if transform == "primal":
128
+ run = jax.jit(solve)
129
+ elif transform == "jvp":
130
+ run = jax.jit(lambda x: jax.jvp(solve, (x,), (tangent,)))
131
+ else:
132
+ run = jax.jit(lambda x: jax.grad(lambda value: tree_sum(solve(value)))(x))
133
+
134
+ compiled = run.lower(initial).compile()
135
+ jax.block_until_ready(compiled(initial))
136
+ benchmark.pedantic(
137
+ lambda: jax.block_until_ready(compiled(initial)), rounds=20, iterations=10
138
+ )
@@ -4,6 +4,10 @@
4
4
 
5
5
  ::: tinydiffeq.solve_ode
6
6
 
7
+ ::: tinydiffeq.solve_semi_explicit_dae
8
+
9
+ ::: tinydiffeq.solve_semi_explicit_sdae
10
+
7
11
  ::: tinydiffeq.solve_sde
8
12
 
9
13
  ## Solvers
@@ -28,6 +32,10 @@
28
32
 
29
33
  ::: tinydiffeq.SaveAt
30
34
 
35
+ ::: tinydiffeq.LMRootSolver
36
+
37
+ ::: tinydiffeq.DAESolution
38
+
31
39
  ::: tinydiffeq.Solution
32
40
 
33
41
  ## Utilities