From 419250a761a17927b4ba977f67881eb0bdc79368 Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Wed, 19 Aug 2026 21:25:22 -0700 Subject: [PATCH 1/5] feat: structured BABD factorization module 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) Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g --- src/tinydiffeq/babd.py | 298 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 298 insertions(+) create mode 100644 src/tinydiffeq/babd.py diff --git a/src/tinydiffeq/babd.py b/src/tinydiffeq/babd.py new file mode 100644 index 0000000..51ca495 --- /dev/null +++ b/src/tinydiffeq/babd.py @@ -0,0 +1,298 @@ +"""Bordered almost block diagonal (BABD) linear systems. + +A BABD system couples chain unknowns ``x_0 .. x_E`` (size ``n`` each) and +border unknowns ``z`` (size ``k``) through ``E`` staircase equations, each +involving two adjacent chain unknowns, plus ``n + k`` boundary equations +involving the first and last chain unknowns: + + block_left[i] x_i + block_right[i] x_{i+1} + block_border[i] z = f_i + boundary_first x_0 + boundary_last x_E + boundary_border z = g + +The block container is a dict with those six keys (``block_border`` and +``boundary_border`` are ``None`` when ``k == 0``). Unknowns and equations are +flattened in that order. This is the structure of collocation and multiple +shooting discretizations of two-point boundary value problems. +""" + +import jax +import jax.numpy as jnp +import numpy as np +from jax.scipy import linalg as jsp_linalg + + +def babd_dimensions(blocks): + chain, n = blocks["block_left"].shape[:2] + k = 0 if blocks["block_border"] is None else blocks["block_border"].shape[-1] + return chain, n, k + + +def babd_indices(n, m, k): + # Row/column indices of the six blocks in the dense embedding, matching + # scipy's compute_jac_indices with node-leading value ravels. The combined + # index set has no duplicates, so a scatter set matches sparse summing. + i_col = np.repeat(np.arange((m - 1) * n), n) + j_col = np.tile(np.arange(n), n * (m - 1)) + np.repeat(np.arange(m - 1) * n, n**2) + + i_bc = np.repeat(np.arange((m - 1) * n, m * n + k), n) + j_bc = np.tile(np.arange(n), n + k) + + i_p_col = np.repeat(np.arange((m - 1) * n), k) + j_p_col = np.tile(np.arange(m * n, m * n + k), (m - 1) * n) + + i_p_bc = np.repeat(np.arange((m - 1) * n, m * n + k), k) + j_p_bc = np.tile(np.arange(m * n, m * n + k), n + k) + + i = np.hstack((i_col, i_col, i_bc, i_bc, i_p_col, i_p_bc)) + j = np.hstack((j_col, j_col + n, j_bc, j_bc + (m - 1) * n, j_p_col, j_p_bc)) + return i, j + + +def babd_dense(blocks): + # Scatter the blocks into the dense square matrix. + chain, n, k = babd_dimensions(blocks) + dtype = blocks["block_left"].dtype + i_jac, j_jac = babd_indices(n, chain + 1, k) + values = [ + blocks["block_left"].reshape(-1), + blocks["block_right"].reshape(-1), + blocks["boundary_first"].reshape(-1), + blocks["boundary_last"].reshape(-1), + ] + if k > 0: + values.extend( + [blocks["block_border"].reshape(-1), blocks["boundary_border"].reshape(-1)] + ) + size = (chain + 1) * n + k + flat = jnp.concatenate(values) + return jnp.zeros((size, size), dtype).at[i_jac, j_jac].set(flat) + + +def babd_matvec(blocks, u): + # Structured product with a flat vector ordered (x_0 .. x_E, z). + chain, n, k = babd_dimensions(blocks) + x = u[: (chain + 1) * n].reshape(chain + 1, n) + z = u[(chain + 1) * n :] + staircase = jnp.einsum("eij,ej->ei", blocks["block_left"], x[:-1]) + jnp.einsum( + "eij,ej->ei", blocks["block_right"], x[1:] + ) + boundary = blocks["boundary_first"] @ x[0] + blocks["boundary_last"] @ x[-1] + if k > 0: + staircase = staircase + jnp.einsum("eik,k->ei", blocks["block_border"], z) + boundary = boundary + blocks["boundary_border"] @ z + return jnp.concatenate([staircase.reshape(-1), boundary]) + + +def structured_qr_factor(blocks): + # Cyclic reduction with orthogonal eliminations (the level-parallel + # ordering of Wright's structured orthogonal factorization): each level + # pairs adjacent staircase equations and eliminates every shared unknown + # at once through one batched complete QR of the stacked (2n, n) middle + # columns, so a chain of length E costs ~log2(E) batched calls instead of + # E sequential ones. Orthogonal eliminations keep every stored row + # bounded by the original row norms, so saddle-path dichotomies cannot + # overflow the way naive condensation does. What remains couples + # (x_0, x_E, z) in one dense (2n + k)-square system. + chain, n, k = babd_dimensions(blocks) + dtype = blocks["block_left"].dtype + a = blocks["block_left"] + b = blocks["block_right"] + c = blocks["block_border"] + if c is None: + c = jnp.zeros((chain, n, 0), dtype) + boundary_border = blocks["boundary_border"] + if boundary_border is None: + boundary_border = jnp.zeros((blocks["boundary_first"].shape[0], 0), dtype) + + levels = [] + ok = jnp.asarray(True) + length = chain + while length > 1: + pairs = length // 2 + odd = length - 2 * pairs + a_even = a[0 : 2 * pairs : 2] + b_even = b[0 : 2 * pairs : 2] + c_even = c[0 : 2 * pairs : 2] + a_odd = a[1 : 2 * pairs : 2] + b_odd = b[1 : 2 * pairs : 2] + c_odd = c[1 : 2 * pairs : 2] + q, r_full = jnp.linalg.qr( + jnp.concatenate([b_even, a_odd], axis=1), mode="complete" + ) + q_t = jnp.swapaxes(q, -1, -2) + zeros_pair = jnp.zeros_like(a_even) + left_pair = q_t @ jnp.concatenate([a_even, zeros_pair], axis=1) + right_pair = q_t @ jnp.concatenate([zeros_pair, b_odd], axis=1) + border_pair = q_t @ jnp.concatenate([c_even, c_odd], axis=1) + r = r_full[:, :n] + diag_r = jnp.diagonal(r, axis1=-2, axis2=-1) + ok = ( + ok + & jnp.all(jnp.isfinite(diag_r)) + & (jnp.min(jnp.abs(diag_r), initial=jnp.inf) > 0.0) + ) + levels.append( + dict( + q=q, + r=r, + left=left_pair[:, :n], + right=right_pair[:, :n], + border=border_pair[:, :n], + ) + ) + a_next = left_pair[:, n:] + b_next = right_pair[:, n:] + c_next = border_pair[:, n:] + if odd: + a_next = jnp.concatenate([a_next, a[-1:]], axis=0) + b_next = jnp.concatenate([b_next, b[-1:]], axis=0) + c_next = jnp.concatenate([c_next, c[-1:]], axis=0) + a, b, c = a_next, b_next, c_next + length = pairs + odd + + final = jnp.concatenate( + [ + jnp.concatenate([a[0], b[0], c[0]], axis=1), + jnp.concatenate( + [blocks["boundary_first"], blocks["boundary_last"], boundary_border], + axis=1, + ), + ], + axis=0, + ) + lu, pivots = jsp_linalg.lu_factor(final, check_finite=False) + ok = jax.lax.stop_gradient( + ok & jnp.all(jnp.isfinite(lu)) & jnp.all(jnp.abs(jnp.diag(lu)) > 0.0) + ) + # An unusable factor becomes the identity so batched lanes stay finite. + size = final.shape[0] + state = dict( + levels=tuple( + dict( + q=jnp.where(ok, level["q"], jnp.eye(2 * n, dtype=dtype)), + r=jnp.where(ok, level["r"], jnp.eye(n, dtype=dtype)), + left=jnp.where(ok, level["left"], jnp.zeros_like(level["left"])), + right=jnp.where(ok, level["right"], jnp.zeros_like(level["right"])), + border=jnp.where(ok, level["border"], jnp.zeros_like(level["border"])), + ) + for level in levels + ), + template=jnp.zeros(n, dtype), + final_lu=jnp.where(ok, lu, jnp.eye(size, dtype=dtype)), + final_pivots=jnp.where(ok, pivots, jnp.arange(size, dtype=pivots.dtype)), + ) + return state, ok + + +def structured_qr_level_sizes(state, chain): + sizes = [] + length = chain + for level in state["levels"]: + pairs = level["q"].shape[0] + odd = length - 2 * pairs + sizes.append((pairs, odd)) + length = pairs + odd + return sizes + + +def structured_qr_solve(state, rhs): + n = state["template"].shape[0] + k = state["final_lu"].shape[0] - 2 * n + chain = (rhs.shape[0] - k) // n - 1 + f = rhs[: chain * n].reshape(chain, n) + boundary_rhs = rhs[chain * n :] + sizes = structured_qr_level_sizes(state, chain) + + tops = [] + for level, (pairs, odd) in zip(state["levels"], sizes, strict=True): + stacked = jnp.concatenate([f[0 : 2 * pairs : 2], f[1 : 2 * pairs : 2]], axis=1) + transformed = jnp.einsum("pji,pj->pi", level["q"], stacked) + tops.append(transformed[:, :n]) + f_next = transformed[:, n:] + if odd: + f_next = jnp.concatenate([f_next, f[-1:]], axis=0) + f = f_next + + boundary = jsp_linalg.lu_solve( + (state["final_lu"], state["final_pivots"]), + jnp.concatenate([f[0], boundary_rhs]), + trans=0, + check_finite=False, + ) + x_first, x_last, z = boundary[:n], boundary[n : 2 * n], boundary[2 * n :] + + x_nodes = jnp.stack([x_first, x_last], axis=0) + for level, top, (pairs, _) in reversed( + list(zip(state["levels"], tops, sizes, strict=True)) + ): + rhs_mid = ( + top + - jnp.einsum("pij,pj->pi", level["left"], x_nodes[:pairs]) + - jnp.einsum("pij,pj->pi", level["right"], x_nodes[1 : pairs + 1]) + - jnp.einsum("pik,k->pi", level["border"], z) + ) + x_mid = jsp_linalg.solve_triangular( + level["r"], rhs_mid[..., None], lower=False, check_finite=False + )[..., 0] + interleaved = jnp.stack([x_nodes[:pairs], x_mid], axis=1).reshape(2 * pairs, n) + x_nodes = jnp.concatenate([interleaved, x_nodes[pairs:]], axis=0) + return jnp.concatenate([x_nodes.reshape(-1), z]) + + +def structured_qr_transpose_solve(state, rhs): + # J = QU from the factorization, so J^T y = b is the lower-triangular + # sweep U^T v = b down the levels followed by y = Q v back up. + n = state["template"].shape[0] + dtype = state["template"].dtype + k = state["final_lu"].shape[0] - 2 * n + nodes = (rhs.shape[0] - k) // n + chain = nodes - 1 + b_nodes = rhs[: nodes * n].reshape(nodes, n) + b_z = rhs[nodes * n :] + sizes = structured_qr_level_sizes(state, chain) + + sum_border = jnp.zeros(k, dtype) + adjoints = [] + for level, (pairs, _) in zip(state["levels"], sizes, strict=True): + v = jsp_linalg.solve_triangular( + level["r"], + b_nodes[1 : 2 * pairs : 2][..., None], + lower=False, + trans=1, + check_finite=False, + )[..., 0] + adjoints.append(v) + sum_border = sum_border + jnp.einsum("pik,pi->k", level["border"], v) + contrib_left = jnp.einsum("pij,pi->pj", level["left"], v) + contrib_right = jnp.einsum("pij,pi->pj", level["right"], v) + reduced = jnp.concatenate( + [b_nodes[0 : 2 * pairs : 2], b_nodes[2 * pairs :]], axis=0 + ) + tail = reduced.shape[0] - pairs + reduced = reduced - jnp.concatenate( + [contrib_left, jnp.zeros((tail, n), dtype)], axis=0 + ) + reduced = reduced - jnp.concatenate( + [jnp.zeros((1, n), dtype), contrib_right, jnp.zeros((tail - 1, n), dtype)], + axis=0, + ) + b_nodes = reduced + + boundary = jsp_linalg.lu_solve( + (state["final_lu"], state["final_pivots"]), + jnp.concatenate([b_nodes[0], b_nodes[1], b_z - sum_border]), + trans=1, + check_finite=False, + ) + + y_rows = boundary[:n][None] + for level, v, (pairs, _) in reversed( + list(zip(state["levels"], adjoints, sizes, strict=True)) + ): + pair = jnp.einsum( + "pij,pj->pi", level["q"], jnp.concatenate([v, y_rows[:pairs]], axis=1) + ) + interleaved = jnp.stack([pair[:, :n], pair[:, n:]], axis=1).reshape( + 2 * pairs, n + ) + y_rows = jnp.concatenate([interleaved, y_rows[pairs:]], axis=0) + return jnp.concatenate([y_rows.reshape(-1), boundary[n:]]) From 008b06fbe34ac06f7d465094bdcb46103c8a54b4 Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Wed, 19 Aug 2026 21:25:38 -0700 Subject: [PATCH 2/5] feat: port scipy solve_bvp collocation solver 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) Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g --- docs/api.md | 6 + docs/bvp.md | 222 ++++++++ docs/dae.md | 5 +- docs/llms.txt | 3 +- mkdocs.yml | 1 + src/tinydiffeq/__init__.py | 22 +- src/tinydiffeq/_bvp_core.py | 121 +++++ src/tinydiffeq/bvp.py | 936 ++++++++++++++++++++++++++++++++ src/tinydiffeq/interpolation.py | 38 ++ src/tinydiffeq/solution.py | 32 ++ tests/test_bvp.py | 623 +++++++++++++++++++++ tests/test_bvp_ad.py | 254 +++++++++ tests/test_bvp_vmap.py | 84 +++ 13 files changed, 2336 insertions(+), 11 deletions(-) create mode 100644 docs/bvp.md create mode 100644 src/tinydiffeq/_bvp_core.py create mode 100644 src/tinydiffeq/bvp.py create mode 100644 tests/test_bvp.py create mode 100644 tests/test_bvp_ad.py create mode 100644 tests/test_bvp_vmap.py diff --git a/docs/api.md b/docs/api.md index 8f4a596..5e89423 100644 --- a/docs/api.md +++ b/docs/api.md @@ -4,6 +4,8 @@ ::: tinydiffeq.solve_ode +::: tinydiffeq.solve_bvp + ::: tinydiffeq.solve_semi_explicit_dae ::: tinydiffeq.solve_semi_explicit_sdae @@ -76,6 +78,8 @@ ::: tinydiffeq.DAESolution +::: tinydiffeq.BVPSolution + ::: tinydiffeq.Solution ::: tinydiffeq.MarkovDistribution @@ -86,4 +90,6 @@ ::: tinydiffeq.hermite_interpolate +::: tinydiffeq.hermite_derivative + ::: tinydiffeq.cumulative_trapezoid diff --git a/docs/bvp.md b/docs/bvp.md new file mode 100644 index 0000000..65a6688 --- /dev/null +++ b/docs/bvp.md @@ -0,0 +1,222 @@ +# Boundary Value Problems + +`solve_bvp` solves two-point boundary value problems of the form + +$$ +\frac{dy}{dt} = f(t, y, z, \mathrm{args}, p) + \frac{S\,y}{t - t_a}, +\qquad +\mathrm{bc}\big(y(t_a),\, y(t_b),\, z, \mathrm{args}, p\big) = 0, +$$ + +on $t \in [t_a, t_b]$. It is a faithful JAX port of +[`scipy.integrate.solve_bvp`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_bvp.html): +the same 4th-order Lobatto IIIA collocation, the same damped Newton method +with an affine-invariant criterion, the same 5-point Lobatto residual +estimator and insert-1/insert-2 mesh refinement, and the same constants and +default tolerances. The optional singular term ($S$ an $n \times n$ matrix on +the flattened state, requiring $S\,y(t_a) = 0$) covers Lane–Emden style +problems. + +## `z` versus `p` + +The one distinction to internalize: + +- **`z` are scipy's *unknown parameters*** — solved jointly with $y$ + (an eigenvalue, a free constant chosen by an extra boundary condition). + You pass a guess `z_0` (any pytree) and read the solved value from + `sol.z`. When present, `bc` must return `n + size(z)` residuals. Like + every guess, `z_0` is differentiation-inert. +- **`p` are *known differentiable parameters*** — the only AD input, as + everywhere else in tinydiffeq. JVP and VJP rules differentiate `sol.y`, + `sol.yp`, `sol.z`, and `sol.aux` with respect to `p` implicitly at the + solution, never through the iterations. +- **`args` is inert pass-through data.** + +scipy names the unknowns `p`; they are renamed here because in tinydiffeq +`p` always means the differentiable input. The solved unknowns are exactly +the DAE interface's unknown-with-a-guess role, hence `z`. + +## Interface + +```python +solve_bvp(fun, bc, t, y_0, z_0=None, *, p=None, args=None, S=None, + fun_jac_ad="auto", bc_jac_ad="auto", tol=1e-3, bc_tol=None, + max_nodes=128, has_aux=None) +``` + +`fun` and `bc` are **pointwise** — a scalar `t` and a single node's state +pytree, vmapped over the mesh internally (scipy instead passes the whole +`(n, m)` mesh; port scipy code by deleting the vectorization). Both may take +two to five positional arguments, always in this order: + +```python +fun(t, y) | fun(t, y, z) | fun(t, y, z, args) | fun(t, y, z, args, p) +bc(ya, yb) | bc(ya, yb, z) | bc(ya, yb, z, args) | bc(ya, yb, z, args, p) +``` + +`t` is the initial mesh (scipy's `x`): strictly increasing, at least two +nodes, at most `max_nodes`. `y_0` is the initial guess — any pytree with +leading axis `len(t)` on every leaf and one shared real floating dtype, +which becomes the working dtype. `fun` may return `(value, aux)`; aux is +evaluated once at the solution and participates in AD. + +```python +import jax.numpy as jnp +from tinydiffeq import solve_bvp + +# Sturm–Liouville: y'' = -z^2 y, y(0) = y(pi) = 0, y'(0) = z. +def fun(t, y, z): + return jnp.array([y[1], -z[0]**2 * y[0]]) + +def bc(ya, yb, z): + return jnp.array([ya[0], yb[0], ya[1] - z[0]]) + +t = jnp.linspace(0.0, jnp.pi, 5) +sol = solve_bvp(fun, bc, t, jnp.ones((5, 2)), jnp.array([0.5])) +sol.z # the eigenvalue, ~1.0 +sol.num_nodes # active nodes on the refined mesh +``` + +## Static shapes + +The mesh grows under refinement, so every returned array is padded to the +static `max_nodes` (default 128 — smaller than scipy's 1000; every loop and +the factorization run over all `max_nodes` padded intervals, so cost grows +linearly in the budget): + +- `sol.t` has shape `(max_nodes,)`; entries past `sol.num_nodes` repeat + $t_b$ exactly. +- `sol.y` and `sol.yp` leaves have leading axis `max_nodes`; tail rows + repeat the last active row bitwise. `sol.yp` holds the (singular-term + corrected) right-hand side at the nodes. +- `sol.rms_residuals` has shape `(max_nodes - 1,)` and is exactly zero on + inactive intervals. +- `sol.aux` rows past the last active node duplicate the endpoint value. + +`max_nodes` is static; changing it recompiles, and `fun` and `bc` key the +compilation cache by object identity — define them at module scope rather +than rebuilding closures at a hot call site. +Everything else — mesh values, guesses, `p`, `args`, `S` values, `tol`, and +`bc_tol` — is traced data and never retraces. The compiled solve is also +shared across initial mesh lengths (inputs are padded to `max_nodes` before +dispatch), though under an outer `jit` a changed input length retraces that +outer function, as any shape change does. + +## Dense output + +There is no callable solution object. The padded tails are what make the +plain arrays sufficient: `hermite_interpolate(ts, sol.t, sol.y, sol.yp)` +evaluates **exactly** the C1 cubic spline scipy returns as `sol.sol(ts)` +(scipy's `create_spline` is the cubic Hermite interpolant of `(y, yp)`), +and `hermite_derivative(ts, sol.t, sol.y, sol.yp)` is `sol.sol(ts, 1)`. +Queries outside $[t_a, t_b]$ clamp to the endpoint values (scipy's `PPoly` +extrapolates the cubic instead); derivatives outside the span are zero. + +## Statuses and failure + +The solve never raises inside traced code: failures are reported as data. +A failed status returns the last iterate, which — as in scipy — may be +non-finite when the final Newton candidate diverged; check `sol.ok` before +trusting values. `sol.status` carries scipy's codes: + +- `0` — converged to the desired accuracy (`sol.ok`). +- `1` — the refinement wanted more than `max_nodes` nodes; the reported + mesh and solution are the last completed iteration's. +- `2` — a singular collocation Jacobian; detected as a non-finite or + exactly rank-deficient LU factor (scipy's `splu` raises here), with the + last iterate returned. +- `3` — the boundary-condition tolerance was not satisfied within 10 + iterations after the mesh stopped refining. + +`tol` is floored at `100 * eps` of the working dtype silently (scipy warns; +`tol` may be a tracer here). In float32 that floor is ~1.2e-5, so tighten +tolerances only as far as the dtype supports. + +## AD contract + +The whole solve sits behind one `custom_jvp`: the implicit function theorem +applied to the collocation system $F(Y, z; p) = 0$ on the frozen final mesh, +using the same assembled Jacobian the Newton method factors. Reverse mode is +JAX's transposition of that rule — there is no separate VJP rule, and the +iteration count, damping, and mesh are never differentiated through. + +- Only `p` is an AD input. `sol.y`, `sol.yp`, `sol.z`, and `sol.aux` carry + tangents; `sol.t`, `sol.rms_residuals`, and every counter and status are + differentiation-inert with exact-zero tangents. +- The guesses `t`, `y_0`, `z_0` and the inert `args` and `S` have exact-zero + gradients by contract. A `p`-dependent singular term belongs inside `fun`. +- Higher-order derivatives (hessians, reverse-over-forward) are exact on the + frozen mesh: the rule leaves the solution and Jacobian differentiable, so + outer transforms recurse through the same implicit rule. +- A failed solve (`status != 0`) has exact-zero, finite tangents; under + `vmap`, a failed lane's tangent program is evaluated at the inert initial + guess so it cannot poison successful lanes. The one loud exception: a + *converged* solve whose final-mesh Jacobian fails to refactor inside the + AD rule has no computable derivative and reports NaN tangents, + lane-locally. +- Wrap gradient computations in `jax.jit` — op-by-op assembly and + factorization of the collocation Jacobian is an order of magnitude slower. +- To differentiate with respect to the endpoints $t_a, t_b$, rescale the + problem to a fixed interval and put the endpoints in `p`. + +## Jacobians and the linear solve + +Local Jacobians of `fun` and `bc` come from AD, not scipy's forward +differences — `fun_jac_ad` and `bc_jac_ad` select `"jvp"` (`jacfwd`), +`"vjp"` (`jacrev`), or `"auto"` (forward when square or tall, reverse when +strictly fat, block by block). There is no analytic-Jacobian argument. The +finite-difference parameter Jacobians are the piece of scipy most prone to +pushing a marginal collocation system singular; AD removes that failure +mode. + +The collocation Jacobian is bordered almost block diagonal — a staircase of +`n`-square blocks coupling adjacent nodes, a dense column border for `z`, +and boundary rows tying the two endpoints — a structure fixed by the +discretization, not the problem. Where scipy hands the assembled sparse +matrix to SuperLU, each Newton refresh here runs a structured orthogonal +factorization (`tinydiffeq.babd`): cyclic reduction eliminates all pair +midpoints per level through one batched complete QR, ~`log2(max_nodes)` +batched calls in total, leaving one dense `(2n + size(z))`-square boundary +system. Factorization costs `O(max_nodes n^3)` and each solve +`O(max_nodes n^2)`, so padding to the static `max_nodes` is nearly free, and +orthogonal eliminations are stable on saddle-path dichotomies where naive +condensation overflows. The factorization is reused across the backtracking +line search and fixed-Jacobian iterations, exactly as scipy reuses its +`splu` object, and the AD rule reuses it through `lax.custom_linear_solve` +with its transpose solve. + +On GPUs in float32, set +`jax.config.update("jax_default_matmul_precision", "highest")` as for the +[ODE solvers](ode.md#stiff-odes-rodas5p): the Jacobian assembly and the +singular-term products are matmuls that XLA otherwise serves from TF32. + +The whole solve is compiled, with `lax.while_loop` outer, Newton, and +backtracking loops. For repeated solves, put the `solve_bvp` call inside +your own `jax.jit` (or `vmap`): calling it from un-jitted Python re-runs +the wrapper's validation, flattening, and dispatch every call — a few +milliseconds that dwarf a small compiled solve — while inside a jitted +function that work happens once at trace time and warm calls run at +compiled speed (faster than scipy even on 5-node problems; see +`benchmarks/results/`). + +## Credit and deviations + +The algorithm is a direct port of +[scipy's `_bvp.py`](https://github.com/scipy/scipy/blob/v1.18.0/scipy/integrate/_bvp.py) +(BSD-3), which implements the residual-control collocation method of +Kierzenka and Shampine, *A BVP Solver Based on Residual Control and the +MATLAB PSE* (ACM TOMS 27(3), 2001), with the damped Newton method of +Ascher, Mattheij, and Russell, *Numerical Solution of Boundary Value +Problems for ODEs* (SIAM, 1995). Same collocation residuals, Jacobian +blocks, Newton constants, Lobatto quadrature weights, refinement thresholds, +and status semantics; regression tests cross-check meshes, iterates, and +residuals against scipy run with analytic Jacobians. + +Deliberate deviations: AD local Jacobians replace `fun_jac`/`bc_jac` and +the finite-difference estimators; `fun`/`bc` are pointwise; scipy's unknown +parameters are `z` and the differentiable parameters `p`; outputs are +padded to a static `max_nodes` (default 128, not 1000); statuses are data +rather than exceptions and there is no `verbose`; the `tol` floor is +silent; real dtypes only (split complex problems into real and imaginary +parts); no dense-output object — reuse `hermite_interpolate`; extrapolation +clamps. Node removal is not implemented, as in scipy. diff --git a/docs/dae.md b/docs/dae.md index ea5c966..00bd736 100644 --- a/docs/dae.md +++ b/docs/dae.md @@ -256,6 +256,7 @@ Only the internally constructed constant block mass matrix `diag(I_y, 0_z)` is supported; there is no general mass-matrix or fully implicit residual API. Rodas5P uses dense Jacobians and dense pivoted LU. Higher-index constraints and automatic index reduction are unsupported. This -is an initial-value solver: it does not determine unknown initial costates -or solve boundary-value problems, and jumps between multiple root branches +is an initial-value solver: it does not determine unknown initial costates — +for two-point boundary conditions see +[Boundary Value Problems](bvp.md) — and jumps between multiple root branches are not differentiable. diff --git a/docs/llms.txt b/docs/llms.txt index 72362da..55bef4f 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1,6 +1,6 @@ # tinydiffeq -> Tiny differentiable 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 only. max_steps is an attempt budget, Solution.num_steps records actual attempts, and fixed-step times are arithmetic and budget-invariant. SaveAt picks the endpoint, a fixed interpolation grid (exact=True gathers aligned knots for explicit constant-step ODEs), or the padded accepted-step prefix. Adaptive AD freezes the internal mesh (stop-gradiented controllers) and omits mesh motion. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, and Steinebach's stiff-aware fourth-order dense output, following SciML's OrdinaryDiffEqRosenbrock. RK4/Tsit5 DAEs delegate stage roots and implicit derivatives to nlls-gram; LMRootSolver requires residual-only stopping (gtol=xtol=0) and accepts only CONVERGED roots with Euclidean residual norm below the root atol, with previous-root or opt-in secant stage predictors. solve_sde integrates diagonal-noise Ito SDEs with EulerMaruyama (strong order 0.5), Milstein (1.0, commutative diagonal), or SRA1 (1.5, additive); each solver declares its per-step noise via sample_noise, drawn from a key or passed explicitly through noise= as validated, differentiable data. solve_semi_explicit_sdae applies EulerMaruyama or SRA1 to the reduced index-1 stochastic system with root-restored consistency; SaveAt(ts=...) raises for SDEs/SDAEs because interpolating rough paths is wrong. Fields and stochastic drifts may return (value, saved_aux); DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary pytrees of same-dtype real floating arrays, and the library never sets jax_enable_x64. Finite-state DTMC/CTMC sampling is primal-only; 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. +> Tiny differentiable 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 only. max_steps is an attempt budget, Solution.num_steps records actual attempts, and fixed-step times are arithmetic and budget-invariant. SaveAt picks the endpoint, a fixed interpolation grid (exact=True gathers aligned knots for explicit constant-step ODEs), or the padded accepted-step prefix. Adaptive AD freezes the internal mesh (stop-gradiented controllers) and omits mesh motion. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, and Steinebach's stiff-aware fourth-order dense output, following SciML's OrdinaryDiffEqRosenbrock. RK4/Tsit5 DAEs delegate stage roots and implicit derivatives to nlls-gram; LMRootSolver requires residual-only stopping (gtol=xtol=0) and accepts only CONVERGED roots with Euclidean residual norm below the root atol, with previous-root or opt-in secant stage predictors. solve_sde integrates diagonal-noise Ito SDEs with EulerMaruyama (strong order 0.5), Milstein (1.0, commutative diagonal), or SRA1 (1.5, additive); each solver declares its per-step noise via sample_noise, drawn from a key or passed explicitly through noise= as validated, differentiable data. solve_semi_explicit_sdae applies EulerMaruyama or SRA1 to the reduced index-1 stochastic system with root-restored consistency; SaveAt(ts=...) raises for SDEs/SDAEs because interpolating rough paths is wrong. solve_bvp is a faithful port of scipy's collocation solve_bvp (4th-order Lobatto IIIA, damped Newton, residual-controlled insert-1/insert-2 mesh refinement, optional singular term S y/(t-a)) with pointwise fun(t, y, z, args, p) and bc(ya, yb, z, args, p): z are scipy's unknown parameters solved from a z_0 guess, p the differentiable inputs, outputs padded to a static max_nodes (default 128) with num_nodes marking the active prefix, scipy's 0/1/2/3 statuses as data, AD-computed local Jacobians replacing fun_jac/bc_jac and finite differences, one custom_jvp implicit-function-theorem rule at the solution (reverse mode by transposition; failed solves and all guesses get exact-zero tangents), and dense output via hermite_interpolate/hermite_derivative which evaluate exactly scipy's returned cubic spline. Fields and stochastic drifts may return (value, saved_aux); DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary pytrees of same-dtype real floating arrays, and the library never sets jax_enable_x64. Finite-state DTMC/CTMC sampling is primal-only; 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 @@ -9,6 +9,7 @@ - [SDEs — EulerMaruyama/Milstein/SRA1 orders and noise contracts, explicit noise= as differentiable data, fixed-key semantics, shared-path convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/) - [Semi-explicit DAEs — index-1 contract, nlls-gram LM roots and predictors, implicit AD, aux contracts, Rodas5P mass-matrix path, dense output, failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) - [Semi-explicit SDAEs — reduced-SDE EulerMaruyama and SRA1, algebraic roots, aux, convergence assumptions, pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) +- [Boundary value problems — scipy solve_bvp port, z unknowns vs differentiable p, static max_nodes padding, statuses, implicit AD contract, dense output via hermite_interpolate, credit and deviations](https://highdimensionaleconlab.github.io/tinydiffeq/bvp/) - [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/) - [Linear exponential solves — dense expm and matrix-free Arnoldi/Krylov actions for fixed linear array or pytree operators, with traced and hand-coded initial-state JVP/VJP](https://highdimensionaleconlab.github.io/tinydiffeq/exponential/) - [API reference — solve functions, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/) diff --git a/mkdocs.yml b/mkdocs.yml index 11db43a..34d293e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -47,6 +47,7 @@ nav: - SDEs: sde.md - Semi-Explicit DAEs: dae.md - Semi-Explicit SDAEs: sdae.md + - Boundary Value Problems: bvp.md - Markov Chains: markov_chains.md - Linear Exponential Solves: exponential.md - API Reference: api.md diff --git a/src/tinydiffeq/__init__.py b/src/tinydiffeq/__init__.py index c347abc..097cb36 100644 --- a/src/tinydiffeq/__init__.py +++ b/src/tinydiffeq/__init__.py @@ -6,14 +6,17 @@ diagonal-noise Ito SDEs (EulerMaruyama, Milstein, SRA1) from a PRNG key or an explicit, differentiable noise pytree. solve_semi_explicit_dae and solve_semi_explicit_sdae handle index-1 systems, delegating algebraic roots -and their implicit derivatives to nlls-gram. solve_linear_ode applies dense or -matrix-free Krylov exponential actions to fixed homogeneous linear systems, -and the Markov tools simulate and forecast finite-state chains. States may be -arrays or pytrees of same-dtype real floating arrays. Fully implicit solvers, -general mass matrices, events, continuous interpolation objects, and adjoint -methods are non-goals. +and their implicit derivatives to nlls-gram. solve_bvp ports scipy's +collocation two-point boundary-value solver with implicit differentiation of +the solution (and any unknown parameters) with respect to known parameters. +solve_linear_ode applies dense or matrix-free Krylov exponential actions to +fixed homogeneous linear systems, and the Markov tools simulate and forecast +finite-state chains. States may be arrays or pytrees of same-dtype real +floating arrays. Fully implicit solvers, general mass matrices, events, +continuous interpolation objects, and adjoint methods are non-goals. """ +from tinydiffeq.bvp import solve_bvp from tinydiffeq.controllers import ConstantStepSize, IController, PIController from tinydiffeq.dae import LMRootSolver, solve_semi_explicit_dae from tinydiffeq.exponential import ( @@ -24,7 +27,7 @@ solve_linear_ode, vjp_linear_ode, ) -from tinydiffeq.interpolation import hermite_interpolate +from tinydiffeq.interpolation import hermite_derivative, hermite_interpolate from tinydiffeq.markov import ( AssociativeMarkov, ContinuousTimeMarkovChain, @@ -43,7 +46,7 @@ from tinydiffeq.save_at import SaveAt from tinydiffeq.sdae import solve_semi_explicit_sdae from tinydiffeq.sde import solve_sde -from tinydiffeq.solution import DAESolution, Solution +from tinydiffeq.solution import BVPSolution, DAESolution, Solution from tinydiffeq.solvers import ( RK4, SRA1, @@ -57,6 +60,7 @@ __all__ = [ "solve_ode", + "solve_bvp", "solve_semi_explicit_dae", "solve_sde", "solve_semi_explicit_sdae", @@ -91,7 +95,9 @@ "SaveAt", "Solution", "DAESolution", + "BVPSolution", "LMRootSolver", "hermite_interpolate", + "hermite_derivative", "cumulative_trapezoid", ] diff --git a/src/tinydiffeq/_bvp_core.py b/src/tinydiffeq/_bvp_core.py new file mode 100644 index 0000000..c3bed57 --- /dev/null +++ b/src/tinydiffeq/_bvp_core.py @@ -0,0 +1,121 @@ +"""Pure numerics for the collocation BVP solver, ported from scipy _bvp.py.""" + +import jax +import jax.numpy as jnp + +# Constants from scipy.integrate._bvp (v1.18.0), same names where they exist. +MAX_ITERATION = 10 +MAX_NEWTON_ITERATIONS = 8 +MAX_NJEV = 4 +SIGMA = 0.2 +TAU = 0.5 +N_TRIAL = 4 +TOL_R_FACTOR = 2.0 / 3.0 * 5e-2 +REFINE_FACTOR = 100.0 +TOL_FLOOR_FACTOR = 100.0 +# 5-point Lobatto quadrature on [0, 1]: interior nodes 0.5 +- sqrt(3/7)/2. +LOBATTO_OFFSET = 0.5 * (3.0 / 7.0) ** 0.5 +LOBATTO_WEIGHT_MIDDLE = 32.0 / 45.0 +LOBATTO_WEIGHT_SIDE = 49.0 / 90.0 + +RUNNING = -1 +STATUS_CONVERGED = 0 +STATUS_MAX_NODES = 1 +STATUS_SINGULAR = 2 +STATUS_BC_TOL = 3 + + +def collocation_jacobian_blocks( + h, + df_dy, + df_dy_middle, + df_dz, + df_dz_middle, + dbc_dya, + dbc_dyb, + dbc_dz, +): + # The six BABD blocks of the global Jacobian, scipy's formulas. On an + # inactive padded interval h == 0 exactly, so the staircase blocks reduce + # to -I and +I: a copy chain that propagates the last active node to the + # static last column block where dbc_dyb sits. The padded system is an + # exact algebraic embedding of the active one. + n = df_dy.shape[-1] + dtype = df_dy.dtype + hb = h[:, None, None] + eye = jnp.eye(n, dtype=dtype) + dphi_dy_0 = ( + -eye + - hb / 6.0 * (df_dy[:-1] + 2.0 * df_dy_middle) + - hb**2 / 12.0 * (df_dy_middle @ df_dy[:-1]) + ) + dphi_dy_1 = ( + eye + - hb / 6.0 * (df_dy[1:] + 2.0 * df_dy_middle) + + hb**2 / 12.0 * (df_dy_middle @ df_dy[1:]) + ) + if df_dz is None: + dphi_dz = None + else: + correction = df_dy_middle @ (df_dz[:-1] - df_dz[1:]) + dphi_dz = ( + -hb + / 6.0 + * (df_dz[:-1] + df_dz[1:] + 4.0 * (df_dz_middle + 0.125 * hb * correction)) + ) + return dict( + block_left=dphi_dy_0, + block_right=dphi_dy_1, + block_border=dphi_dz, + boundary_first=dbc_dya, + boundary_last=dbc_dyb, + boundary_border=dbc_dz, + ) + + +def pad_tail(values, num_active): + # Repeat row num_active - 1 over the inactive tail, bitwise. The tail + # invariant lets the boundary condition read the static last row and + # keeps duplicate-knot interpolation exact at the right endpoint. + last = jax.lax.dynamic_index_in_dim(values, num_active - 1, 0, keepdims=False) + mask = jnp.arange(values.shape[0]) < num_active + return jnp.where(mask.reshape((-1,) + (1,) * (values.ndim - 1)), values, last) + + +def hermite_pair(tau, h, h_safe, y_left, y_right, f_left, f_right): + # Cubic Hermite value and derivative at fixed local coordinate tau — + # the same C1 spline scipy builds in create_spline. The derivative + # divides by h_safe so inactive zero-width intervals stay finite (their + # rows are masked by the caller). + h00 = (1.0 + 2.0 * tau) * (1.0 - tau) ** 2 + h10 = tau * (1.0 - tau) ** 2 + h01 = tau**2 * (3.0 - 2.0 * tau) + h11 = tau**2 * (tau - 1.0) + d00 = 6.0 * tau**2 - 6.0 * tau + d10 = 3.0 * tau**2 - 4.0 * tau + 1.0 + d11 = 3.0 * tau**2 - 2.0 * tau + hb = h[:, None] + value = h00 * y_left + h10 * hb * f_left + h01 * y_right + h11 * hb * f_right + # d01 = -d00 exactly, so equal endpoint rows cancel bitwise on the tail. + slope = d00 * (y_left - y_right) / h_safe[:, None] + return value, slope + d10 * f_left + d11 * f_right + + +def refined_mesh(t, num_nodes, insert_1, insert_2): + # Sorted refined mesh, padded back to static length with t_b. Inactive + # candidate slots sort to the end as +inf and are replaced by the right + # endpoint, so the first num_nodes + nodes_added entries are exactly + # scipy's modify_mesh output. + max_nodes = t.shape[0] + dtype = t.dtype + inf = jnp.asarray(jnp.inf, dtype) + keep = jnp.where(jnp.arange(max_nodes) < num_nodes, t, inf) + middles = jnp.where(insert_1, 0.5 * (t[:-1] + t[1:]), inf) + thirds_left = jnp.where(insert_2, (2.0 * t[:-1] + t[1:]) / 3.0, inf) + thirds_right = jnp.where(insert_2, (t[:-1] + 2.0 * t[1:]) / 3.0, inf) + candidates = jnp.sort(jnp.concatenate([keep, middles, thirds_left, thirds_right])) + nodes_added = jnp.sum(insert_1) + 2 * jnp.sum(insert_2) + num_new = num_nodes + nodes_added.astype(jnp.int32) + t_new = candidates[:max_nodes] + t_new = jnp.where(jnp.arange(max_nodes) < num_new, t_new, t[-1]) + return t_new, num_new diff --git a/src/tinydiffeq/bvp.py b/src/tinydiffeq/bvp.py new file mode 100644 index 0000000..254034e --- /dev/null +++ b/src/tinydiffeq/bvp.py @@ -0,0 +1,936 @@ +"""Collocation boundary value solver ported from scipy.integrate._bvp.""" + +import inspect +from dataclasses import dataclass, field +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +from jax.flatten_util import ravel_pytree + +from tinydiffeq._aux import resolve_field_aux, split_field_output +from tinydiffeq._bvp_core import ( + LOBATTO_OFFSET, + LOBATTO_WEIGHT_MIDDLE, + LOBATTO_WEIGHT_SIDE, + MAX_ITERATION, + MAX_NEWTON_ITERATIONS, + MAX_NJEV, + N_TRIAL, + REFINE_FACTOR, + RUNNING, + SIGMA, + STATUS_BC_TOL, + STATUS_CONVERGED, + STATUS_MAX_NODES, + STATUS_SINGULAR, + TAU, + TOL_FLOOR_FACTOR, + TOL_R_FACTOR, + collocation_jacobian_blocks, + hermite_pair, + pad_tail, + refined_mesh, +) +from tinydiffeq._tree import asarray_state, assert_same_structure, zero_tangent +from tinydiffeq._unvmap import unvmap_all +from tinydiffeq.babd import ( + babd_dense, + babd_matvec, + structured_qr_factor, + structured_qr_solve, + structured_qr_transpose_solve, +) +from tinydiffeq.interpolation import hermite_interpolate +from tinydiffeq.solution import BVPSolution + + +def bvp_arity(f, name, forms): + # Count positional parameters; *args or uninspectable means the full form. + try: + signature = inspect.signature(f) + except (TypeError, ValueError): + return 5 + arity = 0 + for parameter in signature.parameters.values(): + if parameter.kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + arity += 1 + elif parameter.kind == inspect.Parameter.VAR_POSITIONAL: + return 5 + if arity < 2 or arity > 5: + raise ValueError(f"{name} must take 2 to 5 positional arguments: {forms}") + return arity + + +def canonicalize_bvp_fun(f, arity): + if arity == 2: + return lambda t, y, z, args, p: f(t, y) + if arity == 3: + return lambda t, y, z, args, p: f(t, y, z) + if arity == 4: + return lambda t, y, z, args, p: f(t, y, z, args) + return f + + +def canonicalize_bvp_bc(f, arity): + if arity == 2: + return lambda ya, yb, z, args, p: f(ya, yb) + if arity == 3: + return lambda ya, yb, z, args, p: f(ya, yb, z) + if arity == 4: + return lambda ya, yb, z, args, p: f(ya, yb, z, args) + return f + + +@dataclass(frozen=True) +class _BVPConfig: + fun: Any + bc: Any + fun_arity: int + bc_arity: int + max_nodes: int + n: int + k: int + has_aux: bool + fun_jac_ad: str + bc_jac_ad: str + has_singular_term: bool + y_treedef: Any + y_leaf_specs: tuple + z_treedef: Any + z_leaf_specs: Any + # Derived per call and fully determined by the compared fields; excluding + # them keeps equal configurations sharing one jit compilation. + fun_canon: Any = field(compare=False, repr=False) + bc_canon: Any = field(compare=False, repr=False) + unravel_y: Any = field(compare=False, repr=False) + unravel_z: Any = field(compare=False, repr=False) + + +def resolve_jacobian_mode(mode, num_inputs, num_outputs): + # "auto" seeds the small side: forward when tall or square, reverse when + # strictly fat (the nlls-gram convention). + if mode != "auto": + return mode + return "jvp" if num_inputs <= num_outputs else "vjp" + + +def jac_transform(mode): + return jax.jacfwd if mode == "jvp" else jax.jacrev + + +def build_system(cfg, args, S, t_a): + """Closures evaluating the wrapped field, bc, and Jacobians on flat rows. + + Shared by the primal loops and the implicit-differentiation rule so the + Newton Jacobian and the AD Jacobian are the same assembly. + """ + n, k = cfg.n, cfg.k + if cfg.has_singular_term: + S = jax.lax.stop_gradient(S) + eye = jnp.eye(n, dtype=S.dtype) + # rtol pins numpy's pinv cutoff so rank decisions agree with scipy. + B = eye - jnp.linalg.pinv(S, rtol=1e-15) @ S + D = jnp.linalg.pinv(eye - S, rtol=1e-15) + else: + B = D = None + + def raw_node_field(t_j, y_flat, z_flat, p): + y = cfg.unravel_y(y_flat) + out = cfg.fun_canon(t_j, y, cfg.unravel_z(z_flat), args, p) + value, _ = split_field_output(out, cfg.has_aux) + value, value_dtype = asarray_state(value, "fun(t, y, ...)") + assert_same_structure(y, value, "fun(t, y, ...)") + if value_dtype != y_flat.dtype: + raise TypeError("fun(t, y, ...) must preserve the state dtype") + return ravel_pytree(value)[0] + + def node_field(t_j, y_flat, z_flat, p, is_left): + value = raw_node_field(t_j, y_flat, z_flat, p) + if not cfg.has_singular_term: + return value + # double-where: the left-endpoint denominator never reaches the ratio. + denominator = jnp.where(is_left, jnp.ones_like(t_j), t_j - t_a) + interior = value + (S @ y_flat) / denominator + return jnp.where(is_left, D @ value, interior) + + def left_mask(size, at_left): + if at_left: + return jnp.arange(size) == 0 + return jnp.zeros(size, bool) + + def field_values(t_vec, y_rows, z_flat, p, at_left): + is_left = left_mask(t_vec.shape[0], at_left) + return jax.vmap(node_field, in_axes=(0, 0, None, None, 0))( + t_vec, y_rows, z_flat, p, is_left + ) + + def collocation_parts(t, h, Y, Z, p): + f = field_values(t, Y, Z, p, True) + y_middle = 0.5 * (Y[1:] + Y[:-1]) - 0.125 * h[:, None] * (f[1:] - f[:-1]) + f_middle = field_values(t[:-1] + 0.5 * h, y_middle, Z, p, False) + col_res = Y[1:] - Y[:-1] - h[:, None] / 6.0 * (f[:-1] + f[1:] + 4.0 * f_middle) + return col_res, y_middle, f, f_middle + + def bc_values(ya_flat, yb_flat, z_flat, p): + residual = cfg.bc_canon( + cfg.unravel_y(ya_flat), + cfg.unravel_y(yb_flat), + cfg.unravel_z(z_flat), + args, + p, + ) + residual = jnp.asarray(residual) + if residual.shape != (n + k,): + raise ValueError(f"bc must return {n + k} residuals, got {residual.shape}") + if residual.dtype != ya_flat.dtype: + raise TypeError("bc(ya, yb, ...) must preserve the state dtype") + return residual + + mode_fy = resolve_jacobian_mode(cfg.fun_jac_ad, n, n) + mode_fz = resolve_jacobian_mode(cfg.fun_jac_ad, k, n) + # bc is differentiated jointly with respect to (ya, yb, z). + mode_bc = resolve_jacobian_mode(cfg.bc_jac_ad, 2 * n + k, n + k) + + def field_jacobians(t_vec, y_rows, z_flat, p, at_left): + is_left = left_mask(t_vec.shape[0], at_left) + if k == 0: + + def single(t_j, y_flat, il): + def fn(yy): + return node_field(t_j, yy, z_flat, p, il) + + return jac_transform(mode_fy)(fn)(y_flat) + + return jax.vmap(single)(t_vec, y_rows, is_left), None + if mode_fy == mode_fz: + + def single(t_j, y_flat, il): + def fn(yy, zz): + return node_field(t_j, yy, zz, p, il) + + return jac_transform(mode_fy)(fn, argnums=(0, 1))(y_flat, z_flat) + + return jax.vmap(single)(t_vec, y_rows, is_left) + + def single_split(t_j, y_flat, il): + def fn(yy, zz): + return node_field(t_j, yy, zz, p, il) + + df_dy = jac_transform(mode_fy)(fn, argnums=0)(y_flat, z_flat) + df_dz = jac_transform(mode_fz)(fn, argnums=1)(y_flat, z_flat) + return df_dy, df_dz + + return jax.vmap(single_split)(t_vec, y_rows, is_left) + + def bc_jacobians(ya_flat, yb_flat, z_flat, p): + if k == 0: + + def fn(a, b): + return bc_values(a, b, z_flat, p) + + dbc_dya, dbc_dyb = jac_transform(mode_bc)(fn, argnums=(0, 1))( + ya_flat, yb_flat + ) + return dbc_dya, dbc_dyb, None + + def fn(a, b, zz): + return bc_values(a, b, zz, p) + + return jac_transform(mode_bc)(fn, argnums=(0, 1, 2))(ya_flat, yb_flat, z_flat) + + def jacobian_blocks(t, h, Y, Z, y_middle, p): + df_dy, df_dz = field_jacobians(t, Y, Z, p, True) + df_dy_middle, df_dz_middle = field_jacobians( + t[:-1] + 0.5 * h, y_middle, Z, p, False + ) + dbc_dya, dbc_dyb, dbc_dz = bc_jacobians(Y[0], Y[-1], Z, p) + return collocation_jacobian_blocks( + h, + df_dy, + df_dy_middle, + df_dz, + df_dz_middle, + dbc_dya, + dbc_dyb, + dbc_dz, + ) + + def jacobian_at(t, h, Y, Z, y_middle, p): + return babd_dense(jacobian_blocks(t, h, Y, Z, y_middle, p)) + + def project(Y): + if not cfg.has_singular_term: + return Y + return Y.at[0].set(B @ Y[0]) + + def aux_values(t, Y, Z, p): + def node_aux(t_j, y_flat): + out = cfg.fun_canon(t_j, cfg.unravel_y(y_flat), cfg.unravel_z(Z), args, p) + return split_field_output(out, True)[1] + + return jax.vmap(node_aux)(t, Y) + + return dict( + field_values=field_values, + collocation_parts=collocation_parts, + bc_values=bc_values, + jacobian_blocks=jacobian_blocks, + jacobian_at=jacobian_at, + project=project, + aux_values=aux_values, + ) + + +def estimate_rms_residuals(system, t, h, h_safe, active, Y, Z, f, f_middle, col_res, p): + # 5-point Lobatto quadrature of the relative residual over each interval; + # endpoint terms vanish because the spline collocates at the nodes. + r_middle = 1.5 * col_res / h_safe[:, None] + tau_right = 0.5 + LOBATTO_OFFSET + tau_left = 0.5 - LOBATTO_OFFSET + y1, y1_prime = hermite_pair(tau_right, h, h_safe, Y[:-1], Y[1:], f[:-1], f[1:]) + y2, y2_prime = hermite_pair(tau_left, h, h_safe, Y[:-1], Y[1:], f[:-1], f[1:]) + f1 = system["field_values"](t[:-1] + tau_right * h, y1, Z, p, False) + f2 = system["field_values"](t[:-1] + tau_left * h, y2, Z, p, False) + r1 = (y1_prime - f1) / (1.0 + jnp.abs(f1)) + r2 = (y2_prime - f2) / (1.0 + jnp.abs(f2)) + r_middle = r_middle / (1.0 + jnp.abs(f_middle)) + rms = jnp.sqrt( + 0.5 + * ( + LOBATTO_WEIGHT_MIDDLE * jnp.sum(r_middle * r_middle, axis=1) + + LOBATTO_WEIGHT_SIDE + * (jnp.sum(r1 * r1, axis=1) + jnp.sum(r2 * r2, axis=1)) + ) + ) + # The padded cubic has derivative f_b / 7 at the Lobatto points, not zero, + # so inactive intervals must be masked rather than trusted to vanish. + return jnp.where(active, rms, 0.0) + + +def solve_newton(cfg, system, t, h, Y, Z, num_nodes, p, tol, bc_tol, live): + max_nodes, n, k = cfg.max_nodes, cfg.n, cfg.k + size = max_nodes * n + k + dtype = Y.dtype + active = jnp.arange(max_nodes - 1) < num_nodes - 1 + tol_r = TOL_R_FACTOR * h * tol + # The copy chain duplicates the last active node's step over the tail, so + # the affine-invariant cost must count active variables only. + active_variable = jnp.concatenate( + [jnp.repeat(jnp.arange(max_nodes) < num_nodes, n), jnp.ones(k, bool)] + ) + + def masked_cost(step): + return jnp.sum(jnp.where(active_variable, step * step, 0.0)) + + def stack_residual(col_res, bc_res): + return jnp.concatenate([col_res.reshape(-1), bc_res]) + + col_res, y_middle, _, f_middle = system["collocation_parts"](t, h, Y, Z, p) + bc_res = system["bc_values"](Y[0], Y[-1], Z, p) + # The carried solver state is whatever pytree the linear solver's init + # returns; its structure is recovered abstractly for the placeholder. + state_shapes = jax.eval_shape( + lambda tt, hh, YY, ZZ, ym, pp: structured_qr_factor( + system["jacobian_blocks"](tt, hh, YY, ZZ, ym, pp) + )[0], + t, + h, + Y, + Z, + y_middle, + p, + ) + carry = dict( + Y=Y, + Z=Z, + y_middle=y_middle, + col_res=col_res, + f_middle=f_middle, + bc_res=bc_res, + res=stack_residual(col_res, bc_res), + state=jax.tree.map(lambda s: jnp.zeros(s.shape, s.dtype), state_shapes), + step=jnp.zeros(size, dtype), + cost=jnp.zeros((), dtype), + recompute=jnp.asarray(True), + njev=jnp.zeros((), jnp.int32), + iteration=jnp.zeros((), jnp.int32), + singular=jnp.asarray(False), + # A lane whose outer iteration already terminated starts stopped, so + # a batched Newton loop only runs until the live lanes finish. + stop=~live, + ) + + def newton_cond(c): + return (~c["stop"]) & (c["iteration"] < MAX_NEWTON_ITERATIONS) + + def newton_body(c): + def refactor(c): + blocks = system["jacobian_blocks"](t, h, c["Y"], c["Z"], c["y_middle"], p) + state, ok = structured_qr_factor(blocks) + step = structured_qr_solve(state, c["res"]) + return state, ok, step, masked_cost(step), c["njev"] + 1 + + def reuse(c): + return c["state"], jnp.asarray(True), c["step"], c["cost"], c["njev"] + + def per_lane(c): + return jax.lax.cond(c["recompute"], refactor, reuse, c) + + # A batched cond runs both branches as a select; the unvmap_all gate + # keeps a scalar predicate so an all-reusing batch skips the + # factorization for real (see _unvmap). Stopped lanes must not veto: + # a lane that finished on a damped step carries recompute=True + # forever, and its frozen carry would otherwise pin the gate false. + solver_state, factor_ok, step, cost, njev = jax.lax.cond( + unvmap_all(~c["recompute"] | c["stop"]), reuse, per_lane, c + ) + singular = ~factor_ok + + def trial_cond(tc): + return (~tc["accepted"]) & (tc["trial"] <= N_TRIAL) + + def trial_body(tc): + # alpha shrinks by exact multiplication, matching scipy's + # ``alpha *= tau`` (powers of two, bitwise). + alpha = tc["alpha"] + y_candidate = c["Y"] - alpha * step[: max_nodes * n].reshape(max_nodes, n) + y_candidate = system["project"](y_candidate) + y_candidate = pad_tail(y_candidate, num_nodes) + z_candidate = c["Z"] - alpha * step[max_nodes * n :] + col_res, y_middle, _, f_middle = system["collocation_parts"]( + t, h, y_candidate, z_candidate, p + ) + bc_res = system["bc_values"]( + y_candidate[0], y_candidate[-1], z_candidate, p + ) + res = stack_residual(col_res, bc_res) + step_new = structured_qr_solve(solver_state, res) + cost_new = masked_cost(step_new) + return dict( + trial=tc["trial"] + 1, + trial_used=tc["trial"], + alpha=alpha * TAU, + accepted=cost_new < (1.0 - 2.0 * alpha * SIGMA) * cost, + Y=y_candidate, + Z=z_candidate, + y_middle=y_middle, + col_res=col_res, + f_middle=f_middle, + bc_res=bc_res, + res=res, + step_new=step_new, + cost_new=cost_new, + ) + + trial = jax.lax.while_loop( + trial_cond, + trial_body, + dict( + trial=jnp.zeros((), jnp.int32), + trial_used=jnp.zeros((), jnp.int32), + alpha=jnp.ones((), dtype), + # A singular factorization takes no trials, as scipy's break. + accepted=singular, + Y=c["Y"], + Z=c["Z"], + y_middle=c["y_middle"], + col_res=c["col_res"], + f_middle=c["f_middle"], + bc_res=c["bc_res"], + res=c["res"], + step_new=step, + cost_new=cost, + ), + ) + + # A singular factorization stops before taking any step, as scipy does. + def keep_current(tr): + return ( + c["Y"], + c["Z"], + c["y_middle"], + c["col_res"], + c["f_middle"], + c["bc_res"], + c["res"], + ) + + def take_candidate(tr): + return ( + tr["Y"], + tr["Z"], + tr["y_middle"], + tr["col_res"], + tr["f_middle"], + tr["bc_res"], + tr["res"], + ) + + Y_next, Z_next, y_middle, col_res, f_middle, bc_res, res = jax.lax.cond( + singular, keep_current, take_candidate, trial + ) + + # Inactive intervals have tol_r == 0 == col_res, so mask them true. + col_ok = jnp.all( + jnp.where( + active[:, None], + jnp.abs(col_res) < tol_r[:, None] * (1.0 + jnp.abs(f_middle)), + True, + ) + ) + bc_ok = jnp.all(jnp.abs(bc_res) < bc_tol) + stop = singular | (njev == MAX_NJEV) | (col_ok & bc_ok) + # A full step keeps the frozen Jacobian; a damped one forces refresh. + full_step = trial["accepted"] & (trial["trial_used"] == 0) + return dict( + Y=Y_next, + Z=Z_next, + y_middle=y_middle, + col_res=col_res, + f_middle=f_middle, + bc_res=bc_res, + res=res, + state=solver_state, + step=jnp.where(full_step, trial["step_new"], step), + cost=jnp.where(full_step, trial["cost_new"], cost), + recompute=~full_step, + njev=njev, + iteration=c["iteration"] + 1, + singular=singular, + stop=stop, + ) + + final = jax.lax.while_loop(newton_cond, newton_body, carry) + return final["Y"], final["Z"], final["singular"] + + +def _solve_bvp_impl(cfg, t, Y, Z, num_nodes, p, args, S, tol, bc_tol): + max_nodes, n = cfg.max_nodes, cfg.n + dtype = Y.dtype + tol = jnp.maximum(jnp.asarray(tol, dtype), TOL_FLOOR_FACTOR * jnp.finfo(dtype).eps) + bc_tol = jnp.asarray(bc_tol, dtype) + bc_tol = jnp.where(jnp.isnan(bc_tol), tol, bc_tol) + system = build_system(cfg, args, S, t[0]) + Y = system["project"](Y) + + def outer_cond(c): + return (c["status"] == RUNNING) & (c["iteration"] < max_nodes + MAX_ITERATION) + + def outer_body(c): + t_c, num_c = c["t"], c["num_nodes"] + h = t_c[1:] - t_c[:-1] + active = jnp.arange(max_nodes - 1) < num_c - 1 + h_safe = jnp.where(active, h, jnp.ones_like(h)) + Y_c, Z_c, singular = solve_newton( + cfg, + system, + t_c, + h, + c["Y"], + c["Z"], + num_c, + p, + tol, + bc_tol, + c["status"] == RUNNING, + ) + iteration = c["iteration"] + 1 + col_res, _, f, f_middle = system["collocation_parts"](t_c, h, Y_c, Z_c, p) + bc_res = system["bc_values"](Y_c[0], Y_c[-1], Z_c, p) + max_bc_res = jnp.max(jnp.abs(bc_res)) + rms = estimate_rms_residuals( + system, t_c, h, h_safe, active, Y_c, Z_c, f, f_middle, col_res, p + ) + insert_1 = active & (rms > tol) & (rms < REFINE_FACTOR * tol) + insert_2 = active & (rms >= REFINE_FACTOR * tol) + nodes_added = (jnp.sum(insert_1) + 2 * jnp.sum(insert_2)).astype(jnp.int32) + # The max_nodes test precedes any modification: a status-1 result + # reports the unmodified mesh, exactly as scipy does. + overflow = num_c + nodes_added > max_nodes + live = ~singular & ~overflow + refine = live & (nodes_added > 0) + converged = live & (nodes_added == 0) & (max_bc_res <= bc_tol) + # ~(<=) rather than (>): a NaN boundary residual must follow scipy's + # elif chain to status 3 instead of looping to the safety bound. + stalled = ( + live + & (nodes_added == 0) + & ~(max_bc_res <= bc_tol) + & (iteration >= MAX_ITERATION) + ) + status = jnp.where( + singular, + STATUS_SINGULAR, + jnp.where( + overflow, + STATUS_MAX_NODES, + jnp.where( + converged, + STATUS_CONVERGED, + jnp.where(stalled, STATUS_BC_TOL, RUNNING), + ), + ), + ).astype(jnp.int32) + # Refinement is computed unconditionally and selected: a batched cond + # would run both branches under vmap anyway, and the sort is cheap. + t_refined, num_refined = refined_mesh(t_c, num_c, insert_1, insert_2) + y_refined = pad_tail(hermite_interpolate(t_refined, t_c, Y_c, f), num_refined) + return dict( + t=jnp.where(refine, t_refined, t_c), + Y=jnp.where(refine, y_refined, Y_c), + Z=Z_c, + num_nodes=jnp.where(refine, num_refined, num_c), + iteration=iteration, + status=status, + f=f, + rms=rms, + ) + + final = jax.lax.while_loop( + outer_cond, + outer_body, + dict( + t=t, + Y=Y, + Z=Z, + num_nodes=num_nodes, + iteration=jnp.zeros((), jnp.int32), + status=jnp.asarray(RUNNING, jnp.int32), + f=jnp.zeros((max_nodes, n), dtype), + rms=jnp.zeros(max_nodes - 1, dtype), + ), + ) + aux = ( + system["aux_values"](final["t"], final["Y"], final["Z"], p) + if cfg.has_aux + else None + ) + return ( + final["t"], + final["Y"], + final["Z"], + final["f"], + final["rms"], + final["num_nodes"], + final["iteration"], + final["status"], + aux, + ) + + +_solve_bvp_jit = jax.jit(_solve_bvp_impl, static_argnums=(0,)) + + +def mask_tangent(condition, tangent, zero): + def mask(value, zero_value): + if getattr(value, "dtype", None) == jax.dtypes.float0: + return value + return jnp.where(condition, value, zero_value) + + return jax.tree.map(mask, tangent, zero) + + +def _run_with_ad(cfg, t, Y, Z, num_nodes, p, args, S, tol, bc_tol): + def primal(t, Y, Z, num_nodes, p, args, S, tol, bc_tol): + return _solve_bvp_jit(cfg, t, Y, Z, num_nodes, p, args, S, tol, bc_tol) + + run = jax.custom_jvp(primal) + + @run.defjvp + def run_jvp(primals, tangents): + result = run(*primals) + zeros = zero_tangent(result) + t_in, Y_in, Z_in, _, p_in, args_in, S_in, _, _ = primals + p_dot = tangents[4] + if p_in is None: + return result, zeros + t_f, Y_f, Z_f, _, _, num_f, _, status, _ = result + + # Implicit function theorem at the solution on the frozen final mesh; + # a failed lane's tangent program runs at the inert initial guess so + # batched JVPs and transposed VJPs stay finite, then masks to zero. + # The solution, p, and the Jacobian stay differentiable: an outer + # transform of this rule recurses through the same custom_jvp, so + # higher-order derivatives (hessians, reverse-over-forward) are exact + # on the frozen mesh. Only the mesh and the failed-lane substitutes + # are stop-gradiented. + ad_ok = status == STATUS_CONVERGED + t_c = jax.lax.stop_gradient(t_f) + num_c = jax.lax.stop_gradient(num_f) + Y_sol = jnp.where(ad_ok, Y_f, jax.lax.stop_gradient(Y_in)) + Z_sol = jnp.where(ad_ok, Z_f, jax.lax.stop_gradient(Z_in)) + p_c = p_in + args_c = jax.lax.stop_gradient(args_in) + S_c = jax.lax.stop_gradient(S_in) if cfg.has_singular_term else None + system = build_system(cfg, args_c, S_c, t_c[0]) + h_c = t_c[1:] - t_c[:-1] + _, y_middle, _, _ = system["collocation_parts"](t_c, h_c, Y_sol, Z_sol, p_c) + blocks = system["jacobian_blocks"](t_c, h_c, Y_sol, Z_sol, y_middle, p_c) + # The factorization is solver state, not a differentiation path: + # tangents of the matrix flow through custom_linear_solve's matvec. + state, factor_ok = structured_qr_factor(jax.lax.stop_gradient(blocks)) + p_dot_masked = mask_tangent(ad_ok, p_dot, zero_tangent(p_in)) + + def residual_of_p(p_value): + col_res, _, _, _ = system["collocation_parts"]( + t_c, h_c, Y_sol, Z_sol, p_value + ) + bc_res = system["bc_values"](Y_sol[0], Y_sol[-1], Z_sol, p_value) + return jnp.concatenate([col_res.reshape(-1), bc_res]) + + rhs_dot = jax.jvp(residual_of_p, (p_c,), (p_dot_masked,))[1] + delta = jax.lax.custom_linear_solve( + lambda u: babd_matvec(blocks, u), + -rhs_dot, + lambda _, b: structured_qr_solve(state, b), + transpose_solve=lambda _, b: structured_qr_transpose_solve(state, b), + ) + y_dot = delta[: cfg.max_nodes * cfg.n].reshape(cfg.max_nodes, cfg.n) + y_dot = system["project"](y_dot) + y_dot = pad_tail(y_dot, num_c) + z_dot = delta[cfg.max_nodes * cfg.n :] + + # A failed solve has exact-zero tangents: the where-select (not a + # multiplication) kills any non-finite garbage from the inert-guess + # tangent program. A converged solve whose final-mesh Jacobian fails + # to factor has no computable derivative: scale by NaN rather than + # substituting a constant NaN, so the failure survives transposition + # into reverse mode instead of dropping to a silent zero cotangent. + def finalize(dot): + dot = jnp.where(ad_ok, dot, jnp.zeros_like(dot)) + scale = jnp.where(ad_ok & ~factor_ok, jnp.nan, 1.0) + return dot * jnp.asarray(scale, dot.dtype) + + y_dot = finalize(y_dot) + z_dot = finalize(z_dot) + f_dot = jax.jvp( + lambda yv, zv, pv: system["field_values"](t_c, yv, zv, pv, True), + (Y_sol, Z_sol, p_c), + (y_dot, z_dot, p_dot_masked), + )[1] + f_dot = finalize(f_dot) + if cfg.has_aux: + aux_dot = jax.jvp( + lambda yv, zv, pv: system["aux_values"](t_c, yv, zv, pv), + (Y_sol, Z_sol, p_c), + (y_dot, z_dot, p_dot_masked), + )[1] + aux_dot = jax.tree.map(finalize, aux_dot) + else: + aux_dot = None + tangent = ( + zeros[0], + y_dot, + z_dot, + f_dot, + zeros[4], + zeros[5], + zeros[6], + zeros[7], + aux_dot, + ) + return result, tangent + + return run(t, Y, Z, num_nodes, p, args, S, tol, bc_tol) + + +def _unravel_empty(z_flat): + return None + + +def solve_bvp( + fun, + bc, + t, + y_0, + z_0=None, + *, + p=None, + args=None, + S=None, + fun_jac_ad="auto", + bc_jac_ad="auto", + tol=1e-3, + bc_tol=None, + max_nodes=128, + has_aux=None, +): + """Solve ``dy/dt = fun(t, y, z, args, p) + S y / (t - t[0])`` with + two-point boundary conditions ``bc(y(t_a), y(t_b), z, args, p) = 0``. + + A faithful JAX port of :func:`scipy.integrate.solve_bvp` (4th-order Lobatto + IIIA collocation with residual-controlled mesh refinement and a damped + Newton method), with scipy's algorithm, constants, and default tolerances. + ``fun`` and ``bc`` are pointwise — a scalar ``t`` and one node's state + pytree — and may be declared with 2 to 5 positional arguments in the + orders above. ``z_0`` is the guess for scipy's unknown parameters (any + pytree), solved jointly with ``y`` and returned as ``sol.z``; ``bc`` must + then return ``n + size(z)`` residuals as a 1-D array. ``p`` holds known + differentiable parameters — the only AD input: JVP/VJP rules (composing + to higher order) differentiate ``sol.y``, ``sol.yp``, ``sol.z``, and + ``sol.aux`` with respect to ``p`` implicitly at the solution, never + through the iterations, and the guesses ``t``, ``y_0``, ``z_0`` (and + ``args``, ``S``) are differentiation-inert. + ``args`` is inert pass-through data. Local Jacobians come from AD instead + of scipy's finite differences; ``fun_jac_ad``/``bc_jac_ad`` choose + ``"jvp"``, ``"vjp"``, or ``"auto"`` (forward when square or tall, reverse + when strictly fat). ``max_nodes`` (static, default 128) fixes the padded + output length: the mesh ``t`` starts from the given guess and grows under + refinement, the returned tail repeats ``t[-1]`` and the last active rows, + and ``sol.num_nodes`` counts active nodes, so + ``hermite_interpolate(ts, sol.t, sol.y, sol.yp)`` evaluates exactly + scipy's returned C1 cubic spline (``hermite_derivative`` its derivative). + ``fun`` may return ``(value, aux)``; aux is evaluated once at the solution + over all padded nodes and participates in AD. ``tol`` is floored at + ``100 * eps`` of the working dtype (taken from ``y_0``), silently. + Failures never raise inside the solve: ``sol.status`` carries scipy's + codes and a singular collocation Jacobian is reported as status 2 with the + last iterate returned. The collocation system is factored once per Newton + refresh by a structured orthogonal factorization of its bordered + almost-block-diagonal form (``O(max_nodes)`` instead of the dense cubic, + where scipy uses sparse LU). The whole solve is compiled with + ``lax.while_loop`` loops, keyed on the identity of ``fun`` and ``bc`` + (reuse module-level functions rather than rebuilding closures per call); + for repeated solves call it inside an outer ``jax.jit`` so the wrapper's + per-call validation and dispatch trace away. + """ + if not isinstance(max_nodes, int) or isinstance(max_nodes, bool) or max_nodes < 2: + raise ValueError("max_nodes must be a static int of at least 2") + for name, mode in (("fun_jac_ad", fun_jac_ad), ("bc_jac_ad", bc_jac_ad)): + if mode not in ("auto", "jvp", "vjp"): + raise ValueError(f'{name} must be "auto", "jvp", or "vjp"') + + y_0, dtype = asarray_state(y_0, "y_0") + t = jnp.asarray(t, dtype) + if t.ndim != 1: + raise ValueError("t must be 1-dimensional") + m_0 = t.shape[0] + if m_0 < 2: + raise ValueError("t must contain at least two nodes") + if m_0 > max_nodes: + raise ValueError(f"t has {m_0} nodes, more than max_nodes={max_nodes}") + if not isinstance(t, jax.core.Tracer): + if np.any(np.diff(np.asarray(t)) <= 0): + raise ValueError("t must be strictly increasing") + for leaf in jax.tree.leaves(y_0): + if leaf.shape[0] != m_0: + raise ValueError("y_0 leaves must have leading axis len(t)") + + node_template = jax.tree.map(lambda leaf: leaf[0], y_0) + flat_node, unravel_y = ravel_pytree(node_template) + n = flat_node.size + Y_0 = jax.vmap(lambda node: ravel_pytree(node)[0])(y_0) + + if z_0 is None: + k = 0 + Z_0 = jnp.zeros((0,), dtype) + unravel_z = _unravel_empty + z_treedef = None + z_leaf_specs = None + else: + z_0, z_dtype = asarray_state(z_0, "z_0") + if z_dtype != dtype: + raise TypeError("z_0 must have the same dtype as y_0") + Z_0, unravel_z = ravel_pytree(z_0) + k = Z_0.size + z_treedef = jax.tree.structure(z_0) + z_leaf_specs = tuple( + (leaf.shape, str(leaf.dtype)) for leaf in jax.tree.leaves(z_0) + ) + + if S is not None: + S = jnp.asarray(S, dtype) + if S.shape != (n, n): + raise ValueError(f"S must have shape {(n, n)}, got {S.shape}") + + fun_arity = bvp_arity( + fun, "fun", "(t, y), (t, y, z), (t, y, z, args), or (t, y, z, args, p)" + ) + bc_arity = bvp_arity( + bc, "bc", "(ya, yb), (ya, yb, z), (ya, yb, z, args), or (ya, yb, z, args, p)" + ) + # Silently dropping p or args would make derivatives silently zero. + if p is not None and fun_arity < 5 and bc_arity < 5: + raise ValueError("p was passed but neither fun nor bc takes it") + if args is not None and fun_arity < 4 and bc_arity < 4: + raise ValueError("args was passed but neither fun nor bc takes it") + if z_0 is not None and fun_arity < 3 and bc_arity < 3: + raise ValueError("z_0 was passed but neither fun nor bc takes it") + fun_canon = canonicalize_bvp_fun(fun, fun_arity) + bc_canon = canonicalize_bvp_bc(bc, bc_arity) + + has_aux, _ = resolve_field_aux( + fun_canon, + (t[0], node_template, z_0, args, p), + jax.tree.structure(node_template), + has_aux, + name="has_aux", + ) + bc_shape = jax.eval_shape( + lambda *operands: jnp.asarray(bc_canon(*operands)), + node_template, + node_template, + z_0, + args, + p, + ) + if bc_shape.shape != (n + k,): + raise ValueError(f"bc must return a 1-D array of {n + k} residuals") + + cfg = _BVPConfig( + fun=fun, + bc=bc, + fun_arity=fun_arity, + bc_arity=bc_arity, + max_nodes=max_nodes, + n=n, + k=k, + has_aux=has_aux, + fun_jac_ad=fun_jac_ad, + bc_jac_ad=bc_jac_ad, + has_singular_term=S is not None, + y_treedef=jax.tree.structure(node_template), + y_leaf_specs=tuple( + (leaf.shape, str(leaf.dtype)) for leaf in jax.tree.leaves(node_template) + ), + z_treedef=z_treedef, + z_leaf_specs=z_leaf_specs, + fun_canon=fun_canon, + bc_canon=bc_canon, + unravel_y=unravel_y, + unravel_z=unravel_z, + ) + + pad = max_nodes - m_0 + t_pad = jnp.concatenate([t, jnp.full((pad,), t[-1], dtype)]) + Y_pad = jnp.concatenate([Y_0, jnp.broadcast_to(Y_0[-1], (pad, n))]) + bc_tol_value = jnp.asarray(jnp.nan if bc_tol is None else bc_tol, dtype) + t_f, Y_f, Z_f, f_f, rms, num_nodes, num_iterations, status, aux = _run_with_ad( + cfg, + t_pad, + Y_pad, + Z_0, + jnp.asarray(m_0, jnp.int32), + p, + args, + S, + jnp.asarray(tol, dtype), + bc_tol_value, + ) + return BVPSolution( + t=t_f, + y=jax.vmap(unravel_y)(Y_f), + yp=jax.vmap(unravel_y)(f_f), + z=unravel_z(Z_f), + rms_residuals=rms, + num_nodes=num_nodes, + num_iterations=num_iterations, + status=status, + ok=status == STATUS_CONVERGED, + aux=aux, + ) diff --git a/src/tinydiffeq/interpolation.py b/src/tinydiffeq/interpolation.py index 3f073d9..9871dd2 100644 --- a/src/tinydiffeq/interpolation.py +++ b/src/tinydiffeq/interpolation.py @@ -60,6 +60,44 @@ def bc(a): return jax.tree.map(interpolate_leaf, knot_xs, knot_fs) +def hermite_derivative(ts_query, knot_ts, knot_xs, knot_fs): + """Derivative of the C1 cubic Hermite interpolant at ``ts_query``. + + Uses ``searchsorted(side="left")`` so a query at a repeated right-endpoint + knot lands on the last positive-width bracket and returns the knot + derivative instead of the degenerate bracket's zero. Queries outside the + knot span return zero, the derivative of the clamped value extension. + """ + n = knot_ts.shape[0] + idx = jnp.clip(jnp.searchsorted(knot_ts, ts_query, side="left") - 1, 0, n - 2) + t_left, t_right = knot_ts[idx], knot_ts[idx + 1] + width = t_right - t_left + degenerate = width <= 0.0 + width_safe = jnp.where(degenerate, 1.0, width) + s = jnp.clip((ts_query - t_left) / width_safe, 0.0, 1.0) + outside = (ts_query < knot_ts[0]) | (ts_query > knot_ts[n - 1]) + + def derivative_leaf(xs, fs): + x_left, x_right = xs[idx], xs[idx + 1] + f_left, f_right = fs[idx], fs[idx + 1] + extra = xs.ndim - 1 + + def bc(a): + return a.reshape(a.shape + (1,) * extra) + + s_leaf = s.astype(xs.dtype) + width_leaf = width_safe.astype(xs.dtype) + s_, w_, deg_, out_ = bc(s_leaf), bc(width_leaf), bc(degenerate), bc(outside) + d00 = 6.0 * s_**2 - 6.0 * s_ + d10 = 3.0 * s_**2 - 4.0 * s_ + 1.0 + d11 = 3.0 * s_**2 - 2.0 * s_ + value = d00 * (x_left - x_right) / w_ + d10 * f_left + d11 * f_right + value = jnp.where(deg_, f_left, value) + return jnp.where(out_, jnp.zeros_like(value), value) + + return jax.tree.map(derivative_leaf, knot_xs, knot_fs) + + def rodas_interpolate(ts_query, knot_ts, knot_xs, interval_coefficients): """Evaluate the Rodas5P continuous extension over raw attempt rows. diff --git a/src/tinydiffeq/solution.py b/src/tinydiffeq/solution.py index 7a7a6dd..ac4e81d 100644 --- a/src/tinydiffeq/solution.py +++ b/src/tinydiffeq/solution.py @@ -28,6 +28,38 @@ class Solution: num_steps: jax.Array | None = None +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class BVPSolution: + """Result of ``solve_bvp``. + + Arrays are padded to the static ``max_nodes``: the ``t`` tail repeats the + right endpoint and the ``y``/``yp`` tails repeat the last active row, so + ``hermite_interpolate(ts, sol.t, sol.y, sol.yp)`` evaluates exactly the C1 + cubic spline scipy's ``solve_bvp`` returns and ``hermite_derivative`` its + derivative. ``z`` holds the solved unknown parameters (``None`` when the + problem has none), ``rms_residuals`` is zero on inactive intervals, + ``num_nodes`` counts active mesh nodes, and ``num_iterations`` is scipy's + ``niter``. ``status`` uses scipy's codes (0 converged, 1 ``max_nodes`` + exceeded, 2 singular Jacobian, 3 boundary-condition tolerance unsatisfied) + and ``ok`` is ``status == 0``; a failed status returns the last iterate, + which may be non-finite. Under AD only ``y``, ``yp``, ``z``, and ``aux`` + carry tangents with respect to ``p``; every other field is + differentiation-inert with exact-zero tangents. + """ + + t: jax.Array + y: Any + yp: Any + z: Any + rms_residuals: jax.Array + num_nodes: jax.Array + num_iterations: jax.Array + status: jax.Array + ok: jax.Array + aux: Any = None + + @jax.tree_util.register_dataclass @dataclass(frozen=True) class DAESolution: diff --git a/tests/test_bvp.py b/tests/test_bvp.py new file mode 100644 index 0000000..f56021c --- /dev/null +++ b/tests/test_bvp.py @@ -0,0 +1,623 @@ +import jax +import jax.numpy as jnp +import numpy as np +from jax.flatten_util import ravel_pytree +from scipy.integrate import solve_bvp as scipy_solve_bvp +from scipy.special import erf + +from tinydiffeq import bvp as bvp_module +from tinydiffeq import ( + hermite_derivative, + hermite_interpolate, + solve_bvp, +) +from tinydiffeq._bvp_core import refined_mesh +from tinydiffeq.babd import ( + babd_matvec, + structured_qr_factor, + structured_qr_solve, + structured_qr_transpose_solve, +) + + +def exp_fun(t, y): + return jnp.array([y[1], y[0]]) + + +def exp_bc(ya, yb): + return jnp.array([ya[0] - 1.0, yb[0]]) + + +def exp_sol(t): + return (np.exp(-t) - np.exp(t - 2.0)) / (1.0 - np.exp(-2.0)) + + +def sl_fun(t, y, z): + return jnp.array([y[1], -(z[0] ** 2) * y[0]]) + + +def sl_bc(ya, yb, z): + return jnp.array([ya[0], yb[0], ya[1] - z[0]]) + + +def emden_fun(t, y): + return jnp.array([y[1], -(y[0] ** 5)]) + + +def emden_bc(ya, yb): + return jnp.array([ya[1], yb[0] - (3.0 / 4.0) ** 0.5]) + + +EMDEN_S = jnp.array([[0.0, 0.0], [0.0, -2.0]]) + + +def emden_sol(t): + return (1.0 + t**2 / 3.0) ** -0.5 + + +def make_config(fun, bc, max_nodes, n, k, has_singular_term=False): + node = jnp.zeros(n) + _, unravel_y = ravel_pytree(node) + if k > 0: + z_node = jnp.zeros(k) + _, unravel_z = ravel_pytree(z_node) + z_treedef = jax.tree.structure(z_node) + z_leaf_specs = ((z_node.shape, str(z_node.dtype)),) + else: + unravel_z = bvp_module._unravel_empty + z_treedef = None + z_leaf_specs = None + fun_arity = bvp_module.bvp_arity(fun, "fun", "") + bc_arity = bvp_module.bvp_arity(bc, "bc", "") + return bvp_module._BVPConfig( + fun=fun, + bc=bc, + fun_arity=fun_arity, + bc_arity=bc_arity, + max_nodes=max_nodes, + n=n, + k=k, + has_aux=False, + fun_jac_ad="auto", + bc_jac_ad="auto", + has_singular_term=has_singular_term, + y_treedef=jax.tree.structure(node), + y_leaf_specs=((node.shape, str(node.dtype)),), + z_treedef=z_treedef, + z_leaf_specs=z_leaf_specs, + fun_canon=bvp_module.canonicalize_bvp_fun(fun, fun_arity), + bc_canon=bvp_module.canonicalize_bvp_bc(bc, bc_arity), + unravel_y=unravel_y, + unravel_z=unravel_z, + ) + + +def relative_residual_norm(fun_values, derivative_values): + residual = (derivative_values - fun_values) / (1.0 + np.abs(fun_values)) + return np.sum(residual**2, axis=1) ** 0.5 + + +def test_modify_mesh_matches_scipy(): + t = jnp.array([0.0, 1.0, 3.0, 9.0, 9.0, 9.0, 9.0]) + insert_1 = jnp.array([True, False, False, False, False, False]) + insert_2 = jnp.array([False, False, True, False, False, False]) + t_new, num_new = refined_mesh(t, jnp.asarray(4, jnp.int32), insert_1, insert_2) + assert int(num_new) == 7 + assert jnp.array_equal(t_new, jnp.array([0.0, 0.5, 1.0, 3.0, 5.0, 7.0, 9.0])) + + t = jnp.concatenate([jnp.array([-6.0, -3.0, 0.0, 3.0, 6.0]), jnp.full(7, 6.0)]) + insert_1 = jnp.zeros(11, bool).at[1].set(True) + insert_2 = jnp.zeros(11, bool).at[jnp.array([0, 2, 3])].set(True) + t_new, num_new = refined_mesh(t, jnp.asarray(5, jnp.int32), insert_1, insert_2) + assert int(num_new) == 12 + expected = jnp.array( + [-6.0, -5.0, -4.0, -3.0, -1.5, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + assert jnp.array_equal(t_new, expected) + + +def test_exponential_no_refinement(): + t = jnp.linspace(0.0, 1.0, 5) + sol = solve_bvp(exp_fun, exp_bc, t, jnp.zeros((5, 2)), max_nodes=32) + assert int(sol.status) == 0 + assert bool(sol.ok) + assert int(sol.num_nodes) == 5 + num = int(sol.num_nodes) + + t_test = jnp.linspace(0.0, 1.0, 100) + y_test = hermite_interpolate(t_test, sol.t, sol.y, sol.yp) + np.testing.assert_allclose( + np.asarray(y_test[:, 0]), exp_sol(np.asarray(t_test)), atol=1e-5 + ) + assert np.all(np.asarray(sol.rms_residuals[: num - 1]) < 1e-3) + + # The returned arrays are exactly the spline's knots and knot derivatives. + np.testing.assert_allclose( + np.asarray(hermite_interpolate(sol.t, sol.t, sol.y, sol.yp)), + np.asarray(sol.y), + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + np.asarray(hermite_derivative(sol.t, sol.t, sol.y, sol.yp)), + np.asarray(sol.yp), + rtol=1e-10, + atol=1e-10, + ) + + yp_test = hermite_derivative(t_test, sol.t, sol.y, sol.yp) + f_test = jax.vmap(exp_fun, in_axes=(0, 0))(t_test, y_test) + norm_res = relative_residual_norm(np.asarray(f_test), np.asarray(yp_test)) + assert np.all(norm_res < 1e-3) + + +def test_exponential_matches_scipy(): + def np_fun(x, y): + return np.vstack((y[1], y[0])) + + def np_bc(ya, yb): + return np.array([ya[0] - 1.0, yb[0]]) + + def np_fun_jac(x, y): + df_dy = np.zeros((2, 2, x.shape[0])) + df_dy[0, 1] = 1.0 + df_dy[1, 0] = 1.0 + return df_dy + + def np_bc_jac(ya, yb): + return np.array([[1.0, 0.0], [0.0, 0.0]]), np.array([[0.0, 0.0], [1.0, 0.0]]) + + reference = scipy_solve_bvp( + np_fun, + np_bc, + np.linspace(0, 1, 5), + np.zeros((2, 5)), + fun_jac=np_fun_jac, + bc_jac=np_bc_jac, + ) + sol = solve_bvp( + exp_fun, exp_bc, jnp.linspace(0, 1, 5), jnp.zeros((5, 2)), max_nodes=32 + ) + num = int(sol.num_nodes) + assert num == reference.x.size + assert int(sol.num_iterations) == reference.niter + np.testing.assert_allclose(np.asarray(sol.t[:num]), reference.x, rtol=1e-12) + np.testing.assert_allclose( + np.asarray(sol.y[:num]).T, reference.y, rtol=1e-8, atol=1e-14 + ) + np.testing.assert_allclose( + np.asarray(sol.yp[:num]).T, reference.yp, rtol=1e-8, atol=1e-14 + ) + np.testing.assert_allclose( + np.asarray(sol.rms_residuals[: num - 1]), + reference.rms_residuals, + rtol=1e-6, + atol=1e-14, + ) + + +def test_sturm_liouville_unknown_parameter(): + t = jnp.linspace(0.0, jnp.pi, 5) + sol = solve_bvp(sl_fun, sl_bc, t, jnp.ones((5, 2)), jnp.array([0.5]), max_nodes=32) + assert int(sol.status) == 0 + num = int(sol.num_nodes) + assert num < 10 + np.testing.assert_allclose(np.asarray(sol.z), [1.0], rtol=1e-4) + t_active = np.asarray(sol.t[:num]) + np.testing.assert_allclose( + np.asarray(sol.y[:num, 0]), np.sin(t_active), rtol=1e-4, atol=1e-4 + ) + assert np.all(np.asarray(sol.rms_residuals[: num - 1]) < 1e-3) + + def np_fun(x, y, p): + return np.vstack((y[1], -(p[0] ** 2) * y[0])) + + def np_bc(ya, yb, p): + return np.array([ya[0], yb[0], ya[1] - p[0]]) + + def np_fun_jac(x, y, p): + df_dy = np.zeros((2, 2, x.shape[0])) + df_dy[0, 1] = 1.0 + df_dy[1, 0] = -(p[0] ** 2) + df_dp = np.zeros((2, 1, x.shape[0])) + df_dp[1, 0] = -2.0 * p[0] * y[0] + return df_dy, df_dp + + def np_bc_jac(ya, yb, p): + dbc_dya = np.array([[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]]) + dbc_dyb = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 0.0]]) + dbc_dp = np.array([[0.0], [0.0], [-1.0]]) + return dbc_dya, dbc_dyb, dbc_dp + + reference = scipy_solve_bvp( + np_fun, + np_bc, + np.linspace(0, np.pi, 5), + np.ones((2, 5)), + p=[0.5], + fun_jac=np_fun_jac, + bc_jac=np_bc_jac, + ) + assert num == reference.x.size + np.testing.assert_allclose(np.asarray(sol.z), reference.p, rtol=1e-9) + np.testing.assert_allclose( + np.asarray(sol.y[:num]).T, reference.y, rtol=1e-8, atol=1e-12 + ) + + +def test_singular_term_emden(): + t = jnp.linspace(0.0, 1.0, 10) + y_0 = jnp.stack([jnp.full(10, (3.0 / 4.0) ** 0.5), jnp.full(10, 1e-4)], axis=1) + sol = solve_bvp(emden_fun, emden_bc, t, y_0, S=EMDEN_S, max_nodes=32) + assert int(sol.status) == 0 + assert int(sol.num_nodes) == 10 + + t_test = jnp.linspace(0.05, 1.0, 100) + y_test = hermite_interpolate(t_test, sol.t, sol.y, sol.yp) + np.testing.assert_allclose( + np.asarray(y_test[:, 0]), emden_sol(np.asarray(t_test)), atol=1e-5 + ) + + yp_test = hermite_derivative(t_test, sol.t, sol.y, sol.yp) + f_raw = jax.vmap(emden_fun, in_axes=(0, 0))(t_test, y_test) + f_test = ( + np.asarray(f_raw) + + np.asarray(y_test) @ np.asarray(EMDEN_S).T / (np.asarray(t_test)[:, None]) + ) + norm_res = relative_residual_norm(f_test, np.asarray(yp_test)) + assert np.all(norm_res < 1e-3) + + +def test_shock_layer(): + eps = 1e-3 + + def shock_fun(t, y): + return jnp.array( + [ + y[1], + -( + t * y[1] + + eps * jnp.pi**2 * jnp.cos(jnp.pi * t) + + jnp.pi * t * jnp.sin(jnp.pi * t) + ) + / eps, + ] + ) + + def shock_bc(ya, yb): + return jnp.array([ya[0] + 2.0, yb[0]]) + + t = jnp.linspace(-1.0, 1.0, 5) + sol = solve_bvp(shock_fun, shock_bc, t, jnp.zeros((5, 2)), max_nodes=128) + assert int(sol.status) == 0 + num = int(sol.num_nodes) + assert num < 110 + t_active = np.asarray(sol.t[:num]) + exact = np.cos(np.pi * t_active) + erf(t_active / np.sqrt(2.0 * eps)) / erf( + 1.0 / np.sqrt(2.0 * eps) + ) + np.testing.assert_allclose(np.asarray(sol.y[:num, 0]), exact, rtol=1e-5, atol=1e-5) + + +def test_big_problem_with_parameters(): + def big_fun(t, y, z): + f = jnp.zeros_like(y) + f = f.at[::2].set(y[1::2]) + f = f.at[1::4].set(-(z[0] ** 2) * y[::4]) + f = f.at[3::4].set(-(z[1] ** 2) * y[2::4]) + return f + + def big_bc(ya, yb, z): + return jnp.concatenate( + [ya[::2], yb[::2], jnp.array([ya[1] - z[0], ya[3] - z[1]])] + ) + + t = jnp.linspace(0.0, jnp.pi, 5) + sol = solve_bvp( + big_fun, big_bc, t, jnp.ones((5, 60)), jnp.array([0.5, 0.5]), max_nodes=24 + ) + assert int(sol.status) == 0 + num = int(sol.num_nodes) + np.testing.assert_allclose(np.asarray(sol.z), [1.0, 1.0], rtol=1e-4) + t_active = np.asarray(sol.t[:num]) + np.testing.assert_allclose( + np.asarray(sol.y[:num, 0]), np.sin(t_active), rtol=1e-4, atol=1e-4 + ) + np.testing.assert_allclose( + np.asarray(sol.y[:num, 2]), np.sin(t_active), rtol=1e-4, atol=1e-4 + ) + + +def test_failures(): + sol = solve_bvp( + exp_fun, exp_bc, jnp.linspace(0, 1, 2), jnp.zeros((2, 2)), tol=1e-5, max_nodes=5 + ) + assert int(sol.status) == 1 + assert not bool(sol.ok) + + def undefined_fun(t, y): + return jnp.zeros(2) + + def undefined_bc(ya, yb): + return jnp.array([ya[0], yb[0] - 1.0]) + + sol = solve_bvp( + undefined_fun, + undefined_bc, + jnp.linspace(0, 1, 5), + jnp.zeros((5, 2)), + max_nodes=8, + ) + assert int(sol.status) == 2 + assert not bool(sol.ok) + assert bool(jnp.all(jnp.isfinite(sol.y))) + + # A NaN boundary residual follows scipy's elif chain to status 3 at the + # tenth iteration instead of looping to the safety bound. + def nan_bc(ya, yb): + return jnp.array([ya[0] - 1.0, yb[0] + jnp.where(yb[0] < 0.0, jnp.nan, 0.0)]) + + sol = solve_bvp( + exp_fun, nan_bc, jnp.linspace(0, 1, 5), -jnp.ones((5, 2)), max_nodes=16 + ) + assert int(sol.status) == 3 + assert int(sol.num_iterations) == 10 + + +def scipy_sl_j_true(t, h, y, z): + m = t.shape[0] + n = 2 + + def j_block(h_i, z_0): + return np.array( + [ + [ + h_i**2 * z_0**2 / 12 - 1, + -0.5 * h_i, + -(h_i**2) * z_0**2 / 12 + 1, + -0.5 * h_i, + ], + [ + 0.5 * h_i * z_0**2, + h_i**2 * z_0**2 / 12 - 1, + 0.5 * h_i * z_0**2, + 1 - h_i**2 * z_0**2 / 12, + ], + ] + ) + + j_true = np.zeros((m * n + 1, m * n + 1)) + for i in range(m - 1): + j_true[i * n : (i + 1) * n, i * n : (i + 2) * n] = j_block(h[i], z[0]) + j_true[: (m - 1) * n : 2, -1] = z * h**2 / 6 * (y[0, :-1] - y[0, 1:]) + j_true[1 : (m - 1) * n : 2, -1] = z * ( + h * (y[0, :-1] + y[0, 1:]) + h**2 / 6 * (y[1, :-1] - y[1, 1:]) + ) + j_true[(m - 1) * n, 0] = 1 + j_true[(m - 1) * n + 1, (m - 1) * n] = 1 + j_true[(m - 1) * n + 2, 1] = 1 + j_true[(m - 1) * n + 2, -1] = -1 + return j_true + + +def test_global_jacobian_matches_closed_form(): + t = jnp.linspace(0.0, 1.0, 5) + h = t[1:] - t[:-1] + y_rows = jnp.stack([jnp.sin(jnp.pi * t), jnp.pi * jnp.cos(jnp.pi * t)], axis=1) + z = jnp.array([3.0]) + cfg = make_config(sl_fun, sl_bc, 5, 2, 1) + system = bvp_module.build_system(cfg, None, None, t[0]) + _, y_middle, _, _ = system["collocation_parts"](t, h, y_rows, z, None) + jac = system["jacobian_at"](t, h, y_rows, z, y_middle, None) + j_true = scipy_sl_j_true( + np.asarray(t), np.asarray(h), np.asarray(y_rows).T, np.asarray(z) + ) + np.testing.assert_allclose(np.asarray(jac), j_true, rtol=1e-10, atol=1e-14) + + +def test_global_jacobian_padding_is_an_exact_embedding(): + t5 = jnp.linspace(0.0, 1.0, 5) + y5 = jnp.stack([jnp.sin(jnp.pi * t5), jnp.pi * jnp.cos(jnp.pi * t5)], axis=1) + z = jnp.array([3.0]) + cfg5 = make_config(sl_fun, sl_bc, 5, 2, 1) + system5 = bvp_module.build_system(cfg5, None, None, t5[0]) + h5 = t5[1:] - t5[:-1] + _, y_middle5, _, _ = system5["collocation_parts"](t5, h5, y5, z, None) + jac5 = np.asarray(system5["jacobian_at"](t5, h5, y5, z, y_middle5, None)) + + t8 = jnp.concatenate([t5, jnp.full(3, t5[-1])]) + y8 = jnp.concatenate([y5, jnp.broadcast_to(y5[-1], (3, 2))]) + cfg8 = make_config(sl_fun, sl_bc, 8, 2, 1) + system8 = bvp_module.build_system(cfg8, None, None, t8[0]) + h8 = t8[1:] - t8[:-1] + _, y_middle8, _, _ = system8["collocation_parts"](t8, h8, y8, z, None) + jac8 = np.asarray(system8["jacobian_at"](t8, h8, y8, z, y_middle8, None)) + + # Active collocation rows are identical, with the z column relocated. + np.testing.assert_allclose(jac8[:8, :10], jac5[:8, :10], rtol=1e-14) + np.testing.assert_allclose(jac8[:8, 16], jac5[:8, 10], rtol=1e-14) + # Padded interval rows are the exact copy chain [-I | I]. + eye = np.eye(2) + for block in range(4, 7): + rows = jac8[block * 2 : (block + 1) * 2] + expected = np.zeros((2, 17)) + expected[:, block * 2 : (block + 1) * 2] = -eye + expected[:, (block + 1) * 2 : (block + 2) * 2] = eye + assert np.array_equal(rows, expected) + # Boundary rows move to the padded last node block. + np.testing.assert_allclose(jac8[14:, :2], jac5[8:, :2], rtol=1e-14) + np.testing.assert_allclose(jac8[14:, 14:16], jac5[8:, 8:10], rtol=1e-14) + np.testing.assert_allclose(jac8[14:, 16], jac5[8:, 10], rtol=1e-14) + # The embedding preserves the determinant magnitude. + sign5, logdet5 = np.linalg.slogdet(jac5) + sign8, logdet8 = np.linalg.slogdet(jac8) + np.testing.assert_allclose(logdet8, logdet5, rtol=1e-10) + + +def test_global_jacobian_matches_autodiff(): + # Sturm-Liouville with an unknown parameter, padded mesh. + t = jnp.concatenate([jnp.linspace(0.0, 1.0, 5), jnp.full(3, 1.0)]) + y_rows = jnp.stack([jnp.sin(jnp.pi * t), jnp.pi * jnp.cos(jnp.pi * t)], axis=1) + z = jnp.array([3.0]) + cfg = make_config(sl_fun, sl_bc, 8, 2, 1) + system = bvp_module.build_system(cfg, None, None, t[0]) + h = t[1:] - t[:-1] + _, y_middle, _, _ = system["collocation_parts"](t, h, y_rows, z, None) + jac = system["jacobian_at"](t, h, y_rows, z, y_middle, None) + + def residual_vector(u): + y_flat = u[:16].reshape(8, 2) + z_flat = u[16:] + col_res, _, _, _ = system["collocation_parts"](t, h, y_flat, z_flat, None) + bc_res = system["bc_values"](y_flat[0], y_flat[-1], z_flat, None) + return jnp.concatenate([col_res.reshape(-1), bc_res]) + + u_0 = jnp.concatenate([y_rows.reshape(-1), z]) + jac_ad = jax.jacfwd(residual_vector)(u_0) + np.testing.assert_allclose( + np.asarray(jac), np.asarray(jac_ad), rtol=1e-10, atol=1e-12 + ) + + # Emden with the singular term: the D/S wrapping must differentiate + # consistently with the assembled blocks. + t = jnp.linspace(0.0, 1.0, 10) + y_rows = jnp.stack([emden_sol(t), jnp.full(10, 0.1)], axis=1) + cfg = make_config(emden_fun, emden_bc, 10, 2, 0, has_singular_term=True) + system = bvp_module.build_system(cfg, None, EMDEN_S, t[0]) + h = t[1:] - t[:-1] + _, y_middle, _, _ = system["collocation_parts"](t, h, y_rows, None, None) + jac = system["jacobian_at"](t, h, y_rows, None, y_middle, None) + + def emden_residual(u): + y_flat = u.reshape(10, 2) + col_res, _, _, _ = system["collocation_parts"](t, h, y_flat, None, None) + bc_res = system["bc_values"](y_flat[0], y_flat[-1], None, None) + return jnp.concatenate([col_res.reshape(-1), bc_res]) + + jac_ad = jax.jacfwd(emden_residual)(y_rows.reshape(-1)) + np.testing.assert_allclose( + np.asarray(jac), np.asarray(jac_ad), rtol=1e-10, atol=1e-12 + ) + + +def test_structured_qr_matches_dense_linear_algebra(): + rng = np.random.default_rng(0) + + # Padded Sturm-Liouville system with an unknown-parameter border (k = 1). + t = jnp.concatenate([jnp.linspace(0.0, 1.0, 5), jnp.full(3, 1.0)]) + y_rows = jnp.stack([jnp.sin(jnp.pi * t), jnp.pi * jnp.cos(jnp.pi * t)], axis=1) + z = jnp.array([3.0]) + cfg = make_config(sl_fun, sl_bc, 8, 2, 1) + system = bvp_module.build_system(cfg, None, None, t[0]) + h = t[1:] - t[:-1] + _, y_middle, _, _ = system["collocation_parts"](t, h, y_rows, z, None) + blocks = system["jacobian_blocks"](t, h, y_rows, z, y_middle, None) + + # Emden with the singular term, no border (k = 0). + t_e = jnp.linspace(0.0, 1.0, 10) + y_e = jnp.stack([emden_sol(t_e), jnp.full(10, 0.1)], axis=1) + cfg_e = make_config(emden_fun, emden_bc, 10, 2, 0, has_singular_term=True) + system_e = bvp_module.build_system(cfg_e, None, EMDEN_S, t_e[0]) + h_e = t_e[1:] - t_e[:-1] + _, y_middle_e, _, _ = system_e["collocation_parts"](t_e, h_e, y_e, None, None) + blocks_e = system_e["jacobian_blocks"](t_e, h_e, y_e, None, y_middle_e, None) + + for case_blocks in (blocks, blocks_e): + dense = np.asarray(bvp_module.babd_dense(case_blocks)) + size = dense.shape[0] + rhs = jnp.asarray(rng.standard_normal(size)) + np.testing.assert_allclose( + np.asarray(babd_matvec(case_blocks, rhs)), + dense @ np.asarray(rhs), + rtol=1e-13, + atol=1e-13, + ) + state, ok = structured_qr_factor(case_blocks) + assert bool(ok) + np.testing.assert_allclose( + np.asarray(structured_qr_solve(state, rhs)), + np.linalg.solve(dense, np.asarray(rhs)), + rtol=1e-11, + atol=1e-12, + ) + np.testing.assert_allclose( + np.asarray(structured_qr_transpose_solve(state, rhs)), + np.linalg.solve(dense.T, np.asarray(rhs)), + rtol=1e-11, + atol=1e-12, + ) + + +def test_pytree_state_matches_flat(): + def tree_fun(t, y): + return {"a": y["b"][0], "b": jnp.array([y["a"]])} + + def tree_bc(ya, yb): + return jnp.array([ya["a"] - 1.0, yb["a"]]) + + t = jnp.linspace(0.0, 1.0, 5) + tree_sol = solve_bvp( + tree_fun, + tree_bc, + t, + {"a": jnp.zeros(5), "b": jnp.zeros((5, 1))}, + max_nodes=16, + ) + flat_sol = solve_bvp(exp_fun, exp_bc, t, jnp.zeros((5, 2)), max_nodes=16) + assert int(tree_sol.status) == int(flat_sol.status) + assert int(tree_sol.num_nodes) == int(flat_sol.num_nodes) + assert jnp.array_equal(tree_sol.t, flat_sol.t) + assert jnp.array_equal(tree_sol.y["a"], flat_sol.y[:, 0]) + assert jnp.array_equal(tree_sol.y["b"][:, 0], flat_sol.y[:, 1]) + assert jnp.array_equal(tree_sol.yp["a"], flat_sol.yp[:, 0]) + assert jnp.array_equal(tree_sol.rms_residuals, flat_sol.rms_residuals) + + +def test_padded_tail_is_exact(): + sol = solve_bvp( + exp_fun, exp_bc, jnp.linspace(0, 1, 5), jnp.zeros((5, 2)), max_nodes=16 + ) + num = int(sol.num_nodes) + assert jnp.array_equal(sol.t[num:], jnp.full(16 - num, sol.t[num - 1])) + assert jnp.array_equal(sol.y[num:], jnp.broadcast_to(sol.y[num - 1], (16 - num, 2))) + assert jnp.array_equal( + sol.yp[num:], jnp.broadcast_to(sol.yp[num - 1], (16 - num, 2)) + ) + assert jnp.array_equal(sol.rms_residuals[num - 1 :], jnp.zeros(16 - num)) + + +def test_aux_at_solution(): + def aux_fun(t, y, z, args, p): + return jnp.array([y[1], y[0]]), p * y[0] + + def aux_bc(ya, yb, z, args, p): + return jnp.array([ya[0] - 1.0, yb[0]]) + + t = jnp.linspace(0.0, 1.0, 5) + sol = solve_bvp(aux_fun, aux_bc, t, jnp.zeros((5, 2)), p=2.0, max_nodes=16) + assert jnp.array_equal(sol.aux, 2.0 * sol.y[:, 0]) + explicit = solve_bvp( + aux_fun, aux_bc, t, jnp.zeros((5, 2)), p=2.0, max_nodes=16, has_aux=True + ) + assert jnp.array_equal(explicit.aux, sol.aux) + disabled = solve_bvp( + exp_fun, exp_bc, t, jnp.zeros((5, 2)), max_nodes=16, has_aux=False + ) + assert disabled.aux is None + + +def test_float32(): + t = jnp.linspace(0.0, 1.0, 5, dtype=jnp.float32) + y_0 = jnp.zeros((5, 2), jnp.float32) + sol = solve_bvp(exp_fun, exp_bc, t, y_0, max_nodes=16) + assert sol.y.dtype == jnp.float32 + assert sol.t.dtype == jnp.float32 + assert int(sol.status) == 0 + num = int(sol.num_nodes) + assert num == 5 + np.testing.assert_allclose( + np.asarray(sol.y[:num, 0]), exp_sol(np.asarray(sol.t[:num])), atol=1e-4 + ) + # A tolerance below the float32 floor clamps to 100 * eps and converges. + clamped = solve_bvp(exp_fun, exp_bc, t, y_0, tol=1e-12, max_nodes=32) + assert int(clamped.status) == 0 diff --git a/tests/test_bvp_ad.py b/tests/test_bvp_ad.py new file mode 100644 index 0000000..ccca84d --- /dev/null +++ b/tests/test_bvp_ad.py @@ -0,0 +1,254 @@ +import jax +import jax.numpy as jnp +import numpy as np + +from tinydiffeq import solve_bvp + +# y'' = -z^2 p y on [0, 1] with bc [y(0), y(1), y'(0) - z]: the eigenvalue is +# z*(p) = pi / sqrt(p) and y(t) = sin(pi t) / sqrt(p), so at p = 1 +# dz*/dp = -pi/2, dy/dp = -sin(pi t)/2, dy'/dp = -pi cos(pi t)/2. + + +def eigen_fun(t, y, z, args, p): + return jnp.array([y[1], -(z[0] ** 2) * p * y[0]]) + + +def eigen_bc(ya, yb, z, args, p): + return jnp.array([ya[0], yb[0], ya[1] - z[0]]) + + +T_GRID = jnp.linspace(0.0, 1.0, 9) +Y_GUESS = jnp.stack( + [jnp.sin(jnp.pi * T_GRID), jnp.pi * jnp.cos(jnp.pi * T_GRID)], axis=1 +) +MAX_NODES = 128 + + +def eigen_solve(p): + return solve_bvp( + eigen_fun, + eigen_bc, + T_GRID, + Y_GUESS, + jnp.array([3.0]), + p=p, + tol=1e-6, + max_nodes=MAX_NODES, + ) + + +def test_jvp_matches_closed_form(): + sol, tangent = jax.jvp(eigen_solve, (1.0,), (1.0,)) + assert int(sol.status) == 0 + num = int(sol.num_nodes) + np.testing.assert_allclose(float(tangent.z[0]), -np.pi / 2, rtol=1e-6) + t_active = np.asarray(sol.t[:num]) + np.testing.assert_allclose( + np.asarray(tangent.y[:num, 0]), + -np.sin(np.pi * t_active) / 2, + rtol=1e-5, + atol=1e-8, + ) + np.testing.assert_allclose( + np.asarray(tangent.yp[:num, 0]), + -np.pi * np.cos(np.pi * t_active) / 2, + rtol=1e-5, + atol=1e-7, + ) + assert jnp.array_equal(tangent.t, jnp.zeros(MAX_NODES)) + assert jnp.array_equal(tangent.rms_residuals, jnp.zeros(MAX_NODES - 1)) + + +def test_vjp_matches_closed_form(): + gradient = jax.grad(lambda p: eigen_solve(p).z[0])(1.0) + np.testing.assert_allclose(float(gradient), -np.pi / 2, rtol=1e-6) + + gradient = jax.grad(lambda p: jnp.sum(eigen_solve(p).y[:, 0]))(1.0) + sol = eigen_solve(1.0) + num = int(sol.num_nodes) + t_active = np.asarray(sol.t[:num]) + # The padded tail repeats the last active row, so its tangent repeats too. + expected = np.sum(-np.sin(np.pi * t_active) / 2) + (MAX_NODES - num) * ( + -np.sin(np.pi * t_active[-1]) / 2 + ) + np.testing.assert_allclose(float(gradient), expected, rtol=1e-5, atol=1e-8) + + +def test_second_order_matches_closed_form(): + # z*(p) = pi p^(-1/2): d2z/dp2 = (3 pi / 4) p^(-5/2) = 3 pi / 4 at p = 1. + hessian = jax.hessian(lambda p: eigen_solve(p).z[0])(1.0) + np.testing.assert_allclose(float(hessian), 3 * np.pi / 4, rtol=1e-5) + reverse_over_forward = jax.grad( + lambda p: jax.jvp(lambda q: eigen_solve(q).z[0], (p,), (1.0,))[1] + )(1.0) + np.testing.assert_allclose(float(reverse_over_forward), 3 * np.pi / 4, rtol=1e-5) + + +def test_vjp_is_the_transpose_of_jvp(): + cotangent = jax.random.normal(jax.random.key(0), (MAX_NODES, 2)) + + def scalar(p): + sol = eigen_solve(p) + return jnp.sum(cotangent * sol.y) + 0.7 * sol.z[0] + + _, jvp_value = jax.jvp(scalar, (1.0,), (1.0,)) + vjp_value = jax.grad(scalar)(1.0) + np.testing.assert_allclose(float(vjp_value), float(jvp_value), rtol=1e-12) + + +def test_aux_tangent(): + def aux_fun(t, y, z, args, p): + return jnp.array([y[1], -(z[0] ** 2) * p * y[0]]), p * y[0] + + def solve(p): + return solve_bvp( + aux_fun, + eigen_bc, + T_GRID, + Y_GUESS, + jnp.array([3.0]), + p=p, + tol=1e-6, + max_nodes=MAX_NODES, + ) + + sol, tangent = jax.jvp(solve, (1.0,), (1.0,)) + num = int(sol.num_nodes) + t_active = np.asarray(sol.t[:num]) + # aux = p y(t): d aux/dp = y + p dy/dp = sin(pi t)/2 at p = 1. + np.testing.assert_allclose( + np.asarray(tangent.aux[:num]), + np.sin(np.pi * t_active) / 2, + rtol=1e-5, + atol=1e-8, + ) + + +def test_inert_inputs_have_exactly_zero_gradients(): + def with_args_fun(t, y, z, args, p): + return jnp.array([y[1], -(z[0] ** 2) * p * args * y[0]]) + + def objective(t, y_0, z_0, args): + sol = solve_bvp( + with_args_fun, + eigen_bc, + t, + y_0, + z_0, + p=1.0, + args=args, + tol=1e-6, + max_nodes=MAX_NODES, + ) + return sol.z[0] + + gradients = jax.grad(objective, argnums=(0, 1, 2, 3))( + T_GRID, Y_GUESS, jnp.array([3.0]), 1.0 + ) + for gradient in gradients: + assert bool(jnp.all(gradient == 0.0)) + + +def test_failed_solve_has_zero_finite_tangents(): + def solve(p): + # tol = 1e-8 needs far more than 64 nodes: status 1. + return solve_bvp( + eigen_fun, + eigen_bc, + T_GRID, + Y_GUESS, + jnp.array([3.0]), + p=p, + tol=1e-8, + max_nodes=64, + ) + + sol, tangent = jax.jvp(solve, (1.0,), (1.0,)) + assert int(sol.status) == 1 + assert jnp.array_equal(tangent.z, jnp.zeros(1)) + assert jnp.array_equal(tangent.y, jnp.zeros((64, 2))) + gradient = jax.grad(lambda p: solve(p).z[0])(1.0) + assert float(gradient) == 0.0 + + +def test_failed_solve_with_nonfinite_tangent_program_stays_zero(): + # d/dp sqrt(y + p) is infinite at the inert guess y = 0, p = 0, so the + # failed lane's tangent program produces non-finite intermediates that + # the zero-tangent contract must mask by selection, not multiplication. + def sqrt_fun(t, y, z, args, p): + return jnp.array([y[1], jnp.sqrt(y[0] + p)]) + + def sqrt_bc(ya, yb, z, args, p): + return jnp.array([ya[0] + 2.0, yb[0]]) + + t = jnp.linspace(0.0, 1.0, 5) + + def solve(p): + return solve_bvp( + sqrt_fun, sqrt_bc, t, jnp.zeros((5, 2)), p=p, tol=1e-10, max_nodes=8 + ) + + sol, tangent = jax.jvp(solve, (0.0,), (1.0,)) + assert int(sol.status) != 0 + assert jnp.array_equal(tangent.y, jnp.zeros((8, 2))) + gradient = jax.grad(lambda p: solve(p).y[0, 0])(0.0) + assert float(gradient) == 0.0 + + +def test_singular_solution_jacobian_reports_nan_in_both_modes(): + # y' = 0 with bc ya^2 - ya^3/2 - p: p = 0 converges to the double root + # ya = 0, where the bc Jacobian (so the AD-rule refactor) is singular. + def flat_fun(t, y, z, args, p): + return jnp.zeros_like(y) + + def root_bc(ya, yb, z, args, p): + return jnp.array([ya[0] ** 2 - 0.5 * ya[0] ** 3 - p]) + + t = jnp.linspace(0.0, 1.0, 3) + + def solve(p): + return solve_bvp(flat_fun, root_bc, t, jnp.ones((3, 1)), p=p, max_nodes=8) + + sol = solve(0.0) + assert int(sol.status) == 0 + assert float(sol.y[0, 0]) == 0.0 + _, tangent = jax.jvp(solve, (0.0,), (1.0,)) + assert bool(jnp.all(jnp.isnan(tangent.y))) + gradient = jax.grad(lambda p: solve(p).y[0, 0])(0.0) + assert bool(jnp.isnan(gradient)) + + +def test_jit_and_ad_compose(): + gradient = jax.grad(lambda p: eigen_solve(p).z[0])(1.0) + jitted = jax.jit(jax.grad(lambda p: eigen_solve(p).z[0]))(1.0) + np.testing.assert_allclose(float(jitted), float(gradient), rtol=1e-13) + _, jvp_eager = jax.jvp(lambda p: eigen_solve(p).z[0], (1.0,), (1.0,)) + jvp_jitted = jax.jit( + lambda p: jax.jvp(lambda q: eigen_solve(q).z[0], (p,), (1.0,))[1] + )(1.0) + np.testing.assert_allclose(float(jvp_jitted), float(jvp_eager), rtol=1e-12) + + +def test_singular_term_gradient_matches_finite_differences(): + def emden_fun(t, y, z, args, p): + return jnp.array([y[1], -(y[0] ** 5)]) + + def emden_bc(ya, yb, z, args, p): + return jnp.array([ya[1], yb[0] - p * (3.0 / 4.0) ** 0.5]) + + S = jnp.array([[0.0, 0.0], [0.0, -2.0]]) + t = jnp.linspace(0.0, 1.0, 10) + y_0 = jnp.stack([jnp.full(10, (3.0 / 4.0) ** 0.5), jnp.full(10, 1e-4)], axis=1) + + def solve(p): + return solve_bvp(emden_fun, emden_bc, t, y_0, p=p, S=S, tol=1e-6, max_nodes=64) + + value = solve(1.0) + assert int(value.status) == 0 + gradient = jax.grad(lambda p: solve(p).y[5, 0])(1.0) + step = 1e-6 + upper, lower = solve(1.0 + step), solve(1.0 - step) + # Central differences are only valid on an unchanged mesh. + assert int(upper.num_nodes) == int(lower.num_nodes) == int(value.num_nodes) + finite_difference = (float(upper.y[5, 0]) - float(lower.y[5, 0])) / (2 * step) + np.testing.assert_allclose(float(gradient), finite_difference, rtol=1e-4) diff --git a/tests/test_bvp_vmap.py b/tests/test_bvp_vmap.py new file mode 100644 index 0000000..ee0786b --- /dev/null +++ b/tests/test_bvp_vmap.py @@ -0,0 +1,84 @@ +import jax +import jax.numpy as jnp +import numpy as np + +from tinydiffeq import solve_bvp + + +def scaled_fun(t, y, z, args, p): + return jnp.array([y[1], y[0]]) + + +def scaled_bc(ya, yb, z, args, p): + return jnp.array([ya[0] - p, yb[0]]) + + +T_GRID = jnp.linspace(0.0, 1.0, 5) + + +def scaled_solve(p): + return solve_bvp( + scaled_fun, scaled_bc, T_GRID, jnp.zeros((5, 2)), p=p, max_nodes=16 + ) + + +def test_vmap_matches_sequential(): + p_values = jnp.array([0.5, 1.0, 2.0]) + batched = jax.vmap(scaled_solve)(p_values) + for lane, p in enumerate([0.5, 1.0, 2.0]): + single = scaled_solve(p) + assert int(batched.status[lane]) == int(single.status) + assert int(batched.num_nodes[lane]) == int(single.num_nodes) + # Batched linear-algebra kernels round differently than single ones, + # so agreement is at roundoff rather than bitwise. + np.testing.assert_allclose( + np.asarray(batched.y[lane]), np.asarray(single.y), rtol=1e-12, atol=1e-14 + ) + np.testing.assert_allclose( + np.asarray(batched.t[lane]), np.asarray(single.t), rtol=1e-12 + ) + + batched_gradients = jax.vmap(jax.grad(lambda p: scaled_solve(p).y[2, 0]))(p_values) + for lane, p in enumerate([0.5, 1.0, 2.0]): + single = jax.grad(lambda p: scaled_solve(p).y[2, 0])(p) + np.testing.assert_allclose( + float(batched_gradients[lane]), float(single), rtol=1e-10 + ) + + +def test_vmap_mixed_success_and_failure(): + def stiff_solve(eps): + def fun(t, y, z, args, p): + return jnp.array( + [ + y[1], + -( + t * y[1] + + p * jnp.pi**2 * jnp.cos(jnp.pi * t) + + jnp.pi * t * jnp.sin(jnp.pi * t) + ) + / p, + ] + ) + + def bc(ya, yb, z, args, p): + return jnp.array([ya[0] + 2.0, yb[0]]) + + return solve_bvp( + fun, bc, jnp.linspace(-1.0, 1.0, 5), jnp.zeros((5, 2)), p=eps, max_nodes=24 + ) + + eps_values = jnp.array([1e-1, 1e-4]) + batched = jax.vmap(stiff_solve)(eps_values) + assert int(batched.status[0]) == 0 + assert int(batched.status[1]) == 1 + single = stiff_solve(1e-1) + np.testing.assert_allclose( + np.asarray(batched.y[0]), np.asarray(single.y), rtol=1e-12, atol=1e-14 + ) + + gradients = jax.vmap(jax.grad(lambda eps: stiff_solve(eps).y[3, 0]))(eps_values) + assert bool(jnp.all(jnp.isfinite(gradients))) + assert float(gradients[1]) == 0.0 + single_gradient = jax.grad(lambda eps: stiff_solve(eps).y[3, 0])(1e-1) + np.testing.assert_allclose(float(gradients[0]), float(single_gradient), rtol=1e-10) From dfdf21b7a930eeb9fe1f676114da23a4048ea040 Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Wed, 19 Aug 2026 21:25:47 -0700 Subject: [PATCH 3/5] test: drop executable-cache-count tests 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) Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g --- AGENTS.md | 28 +++++++++ CLAUDE.md | 1 + tests/test_pytree_states.py | 20 ------ tests/test_recompile.py | 118 ------------------------------------ 4 files changed, 29 insertions(+), 138 deletions(-) create mode 100644 AGENTS.md create mode 100644 CLAUDE.md delete mode 100644 tests/test_recompile.py diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..c55e71c --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,28 @@ +# tinydiffeq — agent notes + +## Recompilation hygiene: a manual audit, not a test + +Do not write tests against `fn._cache_size()`. It reads JAX's globally shared +C++ executable cache (capacity 8192, LRU, shared by every jitted function and +every `jnp` op in the process), so absolute entry-count assertions turn flaky +as the suite grows: entries get evicted before the assert and the count reads +zero. + +Instead, after a big refactor or a substantial new feature, audit +recompilation with the environment-variable protocol from the `jax-project` +skill: two solves in one process with changed data leaves, a flushed marker +before each, and zero trace/compile events after the second marker. + +```bash +JAX_EXPLAIN_CACHE_MISSES=1 JAX_LOG_COMPILES=1 \ + uv run python -m benchmarks.bvp_scaling --cache-audit > /tmp/audit.log 2>&1 +grep -c "Compiling" /tmp/audit.log # after the "=== SOLVE 1" marker: zero +``` + +A per-call `eval_shape` validation trace is expected and compiles nothing. +Leaf-value changes (tolerances, meshes, `p`, `args`, initial states) must +not recompile; only static configuration (attempt budgets, `max_nodes`, +pytree structure, function identity, jacobian modes, solver objects) may. A +changed initial mesh *length* for `solve_bvp` reuses the solve executable +but compiles trivial eager padding ops (`concatenate`, `broadcast_in_dim`) +once per new length — those events are expected in the audit log. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..43c994c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/tests/test_pytree_states.py b/tests/test_pytree_states.py index 90f8fae..3198570 100644 --- a/tests/test_pytree_states.py +++ b/tests/test_pytree_states.py @@ -247,23 +247,3 @@ def test_field_and_project_must_preserve_structure(): dt_0=0.1, project=lambda x: (x["x"],), ) - - -def test_same_pytree_structure_and_shapes_reuse_compilation(): - @jax.jit - def run(x): - return solve_ode( - lambda state: jax.tree.map(lambda leaf: -leaf, state), - RK4(), - 0.0, - 1.0, - x, - dt_0=0.1, - max_steps=10, - ).xs - - run({"a": jnp.ones(2), "b": jnp.ones(3)}) - run({"a": 2 * jnp.ones(2), "b": 3 * jnp.ones(3)}) - assert run._cache_size() == 1 - run({"a": jnp.ones(2), "b": jnp.ones(4)}) - assert run._cache_size() == 2 diff --git a/tests/test_recompile.py b/tests/test_recompile.py deleted file mode 100644 index f694838..0000000 --- a/tests/test_recompile.py +++ /dev/null @@ -1,118 +0,0 @@ -import jax -import jax.numpy as jnp - -from tinydiffeq import IController, PIController, SaveAt, Tsit5, solve_ode - -# Hyperparameters that only change values (tolerances, dt_0, x_0, args) are -# pytree data leaves, and the bounded scan runs exactly max_steps iterations -# regardless of how many are accepted -- so none of these may retrace. - - -def field(x, t, args, p): - return -args * x - - -@jax.jit -def run_steps(x_0, dt_0, controller, args): - return solve_ode( - field, - Tsit5(), - 0.0, - 1.0, - x_0, - args=args, - dt_0=dt_0, - controller=controller, - max_steps=128, - save_at=SaveAt(steps=True), - ) - - -@jax.jit -def run_pi_steps(x_0, controller, args): - return solve_ode( - field, - Tsit5(), - 0.0, - 1.0, - x_0, - args=args, - dt_0=0.1, - controller=controller, - max_steps=128, - save_at=SaveAt(steps=True), - ) - - -def test_one_compilation_across_leaf_changes(): - base = run_steps( - jnp.asarray(1.0), - jnp.asarray(0.1), - IController(rtol=1e-6, atol=1e-8, dt_min=1e-10), - 1.0, - ) - # different curvature -> different accepted count, same compilation - stiff = run_steps( - jnp.asarray(1.0), - jnp.asarray(0.1), - IController(rtol=1e-6, atol=1e-8, dt_min=1e-10), - 40.0, - ) - assert int(stiff.num_accepted) != int(base.num_accepted) - # different tolerances, dt_0, and x_0 - run_steps( - jnp.asarray(2.0), - jnp.asarray(0.02), - IController(rtol=1e-10, atol=1e-12, dt_min=1e-10), - 3.0, - ) - run_steps( - jnp.asarray(0.3), - jnp.asarray(0.5), - IController(rtol=1e-4, atol=1e-6, dt_min=1e-8, safety=0.8), - 1.0, - ) - assert run_steps._cache_size() == 1 - - -def test_pi_coefficients_are_data_leaves(): - run_pi_steps( - jnp.asarray(1.0), - PIController(rtol=1e-6, atol=1e-8, p_coeff=0.4, i_coeff=0.3), - 1.0, - ) - run_pi_steps( - jnp.asarray(2.0), - PIController( - rtol=1e-10, - atol=1e-12, - p_coeff=0.2, - i_coeff=0.4, - safety=0.8, - ), - 30.0, - ) - assert run_pi_steps._cache_size() == 1 - - -def test_one_compilation_ts_mode_same_grid_length(): - @jax.jit - def run_ts(x_0, save_at, args): - return solve_ode( - field, - Tsit5(), - 0.0, - 1.0, - x_0, - args=args, - dt_0=0.1, - controller=IController(rtol=1e-8, atol=1e-10), - max_steps=128, - save_at=save_at, - ).xs - - grid_a = jnp.linspace(0.0, 1.0, 11) - grid_b = jnp.sqrt(jnp.linspace(0.0, 1.0, 11)) # same length, different knots - run_ts(jnp.asarray(1.0), SaveAt(ts=grid_a), 1.0) - run_ts(jnp.asarray(2.0), SaveAt(ts=grid_b), 25.0) - assert run_ts._cache_size() == 1 From 4eb810466dd4767b9b4775cbf378f84e91901367 Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Wed, 19 Aug 2026 21:25:59 -0700 Subject: [PATCH 4/5] bench: solve_bvp scaling benchmark 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) Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g --- benchmarks/bvp_scaling.py | 212 ++++++++++++++++++ .../results/2026-08-19_bvp-scaling-cpu.json | 68 ++++++ .../results/2026-08-19_bvp-scaling-cpu.md | 49 ++++ 3 files changed, 329 insertions(+) create mode 100644 benchmarks/bvp_scaling.py create mode 100644 benchmarks/results/2026-08-19_bvp-scaling-cpu.json create mode 100644 benchmarks/results/2026-08-19_bvp-scaling-cpu.md diff --git a/benchmarks/bvp_scaling.py b/benchmarks/bvp_scaling.py new file mode 100644 index 0000000..b1f0358 --- /dev/null +++ b/benchmarks/bvp_scaling.py @@ -0,0 +1,212 @@ +"""Compile and steady-state timings for solve_bvp, with scipy references. + +Every case calls the solve through a jitted wrapper: cold is +compile-plus-first-run, run is the warm steady state. +""" + +import argparse +import statistics +import timeit + +import jax +import jax.numpy as jnp +import numpy as np + +from tinydiffeq import solve_bvp + +jax.config.update("jax_enable_x64", True) + + +def exp_problem(max_nodes): + def fun(t, y): + return jnp.array([y[1], y[0]]) + + def bc(ya, yb): + return jnp.array([ya[0] - 1.0, yb[0]]) + + t = jnp.linspace(0.0, 1.0, 5) + y_0 = jnp.zeros((5, 2)) + jitted = jax.jit(lambda tt, yy: solve_bvp(fun, bc, tt, yy, max_nodes=max_nodes)) + + def solve(): + return jitted(t, y_0) + + def scipy_solve(): + from scipy.integrate import solve_bvp as scipy_solve_bvp + + return scipy_solve_bvp( + lambda x, y: np.vstack((y[1], y[0])), + lambda ya, yb: np.array([ya[0] - 1.0, yb[0]]), + np.linspace(0, 1, 5), + np.zeros((2, 5)), + ) + + return solve, scipy_solve + + +def shock_problem(max_nodes): + eps = 1e-3 + + def fun(t, y): + return jnp.array( + [ + y[1], + -( + t * y[1] + + eps * jnp.pi**2 * jnp.cos(jnp.pi * t) + + jnp.pi * t * jnp.sin(jnp.pi * t) + ) + / eps, + ] + ) + + def bc(ya, yb): + return jnp.array([ya[0] + 2.0, yb[0]]) + + t = jnp.linspace(-1.0, 1.0, 5) + y_0 = jnp.zeros((5, 2)) + jitted = jax.jit(lambda tt, yy: solve_bvp(fun, bc, tt, yy, max_nodes=max_nodes)) + + def solve(): + return jitted(t, y_0) + + def scipy_solve(): + from scipy.integrate import solve_bvp as scipy_solve_bvp + + def np_fun(x, y): + return np.vstack( + ( + y[1], + -( + x * y[1] + + eps * np.pi**2 * np.cos(np.pi * x) + + np.pi * x * np.sin(np.pi * x) + ) + / eps, + ) + ) + + return scipy_solve_bvp( + np_fun, + lambda ya, yb: np.array([ya[0] + 2.0, yb[0]]), + np.linspace(-1, 1, 5), + np.zeros((2, 5)), + ) + + return solve, scipy_solve + + +def eigen_gradient(max_nodes): + def fun(t, y, z, args, p): + return jnp.array([y[1], -(z[0] ** 2) * p * y[0]]) + + def bc(ya, yb, z, args, p): + return jnp.array([ya[0], yb[0], ya[1] - z[0]]) + + t = jnp.linspace(0.0, 1.0, 9) + y_0 = jnp.stack([jnp.sin(jnp.pi * t), jnp.pi * jnp.cos(jnp.pi * t)], axis=1) + + def objective(p): + return solve_bvp( + fun, + bc, + t, + y_0, + jnp.array([3.0]), + p=p, + tol=1e-6, + max_nodes=max_nodes, + ).z[0] + + gradient = jax.jit(jax.grad(objective)) + + def solve(): + return gradient(1.0) + + return solve, None + + +def vmapped_exp(max_nodes, batch): + def fun(t, y, z, args, p): + return jnp.array([y[1], y[0]]) + + def bc(ya, yb, z, args, p): + return jnp.array([ya[0] - p, yb[0]]) + + t = jnp.linspace(0.0, 1.0, 5) + y_0 = jnp.zeros((5, 2)) + p_values = jnp.linspace(0.5, 2.0, batch) + + def one(p): + return solve_bvp(fun, bc, t, y_0, p=p, max_nodes=max_nodes) + + batched = jax.jit(lambda ps: jax.vmap(one)(ps)) + + def solve(): + return batched(p_values) + + return solve, None + + +def cold_and_warm_ms(run, repeat): + def timed(): + jax.block_until_ready(jax.tree.leaves(run())) + + # Each factory builds fresh callables, so the first call in this process + # traces and compiles: compile-plus-first-run, reported separately from + # the warm steady state. + cold_ms = 1e3 * timeit.timeit(timed, number=1) + warm_ms = 1e3 * statistics.median(timeit.repeat(timed, repeat=repeat, number=1)) + return cold_ms, warm_ms + + +def cache_audit(max_nodes): + """Two solves with changed data leaves; run under + JAX_EXPLAIN_CACHE_MISSES=1 JAX_LOG_COMPILES=1 and count events after the + second marker — the target is zero.""" + + def fun(t, y, z, args, p): + return jnp.array([y[1], p * y[0]]) + + def bc(ya, yb, z, args, p): + return jnp.array([ya[0] - 1.0, yb[0]]) + + for index, (p, span) in enumerate([(1.0, 1.0), (2.5, 1.5)]): + print(f"=== SOLVE {index} p={p} ===", flush=True) + t = jnp.linspace(0.0, span, 5) + sol = solve_bvp(fun, bc, t, jnp.zeros((5, 2)), p=p, max_nodes=max_nodes) + jax.block_until_ready(sol.y) + print(f"=== DONE {index} status={int(sol.status)} ===", flush=True) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--max-nodes", nargs="+", type=int, default=[32, 128]) + parser.add_argument("--batch", type=int, default=32) + parser.add_argument("--repeat", type=int, default=5) + parser.add_argument("--cache-audit", action="store_true") + args = parser.parse_args() + + if args.cache_audit: + cache_audit(args.max_nodes[0]) + return + + print("case,max_nodes,cold_ms,run_ms,scipy_ms") + for max_nodes in args.max_nodes: + for name, factory in (("exp", exp_problem), ("shock", shock_problem)): + solve, scipy_solve = factory(max_nodes) + cold_ms, run_ms = cold_and_warm_ms(solve, args.repeat) + scipy_ms = 1e3 * statistics.median( + timeit.repeat(scipy_solve, repeat=args.repeat, number=1) + ) + print(f"{name},{max_nodes},{cold_ms:.1f},{run_ms:.3f},{scipy_ms:.3f}") + solve, _ = eigen_gradient(max_nodes) + cold_ms, run_ms = cold_and_warm_ms(solve, args.repeat) + print(f"eigen_grad,{max_nodes},{cold_ms:.1f},{run_ms:.3f},") + solve, _ = vmapped_exp(max_nodes, args.batch) + cold_ms, run_ms = cold_and_warm_ms(solve, args.repeat) + print(f"vmap_exp_b{args.batch},{max_nodes},{cold_ms:.1f},{run_ms:.3f},") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/results/2026-08-19_bvp-scaling-cpu.json b/benchmarks/results/2026-08-19_bvp-scaling-cpu.json new file mode 100644 index 0000000..d50e4e3 --- /dev/null +++ b/benchmarks/results/2026-08-19_bvp-scaling-cpu.json @@ -0,0 +1,68 @@ +{ + "date": "2026-08-19", + "machine": "i9-10900K (20 threads), CPU, float64", + "command": "JAX_PLATFORMS=cpu uv run python -m benchmarks.bvp_scaling --repeat 7", + "methodology": "every case calls the solve through a jitted wrapper: cold_ms is compile-plus-first-run, run_ms is the median warm call", + "linear_solve": "structured orthogonal factorization (tinydiffeq.babd cyclic reduction)", + "cache_audit": "JAX_EXPLAIN_CACHE_MISSES=1 JAX_LOG_COMPILES=1 --cache-audit: 0 XLA compiles and 0 executable-cache misses after the second solve; the single tracing event is the per-call eval_shape validation trace", + "growth_baseline": "kernels neoclassical growth (801->1489 nodes, tol=1e-10, max_nodes=1536): 17 ms warm as a direct (un-jitted) call vs scipy 17.4 ms, solutions agree to 4e-16; dense LU on the same problem took 304 ms", + "gpu_float32": "RTX 3090, 64-lane vmapped exp ensemble at max_nodes=32: 0.79 ms (~12 us/lane), all converged", + "rows": [ + { + "case": "exp", + "max_nodes": "32", + "cold_ms": "821.8", + "run_ms": "0.250", + "scipy_ms": "0.601" + }, + { + "case": "shock", + "max_nodes": "32", + "cold_ms": "681.9", + "run_ms": "0.450", + "scipy_ms": "6.323" + }, + { + "case": "eigen_grad", + "max_nodes": "32", + "cold_ms": "1108.4", + "run_ms": "0.532", + "scipy_ms": "" + }, + { + "case": "vmap_exp_b32", + "max_nodes": "32", + "cold_ms": "1462.9", + "run_ms": "1.459", + "scipy_ms": "" + }, + { + "case": "exp", + "max_nodes": "128", + "cold_ms": "830.9", + "run_ms": "0.470", + "scipy_ms": "0.593" + }, + { + "case": "shock", + "max_nodes": "128", + "cold_ms": "818.8", + "run_ms": "2.072", + "scipy_ms": "6.300" + }, + { + "case": "eigen_grad", + "max_nodes": "128", + "cold_ms": "1306.8", + "run_ms": "1.603", + "scipy_ms": "" + }, + { + "case": "vmap_exp_b32", + "max_nodes": "128", + "cold_ms": "1756.4", + "run_ms": "6.373", + "scipy_ms": "" + } + ] +} diff --git a/benchmarks/results/2026-08-19_bvp-scaling-cpu.md b/benchmarks/results/2026-08-19_bvp-scaling-cpu.md new file mode 100644 index 0000000..e6faac4 --- /dev/null +++ b/benchmarks/results/2026-08-19_bvp-scaling-cpu.md @@ -0,0 +1,49 @@ +# solve_bvp CPU scaling (2026-08-19) + +i9-10900K, float64, `JAX_PLATFORMS=cpu uv run python -m benchmarks.bvp_scaling --repeat 7`, +with the collocation system solved by the structured orthogonal +factorization in `tinydiffeq.babd` (level-batched cyclic reduction). +Every case calls the solve through a jitted wrapper: `cold_ms` is +compile-plus-first-run (fresh callables per case), `run_ms` is the median +of 7 warm runs with `jax.block_until_ready` inside the timed function. +scipy reference uses its finite-difference Jacobians. Warm-cache reuse +verified per the cache-audit protocol (`--cache-audit` under +`JAX_EXPLAIN_CACHE_MISSES=1 JAX_LOG_COMPILES=1`): changing `p` and the +mesh values triggers 0 XLA compiles on the second solve. + +`shock` at max_nodes=32 exhausts the node budget (status 1) and times the +run-to-failure path; `vmap_exp_b32` is a jitted 32-lane `vmap` of the +full solve, and `eigen_grad` is `jit(grad)` of an unknown-eigenvalue +solve via the implicit rule. + +| case | max_nodes | cold (ms) | warm (ms) | scipy (ms) | +|---|---|---|---|---| +| exp | 32 | 821.8 | 0.250 | 0.601 | +| shock | 32 | 681.9 | 0.450 | 6.323 | +| eigen_grad | 32 | 1108.4 | 0.532 | — | +| vmap_exp_b32 | 32 | 1462.9 | 1.459 | — | +| exp | 128 | 830.9 | 0.470 | 0.593 | +| shock | 128 | 818.8 | 2.072 | 6.300 | +| eigen_grad | 128 | 1306.8 | 1.603 | — | +| vmap_exp_b32 | 128 | 1756.4 | 6.373 | — | + +Warm jitted calls beat scipy on every case, including the tiny 5-node +exponential problem (0.25 ms vs 0.60 ms). Calling `solve_bvp` directly +from un-jitted Python instead adds ~3 ms of per-call wrapper work +(validation traces, pytree flattening, dispatch) — the compilation cache +still hits, but the wrapper Python re-runs each call. On an RTX 3090 in +float32, a 64-lane batch at max_nodes=32 solves in ~0.79 ms (~12 +us/lane). + +## Real-workload check: kernels neoclassical growth baseline + +The `kernels` repo's `neoclassical_growth_benchmark` (detrended +saddle-path BVP, `tol=1e-10`, initial mesh `linspace(0, 200, 801)`, +refined to 1489 nodes) rewired to this solver at `max_nodes=1536` +(float64, CPU): mesh evolution is node-for-node identical to scipy (1489 +nodes, 5 iterations, max rms 1.0e-10), paths agree to 4e-16, and the +warm solve — timed as a direct un-jitted call, so including the ~3 ms +wrapper overhead — takes **17 ms vs scipy's 17.4 ms**. A dense LU +factorization of the same padded system took 304 ms; the structured +factorization is what closes that gap. Compile is ~3 s, paid once; +changed calibrations reuse the compilation at warm speed. From ef4f9e683268a4d8caa46bfec5dbdbf9deda7321 Mon Sep 17 00:00:00 2001 From: Jesse Perla Date: Wed, 19 Aug 2026 21:25:59 -0700 Subject: [PATCH 5/5] docs: mark repo as AI-generated research code 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) Claude-Session: https://claude.ai/code/session_01UwxGiVZSEDw3pzMA7Zzb1g --- README.md | 189 ++++++++++++++++---------------------------------- docs/index.md | 36 ++++++---- 2 files changed, 83 insertions(+), 142 deletions(-) diff --git a/README.md b/README.md index c716e9a..1b402fd 100644 --- a/README.md +++ b/README.md @@ -7,28 +7,29 @@ [![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE) [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) -Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4, -adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs, and -fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and semi-explicit -index-1 SDAEs. Solves run in bounded `lax.scan` loops with static shapes and -compose with `jit`, `vmap`, forward mode, reverse mode, and -reverse-over-forward. Finite-state Markov simulation, probability forecasts, -and general fixed homogeneous linear solves (dense or matrix-free Krylov -exponential actions, after SciML's -[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/)) -round out the package. - -This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX -adaptation of Steinebach's method following SciML's -[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock) -implementation, and DAE algebraic roots delegate both the primal solve and -the implicit derivative to -[`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/). Use -[diffrax](https://docs.kidger.site/diffrax/) or +tinydiffeq is an unsupported research repo of vibe-coded ports of +well-established ODE, DAE, and SDE algorithms to JAX. Heavily AI-generated — +but the algorithms are well established, with +[SciML](https://docs.sciml.ai/DiffEqDocs/stable/), +[scipy](https://docs.scipy.org/doc/scipy/reference/integrate.html), and +[diffrax](https://docs.kidger.site/diffrax/) as the reference +implementations — so correctness and performance are often reasonable. The +method set is intentionally minimal, though the package is no longer +especially tiny. + +Fixed-step Euler/RK4, adaptive Tsit5, and linearly implicit Rodas5P for +stiff ODEs and index-1 DAEs; Euler–Maruyama, Milstein, and SRA1 for Itô SDEs +and SDAEs; a port of scipy's collocation solver for two-point BVPs with +unknown parameters; finite-state Markov chains and dense or Krylov linear +exponential actions. Every solve runs in bounded `lax` loops with static +shapes and composes with `jit`, `vmap`, forward mode, reverse mode, and +reverse-over-forward; iterative solves (BVP, DAE roots) differentiate +implicitly at the solution, never through the iterations. + +Use [diffrax](https://docs.kidger.site/diffrax/) or [SciML](https://docs.sciml.ai/DiffEqDocs/stable/) if you need general mass -matrices, fully implicit or higher-index DAEs, adaptive SDE stepping, events, -continuous solution objects, sparse/Krylov ODE/DAE stages, or specialized -adjoints. +matrices, fully implicit or higher-index DAEs, adaptive SDE stepping, +events, continuous solution objects, or specialized adjoints. ## Install @@ -36,19 +37,15 @@ adjoints. uv add tinydiffeq ``` -For GPU use, install the JAX accelerator build that matches your hardware, -for example: +For GPU use, add the JAX accelerator build matching your hardware, for +example `uv add tinydiffeq "jax[cuda13]"`. -```bash -uv add tinydiffeq "jax[cuda13]" -``` - -## Minimal example +## Example The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or -`(x, t, args, p)` — always in that order. `args` is pass-through data (not an -AD target by convention); `p` holds differentiable parameters, and the state -may be any pytree of same-dtype real floating arrays. +`(x, t, args, p)` — always in that order. `args` is pass-through data (not +an AD target by convention); `p` holds differentiable parameters, and the +state may be any pytree of same-dtype real floating arrays. ```python import jax @@ -68,125 +65,57 @@ sol = solve_ode( dt_0=0.1, controller=IController(rtol=1e-8, atol=1e-10), max_steps=512, - save_at=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)), # fixed output shape, -) # however many steps adapt -print(sol.xs) # states on the grid -print(sol.ok) # reached t_1 with every requested output valid? -``` - -`max_steps` is the internal attempt budget (accepted plus rejected steps), -not the number of returned times: `SaveAt` picks the endpoint, a fixed -interpolation grid, or the padded accepted-step prefix, so output shapes -never depend on how many steps the controller took. Omitted controller -tolerances follow the state dtype (`1e-4`/`1e-6` in float32, -`1e-7`/`1e-9` in float64). - -## SDEs with first-class noise - -`solve_sde` integrates diagonal-noise Itô SDEs with `EulerMaruyama` (strong -order 0.5), `Milstein` (1.0, commutative diagonal noise), or `SRA1` (1.5, -additive noise). An Ornstein–Uhlenbeck process under SRA1: - -```python -from tinydiffeq import solve_sde, SRA1 - -theta, sigma, n = 1.0, 0.5, 256 - - -def ou_drift(x): - return -theta * x - - -def ou_diffusion(x): - return sigma * jnp.ones_like(x) - - -sol = solve_sde( - ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, jnp.asarray(1.0), - key=jax.random.key(0), n_steps=n, + save_at=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)), ) +sol.xs # states on the grid, however many internal steps were taken +sol.ok # False if integration or a requested output failed ``` -The noise realization can also be passed explicitly — the same pytree -`sample_noise` would draw, now inspectable, storable data that is -differentiable like any other input: +`max_steps` bounds attempted internal steps (accepted plus rejected); +`SaveAt` fixes the output shape regardless of how many steps the controller +takes. Gradients go straight through the solve: ```python -x_0 = jnp.asarray(1.0) -noise = SRA1().sample_noise(x_0, jax.random.key(0), n, jnp.asarray(1.0 / n), x_0.dtype) -same_sol = solve_sde( - ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n -) # bit-identical to the key= call -d_endpoint_d_noise = jax.grad( - lambda noise: solve_sde( - ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n +def endpoint(p): + return solve_ode( + f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p, + dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12), + max_steps=512, ).xs -)(noise) -``` - -A fixed key (or fixed noise) pins the whole path, so gradients with respect -to `x_0`, `p`, and `noise` are pathwise derivatives under common random -numbers — the setup simulation-based estimators want. `vmap` over -trajectories with per-trajectory `x_0` and noise composes with `jit` and -`grad`. -## Semi-explicit DAEs +jax.grad(endpoint)(jnp.asarray(1.3)) # reverse mode +jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),)) # forward mode +``` -For a square index-1 system `dy/dt = f(y, z, t, args, p)` and -`0 = g(y, z, t, args, p)`: +An SDE ensemble — per-key noise, `vmap` over trajectories: ```python -from tinydiffeq import solve_semi_explicit_dae - - -def dae_f(y, z, t, args, p): - dy = p * z - return dy, {"flow": dy} +from tinydiffeq import solve_sde, SRA1 -def dae_g(y, z, t, args, p): - return z - y +def ou_drift(x): + return -x -dae_sol = solve_semi_explicit_dae( - dae_f, dae_g, Tsit5(), 0.0, 1.0, - jnp.asarray(1.0), jnp.asarray(0.5), - p=jnp.asarray(2.0), dt_0=0.1, - controller=IController(), max_steps=128, -) -print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"]) -``` - -`z_0` is a guess and is made consistent automatically. RK4 and Tsit5 restore -the algebraic root at every stage through `nlls-gram`, which also supplies -the root's implicit derivative; `Rodas5P()` instead performs one initial -consistency solve and then advances the block mass-matrix system with one -reused LU factorization per attempt — the stiff path. Stochastic -semi-explicit systems use `solve_semi_explicit_sdae` with `EulerMaruyama` or -`SRA1`. See the -[DAE](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) and -[SDAE](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) docs. +def ou_diffusion(x): + return 0.5 * jnp.ones_like(x) -## Gradients through the solve -```python -def endpoint(p): - return solve_ode( - f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p, - dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12), - max_steps=512, +def ou_path(key): + return solve_sde( + ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, jnp.asarray(1.0), + key=key, n_steps=256, save_at=SaveAt(steps=True), ).xs -jax.grad(endpoint)(jnp.asarray(1.3)) # reverse mode -jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),)) # forward mode + +paths = jax.vmap(ou_path)(jax.random.split(jax.random.key(0), 1000)) # (1000, 257) ``` -The step-size controller is wrapped in `stop_gradient` (accept/reject is -non-differentiable either way); states differentiate through the solver -stages on the realized, frozen mesh. See the -[docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design -contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE -noise semantics, and the package API. +SDEs with first-class differentiable noise, semi-explicit DAEs and SDAEs, +two-point BVPs, Markov chains, linear exponential solves, and the design +contracts (static shapes and `SaveAt`, AD through adaptive stepping, +failure-as-data) are in the +[docs](https://highdimensionaleconlab.github.io/tinydiffeq/). ## License diff --git a/docs/index.md b/docs/index.md index 9a19dfc..bfdf0e6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,16 +1,25 @@ # tinydiffeq -`tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE/DAE/SDAE -integrators and finite-state Markov simulators for JAX: fixed-step Euler and -RK4, adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 -DAEs, and fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and -SDAEs. Solves run in bounded `lax.scan` loops with static shapes and support -forward mode, reverse mode, and reverse-over-forward; adaptive ODE/DAE solves -can opt into a dynamic actual-work loop (`adaptive_loop="forward"`, no -reverse mode). Probability forecasts and general fixed homogeneous linear -solves use matrix powers, dense exponentials, or matrix-free Krylov actions; -see [Markov Chains](markov_chains.md) and -[Linear Exponential Solves](exponential.md). +tinydiffeq is an unsupported research repo of vibe-coded ports of +well-established ODE, DAE, and SDE algorithms to JAX. Heavily AI-generated — +but the algorithms are well established, with +[SciML](https://docs.sciml.ai/DiffEqDocs/stable/), +[scipy](https://docs.scipy.org/doc/scipy/reference/integrate.html), and +[diffrax](https://docs.kidger.site/diffrax/) as the reference +implementations — so correctness and performance are often reasonable. The +method set is intentionally minimal, though the package is no longer +especially tiny. + +Fixed-step Euler and RK4, adaptive Tsit5, and linearly implicit Rodas5P for +stiff [ODEs](ode.md) and index-1 [DAEs](dae.md); Euler–Maruyama, Milstein, +and SRA1 for Itô [SDEs](sde.md) and [SDAEs](sdae.md); a faithful port of +scipy's collocation [boundary-value solver](bvp.md) for two-point BVPs with +unknown parameters; [Markov chains](markov_chains.md) and +[linear exponential solves](exponential.md). Every solve runs in bounded +`lax` loops with static shapes and composes with `jit`, `vmap`, forward +mode, reverse mode, and reverse-over-forward; iterative solves (BVP, DAE +roots) differentiate implicitly at the solution, never through the +iterations. **Use [SciML](https://docs.sciml.ai/DiffEqDocs/stable/) or [diffrax](https://docs.kidger.site/diffrax/) instead if you need any of:** @@ -20,6 +29,8 @@ see [Markov Chains](markov_chains.md) and - adaptive SDE stepping (Brownian-bridge noise), full PID step-size control - events, root-finding, or backward-time integration - dense output objects or checkpointed/backsolve adjoints for long horizons +- multipoint or complex-valued boundary value problems, or BVP meshes beyond + a few hundred nodes ## Install @@ -122,6 +133,7 @@ jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),)) # forward mode are data leaves, so changing them never recompiles. Read next: [ODEs](ode.md), [SDEs](sde.md), [Semi-Explicit DAEs](dae.md), -[SDAEs](sdae.md), [Markov Chains](markov_chains.md), +[SDAEs](sdae.md), [Boundary Value Problems](bvp.md), +[Markov Chains](markov_chains.md), [Linear Exponential Solves](exponential.md), and the [API Reference](api.md).