Optimization of parameters of a dynamic problem with ModelingToolkit

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