aadc-ode

aadc-ode

ODE integration with exact adjoints on aadc-ng.

Two families of solvers:

C++ kernel replay (on tape) — the RHS f(y, p) is recorded once into its own tape, and each step records a compute block that replays it.

Python discrete adjoint — adaptive integration forward, then a backward sweep computes exact dJ/dp via vector-Jacobian products from the AAD kernel.

What "exact" means here, because it is the one place a knowledgeable user could be misled: exact for the discrete map with the recorded step sequence. The accepted steps h_list are frozen and the adaptive controller is not differentiated — that is the semantics of the cited paper, not an approximation introduced by this port. It is also why the backward sweep needs the forward solve's time grid and cannot invent one:

t_list, x_list, h_list, k_stages = aadc_ode.adaptive_rk45(f, x0, (t0, tf))
dJdp = aadc_ode.discrete_adjoint(aad_rhs, x_list, h_list, k_stages,
                                 aadc_ode.DORMAND_PRINCE_45, p,
                                 dJdx_T=dJdx_T,
                                 t_list=t_list)     # <- pass it

Pass t_list (or t0= if that is all you have). Omitting both reconstructs the stage times from 0.0, which is correct only for an autonomous RHS or a solve that started at t=0 — and wrong, silently, otherwise. It now warns; t0=0.0 is how an autonomous RHS says so and silences it.

AadRhs records the RHS once and replays that one recording at every stage of every step, so a plain if x[0] > 0: inside the RHS freezes its branch at the recording point and uses it at states arbitrarily far away — wrong values and flipped derivatives. Recording therefore raises if the tape reports any active-to-passive conversion; use iif / a masked blend, or allow_passive=True if the branch really is constant over the whole trajectory.

The discrete adjoint is the method from:

R. Martins, E. Lakshtanov. A C++ implementation of the discrete adjoint sensitivity analysis method for explicit adaptive Runge-Kutta methods enabled by automatic adjoint differentiation and SIMD vectorization. Applied Mathematics and Computation, 2025. arXiv:2410.01911

src/bdf.{hpp,cpp}            C++ library: KernelBlock, LUSolveBlock, drivers
src/bdf_selftest.cpp         native checks, incl. the two Python cannot express
bindings/aadcbdf_pybind.cpp  ONE type_caster — where gradients are won or lost
packaging/aadc_ode/          wheel: C++ extension + Python solvers (rk, adjoint)
tests/                       every assertion is on a DERIVATIVE
ci/build.sh · ci/build.ps1   compile + link (no code generation step)

Build and test

. ci/pins.env && SDK_VER="$SDK_VERSION" SDK_SRC_VERSION="$SDK_SRC_VERSION" \
  CORE_WHEEL_VERSION="$CORE_WHEEL_VERSION" bash ci/build.sh python3.10 python3.11 python3.12 python3.13 python3.14
pip install dist/aadc_ode-*.whl
pytest tests/ -q

A requested interpreter that is not installed is a fatal error (matching ci/build.ps1). For a local run that should just build whatever is present, set STRICT=0 to skip missing interpreters instead.

Windows: powershell -File ci\build.ps1 -Pythons 3.10,3.11,3.12,3.13,3.14, from inside an MSVC environment. CI covers linux-x86_64, linux-aarch64, macos-arm64 and windows-x64.

Use

import aadc, aadc_ode

lam = [1.0, 10.0, 100.0, 1000.0]

# 1. Record the RHS once, in ordinary Python.
k = aadc_ode.Kernel()
states, _ = k.begin(n_states=4)
k.end([-lam[i] * states[i] for i in range(4)])

# 2. Record the loop onto an outer tape. 100 steps -> 100 blocks, ONE kernel.
f = aadc.Functions()
f.start_recording()
y0   = [aadc.idouble(1.0) for _ in range(4)]
args = [v.mark_as_input() for v in y0]
xT   = aadc_ode.bdf_solve(k, y0, dt=1e-3, n_steps=100, jac_lag=10)
cost = sum((x * x for x in xT[1:]), xT[0] * xT[0])
out  = cost.mark_as_output()
f.stop_recording()

# 3. Replay and differentiate, at any point — not just the recording point.
ws = f.create_workspace()
for a in args:
    ws.set_val(a, 1.0)
ws.forward()
ws.set_diff(out, 1.0)
ws.reverse()
print([ws.diff(a) for a in args])

begin hands back copies of the library-owned input variables, and that direction is the whole design. An aadc-ng tape slot is keyed by the address of its variable, and this binding marshals Real by value — so a mark_states(y) taking your list would mark throwaway copies while you kept using the originals. Zero gradients, nothing raised. With begin, your expressions sit downstream of the real inputs and the adjoint reaches them.

Design notes

Four properties that affect how you can use this package, rather than how it came to be:

Lane-generic blocks. The compute blocks are ComputeBlockT, so a tape containing them runs on a vector workspace as well as a scalar one. The avx2_lanes self test drives four AVX2 lanes with four different initial conditions and requires them to reproduce four scalar runs bitwise:

PASS avx2_lanes    4 lanes vs 4 scalar runs, worst |delta| = 0.000e+00

One reverse sweep per step, not one per state. A kernel block seeds every output with the outer adjoint and reverses once, computing x̄ += Jᵀȳ directly instead of assembling a full Jacobian. On a 27-state model that is one sweep per step rather than 27.

Thread-safe replay. Each thread gets its own kernel replay workspace. The concurrent_replay self test runs eight threads against one tape and requires bitwise agreement with the sequential run.

The frozen Jacobian is an option, and it is an approximation. The semi-implicit step's cached 1 - dt·J_ii is a record-time constant: replaying with bumped inputs reuses the Jacobian captured while recording. That is the standard semi-implicit trade — the Jacobian is a preconditioner, not part of the answer — but it is an approximation to the derivative of the scheme. freeze_jacobian=False turns it off; jac_lag=1 refreshes it every step.

What is not differentiated

luSolve takes the LU factorisation and pivot vector as passive data and differentiates only the right-hand side. That is the right split for a BDF step — the iteration matrix is a preconditioner rebuilt from the cached Jacobian, not a differentiable input — but it is a real limitation for anyone wanting dx/dA. Adding it means differentiating the factorisation itself (Ā = -x̄ xᵀ and friends), which is a separate piece of work.

Testing philosophy

Every test asserts a derivative. A binding that passivates Real, a second libaadc-ng image, or a block whose reverse never reaches the outer tape all produce values correct to the last digit and gradients of exactly zero — so a value check proves nothing.

The stiff adjoints are checked against a closed form, not finite differences. By T = 0.1 the λ=1000 component has decayed to ~1e-60, and a central difference of a quantity that small is pure cancellation noise; FD gets that adjoint wrong by 0.15%, which would either fail a correct implementation or force a tolerance loose enough to accept a wrong one.

The avx2_lanes and concurrent_replay checks are native C++ and bound as aadc_ode.selftest(). They have to be: aadc_ng's own Python external-function block is scalar-only and its callbacks are Jacobian-supplying, so a Python-only suite could not exercise either — and those are precisely the capabilities this package adds.

Packaging

Per-interpreter (pybind11 is not abi3), on top of the abi3 aadc_ng, bridged by the _aadc_ng._C_API capsule. This wheel ships no libaadc-ng: it links the copy inside the installed aadc_ng wheel and finds it again by RPATH ($ORIGIN/../aadc_ng, @loader_path/../aadc_ng, and on Windows by import order — delvewheel is deliberately not run).

Recording state is thread-local inside libaadc-ng, so a second image gives the extension and the aadc module private recording contexts. For this package that failure would be total: every entry point either opens a recording or records a block onto one someone else opened. test_exactly_one_library_image asserts it directly.

Unlike aadc-example-pybind11 this build does not define AADCNG_ALLOW_TO_PASSIVE_BOOL. That macro enables active-to-passive bool conversions across the whole translation unit, and a package whose entire purpose is keeping an exact adjoint through a stiff solver should not switch on implicit passivation to compile itself. bdf.hpp includes only <aadcNG/idoubleNG.h> — never the <aadcNG/aadcNG.h> umbrella — which keeps the deleted iboolNG::operator bool() doing its job.

Licence and support

aadc-ode is distributed under the MatLogica EULA and depends on the aadc core, which runs in Community Edition unless licensed — non-commercial and academic use only. The full licence and the third-party notices travel inside the wheel, under aadc_ode-<version>.dist-info/licenses/.

For commercial licensing, advanced features or support: matlogica.com · info@matlogica.com