Optimal control and parameter estimation on GPUs¶
Tutorial 1 simulated dynamics forward. This tutorial solves the optimal control problem associated with the same dynamics: choose the control so the system arrives where we want it, at least cost, as a nonlinear program.
ExaModels.jl turns the model into GPU-ready evaluation and derivatives; MadNLP.jl solves it. Every model is written out in full: formulating is what this tutorial teaches.
Setup¶
On the CPU we solve with Ipopt, through NLPModelsIpopt, because Ipopt
is the reference interior-point code in this field and is the baseline a
CPU number should be measured against. On the GPU we use MadNLP. Both
consume the same ExaModels model, and both report objective, iter and a
status. The device comparison later is therefore a comparison of two
solvers as well as two devices.
using CUDA, CUDSS, ExaModels, MadNLP, MadNLPGPU, NLPModelsIpopt
using Plots
# Vector output: plots stay sharp at any zoom, and the notebook stays
# readable when rendered to HTML.
gr(fmt = :svg)
Plots.GRBackend()
A warm-up solve¶
Run the next cell now, then keep listening. The first interior-point solve on the GPU in a fresh session compiles the entire solver stack, which takes roughly two to three minutes. Tutorial 1 showed the small version of this effect. It is paid once per session, and the second solve takes a fraction of a second. The problem below is deliberately trivial. Its only job is to make the compiler do all of its work now, rather than in the middle of an exercise. Each new model structure later adds a few seconds of compilation on its first solve. That part is unavoidable, but it is seconds rather than minutes.
function warmup(backend)
wc = ExaCore(; backend = backend)
@add_var(wc, x, 100; start = 0.5)
@add_obj(wc, (x[i] - 1.0)^2 for i = 1:100)
@add_con(wc, x[i] + x[i+1] for i = 1:99; lcon = -10.0, ucon = 10.0)
m = ExaModel(wc)
return backend === nothing ? ipopt(m; max_iter = 600, print_level = 0).status :
madnlp(m; max_iter = 600).status
end
(cpu = warmup(nothing), gpu = warmup(CUDABackend()))
******************************************************************************
This program contains Ipopt, a library for large-scale nonlinear optimization.
Ipopt is released as open source code under the Eclipse Public License (EPL).
For more information visit https://github.com/coin-or/Ipopt
******************************************************************************
┌ Warning: Your Quadro GV100 GPU (compute capability 7.0) is not fully supported by CUDA 12.9.0.
│ Some functionality may be broken. Ensure you are using the latest version of CUDA.jl in combination with an up-to-date NVIDIA driver.
│ If that does not help, please file an issue to add support for the latest CUDA toolkit.
└ @ CUDACore ~/.julia/packages/CUDACore/OIFhX/lib/cudadrv/state.jl:231
This is MadNLP version v0.10.1, running with cuDSS v0.8.0
Number of nonzeros in constraint Jacobian............: 198
Number of nonzeros in Lagrangian Hessian.............: 100
Total number of variables............................: 100
variables with only lower bounds: 0
variables with lower and upper bounds: 0
variables with only upper bounds: 0
Total number of equality constraints.................: 0
Total number of inequality constraints...............: 99
inequality constraints with only lower bounds: 0
inequality constraints with lower and upper bounds: 99
inequality constraints with only upper bounds: 0
iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls
0 2.5000000e+01 0.00e+00 1.00e+00 1.10e+01 -1.0 - 0.00e+00 0 0
1 2.0845493e+00 2.22e-16 1.05e-02 1.29e+00 -1.0 - 1.00e+00 1 1f
2 7.6700630e-03 8.88e-16 2.22e-16 1.37e-01 -1.0 - 1.00e+00 1 1f
3 2.6844175e-06 6.66e-16 1.06e-16 3.12e-03 -2.5 - 1.00e+00 1 1h
4 1.5812159e-11 6.66e-16 1.10e-16 9.22e-06 -5.0 - 1.00e+00 1 1h
Number of Iterations....: 4
(scaled) (unscaled)
Objective...............: 1.5812159484987419e-11 1.5812159484987419e-11
Dual infeasibility......: 1.1007362556159084e-16 1.1007362556159084e-16
Constraint violation....: 6.6613381477509392e-16 6.6613381477509392e-16
Complementarity.........: 9.2190099817389014e-06 9.2190099817389014e-06
Overall NLP error.......: 9.2190099817389014e-06 9.2190099817389014e-06
Number of objective function evaluations = 5
Number of objective gradient evaluations = 5
Number of constraint evaluations = 5
Number of constraint Jacobian evaluations = 5
Number of Lagrangian Hessian evaluations = 4
Number of KKT factorizations = 4
Number of KKT backsolves = 4
Total wall secs in initialization = 51.716 s
Total wall secs in linear solver = unavailable
Total wall secs in NLP function evaluations = 1.521 s
Total wall secs in solver (w/o init./fun./lin. alg.) = unavailable
Total wall secs = 141.799 s
EXIT: Optimal Solution Found (tol = 1.0e-04).
(cpu = :first_order, gpu = MadNLP.SOLVE_SUCCEEDED)
Those six lines are the whole ExaModels API: variables, objective, constraints, model.
A first optimal control problem¶
The pendulum from tutorial 1, now with a motor at the pivot. Which torque
u(t) swings it from hanging down (θ = 0) to upright (θ = π) with least
effort? The torque is capped well below gravity, so the answer has to
pump the swing.
Direct transcription, the simultaneous approach of the lecture, turns
the optimal control problem into an NLP. Lay
a uniform grid of N intervals over the horizon, make the states θ, ω
and the control u at each grid point decision variables, and impose the
dynamics as constraints linking neighbours, here with the trapezoidal rule:
$$ \begin{aligned} \min_{\theta, \omega, u}\quad & \sum_{i=0}^{N} u_i^2\, \Delta t \\ \text{s.t.}\quad & \theta_i - \theta_{i-1} - \tfrac{\Delta t}{2}(\omega_i + \omega_{i-1}) = 0, && i = 1, \dots, N, \\ & \omega_i - \omega_{i-1} - \tfrac{\Delta t}{2}\big(f(\theta_i,\omega_i,u_i) + f(\theta_{i-1},\omega_{i-1},u_{i-1})\big) = 0, && i = 1, \dots, N, \\ & -u_{\max} \le u_i \le u_{\max}, && i = 0, \dots, N, \\ & \theta_0 = 0,\quad \omega_0 = 0,\quad \theta_N = \pi,\quad \omega_N = 0, \\[2pt] & \text{where } f(\theta, \omega, u) = u - 9.81 \sin\theta - 0.1\, \omega . \end{aligned} $$
Each line is one pattern repeated over its index, and each becomes a
single ExaModels generator: the two
dynamics families are N copies of one expression, the bound is N+1
copies of another, and only the four boundary conditions are one-offs.
What comes out is a structured NLP: large, sparse, block-banded, and
above all repetitive, which is exactly what ExaModels is built to exploit.
Every constraint is written as a generator over the time grid, the SIMD pattern from the lecture: one
expression, N instances. ExaModels compiles that into a single GPU kernel,
and differentiates it on the GPU as well. We wrap the model in a function
taking a backend argument. That one argument is the whole difference
between CPU and GPU below.
function swingup_model(N; T = 5.0, u_max = 4.0, backend = nothing)
dt = T / N
c = ExaCore(; backend = backend)
@add_var(c, θ, 0:N; start = [π * i / N for i = 0:N])
@add_var(c, ω, 0:N; start = fill(π / T, N + 1))
@add_var(c, u, 0:N; lvar = -u_max, uvar = u_max, start = zeros(N + 1))
@add_obj(c, u[i]^2 * dt for i = 0:N)
@add_con(c, θ[i] - θ[i-1] - dt / 2 * (ω[i] + ω[i-1]) for i = 1:N)
@add_con(c,
ω[i] - ω[i-1] - dt / 2 * ((u[i] - 9.81 * sin(θ[i]) - 0.1 * ω[i]) +
(u[i-1] - 9.81 * sin(θ[i-1]) - 0.1 * ω[i-1]))
for i = 1:N)
@add_con(c, θ[0])
@add_con(c, ω[0])
@add_con(c, θ[N] - π)
@add_con(c, ω[N])
return ExaModel(c)
end
swingup_model (generic function with 1 method)
Three formulation choices here are deliberate. The boundary conditions are
constraints like any other. ExaModels constraints are equalities unless you
pass lcon/ucon, and the torque limit went on the variable instead.
The initial guess is dynamically consistent: θ ramps to π and ω starts at
the matching average rate π/T. This is optional. An interior-point method
does not need a feasible starting point and copes perfectly well with an
inconsistent one, so start could be left at zero throughout. Supplying a
guess that roughly satisfies the dynamics simply gives the solver less
distance to travel, which is worth doing when it costs one line.
Build it and solve on the CPU:
N = 1000
model_cpu = swingup_model(N)
result_cpu = ipopt(model_cpu; max_iter = 600)
result_cpu.status
This is Ipopt version 3.14.19, running with linear solver MUMPS 5.9.0.
Number of nonzeros in equality constraint Jacobian...: 10004
Number of nonzeros in inequality constraint Jacobian.: 0
Number of nonzeros in Lagrangian Hessian.............: 3001
Total number of variables............................: 3003
variables with only lower bounds: 0
variables with lower and upper bounds: 1001
variables with only upper bounds: 0
Total number of equality constraints.................: 2004
Total number of inequality constraints...............: 0
inequality constraints with only lower bounds: 0
inequality constraints with lower and upper bounds: 0
inequality constraints with only upper bounds: 0
iter objective inf_pr inf_du lg(mu) ||d|| lg(rg) alpha_du alpha_pr ls
0 0.0000000e+00 6.28e-01 0.00e+00 -1.0 0.00e+00 - 0.00e+00 0.00e+00 0
1 1.9954762e+01 3.20e-01 8.93e+00 -1.0 8.08e+00 - 3.31e-01 4.90e-01h 1
2 2.4592804e+01 2.39e-01 7.09e+00 -1.0 3.27e+00 0.0 6.31e-01 2.52e-01h 1
3 2.7100993e+01 2.17e-01 8.10e+00 -1.0 5.13e+01 -0.5 4.29e-02 9.42e-02h 1
4 3.1567053e+01 1.10e-01 4.17e+00 -1.0 1.47e+00 -0.1 1.00e+00 4.94e-01h 1
5 3.6117812e+01 6.75e-03 3.21e+00 -1.0 1.51e+00 -0.5 8.26e-01 9.38e-01h 1
6 3.5675558e+01 2.41e-03 2.22e+00 -1.0 3.00e+00 -1.0 5.70e-01 8.19e-01f 1
7 3.5054851e+01 4.23e-03 1.80e+00 -1.0 7.65e+00 -1.5 3.01e-01 4.57e-01f 1
8 3.2335780e+01 1.55e-04 1.14e-01 -1.0 1.30e+00 -1.1 1.00e+00 1.00e+00f 1
9 3.1866412e+01 1.02e-03 3.02e-01 -1.7 1.99e+00 - 6.52e-01 1.00e+00f 1
iter objective inf_pr inf_du lg(mu) ||d|| lg(rg) alpha_du alpha_pr ls
10 3.1769836e+01 1.46e-03 3.42e-01 -1.7 3.25e+00 - 3.58e-01 8.60e-01F 1
11 2.9577556e+01 1.77e-04 3.17e-02 -1.7 1.58e+00 - 1.00e+00 1.00e+00f 1
12 2.9471746e+01 2.18e-04 2.75e-02 -2.5 2.15e+00 - 4.91e-01 2.50e-01f 3
13 2.9151897e+01 5.71e-06 6.19e-03 -3.8 2.19e-01 -1.5 8.86e-01 1.00e+00f 1
14 2.8992573e+01 6.48e-04 9.24e-03 -3.8 1.51e+00 - 4.96e-01 5.00e-01f 2
15 2.8940885e+01 8.53e-04 1.16e-02 -3.8 1.11e+00 - 5.80e-01 5.00e-01f 2
16 2.8870109e+01 7.85e-04 1.02e-02 -3.8 8.48e-01 - 6.63e-01 5.00e-01f 2
17 2.8785724e+01 6.29e-04 7.91e-03 -3.8 6.52e-01 - 8.40e-01 5.00e-01h 2
18 2.8783020e+01 6.14e-04 7.65e-03 -3.8 1.18e+00 - 5.55e-01 1.47e-01h 3
19 2.8453879e+01 1.49e-06 6.01e-04 -3.8 6.16e-02 -2.0 1.00e+00 1.00e+00f 1
iter objective inf_pr inf_du lg(mu) ||d|| lg(rg) alpha_du alpha_pr ls
20 2.8437061e+01 1.74e-06 8.83e-05 -5.7 7.38e-02 - 8.88e-01 1.00e+00f 1
21 2.8427390e+01 1.39e-05 6.68e-05 -5.7 8.04e-02 - 9.53e-01 1.00e+00H 1
22 2.8433855e+01 5.77e-07 1.99e-04 -5.7 2.07e-01 - 5.69e-01 2.85e-01h 2
23 2.8428933e+01 9.87e-06 5.45e-06 -5.7 5.90e-02 - 1.00e+00 1.00e+00H 1
24 2.8435287e+01 4.48e-06 8.67e-05 -8.6 7.20e-02 - 5.68e-01 6.93e-01h 1
25 2.8433506e+01 7.04e-07 1.32e-05 -8.6 9.55e-03 - 7.80e-01 8.90e-01h 1
26 2.8433183e+01 4.18e-09 5.84e-07 -8.6 2.91e-03 - 9.85e-01 1.00e+00f 1
27 2.8433180e+01 5.10e-11 5.46e-10 -8.6 1.54e-03 - 1.00e+00 1.00e+00h 1
28 2.8433180e+01 1.66e-11 1.82e-10 -9.0 7.80e-04 - 1.00e+00 1.00e+00h 1
Number of Iterations....: 28
(scaled) (unscaled)
Objective...............: 2.8433180175975917e+01 2.8433180175975917e+01
Dual infeasibility......: 1.8248602629000743e-10 1.8248602629000743e-10
Constraint violation....: 1.6608554809227627e-11 1.6608554809227627e-11
Variable bound violation: 0.0000000000000000e+00 0.0000000000000000e+00
Complementarity.........: 6.9434847817759740e-09 6.9434847817759740e-09
Overall NLP error.......: 6.9434847817759740e-09 6.9434847817759740e-09
Number of objective function evaluations = 52
Number of objective gradient evaluations = 29
Number of equality constraint evaluations = 52
Number of inequality constraint evaluations = 0
Number of equality constraint Jacobian evaluations = 29
Number of inequality constraint Jacobian evaluations = 0
Number of Lagrangian Hessian evaluations = 28
Total seconds in IPOPT = 2.416
EXIT: Optimal Solution Found.
:first_order
Each row of the log is one interior-point
iteration. inf_pr and inf_du are the primal and dual infeasibilities
driving to zero, and lg(mu) is the barrier parameter being driven down.
Convergence in a few dozen iterations on a 3000-variable nonconvex problem
is what interior-point methods are good at.
(objective = result_cpu.objective, iterations = result_cpu.iter)
(objective = 28.433180175975917, iterations = 28)
The solution comes back through solution, addressed by the variable
handles our model function returned:
t_grid = range(0, 5.0; length = N + 1)
θ_sol = Array(solution(result_cpu, model_cpu.refs.θ))
u_sol = Array(solution(result_cpu, model_cpu.refs.u))
1001-element Vector{Float64}:
0.09532692795142414
0.22191775043853534
0.2844301221533846
0.3469039862959328
0.4093239990009578
0.47167482193947524
0.5339411260840587
0.5961075954738733
0.6581589309785607
0.7200798540601324
⋮
-0.151481343795315
-0.15416120810705383
-0.1568802325780955
-0.15963910384311006
-0.1624385186556306
-0.1652791840640551
-0.16816181759026022
-0.171087147410873
-0.08628028904561526
Animating the trajectory shows the maneuver directly. The left panel is the pendulum, drawn from the pivot, with θ = 0 hanging down and θ = π upright. The right panel is the torque, with a cursor at the current time:
anim = @animate for k in 1:20:(N + 1)
pend = plot([0, sin(θ_sol[k])], [0, -cos(θ_sol[k])];
xlim = (-1.4, 1.4), ylim = (-1.45, 1.35), aspect_ratio = 1,
lw = 4, marker = :circle, ms = 8, legend = false,
titlefontsize = 10,
title = "pendulum at t = $(round(t_grid[k]; digits = 2)) s")
trace = plot(t_grid, u_sol;
legend = false, lw = 2, color = :gray, titlefontsize = 10,
xlabel = "t [s]", title = "torque u(t), |u| ≤ 4")
vline!(trace, [t_grid[k]]; color = :red, lw = 2)
plot(pend, trace; layout = (1, 2), size = (820, 380),
left_margin = 4Plots.mm, bottom_margin = 4Plots.mm, top_margin = 3Plots.mm)
end
gif(anim, fps = 25)
[ Info: Saved animation to /tmp/jl_vQ7WRobgXE.gif
The pumping is visible: the torque swings back and forth to build energy, then catches the pendulum at the top.
The device switch¶
Everything above ran on the CPU. Here is the entire porting effort:
model_gpu = swingup_model(N; backend = CUDABackend())
result_gpu = madnlp(model_gpu; max_iter = 600, tol = 1e-8)
(status = result_gpu.status, objective = result_gpu.objective,
cpu_objective = result_cpu.objective)
This is MadNLP version v0.10.1, running with cuDSS v0.8.0 Number of nonzeros in constraint Jacobian............: 10004 Number of nonzeros in Lagrangian Hessian.............: 3001 Total number of variables............................: 3003 variables with only lower bounds: 0 variables with lower and upper bounds: 1001 variables with only upper bounds: 0 Total number of equality constraints.................: 0 Total number of inequality constraints...............: 2004 inequality constraints with only lower bounds: 0 inequality constraints with lower and upper bounds: 2004 inequality constraints with only upper bounds: 0 iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 0 0.0000000e+00 6.28e-01 0.00e+00 4.00e+00 -1.0 - 0.00e+00 0 0 1 3.2615633e-11 6.28e-01 4.97e-01 2.36e-03 -1.0 - 2.00e-07 1 2h 2 1.9954762e+01 3.20e-01 2.38e-01 1.14e-03 -1.7 - 4.90e-01 2 1h 3 2.5103522e+01 2.29e-01 2.34e-01 3.36e-04 -1.7 0.0 2.85e-01 1 1h 4 2.9368662e+01 1.66e-01 2.98e-01 2.74e-04 -1.7 -0.5 2.76e-01 1 1h 5 3.2207321e+01 1.03e-01 8.69e-02 1.31e-04 -1.7 -0.1 3.78e-01 1 1h 6 3.5719699e+01 2.87e-02 7.12e-02 1.30e-04 -2.5 -0.5 7.21e-01 1 1h 7 3.5791119e+01 1.09e-02 1.72e-01 1.22e-04 -2.5 -1.0 6.21e-01 1 1h 8 3.5162404e+01 2.57e-03 2.69e-01 1.20e-04 -2.5 -1.5 1.00e+00 1 1h 9 3.3546657e+01 1.77e-04 1.20e-01 6.98e-05 -2.5 - 1.00e+00 3 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 10 3.2923687e+01 7.20e-04 9.28e-02 5.24e-05 -2.5 - 1.00e+00 3 1h 11 3.2099299e+01 1.29e-04 6.30e-02 3.67e-05 -2.5 - 1.00e+00 3 1h 12 3.1543161e+01 2.71e-04 4.58e-02 2.74e-05 -2.5 - 1.00e+00 3 1h 13 3.0998028e+01 3.66e-04 3.27e-02 2.15e-05 -2.5 - 1.00e+00 3 1h 14 3.0495952e+01 5.15e-04 2.27e-02 1.61e-05 -2.5 - 1.00e+00 3 1h 15 2.9898428e+01 3.78e-04 2.75e-02 1.75e-05 -3.8 - 1.00e+00 3 1h 16 2.9511210e+01 4.06e-04 2.91e-02 1.80e-05 -3.8 - 1.00e+00 3 1h 17 2.9250518e+01 5.43e-04 3.29e-02 1.80e-05 -3.8 - 1.00e+00 3 1h 18 2.9047187e+01 6.03e-04 4.24e-02 1.73e-05 -3.8 - 1.00e+00 3 1h 19 2.8896489e+01 6.08e-04 5.34e-02 1.57e-05 -3.8 - 1.00e+00 3 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 20 2.8780014e+01 5.51e-04 6.35e-02 1.34e-05 -3.8 - 1.00e+00 2 1h 21 2.8680650e+01 4.35e-04 6.75e-02 1.04e-05 -3.8 - 1.00e+00 3 1h 22 2.8598388e+01 3.02e-04 5.71e-02 7.09e-06 -3.8 - 1.00e+00 3 1h 23 2.8539046e+01 1.91e-04 2.08e-02 3.98e-06 -3.8 - 1.00e+00 3 1h 24 2.8496636e+01 1.03e-04 3.14e-05 5.00e-06 -3.8 - 1.00e+00 4 1h 25 2.8443172e+01 1.27e-05 1.37e-02 6.54e-06 -5.7 - 1.00e+00 3 1h 26 2.8434319e+01 6.36e-07 1.81e-03 1.19e-05 -5.7 - 1.00e+00 2 1h 27 2.8433449e+01 3.62e-09 1.32e-07 3.91e-06 -5.7 - 1.00e+00 2 1h 28 2.8433159e+01 1.25e-10 1.49e-03 1.07e-06 -8.6 - 1.00e+00 2 1h 29 2.8432993e+01 1.23e-10 1.34e-03 3.87e-07 -8.6 - 1.00e+00 1 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 30 2.8432925e+01 1.95e-11 4.57e-04 1.22e-07 -8.6 - 1.00e+00 2 1h 31 2.8432908e+01 9.41e-13 7.52e-05 3.45e-08 -8.6 - 1.00e+00 2 1h 32 2.8432906e+01 3.76e-12 3.95e-11 9.88e-09 -8.6 - 1.00e+00 2 1h Number of Iterations....: 32 (scaled) (unscaled) Objective...............: 2.8432905769172695e+01 2.8432905769172695e+01 Dual infeasibility......: 3.9452524580596560e-11 3.9452524580596560e-11 Constraint violation....: 3.7582667023195239e-12 3.7582667023195239e-12 Complementarity.........: 9.8793443721915602e-09 9.8793443721915602e-09 Overall NLP error.......: 9.8793443721915602e-09 9.8793443721915602e-09 Number of objective function evaluations = 34 Number of objective gradient evaluations = 33 Number of constraint evaluations = 34 Number of constraint Jacobian evaluations = 33 Number of Lagrangian Hessian evaluations = 32 Number of KKT factorizations = 41 Number of KKT backsolves = 73 Total wall secs in initialization = 6.624 s Total wall secs in linear solver = unavailable Total wall secs in NLP function evaluations = 2.745 s Total wall secs in solver (w/o init./fun./lin. alg.) = unavailable Total wall secs = 26.960 s EXIT: Optimal Solution Found (tol = 1.0e-08).
(status = MadNLP.SOLVE_SUCCEEDED, objective = 28.432905769172695, cpu_objective = 28.433180175975917)
One argument moved evaluation, derivatives and the KKT factorizations to the GPU, and the two devices agree.
Note tol = 1e-8. MadNLP's GPU default is a loose 1e-4, which suits the
condensed-space method but not these problems. Pass tol = 1e-8 to every
GPU solve.
(*) Scaling with problem size¶
Optional
A (*) marks material we may skip in the live session, depending on time. Nothing later depends on it, and the cells run on their own.
Grow the horizon and time both devices. Both columns are reported per iteration as well as in total, and the two tell different stories.
An interior-point method's cost is per iteration, and the two solvers do
not take the same number of them. Ipopt's count climbs steeply with the grid
while MadNLP's stays near thirty, so the raw ratio mixes that difference into
the device difference and flatters the GPU. The per_iter column divides it
out and is the one to read as a statement about hardware.
Both devices are warmed up before the loop, so compilation is not being timed. Everything that depends on the problem size, including model construction, is still inside the measurement.
# warm up both devices so the first row does not pay for compilation
ipopt(swingup_model(625); max_iter = 600, print_level = 0)
madnlp(swingup_model(625; backend = CUDABackend()); max_iter = 600, tol = 1e-8,
print_level = MadNLP.ERROR)
sizes = [1_250, 2_500, 5_000, 10_000, 20_000]
times = map(sizes) do n
m_c = swingup_model(n)
t_c = @elapsed r_c = ipopt(m_c; max_iter = 600, print_level = 0)
m_g = swingup_model(n; backend = CUDABackend())
t_g = @elapsed r_g = madnlp(m_g; max_iter = 600, tol = 1e-8,
print_level = MadNLP.ERROR)
(N = n, cpu_s = round(t_c; digits = 2), cpu_iters = r_c.iter,
gpu_s = round(t_g; digits = 2), gpu_iters = r_g.iter,
speedup = round(t_c / t_g; digits = 1),
per_iter = round((t_c / r_c.iter) / (t_g / r_g.iter); digits = 1))
end
times
5-element Vector{@NamedTuple{N::Int64, cpu_s::Float64, cpu_iters::Int64, gpu_s::Float64, gpu_iters::Int64, speedup::Float64, per_iter::Float64}}:
(N = 1250, cpu_s = 0.29, cpu_iters = 26, gpu_s = 0.28, gpu_iters = 35, speedup = 1.0, per_iter = 1.4)
(N = 2500, cpu_s = 0.67, cpu_iters = 25, gpu_s = 0.23, gpu_iters = 27, speedup = 2.9, per_iter = 3.2)
(N = 5000, cpu_s = 1.97, cpu_iters = 30, gpu_s = 0.28, gpu_iters = 30, speedup = 7.1, per_iter = 7.1)
(N = 10000, cpu_s = 3.34, cpu_iters = 28, gpu_s = 0.48, gpu_iters = 33, speedup = 6.9, per_iter = 8.1)
(N = 20000, cpu_s = 43.9, cpu_iters = 110, gpu_s = 0.4, gpu_iters = 31, speedup = 110.0, per_iter = 31.0)
The same numbers as a curve, with the break-even line marked:
plot([t.N for t in times], [t.speedup for t in times];
xscale = :log10, marker = :circle, lw = 2, label = "total time")
plot!([t.N for t in times], [t.per_iter for t in times];
marker = :square, lw = 2, ls = :dash, label = "per iteration",
xlabel = "N (horizon length)", ylabel = "CPU / GPU",
title = "GPU speedup against problem size")
hline!([1.0]; ls = :dot, color = :black, label = "break even")
Slower or level at small sizes, pulling away as the problem grows: the same crossover as tutorial 1's batch experiment, now with sparse factorizations in the mix. The gap between the two curves is the iteration-count difference, and it widens with the grid.
The CPU side is what stops the table here. Past this point Ipopt takes minutes while the GPU is still under a second, so we grow the problem on the GPU alone:
for n in (40_000, 80_000, 160_000)
m = swingup_model(n; backend = CUDABackend())
t = @elapsed r = madnlp(m; max_iter = 600, tol = 1e-8,
print_level = MadNLP.ERROR)
println(rpad("N = $n", 14), rpad("$(round(t; digits = 2)) s", 10),
rpad("$(r.iter) iters", 12), r.status)
end
N = 40000 0.9 s 37 iters SOLVE_SUCCEEDED N = 80000 0.82 s 36 iters SOLVE_SUCCEEDED N = 160000 5.69 s 134 iters SOLVE_SUCCEEDED
The iteration count holds near thirty-something out to 80,000 and then rises at 160,000, where the solve still finishes in a few seconds. That is the property worth noticing: the work inside an iteration parallelizes, so the time per iteration grows far more slowly than the problem does. How many iterations are needed is a question about the problem and the algorithm, not about the hardware, and it is the part a faster device cannot help with.
One caveat to carry away, and it is the same lesson as Goddard's tolerance.
At equal tol the two devices do not deliver equal accuracy. The GPU path
condenses the KKT system, and on this problem it lands at a constraint
violation near 2e-8 where Ipopt reaches far tighter. Because tol bounds
each constraint separately, a finer grid means more constraints and a larger
accumulated difference, so the GPU objective drifts slightly below the CPU
one as N grows. Tightening to tol = 1e-10 closes it to five digits, at
the cost of several times more iterations. Read a speedup as a statement
about a given accuracy, never on its own.
✏️ Exercise 1: minimum time¶
Your turn. The horizon above was fixed at 5 seconds and we minimized control effort. Ask the opposite question: with a stronger motor,
u_max = 10, what is the fastest swing-up?The final time has to become part of the problem. Goddard's rocket below does exactly this, so look at how it handles a free final time and apply the same idea here.
Solve on the GPU, plot θ and u, and compare the torque with the minimum-effort solution.
Solution: minimum-time-solution.
Goddard's rocket¶
Maximize the final altitude of a vertically launched rocket, with bounded thrust and drag that grows with speed. With altitude $h$, velocity $v$, mass $m$ and thrust $T$:
$$ \begin{aligned} \max_{T,\, t_f}\quad & h(t_f) \\ \text{s.t.}\quad & \dot h = v, \\ & \dot v = \frac{T - D(h, v) - m\, g(h)}{m}, \qquad \dot m = -\frac{T}{c}, \\ & D(h,v) = D_c\, v^2 \exp\!\big(-h_c (h - h_0)/h_0\big), \qquad g(h) = g_0 \left(\frac{h_0}{h}\right)^{2}, \\ & h(0) = h_0,\quad v(0) = 0,\quad m(0) = m_0,\quad m(t_f) = m_f, \\ & 0 \le T \le T_{\max}, \qquad h \ge h_0 . \end{aligned} $$
Two formulation ideas beyond the swing-up. The final time $t_f$ is free,
so the step length becomes a decision variable, and we maximize by
passing minimize = false.
The third idea is the transcription itself. Everything so far used the trapezoidal rule, which is second order: accuracy comes only from taking more steps. Here we use orthogonal collocation on finite elements, the method from Part II of the lecture, which transcribes at high order, stably for stiff systems, and purely algebraically.
Split the horizon into $N$ elements. On element $i$ of length $h_i$, with a local coordinate $\tau \in [0,1]$, represent each state as a degree-$K$ Lagrange polynomial through its values at $\tau_0 = 0 < \tau_1 < \dots < \tau_K$:
$$ z^K(t) = \sum_{j=0}^{K} \ell_j(\tau)\, z_{ij}, \qquad \ell_j(\tau) = \prod_{k \ne j} \frac{\tau - \tau_k}{\tau_j - \tau_k} . $$
Since $\ell_j(\tau_k) = \delta_{jk}$, the $z_{ij}$ are the state values at the nodes, and they are the decision variables. We cannot ask $\dot z^K = f(z^K, u)$ to hold at every $t$ with finitely many unknowns, so we impose it at the $K$ collocation points, which gives the collocation equations:
$$ \sum_{j=0}^{K} D_{kj}\, z_{ij} \;=\; h_i\, f\big(z_{ik},\, u_{ik}\big), \qquad i = 1, \dots, N, \quad k = 1, \dots, K, \qquad D_{kj} = \frac{d\ell_j}{d\tau}(\tau_k) . $$
Two indices this time, so the pattern repeats $N \times K$ times, and it is still one generator.
$D$ is a table of precomputed constants; the $z_{ij}$ and $h_i$ are the unknowns. Each equation is therefore a linear combination of unknowns set equal to $h_i f(\cdot)$, which is exactly the shape an NLP solver wants.
We use Radau points, whose last node sits at $\tau_K = 1$. That choice gives stiff decay and makes the continuity condition $z_{i+1,0} = \sum_j \ell_j(1) z_{ij}$ collapse to $z_{i+1,0} = z_{iK}$, since $\ell_j(1) = \delta_{jK}$: consecutive elements simply share the node.
$D$ follows from the nodes alone. Differentiating the Lagrange basis is one small linear solve:
# Radau IIA nodes for K = 3, giving order 2K-1 = 5 accuracy per element.
const RADAU_C = [0.15505102572168220, 0.64494897427831780, 1.0]
# D[k,j] = dl_j/dtau at tau_k, over the K+1 nodes [0, c...]. Obtained by
# writing the basis in monomials: values at the nodes are V, derivatives at
# the collocation points are Vp, so D = Vp / V.
function differentiation_matrix(c)
K = length(c)
τ = vcat(0.0, c)
V = [τ[i]^(j - 1) for i = 1:K+1, j = 1:K+1]
Vp = [(j - 1) * c[k]^(max(j - 2, 0)) * (j >= 2) for k = 1:K, j = 1:K+1]
return Vp / V
end
D_RADAU = differentiation_matrix(RADAU_C)
3×4 Matrix{Float64}:
-4.13939 3.22474 1.16784 -0.253197
1.73939 -3.56784 0.775255 1.0532
-3.0 5.53197 -7.53197 5.0
Worth a sanity check before trusting it: a constant state has zero derivative at every collocation point, and the state $z = \tau$ has derivative one.
(constant = D_RADAU * ones(4), linear = D_RADAU * vcat(0.0, RADAU_C))
(constant = [2.220446049250313e-16, 0.0, 0.0], linear = [1.0, 0.9999999999999999, 0.9999999999999999])
Now the model. Note what the decision variables are: the state at each element start, the state at each collocation point, one thrust per element, and the element length itself, since the final time is free.
function goddard_model(nh; K = 3, backend = nothing)
c = RADAU_C
D = differentiation_matrix(c)
h_0, v_0, m_0, g_0 = 1.0, 0.0, 1.0, 1.0
T_c, h_c, v_c, m_c = 3.5, 500.0, 620.0, 0.6
cex = 0.5 * sqrt(g_0 * h_0)
m_f = m_c * m_0
D_c = 0.5 * v_c * (m_0 / g_0)
T_max = T_c * m_0 * g_0
core = ExaCore(; backend = backend, minimize = false)
@add_var(core, step, 1; start = 1 / nh, lvar = 0.0) # element length h_i
# z_ij: state at node j of element i, with j = 0 the element start
@add_var(core, h, 1:nh, 0:K; start = h_0, lvar = h_0)
@add_var(core, v, 1:nh, 0:K; start = 0.1, lvar = v_0)
@add_var(core, m, 1:nh, 0:K; start = m_0, lvar = m_f, uvar = m_0)
@add_var(core, Th, 1:nh; start = T_max / 2, lvar = 0.0, uvar = T_max)
@add_obj(core, h[nh, K])
# One collocation row per (element, collocation point): the right-hand
# side h_i f(z_ik), then the sum over j of D[k,j] z_ij added to it. Both
# index rows through the same formula.
pts = [(n = (i - 1) * K + k, i = i, k = k) for i = 1:nh for k = 1:K]
trm = [(n = (i - 1) * K + k, i = i, j = j, d = D[k, j+1])
for i = 1:nh for k = 1:K for j = 0:K]
@add_con(core, ch, -step[1] * v[e.i, e.k] for e in pts)
@add_con!(core, ch[e.n] += e.d * h[e.i, e.j] for e in trm)
@add_con(core, cv, -step[1] *
((Th[e.i] - D_c * v[e.i, e.k]^2 * exp(-h_c * (h[e.i, e.k] - h_0) / h_0) -
m[e.i, e.k] * g_0 * (h_0 / h[e.i, e.k])^2) / m[e.i, e.k]) for e in pts)
@add_con!(core, cv[e.n] += e.d * v[e.i, e.j] for e in trm)
@add_con(core, cm, step[1] * Th[e.i] / cex for e in pts)
@add_con!(core, cm[e.n] += e.d * m[e.i, e.j] for e in trm)
# continuity: the last node of an element is the first node of the next
@add_con(core, h[i, K] - h[i+1, 0] for i = 1:nh-1)
@add_con(core, v[i, K] - v[i+1, 0] for i = 1:nh-1)
@add_con(core, m[i, K] - m[i+1, 0] for i = 1:nh-1)
@add_con(core, h[1, 0] - h_0)
@add_con(core, v[1, 0] - v_0)
@add_con(core, m[1, 0] - m_0)
@add_con(core, m[nh, K] - m_f)
return ExaModel(core)
end
goddard_model (generic function with 1 method)
The published optimal final altitude for this problem is 1.01283 in normalized units (COPS 3.0), which gives us something independent to check the answer against:
# 200 intervals here, against the 1000 the trapezoidal swing-up needed:
# higher order buys accuracy that a finer grid would otherwise have to.
nh = 200
rocket = goddard_model(nh; backend = CUDABackend())
rocket_result = madnlp(rocket; max_iter = 600, tol = 1e-8)
(status = rocket_result.status, final_altitude = rocket_result.objective,
published_optimum = 1.01283)
This is MadNLP version v0.10.1, running with cuDSS v0.8.0 Number of nonzeros in constraint Jacobian............: 13798 Number of nonzeros in Lagrangian Hessian.............: 13200 Total number of variables............................: 2601 variables with only lower bounds: 1601 variables with lower and upper bounds: 1000 variables with only upper bounds: 0 Total number of equality constraints.................: 0 Total number of inequality constraints...............: 2401 inequality constraints with only lower bounds: 0 inequality constraints with lower and upper bounds: 2401 inequality constraints with only upper bounds: 0 iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 0 -1.0100000e+00 3.96e-01 2.00e+00 1.75e+00 -1.0 - 0.00e+00 0 0 1 -1.0100019e+00 3.96e-01 4.53e+01 3.76e-02 -1.0 - 3.98e-07 2 1h 2 -1.0105262e+00 3.94e-01 5.15e+00 4.65e-03 -1.0 - 4.25e-03 2 1h 3 -1.0144772e+00 3.83e-01 7.47e-01 9.26e-04 -1.0 - 2.95e-02 3 1h 4 -1.0111270e+00 2.55e-01 8.04e+00 1.34e-03 -1.7 - 3.35e-01 3 1h 5 -1.0161009e+00 2.19e-01 5.57e+00 1.00e-03 -1.7 - 1.40e-01 4 1h 6 -1.0693253e+00 1.36e-01 7.48e+00 7.85e-04 -1.7 - 3.79e-01 6 1h 7 -1.0538364e+00 8.77e-02 9.13e+00 5.84e-04 -1.7 2.0 3.55e-01 1 1h 8 -1.0298894e+00 2.36e-02 4.69e+00 2.75e-04 -1.7 1.5 7.31e-01 1 1h 9 -1.0217542e+00 1.81e-02 2.57e+00 1.30e-04 -1.7 1.0 2.30e-01 1 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 10 -1.0054530e+00 7.14e-03 2.35e+00 1.09e-04 -1.7 0.6 6.06e-01 1 1h 11 -1.0066419e+00 7.93e-04 2.77e+00 6.20e-05 -1.7 - 9.89e-01 3 1h 12 -1.0071486e+00 1.18e-04 7.39e+00 1.57e-05 -1.7 - 9.90e-01 3 1h 13 -1.0087303e+00 5.87e-04 1.62e+01 6.03e-06 -1.7 - 1.00e+00 3 1h 14 -1.0086710e+00 5.58e-05 1.16e-01 9.66e-06 -1.7 - 1.00e+00 2 1h 15 -1.0086565e+00 8.46e-08 2.82e-02 6.27e-06 -2.5 - 1.00e+00 2 1h 16 -1.0086572e+00 1.64e-09 9.09e-02 5.38e-06 -3.8 - 1.00e+00 2 1h 17 -1.0086706e+00 5.30e-07 1.70e-04 3.56e-06 -3.8 - 1.00e+00 3 1h 18 -1.0087672e+00 8.05e-07 3.02e+02 6.55e-06 -5.7 - 1.00e+00 2 1h 19 -1.0098460e+00 1.42e-04 2.30e+02 5.00e-06 -5.7 - 1.00e+00 3 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 20 -1.0108674e+00 1.54e-04 5.08e+01 3.19e-06 -5.7 - 7.68e-01 3 1h 21 -1.0117426e+00 1.51e-04 1.63e-01 4.71e-06 -5.7 - 1.00e+00 3 1h 22 -1.0117466e+00 1.58e-06 6.95e-05 2.12e-06 -5.7 - 1.00e+00 3 1h 23 -1.0117471e+00 2.96e-09 3.28e-06 1.85e-06 -5.7 - 1.00e+00 3 1h 24 -1.0126150e+00 9.93e-05 4.25e+01 6.99e-07 -8.6 - 9.63e-01 3 1h 25 -1.0127833e+00 3.88e-05 1.19e+01 2.73e-07 -8.6 - 7.72e-01 2 1h 26 -1.0128274e+00 3.01e-05 4.02e+00 1.13e-07 -8.6 - 7.78e-01 2 1h 27 -1.0128394e+00 1.59e-05 1.20e+00 4.46e-08 -8.6 - 8.40e-01 2 1h 28 -1.0128457e+00 5.84e-06 1.67e-01 1.43e-08 -8.6 - 1.00e+00 2 1h 29 -1.0128486e+00 2.34e-06 1.66e-05 4.64e-09 -8.6 - 1.00e+00 2 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 30 -1.0128489e+00 1.38e-07 7.22e-07 2.75e-09 -8.6 - 1.00e+00 2 1h 31 -1.0128489e+00 1.04e-09 3.35e-09 2.51e-09 -8.6 - 1.00e+00 2 1h Number of Iterations....: 31 (scaled) (unscaled) Objective...............: -1.0128488516553145e+00 -1.0128488516553145e+00 Dual infeasibility......: 3.3489321135037688e-09 3.3489321135037688e-09 Constraint violation....: 1.0404955959342908e-09 1.0404955959342908e-09 Complementarity.........: 2.5108658752423813e-09 2.5108658752423813e-09 Overall NLP error.......: 3.3489321135037688e-09 3.3489321135037688e-09 Number of objective function evaluations = 32 Number of objective gradient evaluations = 32 Number of constraint evaluations = 32 Number of constraint Jacobian evaluations = 32 Number of Lagrangian Hessian evaluations = 31 Number of KKT factorizations = 38 Number of KKT backsolves = 76 Total wall secs in initialization = 7.303 s Total wall secs in linear solver = unavailable Total wall secs in NLP function evaluations = 3.542 s Total wall secs in solver (w/o init./fun./lin. alg.) = unavailable Total wall secs = 24.818 s EXIT: Optimal Solution Found (tol = 1.0e-08).
(status = MadNLP.SOLVE_SUCCEEDED, final_altitude = 1.0128488516553145, published_optimum = 1.01283)
The altitude matches the literature. Checking optimizer output against
something independent matters: SOLVE_SUCCEEDED means the algorithm met
its own criterion, not that you modelled the right problem.
And the trajectories:
h_traj = Array(solution(rocket_result, rocket.refs.h))[:, 1] # node j=0 of each element
Th_traj = Array(solution(rocket_result, rocket.refs.Th)) # nh interval controls
p1 = plot(range(0, 1; length = nh), h_traj; lw = 2, legend = false,
title = "Goddard: altitude")
p2 = plot(range(0, 1; length = nh), Th_traj; lw = 2, legend = false,
title = "thrust (note the singular arc)", xlabel = "normalized time")
plot(p1, p2; layout = (2, 1), size = (760, 460),
left_margin = 4Plots.mm, bottom_margin = 4Plots.mm)
The thrust profile has three phases: full burn, then a singular arc, a sustained intermediate thrust balancing drag against gravity, then cutoff.
Particle steering¶
The particle steering problem (COPS calls it that; you will also see it as "rocket steering") is a minimum-time problem: steer a constant-magnitude thrust vector so a particle reaches a given height with a given terminal velocity as fast as possible. With position $(x_1, x_2)$, velocity $(x_3, x_4)$, steering angle $u$ and acceleration magnitude $a = 100$:
$$ \begin{aligned} \min_{u,\, t_f}\quad & t_f \\ \text{s.t.}\quad & \dot x_1 = x_3, \qquad \dot x_2 = x_4, \\ & \dot x_3 = a \cos u, \qquad \dot x_4 = a \sin u, \\ & x(0) = 0, \qquad x_2(t_f) = 5,\ x_3(t_f) = 45,\ x_4(t_f) = 0, \\ & -\tfrac{\pi}{2} \le u(t) \le \tfrac{\pi}{2} . \end{aligned} $$
Free final time again, here as an explicit tf variable that is the
objective, with states in a 2-D matrix variable, indexed like an array.
A particle starts at rest at the origin. Its engine pushes with a fixed magnitude, and the only thing you choose is the direction of that push at each instant. The target is a state, not a place: be at height 5, travelling horizontally at 45, with no vertical motion left. Reach it as quickly as possible.
The arrows are the decision. Point them too steeply and the particle gains height quickly but arrives with vertical speed it must then cancel; point them too flat and it never reaches the height. The optimizer finds the schedule that trades these against each other in the least time.
Sources
Classical problem: A. E. Bryson and Y.-C. Ho, Applied Optimal Control, Wiley, 1975, pp. 59–62. The thrust magnitude and terminal conditions used here follow J. Betts, S. Eldersveld and W. Huffman, Sparse nonlinear programming test problems (Release 1.0), Boeing Computer Services technical report BCSTECH-93-047, 1993, as collected in COPS 3.0 (problem 9): Dolan, Moré and Munson, ANL/MCS-TM-273, 2004.
✏️ Exercise 2: transcribe it yourself¶
Your turn. You have the continuous-time statement above and nothing else. Transcribe it into an ExaModels model and solve it on the GPU with
tol = 1e-8.The discretization is your choice: implicit Euler, the trapezoidal rule, or collocation. They differ in accuracy per grid point and in how many variables they cost, so the method you pick changes both the answer and the size of the problem.
Report the status, the minimal flight time, and the iteration count. The published optimum is 0.55457; COPS reports 0.554577 at
nh= 200 and 0.554571 atnh= 800, so the discretized optimum shifts slightly as the mesh is refined. Then plot the steering angleuover time and interpret what the rocket is doing.
# your code here
Solution: particle-steering-solution.
Parameter estimation for dynamic models¶
Optional material
This section is extra material. We may not work through it live, depending on time. It is written to be read and run on your own afterwards, and nothing earlier in the workshop depends on it.
The problems above chose a control to reach a target. Estimation inverts the question: a system of ODEs carries unknown parameters, measurements of it are noisy, and the task is to find the parameters that best explain the data.
The machinery is the one you have already used. We discretize the dynamics into constraints as before, the states on the grid become decision variables as before, and the only changes are that the unknown parameters join the variable list and the objective measures misfit against data instead of control cost. This is the simultaneous (or "all-at-once") approach to dynamic optimization, and because it produces one large structured NLP, everything the GPU did for us above applies unchanged.
The system¶
The JAK2/STAT5 signaling model of Boehm et al. (2014), a standard benchmark for parameter estimation.
STAT5 comes in two isoforms. A stimulus that decays over time phosphorylates them, they pair into three dimers, and those are imported into the nucleus and later exported back as monomers. Mass spectrometry measures the phosphorylated forms over four hours. The rate constants $k_{\text{phos}}$, $k_{\text{imp}}$, $k_{\text{exp}}$ and $k_{\text{deg}}$ are what we estimate.
For our purposes it is an 8-state nonlinear ODE with 9 unknown parameters and 48 noisy measurements: the shape of estimation problems in process control and systems biology alike.
Sources
Model and data: M. E. Boehm, L. Adlung, M. Schilling, S. Roth, U. Klingmüller and W. D. Lehmann, "Identification of Isoform-Specific Dynamics in Phosphorylation-Dependent STAT5 Dimerization by Quantitative Mass Spectrometry and Mathematical Modeling", Journal of Proteome Research 13(12):5685–5694, 2014 (doi:10.1021/pr5006923). Distributed as problem
Boehm_JProteomeRes2014of the PEtab benchmark collection (H. Hass et al., "Benchmark problems for dynamic modeling of intracellular processes", Bioinformatics 35:3073–3082, 2019 doi:10.1093/bioinformatics/btz020), in the PEtab format of L. Schmiester et al., PLOS Computational Biology 17(1):e1008646, 2021 (doi:10.1371/journal.pcbi.1008646). We take the problem from that collection and write it out directly, as earlier in this tutorial.
The dynamics¶
Eight species: cytosolic STAT5A, STAT5B, the three phosphorylated
dimers pApA, pApB, pBpB, and their nuclear counterparts nucpApA,
nucpApB, nucpBpB. Nine mass-action reactions, namely phosphorylation-
driven dimerization (3), nuclear import (3) and nuclear export (3), with
cytosolic
and nuclear compartment volumes $V_c = 1.4$, $V_n = 0.45$:
$$ \begin{aligned} v_1 &= V_c\, E(t)\, k_{\text{phos}}\, \mathrm{STAT5A}^2, & v_4 &= V_c\, k_{\text{imp,homo}}\, \mathrm{pApA}, & v_7 &= V_n\, k_{\text{exp,homo}}\, \mathrm{nucpApA}, \\ v_2 &= V_c\, E(t)\, k_{\text{phos}}\, \mathrm{STAT5A}\,\mathrm{STAT5B}, & v_5 &= V_c\, k_{\text{imp,hetero}}\, \mathrm{pApB}, & v_8 &= V_n\, k_{\text{exp,hetero}}\, \mathrm{nucpApB}, \\ v_3 &= V_c\, E(t)\, k_{\text{phos}}\, \mathrm{STAT5B}^2, & v_6 &= V_c\, k_{\text{imp,homo}}\, \mathrm{pBpB}, & v_9 &= V_n\, k_{\text{exp,homo}}\, \mathrm{nucpBpB}, \end{aligned} $$
where the stimulus decays exponentially, $E(t) = 1.25\times10^{-7} \exp(-k_{\text{deg}}\, t)$. The balances (note the stoichiometric 2s: two monomers make one homodimer) are
$$ \begin{aligned} V_c\, \tfrac{d}{dt}\mathrm{STAT5A} &= -2v_1 - v_2 + 2v_7 + v_8, & V_c\, \tfrac{d}{dt}\mathrm{pApA} &= v_1 - v_4, & V_n\, \tfrac{d}{dt}\mathrm{nucpApA} &= v_4 - v_7, \\ V_c\, \tfrac{d}{dt}\mathrm{STAT5B} &= -2v_3 - v_2 + 2v_9 + v_8, & V_c\, \tfrac{d}{dt}\mathrm{pApB} &= v_2 - v_5, & V_n\, \tfrac{d}{dt}\mathrm{nucpApB} &= v_5 - v_8, \\ & & V_c\, \tfrac{d}{dt}\mathrm{pBpB} &= v_3 - v_6, & V_n\, \tfrac{d}{dt}\mathrm{nucpBpB} &= v_6 - v_9 . \end{aligned} $$
What is measured¶
The instrument does not see the states. It sees three relative quantities the fraction of STAT5A that is phosphorylated, the same for STAT5B, and the relative amount of STAT5A overall. Each is a nonlinear function of the states involving a known isoform-specificity constant $c_{17} = 0.107$:
$$ \begin{aligned} y_1 &= \frac{100\,\mathrm{pApB} + 200\,\mathrm{pApA}\, c_{17}} {\mathrm{pApB} + \mathrm{STAT5A}\, c_{17} + 2\,\mathrm{pApA}\, c_{17}}, \\[2pt] y_2 &= \frac{-100\,\mathrm{pApB} + 200\,\mathrm{pBpB}\,(c_{17}-1)} {\mathrm{STAT5B}\,(c_{17}-1) - \mathrm{pApB} + 2\,\mathrm{pBpB}\,(c_{17}-1)}, \\[2pt] y_3 &= \frac{100\,\mathrm{pApB} + 100\,\mathrm{STAT5A}\, c_{17} + 200\,\mathrm{pApA}\, c_{17}} {2\,\mathrm{pApB} + \mathrm{STAT5A}\, c_{17} + 2\,\mathrm{pApA}\, c_{17} - \mathrm{STAT5B}\,(c_{17}-1) - 2\,\mathrm{pBpB}\,(c_{17}-1)} . \end{aligned} $$
Observation functions like these are the norm rather than the exception, because you almost never measure the state directly. Note they are rational, so the objective is not a least-squares problem in the standard sense but a nonconvex one.
The data, 16 time points from 0 to 240 minutes and three observables, ships next to this notebook as literal arrays:
include("boehm_data.jl")
(n_timepoints = length(T_DATA), t_final = T_DATA[end], n_measurements = 3 * length(T_DATA))
(n_timepoints = 16, t_final = 240.0, n_measurements = 48)
Estimation as a nonlinear program¶
Three choices define it.
The dynamics become constraints, by the same orthogonal collocation used for Goddard. Nuclear import here is fast enough to make the system stiff, which rules out the trapezoidal rule: it is A-stable but not L-stable, so it rings on fast modes rather than damping them. Radau collocation damps them and is high order, so it costs nothing to prefer it. The elements are 2.5 minutes long, which puts every one of the 16 measurement times on an element boundary.
The parameters become variables. The rate constants span nine orders of magnitude, so we estimate $\theta = \log_{10} k$ and write $k = 10^{\theta}$ in the dynamics. Positivity is then automatic and the search space is well scaled.
The objective is a likelihood. With Gaussian noise of unknown standard deviation $\sigma_j$ per observable, dropping constants:
$$ \min_{\theta,\, \sigma,\, x(\cdot)} \quad \sum_{j=1}^{3} \sum_{i=1}^{16} \left[ \frac{\big(y_j(t_i) - \bar y_{j,i}\big)^2}{2\sigma_j^2} + \log \sigma_j \right] . $$
The $\log \sigma_j$ term stops the optimizer explaining the data by declaring the noise enormous. We estimate $\log_{10}\sigma$ for the same reason as the rate constants.
In ExaModels:
function boehm_model(N; K=3, backend=nothing)
D = differentiation_matrix(RADAU_C)
Vc, Vn = 1.4, 0.45; c17, ratio = 0.107, 0.693
A0, B0 = 207.6*ratio, 207.6*(1-ratio)
tf = T_DATA[end]; h = tf/N
elem = [round(Int, t/h) for t in T_DATA] # element index whose END is t (0 -> start)
data = [PSTAT5A_REL, PSTAT5B_REL, RSTAT5A_REL]
c = ExaCore(; backend=backend)
@add_var(c, θ, 1:6; start=0.0, lvar=-5.0, uvar=5.0)
@add_var(c, logσ, 1:3; start=0.5, lvar=-3.0, uvar=3.0)
# z_ij: state at node j of element i; 8 species
x0 = zeros(N, K+1, 8); x0[:,:,1] .= A0; x0[:,:,2] .= B0
@add_var(c, x, 1:N, 0:K, 1:8; start=x0, lvar=0.0)
# observables at a measurement time: node K of that element (node 0 of element 1 for t=0)
idx = [(e == 0 ? (1,0) : (e,K)) for e in elem]
yA(i,j) = (100*x[i,j,4] + 200*x[i,j,3]*c17) / (x[i,j,4] + x[i,j,1]*c17 + 2*x[i,j,3]*c17)
yB(i,j) = (-100*x[i,j,4] + 200*x[i,j,5]*(c17-1)) / (x[i,j,2]*(c17-1) - x[i,j,4] + 2*x[i,j,5]*(c17-1))
yR(i,j) = (100*x[i,j,4] + 100*x[i,j,1]*c17 + 200*x[i,j,3]*c17) /
(2*x[i,j,4] + x[i,j,1]*c17 + 2*x[i,j,3]*c17 - x[i,j,2]*(c17-1) - 2*x[i,j,5]*(c17-1))
obsA = [(i=p[1], j=p[2], d=d) for (p,d) in zip(idx, data[1])]
obsB = [(i=p[1], j=p[2], d=d) for (p,d) in zip(idx, data[2])]
obsR = [(i=p[1], j=p[2], d=d) for (p,d) in zip(idx, data[3])]
@add_obj(c, (yA(e.i,e.j)-e.d)^2/(2*(10.0^logσ[1])^2) for e in obsA)
@add_obj(c, (yB(e.i,e.j)-e.d)^2/(2*(10.0^logσ[2])^2) for e in obsB)
@add_obj(c, (yR(e.i,e.j)-e.d)^2/(2*(10.0^logσ[3])^2) for e in obsR)
@add_obj(c, length(T_DATA)*log(10.0)*logσ[q] for q=1:3)
# dynamics at each collocation point
# the time of each collocation point is data, not a symbolic index
E(e) = 1.25e-7*exp(-(10.0^θ[1])*e.t)
v1(e)=Vc*E(e)*(10.0^θ[2])*x[e.i,e.k,1]^2; v2(e)=Vc*E(e)*(10.0^θ[2])*x[e.i,e.k,1]*x[e.i,e.k,2]
v3(e)=Vc*E(e)*(10.0^θ[2])*x[e.i,e.k,2]^2; v4(e)=Vc*(10.0^θ[3])*x[e.i,e.k,3]
v5(e)=Vc*(10.0^θ[4])*x[e.i,e.k,4]; v6(e)=Vc*(10.0^θ[3])*x[e.i,e.k,5]
v7(e)=Vn*(10.0^θ[5])*x[e.i,e.k,6]; v8(e)=Vn*(10.0^θ[6])*x[e.i,e.k,7]; v9(e)=Vn*(10.0^θ[5])*x[e.i,e.k,8]
f = (e->(-2v1(e)-v2(e)+2v7(e)+v8(e))/Vc, e->(-2v3(e)-v2(e)+2v9(e)+v8(e))/Vc,
e->(v1(e)-v4(e))/Vc, e->(v2(e)-v5(e))/Vc, e->(v3(e)-v6(e))/Vc,
e->(v4(e)-v7(e))/Vn, e->(v5(e)-v8(e))/Vn, e->(v6(e)-v9(e))/Vn)
pts = [(n=(i-1)*K+k, i=i, k=k, t=(i-1)*h + RADAU_C[k]*h) for i=1:N for k=1:K]
trm = [(n=(i-1)*K+k, i=i, j=j, d=D[k,j+1]) for i=1:N for k=1:K for j=0:K]
@add_con(c, cc1, -h*f[1](e) for e in pts)
@add_con!(c, cc1[e.n] += e.d*x[e.i,e.j,1] for e in trm)
@add_con(c, cc2, -h*f[2](e) for e in pts)
@add_con!(c, cc2[e.n] += e.d*x[e.i,e.j,2] for e in trm)
@add_con(c, cc3, -h*f[3](e) for e in pts)
@add_con!(c, cc3[e.n] += e.d*x[e.i,e.j,3] for e in trm)
@add_con(c, cc4, -h*f[4](e) for e in pts)
@add_con!(c, cc4[e.n] += e.d*x[e.i,e.j,4] for e in trm)
@add_con(c, cc5, -h*f[5](e) for e in pts)
@add_con!(c, cc5[e.n] += e.d*x[e.i,e.j,5] for e in trm)
@add_con(c, cc6, -h*f[6](e) for e in pts)
@add_con!(c, cc6[e.n] += e.d*x[e.i,e.j,6] for e in trm)
@add_con(c, cc7, -h*f[7](e) for e in pts)
@add_con!(c, cc7[e.n] += e.d*x[e.i,e.j,7] for e in trm)
@add_con(c, cc8, -h*f[8](e) for e in pts)
@add_con!(c, cc8[e.n] += e.d*x[e.i,e.j,8] for e in trm)
@add_con(c, x[i,K,s] - x[i+1,0,s] for i=1:N-1, s=1:8)
@add_con(c, x[1,0,1] - A0); @add_con(c, x[1,0,2] - B0)
@add_con(c, x[1,0,s] for s=3:8)
return ExaModel(c)
end
boehm_model (generic function with 1 method)
Solve it on the GPU, with one departure from the rule above: tol = 1e-6
rather than 1e-8. This problem is the exception. At 1e-8 it converged
in one run out of six we tried, the others wandering for thousands of
iterations toward a worse local optimum; at 1e-6 it converged in all
six, in 49 to 193 iterations.
N = 96 # 96 elements of 2.5 min: every measurement time is a boundary
model = boehm_model(N; backend = CUDABackend())
result = madnlp(model; max_iter = 600, tol = 1e-6)
(status = result.status, objective = result.objective, iterations = result.iter)
This is MadNLP version v0.10.1, running with cuDSS v0.8.0 Number of nonzeros in constraint Jacobian............: 23416 Number of nonzeros in Lagrangian Hessian.............: 46864 Total number of variables............................: 3081 variables with only lower bounds: 3072 variables with lower and upper bounds: 9 variables with only upper bounds: 0 Total number of equality constraints.................: 0 Total number of inequality constraints...............: 3072 inequality constraints with only lower bounds: 0 inequality constraints with lower and upper bounds: 3072 inequality constraints with only upper bounds: 0 iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 0 7.4432489e+03 5.28e-02 1.00e+02 1.44e+02 -1.0 - 0.00e+00 0 0 1 7.4429561e+03 5.28e-02 1.07e+00 7.57e-01 -1.0 - 4.00e-05 1 1h 2 6.3148059e+03 7.93e-02 7.61e-01 5.63e-01 -1.0 - 1.67e-01 2 1h 3 6.0432776e+03 7.57e-02 6.99e+00 2.58e-01 -1.7 - 4.48e-02 1 1h 4 5.7329162e+03 6.01e-02 1.92e+01 2.07e-01 -1.7 - 5.37e-02 1 1h 5 5.1063623e+03 5.31e-02 2.27e+01 1.84e-01 -1.7 - 1.18e-01 1 1h 6 2.8101393e+03 1.15e-01 2.17e+02 1.84e-01 -1.7 - 6.16e-01 1 1h 7 1.8283912e+03 2.98e-01 2.56e+02 1.81e-01 -1.7 - 4.56e-01 1 1h 8 8.5090339e+02 2.78e-01 3.50e+02 1.55e-01 -1.7 0.0 8.79e-01 1 1h 9 4.1702242e+02 5.63e-03 2.16e+01 1.46e-01 -1.7 -0.5 9.97e-01 1 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 10 2.6521423e+02 7.09e-01 8.86e+01 6.45e-02 -1.7 - 1.00e+00 1 1h 11 2.2201159e+02 9.35e-02 1.07e+00 5.71e-03 -1.7 -0.1 1.00e+00 1 1h 12 2.0644652e+02 1.60e+03 1.80e+04 5.83e-04 -1.7 - 1.00e+00 2 1h 13 2.0643885e+02 5.87e+02 6.59e+03 3.40e-04 -1.7 3.1 1.00e+00 2 1h 14 2.0642537e+02 2.16e+02 2.44e+03 3.01e-04 -1.7 2.6 1.00e+00 2 1h 15 2.0638501e+02 7.95e+01 9.00e+02 3.00e-04 -1.7 2.1 1.00e+00 2 1h 16 2.0627049e+02 2.92e+01 3.31e+02 3.00e-04 -1.7 1.7 1.00e+00 2 1h 17 2.0597353e+02 1.07e+01 1.22e+02 3.00e-04 -1.7 1.2 1.00e+00 2 1h 18 2.0536529e+02 3.94e+00 4.48e+01 3.00e-04 -1.7 0.7 1.00e+00 2 1h 19 2.0458786e+02 1.44e+00 1.65e+01 3.00e-04 -1.7 0.2 1.00e+00 2 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 20 2.0406362e+02 5.20e-01 6.07e+00 3.00e-04 -1.7 -0.3 1.00e+00 2 1h 21 2.0375843e+02 1.80e-01 2.23e+00 3.00e-04 -1.7 -0.7 1.00e+00 2 1h 22 2.0335493e+02 5.14e-02 7.87e-01 3.01e-04 -1.7 -1.2 1.00e+00 2 1h 23 2.0279511e+02 2.64e-01 1.71e-01 3.00e-04 -1.7 -1.7 1.00e+00 2 1h 24 1.9993551e+02 1.00e+01 1.24e+01 3.05e-04 -1.7 - 1.00e+00 2 1h 25 1.9988944e+02 3.03e+00 7.97e-02 3.09e-04 -1.7 -0.4 1.00e+00 2 1h 26 1.9985459e+02 1.71e+00 6.14e-03 3.00e-04 -1.7 -0.8 1.00e+00 3 1h 27 1.9382678e+02 1.97e+01 1.76e-01 3.00e-04 -1.7 - 1.00e+00 5 1h 28 1.9379427e+02 3.21e-01 8.68e-03 3.06e-04 -1.7 -0.4 1.00e+00 3 1h 29 1.9379193e+02 4.32e-01 2.06e-03 3.00e-04 -1.7 0.0 1.00e+00 3 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 30 1.8067948e+02 4.01e+01 1.09e-01 3.00e-04 -1.7 - 1.00e+00 6 1h 31 1.8056098e+02 1.90e+00 3.88e-01 3.06e-04 -1.7 0.4 1.00e+00 2 1h 32 1.8055880e+02 2.47e-02 2.53e-02 3.00e-04 -1.7 -0.0 1.00e+00 2 1h 33 1.8055051e+02 6.38e-05 3.45e+01 5.91e-04 -3.8 -0.5 1.00e+00 1 1h 34 1.7845866e+02 1.19e-01 3.60e+03 3.13e-04 -3.8 - 1.00e+00 2 1h 35 1.7479766e+02 4.60e-01 1.11e+03 3.92e-04 -3.8 - 9.41e-01 7 1h 36 1.7478378e+02 7.18e-03 2.88e-03 1.51e-04 -3.8 -1.0 1.00e+00 2 1h 37 1.7478273e+02 1.13e-03 6.77e-04 1.52e-04 -3.8 -1.5 1.00e+00 4 1H 38 1.7478109e+02 7.16e-04 5.41e-01 4.02e-06 -5.7 -1.9 1.00e+00 1 1h 39 1.7477611e+02 2.02e-02 3.63e-05 2.04e-06 -5.7 -2.4 1.00e+00 2 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 40 1.7476115e+02 1.52e-01 3.23e-05 1.85e-06 -5.7 -2.9 1.00e+00 3 1h 41 1.7471612e+02 1.19e+00 3.30e-05 1.84e-06 -5.7 -3.4 1.00e+00 3 1h 42 1.7458016e+02 5.77e+00 5.79e-05 1.84e-06 -5.7 -3.9 1.00e+00 4 1h 43 1.7417639e+02 6.02e+00 5.51e-05 1.85e-06 -5.7 -4.3 1.00e+00 2 1h 44 1.7299116e+02 1.08e+01 1.77e-04 1.84e-06 -5.7 -4.8 1.00e+00 3 1h 45 1.6960273e+02 8.15e+00 1.17e-03 1.84e-06 -5.7 -5.3 1.00e+00 2 1h 46 1.5912666e+02 8.45e+00 7.27e-03 1.84e-06 -5.7 -5.8 1.00e+00 2 1h 47 1.3522879e+02 1.14e+01 5.45e-02 1.85e-06 -5.7 -6.2 2.80e-01 2 1h 48 1.1639914e+02 8.86e+00 6.73e-02 8.01e-04 -5.7 -5.8 5.00e-01 2 2h 49 1.1644044e+02 3.04e+00 9.21e-02 2.37e-06 -5.7 -5.4 1.00e+00 2 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 50 1.0726088e+02 2.98e+00 1.16e+00 1.96e-06 -5.7 -5.9 1.00e+00 2 1h 51 1.1219142e+02 2.36e+00 8.04e-01 1.95e-06 -5.7 -6.3 3.14e-01 2 1h 52 1.1257089e+02 1.19e+00 4.05e-01 2.51e-05 -5.7 -6.8 1.00e+00 2 1h 53 1.1059983e+02 1.15e+00 3.66e-01 9.57e-06 -5.7 -5.5 1.13e-01 2 1H 54 1.0422606e+02 4.28e+00 1.69e-01 1.36e-05 -5.7 - 1.00e+00 2 1h 55 1.0131400e+02 1.54e+00 6.23e-02 7.40e-06 -5.7 - 1.00e+00 2 1h 56 1.0061686e+02 1.28e+00 5.10e-02 3.33e-06 -5.7 - 1.95e-01 2 1h 57 1.0078425e+02 5.64e-01 1.86e-02 6.07e-05 -5.7 - 1.00e+00 2 1h 58 1.0127542e+02 2.08e-01 6.76e-03 6.44e-06 -5.7 - 1.00e+00 2 1h 59 1.0166847e+02 4.33e-02 2.48e-03 1.85e-06 -5.7 - 1.00e+00 2 1h iter objective inf_pr inf_du inf_compl lg(mu) lg(rg) alpha_pr ir ls 60 1.0164572e+02 1.63e-02 9.12e-04 1.84e-06 -5.7 - 1.00e+00 2 1h 61 1.0164577e+02 5.99e-03 3.34e-04 1.84e-06 -5.7 - 1.00e+00 1 1h 62 1.0164619e+02 2.18e-03 1.22e-04 1.84e-06 -5.7 - 9.81e-01 1 1h 63 1.0166430e+02 4.27e-04 2.38e-05 4.83e-05 -5.7 - 1.00e+00 2 1h 64 1.0165877e+02 2.43e-04 1.36e-05 1.13e-05 -5.7 - 1.00e+00 2 1h 65 1.0165797e+02 2.70e-05 1.51e-06 2.38e-06 -5.7 - 1.00e+00 2 1h 66 1.0165799e+02 1.04e-07 5.79e-09 1.84e-06 -5.7 - 1.00e+00 1 1h 67 1.0165178e+02 3.92e-03 1.02e-04 2.78e-07 -7.0 - 1.00e+00 2 1h 68 1.0165157e+02 9.12e-05 1.79e-06 9.30e-08 -7.0 - 1.00e+00 1 1h 69 1.0165156e+02 4.18e-09 6.41e-11 9.09e-08 -7.0 - 1.00e+00 1 1h Number of Iterations....: 69 (scaled) (unscaled) Objective...............: 4.9783606201825115e-01 1.0165156234542587e+02 Dual infeasibility......: 6.4130406845755009e-11 1.3094583834065040e-08 Constraint violation....: 4.1791979974049499e-09 4.1791979974049499e-09 Complementarity.........: 4.4522547895846272e-10 9.0909174696265920e-08 Overall NLP error.......: 9.0909174696265920e-08 9.0909174696265920e-08 Number of objective function evaluations = 74 Number of objective gradient evaluations = 70 Number of constraint evaluations = 74 Number of constraint Jacobian evaluations = 70 Number of Lagrangian Hessian evaluations = 69 Number of KKT factorizations = 124 Number of KKT backsolves = 171 Total wall secs in initialization = 14.361 s Total wall secs in linear solver = unavailable Total wall secs in NLP function evaluations = 12.420 s Total wall secs in solver (w/o init./fun./lin. alg.) = unavailable Total wall secs = 45.895 s EXIT: Optimal Solution Found (tol = 1.0e-06).
(status = MadNLP.SOLVE_SUCCEEDED, objective = 101.65156234542587, iterations = 69)
The fitted parameters¶
Back on the natural scale, since these are $10^{\theta}$:
θ̂ = Array(solution(result, model.refs.θ))
names = ["k_deg", "k_phos", "k_imp_homo", "k_imp_hetero", "k_exp_homo", "k_exp_hetero"]
[(names[i], 10.0^θ̂[i]) for i = 1:6]
6-element Vector{Tuple{String, Float64}}:
("k_deg", 0.02997564974812803)
("k_phos", 8731.322742859838)
("k_imp_homo", 0.02113558337604944)
("k_imp_hetero", 0.014763496205740198)
("k_exp_homo", 0.009264690564158696)
("k_exp_hetero", 1.0015167681662182e-5)
The fit¶
The basic check on any estimation is to overlay the fitted model's predictions on the data it was fitted to:
# x̂ is indexed (element, node, species); the element boundaries are node 0
# of the first element followed by node K of each one.
x̂ = Array(solution(result, model.refs.x))
K = 3
c17 = 0.107
bnd = vcat([x̂[1, 1, :]], [x̂[i, K + 1, :] for i = 1:N])
yA = [(100 * z[4] + 200 * z[3] * c17) /
(z[4] + z[1] * c17 + 2 * z[3] * c17) for z in bnd]
yB = [(-100 * z[4] + 200 * z[5] * (c17 - 1)) /
(z[2] * (c17 - 1) - z[4] + 2 * z[5] * (c17 - 1)) for z in bnd]
t_grid = range(0, T_DATA[end]; length = N + 1)
plt = plot(t_grid, [yA yB]; lw = 2, label = ["pSTAT5A_rel (fit)" "pSTAT5B_rel (fit)"],
xlabel = "time [min]", ylabel = "relative phosphorylation [%]",
title = "Boehm STAT5 model: fit vs. data")
scatter!(plt, T_DATA, PSTAT5A_REL; label = "pSTAT5A_rel (data)", ms = 5)
scatter!(plt, T_DATA, PSTAT5B_REL; label = "pSTAT5B_rel (data)", ms = 5)
This notebook was generated using Literate.jl.