From 41da6348d5d5ee0887cfc030dbfab82def4925dd Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Fri, 31 Jul 2026 09:30:56 -0700 Subject: [PATCH] feat: release tinydiffeq 2.4.0 --- .github/workflows/ci.yml | 17 +- README.md | 69 +++- docs/adaptive_ad.md | 158 ++++++--- docs/aux.md | 29 +- docs/dae.md | 121 +++++-- docs/exponential.md | 17 +- docs/index.md | 56 ++- docs/llms.txt | 9 +- docs/sdae.md | 15 +- docs/static_shapes.md | 56 +-- pyproject.toml | 2 +- src/tinydiffeq/__init__.py | 21 +- src/tinydiffeq/_loop.py | 112 ++++++ src/tinydiffeq/dae.py | 545 +++++++++++++++++++++++------- src/tinydiffeq/exponential.py | 67 ++-- src/tinydiffeq/markov.py | 54 ++- src/tinydiffeq/ode.py | 296 +++++++++++++--- src/tinydiffeq/save_at.py | 7 + src/tinydiffeq/sdae.py | 45 ++- src/tinydiffeq/sde.py | 27 +- src/tinydiffeq/solution.py | 16 +- tests/test_adaptive.py | 218 ++++++++++++ tests/test_dae.py | 353 ++++++++++++++++++- tests/test_exponential.py | 130 +++++++ tests/test_float64_subprocess.py | 318 +++++++++++++++++ tests/test_markov.py | 103 +++++- tests/test_markov_distribution.py | 42 +++ tests/test_rodas5p.py | 75 ++++ tests/test_save_at.py | 337 +++++++++++++++++- tests/test_sdae.py | 5 + tests/test_sde.py | 1 + tests/test_solution.py | 59 ++++ tests/test_solvers_fixed.py | 63 ++++ uv.lock | 2 +- 34 files changed, 3079 insertions(+), 366 deletions(-) create mode 100644 src/tinydiffeq/_loop.py create mode 100644 tests/test_solution.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e827515..01199f9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,12 +18,27 @@ jobs: - run: uv run ruff format --check . test: + name: Test (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v7 - uses: astral-sh/setup-uv@v7 with: enable-cache: true - python-version: "3.13" + python-version: ${{ matrix.python-version }} - run: uv sync --locked - run: uv run pytest -q + + package: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + - run: uv build + - run: uvx twine check dist/* diff --git a/README.md b/README.md index 07bc62f..f8e5296 100644 --- a/README.md +++ b/README.md @@ -18,14 +18,13 @@ pytrees. The same dense and matrix-free backends are available directly through `solve_linear_ode` for any fixed homogeneous linear array or pytree operator; `jvp_linear_ode` and `vjp_linear_ode` apply the exact initial-state tangent and adjoint exponential actions without differentiating Arnoldi orthogonalization. -Bounded `lax.scan` loops provide exactly `max_steps` attempt slots for fixed and -adaptive stepping, so shapes are static, nothing recompiles as tolerances or -curvature change, and every solve is differentiable in **both** forward and -reverse mode — including reverse-over-forward, the pattern a -Levenberg–Marquardt optimizer with geodesic acceleration needs when it -differentiates through a rollout. Adaptive attempts are grouped into static -chunks, allowing one `lax.cond` to -skip solver and controller work for an entire padded chunk. +Fixed stepping and the default adaptive path use bounded `lax.scan` loops with +exactly `max_steps` attempt slots. Shapes stay static as tolerances or curvature +change, and these solves support forward mode, reverse mode, and +reverse-over-forward. Adaptive ODE and DAE solves may instead select +`adaptive_loop="forward"`: a dynamic `lax.while_loop` that executes only actual +attempts and supports JVP and nested forward AD, but not reverse mode. A vmapped +forward loop runs until its slowest lane finishes. This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX adaptation of Steinebach's method and follows SciML's @@ -35,8 +34,12 @@ implementation. Use [diffrax](https://docs.kidger.site/diffrax/) or matrices, fully implicit or higher-index DAEs, events, continuous solution objects, sparse/Krylov linear solvers for ODE/DAE stages, or specialized adjoints. Initial DAE consistency and explicit DAE stages use `nlls-gram`; -`LMRootSolver` accepts a `MAX_STEPS` iterate by default and offers strict -`CONVERGED`-only status via `max_steps_is_success=False`. +the same nlls solve supplies the square root's implicit derivative, whose +default is a direct nonsymmetric `LU()` solve. `LMRootSolver` requires +residual-only stopping (`gtol=xtol=0`) and accepts only `CONVERGED` roots whose +Euclidean residual norm is below the root `atol`. Its +`max_steps_is_success` field remains for source compatibility but does not make +`MAX_STEPS` a valid DAE root. The linear exponential-action API follows SciML [`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/). @@ -46,6 +49,31 @@ SciML's [`ExponentialIntegrators.jl`](https://docs.sciml.ai/ExponentialIntegrators/stable/) is the reference for the broader nonlinear exponential-integrator family. +## 2.4.0 migration note + +- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit + fixed-step ODEs. Every query must align with a knot; adaptive methods, + Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode. +- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts, + including rejections. DAE results additionally expose `num_root_solves` and + `num_root_steps`; `num_accepted` retains its existing meaning. +- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an + actual-work loop. It supports primal, JVP, and nested forward AD but not + reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable + default. Under `vmap`, the forward loop runs to the slowest lane. +- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for + locally unique algebraic branches; `predictor="previous"` remains the + default. + +DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root +solve and implicit derivative; square implicit AD defaults to direct `LU()`. +Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so +`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source +compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid +root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they +relied on budget exhaustion, increase the root budget or adjust the residual +tolerance instead. + ## Install ```bash @@ -107,11 +135,16 @@ returns one time/state, `SaveAt(ts=...)` returns the requested grid, and a contiguous prefix of `max_steps + 1` rows. The remaining rows repeat the last accepted state by default; `sol.accepted` distinguishes data from padding. Rejected attempts never appear in the returned trajectory. +`sol.num_steps` reports the number of attempts actually made, while +`sol.num_accepted` excludes rejections. `SaveAt(ts=...)` also accepts a Python sequence. These are observation times: the adaptive controller still chooses its own internal mesh. Explicit methods use cubic Hermite interpolation; Rodas5P uses its published stiff-aware -fourth-order continuous extension. +fourth-order continuous extension. For an explicit fixed-step ODE, +`SaveAt(ts=..., exact=True)` instead requires every requested time to be an +internal knot and gathers the stored state directly. Exact mode does not apply +to adaptive ODEs, Rodas5P, DAEs, SDEs, or SDAEs. ## Semi-explicit DAEs @@ -154,11 +187,17 @@ initialization: it advances the corresponding block mass-matrix system using one reused LU factorization per attempt. Differential fields may return a floating saved-aux pytree stored at accepted nodes and interpolated on requested deterministic grids. Algebraic equations may separately return internal context -passed to the dynamics. JVP, VJP, and reverse-over-forward propagate through both -implicit initialization and the time integrator. See the +passed to the dynamics. On the default bounded path, JVP, VJP, and +reverse-over-forward propagate through both implicit initialization and the +time integrator. See the [DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) for root controls, `SaveAt`, and scope limits. +DAE solutions expose `num_steps`, `num_root_solves`, and `num_root_steps` as +logical per-trajectory work counters. Explicit methods default to reusing the +previous algebraic root as the next stage guess; `LMRootSolver(predictor="secant")` +is an opt-in continuation predictor for locally unique root branches. + 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/). @@ -182,7 +221,9 @@ jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])( The step-size controller is wrapped in `stop_gradient` (accept/reject is non-differentiable either way, and the error-ratio power blows up at exactly -zero error); the states differentiate fully through the solver stages. See the +zero error); states differentiate through the solver stages on the realized, +frozen mesh. In particular, adaptive `SaveAt(steps=True)` does not include mesh +motion in its time or state derivatives. See the [docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key semantics, and the package API. diff --git a/docs/adaptive_ad.md b/docs/adaptive_ad.md index 50da505..591228e 100644 --- a/docs/adaptive_ad.md +++ b/docs/adaptive_ad.md @@ -1,9 +1,11 @@ # Adaptive Stepping and AD -tinydiffeq's central claim is that one adaptive solve is safely -differentiable in forward mode, reverse mode, and reverse-over-forward. That -holds because of three deliberate choices about **what is not -differentiated**. +The default `adaptive_loop="bounded"` supports forward mode, reverse mode, and +reverse-over-forward through one adaptive solve. The alternative +`adaptive_loop="forward"` executes a data-dependent number of attempts and +supports primal evaluation, JVP, and nested forward mode, but not ordinary +reverse mode. Both paths use the same deliberate frozen-controller derivative +convention described below. ## Default tolerances follow the state precision @@ -29,16 +31,16 @@ growth without introducing infinities. ## The controller is stop-gradiented The adaptive step-size controllers (`IController` and `PIController`) compute -their scaled error norms and next-step factors inside `stop_gradient`: +their scaled error norms, decisions, and next-step factors inside +`stop_gradient`. Accept/reject is a discrete branch, and differentiating +`E**(-1/order)` is singular at the exact-zero error of a flat-start policy. -> accept/reject is a non-differentiable branch either way, the gradient of -> `E**(-1/order)` blows up at the exact-zero error of a flat-start policy, -> and the `d(dt)/dtheta` term only slides sample points along the visited -> trajectory — irrelevant to a residual that must vanish at every state. - -The states themselves remain fully differentiable through the solver stages: the -gradient you get is the derivative of the numerical flow map *for the step -pattern actually taken*. This is the same convention diffrax uses. +States remain differentiable through solver stages, but the derivative holds +the realized step sizes and accept/reject pattern fixed. It is the derivative +of the discrete flow on a **frozen mesh**, not the total derivative of a mesh +that moves with parameters. In particular, for parameter-only differentiation, +the returned adaptive knot times in `SaveAt(steps=True)` have exactly zero +tangent. The `E**(-1/5)` blow-up is not hypothetical: a policy initialized flat gives an exactly-zero error estimate on the first step, and without the @@ -56,9 +58,14 @@ safety * E_n**(-(p_coeff + i_coeff) / order) and `E_prev` changes only after acceptance. The whole recurrence is controller-internal and stop-gradiented. Setting `p_coeff=0, i_coeff=1` -reproduces `IController` bit for bit; the default `p_coeff=0.4, i_coeff=0.3` -is less sensitive to oscillatory error estimates and typically rejects fewer -attempts on harder problems. +reproduces `IController` bit for bit. + +`IController` remains the default. On the smooth ODE and DAE screening +problems, the default PI coefficients did not reduce rejections, and +matched-work accuracy was comparable or favored I, except that PI sometimes +improved common-grid ODE interpolation. PI remains useful as an opt-in for +genuinely rejection-prone or step-size-oscillatory problems; those cases were +not covered by this screen. ## The horizon clip is the growth guard @@ -76,19 +83,56 @@ a few accepted steps; residuals sampled from three or four giant steps make the optimizer's linear model useless. Clipping first means the proposal can never exceed `factor_max × remaining horizon`. -One refinement over a bare `min(dt, remaining)`: when the remaining horizon -is within `max_steps × eps` of the desired step, the step is stretched to -land on `t_1` exactly. Summing `n` rounded steps of `(t_1 - t_0)/n` can leave -`t` one accumulated ulp short of `t_1`, and without the stretch that sliver -would cost an extra iteration a `max_steps = n` budget doesn't have. +Adaptive steps use `min(dt, remaining)` and are never enlarged to consume a +floating-point sliver. Fixed-step times are instead formed arithmetically as +`t_0 + i * dt_0`, avoiding drift from repeated addition, with a small local +endpoint snap capped below a quarter step. No time tolerance scales with +`max_steps`: changing a nonbinding attempt budget must not stretch a step or +change the numerical method. + +## Frozen mesh versus a moving-mesh derivative + +Let a parameter-dependent controller choose an internal knot +\(t_k(\theta)\), and let \(q(x,t,\theta)\) be a saved quantity. Along the exact +trajectory, its total derivative contains + +$$ +\frac{d q_k}{d\theta} +=q_x\left(x_\theta\rvert_t+f\,t_{k,\theta}\right) + +q_t t_{k,\theta}+q_\theta. +$$ + +tinydiffeq's frozen-mesh rule returns the terms holding \(t_k\) fixed. The +omitted mesh-motion term is + +$$ +\left(q_t+q_x f\right)t_{k,\theta} +=\frac{d q}{dt}\,t_{k,\theta}. +$$ + +For a collocation residual \(r(x_k,t_k,\theta)\), this omission is the total +trajectory derivative \((d r/dt)t_{k,\theta}\). If the learned policy approaches +a smooth root for which the residual vanishes along the trajectory, then +\(d r/dt\) vanishes with it locally; under that regularity, the omitted term is +\(O(\lVert r\rVert)\), so the frozen-mesh residual Jacobian is asymptotically +exact near the root. Pointwise cancellation at a few knots alone is not enough +for that conclusion. + +Far from a root the omitted term can be material, so a Gauss--Newton or LM +predicted reduction and trust ratio may be misleading. A targeted +finite-difference audit found that adaptive `SaveAt(steps=True)` state +derivatives need not approach the moving-mesh derivative under tolerance +refinement. Endpoint output at fixed `t_1` and fixed requested-grid output do +not have this particular moving-output-time ambiguity and did converge in the +same audit. ## Interpolation knots are non-differentiable `SaveAt(ts=...)` brackets each query with `searchsorted` — integer indices, -no gradient. This is consistent with the stop-gradiented controller: the -term excluded is again "the knots slide along the trajectory as parameters -change". Values differentiate fully through the bracketing states and -derivatives (`xL`, `xR`, `fL`, `fR`). +no gradient. Values differentiate through the bracketing states and +derivatives (`xL`, `xR`, `fL`, `fR`) while internal knot locations remain +frozen. The requested times themselves are fixed outputs, unlike +`SaveAt(steps=True)`'s adaptive internal times. Zero-width brackets (duplicate rows from rejections and the frozen tail) use the **double-where trick**: the divisor is replaced by 1 *before* dividing, @@ -105,13 +149,30 @@ on the output is not enough — reverse mode differentiates both branches, and ## What this buys you -- `jax.grad`, `jax.jvp`, and `jax.grad(jax.jvp(...))` (the +- With the default bounded loop, `jax.grad`, `jax.jvp`, and + `jax.grad(jax.jvp(...))` (the Levenberg–Marquardt geodesic-acceleration pattern) all work through adaptive solves and interpolated output, verified against closed forms in `tests/test_ad.py`. +- With `adaptive_loop="forward"`, primal evaluation, `jax.jvp`, and nested + forward mode work; `jax.vjp`, `jax.grad`, and reverse-over-forward do not. - `jax.vmap` over `x_0` or `p` gives genuinely per-lane adaptivity: each lane - accepts/rejects independently through the masked scan, and batched results - equal the individual solves exactly. + accepts/rejects independently. Execution continues until the slowest lane + finishes. + +The dynamic loop is most useful when actual attempts are much fewer than +`max_steps`. When an integration uses most of its budget, primal performance +can be similar or slower than the bounded scan; compilation and forward-mode +paths may still improve. Choose from representative end-to-end measurements, +not from the loop form alone. + +On float32 GPUs, different loop lowerings can also fuse neural-network matrix +products differently. With accelerator-default reduced-precision matrix +multiplication, that roundoff was large enough in one width-32 policy test to +move the adaptive mesh even though both loop contracts were correct. Setting +`jax_default_matmul_precision="highest"` restored matching paths. tinydiffeq +does not change this application-global JAX setting; set it in reproducibility +entry points when adaptive decisions depend on neural-network outputs. ## Reusing a linearization @@ -131,9 +192,14 @@ cotangent_batch = jax.jit( The setup computes and stores residuals for that one primal point. Reuse the pushforward or pullback only while every primal input is unchanged; otherwise -linearize again. This is particularly important for Rodas5P and implicit DAE -roots: their custom linear-solve rules retain primal LU/Cholesky factors, so -multiple directions reuse those factors rather than refactoring independently. +linearize again. Rodas5P's custom linear-solve rule explicitly retains and +reuses its primal pivoted LU factors. Semi-explicit DAE roots have a different +boundary: nlls-gram supplies both the primal LM root and its implicit derivative, +selecting direct `LU()` by default for a square residual system. The primal LM +Cholesky factors are not reused because they represent damped normal equations, +not the accepted root's nonsymmetric Jacobian. `jax.linearize` still avoids +repeating the primal trajectory, but the API makes no promise that the DAE +root's direct factorization is cached across pushforward directions. On the 256-state fixed Tsit5 benchmark, cached pushforwards were about 8–18% faster than a fused `vmap(jvp)` on CPU after excluding setup. On the RTX 3090, @@ -148,26 +214,36 @@ lanes. For one direction at a new primal point, ordinary `jax.jvp` or **discontinuous** in the inputs: its length changes when an accept flips to a reject. Use `sol.accepted` when padding must not contribute. Because the default tail repeats the endpoint, `xs[-1]` remains the reached final state. +- Adaptive `SaveAt(steps=True)` derivatives hold the internal times fixed and + omit mesh motion. Use a fixed requested grid when the residual definition + requires fixed sample times, or treat the frozen-mesh Jacobian explicitly as + a near-root approximation. - Finite-difference checks of adaptive solves are noisy for the same reason; compare AD against closed forms or use fixed-step solvers for FD tests. ## Custom-rule audit -tinydiffeq uses hand-coded derivative boundaries only where they remove -iteration or factorization work without changing the mathematical derivative: +tinydiffeq uses hand-coded derivative boundaries where they enforce the intended +mathematical derivative or avoid differentiating iteration/factorization work: + +There is no custom reverse rule that replays an adaptive solve on a recorded +mesh. The bounded path uses ordinary traced AD through its scan, subject to the +controller `stop_gradient`; the dynamic forward path deliberately exposes +JAX's no-reverse boundary. A replay/custom-VJP design would be a distinct future +API and would not restore the omitted moving-mesh term by itself. - Rodas5P's factored linear solve has a custom JVP \(\delta x=A^{-1}(\delta b-\delta A\,x)\). Every stage tangent and the transposed VJP reuse the attempt's pivoted LU factors; pivot selection is not differentiated. -- Semi-explicit DAE roots use the implicit-function rule - \(\delta z=-g_z^{-1}(g_y\delta y+g_t\delta t+g_p\delta p)\). The exact - constraint Jacobian generally cannot reuse the primal LM factor because the - latter is damped and may be based on an earlier iterate. A cached - `jax.linearize`/`jax.vjp` does reuse the implicit factor across directions. -- nlls-gram's public solve boundary already supplies implicit Cholesky/CG - derivative rules, including aux derivatives. tinydiffeq does not - differentiate its optimizer iterations. +- Semi-explicit DAE roots delegate both the primal solve and implicit AD to + nlls-gram. Its implicit-function rule is + \(\delta z=-g_z^{-1}(g_y\delta y+g_t\delta t+g_p\delta p)\); the default + square `LU()` implementation differentiates the defining constraint rather + than the optimizer iterations, and reverse mode transposes that rule. + `LMRootSolver` requires residual-only convergence (`gtol=xtol=0`) and never + accepts `MAX_STEPS` as a differentiable root. `ad_solver` remains an + nlls-owned option forwarded through `solver_options`. - Dense linear exponential actions use a Fréchet custom JVP for active matrix or time tangents and reuse the matrix exponential when only the initial state varies. diff --git a/docs/aux.md b/docs/aux.md index b3adad6..a0cfc1c 100644 --- a/docs/aux.md +++ b/docs/aux.md @@ -115,20 +115,29 @@ not valid for rough paths. `SaveAt(steps=True)` stores it at the initial and accepted nodes and applies the same prefix/padding mask as the state. -For `SaveAt(ts=grid)`, ODE and root-restored DAE aux uses normalized cubic -Hermite interpolation. Endpoint aux slopes are JVPs along the full solution -velocity, so they include direct parameter dependence and indirect dependence -through the state and any implicit algebraic root. Rodas5P obtains endpoint -state velocities from its published stiff-aware continuous extension, then -uses the same JVP construction for aux. No requested-time algebraic root or -aux recomputation is performed. See [Rodas5P](rodas5p.md) and -[Semi-Explicit DAEs](dae.md) for the dense-output details and links to SciML's -implementation. +By default, `SaveAt(ts=grid)` uses normalized cubic Hermite interpolation for +ODE and root-restored DAE aux. Endpoint aux slopes are JVPs along the full +solution velocity, so they include direct parameter dependence and indirect +dependence through the state and any implicit algebraic root. Rodas5P obtains +endpoint state velocities from its published stiff-aware continuous extension, +then uses the same JVP construction for aux. No requested-time algebraic root +or aux recomputation is performed on this dense-output path. See +[Rodas5P](rodas5p.md) and [Semi-Explicit DAEs](dae.md) for the dense-output +details and links to SciML's implementation. + +For an explicit fixed-step ODE with `SaveAt(ts=grid, exact=True)`, every +requested time must instead coincide with a realized solver knot. The solver +selects those states without Hermite interpolation and evaluates aux directly +at the requested knots. Saved-aux validity is therefore checked only at those +requested knots on this exact path. Ordinary JAX transformations compose through every saved or interpolated aux leaf. This includes `jax.jvp`, reverse-mode VJP/`jax.grad`, `vmap`, and reverse-over-forward. Adaptive accept/reject decisions and mesh selection -retain the package's frozen-controller derivative convention. +retain the package's frozen-controller derivative convention. Consequently, +adaptive `SaveAt(steps=True)` aux derivatives hold internal times fixed and +omit mesh motion; fixed requested-grid aux does not have a moving output-time +axis. See [Adaptive Stepping and AD](adaptive_ad.md#frozen-mesh-versus-a-moving-mesh-derivative). ## Failure behavior diff --git a/docs/dae.md b/docs/dae.md index db599b9..8aff8ec 100644 --- a/docs/dae.md +++ b/docs/dae.md @@ -16,10 +16,13 @@ linearly implicit Rodas5P with fixed or adaptive control. The algebraic solve uses [`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/)'s -general Levenberg–Marquardt solver. The default `linear_solver="auto"` selects -the dense normal-Cholesky form for the square algebraic system, while -`ad_solver="auto"` resolves every square system to a general nonsymmetric -direct solve of the algebraic Jacobian. No implicit ridge is added by default. +general Levenberg–Marquardt solver for both the primal root and its implicit +derivative. The default `Cholesky()` primal solver selects the dense normal form +for the square algebraic system, while nlls's default `ad_solver` selects the +direct nonsymmetric `LU()` rule for that square system. The implicit rule +differentiates the defining equation rather than the LM iterations; no implicit +ridge is added. Tinydiffeq applies the DAE validity policy around that solve but +does not implement a separate square-system IFT. Rodas5P is a JAX adaptation of Steinebach's method following SciML's [`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock) @@ -94,10 +97,11 @@ from tinydiffeq import LMRootSolver root_solver = LMRootSolver( max_steps=8, - max_steps_is_success=True, + max_steps_is_success=False, # compatibility field; MAX_STEPS stays invalid atol=None, # 1e-6 float32, 1e-10 float64 - gtol=0.0, # disabled - xtol=0.0, # disabled + gtol=0.0, # required: residual stopping only + xtol=0.0, # required: residual stopping only + predictor="previous", solver_options=(), # nlls-gram constructor defaults ) ``` @@ -107,16 +111,25 @@ rejections. `root_solver.max_steps` separately bounds one algebraic root. For Rodas5P it affects only initial consistency; the method's later stages reuse one dense LU factorization per attempted time step. -Root tolerances are independent of the outer controller tolerances. The -`atol`, `gtol`, and `xtol` fields pass to the nlls solve; zero disables the -corresponding test. By default, exhausting `root_solver.max_steps` accepts the -last root iterate and uses its implicit derivative even though nlls retains the -diagnostic `MAX_STEPS` status. Set `max_steps_is_success=False` to require -`CONVERGED`; the failed root then has zero implicit tangent. Everything -algorithmic runs at nlls-gram's own defaults — a dense `Cholesky()` forward -solve and `ad_solver=None` for the implicit derivative. For the rare root that -needs to depart from them, `solver_options` forwards keyword arguments -verbatim to the `LevenbergMarquardt` constructor: +Root tolerances are independent of the outer controller tolerances. Explicit +`atol` must be positive; `None` selects the dtype default, and `atol=0` is +invalid. `gtol` and `xtol` must both equal zero; `LMRootSolver` rejects nonzero +values so `CONVERGED` can only come from the residual test. Every accepted +algebraic root must report `CONVERGED` and have Euclidean residual norm +`sqrt(sum(residual**2))` strictly below the root `atol`. +`max_steps_is_success` remains in the configuration for source compatibility, +but Tinydiffeq always asks nlls to treat `MAX_STEPS` as a failed implicit solve; +setting the field to `True` does not broaden DAE root acceptance. + +The primal nonlinear solve uses nlls-gram's dense `Cholesky()` normal-equation +default. For a successful square root, nlls's implicit rule forms `dg/dz`, +forms the right-hand side with respect to `(y, t, p)`, and applies its direct +`LU()` square solve. Transposing that nlls rule supplies the VJP. The primal LM +factorization is not reused because it represents a damped normal equation, +not the accepted root's generally nonsymmetric Jacobian. A failed nlls status +has zero implicit tangent by the nlls solve contract. For the rare root solve +that needs to depart from the defaults, `solver_options` forwards constructor +arguments for either the primal `linear_solver` or the implicit `ad_solver`: ```python from nlls_gram import QR @@ -130,10 +143,24 @@ to a sorted tuple so equal configurations remain hashable and share one compiled solver. Algebraic roots fix `cache_jacobian=False` and `geodesic_acceleration=False` — each DAE stage changes the root problem and the intended path is the ordinary dense LM step — and `solver_options` rejects -both rather than silently honoring an override. - -Every nonlinear root passes `(y, t, p)` through nlls-gram's differentiated parameter -pytree. Thus it differentiates the defining constraint, +those two options rather than silently honoring an override. `ad_solver` +remains nlls-owned and may be supplied explicitly; the square default is +`LU()`. + +`predictor="previous"` is the default: each explicit RK stage starts from the +most recent successful algebraic root. `predictor="secant"` extrapolates from +the accepted-step root through the most recent successful stage at a later +time. Duplicate RK4 stage times, non-forward targets, and failed stages fall +back to the previous root. The predictor is stop-gradiented, so a successful +root still uses the same implicit derivative. Its time-derived extrapolation +scale is cast separately to each algebraic leaf's dtype, preserving the `z` +dtype when the differential/time and algebraic dtypes differ. Secant prediction +assumes the continued branch is locally unique; with multiple roots or a tight +finite iteration budget, changing the guess can change the selected branch, +value, or status. + +Every nonlinear root passes `(y, t, p)` to nlls-gram. Its implicit rule then +differentiates the defining constraint, $$ \dot z = -g_z^{-1} @@ -143,8 +170,23 @@ $$ rather than differentiating the LM iterations. The warm-start guess has zero derivative by design. Rodas5P differentiates through its exact JAX Jacobian, time derivative, LU factorization, and linear stage solves. `args` is fixed -data; put every differentiated model quantity in `p`. JVP, VJP, `vmap`, and -reverse-over-forward compose through the complete DAE solve. +data; put every differentiated model quantity in `p`. On the default bounded +integration path, JVP, VJP, `vmap`, and reverse-over-forward compose through +the complete DAE solve. + +`sol.num_steps` counts logical time-step attempts, including adaptive +rejections. `sol.num_root_solves` counts active nonlinear root calls, including +the initial consistency solve and failed calls, and `sol.num_root_steps` sums +their LM update counts. Rodas5P therefore reports one root solve regardless of +its time-step count; later linear stages are not nonlinear roots. These +counters have exact-zero tangents. Under `vmap`, a masked lane may still execute +physically while remaining absent from its logical counters. + +The default `adaptive_loop="bounded"` uses a reverse-mode-capable bounded scan. +`adaptive_loop="forward"` uses a dynamic actual-work loop for adaptive Tsit5 +and Rodas5P. It supports primal evaluation, JVP, and nested forward mode, but +not ordinary reverse mode. A vmapped forward loop runs until the slowest lane +finishes. The differential field may return `(dy, saved_aux)`. Saved aux is a nonempty pytree of nonempty real floating arrays; different leaves may use different @@ -167,15 +209,16 @@ the initial and accepted nodes, so an invalid value freezes the previous valid prefix. Endpoint mode evaluates saved aux only after integration; an invalid final value retains the endpoint state, returns zero aux, and sets `ok=False`. -With strict root status enabled, an adaptive stage-root failure rejects the -time-step attempt and asks the controller for a smaller step; a fixed-step -failure terminates. Rodas5P linear failures follow the same controller policy. +An adaptive stage-root failure rejects the time-step attempt and asks the +controller for a smaller step; a fixed-step failure terminates. Rodas5P linear +failures follow the same controller policy. In either case `sol.ok` is false if the endpoint is not reached with valid -algebraic states. tinydiffeq delegates each root and its implicit JVP/VJP to -nlls-gram; statuses rejected by the configured policy receive a zero implicit -tangent, and aux at a failed initial root is a zero pytree of the declared -shape. Callers that want to retain successful-lane JVPs/VJPs after another -lane has already become inactive should pass +algebraic states. nlls supplies the primal LM iterate, diagnostics, and the +implicit JVP/VJP. Tinydiffeq applies the DAE's residual-and-status acceptance +check and returns the differentiation-inert warm-start guess when that check +fails; aux at a failed initial root is a zero pytree of the declared shape. +Callers that want to retain successful-lane JVPs/VJPs after another lane has +already become inactive should pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)`, choosing a point where the residual, context, and saved-aux maps are finite and differentiable. tinydiffeq substitutes this point into an already-inactive root call before @@ -201,7 +244,11 @@ All `SaveAt` modes are supported: stiff-aware continuous extension for `(y, z)`. Aux uses cubic Hermite in both cases. It performs no query-time nonlinear solves. -The result is `DAESolution(ts, ys, zs, ok, num_accepted, accepted, aux)`. +`SaveAt(ts=..., exact=True)` is not a DAE mode; exact knot gathering is limited +to explicit fixed-step ODEs. + +The result is a `DAESolution` with `ts`, `ys`, `zs`, `ok`, `num_accepted`, +`accepted`, `aux`, `num_steps`, `num_root_solves`, and `num_root_steps` fields. For pytree states, saved rows are a leading axis on every state and aux leaf; the one `accepted` mask applies to the complete output. @@ -252,10 +299,12 @@ query. Rodas5P accepted knots are not root-restored: their constraint defect, and that of dense output, is controlled by integration accuracy rather than `LMRootSolver.atol`. -Knot selection and adaptive step sizes remain non-differentiable, consistent -with the frozen-controller convention. Values, implicit slopes, and aux are -fully differentiated. If `sol.ok` is false, neither outputs nor their -derivatives should be treated as a valid solution. +Knot selection and adaptive step sizes remain differentiation-inert under the +frozen-controller convention. Values, implicit slopes, and aux differentiate +on that realized mesh, but adaptive `SaveAt(steps=True)` omits mesh motion. +See [Frozen mesh versus a moving-mesh derivative](adaptive_ad.md#frozen-mesh-versus-a-moving-mesh-derivative). +If `sol.ok` is false, neither outputs nor their derivatives should be treated +as a valid solution. ## Deliberate limits diff --git a/docs/exponential.md b/docs/exponential.md index 684e944..1ed443b 100644 --- a/docs/exponential.md +++ b/docs/exponential.md @@ -41,7 +41,8 @@ is useful as a correctness baseline for small systems. ## Matrix-free pytree operator `KrylovExponential` applies the exponential without constructing `A` or -`exp(A)`. The callable sees and returns the original pytree: +`exp(A)` when `krylov_dim` is smaller than the flattened state dimension. The +callable sees and returns the original pytree: ```python from tinydiffeq import KrylovExponential @@ -76,6 +77,13 @@ Krylov subspace has closed. `solution.ok` combines finite-output checks with a leading-term Arnoldi error estimate. Increase `krylov_dim` or `num_substeps` if it is false. +When `krylov_dim` reaches the full state dimension, the Krylov space is the +whole linear space and the result is exactly the dense exponential action. In +that small-system case Tinydiffeq materializes the operator and uses the dense +action, avoiding the coordinate singularity of an Arnoldi basis at an early +happy breakdown. The genuinely matrix-free regime remains +`krylov_dim < ravel(x_0).size`. + ## Adaptive matrix-free propagation Use `AdaptiveKrylovExponential` when the appropriate number of internal @@ -175,6 +183,13 @@ Arnoldi computation. Because Arnoldi normalizes its starting vector, that path assumes the initial state has nonzero norm. At exactly zero, use the hand-coded initial-state functions above or `DenseExponential`. +At an exact early happy breakdown in a *truncated* Krylov space, the normalized +basis itself changes nonsmoothly for directions outside the closed invariant +subspace. Use `jvp_linear_ode` or `vjp_linear_ode` for the mathematical +fixed-operator initial-state derivative in that case. Final-vector breakdowns +have finite ordinary AD, and full-dimensional Krylov delegates to the exact +dense derivative as described above. + For the adaptive method, ordinary AD differentiates the numerical computation on the realized accepted/rejected path; the discrete controller decisions are locally constant. The hand-coded initial-state JVP/VJP instead apply an diff --git a/docs/index.md b/docs/index.md index 62cf27c..2b74c08 100644 --- a/docs/index.md +++ b/docs/index.md @@ -4,10 +4,11 @@ finite-state Markov simulators for JAX: fixed-step Euler and RK4, adaptive Tsit5 with integral or proportional-integral step-size control, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs, and fixed-step Euler–Maruyama for Itô SDEs and SDAEs. -Time integration uses bounded `lax.scan` loops 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 -geodesic acceleration needs when it differentiates through a rollout. +Fixed stepping and the default adaptive path use bounded `lax.scan` loops with +static shapes. These solves support forward mode, reverse mode, and +reverse-over-forward. Adaptive ODE and DAE solves also offer +`adaptive_loop="forward"`, a dynamic actual-work loop for primal, JVP, and +nested forward AD; JAX cannot reverse-transpose this path. Finite-state DTMC/CTMC simulation is primal-only and offers chronological scan and associative parallel-prefix methods. Fixed-chain probability forecasts use binary matrix powers for DTMC endpoints, dense exponentials for small CTMCs, or @@ -28,9 +29,36 @@ implementation and Steinebach's published method. **Use SciML or - dense output / continuous interpolation objects - checkpointed or backsolve adjoints for long horizons -tinydiffeq ships only the O(max_steps)-memory bounded-scan approach, because -that is the one that composes cleanly with `jax.jvp`, `jax.vjp`, `jax.vmap`, -and reverse-over-forward without custom adjoint machinery. +Use the default bounded loop when reverse mode is required. Use the forward +loop when the actual adaptive attempt count is much smaller than `max_steps` +and the caller needs only primal or forward-mode execution. Both retain static +public output shapes; under `vmap`, a dynamic loop still runs until the slowest +lane finishes. + +## 2.4.0 migration note + +- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit + fixed-step ODEs. Every query must align with a knot; adaptive methods, + Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode. +- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts, + including rejections. DAE results additionally expose `num_root_solves` and + `num_root_steps`; `num_accepted` retains its existing meaning. +- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an + actual-work loop. It supports primal, JVP, and nested forward AD but not + reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable + default. Under `vmap`, the forward loop runs to the slowest lane. +- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for + locally unique algebraic branches; `predictor="previous"` remains the + default. + +DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root +solve and implicit derivative; square implicit AD defaults to direct `LU()`. +Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so +`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source +compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid +root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they +relied on budget exhaustion, increase the root budget or adjust the residual +tolerance instead. ## Install @@ -136,6 +164,18 @@ jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])( `SaveAt(steps=True)`, which returns `max_steps + 1` padded rows including the initial state. Accepted steps form a contiguous prefix; rejected attempts are omitted and the tail repeats the last accepted state. + `sol.num_steps` reports actual attempts and `sol.num_accepted` reports + successful advances. +- **Fixed-step times do not depend on the attempt budget.** They are formed + arithmetically from the accepted-step index, with only a small local endpoint + snap. Increasing a nonbinding `max_steps` therefore does not change the + numerical method. +- **Exact fixed-step output is explicit-ODE-only.** + `SaveAt(ts=grid, exact=True)` gathers internal knots without interpolation; + every query must align with a realized knot. +- **Adaptive loop choice is an AD choice.** `adaptive_loop="bounded"` is the + reverse-mode-capable default. `adaptive_loop="forward"` executes actual + attempts but supports only primal, JVP, and nested forward mode. - **Forward time only**: `t_1 > t_0`. - **Never poisons.** `sol.ok` reports whether `t_1` was reached and every requested output was valid; callers that want diverging residuals can map @@ -148,7 +188,7 @@ jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])( registered as pytrees: numeric fields (tolerances, grids, `dt_0`, `x_0`) are data leaves, so changing them never recompiles. -Read next: [Static Shapes](static_shapes.md) for the bounded-scan design and +Read next: [Static Shapes](static_shapes.md) for the loop and `SaveAt`, [Adaptive Stepping and AD](adaptive_ad.md) for what is and is not differentiated, [Auxiliary Outputs](aux.md), [Rodas5P](rodas5p.md) for the SciML-derived linearly implicit method, [DAEs](dae.md), [SDEs](sde.md), [SDAEs](sdae.md), and the [API diff --git a/docs/llms.txt b/docs/llms.txt index 1bda9b6..b2cb761 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1,14 +1,15 @@ # tinydiffeq -> Tiny ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Euler, RK4, adaptive Tsit5, and the SciML-derived linearly implicit Rodas5P use bounded lax.scan loops with static shapes and forward/reverse AD, including reverse-over-forward. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, an embedded estimator, and its published stiff-aware fourth-order dense output. See https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock and https://doi.org/10.1007/s10543-023-00967-x. RK4/Tsit5 DAEs use nlls-gram LM roots at every stage; the square primal defaults to normal Cholesky, implicit AD defaults to a direct nonsymmetric solve, and MAX_STEPS is accepted unless LMRootSolver.max_steps_is_success is false. Rodas5P uses LM only for initial consistency. Finite-state DTMC/CTMC sampling is primal-only and offers sequential and associative parallel-prefix methods. Deterministic distribution forecasts are differentiable in their initial mass: DTMC endpoints use binary matrix powers, dense CTMCs use matrix exponentials, and matrix-free CTMCs use fixed or residual-controlled adaptive Arnoldi time slicing over array or pytree probabilities. The adaptive method keeps the Krylov dimension static and bounds accepted plus rejected attempts with max_steps. The Krylov design follows SciML ExponentialUtilities.expv, stores basis vectors as contiguous rows, defaults to stable two-pass reorthogonalization, and offers an explicitly selected one-pass performance mode. ODE fields and stochastic drifts may return (value, saved_aux), a real-floating pytree saved with the solution and differentiated through JVP/VJP. DAE/SDAE algebraic functions may separately return (residual, algebraic_aux); that internal array pytree is passed to the differential fields but is not itself stored or interpolated. Requested-grid deterministic saved aux uses cubic Hermite interpolation and no query-time nonlinear solves. Fixed-step solve_semi_explicit_sdae applies Euler-Maruyama to the reduced Ito SDE. States may be arrays or arbitrary JAX pytrees; state leaves share one real floating dtype and saved-aux leaves may use different floating dtypes. Adaptive ODE/DAE defaults are rtol=1e-4, atol=1e-6 for float32 and rtol=1e-7, atol=1e-9 for float64; adaptive Krylov defaults are rtol=1e-5, atol=1e-7 and rtol=1e-10, atol=1e-12 respectively. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints. +> Tiny ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Fixed stepping and the default adaptive path use bounded lax.scan loops with static output shapes and forward/reverse AD, including reverse-over-forward. Adaptive Tsit5 and Rodas5P ODE/DAE solves can instead select adaptive_loop="forward", an actual-work lax.while_loop for primal, JVP, and nested forward AD; reverse mode is unsupported, and vmap runs to the slowest lane. max_steps is an attempt budget, while Solution.num_steps records actual attempts. Fixed-step times are arithmetic and budget-invariant. SaveAt(ts=..., exact=True) gathers aligned internal knots for explicit constant-step ODEs only; other requested-grid deterministic output uses Hermite or Rodas5P interpolation. Adaptive SaveAt(steps=True) AD freezes the internal mesh and omits mesh motion, an asymptotically exact residual Jacobian only near a smooth residual root. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, an embedded estimator, and its published stiff-aware fourth-order dense output. RK4/Tsit5 DAEs delegate both their stage roots and implicit derivatives to nlls-gram; the square primal defaults to normal Cholesky and the square implicit rule defaults to direct nonsymmetric LU. LMRootSolver requires residual-only stopping with gtol=xtol=0: every accepted algebraic root must report CONVERGED and have Euclidean residual norm below root atol. max_steps_is_success remains for source compatibility but never makes MAX_STEPS a valid DAE root. LMRootSolver defaults to the previous-root predictor and offers an opt-in secant predictor for locally unique branches. DAESolution reports logical time attempts, nonlinear root calls, and LM update counts; Rodas5P uses LM only for initial consistency. ODE fields and stochastic drifts may return (value, saved_aux), while DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary JAX pytrees. IController remains the adaptive default; PIController is opt-in. Finite-state DTMC/CTMC sampling is primal-only and offers sequential and associative parallel-prefix methods. Deterministic distribution forecasts are differentiable in their initial mass through matrix powers, dense exponentials, or matrix-free Arnoldi/Krylov actions. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints. ## 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/) +- [Static shapes — bounded and actual-work loops, budget-invariant fixed times, SaveAt modes, work counters](https://highdimensionaleconlab.github.io/tinydiffeq/static_shapes/) +- [Adaptive stepping and AD — frozen-mesh derivative contract, loop-mode AD boundary, controller and horizon logic](https://highdimensionaleconlab.github.io/tinydiffeq/adaptive_ad/) +- [Auxiliary outputs — saved versus algebraic aux, interpolation, AD, and failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/aux/) - [Rodas5P — Steinebach method, direct SciML implementation links, stiff ODE/DAE formulation, dense output, and AD](https://highdimensionaleconlab.github.io/tinydiffeq/rodas5p/) -- [Semi-explicit DAEs — index-1 contract, algebraic LM roots, implicit AD, failure behavior, and SaveAt](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) +- [Semi-explicit DAEs — index-1 contract, algebraic LM roots and predictors, counters, 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/) - [Semi-explicit SDAEs — reduced-SDE Euler-Maruyama, algebraic roots, aux, convergence assumptions, and pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) - [Finite-state Markov chains — sampling, deterministic PMF forecasts, matrix-free Krylov CTMC actions, pytrees, vmap, and AD scope](https://highdimensionaleconlab.github.io/tinydiffeq/markov_chains/) diff --git a/docs/sdae.md b/docs/sdae.md index 3f0330b..fd65636 100644 --- a/docs/sdae.md +++ b/docs/sdae.md @@ -86,9 +86,11 @@ random numbers. The key is not differentiable. Algebraic solves use the same `LMRootSolver` configuration and implicit-AD contract as deterministic DAEs; see [Nonlinear-solve and AD contract](dae.md#nonlinear-solve-and-ad-contract). -`MAX_STEPS` is accepted by default; use -`LMRootSolver(max_steps_is_success=False)` when every stochastic node must -report nlls `CONVERGED`. +nlls-gram supplies both the primal LM root and its implicit derivative, using +direct `LU()` by default for the square derivative. `gtol` and `xtol` must +remain zero, and every accepted iterate must report `CONVERGED` with Euclidean +residual norm strictly below the root `atol`. `max_steps_is_success` is retained +for source compatibility but does not make `MAX_STEPS` valid. The algebraic function may return `(residual, algebraic_aux)`; that internal context is passed to both drift and diffusion but is not stored. The drift may return `(drift_value, saved_aux)`, and only that saved aux becomes `sol.aux`. @@ -101,10 +103,13 @@ stochastic step. In steps mode, invalid saved aux terminates at the previous consistent node. In endpoint mode, invalid final saved aux retains the endpoint state, returns zero aux, and sets `ok=False`. +`sol.num_root_solves` counts logical active root calls, including the initial +consistency solve and failures, while `sol.num_root_steps` sums their LM update +counts. Both are path diagnostics with exact-zero tangents. + Only `SaveAt(t_1=True)` and `SaveAt(steps=True)` are supported. Stochastic paths are rough, so deterministic dense interpolation between nodes would be -mathematically wrong. A root status rejected by the configured policy freezes -the last consistent prefix, +mathematically wrong. A failed root freezes the last consistent prefix, sets `ok=False`, and pads the remaining static buffer. Failed roots have zero implicit tangents, and aux at a failed initial root is zero-filled, so masked lanes can preserve successful JVPs or VJPs under `vmap`. For that contract, diff --git a/docs/static_shapes.md b/docs/static_shapes.md index 7ce5f7b..b83f60f 100644 --- a/docs/static_shapes.md +++ b/docs/static_shapes.md @@ -1,11 +1,13 @@ # Static Shapes -JAX jits fixed-shape programs. An adaptive integrator is naturally -dynamic — the number of steps depends on the data — so something must give. -tinydiffeq's answer is a **bounded scan**: `solve_ode` always provides exactly -`max_steps` static attempt slots, whatever the controller does. Adaptive -attempts are grouped into small nested-scan chunks so a completed solve skips -whole padded chunks instead of visiting every unused attempt individually. +JAX jits fixed-shape programs. An adaptive integrator is naturally dynamic — +the number of steps depends on the data — so tinydiffeq offers two execution +contracts. The default `adaptive_loop="bounded"` provides exactly `max_steps` +static attempt slots in a bounded scan. Adaptive attempts are grouped into +small nested-scan chunks so a completed solve skips whole padded chunks. +`adaptive_loop="forward"` uses a dynamic `lax.while_loop` and executes only +actual attempts, while retaining the same static public output shapes and a +`max_steps` row buffer when the selected `SaveAt` mode requires one. Each iteration attempts one step: @@ -19,17 +21,22 @@ The raw internal scan buffer contains repeated rows for rejected and frozen iterations. That buffer is an implementation detail used by interpolation; step output compacts it into accepted rows plus tail padding. -The bounded loops always contain `max_steps` attempt slots, preserving static -shapes and reverse-mode AD. Chunk-level and attempt-level `lax.cond` branches -keep the expensive vector-field, solver-stage, and controller computations -out of the frozen tail. Under `vmap`, JAX may turn each lane's conditional -into selection; batched work can therefore continue until the slowest lane -finishes. +The bounded loop preserves reverse-mode AD. Chunk-level and attempt-level +`lax.cond` branches keep expensive field, stage, and controller computations +out of its frozen tail. The forward loop supports primal evaluation, JVP, and +nested forward mode, but JAX cannot transpose its data-dependent while loop. +Under `vmap`, both strategies advance lanes together until the slowest lane +finishes; the forward strategy stops at that lane's actual attempt count rather +than always reaching `max_steps`. Fixed-step integration uses a smaller specialized scan without adaptive controller or embedded-error work. `ConstantStepSize` accepts every attempt, so `dt_0 = (t_1 - t_0)/n` with `max_steps = n` reproduces a fixed grid -exactly. +exactly. Times are formed arithmetically as `t_0 + i * dt_0`, rather than by +repeatedly accumulating rounded steps. A small local tolerance snaps the +nominal last point to `t_1`; it is capped relative to `dt_0` and never scales +with `max_steps`. Consequently, increasing a nonbinding attempt budget cannot +stretch an earlier step or otherwise change the numerical method. If the budget runs out before `t_1`, `sol.ok` is `False` and the outputs hold the reached prefix. The package never poisons values; the caller decides: @@ -38,6 +45,10 @@ the reached prefix. The package never poisons values; the caller decides: xs = jnp.where(sol.ok, sol.xs, jnp.inf) # kernels-style rejection ``` +`sol.num_steps` counts attempts actually made, including rejections; +`sol.num_accepted` counts advances. These scalar diagnostics do not change the +fixed output shape. + ## SaveAt is the shape contract Exactly one of three modes: @@ -60,11 +71,16 @@ accepted; times must be nondecreasing and within `[t_0, t_1]`, while repeated times and omitted endpoints are allowed. Changing values without changing the grid length does not recompile. -These are observation times, not mandatory internal stops. The adaptive -controller chooses exactly the same mesh regardless of the requested grid, -then the solver evaluates every requested point through dense interpolation. -Forcing exact internal landing times is a distinct feature and is not part of -this API. +By default these are observation times, not mandatory internal stops. The +adaptive controller chooses the same mesh regardless of the requested grid, +then the solver evaluates each point through dense interpolation. + +For an explicit ODE with `ConstantStepSize`, +`SaveAt(ts=grid, exact=True)` instead requires every query to coincide with a +realized internal endpoint and gathers the stored states directly. It avoids +interpolation and the endpoint-slope work needed by Hermite output. A +misaligned or unreached query makes `sol.ok` false. Exact mode is unavailable +for adaptive ODEs, Rodas5P, DAEs, SDEs, and SDAEs. The interpolation runs directly over the raw padded rows: duplicate knots from rejections or the frozen tail form zero-width brackets, and the @@ -117,8 +133,8 @@ leaf independently. `sol.accepted` is one shared mask for all leaves. - Tolerances and PI coefficients (`IController(...)` / `PIController(...)`), `dt_0`, `t_0`, `t_1`, `x_0`, `args`, `p`, and `SaveAt.ts` are pytree **data leaves**. Only genuine - structure — the solver type, `SaveAt` mode, `fill`, `max_steps`, the - functions themselves — is static. + structure — the solver type, `SaveAt` mode, `fill`, `exact`, `max_steps`, + `adaptive_loop`, and the functions themselves — is static. An omitted tolerance or `dt_min` is represented by `None`, so switching a jitted call between automatic and explicit values changes the controller diff --git a/pyproject.toml b/pyproject.toml index 91e922f..bdbf795 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "tinydiffeq" -version = "2.3.0" +version = "2.4.0" description = "Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX with static shapes and composable AD" readme = "README.md" license = "MIT" diff --git a/src/tinydiffeq/__init__.py b/src/tinydiffeq/__init__.py index b442a06..f16f738 100644 --- a/src/tinydiffeq/__init__.py +++ b/src/tinydiffeq/__init__.py @@ -1,17 +1,24 @@ """Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX. solve_ode integrates dx/dt = f(x, t, args, p) with fixed-step (Euler, RK4), -adaptive explicit Tsit5, or linearly implicit Rodas5P methods -inside bounded lax.scan loops with exactly max_steps attempt slots, so shapes are -static and solves are differentiable in BOTH forward and reverse mode -(including reverse-over-forward) with O(max_steps) memory. SaveAt picks the -output: the endpoint, method-specific dense interpolation onto a fixed grid, or -accepted internal steps with fixed-shape padding. solve_sde is fixed-step +adaptive explicit Tsit5, or linearly implicit Rodas5P methods. Fixed stepping +and the default adaptive path use bounded lax.scan loops with exactly max_steps +attempt slots and support forward/reverse AD, including reverse-over-forward. +Adaptive ODE/DAE solves may instead use an actual-work lax.while_loop for +primal, JVP, and nested forward AD; reverse mode is unsupported on that path. +SaveAt picks the endpoint, dense interpolation onto a fixed grid, accepted +internal steps with padding, or exact knot gathering for explicit fixed-step +ODEs. max_steps bounds attempted steps; Solution.num_steps reports actual work, +and arithmetic fixed-step times are independent of a nonbinding budget. +solve_sde is fixed-step Euler-Maruyama with presampled diagonal noise. solve_semi_explicit_dae handles index-1 systems with either root-restored explicit methods or the stiff Rodas5P mass-matrix formulation, plus differentiable saved differential-field aux and internal algebraic -context. +context. DAESolution also reports nonlinear-root calls and LM update counts; +LMRootSolver delegates primal roots and square implicit derivatives to +nlls-gram, requires residual-only stopping, and offers previous-root and secant +stage predictors. solve_semi_explicit_sdae applies fixed-step Euler-Maruyama to the reduced index-1 stochastic system. solve_linear_ode applies dense, fixed-Krylov, or adaptive matrix-free exponential actions to fixed homogeneous linear systems. diff --git a/src/tinydiffeq/_loop.py b/src/tinydiffeq/_loop.py new file mode 100644 index 0000000..3afc307 --- /dev/null +++ b/src/tinydiffeq/_loop.py @@ -0,0 +1,112 @@ +"""Shared actual-work loop utilities for adaptive differential equations.""" + +import jax +import jax.numpy as jnp + + +def select_step( + t, + t_0, + t_1, + dt, + dt_0, + step_count, + *, + constant, + time_tolerance, +): + """Choose a forward step without making it depend on the attempt budget. + + Constant-step times are formed arithmetically from the accepted-step index, + avoiding drift from repeated addition. A nominal endpoint within a small, + local floating-point tolerance is corrected to ``t_1``. Adaptive steps are + never enlarged: they use ``min(dt, t_1 - t)`` and snap the bookkeeping time + only when that exact remaining interval was integrated. + + In particular, ``time_tolerance`` must not scale with ``max_steps``. Such a + tolerance can turn a nonbinding attempt budget into a change in the + numerical method by stretching an earlier step to the horizon. + """ + dtype = jnp.result_type(t, t_0, t_1, dt, dt_0) + positive_floor = jnp.asarray(jnp.finfo(dtype).tiny, dtype) + if constant: + snap_tolerance = jnp.minimum( + jnp.asarray(time_tolerance, dtype), + 0.25 * jnp.abs(jnp.asarray(dt_0, dtype)), + ) + next_index = jnp.asarray(step_count + 1, dtype) + nominal_next = t_0 + next_index * dt_0 + reaches_horizon = (nominal_next >= t_1) | ( + jnp.abs(nominal_next - t_1) <= snap_tolerance + ) + t_next = jnp.where(reaches_horizon, t_1, nominal_next) + h = jnp.maximum(t_next - t, positive_floor) + return h, t_next, reaches_horizon + + remaining = t_1 - t + reaches_horizon = remaining <= dt + h = jnp.where(reaches_horizon, jnp.maximum(remaining, positive_floor), dt) + t_next = jnp.where(reaches_horizon, t_1, t + h) + return h, t_next, reaches_horizon + + +def forward_adaptive_while( + carry, + *, + attempt_step, + skip_step, + terminated, + max_steps, +): + """Run adaptive attempts until termination, retaining static output shapes. + + ``attempt_step`` and ``skip_step`` share the scan-style contract + ``(carry, output)``. Outputs receive a fixed leading ``max_steps`` buffer, + but only actual attempts execute. Under ``vmap``, JAX's while batching rule + advances lanes in lockstep until every lane is terminated and freezes lanes + that completed earlier. + + Dynamic while loops support primal evaluation and forward-mode AD, but JAX + cannot transpose them. Callers must expose that boundary in their API. + """ + _, initial_output = skip_step(carry) + if initial_output is None: + initial_rows = None + else: + initial_rows = jax.tree.map( + lambda value: jnp.broadcast_to(value, (max_steps,) + value.shape), + initial_output, + ) + + def condition(loop_state): + attempt, loop_carry, _ = loop_state + return (attempt < max_steps) & ~terminated(loop_carry) + + def run_attempt(loop_state): + attempt, loop_carry, loop_rows = loop_state + next_carry, output = attempt_step(loop_carry) + if loop_rows is not None: + loop_rows = jax.tree.map( + lambda rows, value: rows.at[attempt].set(value), + loop_rows, + output, + ) + return attempt + 1, next_carry, loop_rows + + attempt, final_carry, rows = jax.lax.while_loop( + condition, + run_attempt, + (jnp.asarray(0, jnp.int32), carry, initial_rows), + ) + if rows is not None: + _, final_output = skip_step(final_carry) + attempted = jnp.arange(max_steps) < attempt + + def fill_unattempted(values, final_value): + mask = attempted.reshape( + attempted.shape + (1,) * (values.ndim - attempted.ndim) + ) + return jnp.where(mask, values, final_value) + + rows = jax.tree.map(fill_unattempted, rows, final_output) + return final_carry, rows diff --git a/src/tinydiffeq/dae.py b/src/tinydiffeq/dae.py index d371b2a..d26a952 100644 --- a/src/tinydiffeq/dae.py +++ b/src/tinydiffeq/dae.py @@ -7,7 +7,7 @@ import jax.numpy as jnp import numpy as np from jax.flatten_util import ravel_pytree -from nlls_gram import LevenbergMarquardt, LMStatus +from nlls_gram import LevenbergMarquardt, LMState, LMStatus from tinydiffeq._aux import ( make_safe_evaluator, @@ -17,6 +17,7 @@ split_algebraic_output, split_field_output, ) +from tinydiffeq._loop import forward_adaptive_while, select_step from tinydiffeq._rodas5p import rodas5p_step, rodas_dense_endpoint_derivatives from tinydiffeq._tree import ( add_scaled, @@ -88,19 +89,21 @@ class LMRootSolver: """Configuration for algebraic solves in a semi-explicit DAE. The implementation is :class:`nlls_gram.LevenbergMarquardt` at its - defaults: a dense ``Cholesky()`` forward solve, which takes the normal - form for a square DAE constraint, and ``ad_solver=None`` for the implicit - derivative, which matches the forward family. + defaults: dense ``Cholesky()`` for primal LM updates and the direct + nonsymmetric ``LU()`` implicit derivative that nlls selects automatically + for a square residual system. The fields here are the ones this package owns; all of them reach the nlls ``solve`` rather than its constructor. ``max_steps`` counts nonlinear iterations for one algebraic root and is independent of the integration's - time-step ``max_steps``. ``max_steps_is_success=True`` accepts the final - iterate when that budget is exhausted; set it to ``False`` to require an - nlls ``CONVERGED`` status. ``atol=None`` selects ``1e-6`` in float32 and - ``1e-10`` in float64. ``gtol`` and ``xtol`` default to zero (disabled); - root tolerances are deliberately independent of the outer integration - tolerances. + time-step ``max_steps``. Every accepted algebraic root must have Euclidean + residual norm strictly below the root ``atol``. Root solves therefore use + residual stopping only: ``gtol`` and ``xtol`` must remain zero, and a + ``MAX_STEPS`` iterate is never treated as a differentiable root. + ``max_steps_is_success`` remains in the configuration for source + compatibility but does not broaden root acceptance. ``atol=None`` selects + ``1e-6`` in float32 and ``1e-10`` in float64. Root tolerances are + deliberately independent of the outer integration tolerances. ``solver_options`` is the escape hatch for the rare root that needs a non-default algorithm: a mapping (or pairs) forwarded verbatim to the @@ -114,15 +117,28 @@ class LMRootSolver: step. ``cache_jacobian`` and ``geodesic_acceleration`` are fixed to ``False`` and rejected here: each DAE stage changes the root problem, and the intended path is the ordinary dense LM step. Algebraic residuals do - not expose nlls aux. + not expose nlls aux. ``ad_solver`` remains an nlls-owned option; its default + is the direct square ``LU()`` rule. + + ``predictor="previous"`` warm-starts every explicit RK stage from the most + recent successful root. ``predictor="secant"`` instead extrapolates from + the accepted-step root through the most recent successful stage at a later + time. The secant is used only for a strictly later target; duplicate RK4 + stage times and failed stages fall back to the previous root. Predictor + values are differentiation-inert, so successful roots retain the same + implicit derivative. With multiple algebraic roots, however, a different + warm start can select a different branch; the secant mode is intended only + when the continued root is locally unique. A finite root-iteration budget + can also make predictor choice affect values or success status. """ max_steps: int = field(default=8, metadata=dict(static=True)) - max_steps_is_success: bool = field(default=True, metadata=dict(static=True)) + max_steps_is_success: bool = field(default=False, metadata=dict(static=True)) atol: float | None = field(default=None, metadata=dict(static=True)) gtol: float = field(default=0.0, metadata=dict(static=True)) xtol: float = field(default=0.0, metadata=dict(static=True)) solver_options: Any = field(default=(), metadata=dict(static=True)) + predictor: str = field(default="previous", metadata=dict(static=True)) def __post_init__(self): if not isinstance(self.max_steps, int) or isinstance(self.max_steps, bool): @@ -131,12 +147,18 @@ def __post_init__(self): raise ValueError("LMRootSolver.max_steps must be a positive int") if not isinstance(self.max_steps_is_success, bool): raise TypeError("LMRootSolver.max_steps_is_success must be a bool") - if self.atol is not None and self.atol < 0: - raise ValueError("LMRootSolver.atol must be nonnegative or None") - if self.gtol < 0: - raise ValueError("LMRootSolver.gtol must be nonnegative") - if self.xtol < 0: - raise ValueError("LMRootSolver.xtol must be nonnegative") + if self.atol is not None and self.atol <= 0: + raise ValueError("LMRootSolver.atol must be positive or None") + if self.gtol != 0: + raise ValueError("LMRootSolver.gtol must be zero for DAE root solves") + if self.xtol != 0: + raise ValueError("LMRootSolver.xtol must be zero for DAE root solves") + if not isinstance(self.predictor, str): + raise TypeError("LMRootSolver.predictor must be a string") + if self.predictor not in ("previous", "secant"): + raise ValueError( + 'LMRootSolver.predictor must be either "previous" or "secant"' + ) try: options = tuple(sorted(dict(self.solver_options).items())) except (TypeError, ValueError) as error: @@ -255,7 +277,6 @@ def _get_algebraic_solver(g, config, has_algebraic_aux=False): return _build_algebraic_solver(g, config, has_algebraic_aux) -@jax.custom_jvp def _inactive_safe_inputs(active, inputs, reference): return jax.tree.map( lambda value, reference_value: jnp.where(active, value, reference_value), @@ -264,13 +285,6 @@ def _inactive_safe_inputs(active, inputs, reference): ) -@_inactive_safe_inputs.defjvp -def _inactive_safe_inputs_jvp(primals, tangents): - active, inputs, reference = primals - _, inputs_dot, _ = tangents - return _inactive_safe_inputs(active, inputs, reference), inputs_dot - - def _prepare_failure_ad_reference(reference, y, z, t, p): """Validate a derivative-only reference point for inactive vmap lanes.""" @@ -331,7 +345,7 @@ def _make_implicit_root_solver( args, has_algebraic_aux, ): - """Delegate root solving and status-safe implicit AD to nlls-gram.""" + """Delegate the primal root and square implicit derivative to nlls-gram.""" root_atol = root_solver.atol if root_atol is None: @@ -358,22 +372,30 @@ def solve_root(y, t, z_guess, p, active, failure_ad_reference): (y, z_guess, t, p), failure_ad_reference, ) + root_p = (y_initial, t_initial, p_initial) + lm_state = LMState( + jnp.asarray(algebraic_solver.init_damping, z_dtype), + hyper=algebraic_solver.hyperparams(z_dtype), + ) + root_tolerance = jnp.asarray(root_atol, z_dtype) + zero_tolerance = jnp.zeros((), z_dtype) result = algebraic_solver.solve( z_initial, args, - p=(y_initial, t_initial, p_initial), + p=root_p, + lm_state=lm_state, max_steps=root_solver.max_steps, - max_steps_is_success=root_solver.max_steps_is_success, - atol=root_atol, - gtol=root_solver.gtol, - xtol=root_solver.xtol, + max_steps_is_success=False, + atol=root_tolerance, + gtol=zero_tolerance, + xtol=zero_tolerance, ) - ok = result.status == jnp.asarray(LMStatus.CONVERGED, result.status.dtype) - if root_solver.max_steps_is_success: - ok = ok | ( - result.status == jnp.asarray(LMStatus.MAX_STEPS, result.status.dtype) - ) - ok = active & ok + converged = result.status == jnp.asarray( + LMStatus.CONVERGED, result.status.dtype + ) + residual_norm = jnp.sqrt(result.info.loss) + residual_ok = residual_norm < jnp.asarray(root_atol, residual_norm.dtype) + ok = active & residual_ok & converged value, dtype = asarray_state(result.x, "algebraic root") assert_same_structure(z_reference, value, "algebraic root") if dtype != z_dtype: @@ -381,8 +403,15 @@ def solve_root(y, t, z_guess, p, active, failure_ad_reference): # A failed root is not a solution. Returning the finite warm start # keeps the static failure prefix usable while `ok` carries validity. # The guess is never a differentiable algebraic state, including on - # failure; successful tangents come entirely from nlls implicit AD. - return where(ok, value, jax.lax.stop_gradient(z_guess)), ok + # failure; successful tangents come from nlls's direct square LU rule. + num_root_solves = jnp.asarray(active, jnp.int32) + num_root_steps = jnp.where(active, result.steps, jnp.asarray(0, jnp.int32)) + return ( + where(ok, value, jax.lax.stop_gradient(z_guess)), + ok, + num_root_solves, + num_root_steps, + ) return solve_root, residual, algebraic_auxiliary @@ -396,6 +425,8 @@ def _solve_rodas5p_dae( z_initial, aux_initial, initial_ok, + num_root_solves, + num_root_steps, p, failure_ad_reference, dt_0, @@ -405,11 +436,11 @@ def _solve_rodas5p_dae( time_dtype, z_dtype, has_aux, + adaptive_loop, ): """Integrate a semi-explicit index-1 DAE with native Rodas5P stages.""" - positive_time_floor = jnp.asarray(jnp.finfo(time_dtype).tiny, time_dtype) - t_eps = 4.0 * jnp.finfo(time_dtype).eps * jnp.maximum(1.0, jnp.abs(t_1)) - t_slack = max_steps * t_eps + time_scale = jnp.maximum(jnp.maximum(1.0, jnp.abs(t_0)), jnp.abs(t_1)) + t_eps = 4.0 * jnp.finfo(time_dtype).eps * time_scale y_flat, _ = ravel_pytree(y_0) z_flat, _ = ravel_pytree(z_initial) mass_diagonal = jnp.concatenate([jnp.ones_like(y_flat), jnp.zeros_like(z_flat)]) @@ -472,13 +503,18 @@ def attempt_step(carry): reached, failed, num_accepted, + num_steps, controller_state, ) = carry - remaining = t_1 - t - h = jnp.where( - remaining <= dt + t_slack, - jnp.maximum(remaining, positive_time_floor), + h, proposed_t, reaches_horizon = select_step( + t, + t_0, + t_1, dt, + dt_0, + num_steps, + constant=not controller.uses_error_estimate, + time_tolerance=t_eps, ) state = (y, z) field_active = ~reached & ~failed @@ -513,7 +549,7 @@ def accepted_auxiliary(): aux_candidate, aux_ok = evaluate_aux( y_1, z_1, - t + h, + proposed_t, p, provisional_advance, failure_ad_reference, @@ -526,7 +562,7 @@ def accepted_auxiliary(): y, z, t, left_dot, provisional_advance ) aux_candidate, aux_ok, aux_right_dot = auxiliary_value_and_derivative( - y_1, z_1, t + h, right_dot, provisional_advance + y_1, z_1, proposed_t, right_dot, provisional_advance ) return aux_candidate, aux_ok, aux_left_dot, aux_right_dot @@ -545,7 +581,7 @@ def accepted_auxiliary(): y_new = where(advance, y_1, y) z_new = where(advance, z_1, z) aux_new = where(advance, aux_candidate, aux) if track_aux else None - t_new = jnp.where(advance, t + h, t) + t_new = jnp.where(advance, proposed_t, t) dt_new = jnp.where(reached | failed, dt, dt_next) controller_state_new = jax.tree.map( lambda old, new: jnp.where(step_ok, new, old), @@ -557,8 +593,9 @@ def accepted_auxiliary(): else: failed_new = failed | ~step_ok failed_new = failed_new | (provisional_advance & ~aux_ok) - reached_new = reached | (t_new >= t_1 - t_eps) + reached_new = reached | (advance & reaches_horizon) num_new = num_accepted + advance.astype(jnp.int32) + num_steps_new = num_steps + (~reached & ~failed).astype(jnp.int32) carry_new = ( t_new, y_new, @@ -568,6 +605,7 @@ def accepted_auxiliary(): reached_new, failed_new, num_new, + num_steps_new, controller_state_new, ) if save_at.t_1: @@ -588,7 +626,7 @@ def accepted_auxiliary(): return carry_new, out def skip_step(carry): - t, y, z, aux, _, _, _, _, _ = carry + t, y, z, aux, _, _, _, _, _, _ = carry if save_at.t_1: out = None elif save_at.steps: @@ -618,22 +656,31 @@ def body(carry, _): jnp.asarray(False), ~initial_ok, jnp.asarray(0, jnp.int32), + jnp.asarray(0, jnp.int32), controller_state_initial, ) + if controller.uses_error_estimate and adaptive_loop == "forward": + final_carry, rows = forward_adaptive_while( + carry_0, + attempt_step=attempt_step, + skip_step=skip_step, + terminated=lambda carry: carry[5] | carry[6], + max_steps=max_steps, + ) + else: + final_carry, rows = jax.lax.scan(body, carry_0, None, length=max_steps) ( - ( - t_final, - y_final, - z_final, - aux_final, - _, - reached, - failed, - num_accepted, - _, - ), - rows, - ) = jax.lax.scan(body, carry_0, None, length=max_steps) + t_final, + y_final, + z_final, + aux_final, + _, + reached, + failed, + num_accepted, + num_steps, + _, + ) = final_carry integration_ok = reached & ~failed if save_at.t_1: @@ -655,7 +702,10 @@ def body(carry, _): zs=z_final, ok=integration_ok & aux_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=aux_final, + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) if save_at.steps: @@ -700,12 +750,15 @@ def body(carry, _): zs=fill_rows(compact_zs, accepted, last_z, save_at.fill), ok=integration_ok, num_accepted=num_accepted, + num_steps=num_steps, accepted=accepted, aux=( fill_rows(compact_aux, accepted, last_aux, save_at.fill) if track_aux else None ), + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) query_times = jnp.asarray(save_at.ts, time_dtype) @@ -732,7 +785,10 @@ def body(carry, _): zs=query_zs, ok=integration_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=query_aux, + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) @@ -755,6 +811,7 @@ def solve_semi_explicit_dae( has_algebraic_aux=None, failure_ad_reference=None, max_steps=4096, + adaptive_loop="bounded", ): """Integrate a semi-explicit index-1 DAE. @@ -772,6 +829,9 @@ def solve_semi_explicit_dae( for initial consistency, then advances the block mass-matrix system with reused linear solves; later ``z`` values satisfy the constraint to the integration accuracy rather than the root tolerance. + ``sol.num_root_solves`` counts logical active nonlinear root calls and + ``sol.num_root_steps`` sums their LM update steps. Rodas5P therefore reports + one root call regardless of its number of integration attempts. ``args`` is fixed data by convention. All differentiated model parameters belong in ``p``. Initial consistency and explicit-method roots @@ -796,9 +856,20 @@ def solve_semi_explicit_dae( time-step work. Saved aux is checked at the initial and accepted nodes in prefix/grid modes; endpoint mode checks it only after integration and retains the endpoint state with zero aux if that check fails. + + ``adaptive_loop="bounded"`` keeps the reverse-mode-capable static scan. + ``adaptive_loop="forward"`` executes only actual adaptive attempts and, + under ``vmap``, stops after the slowest lane. It supports JVP and nested + forward mode but not ordinary reverse mode, matching JAX's dynamic-while + differentiation boundary. """ if dt_0 is None: raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)") + if adaptive_loop not in ("bounded", "forward"): + raise ValueError( + 'adaptive_loop must be either "bounded" or "forward", got ' + f"{adaptive_loop!r}" + ) if save_at is None: save_at = SaveAt(t_1=True) if controller is None: @@ -814,6 +885,12 @@ def solve_semi_explicit_dae( f"{type(controller).__name__} needs an embedded error estimate, " f"which {type(solver).__name__} does not provide" ) + if adaptive_loop == "forward" and not controller.uses_error_estimate: + raise ValueError( + 'adaptive_loop="forward" requires an adaptive error controller' + ) + if save_at.exact: + raise ValueError("SaveAt exact=True is only supported by solve_ode") y_0, time_dtype = asarray_state(y_0, "y_0") z_0, z_dtype = asarray_state(z_0, "z_0") @@ -823,9 +900,8 @@ def solve_semi_explicit_dae( failure_ad_reference = _prepare_failure_ad_reference( failure_ad_reference, y_0, z_0, t_0, p ) - positive_time_floor = jnp.asarray(jnp.finfo(time_dtype).tiny, time_dtype) - t_eps = 4.0 * jnp.finfo(time_dtype).eps * jnp.maximum(1.0, jnp.abs(t_1)) - t_slack = max_steps * t_eps + time_scale = jnp.maximum(jnp.maximum(1.0, jnp.abs(t_0)), jnp.abs(t_1)) + t_eps = 4.0 * jnp.finfo(time_dtype).eps * time_scale raw_f = f g_field = _canonicalize_dae_field(g, "g") @@ -973,7 +1049,12 @@ def residual_y_t(y_value, t_value): aux_dot = where(active, aux_dot, zeros_like(aux_dot)) return z_dot, aux_dot - z_initial, initial_root_ok = solve_root(y_0, t_0, z_0, jnp.asarray(True)) + ( + z_initial, + initial_root_ok, + initial_root_solves, + initial_root_steps, + ) = solve_root(y_0, t_0, z_0, jnp.asarray(True)) if has_algebraic_aux: _, initial_context_ok = evaluate_context( y_0, z_initial, t_0, p, initial_root_ok @@ -1054,6 +1135,8 @@ def combined_values(y, z, t, active): z_initial, aux_initial, initial_ok, + initial_root_solves, + initial_root_steps, p, failure_ad_reference, dt_0, @@ -1063,6 +1146,7 @@ def combined_values(y, z, t, active): time_dtype, z_dtype, has_aux, + adaptive_loop, ) need_f = solver.fsal or (save_at.ts is not None) @@ -1081,7 +1165,9 @@ def combined_values(y, z, t, active): def stage(y_stage, t_stage, z_guess, active, need_derivative=True): def evaluate(): - z_stage, root_ok = solve_root(y_stage, t_stage, z_guess, active) + z_stage, root_ok, root_solves, root_steps = solve_root( + y_stage, t_stage, z_guess, active + ) if need_derivative: k_stage, field_ok = jax.lax.cond( root_ok, @@ -1096,28 +1182,174 @@ def evaluate(): ) else: field_ok = root_ok - return z_stage, k_stage, root_ok & field_ok + return z_stage, k_stage, root_ok & field_ok, root_solves, root_steps def skip(): - return z_guess, zeros_like(y_stage), jnp.asarray(False) + zero = jnp.asarray(0, jnp.int32) + return z_guess, zeros_like(y_stage), jnp.asarray(False), zero, zero return jax.lax.cond(active, evaluate, skip) + def predicted_stage( + y_stage, + t_stage, + z_previous, + t_base, + z_base, + t_latest, + z_latest, + has_later_stage, + active, + need_derivative=True, + ): + z_guess = z_previous + if root_solver.predictor == "secant": + distinct = has_later_stage & (t_latest > t_base) + strictly_later = distinct & (t_stage > t_latest) + denominator = jnp.where(distinct, t_latest - t_base, 1.0) + scale = (t_stage - t_base) / denominator + extrapolated = jax.tree.map( + lambda base, latest: ( + base + jnp.asarray(scale, dtype=base.dtype) * (latest - base) + ), + z_base, + z_latest, + ) + z_guess = where(strictly_later, extrapolated, z_previous) + z_guess = jax.tree.map(jax.lax.stop_gradient, z_guess) + + z_stage, k_stage, stage_ok, root_solves, root_steps = stage( + y_stage, t_stage, z_guess, active, need_derivative + ) + if root_solver.predictor == "secant": + update_latest = stage_ok & (t_stage > t_base) + t_latest = jnp.where(update_latest, t_stage, t_latest) + z_latest = where(update_latest, z_stage, z_latest) + has_later_stage = has_later_stage | update_latest + return ( + z_stage, + k_stage, + stage_ok, + root_solves, + root_steps, + t_latest, + z_latest, + has_later_stage, + ) + def rk4_step(t, y, z, h, f_cur): k_1 = differential(y, z, t)[0] if f_cur is None else f_cur - z_2, k_2, ok_2 = stage(add_scaled(y, (0.5 * h, k_1)), t + 0.5 * h, z, True) - z_3, k_3, ok_3 = stage(add_scaled(y, (0.5 * h, k_2)), t + 0.5 * h, z_2, ok_2) - z_4, k_4, ok_4 = stage(add_scaled(y, (h, k_3)), t + h, z_3, ok_2 & ok_3) - y_1 = add_scaled(y, (h / 6.0, weighted_sum((k_1, k_2, k_3, k_4), (1, 2, 2, 1)))) + t_latest = t + z_latest = z + has_later_stage = jnp.asarray(False) + ( + z_2, + k_2, + ok_2, + solves_2, + steps_2, + t_latest, + z_latest, + has_later_stage, + ) = predicted_stage( + add_scaled(y, (0.5 * h, k_1)), + t + 0.5 * h, + z, + t, + z, + t_latest, + z_latest, + has_later_stage, + True, + ) + ( + z_3, + k_3, + ok_3, + solves_3, + steps_3, + t_latest, + z_latest, + has_later_stage, + ) = predicted_stage( + add_scaled(y, (0.5 * h, k_2)), + t + 0.5 * h, + z_2, + t, + z, + t_latest, + z_latest, + has_later_stage, + ok_2, + ) + ( + z_4, + k_4, + ok_4, + solves_4, + steps_4, + t_latest, + z_latest, + has_later_stage, + ) = predicted_stage( + add_scaled(y, (h, k_3)), + t + h, + z_3, + t, + z, + t_latest, + z_latest, + has_later_stage, + ok_2 & ok_3, + ) + y_1 = add_scaled( + y, + (h / 6.0, weighted_sum((k_1, k_2, k_3, k_4), (1, 2, 2, 1))), + ) stage_ok = ok_2 & ok_3 & ok_4 - z_1, f_1, endpoint_ok = stage(y_1, t + h, z_4, stage_ok, need_derivative=need_f) - return y_1, z_1, f_1, None, stage_ok & endpoint_ok + ( + z_1, + f_1, + endpoint_ok, + endpoint_solves, + endpoint_steps, + _, + _, + _, + ) = predicted_stage( + y_1, + t + h, + z_4, + t, + z, + t_latest, + z_latest, + has_later_stage, + stage_ok, + need_derivative=need_f, + ) + root_solves = solves_2 + solves_3 + solves_4 + endpoint_solves + root_steps = steps_2 + steps_3 + steps_4 + endpoint_steps + return ( + y_1, + z_1, + f_1, + None, + stage_ok & endpoint_ok, + root_solves, + root_steps, + ) def tsit5_step(t, y, z, h, f_cur): k_1 = differential(y, z, t)[0] if f_cur is None else f_cur ks = [k_1] z_stage = z stages_ok = jnp.asarray(True) + root_solves = jnp.asarray(0, jnp.int32) + root_steps = jnp.asarray(0, jnp.int32) + t_latest = t + z_latest = z + has_later_stage = jnp.asarray(False) rows = ( ((A_21,), C_2), ((A_31, A_32), C_3), @@ -1127,20 +1359,60 @@ def tsit5_step(t, y, z, h, f_cur): ) for coefficients, stage_time in rows: y_stage = add_scaled(y, (h, weighted_sum(ks, coefficients))) - z_stage, k_stage, root_ok = stage( - y_stage, t + stage_time * h, z_stage, stages_ok + ( + z_stage, + k_stage, + root_ok, + stage_solves, + stage_steps, + t_latest, + z_latest, + has_later_stage, + ) = predicted_stage( + y_stage, + t + stage_time * h, + z_stage, + t, + z, + t_latest, + z_latest, + has_later_stage, + stages_ok, ) stages_ok = stages_ok & root_ok + root_solves = root_solves + stage_solves + root_steps = root_steps + stage_steps ks.append(k_stage) y_1 = add_scaled(y, (h, weighted_sum(ks, (B_1, B_2, B_3, B_4, B_5, B_6)))) - z_1, k_7, endpoint_ok = stage(y_1, t + h, z_stage, stages_ok) + ( + z_1, + k_7, + endpoint_ok, + endpoint_solves, + endpoint_steps, + _, + _, + _, + ) = predicted_stage( + y_1, + t + h, + z_stage, + t, + z, + t_latest, + z_latest, + has_later_stage, + stages_ok, + ) + root_solves = root_solves + endpoint_solves + root_steps = root_steps + endpoint_steps ks.append(k_7) root_ok = stages_ok & endpoint_ok err = jax.tree.map( lambda value: h * value, weighted_sum(ks, (E_1, E_2, E_3, E_4, E_5, E_6, E_7)), ) - return y_1, z_1, k_7, err, root_ok + return y_1, z_1, k_7, err, root_ok, root_solves, root_steps def attempt_step(carry): ( @@ -1155,22 +1427,41 @@ def attempt_step(carry): reached, failed, num_accepted, + num_steps, + num_root_solves, + num_root_steps, controller_state, ) = carry - remaining = t_1 - t - h = jnp.where( - remaining <= dt + t_slack, - jnp.maximum(remaining, positive_time_floor), + h, proposed_t, reaches_horizon = select_step( + t, + t_0, + t_1, dt, + dt_0, + num_steps, + constant=not controller.uses_error_estimate, + time_tolerance=t_eps, ) if isinstance(solver, RK4): - y_1, z_1, f_1, err, root_ok = rk4_step( - t, y, z, h, f_cur if need_f else None - ) + ( + y_1, + z_1, + f_1, + err, + root_ok, + attempt_root_solves, + attempt_root_steps, + ) = rk4_step(t, y, z, h, f_cur if need_f else None) else: - y_1, z_1, f_1, err, root_ok = tsit5_step( - t, y, z, h, f_cur if need_f else None - ) + ( + y_1, + z_1, + f_1, + err, + root_ok, + attempt_root_solves, + attempt_root_steps, + ) = tsit5_step(t, y, z, h, f_cur if need_f else None) if controller.uses_error_estimate: control_err = where(root_ok, err, full_like(err, jnp.inf)) @@ -1186,7 +1477,7 @@ def attempt_step(carry): def accepted_aux(): y_safe = where(provisional_advance, y_1, y) z_safe = where(provisional_advance, z_1, z) - t_safe = jnp.where(provisional_advance, t + h, t) + t_safe = jnp.where(provisional_advance, proposed_t, t) return evaluate_aux( y_safe, z_safe, @@ -1207,7 +1498,7 @@ def accepted_aux(): advance = provisional_advance & aux_ok y_new = where(advance, y_1, y) z_new = where(advance, z_1, z) - t_new = jnp.where(advance, t + h, t) + t_new = jnp.where(advance, proposed_t, t) f_new = where(advance, f_1, f_cur) if need_f else f_cur if track_aux: aux_new = where(advance, aux_candidate, aux) @@ -1218,7 +1509,7 @@ def accepted_aux(): def accepted_derivatives(): y_safe = where(advance, y_1, y) z_safe = where(advance, z_1, z) - t_safe = jnp.where(advance, t + h, t) + t_safe = jnp.where(advance, proposed_t, t) f_safe = where(advance, f_1, f_cur) return algebraic_time_derivatives( y_safe, z_safe, t_safe, f_safe, advance @@ -1242,8 +1533,11 @@ def accepted_derivatives(): else: failed_new = failed | ~root_ok failed_new = failed_new | (provisional_advance & ~aux_ok) - reached_new = reached | (t_new >= t_1 - t_eps) + reached_new = reached | (advance & reaches_horizon) num_new = num_accepted + advance.astype(jnp.int32) + num_steps_new = num_steps + (~reached & ~failed).astype(jnp.int32) + num_root_solves_new = num_root_solves + attempt_root_solves + num_root_steps_new = num_root_steps + attempt_root_steps carry_new = ( t_new, y_new, @@ -1256,6 +1550,9 @@ def accepted_derivatives(): reached_new, failed_new, num_new, + num_steps_new, + num_root_solves_new, + num_root_steps_new, controller_state_next, ) if save_at.t_1: @@ -1276,7 +1573,7 @@ def accepted_derivatives(): return carry_new, out def skip_step(carry): - t, y, z, aux, _, f_cur, z_dot, aux_dot, _, _, _, _ = carry + t, y, z, aux, _, f_cur, z_dot, aux_dot, _, _, _, _, _, _, _ = carry if save_at.t_1: out = None elif save_at.steps: @@ -1310,25 +1607,38 @@ def body(carry, _): jnp.asarray(False), ~initial_ok, jnp.asarray(0, jnp.int32), + jnp.asarray(0, jnp.int32), + initial_root_solves, + initial_root_steps, controller_state_initial, ) + if controller.uses_error_estimate and adaptive_loop == "forward": + final_carry, rows = forward_adaptive_while( + carry_0, + attempt_step=attempt_step, + skip_step=skip_step, + terminated=lambda carry: carry[8] | carry[9], + max_steps=max_steps, + ) + else: + final_carry, rows = jax.lax.scan(body, carry_0, None, length=max_steps) ( - ( - t_final, - y_final, - z_final, - aux_final, - _, - _, - _, - _, - reached, - failed, - num_accepted, - _, - ), - rows, - ) = jax.lax.scan(body, carry_0, None, length=max_steps) + t_final, + y_final, + z_final, + aux_final, + _, + _, + _, + _, + reached, + failed, + num_accepted, + num_steps, + num_root_solves, + num_root_steps, + _, + ) = final_carry integration_ok = reached & ~failed if save_at.t_1: @@ -1350,7 +1660,10 @@ def body(carry, _): zs=z_final, ok=integration_ok & aux_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=aux_final, + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) if save_at.steps: @@ -1387,12 +1700,15 @@ def body(carry, _): zs=fill_rows(compact_zs, accepted, last_z, save_at.fill), ok=integration_ok, num_accepted=num_accepted, + num_steps=num_steps, accepted=accepted, aux=( fill_rows(compact_aux, accepted, last_aux, save_at.fill) if track_aux else None ), + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) fs_all = prepend(f_initial, fs_s) @@ -1415,5 +1731,8 @@ def body(carry, _): zs=query_zs, ok=integration_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=query_aux, + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) diff --git a/src/tinydiffeq/exponential.py b/src/tinydiffeq/exponential.py index 7577537..15de6dd 100644 --- a/src/tinydiffeq/exponential.py +++ b/src/tinydiffeq/exponential.py @@ -169,6 +169,18 @@ def _arnoldi_exponential_action( ): dimension = vector.shape[0] subspace_dim = min(krylov_dim, dimension) + if subspace_dim == dimension: + # A full-dimensional Krylov space is the whole linear state space, so + # the action is exact. Materializing the operator also removes the + # coordinate singularity of an Arnoldi basis at an early happy + # breakdown (for example, an eigenvector initial state). In that case + # the exponential action is smooth even though the normalized Krylov + # basis is not; differentiating the dense action gives its unique + # JVP/VJP. This does not change the matrix-free scaling regime, where + # ``krylov_dim < dimension``. + dense_operator = jax.jacfwd(action)(jnp.zeros_like(vector)) + result = _dense_exponential_action(dense_operator, vector, time) + return result, jnp.asarray(0.0, vector.dtype) beta = jnp.linalg.norm(vector) beta_safe = jnp.where(beta > 0, beta, 1) # Each Krylov vector is a contiguous row. This orientation is materially @@ -189,17 +201,22 @@ def arnoldi_step(index, carry): correction = current_basis @ candidate candidate = candidate - correction @ current_basis coefficients = coefficients + correction - next_norm = jnp.linalg.norm(candidate) breakdown_floor = ( 100 * epsilon * jnp.maximum(jnp.asarray(1.0, vector.dtype), action_norm) ) - continues = next_norm > breakdown_floor - safe_norm = jnp.where(continues, next_norm, 1) + # ``norm`` has an undefined derivative at the exact zero residual of a + # happy breakdown. Masking ``norm(candidate)`` afterwards is too late: + # reverse mode still encounters ``0 / 0`` in the norm pullback. Branch + # on the squared norm first, so the inactive branch reaches ``sqrt`` as + # a constant and therefore has an exact-zero, finite tangent. + squared_norm = jnp.real(jnp.vdot(candidate, candidate)) + continues = squared_norm > breakdown_floor**2 + safe_norm = jnp.sqrt(jnp.where(continues, squared_norm, 1)) next_vector = jnp.where(continues, candidate / safe_norm, 0) current_basis = current_basis.at[index + 1].set(next_vector) current_hessenberg = current_hessenberg.at[:, index].set(coefficients) current_hessenberg = current_hessenberg.at[index + 1, index].set( - jnp.where(continues, next_norm, 0) + jnp.where(continues, safe_norm, 0) ) return current_basis, current_hessenberg @@ -294,7 +311,9 @@ def advance(active_carry): exponent = -one / jnp.asarray(max(method.krylov_dim - 1, 1), dtype) raw_factor = jnp.asarray(method.safety, dtype) * safe_ratio**exponent raw_factor = jnp.where( - jnp.isfinite(raw_factor), raw_factor, method.min_factor + jnp.isfinite(raw_factor), + raw_factor, + jnp.asarray(method.min_factor, dtype), ) accepted_factor = jnp.clip( raw_factor, @@ -356,13 +375,13 @@ def advance(active_carry): def _propagate_krylov(action, vector, time, method): if isinstance(method, AdaptiveKrylovExponential): - value, ok, accepted, _ = _adaptive_krylov_propagate( + value, ok, accepted, rejected = _adaptive_krylov_propagate( action, vector, time, method ) - return value, ok, accepted + return value, ok, accepted, accepted + rejected value, ok = _krylov_propagate(action, vector, time, method) accepted = jnp.where(time > 0, method.num_substeps, 0).astype(jnp.int32) - return value, ok, accepted + return value, ok, accepted, accepted def _prepare_linear_operator(operator, x_0): @@ -421,6 +440,8 @@ def solve_linear_ode(operator, method, t_0, t_1, x_0, *, save_at=None): ) if save_at is None: save_at = SaveAt(t_1=True) + if save_at.exact: + raise ValueError("SaveAt exact=True is only supported by solve_ode") if save_at.steps: raise ValueError("linear exponential solves require endpoint or SaveAt.ts") @@ -443,12 +464,12 @@ def evaluate(time): if isinstance(method, DenseExponential): value = _dense_exponential_action(dense_operator, flat_initial, elapsed) count = jnp.where(elapsed > 0, 1, 0).astype(jnp.int32) - return value, jnp.asarray(True), count + return value, jnp.asarray(True), count, count return _propagate_krylov(flat_action, flat_initial, elapsed, method) if save_at.t_1: times = t_1 - flat_states, method_ok, num_accepted = evaluate(t_1) + flat_states, method_ok, num_accepted, num_steps = evaluate(t_1) states = unravel(flat_states) times_ok = t_1 >= t_0 else: @@ -456,8 +477,9 @@ def evaluate(time): if times.ndim != 1: raise TypeError("SaveAt.ts must be one-dimensional") times_ok = jnp.all((times >= t_0) & (times <= t_1)) - flat_states, method_ok, counts = jax.vmap(evaluate)(times) - num_accepted = jnp.max(counts, initial=jnp.asarray(0, jnp.int32)) + flat_states, method_ok, accepted_counts, step_counts = jax.vmap(evaluate)(times) + num_accepted = jnp.max(accepted_counts, initial=jnp.asarray(0, jnp.int32)) + num_steps = jnp.max(step_counts, initial=jnp.asarray(0, jnp.int32)) states = jax.vmap(unravel)(flat_states) finite = jnp.all(jnp.isfinite(flat_states)) return Solution( @@ -465,6 +487,7 @@ def evaluate(time): xs=states, ok=times_ok & jnp.all(method_ok) & finite, num_accepted=num_accepted, + num_steps=num_steps, ) @@ -502,9 +525,9 @@ def _terminal_value(method, dense_operator, action, initial, elapsed): if isinstance(method, DenseExponential): exponential = jsp_linalg.expm(elapsed * dense_operator) count = jnp.where(elapsed > 0, 1, 0).astype(jnp.int32) - return exponential @ initial, jnp.asarray(True), exponential, count - value, ok, count = _propagate_krylov(action, initial, elapsed, method) - return value, ok, None, count + return exponential @ initial, jnp.asarray(True), exponential, count, count + value, ok, accepted, steps = _propagate_krylov(action, initial, elapsed, method) + return value, ok, None, accepted, steps def jvp_linear_ode( @@ -552,18 +575,18 @@ def jvp_linear_ode( t_0 = jnp.asarray(t_0, dtype) t_1 = jnp.asarray(t_1, dtype) elapsed = t_1 - t_0 - flat_value, primal_ok, exponential, num_accepted = _terminal_value( + flat_value, primal_ok, exponential, num_accepted, num_steps = _terminal_value( method, dense_operator, flat_action, flat_initial, elapsed ) if isinstance(method, DenseExponential): flat_tangent = tangent @ exponential.T if batched else exponential @ tangent tangent_ok = jnp.asarray(True) elif batched: - flat_tangent, tangent_ok, _ = jax.vmap( + flat_tangent, tangent_ok, _, _ = jax.vmap( lambda direction: _propagate_krylov(flat_action, direction, elapsed, method) )(tangent) else: - flat_tangent, tangent_ok, _ = _propagate_krylov( + flat_tangent, tangent_ok, _, _ = _propagate_krylov( flat_action, tangent, elapsed, method ) finite = jnp.all(jnp.isfinite(flat_value)) & jnp.all(jnp.isfinite(flat_tangent)) @@ -572,6 +595,7 @@ def jvp_linear_ode( xs=unravel(flat_value), ok=(t_1 >= t_0) & primal_ok & jnp.all(tangent_ok) & finite, num_accepted=num_accepted, + num_steps=num_steps, ) return solution, restore_tangent(flat_tangent) @@ -621,7 +645,7 @@ def vjp_linear_ode( t_0 = jnp.asarray(t_0, dtype) t_1 = jnp.asarray(t_1, dtype) elapsed = t_1 - t_0 - flat_value, primal_ok, exponential, num_accepted = _terminal_value( + flat_value, primal_ok, exponential, num_accepted, num_steps = _terminal_value( method, dense_operator, flat_action, flat_initial, elapsed ) if isinstance(method, DenseExponential): @@ -636,13 +660,13 @@ def transpose_action(vector): return jax.linear_transpose(flat_action, zero)(vector)[0] if batched: - flat_gradient, gradient_ok, _ = jax.vmap( + flat_gradient, gradient_ok, _, _ = jax.vmap( lambda direction: _propagate_krylov( transpose_action, direction, elapsed, method ) )(flat_cotangent) else: - flat_gradient, gradient_ok, _ = _propagate_krylov( + flat_gradient, gradient_ok, _, _ = _propagate_krylov( transpose_action, flat_cotangent, elapsed, method ) finite = jnp.all(jnp.isfinite(flat_value)) & jnp.all(jnp.isfinite(flat_gradient)) @@ -651,5 +675,6 @@ def transpose_action(vector): xs=unravel(flat_value), ok=(t_1 >= t_0) & primal_ok & jnp.all(gradient_ok) & finite, num_accepted=num_accepted, + num_steps=num_steps, ) return solution, restore_cotangent(flat_gradient) diff --git a/src/tinydiffeq/markov.py b/src/tinydiffeq/markov.py index a56e507..79fe2e9 100644 --- a/src/tinydiffeq/markov.py +++ b/src/tinydiffeq/markov.py @@ -243,6 +243,11 @@ class MarkovDistribution: ok: jax.Array +def _reject_exact_save_at(save_at): + if save_at.exact: + raise ValueError("SaveAt exact=True is only supported by solve_ode") + + def _validate_simulation_inputs(chain, state_0, count, count_name, save_at): if not isinstance(count, int) or isinstance(count, bool): raise TypeError(f"{count_name} must be a static Python int") @@ -250,6 +255,7 @@ def _validate_simulation_inputs(chain, state_0, count, count_name, save_at): raise ValueError(f"{count_name} must be at least 1") if save_at is None: save_at = SaveAt(t_1=True) + _reject_exact_save_at(save_at) if save_at.steps and save_at.fill != "last": raise ValueError( 'Markov integer states require SaveAt(steps=True, fill="last")' @@ -273,6 +279,21 @@ def _alias_sample(chain, state, uniform): ).astype(jnp.int32) +def _unit_uniform(key, shape, dtype): + return jax.random.uniform( + key, + shape, + dtype=dtype, + minval=jnp.asarray(0.0, dtype), + maxval=jnp.asarray(1.0, dtype), + ) + + +def _unit_exponential(key, shape, dtype): + uniform = _unit_uniform(key, shape, dtype) + return -jnp.log1p(-uniform) + + def _random_maps(chain, uniforms): scaled = uniforms * chain.num_states columns = jnp.minimum(scaled.astype(jnp.int32), chain.num_states - 1) @@ -326,20 +347,25 @@ def simulate_markov_chain( state_0, initial_ok, save_at = _validate_simulation_inputs( chain, state_0, num_steps, "num_steps", save_at ) - uniforms = jax.random.uniform( - key, (num_steps,), dtype=chain.transition_matrix.dtype - ) + uniforms = _unit_uniform(key, (num_steps,), chain.transition_matrix.dtype) step_states = _simulate_discrete_states(chain, state_0, uniforms, method) all_states = prepend(state_0, step_states) count = jnp.asarray(num_steps, jnp.int32) if save_at.t_1: - return Solution(ts=count, xs=step_states[-1], ok=initial_ok, num_accepted=count) + return Solution( + ts=count, + xs=step_states[-1], + ok=initial_ok, + num_accepted=count, + num_steps=count, + ) if save_at.steps: return Solution( ts=jnp.arange(num_steps + 1, dtype=jnp.int32), xs=all_states, ok=initial_ok, num_accepted=count, + num_steps=count, accepted=jnp.ones((num_steps + 1,), dtype=bool), ) query_steps = jnp.asarray(save_at.ts) @@ -352,11 +378,13 @@ def simulate_markov_chain( xs=all_states[safe_queries], ok=initial_ok & queries_ok, num_accepted=count, + num_steps=count, ) def _simulate_continuous_events(chain, state_0, exponentials, uniforms, method): - safe_rates = jnp.where(chain.exit_rates > 0, chain.exit_rates, jnp.inf) + infinity = jnp.asarray(jnp.inf, chain.exit_rates.dtype) + safe_rates = jnp.where(chain.exit_rates > 0, chain.exit_rates, infinity) if isinstance(method, SequentialMarkov): def step(carry, random_values): @@ -365,7 +393,7 @@ def step(carry, random_values): holding_time = jnp.where( chain.exit_rates[state] > 0, exponential / safe_rates[state], - jnp.inf, + infinity, ) next_state = _alias_sample(chain, state, uniform) next_time = time + holding_time @@ -382,7 +410,7 @@ def step(carry, random_values): holding_times = jnp.where( chain.exit_rates[None, :] > 0, exponentials[:, None] / safe_rates[None, :], - jnp.inf, + infinity, ) def compose(earlier, later): @@ -433,13 +461,16 @@ def simulate_continuous_time_markov_chain( t_1 = jnp.asarray(t_1, dtype) time_ok = t_1 >= t_0 exponential_key, transition_key = jax.random.split(key) - exponentials = jax.random.exponential(exponential_key, (max_jumps,), dtype=dtype) - uniforms = jax.random.uniform(transition_key, (max_jumps,), dtype=dtype) + exponentials = _unit_exponential(exponential_key, (max_jumps,), dtype) + uniforms = _unit_uniform(transition_key, (max_jumps,), dtype) states_after, elapsed_event_times = _simulate_continuous_events( chain, state_0, exponentials, uniforms, method ) event_times = t_0 + elapsed_event_times num_jumps = jnp.sum(event_times <= t_1, dtype=jnp.int32) + # One event beyond the horizon is needed to establish endpoint coverage; + # a starved event budget has already attempted every available event. + num_steps = jnp.minimum(num_jumps + 1, jnp.asarray(max_jumps, jnp.int32)) all_states = prepend(state_0, states_after) final_state = all_states[num_jumps] covered = event_times[-1] >= t_1 @@ -452,6 +483,7 @@ def simulate_continuous_time_markov_chain( xs=final_state, ok=integration_ok, num_accepted=num_jumps, + num_steps=num_steps, ) if save_at.steps: accepted = jnp.arange(max_jumps + 1) <= num_jumps @@ -463,6 +495,7 @@ def simulate_continuous_time_markov_chain( xs=output_states, ok=integration_ok, num_accepted=num_jumps, + num_steps=num_steps, accepted=accepted, ) query_times = jnp.asarray(save_at.ts, dtype) @@ -478,6 +511,7 @@ def simulate_continuous_time_markov_chain( xs=all_states[event_counts], ok=integration_ok & queries_ok, num_accepted=num_jumps, + num_steps=num_steps, ) @@ -601,6 +635,7 @@ def forecast_markov_chain( raise ValueError("num_steps must be nonnegative") if save_at is None: save_at = SaveAt(t_1=True) + _reject_exact_save_at(save_at) distribution_0, initial_ok = _prepare_distribution(chain, distribution_0) if save_at.t_1: @@ -702,6 +737,7 @@ def forecast_continuous_time_markov_chain( ) if save_at is None: save_at = SaveAt(t_1=True) + _reject_exact_save_at(save_at) if save_at.steps: raise ValueError("CTMC distribution forecasts require endpoint or SaveAt.ts") if is_dense: diff --git a/src/tinydiffeq/ode.py b/src/tinydiffeq/ode.py index a308932..64045bc 100644 --- a/src/tinydiffeq/ode.py +++ b/src/tinydiffeq/ode.py @@ -1,7 +1,9 @@ import inspect +import numbers import jax import jax.numpy as jnp +import numpy as np from jax.flatten_util import ravel_pytree from tinydiffeq._aux import ( @@ -11,6 +13,7 @@ split_field_output, zeros_from_shape, ) +from tinydiffeq._loop import forward_adaptive_while, select_step from tinydiffeq._rodas5p import rodas5p_step, rodas_dense_endpoint_derivatives from tinydiffeq._tree import ( asarray_state, @@ -35,6 +38,45 @@ ADAPTIVE_SCAN_CHUNK_SIZE = 16 +def _has_static_uniform_horizon(t_0, t_1, dt_0, max_steps, dtype): + """Return whether ordinary scalar inputs define the branch-free fixed grid. + + JAX batches a data-dependent ``lax.cond`` by selecting between the results + of both branches. If those branches are complete scans, a ``vmap`` over + horizons therefore performs both integrations. Restrict the branch-free + scan to horizons known before tracing; traced or array-valued horizons use + the single clipped scan below. + """ + values = (t_0, t_1, dt_0) + if not all( + isinstance(value, numbers.Real) and not isinstance(value, bool) + for value in values + ): + return False + try: + numpy_dtype = np.dtype(dtype) + t_0_value, t_1_value, dt_0_value = (numpy_dtype.type(value) for value in values) + if not dt_0_value > 0: + return False + nominal_final = numpy_dtype.type( + t_0_value + numpy_dtype.type(max_steps) * dt_0_value + ) + nominal_penultimate = numpy_dtype.type( + t_0_value + numpy_dtype.type(max_steps - 1) * dt_0_value + ) + scale = max(numpy_dtype.type(1.0), abs(t_0_value), abs(t_1_value)) + tolerance = min( + numpy_dtype.type(4.0) * np.finfo(numpy_dtype).eps * scale, + numpy_dtype.type(0.25) * abs(dt_0_value), + ) + except (TypeError, ValueError): + return False + penultimate_reaches = (nominal_penultimate >= t_1_value) or ( + abs(nominal_penultimate - t_1_value) <= tolerance + ) + return bool(abs(nominal_final - t_1_value) <= tolerance and not penultimate_reaches) + + def identity_project(x): return x @@ -89,6 +131,7 @@ def solve_ode( project=None, has_aux=None, failure_ad_reference=None, + adaptive_loop="bounded", ): """Integrate ``dx/dt = f(x, t, args, p)`` from ``t_0`` to ``t_1 > t_0``. @@ -99,10 +142,16 @@ def solve_ode( parameters (any pytree); jvp/vjp with respect to ``p`` and ``x_0`` are first-class. - Fixed and adaptive stepping use bounded ``lax.scan`` loops with exactly - ``max_steps`` attempt slots, so shapes are static and curvature-dependent - step counts never retrace. Adaptive attempts are grouped into static - chunks so one ``lax.cond`` skips an entire padded chunk after completion. + Fixed stepping and the default ``adaptive_loop="bounded"`` use bounded + ``lax.scan`` loops with exactly ``max_steps`` attempt slots, so shapes are + static and curvature-dependent step counts never retrace. Bounded adaptive + attempts are grouped into static chunks so one ``lax.cond`` skips an entire + padded chunk after completion. ``adaptive_loop="forward"`` instead uses a + true ``lax.while_loop`` and stops after the actual attempts. It supports + primal evaluation, JVP, and nested forward-mode AD; ordinary reverse mode + is unavailable because JAX cannot transpose a dynamic ``while_loop``. + Under ``vmap``, the forward loop advances lanes together until the slowest + live trajectory finishes rather than running every lane to ``max_steps``. ``dt_0`` is required (no auto-initial-step heuristic). Each attempt is clipped to the remaining horizon; the clipped step also feeds the controller's next-step proposal, which doubles as the growth @@ -117,14 +166,21 @@ def solve_ode( detects the form with an abstract trace; ``has_aux=False`` selects the minimal no-aux path without that trace. Saved aux is a nonempty pytree of real floating arrays. It follows ``SaveAt`` and participates in JVP/VJP. - Requested-grid aux uses cubic Hermite interpolation with endpoint slopes - obtained by JVP, including for Rodas5P's dense state path. + By default, requested-grid aux uses cubic Hermite interpolation with + endpoint slopes obtained by JVP, including for Rodas5P's dense state path. + With ``SaveAt(ts=..., exact=True)``, an explicit fixed-step solve instead + selects realized knots and evaluates aux directly at those requested knots. The time dtype follows the state dtype; the library never sets ``jax_enable_x64`` — do that in your application. """ if dt_0 is None: raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)") + if adaptive_loop not in ("bounded", "forward"): + raise ValueError( + 'adaptive_loop must be either "bounded" or "forward", got ' + f"{adaptive_loop!r}" + ) if save_at is None: save_at = SaveAt(t_1=True) if controller is None: @@ -136,21 +192,28 @@ def solve_ode( f"{type(controller).__name__} needs an embedded error estimate, " f"which {type(solver).__name__} does not provide" ) + if adaptive_loop == "forward" and not controller.uses_error_estimate: + raise ValueError( + 'adaptive_loop="forward" requires an adaptive error controller' + ) f = canonicalize_field(f) is_rodas = isinstance(solver, Rodas5P) is_fixed = isinstance(controller, ConstantStepSize) + if save_at.exact and (not is_fixed or is_rodas): + raise ValueError( + "SaveAt exact=True requires an explicit solver with ConstantStepSize" + ) + original_times = (t_0, t_1, dt_0) x_0, time_dtype = asarray_state(x_0, "x_0") + static_uniform_horizon = _has_static_uniform_horizon( + *original_times, max_steps, time_dtype + ) t_0 = jnp.asarray(t_0, time_dtype) t_1 = jnp.asarray(t_1, time_dtype) dt_0 = jnp.asarray(dt_0, time_dtype) - positive_time_floor = jnp.asarray(jnp.finfo(time_dtype).tiny, time_dtype) - t_eps = 4.0 * jnp.finfo(time_dtype).eps * jnp.maximum(1.0, jnp.abs(t_1)) - # Summing ~max_steps rounded steps can leave t short of t_1 by up to - # ~max_steps * eps, so a step whose remaining horizon is within that - # slack of the desired dt is stretched to land on t_1 exactly; otherwise - # dt_0 = (t_1 - t_0)/n with max_steps = n would strand a one-ulp sliver. - t_slack = max_steps * t_eps + time_scale = jnp.maximum(jnp.maximum(1.0, jnp.abs(t_0)), jnp.abs(t_1)) + t_eps = 4.0 * jnp.finfo(time_dtype).eps * time_scale def project_state(x): value, dtype = asarray_state(project(x), "project(x)") @@ -202,9 +265,11 @@ def aux_value_and_derivative(x, t, x_dot, active): evaluate_aux = None zero_aux = None - need_f = not is_rodas and (solver.fsal or (save_at.ts is not None)) + need_f = not is_rodas and ( + solver.fsal or (save_at.ts is not None and not save_at.exact) + ) f_init = g(x_0, t_0) if need_f else zeros_like(x_0) - track_aux = has_aux and not save_at.t_1 + track_aux = has_aux and not save_at.t_1 and not save_at.exact if track_aux: aux_init, aux_init_ok = evaluate_aux( (x_0, t_0, p), jnp.asarray(True), failure_ad_reference @@ -234,13 +299,18 @@ def attempt_step(carry): done, failed, num_accepted, + num_steps, controller_state, ) = carry - remaining = t_1 - t - h = jnp.where( - remaining <= dt + t_slack, - jnp.maximum(remaining, positive_time_floor), + h, proposed_t, reaches_horizon = select_step( + t, + t_0, + t_1, dt, + dt_0, + num_steps, + constant=is_fixed, + time_tolerance=t_eps, ) if is_rodas: x_1, err, dense, step_ok = rodas5p_step( @@ -251,7 +321,7 @@ def attempt_step(carry): step = solver.step_fixed if is_fixed else solver.step x_1, f_1, err = step(g, t, x, h, f_cur if need_f else None, project_state) if need_f and f_1 is None: - f_1 = g(x_1, t + h) + f_1 = g(x_1, proposed_t) dense = None step_ok = jnp.asarray(True) if is_rodas: @@ -271,7 +341,7 @@ def attempt_step(carry): def accepted_auxiliary(): if save_at.ts is None: aux_candidate, aux_ok = evaluate_aux( - (x_1, t + h, p), + (x_1, proposed_t, p), provisional_advance, failure_ad_reference, ) @@ -284,11 +354,11 @@ def accepted_auxiliary(): x, t, left_dot, provisional_advance ) aux_candidate, aux_ok, aux_right_dot = aux_value_and_derivative( - x_1, t + h, right_dot, provisional_advance + x_1, proposed_t, right_dot, provisional_advance ) return aux_candidate, aux_ok, aux_left_dot, aux_right_dot aux_candidate, aux_ok, aux_right_dot = aux_value_and_derivative( - x_1, t + h, f_1, provisional_advance + x_1, proposed_t, f_1, provisional_advance ) return aux_candidate, aux_ok, aux_dot, aux_right_dot @@ -310,7 +380,7 @@ def accepted_auxiliary(): if track_aux and save_at.ts is not None and not is_rodas else aux_dot ) - t_new = jnp.where(advance, t + h, t) + t_new = jnp.where(advance, proposed_t, t) f_new = where(advance, f_1, f_cur) if need_f else f_cur dt_new = jnp.where(done | failed, dt, dt_next) controller_state_new = jax.tree.map( @@ -318,10 +388,11 @@ def accepted_auxiliary(): controller_state, controller_state_next, ) - done_new = done | (t_new >= t_1 - t_eps) + done_new = done | (advance & reaches_horizon) failed_new = failed if controller.uses_error_estimate else failed | ~step_ok failed_new = failed_new | (provisional_advance & ~aux_ok) num_new = num_accepted + advance.astype(jnp.int32) + num_steps_new = num_steps + (~done & ~failed).astype(jnp.int32) carry_new = ( t_new, x_new, @@ -332,6 +403,7 @@ def accepted_auxiliary(): done_new, failed_new, num_new, + num_steps_new, controller_state_new, ) if save_at.t_1: @@ -353,7 +425,7 @@ def accepted_auxiliary(): return carry_new, out def skip_step(carry): - t, x, aux, aux_dot, _, f_cur, _, _, _, _ = carry + t, x, aux, aux_dot, _, f_cur, _, _, _, _, _ = carry if save_at.t_1: out = None elif save_at.steps: @@ -377,23 +449,26 @@ def body(carry, _): def fixed_attempt_step(carry): t, x, f_cur, done, num_accepted = carry - remaining = t_1 - t - h = jnp.where( - remaining <= dt_0 + t_slack, - jnp.maximum(remaining, positive_time_floor), + h, t_1_step, reaches_horizon = select_step( + t, + t_0, + t_1, dt_0, + dt_0, + num_accepted, + constant=True, + time_tolerance=t_eps, ) x_1, f_1, _ = solver.step_fixed( g, t, x, h, f_cur if need_f else None, project_state ) if need_f and f_1 is None: - f_1 = g(x_1, t + h) - t_1_step = t + h - done_1 = t_1_step >= t_1 - t_eps + f_1 = g(x_1, t_1_step) + done_1 = reaches_horizon carry_1 = (t_1_step, x_1, f_1 if need_f else f_cur, done_1, num_accepted + 1) if save_at.t_1: output = None - elif save_at.steps: + elif save_at.steps or save_at.exact: output = (t_1_step, x_1, jnp.asarray(True)) else: output = (t_1_step, x_1, f_1, jnp.asarray(True)) @@ -403,7 +478,7 @@ def fixed_skip_step(carry): t, x, f_cur, _, _ = carry if save_at.t_1: output = None - elif save_at.steps: + elif save_at.steps or save_at.exact: output = (t, x, jnp.asarray(False)) else: output = (t, x, f_cur, jnp.asarray(False)) @@ -412,6 +487,38 @@ def fixed_skip_step(carry): def fixed_body(carry, _): return jax.lax.cond(carry[3], fixed_skip_step, fixed_attempt_step, carry) + def uniform_fixed_body(carry, step_index): + t, x, f_cur, _, num_accepted = carry + h, t_1_step, reaches_horizon = select_step( + t, + t_0, + t_1, + dt_0, + dt_0, + step_index, + constant=True, + time_tolerance=t_eps, + ) + x_1, f_1, _ = solver.step_fixed( + g, t, x, h, f_cur if need_f else None, project_state + ) + if need_f and f_1 is None: + f_1 = g(x_1, t_1_step) + carry_1 = ( + t_1_step, + x_1, + f_1 if need_f else f_cur, + reaches_horizon, + num_accepted + 1, + ) + if save_at.t_1: + output = None + elif save_at.steps or save_at.exact: + output = (t_1_step, x_1, jnp.asarray(True)) + else: + output = (t_1_step, x_1, f_1, jnp.asarray(True)) + return carry_1, output + def bounded_adaptive_scan(carry): chunk_size = min(ADAPTIVE_SCAN_CHUNK_SIZE, max_steps) num_chunks = (max_steps + chunk_size - 1) // chunk_size @@ -467,6 +574,7 @@ def outer(chunk_carry, chunk_valid): jnp.asarray(False), ~aux_init_ok, jnp.asarray(0, jnp.int32), + jnp.asarray(0, jnp.int32), controller_state_init, ) use_fast_fixed = is_fixed and not is_rodas and not track_aux @@ -478,17 +586,69 @@ def outer(chunk_carry, chunk_valid): jnp.asarray(False), jnp.asarray(0, jnp.int32), ) - fixed_final, rows = jax.lax.scan( - fixed_body, fixed_carry_0, None, length=max_steps - ) + if static_uniform_horizon: + step_indices = jnp.arange(max_steps, dtype=jnp.int32) + fixed_final, rows = jax.lax.scan( + uniform_fixed_body, fixed_carry_0, step_indices + ) + else: + fixed_final, rows = jax.lax.scan( + fixed_body, fixed_carry_0, None, length=max_steps + ) t_final, x_final, _, done, num_accepted = fixed_final + num_steps = num_accepted failed = jnp.asarray(False) + elif controller.uses_error_estimate and adaptive_loop == "forward": + final_carry, rows = forward_adaptive_while( + carry_0, + attempt_step=attempt_step, + skip_step=skip_step, + terminated=lambda carry: carry[6] | carry[7], + max_steps=max_steps, + ) + ( + t_final, + x_final, + _, + _, + _, + _, + done, + failed, + num_accepted, + num_steps, + _, + ) = final_carry elif controller.uses_error_estimate: final_carry, rows = bounded_adaptive_scan(carry_0) - (t_final, x_final, _, _, _, _, done, failed, num_accepted, _) = final_carry + ( + t_final, + x_final, + _, + _, + _, + _, + done, + failed, + num_accepted, + num_steps, + _, + ) = final_carry else: final_carry, rows = jax.lax.scan(body, carry_0, None, length=max_steps) - (t_final, x_final, _, _, _, _, done, failed, num_accepted, _) = final_carry + ( + t_final, + x_final, + _, + _, + _, + _, + done, + failed, + num_accepted, + num_steps, + _, + ) = final_carry integration_ok = done & ~failed if save_at.t_1: @@ -504,10 +664,11 @@ def outer(chunk_carry, chunk_valid): xs=x_final, ok=integration_ok & aux_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=aux_final, ) - if save_at.steps: + if save_at.steps or save_at.exact: if use_fast_fixed: ts_s, xs_s, adv_s = rows aux_s = None @@ -532,6 +693,57 @@ def outer(chunk_carry, chunk_valid): ts_s, xs_s, aux_s, fs_s, aux_dots_s, adv_s = rows all_times = jnp.concatenate([t_0[None], ts_s]) all_states = prepend(x_0, xs_s) + if save_at.exact: + query_times = jnp.asarray(save_at.ts, time_dtype) + uniform_indices = jnp.rint((query_times - t_0) / dt_0).astype(jnp.int32) + uniform_indices = jnp.clip(uniform_indices, 0, max_steps) + uniform_times = all_times[uniform_indices] + final_index = jnp.clip(num_accepted, 0, max_steps) + final_times = jnp.broadcast_to(all_times[final_index], query_times.shape) + use_final = jnp.abs(query_times - final_times) < jnp.abs( + query_times - uniform_times + ) + query_indices = jnp.where(use_final, final_index, uniform_indices) + exact_times = all_times[query_indices] + alignment_scale = jnp.maximum( + jnp.maximum(1.0, jnp.abs(t_0)), + jnp.maximum(jnp.abs(query_times), jnp.abs(exact_times)), + ) + alignment_tolerance = jnp.minimum( + 8.0 * jnp.finfo(time_dtype).eps * alignment_scale, + 0.25 * jnp.abs(dt_0), + ) + aligned = ( + (query_times >= t_0) + & (query_times <= t_final) + & (jnp.abs(query_times - exact_times) <= alignment_tolerance) + & (query_indices <= num_accepted) + ) + query_states = take(all_states, query_indices) + if has_aux: + + def exact_auxiliary(x_value, t_value): + return evaluate_aux( + (x_value, t_value, p), + jnp.asarray(True), + failure_ad_reference, + ) + + query_aux, query_aux_ok = jax.vmap(exact_auxiliary)( + query_states, exact_times + ) + aux_ok = jnp.all(query_aux_ok) + else: + query_aux = None + aux_ok = jnp.asarray(True) + return Solution( + ts=query_times, + xs=query_states, + ok=integration_ok & jnp.all(aligned) & aux_ok, + num_accepted=num_accepted, + num_steps=num_steps, + aux=query_aux, + ) all_aux = prepend(aux_init, aux_s) if has_aux else None raw_accepted = jnp.concatenate([jnp.ones((1,), bool), adv_s]) @@ -555,6 +767,7 @@ def outer(chunk_carry, chunk_valid): xs=output_states, ok=integration_ok, num_accepted=num_accepted, + num_steps=num_steps, accepted=accepted, aux=( fill_rows(compact_aux, accepted, last_aux, save_at.fill) @@ -591,5 +804,6 @@ def outer(chunk_carry, chunk_valid): xs=query_states, ok=integration_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=query_aux, ) diff --git a/src/tinydiffeq/save_at.py b/src/tinydiffeq/save_at.py index 24d0407..722665a 100644 --- a/src/tinydiffeq/save_at.py +++ b/src/tinydiffeq/save_at.py @@ -17,6 +17,10 @@ class SaveAt: takes, so changing curvature never changes shapes or recompiles. ``ts`` is a data leaf; a different grid of the same length retraces nothing. ODE, deterministic DAE, and linear exponential solves only. + For an explicit constant-step ODE, ``exact=True`` instead requires every + query to coincide with an internal step endpoint. It gathers those states + directly and avoids both interpolation and its extra endpoint-slope field + evaluation. - ``steps=True``: the initial state and accepted internal steps as a chronological prefix of a ``max_steps + 1`` buffer. Rejected attempts are omitted. ``fill="last"`` (default) pads the tail with the final @@ -30,6 +34,7 @@ class SaveAt: ts: ArrayLike | None = None steps: bool = field(default=False, metadata=dict(static=True)) fill: str = field(default="last", metadata=dict(static=True)) + exact: bool = field(default=False, metadata=dict(static=True)) def __post_init__(self): modes = int(bool(self.t_1)) + int(self.ts is not None) + int(bool(self.steps)) @@ -39,3 +44,5 @@ def __post_init__(self): ) if self.fill not in ("last", "inf"): raise ValueError('SaveAt fill must be "last" or "inf"') + if self.exact and self.ts is None: + raise ValueError("SaveAt exact=True requires ts=...") diff --git a/src/tinydiffeq/sdae.py b/src/tinydiffeq/sdae.py index 1344a47..d510c97 100644 --- a/src/tinydiffeq/sdae.py +++ b/src/tinydiffeq/sdae.py @@ -61,6 +61,8 @@ def solve_semi_explicit_sdae( fixed uniform grid, then an algebraic root solve restores consistency at every node. This is Euler--Maruyama applied to the reduced SDE obtained from the locally unique root ``z = Z(y, t)``. + ``sol.num_root_solves`` counts logical active nonlinear root calls and + ``sol.num_root_steps`` sums their LM update steps. ``drift`` may return ``value`` or ``(value, saved_aux)``. If ``g`` returns ``(residual, algebraic_aux)``, that internal context is passed to both @@ -195,7 +197,12 @@ def checked_value(output, name): raise TypeError(f"{name} must preserve the y dtype") return value - z_initial, initial_root_ok = solve_root(y_0, t_0, z_0, jnp.asarray(True)) + ( + z_initial, + initial_root_ok, + initial_root_solves, + initial_root_steps, + ) = solve_root(y_0, t_0, z_0, jnp.asarray(True)) if has_algebraic_aux: context_initial, initial_context_ok = evaluate_context( y_0, z_initial, t_0, initial_root_ok @@ -224,7 +231,18 @@ def checked_value(output, name): context_reference = None def attempt_step(carry, inputs): - y, z, context, aux, t, failed, num_accepted = carry + ( + y, + z, + context, + aux, + t, + failed, + num_accepted, + num_steps, + num_root_solves, + num_root_steps, + ) = carry t_step, t_next, d_w_step = inputs active = ~failed y_ref, z_ref, t_ref, p_ref = failure_ad_reference @@ -248,7 +266,9 @@ def attempt_step(carry, inputs): (dt, drift_value), (1.0, multiply(diffusion_value, d_w_step)), ) - z_candidate, root_ok = solve_root(y_candidate, t_next, z, active) + z_candidate, root_ok, root_solves, root_steps = solve_root( + y_candidate, t_next, z, active + ) if has_algebraic_aux: context_candidate, context_ok = evaluate_context( y_candidate, z_candidate, t_next, root_ok & active @@ -304,14 +324,17 @@ def accepted_aux(): t_new, failed_new, num_new, + num_steps + jnp.asarray(1, jnp.int32), + num_root_solves + root_solves, + num_root_steps + root_steps, ) out = (t_new, y_new, z_new, aux_new, advance) if save_at.steps else None return carry_new, out def skip_step(carry, _): - y, z, context, aux, t, failed, num_accepted = carry + y, z, context, aux, t, failed, num_accepted, _, _, _ = carry out = (t, y, z, aux, jnp.asarray(False)) if save_at.steps else None - return (y, z, context, aux, t, failed, num_accepted), out + return carry, out def body(carry, inputs): return jax.lax.cond( @@ -329,6 +352,9 @@ def body(carry, inputs): t_0, ~initial_ok, jnp.asarray(0, jnp.int32), + jnp.asarray(0, jnp.int32), + initial_root_solves, + initial_root_steps, ) ( ( @@ -339,6 +365,9 @@ def body(carry, inputs): t_final, failed, num_accepted, + num_steps, + num_root_solves, + num_root_steps, ), rows, ) = jax.lax.scan( @@ -361,7 +390,10 @@ def body(carry, inputs): zs=z_final, ok=ok & aux_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=aux_final, + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) ts_s, ys_s, zs_s, aux_s, advance_s = rows @@ -385,8 +417,11 @@ def body(carry, inputs): zs=fill_rows(all_zs, accepted, last_z, save_at.fill), ok=ok, num_accepted=num_accepted, + num_steps=num_steps, accepted=accepted, aux=( fill_rows(all_aux, accepted, last_aux, save_at.fill) if track_aux else None ), + num_root_solves=num_root_solves, + num_root_steps=num_root_steps, ) diff --git a/src/tinydiffeq/sde.py b/src/tinydiffeq/sde.py index 4a57a89..781be75 100644 --- a/src/tinydiffeq/sde.py +++ b/src/tinydiffeq/sde.py @@ -149,6 +149,7 @@ def body(x, inputs): if save_at.t_1 or not has_aux: x_final, step_states = jax.lax.scan(body, x_0, (time_grid[:-1], d_w)) num_accepted = jnp.asarray(n_steps, jnp.int32) + num_steps = num_accepted ok = jnp.asarray(True) if save_at.t_1: @@ -175,6 +176,7 @@ def auxiliary(inputs): xs=x_final, ok=ok & aux_ok, num_accepted=num_accepted, + num_steps=num_steps, aux=aux_final, ) @@ -191,7 +193,7 @@ def auxiliary(inputs): ) def aux_attempt(carry, inputs): - x, aux, t, failed, count = carry + x, aux, t, failed, count, num_steps = carry t_step, t_next, d_w_step = inputs x_candidate = solver.step( g_drift, @@ -213,16 +215,18 @@ def aux_attempt(carry, inputs): t_new = jnp.where(advance, t_next, t) failed_new = failed | ~aux_ok count_new = count + advance.astype(jnp.int32) + num_steps_new = num_steps + jnp.asarray(1, jnp.int32) return ( x_new, aux_new, t_new, failed_new, count_new, + num_steps_new, ), (t_new, x_new, aux_new, advance) def aux_skip(carry, inputs): - x, aux, t, failed, count = carry + x, aux, t, failed, count, num_steps = carry return carry, (t, x, aux, jnp.asarray(False)) def aux_body(carry, inputs): @@ -239,12 +243,19 @@ def aux_body(carry, inputs): t_0, ~initial_ok, jnp.asarray(0, jnp.int32), + jnp.asarray(0, jnp.int32), ) - (x_final, aux_final, t_final, failed, num_accepted), rows = jax.lax.scan( - aux_body, - carry_0, - (time_grid[:-1], time_grid[1:], d_w), - ) + ( + ( + x_final, + aux_final, + t_final, + failed, + num_accepted, + num_steps, + ), + rows, + ) = jax.lax.scan(aux_body, carry_0, (time_grid[:-1], time_grid[1:], d_w)) ts_s, xs_s, aux_s, advance_s = rows all_times = jnp.concatenate([t_0[None], ts_s]) all_states = prepend(x_0, xs_s) @@ -263,6 +274,7 @@ def aux_body(carry, inputs): xs=fill_rows(all_states, accepted, last_state, save_at.fill), ok=~failed & (num_accepted == n_steps), num_accepted=num_accepted, + num_steps=num_steps, accepted=accepted, aux=fill_rows(all_aux, accepted, last_aux, save_at.fill), ) @@ -274,5 +286,6 @@ def aux_body(carry, inputs): xs=all_states, ok=ok, num_accepted=num_accepted, + num_steps=num_steps, accepted=accepted, ) diff --git a/src/tinydiffeq/solution.py b/src/tinydiffeq/solution.py index ed57c53..4b61c7d 100644 --- a/src/tinydiffeq/solution.py +++ b/src/tinydiffeq/solution.py @@ -18,6 +18,9 @@ class Solution: ``jnp.where(sol.ok, x, jnp.inf)`` over leaves. - ``num_accepted``: number of accepted steps (excluding the initial state). + - ``num_steps``: scalar integer array counting logical attempted steps, + including rejected attempts. Public solver outputs always populate it; + the optional default only preserves direct-construction compatibility. - ``accepted``: ``steps`` mode only (otherwise None): validity mask for the contiguous accepted-step prefix. Row 0 (the initial state) is always True, so ``accepted.sum() == num_accepted + 1``. @@ -32,6 +35,7 @@ class Solution: num_accepted: jax.Array accepted: jax.Array | None = None aux: Any = None + num_steps: jax.Array | None = None @jax.tree_util.register_dataclass @@ -49,7 +53,14 @@ class DAESolution: after its initial consistency root. Requested values are dense interpolants and need not satisfy the constraint exactly. ``ok`` is true only when initialization, all required stage and saved-output operations - succeeded, and the integration reached ``t_1``. + succeeded, and the integration reached ``t_1``. ``num_root_solves`` counts + logical active algebraic root calls, including failed calls and the initial + consistency solve. ``num_root_steps`` sums their nonlinear LM update steps. + ``num_steps`` counts logical attempted integration steps, including rejected + attempts, with the same semantics as :class:`Solution`. + All three are path diagnostics with exact-zero tangents. Under ``vmap``, masked + lanes may still execute physically even though they do not increment these + logical per-lane counters. """ ts: jax.Array @@ -59,3 +70,6 @@ class DAESolution: num_accepted: jax.Array accepted: jax.Array | None = None aux: Any = None + num_steps: jax.Array | None = None + num_root_solves: jax.Array | None = None + num_root_steps: jax.Array | None = None diff --git a/tests/test_adaptive.py b/tests/test_adaptive.py index 4846df6..178578b 100644 --- a/tests/test_adaptive.py +++ b/tests/test_adaptive.py @@ -249,6 +249,224 @@ def endpoint(x_0, max_steps): assert jnp.isfinite(gradient) +@pytest.mark.parametrize( + "save_at", + [ + SaveAt(t_1=True), + SaveAt(steps=True), + SaveAt(ts=jnp.linspace(0.0, 2.0, 11)), + ], +) +def test_forward_adaptive_loop_matches_bounded_primal_and_forward_ad(save_at): + def run(parameter, adaptive_loop): + return solve_ode( + lambda x, t, args, p: -p * x, + Tsit5(), + 0.0, + 2.0, + jnp.asarray(1.0), + p=parameter, + dt_0=0.1, + controller=IController(rtol=1e-9, atol=1e-11), + max_steps=96, + save_at=save_at, + has_aux=False, + adaptive_loop=adaptive_loop, + ) + + parameter = jnp.asarray(0.7) + bounded = run(parameter, "bounded") + forward = run(parameter, "forward") + assert bool(bounded.ok & forward.ok) + assert bounded.num_accepted == forward.num_accepted + if bounded.accepted is not None: + assert jnp.array_equal(bounded.accepted, forward.accepted) + assert jnp.allclose(bounded.ts, forward.ts, rtol=1e-8, atol=1e-10) + assert jnp.allclose(bounded.xs, forward.xs, rtol=1e-9, atol=1e-11) + + def output(value, adaptive_loop): + return jnp.sum(run(value, adaptive_loop).xs) + + tangent = jnp.ones_like(parameter) + bounded_jvp = jax.jvp( + lambda value: output(value, "bounded"), (parameter,), (tangent,) + )[1] + forward_jvp = jax.jvp( + lambda value: output(value, "forward"), (parameter,), (tangent,) + )[1] + + def first_jvp(value, adaptive_loop): + return jax.jvp( + lambda inner: output(inner, adaptive_loop), (value,), (tangent,) + )[1] + + bounded_second = jax.jvp( + lambda value: first_jvp(value, "bounded"), (parameter,), (tangent,) + )[1] + forward_second = jax.jvp( + lambda value: first_jvp(value, "forward"), (parameter,), (tangent,) + )[1] + assert jnp.allclose(bounded_jvp, forward_jvp, rtol=1e-12, atol=1e-12) + assert jnp.allclose(bounded_second, forward_second, rtol=1e-11, atol=1e-11) + + +def test_forward_adaptive_loop_vmap_matches_bounded_lane_results(): + initial_values = jnp.asarray([0.05, 1.0, 50.0]) + + def batched(adaptive_loop): + def one(initial): + return solve_ode( + lambda x: -0.4 * x, + Tsit5(), + 0.0, + 3.0, + initial, + dt_0=0.2, + controller=IController(rtol=1e-8, atol=1e-10), + max_steps=96, + save_at=SaveAt(steps=True), + has_aux=False, + adaptive_loop=adaptive_loop, + ) + + return jax.vmap(one)(initial_values) + + bounded = batched("bounded") + forward = batched("forward") + assert bool(jnp.all(bounded.ok & forward.ok)) + assert len(set(map(int, bounded.num_accepted))) > 1 + assert jnp.array_equal(bounded.num_accepted, forward.num_accepted) + assert jnp.array_equal(bounded.accepted, forward.accepted) + assert jnp.allclose(bounded.ts, forward.ts, rtol=1e-8, atol=1e-10) + assert jnp.allclose(bounded.xs, forward.xs, rtol=1e-9, atol=1e-11) + + +def test_forward_vmap_matches_bounded_with_mixed_budget_exhaustion(): + rates = jnp.asarray([0.01, 0.1, 1.0, 10.0]) + + def batched(adaptive_loop): + def one(rate): + return solve_ode( + lambda x, t, args, p: -p * x, + Tsit5(), + 0.0, + 2.0, + jnp.asarray(1.0), + p=rate, + dt_0=0.2, + controller=IController(rtol=1e-7, atol=1e-9), + max_steps=4, + save_at=SaveAt(steps=True), + adaptive_loop=adaptive_loop, + ) + + return jax.vmap(one)(rates) + + bounded = batched("bounded") + forward = batched("forward") + assert jnp.array_equal(bounded.ok, jnp.asarray([True, True, False, False])) + assert jnp.array_equal(forward.ok, bounded.ok) + assert jnp.array_equal(forward.num_steps, bounded.num_steps) + assert jnp.array_equal(forward.num_accepted, bounded.num_accepted) + assert jnp.array_equal(forward.accepted, bounded.accepted) + assert jnp.array_equal(forward.ts, bounded.ts) + assert jnp.array_equal(forward.xs, bounded.xs) + + +def test_forward_adaptive_loop_documents_reverse_mode_boundary(): + def endpoint(parameter): + return solve_ode( + lambda x, t, args, p: -p * x, + Tsit5(), + 0.0, + 1.0, + jnp.asarray(1.0), + p=parameter, + dt_0=0.1, + controller=IController(rtol=1e-7, atol=1e-9), + max_steps=64, + has_aux=False, + adaptive_loop="forward", + ).xs + + with pytest.raises(ValueError, match="Reverse-mode differentiation"): + jax.grad(endpoint)(jnp.asarray(0.7)) + + +def test_steps_mesh_is_frozen_and_residual_jacobian_is_exact_at_root(): + theta_star = jnp.asarray(1.5) + controller = IController(rtol=1e-8, atol=1e-10) + + def solve(theta): + return solve_ode( + lambda x, t, args, p: p * x, + Tsit5(), + 0.0, + 2.0, + jnp.asarray([1.0]), + p=theta, + dt_0=0.1, + controller=controller, + max_steps=128, + save_at=SaveAt(steps=True), + ) + + solution = solve(theta_star) + assert bool(solution.ok) + n_valid = int(solution.num_accepted) + 1 + + _, mesh_tangent = jax.jvp( + lambda theta: solve(theta).ts, + (theta_star,), + (jnp.ones_like(theta_star),), + ) + assert jnp.array_equal(mesh_tangent, jnp.zeros_like(mesh_tangent)) + + epsilon = jnp.asarray(1e-5) + plus = solve(theta_star + epsilon) + minus = solve(theta_star - epsilon) + assert int(plus.num_accepted) == int(minus.num_accepted) == n_valid - 1 + finite_difference_mesh = (plus.ts - minus.ts) / (2.0 * epsilon) + assert float(jnp.max(jnp.abs(finite_difference_mesh[:n_valid]))) > 0.1 + + def residual(theta): + candidate = solve(theta) + return ( + (theta - theta_star) + * candidate.xs[:, 0] + * candidate.accepted.astype(theta.dtype) + ) + + _, frozen_jacobian = jax.jvp( + residual, + (theta_star,), + (jnp.ones_like(theta_star),), + ) + finite_difference_jacobian = ( + residual(theta_star + epsilon) - residual(theta_star - epsilon) + ) / (2.0 * epsilon) + np.testing.assert_allclose( + np.asarray(frozen_jacobian[:n_valid]), + np.asarray(finite_difference_jacobian[:n_valid]), + rtol=5e-9, + atol=1e-8, + ) + + +def test_invalid_adaptive_loop_is_rejected(): + with pytest.raises(ValueError, match="adaptive_loop"): + solve_ode( + lambda x: -x, + Tsit5(), + 0.0, + 1.0, + jnp.asarray(1.0), + dt_0=0.1, + controller=IController(), + adaptive_loop="unknown", + ) + + def test_parity_tsit5_free(): # Identical tolerances, budget, and dt_0 must reproduce the kernels # free-stepper's accepted trajectory bit-for-bit (project never binds diff --git a/tests/test_dae.py b/tests/test_dae.py index 716a45a..8337797 100644 --- a/tests/test_dae.py +++ b/tests/test_dae.py @@ -1,12 +1,13 @@ import jax import jax.numpy as jnp import pytest -from nlls_gram import QR, Cholesky +from nlls_gram import LU, QR, Cholesky from tinydiffeq import ( RK4, IController, LMRootSolver, + Rodas5P, SaveAt, Tsit5, solve_ode, @@ -45,7 +46,7 @@ def constraint(y, z, t, args, p): return z - y defaults = _build_algebraic_solver(constraint, LMRootSolver(), False) - assert LMRootSolver().max_steps_is_success + assert not LMRootSolver().max_steps_is_success # Everything algorithmic is nlls-gram's default; only the two invariants # this package owns are pinned. assert isinstance(defaults.linear_solver, Cholesky) @@ -76,6 +77,13 @@ def constraint(y, z, t, args, p): assert not configured.cache_jacobian assert not configured.geodesic_acceleration + direct_ad = _build_algebraic_solver( + constraint, + LMRootSolver(solver_options={"ad_solver": LU()}), + False, + ) + assert isinstance(direct_ad.ad_solver, LU) + def test_lm_root_solver_options_normalize_and_reject_fixed_keys(): # A mapping and the equivalent pairs must compare and hash equal, so @@ -94,9 +102,13 @@ def test_lm_root_solver_options_normalize_and_reject_fixed_keys(): LMRootSolver(solver_options={fixed: True}) with pytest.raises(TypeError, match="mapping or key/value pairs"): LMRootSolver(solver_options=5) + with pytest.raises(TypeError, match="predictor must be a string"): + LMRootSolver(predictor=1) + with pytest.raises(ValueError, match="predictor must be either"): + LMRootSolver(predictor="quadratic") -def test_max_steps_policy_and_strict_batched_derivative(): +def test_max_steps_policy_requires_root_residual_and_strict_batched_derivative(): def one_lane(z_0, p, max_steps_is_success): return solve_semi_explicit_dae( lambda y, z: jnp.zeros_like(y), @@ -117,10 +129,13 @@ def one_lane(z_0, p, max_steps_is_success): ) p = jnp.asarray(1.0) - forgiving = one_lane(jnp.asarray(0.0), p, True) + forgiving = one_lane(jnp.asarray(-1e6), p, True) strict = one_lane(jnp.asarray(0.0), p, False) - assert bool(forgiving.ok) + assert not bool(forgiving.ok) assert not bool(strict.ok) + assert int(forgiving.num_root_solves) == 1 + assert int(forgiving.num_root_steps) == 1 + assert forgiving.zs == jnp.asarray(-1e6) def strict_batch(parameter): return jax.vmap(lambda z_0: one_lane(z_0, parameter, False))( @@ -140,6 +155,56 @@ def endpoint(parameter): assert jnp.array_equal(pullback(jnp.ones_like(value))[0], jnp.asarray(1.0)) +def test_dae_roots_reject_nonresidual_stopping_rules(): + with pytest.raises(ValueError, match="gtol must be zero"): + LMRootSolver(gtol=1e-6) + with pytest.raises(ValueError, match="xtol must be zero"): + LMRootSolver(xtol=1e-6) + + +def test_compat_max_steps_success_does_not_broaden_nlls_root_ad(): + guesses = jnp.asarray([1.0, 0.0]) + + def solve_one(parameter, guess): + return solve_semi_explicit_dae( + lambda y, z: jnp.zeros_like(y), + lambda y, z, t, args, p: z**2 + p, + RK4(), + 0.0, + 0.1, + jnp.asarray(0.0), + guess, + p=parameter, + dt_0=0.1, + max_steps=1, + root_solver=LMRootSolver( + max_steps=1, + max_steps_is_success=True, + atol=1e-12, + ), + ) + + def endpoints(parameters): + return jax.vmap(solve_one)(parameters, guesses).zs + + parameters = jnp.asarray([-1.0, 1.0]) + solutions = jax.vmap(solve_one)(parameters, guesses) + value, tangent = jax.jvp( + endpoints, + (parameters,), + (jnp.ones_like(parameters),), + ) + _, pullback = jax.vjp(endpoints, parameters) + expected_derivative = jnp.asarray([-0.5, 0.0]) + assert jnp.array_equal(solutions.ok, jnp.asarray([True, False])) + assert jnp.array_equal(value, guesses) + assert jnp.array_equal(tangent, expected_derivative) + assert jnp.array_equal(pullback(jnp.ones_like(value))[0], expected_derivative) + assert jnp.array_equal( + jax.grad(lambda p: jnp.sum(endpoints(p)))(parameters), expected_derivative + ) + + def test_rk4_initial_consistency_and_endpoint_accuracy(): sol = solve_linear( jnp.asarray(1.0), @@ -152,10 +217,37 @@ def test_rk4_initial_consistency_and_endpoint_accuracy(): ) assert bool(sol.ok) assert int(sol.num_accepted) == 20 + assert sol.num_root_solves.dtype == jnp.int32 + assert int(sol.num_root_solves) == 1 + 4 * int(sol.num_steps) + assert 0 <= int(sol.num_root_steps) <= 8 * int(sol.num_root_solves) assert jnp.abs(sol.ys - jnp.e) < 2e-5 assert jnp.abs(sol.zs - sol.ys) < 2e-6 +@pytest.mark.parametrize("solver", [RK4(), Rodas5P()]) +def test_fixed_dae_large_negative_initial_time_snaps_to_horizon(solver): + # The float32 nominal fifth time is 0.99975586. Its roundoff is set by + # |t_0|, not |t_1|; the local tolerance should snap it without changing h + # appreciably or depending on the attempt budget. + t_0, t_1, num_steps = -10_000.0, 1.0, 5 + solution = solve_semi_explicit_dae( + lambda y, z: jnp.zeros_like(y), + identity_constraint, + solver, + t_0, + t_1, + jnp.asarray(1.0, jnp.float32), + jnp.asarray(1.0, jnp.float32), + dt_0=(t_1 - t_0) / num_steps, + max_steps=num_steps, + root_solver=LMRootSolver(atol=1e-6), + save_at=SaveAt(steps=True), + ) + assert bool(solution.ok) + assert solution.ts[-1] == jnp.asarray(t_1, jnp.float32) + assert int(solution.num_steps) == int(solution.num_accepted) == num_steps + + def test_rk4_fourth_order_convergence(): def error(n): sol = solve_linear( @@ -184,6 +276,8 @@ def test_adaptive_tsit5_and_steps_padding(): controller=IController(rtol=1e-5, atol=1e-7), ) assert bool(sol.ok) + assert int(sol.num_root_solves) == 1 + 6 * int(sol.num_steps) + assert 0 <= int(sol.num_root_steps) <= 8 * int(sol.num_root_solves) assert sol.ts.shape == (65,) assert int(sol.accepted.sum()) == int(sol.num_accepted) + 1 assert bool(jnp.all(sol.ts[1:] >= sol.ts[:-1])) @@ -191,6 +285,93 @@ def test_adaptive_tsit5_and_steps_padding(): assert jnp.abs(sol.ys[-1] - jnp.exp(2.0)) < 2e-4 +def test_dae_rejects_ode_only_exact_save_mode(): + with pytest.raises(ValueError, match="only supported by solve_ode"): + solve_linear( + jnp.asarray(1.0), + jnp.asarray(0.0), + jnp.asarray(1.0), + RK4(), + SaveAt(ts=jnp.asarray([0.0, 0.15, 1.0]), exact=True), + dt_0=0.1, + max_steps=10, + ) + + +@pytest.mark.parametrize("predictor", ["previous", "secant"]) +def test_forward_adaptive_dae_loop_matches_bounded_and_supports_forward_ad(predictor): + def run(parameter, adaptive_loop): + return solve_linear( + parameter, + jnp.asarray(0.7), + jnp.asarray(1.0), + Tsit5(), + SaveAt(steps=True), + dt_0=0.1, + max_steps=96, + controller=IController(rtol=1e-7, atol=1e-9), + root_solver=LMRootSolver(atol=1e-10, predictor=predictor), + adaptive_loop=adaptive_loop, + ) + + parameter = jnp.asarray(1.3) + bounded = run(parameter, "bounded") + forward = run(parameter, "forward") + assert bool(bounded.ok & forward.ok) + assert bounded.num_accepted == forward.num_accepted + assert bounded.num_steps == forward.num_steps + assert bounded.num_root_solves == forward.num_root_solves + assert bounded.num_root_steps == forward.num_root_steps + assert jnp.array_equal(bounded.accepted, forward.accepted) + assert jnp.allclose(bounded.ts, forward.ts, rtol=2e-7, atol=2e-9) + assert jnp.allclose(bounded.ys, forward.ys, rtol=2e-7, atol=2e-9) + assert jnp.allclose(bounded.zs, forward.zs, rtol=2e-7, atol=2e-9) + + tangent = jnp.ones_like(parameter) + + def endpoint(value, adaptive_loop): + return run(value, adaptive_loop).ys[-1] + + bounded_jvp = jax.jvp( + lambda value: endpoint(value, "bounded"), (parameter,), (tangent,) + )[1] + forward_jvp = jax.jvp( + lambda value: endpoint(value, "forward"), (parameter,), (tangent,) + )[1] + + def first_jvp(value, adaptive_loop): + return jax.jvp( + lambda inner: endpoint(inner, adaptive_loop), (value,), (tangent,) + )[1] + + bounded_second = jax.jvp( + lambda value: first_jvp(value, "bounded"), (parameter,), (tangent,) + )[1] + forward_second = jax.jvp( + lambda value: first_jvp(value, "forward"), (parameter,), (tangent,) + )[1] + assert jnp.allclose(bounded_jvp, forward_jvp, rtol=2e-6, atol=2e-8) + assert jnp.allclose(bounded_second, forward_second, rtol=2e-5, atol=2e-7) + + +def test_forward_adaptive_dae_loop_documents_reverse_boundary(): + def endpoint(parameter): + return solve_linear( + parameter, + jnp.asarray(0.7), + jnp.asarray(1.0), + Tsit5(), + SaveAt(t_1=True), + dt_0=0.1, + max_steps=64, + controller=IController(rtol=1e-6, atol=1e-8), + adaptive_loop="forward", + ).ys + + with pytest.raises(ValueError, match="Reverse-mode differentiation"): + jax.grad(endpoint)(jnp.asarray(1.3)) + + def test_interpolated_z_has_small_constraint_defect(): def f(y, z): return z @@ -349,7 +530,7 @@ def from_y_0(y_0): def test_jit_and_vmap(): @jax.jit def endpoint(p, y_0): - return solve_linear( + sol = solve_linear( p, y_0, y_0, @@ -358,12 +539,140 @@ def endpoint(p, y_0): dt_0=0.1, max_steps=64, controller=IController(rtol=1e-5, atol=1e-7), - ).ys + ) + return sol.ys, sol.num_root_solves, sol.num_root_steps ps = jnp.asarray([0.5, 1.0, 1.5]) y_0s = jnp.asarray([0.7, 1.0, 1.3]) - got = jax.vmap(endpoint)(ps, y_0s) + got, root_solves, root_steps = jax.vmap(endpoint)(ps, y_0s) assert jnp.max(jnp.abs(got - y_0s * jnp.exp(ps))) < 2e-4 + assert root_solves.shape == root_steps.shape == ps.shape + assert root_solves.dtype == root_steps.dtype == jnp.int32 + scalar_stats = jnp.stack( + [jnp.stack(endpoint(p, y_0)[1:]) for p, y_0 in zip(ps, y_0s, strict=True)] + ) + assert jnp.array_equal(root_solves, scalar_stats[:, 0]) + assert jnp.array_equal(root_steps, scalar_stats[:, 1]) + + +def test_root_statistics_are_ad_inert_without_changing_state_jvp(): + def full_solution(parameter): + return solve_linear( + parameter, + jnp.asarray(1.0), + jnp.asarray(1.0), + RK4(), + SaveAt(t_1=True), + dt_0=0.05, + max_steps=20, + ) + + parameter = jnp.asarray(0.4) + _, tangent = jax.jvp(full_solution, (parameter,), (jnp.ones_like(parameter),)) + assert jnp.abs(tangent.ys - jnp.exp(parameter)) < 2e-5 + assert tangent.num_root_solves.dtype == jax.dtypes.float0 + assert tangent.num_root_steps.dtype == jax.dtypes.float0 + + +@pytest.mark.parametrize(("solver", "roots_per_step"), [(RK4(), 4), (Tsit5(), 6)]) +def test_secant_predictor_preserves_explicit_solution_and_root_call_count( + solver, roots_per_step +): + def run(predictor): + return solve_linear( + jnp.asarray(0.4), + jnp.asarray(1.0), + jnp.asarray(1.0), + solver, + SaveAt(t_1=True), + dt_0=0.1, + max_steps=10, + root_solver=LMRootSolver(atol=1e-10, predictor=predictor), + ) + + previous = run("previous") + secant = run("secant") + assert bool(previous.ok & secant.ok) + assert int(previous.num_root_solves) == int(secant.num_root_solves) + assert int(secant.num_root_solves) == 1 + roots_per_step * 10 + assert int(secant.num_root_steps) <= int(previous.num_root_steps) + assert jnp.allclose(secant.ys, previous.ys, rtol=1e-9, atol=1e-10) + assert jnp.allclose(secant.zs, previous.zs, rtol=1e-9, atol=1e-10) + + +def test_secant_predictor_preserves_mixed_differential_algebraic_dtypes(): + y_0 = jnp.asarray(1.0, dtype=jnp.float64) + z_0 = jnp.asarray(1.0, dtype=jnp.float32) + + sol = solve_semi_explicit_dae( + lambda y, z: z.astype(y.dtype), + lambda y, z: z - y.astype(z.dtype), + RK4(), + 0.0, + 0.2, + y_0, + z_0, + dt_0=0.1, + max_steps=2, + root_solver=LMRootSolver(atol=1e-6, predictor="secant"), + save_at=SaveAt(steps=True), + ) + + assert bool(sol.ok) + assert sol.ys.dtype == jnp.float64 + assert sol.zs.dtype == jnp.float32 + + +@pytest.mark.parametrize("predictor", ["previous", "secant"]) +def test_explicit_root_predictor_preserves_implicit_jvp_and_vjp(predictor): + def endpoint(parameter): + return solve_linear( + parameter, + jnp.asarray(1.0), + jnp.asarray(1.0), + Tsit5(), + SaveAt(t_1=True), + dt_0=0.1, + max_steps=10, + root_solver=LMRootSolver(atol=1e-10, predictor=predictor), + ).ys + + parameter = jnp.asarray(0.4) + expected = jnp.exp(parameter) + tangent = jax.jvp(endpoint, (parameter,), (jnp.ones_like(parameter),))[1] + cotangent = jax.grad(endpoint)(parameter) + assert jnp.allclose(tangent, expected, rtol=2e-8, atol=2e-10) + assert jnp.allclose(cotangent, expected, rtol=2e-8, atol=2e-10) + + +def test_secant_predictor_continues_a_locally_unique_positive_root(): + def run(predictor): + return solve_semi_explicit_dae( + lambda y, z: -0.2 * z, + lambda y, z: z**2 - y - 2.0, + Tsit5(), + 0.0, + 1.0, + jnp.asarray(1.0), + jnp.sqrt(jnp.asarray(3.0)), + dt_0=0.1, + controller=IController(rtol=1e-8, atol=1e-10), + root_solver=LMRootSolver(atol=1e-10, predictor=predictor), + max_steps=64, + save_at=SaveAt(steps=True), + ) + + previous = run("previous") + secant = run("secant") + assert bool(previous.ok & secant.ok) + assert jnp.all(secant.zs[secant.accepted] > 0.0) + assert int(secant.num_root_steps) <= int(previous.num_root_steps) + assert jnp.allclose( + secant.ys[secant.accepted], + previous.ys[previous.accepted], + rtol=2e-9, + atol=2e-10, + ) def test_kernels_optimal_advertising_system_matches_elimination(): @@ -498,6 +807,8 @@ def test_initial_root_failure_and_time_budget_failure(): ) assert not bool(failed_root.ok) assert int(failed_root.num_accepted) == 0 + assert int(failed_root.num_root_solves) == 1 + assert 0 <= int(failed_root.num_root_steps) <= 8 starved = solve_linear( jnp.asarray(1.0), @@ -510,10 +821,12 @@ def test_initial_root_failure_and_time_budget_failure(): ) assert not bool(starved.ok) assert int(starved.num_accepted) == 3 + assert int(starved.num_root_solves) == 1 + 4 * int(starved.num_steps) assert jnp.abs(starved.ts - 0.3) < 1e-7 -def test_adaptive_stage_root_failure_retries_with_smaller_step(): +@pytest.mark.parametrize("predictor", ["previous", "secant"]) +def test_adaptive_stage_root_failure_retries_with_smaller_step(predictor): # One damped LM iteration cannot meet the root tolerance at dt_0, but it # can after the adaptive controller reduces the step. A fixed controller # would terminate on the same stage-root failure. @@ -522,19 +835,25 @@ def test_adaptive_stage_root_failure_retries_with_smaller_step(): identity_constraint, Tsit5(), 0.0, - 0.005, + 0.0001, jnp.asarray(0.0), jnp.asarray(0.0), - dt_0=0.005, + dt_0=0.0001, controller=IController(), - root_solver=LMRootSolver(max_steps=1, max_steps_is_success=False, atol=1e-6), + root_solver=LMRootSolver( + max_steps=1, + max_steps_is_success=False, + atol=1e-8, + predictor=predictor, + ), max_steps=64, save_at=SaveAt(steps=True), ) assert bool(sol.ok) assert int(sol.num_accepted) > 1 - assert jnp.abs(sol.ys[-1] - 0.005) < 1e-7 - assert jnp.abs(sol.zs[-1] - sol.ys[-1]) < 1e-6 + assert int(sol.num_steps) > int(sol.num_accepted) + assert jnp.abs(sol.ys[-1] - 0.0001) < 1e-9 + assert jnp.abs(sol.zs[-1] - sol.ys[-1]) < 1e-8 def test_masked_failed_lane_has_safe_implicit_root_jvp_and_vjp(): @@ -628,11 +947,13 @@ def test_validation(): ) with pytest.raises(ValueError, match="positive int"): LMRootSolver(max_steps=0) + with pytest.raises(ValueError, match="atol must be positive or None"): + LMRootSolver(atol=0.0) with pytest.raises(TypeError, match="max_steps_is_success must be a bool"): LMRootSolver(max_steps_is_success=1) - with pytest.raises(ValueError, match="gtol must be nonnegative"): + with pytest.raises(ValueError, match="gtol must be zero"): LMRootSolver(gtol=-1.0) - with pytest.raises(ValueError, match="xtol must be nonnegative"): + with pytest.raises(ValueError, match="xtol must be zero"): LMRootSolver(xtol=-1.0) with pytest.raises(ValueError, match="2 to 5 positional"): solve_semi_explicit_dae( diff --git a/tests/test_exponential.py b/tests/test_exponential.py index cb87e11..e22cc63 100644 --- a/tests/test_exponential.py +++ b/tests/test_exponential.py @@ -167,6 +167,15 @@ def test_linear_exponential_input_validation(): solve_linear_ode( jnp.eye(2), DenseExponential(), 0.0, 1.0, x_0, save_at=SaveAt(steps=True) ) + with pytest.raises(ValueError, match="only supported by solve_ode"): + solve_linear_ode( + jnp.eye(2), + DenseExponential(), + 0.0, + 1.0, + x_0, + save_at=SaveAt(ts=jnp.asarray([0.0, 1.0]), exact=True), + ) outside = solve_linear_ode( jnp.eye(2), DenseExponential(), @@ -346,6 +355,7 @@ def test_adaptive_krylov_rejects_then_meets_endpoint_tolerance(dtype): expected = jnp.exp(2 * eigenvalues) assert bool(solution.ok) assert 1 < int(solution.num_accepted) < method.max_steps + assert int(solution.num_steps) > int(solution.num_accepted) assert jnp.linalg.norm(solution.xs - expected) <= 2 * tolerance @@ -366,6 +376,7 @@ def test_adaptive_krylov_attempt_budget_failure_is_fast_and_finite(): ) assert not bool(solution.ok) assert int(solution.num_accepted) == 0 + assert int(solution.num_steps) == 1 assert jnp.all(jnp.isfinite(solution.xs)) assert jnp.array_equal(solution.xs, initial) @@ -432,6 +443,125 @@ def endpoint(vector): assert jnp.allclose(traced_cotangent, exponential.T @ flatten(cotangent), atol=3e-5) +@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64]) +@pytest.mark.parametrize("adaptive", [False, True]) +def test_full_space_krylov_happy_breakdowns_have_dense_ad(dtype, adaptive): + tolerance = 3e-5 if dtype == jnp.float32 else 3e-11 + cases = ( + # A generic two-dimensional Krylov space breaks down only after its + # final useful Arnoldi vector. This was the float32 reverse-mode NaN. + ( + jnp.asarray([[-0.4, 0.2], [0.1, -0.3]], dtype), + jnp.asarray([1.0, -0.2], dtype), + jnp.asarray([0.3, 0.4], dtype), + jnp.asarray([-0.2, 0.7], dtype), + ), + # An eigenvector initial state breaks down after the first vector, + # before the full three-dimensional basis has been constructed. + ( + jnp.diag(jnp.asarray([-0.4, -0.3, -0.6], dtype)), + jnp.asarray([1.0, 0.0, 0.0], dtype), + jnp.asarray([0.3, 0.4, -0.2], dtype), + jnp.asarray([-0.2, 0.7, 0.5], dtype), + ), + ) + + for matrix, initial, tangent, cotangent in cases: + krylov_dim = initial.size + method = ( + AdaptiveKrylovExponential(krylov_dim=krylov_dim, max_steps=8) + if adaptive + else KrylovExponential(krylov_dim=krylov_dim) + ) + + def endpoint(current_initial, matrix=matrix, method=method): + return solve_linear_ode(matrix, method, 0.0, 0.2, current_initial).xs + + def dense_endpoint(current_initial, matrix=matrix): + return solve_linear_ode( + matrix, DenseExponential(), 0.0, 0.2, current_initial + ).xs + + value, output_tangent = jax.jvp(endpoint, (initial,), (tangent,)) + dense_value, dense_tangent = jax.jvp(dense_endpoint, (initial,), (tangent,)) + _, pullback = jax.vjp(endpoint, initial) + _, dense_pullback = jax.vjp(dense_endpoint, initial) + input_cotangent = pullback(cotangent)[0] + dense_cotangent = dense_pullback(cotangent)[0] + batch = jnp.stack([initial, 0.7 * initial]) + batched = jax.jit(jax.vmap(endpoint))(batch) + dense_batched = jax.jit(jax.vmap(dense_endpoint))(batch) + + for actual in (value, output_tangent, input_cotangent, batched): + assert actual.dtype == dtype + assert bool(jnp.all(jnp.isfinite(actual))) + assert jnp.allclose(value, dense_value, rtol=tolerance, atol=tolerance) + assert jnp.allclose( + output_tangent, dense_tangent, rtol=tolerance, atol=tolerance + ) + assert jnp.allclose( + input_cotangent, dense_cotangent, rtol=tolerance, atol=tolerance + ) + assert jnp.allclose(batched, dense_batched, rtol=tolerance, atol=tolerance) + + +@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64]) +@pytest.mark.parametrize("adaptive", [False, True]) +def test_truncated_krylov_final_breakdown_has_finite_ad(dtype, adaptive): + # The first two coordinates form an invariant subspace, so a two-vector + # Arnoldi run in a three-dimensional state breaks down on its final step. + # Keeping tangents and cotangents in that invariant subspace gives an exact + # dense comparison while directly exercising the matrix-free norm guard. + matrix = jnp.asarray([[-0.4, 0.2, 0.0], [0.1, -0.3, 0.0], [0.0, 0.0, -0.7]], dtype) + initial = jnp.asarray([1.0, -0.2, 0.0], dtype) + tangent = jnp.asarray([0.3, 0.4, 0.0], dtype) + cotangent = jnp.asarray([-0.2, 0.7, 0.0], dtype) + method = ( + AdaptiveKrylovExponential(krylov_dim=2, max_steps=8) + if adaptive + else KrylovExponential(krylov_dim=2) + ) + + def endpoint(current_initial): + return solve_linear_ode(matrix, method, 0.0, 0.2, current_initial).xs + + def dense_endpoint(current_initial): + return solve_linear_ode( + matrix, DenseExponential(), 0.0, 0.2, current_initial + ).xs + + value, output_tangent = jax.jvp(endpoint, (initial,), (tangent,)) + dense_value, dense_tangent = jax.jvp(dense_endpoint, (initial,), (tangent,)) + _, pullback = jax.vjp(endpoint, initial) + _, dense_pullback = jax.vjp(dense_endpoint, initial) + input_cotangent = pullback(cotangent)[0] + dense_cotangent = dense_pullback(cotangent)[0] + tolerance = 4e-5 if dtype == jnp.float32 else 3e-11 + + assert bool(jnp.all(jnp.isfinite(output_tangent))) + assert bool(jnp.all(jnp.isfinite(input_cotangent))) + assert jnp.allclose(value, dense_value, rtol=tolerance, atol=tolerance) + assert jnp.allclose(output_tangent, dense_tangent, rtol=tolerance, atol=tolerance) + assert jnp.allclose( + input_cotangent, dense_cotangent, rtol=tolerance, atol=tolerance + ) + + +def test_adaptive_krylov_float32_jaxpr_contains_no_float64(): + # conftest enables x64, so this catches Python-float fallbacks that would + # otherwise be silently float32 when the application leaves x64 disabled. + matrix = jnp.asarray( + [[-0.4, 0.2, 0.0], [0.1, -0.3, 0.1], [0.0, 0.2, -0.5]], + jnp.float32, + ) + initial = jnp.asarray([1.0, -0.2, 0.4], jnp.float32) + method = AdaptiveKrylovExponential(krylov_dim=2, max_steps=8) + jaxpr = jax.make_jaxpr( + lambda state: solve_linear_ode(matrix, method, 0.0, 0.2, state).xs + )(initial) + assert "f64" not in str(jaxpr), jaxpr + + def test_adaptive_krylov_constructor_validation(): with pytest.raises(ValueError, match="max_steps"): AdaptiveKrylovExponential(max_steps=0) diff --git a/tests/test_float64_subprocess.py b/tests/test_float64_subprocess.py index 98dd564..186cdac 100644 --- a/tests/test_float64_subprocess.py +++ b/tests/test_float64_subprocess.py @@ -244,6 +244,193 @@ def endpoint(x_0): """) +def test_default_x64_disabled_representative_paths_stay_float32(): + run_script(r""" +import jax +import jax.numpy as jnp + +from tinydiffeq import ( + ContinuousTimeMarkovChain, + DenseExponential, + DiscreteMarkovChain, + EulerMaruyama, + IController, + SaveAt, + Tsit5, + forecast_continuous_time_markov_chain, + forecast_markov_chain, + simulate_continuous_time_markov_chain, + simulate_markov_chain, + solve_linear_ode, + solve_ode, + solve_sde, + solve_semi_explicit_dae, + solve_semi_explicit_sdae, +) + + +assert not jax.config.x64_enabled +x_0 = jnp.asarray(1.0) +parameter = jnp.asarray(0.2) + + +def check_value_jvp_vjp(function, value): + jaxpr = jax.make_jaxpr(function)(value) + assert "f64" not in str(jaxpr), jaxpr + primal, tangent = jax.jvp( + function, + (value,), + (jax.tree.map(jnp.ones_like, value),), + ) + _, pullback = jax.vjp(function, value) + cotangent = pullback(jax.tree.map(jnp.ones_like, primal))[0] + for leaf in jax.tree.leaves((primal, tangent, cotangent)): + if jnp.issubdtype(leaf.dtype, jnp.inexact): + assert leaf.dtype == jnp.float32, leaf.dtype + assert jnp.all(jnp.isfinite(leaf)) + + +def ode_endpoint(rate): + return solve_ode( + lambda x, t, args, p: p * x, + Tsit5(), + 0.0, + 0.5, + x_0, + p=rate, + dt_0=0.1, + controller=IController(), + max_steps=32, + ).xs + + +def dae_endpoint(rate): + return solve_semi_explicit_dae( + lambda y, z, t, args, p: p * z, + lambda y, z: z - y, + Tsit5(), + 0.0, + 0.5, + x_0, + x_0, + p=rate, + dt_0=0.1, + controller=IController(), + max_steps=32, + ).ys + + +def sde_endpoint(rate): + return solve_sde( + lambda x, t, args, p: p * x, + lambda x, t, args, p: jnp.asarray(0.1, x.dtype) * x, + EulerMaruyama(), + 0.0, + 0.5, + x_0, + p=rate, + key=jax.random.key(1), + n_steps=8, + ).xs + + +def sdae_endpoint(rate): + return solve_semi_explicit_sdae( + lambda y, z, t, args, p: p * z, + lambda y, z, t, args, p: jnp.asarray(0.1, y.dtype) * z, + lambda y, z: z - y, + EulerMaruyama(), + 0.0, + 0.5, + x_0, + x_0, + p=rate, + key=jax.random.key(2), + n_steps=8, + ).ys + + +def exponential_endpoint(rate): + return solve_linear_ode( + lambda x: rate * x, + DenseExponential(), + 0.0, + 0.5, + x_0, + ).xs + + +for endpoint in ( + ode_endpoint, + dae_endpoint, + sde_endpoint, + sdae_endpoint, + exponential_endpoint, +): + check_value_jvp_vjp(endpoint, parameter) + + +discrete = DiscreteMarkovChain(jnp.asarray([[0.8, 0.2], [0.3, 0.7]])) +continuous = ContinuousTimeMarkovChain(jnp.asarray([[-1.0, 1.0], [0.5, -0.5]])) +keys = jax.random.split(jax.random.key(3), 3) + + +def discrete_path(key): + return simulate_markov_chain( + discrete, + jnp.int32(0), + key=key, + num_steps=8, + save_at=SaveAt(steps=True), + ) + + +def continuous_path(key): + return simulate_continuous_time_markov_chain( + continuous, + 0.0, + 2.0, + jnp.int32(0), + key=key, + max_jumps=32, + save_at=SaveAt(steps=True), + ) + + +for path in (discrete_path, continuous_path): + jaxpr = jax.make_jaxpr(jax.vmap(path))(keys) + assert "f64" not in str(jaxpr), jaxpr +discrete_paths = jax.jit(jax.vmap(discrete_path))(keys) +continuous_paths = jax.jit(jax.vmap(continuous_path))(keys) +assert discrete.transition_matrix.dtype == jnp.float32 +assert continuous.generator.dtype == jnp.float32 +assert continuous_paths.ts.dtype == jnp.float32 +assert bool(jnp.all(continuous_paths.ok)) + + +distribution = jnp.asarray([0.4, 0.6]) + + +def discrete_forecast(value): + return forecast_markov_chain( + discrete, value, num_steps=4 + ).probabilities + + +def continuous_forecast(value): + return forecast_continuous_time_markov_chain( + continuous, 0.0, 1.0, value + ).probabilities + + +for forecast in (discrete_forecast, continuous_forecast): + check_value_jvp_vjp(forecast, distribution) + probabilities = jax.jit(forecast)(distribution) + assert probabilities.dtype == jnp.float32 + assert jnp.allclose(jnp.sum(probabilities), 1.0, atol=1e-5) +""") + + def test_pytree_states_preserve_float32_and_float64(): run_script(r""" import jax @@ -393,6 +580,137 @@ def mixed_solve(q): """) +def test_float32_dae_and_sdae_lowerings_contain_no_float64_under_x64(): + run_script(r""" +import os +os.environ["JAX_PLATFORMS"] = "cpu" +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +from tinydiffeq import ( + EulerMaruyama, + IController, + LMRootSolver, + RK4, + Tsit5, + solve_semi_explicit_dae, + solve_semi_explicit_sdae, +) + + +DTYPE = jnp.float32 +T_0 = jnp.asarray(0.0, DTYPE) +T_1 = jnp.asarray(0.2, DTYPE) +DT_0 = jnp.asarray(0.1, DTYPE) +Y_0 = jnp.asarray(1.0, DTYPE) +Z_0 = jnp.asarray(1.0, DTYPE) +P = jnp.asarray(0.2, DTYPE) + + +def differential(y, z, t, args, p): + return p * z + + +def constraint(y, z, t, args, p): + return z - y + + +def dae_endpoint(p, solver, root_solver, adaptive_loop="bounded"): + adaptive_options = {} + if isinstance(solver, Tsit5): + adaptive_options = { + "controller": IController(), + "adaptive_loop": adaptive_loop, + } + return solve_semi_explicit_dae( + differential, + constraint, + solver, + T_0, + T_1, + Y_0, + Z_0, + p=p, + dt_0=DT_0, + max_steps=4, + root_solver=root_solver, + **adaptive_options, + ).ys + + +def sdae_endpoint(p): + return solve_semi_explicit_sdae( + differential, + lambda y, z, t, args, p: jnp.zeros_like(y), + constraint, + EulerMaruyama(), + T_0, + T_1, + Y_0, + Z_0, + p=p, + key=jax.random.key(0), + n_steps=2, + root_solver=LMRootSolver(), + ).ys + + +def with_jvp(function): + return lambda p: jax.jvp(function, (p,), (jnp.ones_like(p),)) + + +def with_vjp(function): + def transformed(p): + value, pullback = jax.vjp(function, p) + return value, pullback(jnp.ones_like(value))[0] + + return transformed + + +def assert_pure_float32(name, function): + value = function(P) + for leaf in jax.tree.leaves(value): + if jnp.issubdtype(leaf.dtype, jnp.inexact): + assert leaf.dtype == DTYPE, (name, leaf.dtype) + assert jnp.all(jnp.isfinite(leaf)), (name, leaf) + + jaxpr = str(jax.make_jaxpr(function)(P)) + assert "f64" not in jaxpr, (name, jaxpr) + stablehlo = str(jax.jit(function).lower(P).compiler_ir("stablehlo")) + offending = [line for line in stablehlo.splitlines() if "f64" in line] + assert not offending, (name, offending[:10]) + + +default_root = LMRootSolver() +configured_root = LMRootSolver( + solver_options={ + "init_damping": 2e-3, + "damping_decrease": 0.4, + "damping_increase": 3.0, + } +) +dae_cases = { + "rk4_primal": lambda p: dae_endpoint(p, RK4(), default_root), + "tsit5_bounded_primal": lambda p: dae_endpoint( + p, Tsit5(), configured_root, "bounded" + ), + "tsit5_forward_primal": lambda p: dae_endpoint( + p, Tsit5(), configured_root, "forward" + ), +} +for name, function in dae_cases.items(): + assert_pure_float32(name, function) + assert_pure_float32(name.replace("primal", "jvp"), with_jvp(function)) + if "forward" not in name: + assert_pure_float32(name.replace("primal", "vjp"), with_vjp(function)) + +assert_pure_float32("sdae_primal", sdae_endpoint) +assert_pure_float32("sdae_jvp", with_jvp(sdae_endpoint)) +assert_pure_float32("sdae_vjp", with_vjp(sdae_endpoint)) +""") + + def test_aux_dense_output_and_sdae_dtype_contracts(): run_script(r""" import os diff --git a/tests/test_markov.py b/tests/test_markov.py index e19ef36..e745bb1 100644 --- a/tests/test_markov.py +++ b/tests/test_markov.py @@ -9,6 +9,8 @@ DiscreteMarkovChain, SaveAt, SequentialMarkov, + forecast_continuous_time_markov_chain, + forecast_markov_chain, simulate_continuous_time_markov_chain, simulate_markov_chain, ) @@ -36,6 +38,7 @@ def test_discrete_deterministic_path_and_save_modes(): assert endpoint.xs == 0 assert jnp.array_equal(selected.xs, jnp.asarray([0, 1, 0])) assert bool(steps.ok) and bool(jnp.all(steps.accepted)) + assert int(steps.num_steps) == int(endpoint.num_steps) == 6 passed_as_pytree = jax.jit( lambda prepared, random_key: ( @@ -71,8 +74,53 @@ def run(key, method): assert jnp.array_equal(sequential, associative) +@pytest.mark.parametrize("method", [SequentialMarkov(), AssociativeMarkov()]) +def test_float32_simulation_jaxprs_are_pure_under_x64(method): + discrete = DiscreteMarkovChain(jnp.asarray([[0.8, 0.2], [0.3, 0.7]], jnp.float32)) + continuous = ContinuousTimeMarkovChain( + jnp.asarray([[-1.0, 1.0], [0.5, -0.5]], jnp.float32) + ) + keys = jax.random.split(jax.random.key(91), 4) + + def discrete_path(key): + return simulate_markov_chain( + discrete, + jnp.int32(0), + key=key, + num_steps=16, + method=method, + save_at=SaveAt(steps=True), + ) + + def continuous_path(key): + return simulate_continuous_time_markov_chain( + continuous, + jnp.asarray(0.0, jnp.float32), + jnp.asarray(2.0, jnp.float32), + jnp.int32(0), + key=key, + max_jumps=32, + method=method, + save_at=SaveAt(steps=True), + ) + + discrete_jaxpr = jax.make_jaxpr(jax.vmap(discrete_path))(keys) + continuous_jaxpr = jax.make_jaxpr(jax.vmap(continuous_path))(keys) + assert "f64" not in str(discrete_jaxpr), discrete_jaxpr + assert "f64" not in str(continuous_jaxpr), continuous_jaxpr + + discrete_paths = jax.jit(jax.vmap(discrete_path))(keys) + continuous_paths = jax.jit(jax.vmap(continuous_path))(keys) + assert discrete_paths.xs.shape == (4, 17) + assert continuous_paths.ts.dtype == jnp.float32 + assert bool(jnp.all(continuous_paths.ok)) + + def test_discrete_one_step_distribution(): - transition = jnp.asarray([[0.1, 0.3, 0.6], [0.2, 0.5, 0.3], [0.7, 0.2, 0.1]]) + transition = jnp.asarray( + [[0.1, 0.3, 0.6], [0.2, 0.5, 0.3], [0.7, 0.2, 0.1]], + jnp.float32, + ) chain = DiscreteMarkovChain(transition) keys = jax.random.split(jax.random.key(2), 30_000) samples = jax.jit( @@ -139,6 +187,7 @@ def test_ctmc_methods_match_states_and_event_times(dtype): tolerance = 2e-4 if dtype == jnp.float32 else 1e-11 assert jnp.allclose(sequential.ts, associative.ts, rtol=tolerance, atol=tolerance) assert sequential.num_accepted == associative.num_accepted + assert sequential.num_steps == associative.num_steps assert bool(sequential.ok) and bool(associative.ok) @@ -193,15 +242,21 @@ def test_ctmc_absorbing_and_starved_contracts(): ) assert bool(absorbed.ok) assert int(absorbed.num_accepted) == 0 + assert int(absorbed.num_steps) == 1 assert jnp.all(absorbed.xs == 0) assert not bool(starved.ok) + assert int(starved.num_steps) == 1 assert starved.ts < 1_000.0 @pytest.mark.parametrize("method", [SequentialMarkov(), AssociativeMarkov()]) def test_two_state_ctmc_endpoint_distribution(method): - rate_01, rate_10, horizon = 2.0, 1.0, 0.8 - chain = ContinuousTimeMarkovChain([[-rate_01, rate_01], [rate_10, -rate_10]]) + rate_01 = jnp.asarray(2.0, jnp.float32) + rate_10 = jnp.asarray(1.0, jnp.float32) + horizon = jnp.asarray(0.8, jnp.float32) + chain = ContinuousTimeMarkovChain( + jnp.asarray([[-rate_01, rate_01], [rate_10, -rate_10]], jnp.float32) + ) keys = jax.random.split(jax.random.key(7), 20_000) solutions = jax.jit( jax.vmap( @@ -257,3 +312,45 @@ def test_markov_static_argument_validation(): num_steps=2, save_at=SaveAt(steps=True, fill="inf"), ) + + +def test_markov_save_at_exact_is_rejected_by_every_public_query_api(): + discrete = DiscreteMarkovChain([[0.75, 0.25], [0.4, 0.6]]) + continuous = ContinuousTimeMarkovChain([[-0.5, 0.5], [0.25, -0.25]]) + discrete_queries = SaveAt(ts=jnp.asarray([0, 2], jnp.int32), exact=True) + continuous_queries = SaveAt(ts=jnp.asarray([0.0, 1.0]), exact=True) + match = "only supported by solve_ode" + + with pytest.raises(ValueError, match=match): + simulate_markov_chain( + discrete, + jnp.int32(0), + key=jax.random.key(0), + num_steps=2, + save_at=discrete_queries, + ) + with pytest.raises(ValueError, match=match): + simulate_continuous_time_markov_chain( + continuous, + 0.0, + 1.0, + jnp.int32(0), + key=jax.random.key(1), + max_jumps=8, + save_at=continuous_queries, + ) + with pytest.raises(ValueError, match=match): + forecast_markov_chain( + discrete, + jnp.asarray([1.0, 0.0]), + num_steps=2, + save_at=discrete_queries, + ) + with pytest.raises(ValueError, match=match): + forecast_continuous_time_markov_chain( + continuous, + 0.0, + 1.0, + jnp.asarray([1.0, 0.0]), + save_at=continuous_queries, + ) diff --git a/tests/test_markov_distribution.py b/tests/test_markov_distribution.py index 2d76822..a039295 100644 --- a/tests/test_markov_distribution.py +++ b/tests/test_markov_distribution.py @@ -112,6 +112,48 @@ def endpoint(distribution): assert jnp.allclose(batched, batch @ transition_power) +def test_float32_forecast_transforms_are_pure_under_x64(): + transition = jnp.asarray([[0.9, 0.1], [0.25, 0.75]], jnp.float32) + generator = jnp.asarray([[-2.0, 2.0], [1.0, -1.0]], jnp.float32) + discrete = DiscreteMarkovChain(transition) + continuous = ContinuousTimeMarkovChain(generator) + distribution = jnp.asarray([0.4, 0.6], jnp.float32) + tangent = jnp.asarray([0.2, -0.2], jnp.float32) + cotangent = jnp.asarray([0.7, -0.3], jnp.float32) + batch = jnp.asarray([[1.0, 0.0], [0.0, 1.0], [0.4, 0.6]], jnp.float32) + + def discrete_endpoint(value): + return forecast_markov_chain(discrete, value, num_steps=12).probabilities + + def continuous_endpoint(value): + return forecast_continuous_time_markov_chain( + continuous, + jnp.asarray(0.0, jnp.float32), + jnp.asarray(1.7, jnp.float32), + value, + ).probabilities + + def transformed(value, direction, weights, values): + discrete_value, discrete_tangent = jax.jvp( + discrete_endpoint, (value,), (direction,) + ) + continuous_value, continuous_pullback = jax.vjp(continuous_endpoint, value) + return ( + discrete_value, + discrete_tangent, + continuous_value, + continuous_pullback(weights)[0], + jax.vmap(discrete_endpoint)(values), + jax.vmap(continuous_endpoint)(values), + ) + + jaxpr = jax.make_jaxpr(transformed)(distribution, tangent, cotangent, batch) + assert "f64" not in str(jaxpr), jaxpr + result = jax.jit(transformed)(distribution, tangent, cotangent, batch) + assert all(leaf.dtype == jnp.float32 for leaf in jax.tree.leaves(result)) + assert all(bool(jnp.all(jnp.isfinite(leaf))) for leaf in jax.tree.leaves(result)) + + def test_invalid_initial_distribution_reports_not_ok_with_finite_output(): chain = DiscreteMarkovChain([[0.8, 0.2], [0.3, 0.7]]) for distribution in ( diff --git a/tests/test_rodas5p.py b/tests/test_rodas5p.py index ab0cab6..16e3ffa 100644 --- a/tests/test_rodas5p.py +++ b/tests/test_rodas5p.py @@ -156,6 +156,9 @@ def test_nonlinear_dae_adaptive_constraint_accuracy_and_initial_root_only(): save_at=SaveAt(steps=True), ) assert bool(sol.ok) + assert int(sol.num_steps) > 1 + assert int(sol.num_root_solves) == 1 + assert 0 <= int(sol.num_root_steps) <= 1 valid = sol.accepted assert jnp.max(jnp.abs(sol.zs[valid] ** 2 - sol.ys[valid] - 2.0)) < 2e-8 @@ -175,6 +178,78 @@ def from_guess(guess): assert jnp.abs(jax.grad(from_guess)(jnp.sqrt(jnp.asarray(3.0)))) < 1e-14 +def test_rodas5p_dae_forward_loop_matches_bounded_primal_and_jvp(): + def run(rate, adaptive_loop): + return solve_semi_explicit_dae( + lambda y, z, t, args, p: p * z, + lambda y, z: z**2 - y - 2.0, + Rodas5P(), + 0.0, + 1.0, + jnp.asarray(1.0), + jnp.sqrt(jnp.asarray(3.0)), + p=rate, + dt_0=0.2, + controller=IController(rtol=1e-8, atol=1e-10), + max_steps=64, + save_at=SaveAt(steps=True), + adaptive_loop=adaptive_loop, + ) + + rate = jnp.asarray(-0.2) + bounded = run(rate, "bounded") + forward = run(rate, "forward") + assert bool(bounded.ok & forward.ok) + assert bounded.num_accepted == forward.num_accepted + assert bounded.num_steps == forward.num_steps + assert bounded.num_root_solves == forward.num_root_solves == 1 + assert bounded.num_root_steps == forward.num_root_steps + assert jnp.array_equal(bounded.accepted, forward.accepted) + assert jnp.allclose(bounded.ts, forward.ts, rtol=1e-8, atol=1e-10) + assert jnp.allclose(bounded.ys, forward.ys, rtol=1e-8, atol=1e-10) + assert jnp.allclose(bounded.zs, forward.zs, rtol=1e-8, atol=1e-10) + + bounded_jvp = jax.jvp( + lambda value: run(value, "bounded").ys[-1], + (rate,), + (jnp.ones_like(rate),), + )[1] + forward_jvp = jax.jvp( + lambda value: run(value, "forward").ys[-1], + (rate,), + (jnp.ones_like(rate),), + )[1] + assert jnp.allclose(bounded_jvp, forward_jvp, rtol=2e-7, atol=2e-9) + + +def test_float32_fixed_dae_is_invariant_to_nonbinding_budget(): + dtype = jnp.float32 + + def run(max_steps): + return solve_semi_explicit_dae( + lambda y, z, t: jnp.sin(jnp.asarray(20.0, dtype) * t), + lambda y, z: z - y, + Rodas5P(), + jnp.asarray(0.0, dtype), + jnp.asarray(1.0, dtype), + jnp.asarray(0.0, dtype), + jnp.asarray(0.0, dtype), + dt_0=jnp.asarray(0.001, dtype), + max_steps=max_steps, + save_at=SaveAt(t_1=True), + ) + + exact_budget = run(1000) + extra_budget = run(4096) + assert bool(exact_budget.ok & extra_budget.ok) + assert int(exact_budget.num_accepted) == 1000 + assert int(extra_budget.num_accepted) == 1000 + assert exact_budget.ts == jnp.asarray(1.0, dtype) + assert extra_budget.ts == jnp.asarray(1.0, dtype) + assert jnp.allclose(exact_budget.ys, extra_budget.ys, rtol=2e-5, atol=2e-6) + assert jnp.allclose(exact_budget.zs, extra_budget.zs, rtol=2e-5, atol=2e-6) + + def test_dae_dense_state_aux_jvp_vjp_and_reverse_over_forward(): grid = jnp.linspace(0.0, 1.0, 17) y_0 = jnp.asarray(1.0) diff --git a/tests/test_save_at.py b/tests/test_save_at.py index 29f4c5a..be5b29b 100644 --- a/tests/test_save_at.py +++ b/tests/test_save_at.py @@ -1,7 +1,16 @@ +import jax import jax.numpy as jnp import pytest -from tinydiffeq import IController, SaveAt, Tsit5, solve_ode +from tinydiffeq import ( + RK4, + ConstantStepSize, + Euler, + IController, + SaveAt, + Tsit5, + solve_ode, +) def solve(save_at, *, rtol=1e-9, atol=1e-12, max_steps=256, dt_0=0.2): @@ -71,6 +80,7 @@ def test_steps_omit_rejections_and_pad_with_last_value(): assert sol.ts[int(sol.num_accepted)] == 2.0 assert bool(jnp.all(sol.ts[n_valid:] == sol.ts[n_valid - 1])) assert bool(jnp.all(sol.xs[n_valid:] == sol.xs[n_valid - 1])) + assert int(sol.num_steps) > int(sol.num_accepted) def test_fill_inf_masks_non_accepted_rows(): @@ -102,6 +112,329 @@ def test_requested_grid_does_not_change_adaptive_step_count(): assert sampled.xs[-1] == endpoint.xs +def test_exact_fixed_grid_gathers_internal_rk4_states_and_supports_ad(): + grid = jnp.linspace(0.0, 2.0, 9) + + def sampled(parameter): + return solve_ode( + lambda x, t, args, p: -p * x, + RK4(), + 0.0, + 2.0, + jnp.asarray(1.0), + p=parameter, + dt_0=0.125, + controller=ConstantStepSize(), + max_steps=16, + save_at=SaveAt(ts=grid, exact=True), + has_aux=False, + ) + + parameter = jnp.asarray(0.4) + solution = sampled(parameter) + steps = solve_ode( + lambda x, t, args, p: -p * x, + RK4(), + 0.0, + 2.0, + jnp.asarray(1.0), + p=parameter, + dt_0=0.125, + controller=ConstantStepSize(), + max_steps=16, + save_at=SaveAt(steps=True), + has_aux=False, + ) + assert bool(solution.ok) + assert jnp.array_equal(solution.xs, steps.xs[::2]) + tangent = jax.jvp( + lambda value: sampled(value).xs, + (parameter,), + (jnp.ones_like(parameter),), + )[1] + gradient = jax.grad(lambda value: jnp.sum(sampled(value).xs))(parameter) + assert bool(jnp.all(jnp.isfinite(tangent))) + assert bool(jnp.isfinite(gradient)) + + +def test_exact_grid_accepts_clipped_final_step_and_rejects_nonknots(): + def run(grid): + return solve_ode( + lambda x: -x, + RK4(), + 0.0, + 1.0, + jnp.asarray(1.0), + dt_0=0.3, + max_steps=4, + save_at=SaveAt(ts=grid, exact=True), + has_aux=False, + ) + + aligned = run(jnp.asarray([0.0, 0.3, 0.6, 0.9, 1.0])) + unaligned = run(jnp.asarray([0.0, 0.3, 0.65, 1.0])) + assert bool(aligned.ok) + assert not bool(unaligned.ok) + assert aligned.xs[-1] == unaligned.xs[-1] + + +def test_float32_long_exact_grid_selects_unique_late_knots(): + dtype = jnp.float32 + horizon = jnp.asarray(100.0, dtype) + max_steps = 4096 + dt = horizon / max_steps + grid = horizon - dt * jnp.asarray([2.0, 1.0, 0.0], dtype) + solution = solve_ode( + lambda x, t: jnp.ones_like(x), + RK4(), + jnp.asarray(0.0, dtype), + horizon, + jnp.asarray(0.0, dtype), + dt_0=dt, + max_steps=max_steps, + save_at=SaveAt(ts=grid, exact=True), + has_aux=False, + ) + assert bool(solution.ok) + assert jnp.all(jnp.diff(solution.xs) > 0.0) + assert jnp.allclose(solution.xs, grid, rtol=2e-6, atol=2e-5) + + +def test_float32_exact_grid_alignment_accounts_for_large_time_offset(): + dtype = jnp.float32 + t_0 = jnp.asarray(-1000.0, dtype) + t_1 = jnp.asarray(1.0, dtype) + max_steps = 100 + dt = (t_1 - t_0) / max_steps + # linspace and t_0 + i * dt are mathematically the same grid, but their + # cancellation error near zero is governed by |t_0|, not the local time. + grid = jnp.linspace(t_0, t_1, max_steps + 1) + solution = solve_ode( + lambda x: jnp.ones_like(x), + RK4(), + t_0, + t_1, + jnp.asarray(0.0, dtype), + dt_0=dt, + max_steps=max_steps, + save_at=SaveAt(ts=grid, exact=True), + has_aux=False, + ) + assert bool(solution.ok) + assert jnp.allclose(solution.xs, grid - t_0, rtol=2e-6, atol=2e-4) + + +def test_exact_and_dense_fixed_grid_outputs_match_at_internal_knots(): + def run(save_at): + return solve_ode( + lambda x, t, args, p: jnp.sin(t) - p * x, + RK4(), + -0.5, + 1.0, + jnp.asarray(0.7), + p=jnp.asarray(0.2), + dt_0=0.125, + max_steps=12, + save_at=save_at, + has_aux=False, + ) + + steps = run(SaveAt(steps=True)) + knots = steps.ts[steps.accepted] + exact = run(SaveAt(ts=knots, exact=True)) + dense = run(SaveAt(ts=knots)) + assert bool(exact.ok & dense.ok) + assert jnp.array_equal(exact.xs, steps.xs[steps.accepted]) + assert jnp.allclose(dense.xs, exact.xs, rtol=1e-14, atol=1e-14) + + +def test_float32_fixed_solution_is_invariant_to_nonbinding_budget(): + dtype = jnp.float32 + horizon = jnp.asarray(100.0, dtype) + dt = jnp.asarray(0.025, dtype) + + def run(max_steps): + return solve_ode( + lambda x, t: jnp.sin(jnp.asarray(20.0, dtype) * t), + RK4(), + jnp.asarray(0.0, dtype), + horizon, + jnp.asarray(0.0, dtype), + dt_0=dt, + max_steps=max_steps, + has_aux=False, + ) + + exact_budget = run(4000) + extra_budget = run(8192) + assert bool(exact_budget.ok & extra_budget.ok) + assert int(exact_budget.num_accepted) == 4000 + assert int(extra_budget.num_accepted) == 4000 + assert exact_budget.ts == horizon + assert extra_budget.ts == horizon + assert jnp.allclose(exact_budget.xs, extra_budget.xs, rtol=1e-6, atol=1e-6) + + +def test_float32_uniform_gate_caps_tolerance_at_large_time_offsets(): + dtype = jnp.float32 + + def run(max_steps): + return solve_ode( + lambda x, t: jnp.where( + t >= jnp.asarray(1_000_000.25, dtype), + jnp.asarray(1e38, dtype), + jnp.zeros_like(x), + ), + Euler(), + 1_000_000.0, + 1_000_000.25, + jnp.asarray(0.0, dtype), + dt_0=0.125, + max_steps=max_steps, + has_aux=False, + ) + + exact_budget = run(2) + extra_budget = run(3) + assert bool(exact_budget.ok & extra_budget.ok) + assert int(exact_budget.num_steps) == int(extra_budget.num_steps) == 2 + assert exact_budget.ts == extra_budget.ts == jnp.asarray(1_000_000.25, dtype) + assert exact_budget.xs == extra_budget.xs == jnp.asarray(0.0, dtype) + + +def test_float32_fixed_horizon_tolerance_accounts_for_large_start_time(): + dtype = jnp.float32 + t_0 = -9_834_724.0 + t_1 = 6.2654047 + max_steps = 410 + dt_0 = (t_1 - t_0) / max_steps + + sol = solve_ode( + lambda x: jnp.zeros_like(x), + Euler(), + t_0, + t_1, + jnp.asarray(0.0, dtype), + dt_0=dt_0, + max_steps=max_steps, + has_aux=False, + ) + + assert bool(sol.ok) + assert int(sol.num_steps) == max_steps + assert sol.ts == jnp.asarray(t_1, dtype) + + +def test_float32_uniform_gate_requires_first_horizon_reach_at_final_slot(): + dtype = jnp.float32 + t_0 = 1_000_000.0 + t_1 = 1_000_000.0625 + dt_0 = (t_1 - t_0) / 10 + + def run(max_steps): + return solve_ode( + lambda x, t: jnp.where( + t >= jnp.asarray(t_1, dtype), + jnp.asarray(1e38, dtype), + jnp.zeros_like(x), + ), + Euler(), + t_0, + t_1, + jnp.asarray(0.0, dtype), + dt_0=dt_0, + max_steps=max_steps, + has_aux=False, + ) + + exact_budget = run(10) + extra_budget = run(11) + loose_budget = run(20) + assert bool(exact_budget.ok & extra_budget.ok & loose_budget.ok) + assert 0 < int(exact_budget.num_steps) < 10 + assert exact_budget.num_steps == extra_budget.num_steps == loose_budget.num_steps + assert ( + exact_budget.ts == extra_budget.ts == loose_budget.ts == jnp.asarray(t_1, dtype) + ) + assert ( + exact_budget.xs == extra_budget.xs == loose_budget.xs == jnp.asarray(0.0, dtype) + ) + + +def test_float32_uniform_gate_does_not_accept_a_different_horizon(): + dtype = jnp.float32 + dt = jnp.asarray(100.0 / 4096, dtype) + + def run(horizon): + return solve_ode( + lambda x: jnp.ones_like(x), + RK4(), + jnp.asarray(0.0, dtype), + jnp.asarray(horizon, dtype), + jnp.asarray(0.0, dtype), + dt_0=dt, + max_steps=4096, + has_aux=False, + ) + + shorter = run(99.9) + longer = run(100.1) + assert bool(shorter.ok) + assert shorter.ts == jnp.asarray(99.9, dtype) + assert not bool(longer.ok) + assert longer.ts < jnp.asarray(100.1, dtype) + + +def test_exact_grid_aux_and_derivatives_match_selected_step_rows(): + grid = jnp.linspace(0.0, 2.0, 9) + + def run(parameter, save_at): + return solve_ode( + lambda x, t, args, p: (-p * x, {"level": p * x + t}), + RK4(), + 0.0, + 2.0, + jnp.asarray(1.0), + p=parameter, + dt_0=0.125, + max_steps=16, + save_at=save_at, + has_aux=True, + ) + + parameter = jnp.asarray(0.4) + exact = run(parameter, SaveAt(ts=grid, exact=True)) + steps = run(parameter, SaveAt(steps=True)) + assert bool(exact.ok) + assert jnp.array_equal(exact.xs, steps.xs[::2]) + assert jnp.array_equal(exact.aux["level"], steps.aux["level"][::2]) + + def selected(value, exact_mode): + save_at = SaveAt(ts=grid, exact=True) if exact_mode else SaveAt(steps=True) + values = run(value, save_at).aux["level"] + return values if exact_mode else values[::2] + + exact_jvp = jax.jvp( + lambda value: selected(value, True), + (parameter,), + (jnp.ones_like(parameter),), + )[1] + step_jvp = jax.jvp( + lambda value: selected(value, False), + (parameter,), + (jnp.ones_like(parameter),), + )[1] + exact_vjp = jax.grad(lambda value: jnp.sum(selected(value, True)))(parameter) + step_vjp = jax.grad(lambda value: jnp.sum(selected(value, False)))(parameter) + assert jnp.allclose(exact_jvp, step_jvp, rtol=1e-12, atol=1e-12) + assert jnp.allclose(exact_vjp, step_vjp, rtol=1e-12, atol=1e-12) + + +def test_exact_grid_requires_fixed_explicit_ode(): + with pytest.raises(ValueError, match="ConstantStepSize"): + solve(SaveAt(ts=[0.0, 1.0, 2.0], exact=True)) + + def test_save_at_exclusivity(): with pytest.raises(ValueError, match="exactly one"): SaveAt() @@ -111,6 +444,8 @@ def test_save_at_exclusivity(): SaveAt(t_1=True, ts=jnp.linspace(0.0, 1.0, 5)) with pytest.raises(ValueError, match="fill"): SaveAt(steps=True, fill="zero") + with pytest.raises(ValueError, match="requires ts"): + SaveAt(steps=True, exact=True) def test_noncanonical_public_spellings_are_rejected(): diff --git a/tests/test_sdae.py b/tests/test_sdae.py index 003238a..ed822e9 100644 --- a/tests/test_sdae.py +++ b/tests/test_sdae.py @@ -68,6 +68,9 @@ def test_sdae_matches_reduced_sde_on_identical_noise_path(): save_at=SaveAt(steps=True), ) assert bool(full.ok) + assert int(full.num_steps) == n + assert int(full.num_root_solves) == n + 1 + assert 0 <= int(full.num_root_steps) <= 8 * int(full.num_root_solves) assert jnp.allclose(full.ys, reduced.xs, atol=2e-6, rtol=2e-6) assert jnp.allclose(full.zs, full.ys, atol=2e-6) assert jnp.allclose(full.aux["scaled"], 2.0 * full.zs) @@ -287,6 +290,7 @@ def bad_diffusion(y, z, t, args, p, algebraic_aux): ) assert not bool(sol.ok) assert int(sol.num_accepted) == 0 + assert int(sol.num_steps) == 0 assert jnp.all(sol.accepted == jnp.asarray([True] + [False] * 8)) assert jnp.all(sol.ys == 1.0) assert jnp.all(sol.aux["bad"] == 0.0) @@ -316,5 +320,6 @@ def zero_diffusion(y, z, t, args, p, algebraic_aux): ) assert not bool(sol.ok) assert int(sol.num_accepted) == 0 + assert int(sol.num_steps) == 1 assert jnp.all(sol.ys == 0.0) assert jnp.all(sol.aux["limited"] == sol.aux["limited"][0]) diff --git a/tests/test_sde.py b/tests/test_sde.py index d830645..f32c455 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -199,5 +199,6 @@ def test_steps_mode_shapes_and_flags(): assert sol.xs.shape == (n + 1, 2) assert bool(sol.ok) assert int(sol.num_accepted) == n + assert int(sol.num_steps) == n assert bool(jnp.all(sol.accepted)) assert sol.ts[0] == 0.0 and sol.ts[-1] == T diff --git a/tests/test_solution.py b/tests/test_solution.py new file mode 100644 index 0000000..0251719 --- /dev/null +++ b/tests/test_solution.py @@ -0,0 +1,59 @@ +import jax +import jax.numpy as jnp + +from tinydiffeq import ( + RK4, + DenseExponential, + DiscreteMarkovChain, + EulerMaruyama, + simulate_markov_chain, + solve_linear_ode, + solve_ode, + solve_sde, +) + + +def test_endpoint_solution_step_counts_have_uniform_array_pytree_structure(): + ode = solve_ode( + lambda x: -x, + RK4(), + 0.0, + 1.0, + jnp.asarray(1.0), + dt_0=0.25, + max_steps=4, + has_aux=False, + ) + sde = solve_sde( + lambda x: -x, + lambda x: jnp.zeros_like(x), + EulerMaruyama(), + 0.0, + 1.0, + jnp.asarray(1.0), + key=jax.random.key(0), + n_steps=4, + has_aux=False, + ) + exponential = solve_linear_ode( + jnp.asarray([[-1.0]]), + DenseExponential(), + 0.0, + 1.0, + jnp.asarray([1.0]), + ) + markov = simulate_markov_chain( + DiscreteMarkovChain([[0.0, 1.0], [1.0, 0.0]]), + jnp.asarray(0, jnp.int32), + key=jax.random.key(1), + num_steps=4, + ) + + solutions = (ode, sde, exponential, markov) + assert tuple(int(solution.num_steps) for solution in solutions) == (4, 4, 1, 4) + expected_structure = jax.tree.structure(ode) + for solution in solutions: + assert jax.tree.structure(solution) == expected_structure + assert isinstance(solution.num_steps, jax.Array) + assert solution.num_steps.shape == () + assert solution.num_steps.dtype == jnp.int32 diff --git a/tests/test_solvers_fixed.py b/tests/test_solvers_fixed.py index d355b26..27e0885 100644 --- a/tests/test_solvers_fixed.py +++ b/tests/test_solvers_fixed.py @@ -24,6 +24,8 @@ def f(x): euler = solve_ode(f, Euler(), 0.0, T, x_0, dt_0=T / n, max_steps=n) rk4 = solve_ode(f, RK4(), 0.0, T, x_0, dt_0=T / n, max_steps=n) assert bool(euler.ok) and bool(rk4.ok) + assert int(euler.num_steps) == int(euler.num_accepted) == n + assert int(rk4.num_steps) == int(rk4.num_accepted) == n assert jnp.max(jnp.abs(euler.xs - exact)) < 2e-3 assert jnp.max(jnp.abs(rk4.xs - exact)) < 1e-12 @@ -91,6 +93,67 @@ def f(x, t): assert evaluation_times == [0.0, 0.25, 0.5, 0.75] +def test_fixed_traced_horizon_uses_one_clipped_scan_with_static_branch_parity(): + def traced_solve(horizon): + return solve_ode( + lambda x: -x, + RK4(), + 0.0, + horizon, + jnp.asarray(1.0), + dt_0=0.25, + max_steps=4, + has_aux=False, + ) + + static = solve_ode( + lambda x: -x, + RK4(), + 0.0, + 1.0, + jnp.asarray(1.0), + dt_0=0.25, + max_steps=4, + has_aux=False, + ) + traced = jax.jit(traced_solve)(jnp.asarray(1.0)) + assert jnp.array_equal(traced.ts, static.ts) + assert jnp.array_equal(traced.xs, static.xs) + assert traced.ok == static.ok + assert traced.num_accepted == static.num_accepted + assert traced.num_steps == static.num_steps + + horizons = jnp.asarray([0.75, 1.0]) + batched = jax.jit(jax.vmap(traced_solve))(horizons) + for index, horizon in enumerate(horizons): + scalar = jax.jit(traced_solve)(horizon) + assert jnp.array_equal(batched.ts[index], scalar.ts) + assert jnp.array_equal(batched.xs[index], scalar.xs) + assert batched.ok[index] == scalar.ok + assert batched.num_accepted[index] == scalar.num_accepted + + # A batched cond over two complete scans executes both branches. The traced + # horizon path must instead contain only the one clipped integration scan. + jaxpr = str(jax.make_jaxpr(jax.vmap(traced_solve))(horizons)) + assert jaxpr.count("scan[") == 1 + + def static_uniform(initial): + return solve_ode( + lambda x: -x, + RK4(), + 0.0, + 1.0, + initial, + dt_0=0.25, + max_steps=4, + has_aux=False, + ).xs + + static_jaxpr = str(jax.make_jaxpr(static_uniform)(jnp.asarray(1.0))) + assert static_jaxpr.count("scan[") == 1 + assert "cond[" not in static_jaxpr + + def test_scalar_vs_vector_shapes(): def f(x): return -x diff --git a/uv.lock b/uv.lock index b978581..df761a5 100644 --- a/uv.lock +++ b/uv.lock @@ -1227,7 +1227,7 @@ wheels = [ [[package]] name = "tinydiffeq" -version = "2.3.0" +version = "2.4.0" source = { editable = "." } dependencies = [ { name = "jax" },