Extend example of optimizing MTK System parameters to also optimize initial states

[bgctw]

Is your feature request related to a problem? Please describe.

I find it hard to translate the help given in FAQ on Optimization/AD and the basic example into code that actually performs optimization of a combination of a subset of parameters and initial states.

Describe the solution you’d like

Modify the basic example where the parameters to optimize comprise only a subset of the tunable parameters of the MTK system but include also a subset of initial parameters of the MTK ODE system.

Describe alternatives you’ve considered

I explored different approaches at [#3952] but got confused with SciMLStructures.Initials() and ODEProblem.u0

[bgctw]

I extended the example. It works. But why is the code so slow?

Moreover, why does setting u0 and p at the same remake call, ignore the u0 during the solution (as seen at the first time step of the solution)?

# modification of https://docs.sciml.ai/ModelingToolkit/dev/examples/remake/
# that optimizes a subset of initial states and a subset of parameters

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using OrdinaryDiffEq

using SymbolicIndexingInterface: parameter_values, state_values
using SymbolicIndexingInterface: SymbolicIndexingInterface as SII
using SciMLStructures: Tunable, Initials, canonicalize, replace, replace!
using PreallocationTools

using Optimization
#using OptimizationOptimJL  # BFGS()
import OptimizationPolyalgorithms as OPA
import SciMLSensitivity

import ForwardDiff
import Zygote

@parameters α β γ δ
@variables x(t) y(t)
eqs = [D(x) ~ (α - β * y) * x
       D(y) ~ (δ * x - γ) * y]
@mtkcompile odesys = System(eqs, t)


odeprob = ODEProblem(
    odesys, [x => 1.1, y => 1.2, α => 1.5, β => 1.0, γ => 3.0, δ => 1.0], (0.0, 10.0))
timesteps = 0.0:0.1:10.0
sol = solve(odeprob, Tsit5(); saveat = timesteps)
sol_true = Array(sol)
# add some random noise
data = sol_true + 0.01 * randn(size(sol_true))

paropt = [α, γ, δ]  # paropt = [α, β, γ, δ]
stateopt = [x]
ptrue = vcat(odeprob[stateopt], odeprob.ps[paropt])
n_stateopt = sum(length.(stateopt))
p0 =  ptrue + 0.2*randn(length(ptrue))

function loss(x, p)  # ForwardDiff
    local odeprob = p[1] # ODEProblem stored as parameters to avoid using global variables
    local ps = parameter_values(odeprob) # obtain the parameter object from the problem
    local diffcache = p[5]
    # get an appropriately typed preallocated buffer to store the `x` values in
    local buffer = get_tmp(diffcache, x)
    # copy the current values to this buffer
    copyto!(buffer, canonicalize(Tunable(), ps)[1])
    # create a copy of the parameter object with the buffer
    local ps = replace(Tunable(), ps, buffer)
    # set the updated values in the parameter object
    local par_setter! = p[4]
    local n_stateopt = p[8]
    par_setter!(ps, x[(n_stateopt+1):end])
    #
    # similar for state, use a diffcache and replace u0
    local diffcache_state = p[7]
    local u0new = get_tmp(diffcache_state, x)
    copyto!(u0new, state_values(odeprob))    
    local state_setter! = p[6]
    state_setter!(u0new, x[1:n_stateopt])
    #
    # remake the problem, passing in our new parameter object
    # when setting u0 together with ps, u0 has no effect?
    #newprob = remake(odeprob; u0 = u0new, p = ps)
    local newprob = remake(odeprob; p = ps)
    local newprob = remake(newprob; u0 = u0new)
    local timesteps = p[2]
    local sol = solve(newprob, AutoTsit5(Rosenbrock23()); saveat = timesteps)
    local truth = p[3]
    local data = Array(sol)
    return sum((truth .- data) .^ 2) / length(truth)
end

# function to set the parameters we are optimizing
setter = SII.setp(odeprob, paropt)
# `DiffCache` to avoid allocations.
# `copy` prevents the buffer stored by `DiffCache` from aliasing the one in
# `parameter_values(odeprob)`.
diffcache = DiffCache(copy(canonicalize(Tunable(), parameter_values(odeprob))[1]))
diffcache_state = DiffCache(copy(state_values(odeprob)))
state_setter = SII.setu(odeprob, stateopt)

parg = (odeprob, timesteps, data, setter, diffcache, 
    state_setter, diffcache_state, n_stateopt);

loss(p0, parg)
ForwardDiff.gradient(x -> loss(x, parg), p0)

# manually create an OptimizationFunction to ensure usage of `ForwardDiff`, which will
# require changing the types of parameters from `Float64` to `ForwardDiff.Dual`
#adtype = Optimization.AutoZygote()
#adtype = Optimization.AutoEnzyme()
adtype = Optimization.AutoForwardDiff()
optfn = OptimizationFunction(loss, adtype)
# parameter object is a tuple, to store differently typed objects together
optprob = OptimizationProblem(
    optfn, p0, parg,
    lb = 0.1zeros(length(p0)), ub = 5ones(length(p0)))
#opt_alg = OPA.PolyOpt()    
opt_alg = OPA.LBFGS()
@time sol = solve(optprob, opt_alg)

hcat(p0, sol.u, ptrue)

[DhairyaLGandhi]

why does setting u0 and p at the same remake call, ignore the u0 during the solution

u0 is always taken from the problem, we can check for any missing kwargs in the remake flow.

Have you tried with reverse mode AD as well?

[bgctw]

u0 is always taken from the problem

That is what I expected. However, when I run above code with the commented line rather than the two subsequent remake lines, the initial solution of the problem does not correspond to newprob.u0 but rather to the initials in newprob.ps that do not reflect the update by u0 (Julia 1.11.4).

Have you tried with reverse mode AD as well?

Eventually I want to optimize using Zygote and strongly would like to see an additional example/tutorial that optimizes a subset of parameters and initial values using Zygote. However, for me currently neither setp nor setter_oop work with Zygote (see discussion in #3952).