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())
2

Ok, I should relly read the docs more carefully… It is explained in more details in the example here Optimizing through an ODE solve and re-creating MTK Problems · ModelingToolkit.jl.

The corrected MWE reads

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
using SymbolicIndexingInterface: parameter_values, state_values
using SciMLStructures: Tunable, canonicalize, replace, replace!
using PreallocationTools

# 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]
non_opt_params = setdiff(parameters(sys), opt_params)
setter = ModelingToolkit.setp(sys, opt_params)
getter = ModelingToolkit.getp(sys, parameters(sys))
getter_nonopt = ModelingToolkit.getp(sys, non_opt_params)



function loss_fun(u, p)
    (prob_ss, setter, diffcache) = p
    ps = parameter_values(prob_ss) # obtain the parameter object from the problem
    # get an appropriately typed preallocated buffer to store the `x` values in
    buffer = get_tmp(diffcache, u)
    # copy the current values to this buffer
    copyto!(buffer, canonicalize(Tunable(), ps)[1])
    # create a copy of the parameter object with the buffer
    ps = replace(Tunable(), ps, buffer)
    # set the updated values in the parameter object
    setter(ps, u)
    # remake the problem, passing in our new parameter object
    prob_ss_opt = remake(prob_ss; p=ps)
    sol_ss_opt = solve(prob_ss_opt, DynamicSS(Rodas5()))
    solval = sol_ss_opt[sys.mass.s] - (-3)
    @info solval
    return abs2(solval)
end

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

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

Yup and I found something on v11 so this code should be good but if you notice any issues, well the tests should be getting fixed up over the next week here sorry.

1 Like