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.
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/PKG-INFO +55 -10
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/README.md +51 -7
- tinydiffeq-1.0.0/benchmarks/README.md +38 -0
- tinydiffeq-1.0.0/benchmarks/__init__.py +1 -0
- tinydiffeq-1.0.0/benchmarks/compile_times.py +64 -0
- tinydiffeq-1.0.0/benchmarks/test_aux_dense_output.py +159 -0
- tinydiffeq-1.0.0/benchmarks/test_pytree_overhead.py +138 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/api.md +8 -0
- tinydiffeq-1.0.0/docs/dae.md +184 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/index.md +20 -12
- tinydiffeq-1.0.0/docs/llms.txt +13 -0
- tinydiffeq-1.0.0/docs/sdae.md +99 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/sde.md +10 -4
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/static_shapes.md +17 -2
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/mkdocs.yml +3 -1
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/pyproject.toml +4 -3
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/__init__.py +15 -5
- tinydiffeq-1.0.0/src/tinydiffeq/_tree.py +122 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/controllers.py +12 -12
- tinydiffeq-1.0.0/src/tinydiffeq/dae.py +872 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/interpolation.py +22 -12
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/ode.py +38 -18
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/save_at.py +1 -1
- tinydiffeq-1.0.0/src/tinydiffeq/sdae.py +248 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/sde.py +50 -13
- tinydiffeq-1.0.0/src/tinydiffeq/solution.py +52 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/solvers.py +52 -22
- tinydiffeq-1.0.0/tests/test_aux.py +319 -0
- tinydiffeq-1.0.0/tests/test_dae.py +524 -0
- tinydiffeq-1.0.0/tests/test_float64_subprocess.py +488 -0
- tinydiffeq-1.0.0/tests/test_gpu.py +211 -0
- tinydiffeq-1.0.0/tests/test_pytree_states.py +269 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_save_at.py +1 -1
- tinydiffeq-1.0.0/tests/test_sdae.py +286 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/uv.lock +18 -2
- tinydiffeq-0.2.0/docs/llms.txt +0 -11
- tinydiffeq-0.2.0/src/tinydiffeq/solution.py +0 -28
- tinydiffeq-0.2.0/tests/test_float64_subprocess.py +0 -244
- tinydiffeq-0.2.0/tests/test_gpu.py +0 -75
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.githooks/pre-commit +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.github/workflows/ci.yml +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.github/workflows/docs.yml +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.github/workflows/publish.yml +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.gitignore +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/.python-version +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/LICENSE +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/docs/adaptive_ad.md +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/src/tinydiffeq/quadrature.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/conftest.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/kernels_reference.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_ad.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_adaptive.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_quadrature.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_recompile.py +0 -0
- {tinydiffeq-0.2.0 → tinydiffeq-1.0.0}/tests/test_sde.py +0 -0
- {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.
|
|
4
|
-
Summary: Tiny differentiable ODE/SDE solvers for JAX
|
|
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 ::
|
|
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
|
[](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
|
|
30
31
|
[](https://github.com/astral-sh/ruff)
|
|
31
32
|
|
|
32
|
-
Tiny differentiable ODE/SDE solvers for JAX: fixed-step Euler/RK4,
|
|
33
|
-
Tsit5 with integral or proportional-integral step-size control,
|
|
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
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
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
|
[](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
|
|
8
8
|
[](https://github.com/astral-sh/ruff)
|
|
9
9
|
|
|
10
|
-
Tiny differentiable ODE/SDE solvers for JAX: fixed-step Euler/RK4,
|
|
11
|
-
Tsit5 with integral or proportional-integral step-size control,
|
|
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
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
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
|