Error with MTKParameters and Dual Numbers

[timknab]

It seems that the new MTKParameters system does not work with Dual Numbers from things like Automatic Differentiation. Trying to follow the advice given in the FAQ for using ModelingToolkit with Optimization/AutomaticDifferentiation results in something akin to MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{OptimizationForwardDiffExt.var"#37#55"{…}, Float64}, Float64, 4})

using ModelingToolkit
using DifferentialEquations
using Plots
using Optimization
using OptimizationOptimJL
using SymbolicIndexingInterface
using SciMLStructures: replace!, Tunable, canonicalize

function lotka_volterra(;name)
    @parameters t
    pars = @parameters α β γ δ
    D = Differential(t)
    vars = @variables x(t) y(t)
    eqs = [
        D(x) ~ (α - β * y) * x, # prey
        D(y) ~ (δ * x - γ) * y # predator
    ]
    tspan = (0.0, 10.0)
    sys = ODESystem(eqs, t, vars, pars; tspan = tspan, name = name)
    return sys
end

@mtkbuild lvsys = lotka_volterra();
prob = ODEProblem(lvsys, [lvsys.x => 1.0, lvsys.y => 1.0], lvsys.tspan, [lvsys.α => 1.5, lvsys.β => 1.0, lvsys.γ => 3.0, lvsys.δ => 1.0]);

# Plot simulation.
plot(solve(prob, Tsit5()))

sol = solve(prob, Tsit5(); saveat=0.1);
data = Array(sol) + 0.01 * randn(size(Array(sol)));

function loss(x, p)
    prob = p[1]
    prob = remake(prob, p = [:α => x[1], :β => x[2], :γ => x[3], :δ => x[4]])
    sol = solve(prob, Tsit5(), saveat=0.1);
    return sum((sol .- data).^2)
end

f = OptimizationFunction(loss, Optimization.AutoForwardDiff())
oprob = OptimizationProblem(f, rand(4), [prob]);
sol = solve(oprob, BFGS())

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{OptimizationForwardDiffExt.var"#37#55"{…}, Float64}, Float64, 4})

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:207
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:792
  Float64(::IrrationalConstants.Invsqrt2π)
   @ IrrationalConstants ~/.julia/packages/IrrationalConstants/vp5v4/src/macro.jl:112
  ...

The other methods in the FAQ (i.e setter! and using SciMLStructures) results in the same issue. As far as I can tell, it seems that passing a Dual Number into MTKParameters causes the problem.

Pkg.status()
  [0c46a032] DifferentialEquations v7.13.0
  [961ee093] ModelingToolkit v9.7.0
  [7f7a1694] Optimization v3.24.3
  [36348300] OptimizationOptimJL v0.2.3
  [91a5bcdd] Plots v1.40.2
  [1ed8b502] SciMLSensitivity v7.56.2
  [53ae85a6] SciMLStructures v1.1.0
  [2efcf032] SymbolicIndexingInterface v0.3.12

[ChrisRackauckas]

This is known and getting fixed.

[bgctw]

The commit (#e476efa9bf776593896a35d0bf26deeb8a58e68a) that completed this issue uses the internal function remake_buffer(sys, ps, vals) in the test. Could you, please, provide an example using the public API in a differentiated cost function?

The MWE above, using remake, still gives the described error with the current master of ModelingToolkit.jl and current master of SymbolicIndexingInterface.jl for me.
And also gives the error for the simpler command:

ForwardDiff.gradient(x -> loss(x,[prob]), rand(4))

[AayushSabharwal]

remake_buffer isn't really internal API, it's part of SymbolicIndexingInterface.jl which can be used directly if required. remake needs to be updated to use remake_buffer and support your use case. This is near completion and can be tracked using SciML/SciMLBase.jl#658.

While waiting for the above PR, you can use remake_buffer to re-create the MTKParameters object and pass that to remake. Alternatively, you can use replace(Tunable(), prob.p, x) (see SciMLStructures.jl) inside your loss function to obtain the new parameter object and remake using that.

[AayushSabharwal]

I'll add a tutorial highlighting this particular workflow soon, and the replace approach is preferred

[bgctw]

Thank you, @AayushSabharwal. I am looking forward to the tutorial, since I was not yet able to implement a working cost function using replace(Tunable(), prob.p, newvals).

The documentation of replace in SciMLStructures.jl does not tell about the type and semantics of newvals, and I could not figure it out from the corresponding code that splits up the argument it into some more complicated substructure.

[AayushSabharwal]

The docs are available on master: https://docs.sciml.ai/ModelingToolkit/dev/examples/remake/

newvals has to be a Vector of whatever types the parameters in that portion have. For tunable, all parameters are numeric

[timknab]

This looks great, the example seems really helpful and thorough. Thanks!

[bgctw]

Thanks again, @AayushSabharwal for the tutorial.

One question that is not answered yet, is how to relate the positions in initial parameters estimates, currently rand(4), to the parameters used in replace. I have the following suggestion for improving the tutorial.

# obtain indices of parameters
#  prob_ind = remake(odeprob) # does not copy
prob_ind = remake(odeprob, p = copy(odeprob.p)) 
replace!(Tunable(), parameter_values(prob_ind), eachindex(parameters(odesys)))
pos_opt = Int.(prob_ind.ps[[α, β, γ, δ]])

function loss(x, p)
       odeprob = p[1] # ODEProblem stored as parameters to avoid using global variables
       ps = parameter_values(odeprob) # obtain the parameter object from the problem
       pos_opt = p[4]
       ps = replace(Tunable(), ps, x[pos_opt]) # create a copy with the values passed to the loss function
       T = eltype(x)
       # we also have to convert the `u0` vector
       u0 = T.(state_values(odeprob))
       # remake the problem, passing in our new parameter object
       newprob = remake(odeprob; u0 = u0, p = ps)
       timesteps = p[2]
       sol = solve(newprob, AutoTsit5(Rosenbrock23()); saveat = timesteps)
       truth = p[3]
       data = Array(sol)
       return sum((truth .- data) .^ 2) / length(truth)
   end

optprob = OptimizationProblem(
       optfn, odeprob.ps[[α, β, γ, δ]], (odeprob, timesteps, data, pos_opt), lb = 0.1zeros(4), ub = 3ones(4))

Note, how

• explicit parameters are stated as initial values,
• positions are passed to the optimization as parameters
• remake uses x[pos_opt]

In addition, I have a proposal for a similar example, where only a subset of the tunable parameters is optimized.
Where is the adequate place/issue to put it?

[AayushSabharwal]

That's a pretty cool trick, thanks for sharing! Note that x[pos_opt] copies, and then replace will (potentially) copy again. Using view(x, pos_opt) should help mitigate this, but I haven't tested what happens when views are passed to replace.

If you have another example, please feel free to open a PR adding it to the docs. Since your proposed example relates to this new one, I recommend adding it in the same file under a distinct header.

[AayushSabharwal]

I should also point out that

replace!(Tunable(), parameter_values(prob_ind), eachindex(parameters(odesys)))

is incorrect. Not all parameters in parameters(odesys) are tunable. E.g. ones modified by callbacks are marked as Discrete. A better way to do this would be

values, repack, alias = SciMLStructures.canonicalize(Tunable(), parameter_values(prob_ind))
values .= eachindex(values)
repack(values)

The parameter object will be modified in-place if alias == true (which is the case for MTKParameters)

[bgctw]

If you have another example, please feel free to open a PR adding it to the docs

The PR it out for more than a week now with only a discouraging comment by Chris.

[ChrisRackauckas]

Discouraging? It already got fixed and already got a tutorial? https://docs.sciml.ai/ModelingToolkit/stable/examples/remake/ Could you please be more specific on what you're looking for beyond that?

[bgctw]

Could you please be more specific on what you're looking for beyond that?

Thanks for checking back.
I detailed what I expect at the specific issue: #2606

An efficient and differentiable updating of only a small subset of parameters is important to me, and I hoped that the new MTK9 infrastructure would aim at supporting this.
Your comment was discouraging to me, because it implied to me, that I need to resort to remake with merging dictionaries. I am a bit afraid to be classified a 'support vampire' when I keep coming back to this topic, despite contributing MWEs and PRs.

[ChrisRackauckas]

I think I was just sending a message to Aayush 😅