Module Jera.Ode

The type for one-step methods with tags 'k over states 'y and times of dtype 't. A method with more tags can be used as one with fewer, by coercion: (Ode.tsit5 :> ([ `Formula ], _, _, _) Ode.t).

euler is the explicit Euler method: order 1, one field evaluation per step.

rk4 is the classical Runge–Kutta method: order 4, four evaluations per step.

ssprk3 is the three-stage strong-stability-preserving method of Shu and Osher (1988): order 3, three evaluations per step. A step is a convex combination of Euler steps, so it keeps any convex bound Euler keeps under the same step size, such as a total-variation bound of a hyperbolic discretisation.

bs3 is Bogacki and Shampine's (1989) method of order 3 with an embedded order 2: three evaluations per step, its last stage the next step's first.

tsit5 is Tsitouras' (2011) method of order 5 with an embedded order 4: six evaluations per step, its last stage the next step's first.

dopri5 is Dormand and Prince's (1980) method of order 5 with an embedded order 4: six evaluations per step, its last stage the next step's first.

tableau ~a ~b ~c is the explicit Runge–Kutta method of Butcher tableau (a, b, c) with s stages: stage i evaluates the field at t + c.(i) h on y + h Σ_(j<i) a.(i).(j) k_j, and the step is y + h Σ_i b.(i) k_i. Its coefficients are given in float64 and rounded once to the time's and each leaf's dtype.

Raises Invalid_argument unless s ≥ 1, b and c have s elements, a has s rows of which row i has i elements, every coefficient is finite, and b sums to 1 within rounding.

kvaerno5 ~mass ~linear derivative is Kværnø's (2004) singly diagonally implicit method of order 5 with an embedded order 4, for stiff fields: derivative t y dy is the field's Jacobian at (t, y) applied to dy, such as snd (Rune.jvp' (f t) y dy). mass is the linear map M of M y' = f t y, the identity by default; a singular M makes the problem differential-algebraic.

Method. Seven stages, the first explicit and the others implicit with the diagonal 0.26. Each stage is solved by simplified Newton on M − 0.26 h J, J the Jacobian at the step's start, which linear prepares once per attempted step: Linear.dense materialises and factors it once, and every stage and Newton iteration reuses the factors. A stage whose corrections do not shrink to 0.03 of tol in ten iterations, or whose linear solve fails, rejects the attempt. Error. The step is the last stage and its error the difference from the sixth, both at c = 1: the embedded formula is stiffly accurate like the step, so the estimate is defined on algebraic components. Stability. L-stable, so the step follows the tolerance on a stiff field instead of its fastest eigenvalue. A DAE is solved when its index is 1, when the equations in M's left null space determine the variables in its null space; a higher index shows as a collapsing step and ends Stalled. y0 must be consistent, f t0 y0 in the range of M: from an inconsistent y0 the first steps project onto the constraint. Cost. Per attempt, one preparation of linear (n applications of derivative and a factorisation for Linear.dense), and one evaluation of the field and one linear solve per Newton iteration. Derivative. In the answer each accepted stage is stated as a zero of its stage equation through Rune.root, solved by linear, so the derivative is the accepted steps' with their stages' implicit derivatives.

The type for times.

The type for fields: f t y is the derivative of y at the scalar time t, a value of y's structure, dtypes and shapes.

march y m ~steps f ~at y0 is the state at each time of at, stacked on a new leading axis of each leaf, y0 first at at.(0). The interval from at.(i) to at.(i+1) takes steps steps of m: step j starts at at.(i) + j h with h = (at.(i+1) − at.(i)) / steps, or at at.(i+1) where that rounds past it. Every step but the last has size h; the last has size at.(i+1) minus its start, never negative, and ends at at.(i+1) itself. at may decrease, to march back.

Error. A method of order p has global error O(h^p) for a smooth field. Stability. An explicit method is stable only while h times the field's eigenvalues lies in its stability region, so a stiff field needs a step far below its accuracy's. Cost. The method's evaluations per step, once each step; a method whose last stage is the next step's first evaluates it once. Derivative. The composition's: in the initial state, in the times of at, and in every tracked value the field reads. Reverse mode keeps one state per time of at and recomputes each interval while it reverses it.

Raises Invalid_argument if steps < 1, if at is not a non-empty 1-D tensor, if it is not strictly monotone, or if f returns a value of another structure, dtype or shape than its state.

solve y m ~tol ~budget f ~t0 ~t1 y0 is the state at t1 of the solution from y0 at t0, scalars; t1 may equal or precede t0. At t1 = t0 it is y0, converged with a zero error.

Raises Invalid_argument if budget < 1, if t0 or t1 is not a scalar, or as march does for a field of another structure.

sample y m ~tol ~budget f ~at y0 is the state at each time of at, stacked on a new leading axis of each leaf, y0 first at at.(0), with budget attempts across all of them. The steps land on every time of at, so no state is interpolated. Its error is stacked like its value: at each time, per component, the sum of the magnitudes of the local estimates of the accepted steps before it. Times that are not strictly monotone end the lane Stalled before any step, and its report names the first.

Raises Invalid_argument if budget < 1, if at is not a non-empty 1-D tensor, and as march does for a field of another structure.

path y m ~tol ~budget f ~t0 ~t1 y0 is the solution from y0 at t0 to t1 as a function of time: a series of budget pieces, one per accepted step in increasing time, each the step's continuous extension. The pieces past the last step are empty at its end, so the path's domain is the span between t0 and t1, whichever way the solve runs. At a step's end the path is that step's state, up to the rounding of its series.

t0 = t1 ends the lane Stalled, since a path needs a piece of positive width.

Error. The error is a series of degree 0 on the same breaks: on each piece, per component, the sum of the magnitudes of the local estimates of the accepted steps through it. Stability. As solve's. Cost. A solve's search, then its accepted steps taken again with the tracked field, whatever the budget, which sizes the pieces' storage. Each piece costs p + 1 combinations of the stages for an extension of order p. Derivative. The accepted steps' and their extensions', through the initial state, t0, t1 and every tracked value the field reads; the breaks move with t0 and t1.

Raises Invalid_argument if budget < 1, if t0 or t1 is not a scalar, if a leaf of the state is not a float tensor, or as march does for a field of another structure.

event y m ~tol ~budget f ~event ~t0 ~t1 y0 is the solve from y0 at t0 toward t1 that ends at the first crossing of an event: the time, the state there, and the flat index of the component of event that crossed. Each component of event t y is one event, so several events are one tensor.

Method. A component's sign at an accepted step's end is its last non-zero one, so a component that touches zero and turns back does not cross, and a zero at t0 has no sign. The first accepted step across which a component's sign changes holds the crossing. In it, a bracketing search, Root.bracket's, finds where each such component takes its new sign on the step's continuous extension (see Paths) to the resolution of the time's dtype, which tol does not set, and the earliest wins. The time returned is the end of the final bracket past the crossing, where the component has its new sign, so a solve restarted there starts on that side, and the component's new sign tells the crossing's direction. A crossing that should not stop the solve, or one in a direction to ignore, is a restart from the returned time, one solve per crossing. To restart on the crossing's old side, as a bounce that reflects a velocity does, set the component to zero: a zero at t0 is no crossing. A lane without a crossing after t0 up to t1 included converges with (t1, y t1, −1): the index tells the two apart, whichever way the solve runs. At t1 = t0 the answer is (t0, y0, −1), converged, so a compiled loop of restarts holds a finished lane where it stopped. A step holding an even number of crossings of a component shows none, so a crossing narrower than the field's steps can be missed or reported out of order: an event whose sign changes once per crossing, or path and a finer search, resolves it.

Error. The time's is half the final bracket, about one unit in the last place of the time; the state's, per component, the sum of the magnitudes of the local estimates of the accepted steps through the one that holds the crossing; the index's is 0. A crossing component whose bracket does not converge ends the lane Stalled. Cost. A solve's search with one evaluation of event per attempted step, then the crossing's search, at most 2b + 1 iterations for a time of b bits, each of which evaluates event once per component; the answer takes the accepted steps again. Derivative. The time is stated as a zero of the crossing component R t = event t (y t) on the tracked step, so its derivative is −∂R/∂θ / ∂R/∂t, through the initial state, t0 and every tracked value the field and event read; the state's is the extension's at that time. Without a crossing, the derivative is solve's.

A compiled function returns the answer as a value of this structure, s the state's:

Solution.ptree
  (Nx.Ptree.v (fun c (t, y, i) ->
       let open Nx.Ptree.Walk in
       let t = index c 0 tensor t in
       let y = index c 1 (structure s) y in
       let i = index c 2 tensor i in
       (t, y, i)))

Raises Invalid_argument if budget < 1, if t0 or t1 is not a scalar, if event has no component, or as march does for a field of another structure.

delay y m ~tol ~budget ~pieces f ~lags ~history ~at y0 is the state at each time of at, increasing, of the solution of y' t = f t (y t) (y (t − τ)) for the constant lags τ of lags, 1-D, from y0 at at.(0), stacked as sample stacks them. f's third argument holds the delayed states stacked on a leading axis, one per lag; history s is the state at the scalar time s ≤ at.(0), which delay calls once per lag and stacks. y0 may differ from history at.(0): the solution then jumps at at.(0), and f reads the delayed state at at.(0) from history in the step that ends there and as y0 in the step that starts there.

Method. An explicit embedded method, which the two tags together mark, whose steps never exceed the smallest lag, so every delayed state is read from an accepted step's continuous extension (see Paths) or from history. The steps land on the breakpoints at.(0) + Σ_j k_j τ_j with 1 ≤ Σ_j k_j ≤ p, p the method's order, where the solution's derivative of order Σ_j k_j can jump, or of order 1 + Σ_j k_j when y0 = history at.(0). The order of a delay solve is min(p, q + 1), q its extension's order (Bellen and Zennaro, 2003). Error. As sample's. A lag that is not positive, or the largest lag reaching back past the last pieces accepted steps, ends the lane Stalled, and its report names the count that would hold it. Times of at that do not increase end the lane Stalled too. Stability. As solve's. Cost. Each attempt costs the method's evaluations, its first stage among them, since the field may jump where one step ends and the next starts. Each stage reads lags delayed states, each a binary search of the pieces pieces and a series of the extension's degree; the carry holds pieces pieces, so a compiled reverse keeps budget × pieces. Derivative. The answer reads its own tracked pieces, so the derivative reaches the delayed states, history, the lags and every tracked value f reads.

Raises Invalid_argument if budget < 1, if pieces < 1, if at is not a non-empty 1-D tensor, if lags is not a non-empty 1-D tensor, if history returns a value of another structure, dtype or shape than y0, or as march does for a field of another structure.