Laguz

Ordinary differential equation (ODE) solvers for OCaml, on Nx tensors.

Laguz solves initial value problems

dy/dt = f(t, y),   y(t0) = y0

where the state y is an Nx tensor of any shape and any float dtype (float32, float64, …) and time t is a plain float. All stepping is pure tensor computation — no mutation, no C bindings.

---

⚠️ Work in progress — not for public consumption

This library is a rough, early prototype. Please do not build on it yet.

• The code is a first implementation pass (generated in one shot and only lightly reviewed since). The numerics have not been independently verified.
• The API is unstable and will change without notice as design feedback lands. There is no semantic versioning, no opam release, no CI.
• There is no benchmark evidence against existing solvers (Owl_ode, sundialsml) yet, and performance has not been a priority.
• Only a single example ships, and the documentation story is still incomplete.

Consider this an exploratory spike. If you are looking for a production ODE solver for OCaml, look elsewhere (for now). Feedback is still welcome — see Contributing.

---

What's here

The full v1 feature set from the design doc is implemented:

Feature	Status
Explicit Runge–Kutta steppers: euler, heun, rk4 (fixed-step), bs23, dopri5, tsit5 (adaptive embedded pairs, FSAL)	✅ implemented
Adaptive step-size control: [Elementary|Pi] controllers, automatic initial step (Hairer's heuristic)	✅ implemented
Dense output (cubic Hermite per accepted step) → eval, ~output:(At ts)`	✅ implemented
Event detection (terminal / non-terminal, direction, bisection refinement)	✅ implemented
Status reporting instead of exceptions (Success, Terminated_by_event, Max_steps_reached, Step_size_too_small, Nan_in_state)	✅ implemented
Step module for manual, interactive stepping loops	✅ implemented
Backward integration (t1 < t0)	✅ implemented
Convergence-order, adaptive-accuracy, dense-output, event, and validation tests (windtrap)	✅ implemented
Example: Lorenz butterfly, batch of 50 initial conditions	✅ implemented
Rune companion (laguz.rune: jit/grad glue, adjoints)	❌ deferred
Stiff solvers, symplectic integrators, SDEs, DAEs	❌ out of scope
opam release, CI, benchmarks vs. Owl_ode	❌ not started

Quick start

open Nx
open Laguz

let y0 = full float32 [| 1 |] 1.0
(* dy/dt = -y *)
let rhs _t y = neg y

let sol = solve ~output:`Every_step ~rhs ~t0:0. ~t1:1. ~y0 ()

let () =
  Printf.printf "y(1) = %g (exact: %g)\n"
    (item [ 0 ] (y_final sol))
    (Float.exp (-1.))

Steppers

Stepper	Order	Embedded	FSAL	Use
euler	1	—	—	demo/teaching, fixed-step
heun	2	yes (Euler)	—	fixed/adaptive
rk4	4	—	—	the classic, fixed-step
bs23	3(2)	yes	yes	adaptive, loose tolerances
dopri5	5(4)	yes	yes	adaptive workhorse (scipy RK45, matlab ode45)
tsit5	5(4)	yes	yes	adaptive default; usually the best efficiency

Default settings: stepper = tsit5, rtol = 1e-6, atol = 1e-9, controller =Pi``, max_steps = 100_000. Use `float64` for tolerances tighter than ~1e-8.

Save modes and dense output

type output = [ `Final | `Every_step | `At of float array ]

• `Final — keep only y(t1): O(1) extra memory.
• `Every_step — record every accepted (t, y) pair.
• `At ts — evaluate at arbitrary times via cubic Hermite dense output (ts may be unsorted).

solve never raises on solver failure: divergence, NaN, or a step-size underflow return a partial trajectory plus a status field. Invalid_argument is reserved for caller bugs (bad tolerances, ~dt missing for fixed-step methods, RHS shape changes, out-of-interval eval, …).

Events

let ground =
  { f = (fun _ y -> item [ 0 ] y); direction = `Falling; terminal = true }

Zero crossings are detected on accepted steps and refined by bisection on the dense interpolant. Recorded in events sol : 'b event_record array, and a terminal event ends the integration with status = Terminated_by_event i.

Manual stepping (Step)

For interactive simulations, RL environments, live plots — a fold-style loop with no output collection:

let st = Step.init ~stepper:dopri5 ~rhs ~t0:0. ~t1:1. ~y0 () in
let rec loop st = if Step.is_done st then st else loop (Step.step st) in
let st = loop st in
Step.y st  (* final state *)

Example: Lorenz butterfly

example/lorenz.ml integrates a batch of 50 random initial conditions of the Lorenz system with dopri5, samples the whole batch at fixed times via dense output, and writes traj{i} files:

dune exec example/lorenz.exe

States are tensors of any shape, so the whole batch is integrated at once as a [|50; 3|] state with a vectorized RHS — a built-in feature of the tensor state representation, not a loop.

Building and testing

Requires OCaml ≥ 5.5 and an Nx checkout in the same dune workspace.

dune build
dune runtest        # windtrap test suite: convergence orders, adaptive
                    # accuracy, dense output, events, statuses, validation
dune exec example/lorenz.exe

Design notes

The architecture borrows from diffrax and OrdinaryDiffEq.jl, scaled down: steppers + controllers + interpolation + save modes as separate concerns (Tableau, Stepper, Controller, Interpolant, Driver — see lib/).

• Pure by construction. Stepping is branch-free tensor arithmetic, so solves are compatible with Rune.grad/Rune.jit down the road (adjoint sensitivities, neural ODEs). The per-step function is written in the jit-friendly (t, y, h) → (y₁, ŷ₁, f₁) shape from day one.
• Fixed-step is a special case. The same engine drives fixed-step and adaptive integration; error control is simply disabled.
• Rune interop is designed, not built. The core stays nx-only; a laguz.rune companion (scan-shaped fixed-step driver, detach-in-step, later continuous adjoints) is deferred.

Contributing

This library is at the prototype stage. If you'd like to help, useful first contributions are: independent verification of the numerics (the convergence-order tests are the place to start), benchmarks against Owl_ode, more examples (van der Pol, Kepler, harmonic oscillator), and the deferred laguz.rune companion. Please open an issue before sending anything larger than a typo fix.

License

ISC — see LICENSE.md.