RasmahRasmah

Differentiating Through ODEs and Callbacks

Overview

Many engineering quantities are the result of solving an ordinary differential equation (ODE) and then post-processing the solution: a cooling curve, a discharged battery, an orbit, a structural transient. When you want to optimize or fit such a system, you need the gradient of a loss with respect to the ODE’s inputs — its initial conditions, its parameters, and, when the system has discrete events, the parameters and timing of those events.

Rasmah makes this easy with solve_ode (a native Runge–Kutta solver) and ode_gradient, which differentiates through the solution with your choice of sensitivity algorithm:

method direction meaning
DiscreteAdjoint reverse “discretize then differentiate” (AD through the solver steps)
ContinuousAdjoint reverse “differentiate then discretize” (backward adjoint ODE)
ForwardSensitivity / ForwardDiffAdjoint forward forward sensitivity / forward-mode AD

Discrete events are handled by ODECallback, and their trigger times and closed-over parameters are differentiable too. Along the way we verify every derivative against the closed-form solution.

The model: Newton cooling

A hot body at temperature u cools at a rate proportional to its temperature,

\[ \frac{du}{dt} = -k\,u, \qquad u(0) = u_0, \]

whose exact solution is u(t) = u₀ e^{-kt}. The ODE “vector field” is a Julia function f(u, p, t) where p holds the parameters (here just k):

using Rasmah
using LinearAlgebra

f(u, p, t) = [-p[1] * u[1]]     # p = [k] : Newton cooling
u0 = [1.0]                       # initial temperature
tspan = (0.0, 2.0)               # cool for 2 seconds
k = [0.5]                        # decay rate
1-element Vector{Float64}:
 0.5

Forward solve

solve_ode integrates with a Runge–Kutta method (here RK4) and returns the whole trajectory; .u[end] is the state at the final time.

sol = solve_ode(f, u0, tspan, k; method = RK4(), steps = 200)
u_end = sol.u[end][1]
println("u(2.0) = ", round(u_end; digits = 5), "   (exact e⁻¹ = ", round(exp(-1.0); digits = 5), ")")
u(2.0) = 0.36788   (exact e⁻¹ = 0.36788)

Gradients: initial condition and parameter

Take the loss L = u(2.0). Its exact gradients are

\[ \frac{\partial L}{\partial u_0} = e^{-kT} = e^{-1}, \qquad \frac{\partial L}{\partial k} = -T u_0 e^{-kT} = -2 e^{-1}. \]

ode_gradient needs the partial ∂L/∂u(T) (here [1.0]) and returns the total gradients w.r.t. u0 and p. Compare two reverse methods against the analytic values:

dL_du_end = [1.0]   # L = u(T)[1]
for sensealg in (DiscreteAdjoint(), ContinuousAdjoint())
    du0, dp = ode_gradient(f, u0, tspan, k, dL_du_end, nothing;
                           sensealg = sensealg, method = RK4(), steps = 200)
    println(nameof(typeof(sensealg)), ":  ∂L/∂u0 = ", round(du0[1]; digits = 5),
            ",  ∂L/∂k = ", round(dp[1]; digits = 5))
end
println("exact:                 ∂L/∂u0 = ", round(exp(-1.0); digits = 5),
        ",  ∂L/∂k = ", round(-2 * exp(-1.0); digits = 5))
DiscreteAdjoint:  ∂L/∂u0 = 0.36788,  ∂L/∂k = -0.73576
ContinuousAdjoint:  ∂L/∂u0 = 0.36788,  ∂L/∂k = -0.73576
exact:                 ∂L/∂u0 = 0.36788,  ∂L/∂k = -0.73576

Both methods reproduce the closed-form gradients. DiscreteAdjoint is exact on the discretized problem; ContinuousAdjoint solves the adjoint ODE λ' = k λ backward and agrees to the integration tolerance.

Adding a discrete event (a callback)

Suppose a heater injects an impulse θ at time t_c, so the temperature jumps u ← u + θ. The solution becomes

\[ u(T) = u_0 e^{-kT} + \theta\, e^{-k(T - t_c)}, \]

and the trigger time itself matters: a later impulse has less time left to decay. An ODECallback bundles the affect (the jump), its closed-over parameters θ, and its trigger time t.

affect(u, θ) = [u[1] + θ[1]]     # add the impulse θ
t_c, θ = 1.0, [0.3]              # fire at t = 1.0
cb = ODECallback(affect, θ, t_c)

u_cb = solve_ode_callbacks(f, u0, tspan, k, [cb]; method = RK4(), steps = 200)[1]
exact_cb = exp(-1.0) + 0.3 * exp(-0.5)
println("u(2.0) with impulse = ", round(u_cb; digits = 5), "   (exact ", round(exact_cb; digits = 5), ")")
u(2.0) with impulse = 0.54984   (exact 0.54984)

Gradients of the callback: parameters and time

The exact sensitivities of L = u(2.0) w.r.t. the impulse and its timing are

\[ \frac{\partial L}{\partial \theta} = e^{-k(T - t_c)} = e^{-0.5}, \qquad \frac{\partial L}{\partial t_c} = k\,\theta\,e^{-k(T - t_c)}. \]

There are two routes. Differentiate through the discrete solve (any AD):

loss(v) = solve_ode_callbacks(f, [v[1]], tspan, [v[2]],
                              [ODECallback(affect, [v[3]], v[4])]; steps = 200)[1]
g = gradient(loss, [1.0, 0.5, 0.3, 1.0])   # [u0, k, θ, t_c]
println("AD: ∂L/∂θ = ", round(g[3]; digits = 5), ",  ∂L/∂t_c = ", round(g[4]; digits = 5))
AD: ∂L/∂θ = 0.60653,  ∂L/∂t_c = 0.09098

Or use the continuous adjoint with jump conditions (ode_gradient_callbacks), which returns the gradient w.r.t. each callback’s parameters and trigger time analytically:

du0, dp, dθs, dts = ode_gradient_callbacks(f, u0, tspan, k, [cb], [1.0];
                                           sensealg = ContinuousAdjoint(), steps = 200)
println("adjoint: ∂L/∂θ = ", round(dθs[1][1]; digits = 5),
        ",  ∂L/∂t_c = ", round(dts[1]; digits = 5))
println("exact:   ∂L/∂θ = ", round(exp(-0.5); digits = 5),
        ",  ∂L/∂t_c = ", round(0.5 * 0.3 * exp(-0.5); digits = 5))
adjoint: ∂L/∂θ = 0.60653,  ∂L/∂t_c = 0.09098
exact:   ∂L/∂θ = 0.60653,  ∂L/∂t_c = 0.09098

The continuous adjoint applies the jump conditions λ(t_c⁻) = ∂affect/∂uᵀ λ(t_c⁺) and accumulates ∂L/∂θ and ∂L/∂t_c as it integrates backward — the “differentiate then discretize” counterpart to the AD route.

Next steps