Optimizing the parameters of a MTK model with automatic differetiation

General Usage
,
1

I’m trying to optimize some parameters of a MTK model to reach some steady-state solution objective, using the recommendations in Frequently Asked Questions · ModelingToolkit.jl. This works fine with finite differences computed gradients (using AutoFiniteDiff) but not with AD.
This previous post seems to be very close, but also dates from 2022. Since then MTK has changed a bit… Is the information there still actual, and setp cannot be used as it mutates the parameters ? Or am I using setp incorrectly ?

Below is a MWE where it’s easy to find the good solution since int basically matches the position of the attachment : basically a spring-damper-mass system.

using ModelingToolkit, DifferentialEquations, ModelingToolkitStandardLibrary
import ModelingToolkit: connect, t_nounits as t, D_nounits as D
import ModelingToolkitStandardLibrary.Mechanical.TranslationalPosition: Mass, Damper, Spring, Fixed
import ModelingToolkitStandardLibrary.Blocks: Sine

# Build model
@named ds = Sine(amplitude=-1, frequency=10)
@named mass = Mass(m=1)
@named spring = Spring(k=4, l=0.)
@named damper = Damper(d=1)
@named fix = Fixed()

eqs = Equation[
    connect(fix.flange, spring.flange_a, damper.flange_a)
    connect(mass.flange, spring.flange_b, damper.flange_b)
]

@named model = ODESystem(eqs, t; systems=[mass, spring, damper, fix])
sys = mtkcompile(model)


# Define transient problem
initial_positions = Dict{Num,Any}([
    mass.s => 1.
    mass.v => 0.
]
)
prob = ODEProblem(sys, initial_positions, (0.0, 100.0))
sol_mass = solve(prob, Rodas5())

# Define steady state problem
prob_ss = SteadyStateProblem(prob)
sol_ss = solve(prob_ss, DynamicSS(Rodas5()), abstol=1e-8, reltol=1e-8)


# Define optimization parameters and functions
opt_params = [sys.fix.s_0]

setter! = ModelingToolkit.setp(sys, opt_params)
function loss_fun(u, p)
    setter!(prob_ss, u)
    loss_sol_ss = solve(prob_ss, DynamicSS(Rodas5()), reltol=1e-8)
    solval = loss_sol_ss[sys.mass.s] - (-3)
    mse = sqrt(sum(abs2, solval))
    return mse
end


# Define optimization problem
u0 = sol_mass.ps[opt_params]
opt_fun = OptimizationFunction(loss_fun, Optimization.AutoForwardDiff()); # Works with AutoFiniteDiff()
opt_prob = OptimizationProblem(opt_fun, u0)
sol_opt = solve(opt_prob, BFGS(), abstol=1e-10, reltol=1e-10)

# Verify solution
setter!(prob, sol_opt.u)
sol_mass = solve(prob, Rodas5())