# Optimization of parameters of a dynamic problem with ModelingToolkit

Hello all !

I have a nonlinear ODE problem that I set up with ModelingToolkit (great job with this awesome package !), for which I can compute some steady-state (with SteadyStateProblem) or transient solutions (with ODEroblem).

Now the solutions I get seem perfectly fine, but I would like to optimize some parameters of my problem to reach some objective, e.g a steady-state value and some oscillation frequency.
Since ModelingTookit supports OptimizationProblems I figure this could be leveraged for this purpose, but how can I combine or wrap an OptimizationProblem around a SteadyStateProblem ?

To make the discussion more practical, I made a very simple mass-spring model to illustrate this on a minimal working example.

The main thing I have trouble with is should these parameters of the ODE problem be declared as parameters or as variables, since they become the variables of the optimization problem ? What are the good practices to wrap the optimization around the ODE problem to avoid complicating the optimization process ?

Thanks.

using GLMakie
using ControlSystemsBase

# Definition of model parameters
@parameters m g l0 k c

# Dict of parameters
ps = Dict([
m => 1,
l0 => -0.5,
g => 9.81,
k => 1,
c => 0.1,
])

# Definition of problem variables
@variables begin
t,
y(t)
yt(t)
ytt(t)
T(t)
end

# Definition of transient problem variables
D = Differential(t)
D2 = D^2

# Definition of transient problem equations
eqs_tot_t = [
T ~ -k*(y-l0),
m * ytt + c * yt ~ -m*g + T,
D2(y) ~ ytt,
D(y) ~ yt,
]

# Definition of the ODE problem including variables, parameters and time span of interest
@named ode = ODESystem(eqs_tot_t, t, [y,yt,ytt,T], [m,c,k,l0,g])
simple_ode = structural_simplify(ode)

# Definition of initial problem conditions for transient problem
guess_t = [
y => 2*ps[l0],
yt => 0,
ytt => 0,
]

# Definition of steady-state problem to solve for steady solution from the ODE problem
# Actually solve for the steady state solution
sol_ss = solve(prob_ss, DynamicSS(Rodas5()), abstol=1e-8, reltol=1e-8)

# Definition of the ODE problem to solve for a transient solution using the guess
prob = ODEProblem(simple_ode, guess_t, (0, 100), ps)
# Actually solve the transient problem
sol = solve(prob, Rodas5(), abstol=1e-8, reltol=1e-8, saveat=0:0.1:100)

# Define iteration observable
i = Observable(1)

# Build the plot for the transient simulation
begin
# Initialize figure
f = Figure()
# Initialize the axis
ax1 = Axis(f[:, 1], autolimitaspect=1)
# Position reference geometry points as needed
PtO = Observable(Point2f(0, 0))

st = @lift(sol.t[$i]) sy = @lift(sol[simple_ode.y][$i])
syt = @lift(sol[simple_ode.yt][$i]) sytt = @lift(sol[simple_ode.ytt][$i])
sT = @lift(sol[simple_ode.T][$i]) PtM = @lift(Point2(0,$sy))
PtM_0 = Point2(0, sol[simple_ode.y])
PtM_ss = Point2(0, sol_ss[simple_ode.y])

widths = Vec2f(0.1,0.1)

rects = @lift([Rect2(Vec2f($PtM - widths/2),widths),]) rects_0 = [Rect2(Vec2f(PtM_0) - widths/2, widths),] rects_ss = [Rect2(Vec2f(PtM_ss) - widths / 2, widths),] bar_0 = [PtO[],PtM_0] bar_ss = [PtO[], PtM_ss] bar = @lift([$PtO, $PtM]) poly!(rects_0,color=:red) poly!(rects_ss, color=:green) poly!(rects,color=:transparent,strokecolor=:blue, strokewidth=5) linesegments!(bar_0,linestyle = :dot) linesegments!(bar) linesegments!(bar_ss, linestyle=:dash) ymin,ymax = extrema(sol[simple_ode.y]) ylims!(2*ymin,2*ymax) g2 = GridLayout(f[1, 2]) ax22 = Axis(g2[1, 1], ylabel="Force magnitude [N]", xlabel="Time [s]") lines!(ax22, sol.t, sol[simple_ode.T]) lines!(ax22, sol.t, fill(ps[m]*ps[g],length(sol.t)) ) vlines!(st, color=:black) scatter!(st, sT, color=:black) g3 = GridLayout(f[1, 3]) ax31 = Axis(g3[1, 1], xlabel="Time [s]", ylabel="Position [m]", aspect=1) lines!(ax31, sol.t, sol[simple_ode.y]) vlines!(ax31, st, color=:black) scatter!(ax31, st, sy, color=:black) ax32 = Axis(g3[2, 1], xlabel="Time [s]", ylabel="Velocity yt [m/s]", aspect=1) lines!(sol.t, sol[simple_ode.yt]) vlines!(st, color=:black) scatter!(st, syt, color=:black) ax33 = Axis(g3[3, 1], xlabel="Time [s]", ylabel="Acceleration ytt [m/s^2]", aspect=1) lines!(sol.t, sol[simple_ode.ytt]) vlines!(st, color=:black) scatter!(st, sytt, color=:black) linkxaxes!(ax31, ax32, ax33) i[] = 1 display(GLMakie.Screen(), f) end # Define parameters for animation framerate = 30 timestamps = range(1, length(sol.t), step=1) for ii in timestamps i[] = ii sleep(0.01) end # Operating point for linearization analysis to retrieve dynamic properties op = merge(Dict(states(simple_ode) .=> sol_ss.u), ps) # Linearize the system (of course already linear in this case) (; A, B, C, D), simplified_sys = linearize(ode, [], [y]; t=last(sol.t), op=op) # Build a steady state model from matrices A, B, C, D state = ss(A, B, C, D) # Retrieve dynamic properties ω, ξ, ps = damp(state) # Get system frequencies ff = ω ./ (2 * pi) ################################################################################################################################################################## #####################################Optmization on the m and k parameters to reach a frequency and steady state position of interest ############################ ################################################################################################################################################################## ################################################################################################################################################################## # Define optimization objectives for y and f y_obj = -2 f_obj = 5 function loss(p) # Solve the steady state problem for a set of parameter values sol_ss = solve(prob_ss, DynamicSS(Rodas5()), abstol=1e-8, reltol=1e-8) y_test = sol_ss[simple_ode.y] # Linearize to retrive the frequency around steady state op = merge(Dict(st .=> sol_ss.u), p) (; A, B, C, D), _ = linearize(ode, [], [y]; t=last(sol.t), op=op) state = ss(A, B, C, D) ω, ξ, ps = damp(state) f_test = maximum(ω) / (2 * pi) # Calculate mean square error mse = (y_test - y_obj)^2 + (f_test - f_obj)^2 return mse end function callback(p, l) println(l, " ", p) return false end # Load additional packages for optimization (maybe a bit much) using Optimization, OptimizationOptimisers, Zygote, Optim, OptimizationOptimJL **EDIT** # Define the system to optimize where k and m are the variables now and l0 and g are considered fixed @named opt = OptimizationSystem(loss,[k,m],[l0,g]) u0 = [k => 1.0,m=>1.0,g=>9.81,l0 => 1.0] opt_prob = OptimizationProblem(opt, u0)  Hey! I think you can solve this problem analytically. The resonance frequency of the mass-spring-damper is \omega = \sqrt{k / m} and at steady state your derivatives will be zero, so you can rearrange the system of equations to arrive at 0 = -mg - k(y-l_0) If you set \omega = 2\pi f_{obj} you can solve the first equation for m m = \dfrac{k}{(2\pi f_{obj})^2} and plug that into the second equation together with y = y_{obj} to solve for k, which you then plug into the last equation to obtain m. This all went pretty fast, so do check all the signs etc. in my derivation before relying on it 2 Likes Thanks for you answer @baggepinnen, though I used the mass-spring model just for the sake of producing a MWE, this cannot hold for my actual problem. But as you pointed out there is an analytical solution to this example problem, so I can actually check that the optimiser provides the correct optimum, once I figure out to set it up. 2 Likes So I managed to get something working, though the problem must be modified a bit taking @baggepinnen approach : m = \frac{k}{2\pi f_{obj}} introducing that in the force equilibrium one gets 0 = - \frac{kg}{2\pi f_{obj}} - k(y-l0) which has the bad behavior of being independent of k, so one cannot optimize this problem for two criteria since there is actually only one independent variable in this case… So in this very simple model I am now trying to only optimize the frequency of the system, directly tractable analytically from \omega = \sqrt{\frac{k}{m}} Here is the modified MWE for those interested using GLMakie using ModelingToolkit, SteadyStateDiffEq, OrdinaryDiffEq using ControlSystemsBase # Definition of model parameters @parameters m g l0 k c # Dict of parameters ps = Dict([ m => 1, l0 => -0.5, g => 9.81, k => 1, c => 0.1, ]) # Definition of problem variables @variables begin t, y(t) yt(t) ytt(t) T(t) end # Definition of transient problem variables D = Differential(t) D2 = D^2 # Definition of transient problem equations eqs_tot_t = [ T ~ -k * (y - l0), m * ytt + c * yt ~ -m * g + T, D2(y) ~ ytt, D(y) ~ yt, ] # Definition of the ODE problem including variables, parameters and time span of interest @named ode = ODESystem(eqs_tot_t, t, [y, yt, ytt, T], [m, c, k, l0, g]) simple_ode = structural_simplify(ode) # Definition of initial problem conditions for transient problem guess_t = [ y => 2 * ps[l0], yt => 0, ytt => 0, ] # Definition of steady-state problem to solve for steady solution from the ODE problem prob_ss = SteadyStateProblem(simple_ode, guess_t, ps) # Actually solve for the steady state solution sol_ss = solve(prob_ss, DynamicSS(Rodas5()), abstol=1e-8, reltol=1e-8) # Definition of the ODE problem to solve for a transient solution using the guess prob = ODEProblem(simple_ode, guess_t, (0, 100), ps) # Actually solve the transient problem sol = solve(prob, Rodas5(), abstol=1e-8, reltol=1e-8, saveat=0:0.1:100) # Define iteration observable i = Observable(1) # Build the plot for the transient simulation begin # Initialize figure f = Figure() # Initialize the axis ax1 = Axis(f[:, 1], autolimitaspect=1) # Position reference geometry points as needed PtO = Observable(Point2f(0, 0)) st = @lift(sol.t[$i])
sy = @lift(sol[simple_ode.y][$i]) syt = @lift(sol[simple_ode.yt][$i])
sytt = @lift(sol[simple_ode.ytt][$i]) sT = @lift(sol[simple_ode.T][$i])
PtM = @lift(Point2(0, $sy)) PtM_0 = Point2(0, sol[simple_ode.y]) PtM_ss = Point2(0, sol_ss[simple_ode.y]) widths = Vec2f(0.1, 0.1) rects = @lift([Rect2(Vec2f($PtM - widths / 2), widths),])
rects_0 = [Rect2(Vec2f(PtM_0) - widths / 2, widths),]
rects_ss = [Rect2(Vec2f(PtM_ss) - widths / 2, widths),]
bar_0 = [PtO[], PtM_0]
bar_ss = [PtO[], PtM_ss]
bar = @lift([$PtO,$PtM])

poly!(rects_0, color=:red)
poly!(rects_ss, color=:green)
poly!(rects, color=:transparent, strokecolor=:blue, strokewidth=5)
linesegments!(bar_0, linestyle=:dot)
linesegments!(bar)
linesegments!(bar_ss, linestyle=:dash)
ymin, ymax = extrema(sol[simple_ode.y])
ylims!(2 * ymin, 2 * ymax)

g2 = GridLayout(f[1, 2])

ax22 = Axis(g2[1, 1], ylabel="Force magnitude [N]", xlabel="Time [s]")
lines!(ax22, sol.t, sol[simple_ode.T])
lines!(ax22, sol.t, fill(ps[m] * ps[g], length(sol.t)))
vlines!(st, color=:black)
scatter!(st, sT, color=:black)

g3 = GridLayout(f[1, 3])
ax31 = Axis(g3[1, 1], xlabel="Time [s]", ylabel="Position [m]", aspect=1)
lines!(ax31, sol.t, sol[simple_ode.y])
vlines!(ax31, st, color=:black)
scatter!(ax31, st, sy, color=:black)

ax32 = Axis(g3[2, 1], xlabel="Time [s]", ylabel="Velocity yt [m/s]", aspect=1)
lines!(sol.t, sol[simple_ode.yt])
vlines!(st, color=:black)
scatter!(st, syt, color=:black)

ax33 = Axis(g3[3, 1], xlabel="Time [s]", ylabel="Acceleration ytt [m/s^2]", aspect=1)
lines!(sol.t, sol[simple_ode.ytt])
vlines!(st, color=:black)
scatter!(st, sytt, color=:black)

i[] = 1
display(GLMakie.Screen(), f)
end

# Define parameters for animation
framerate = 30
timestamps = range(1, length(sol.t), step=1)
for ii in timestamps
i[] = ii
sleep(0.01)
end

# Operating point for linearization analysis to retrieve dynamic properties
op = merge(Dict(states(simple_ode) .=> sol_ss.u), ps)
# Linearize the system (of course already linear in this case)
(; A, B, C, D), simplified_sys = linearize(ode, [], [y]; t=last(sol.t), op=op)

# Build a steady state model from matrices A, B, C, D
state = ss(A, B, C, D)
# Retrieve dynamic properties
ω, ξ, phi = damp(state)
# Get system frequencies
ff = ω ./ (2 * pi)

##################################################################################################################################################################
#####################################Optmization on the m and k parameters to reach a frequency and steady state position of interest ############################
##################################################################################################################################################################
##################################################################################################################################################################

# Define optimization objectives for y and f
y_obj = -2
f_obj = 5

function loss(u, p)
# @info "loss u,p" u, p
prob_ss[simple_ode.k] = u
prob_ss[simple_ode.m] = p
prob_ss[simple_ode.g] = p
prob_ss[simple_ode.l0] = p
# Solve the steady state problem for a set of parameter values
sol_ss = solve(prob_ss, DynamicSS(Rodas5()), abstol=1e-8, reltol=1e-8)
y_test = sol_ss[simple_ode.y]
# Linearize to retrive the frequency around steady state
op = merge(Dict(states(simple_ode) .=> sol_ss.u), Dict(k => prob_ss[simple_ode.k], m => prob_ss[simple_ode.m], g => prob_ss[simple_ode.g], l0 => prob_ss[simple_ode.l0], c => prob_ss[simple_ode.c], v => prob_ss[simple_ode.v]))
(; A, B, C, D), _ = linearize(ode, [], [y]; t=1e5, op=op)
state = ss(A, B, C, D)
ω, ξ, ps = damp(state)
f_test = maximum(ω) / (2 * pi)
# Calculate mean square error
mse = 0*(y_test - y_obj)^2 + (f_test - f_obj)^2
return mse, (f_test, y_test)
end

function callback(p, mse, vals)
# @info p, mse, vals
(f_test,y_test) = vals
@info "f_test : $f_test | y_test :$y_test | mse : \$mse"
return mse < 1e-5
end

using Optimization, OptimizationOptimisers, Zygote, Optim, OptimizationOptimJL, OptimizationBBO

u0 = [k => 1.0, ]
p = [g => 9.81, l0 => -0.5, m => 1.0]
opt_prob = OptimizationProblem(loss, collect(values(Dict(u0))), collect(values(Dict(p))))
sol = solve(opt_prob, NelderMead(),callback = callback)