Optimization of parameters of a dynamic problem with ModelingToolkit

Specific Domains Modelling & Simulations
, ,
1

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 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][1])
    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)
2

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 :slight_smile:

2 Likes
3

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
4

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][1])
    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
ω, ξ, 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[1]
    prob_ss[simple_ode.m] = p[3]
    prob_ss[simple_ode.g] = p[1]
    prob_ss[simple_ode.l0] = p[2]
    # 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)
2 Likes