feat: solve_bvp collocation solver with implicit differentiation - #11
Merged
Conversation
Level-batched cyclic reduction with orthogonal eliminations for bordered almost-block-diagonal systems: one batched complete QR per level, a final dense boundary LU, forward and transpose solves from the same factors, and a dense scatter/matvec reference. Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g
Faithful port of scipy.integrate.solve_bvp: 4th-order Lobatto IIIA collocation, damped Newton with the affine-invariant criterion, 5-point Lobatto residual control, and insert-1/insert-2 mesh refinement, with scipy's constants and status semantics. Local Jacobians come from AD instead of finite differences; the collocation system is factored by the structured BABD cyclic reduction each Newton refresh. The whole solve is jit-compiled with lax.while_loop loops and padded to a static max_nodes, vmaps with per-lane statuses, and carries one custom_jvp: the implicit function theorem at the solution on the frozen final mesh, transposing to reverse mode through lax.custom_linear_solve and recursing for higher order. Failures are data (statuses 0-3, never raises in traced code); failed lanes have exact-zero tangents and a converged solve whose Jacobian cannot be refactored reports NaN tangents in both AD modes. Unknown parameters are z (guess z_0, solved jointly); p is the only AD input; args is inert. Adds BVPSolution, hermite_derivative, docs, and tests cross-checked against scipy. Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g
fn._cache_size() reads JAX's globally shared C++ executable cache (capacity 8192, LRU), so absolute entry-count assertions turn flaky as the suite grows. Recompilation hygiene is instead audited manually with the environment-variable protocol recorded in AGENTS.md. Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g
Cold-compile and warm timings of jitted-wrapper solves against scipy references, plus a cache-audit mode for the AGENTS.md recompilation protocol. Warm jitted calls beat scipy on every recorded case. Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g
Lead the README and docs index with the unsupported-research-repo disclaimer naming SciML, scipy, and diffrax as the reference implementations, and trim the README to one example block with an SDE ensemble snippet. Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ports scipy v1.18's
solve_bvpto JAX astinydiffeq.solve_bvp: the same 4th-order Lobatto IIIA collocation, damped Newton method, 5-point Lobatto residual control, mesh refinement, constants, and status semantics, validated node-for-node against scipy to machine epsilon (including failure paths).Beyond the port: fully jit-compiled with static
max_nodespadding and zero recompiles across changed meshes/guesses/p/tolerances; local Jacobians from AD instead of finite differences; onecustom_jvpimplementing the implicit function theorem on the frozen final mesh (reverse mode via transposition throughlax.custom_linear_solve, exact higher order); vmap with per-lane statuses; float32/GPU support; and a structured orthogonal factorization of the bordered almost-block-diagonal collocation system (tinydiffeq.babd, level-batched cyclic reduction). Unknown parameters arez(solved jointly);pis the only AD input; the singular termSis supported.Warm jitted calls beat scipy on every benchmark case (shock layer 2.1 ms vs 6.3 ms; 5-node problem 0.25 ms vs 0.60 ms); the kernels neoclassical-growth baseline (1489 nodes, tol=1e-10) matches scipy's meshes node-for-node at runtime parity. Also removes the flaky
_cache_size()tests in favor of the AGENTS.md audit protocol, adds the repo disclaimer, and trims the README.Reviewed pre-merge by Codex and an Opus JAX reviewer; all confirmed findings fixed.
🤖 Generated with Claude Code
https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g