Invalid index of type Num from loss function.

[mzhen77]

I have the simple system from the documentation.

module Grid1

export get_model

using ModelingToolkit
using ModelingToolkitStandardLibrary.Electrical
using ModelingToolkitStandardLibrary.Blocks: Constant

function get_model(; resistance::Float64=0.4, capacitance::Float64=0.6, voltage::Float64=1.0)
    @variables t
    @named R = Resistor(R=resistance)
    @named C = Capacitor(C=capacitance)
    @named V = Voltage()
    @named ground = Ground()
    @named source = Constant(k=voltage)

    rc_eqs = [
        connect(V.p, R.p)
        connect(R.n, C.p)
        connect(C.n, V.n, ground.g)
        connect(V.V, source.output)
    ]

    @named _rc_model = ODESystem(rc_eqs, t)

    @named rc_model = compose(_rc_model, [R, C, V, source, ground])
    sys = structural_simplify(rc_model)
    return sys, parameters(sys), states(sys), C 
end

end

I want to run optimization of that system wrt parameters [R and C]. When I have the following loss function

function loss(p::Vector{Float64})
    pred_sol = solve(prob, Tsit5(), p=p, saveat=t_save)
    pred = Array(pred_sol')
    mse = mean((pred - obs_vals) .^ 2)
    return mse
end

Everything works correctly.

But If I have the next loss function:

function loss(p::Vector{Float64})
    pred_sol = solve(prob, Tsit5(), p=p, saveat=t_save)
    pred = Array(pred_sol[target_var.v])
    mse = mean((pred - obs_vals) .^ 2)
    return mse
end

I get the error:

ERROR: ArgumentError: invalid index: C₊v(t) of type Num

Where:

julia> target_var
Model C with 4 (6) equations
States (6):
  v(t) [defaults to 0.0]
  i(t) [defaults to 0.0]
  p₊v(t) [defaults to 1.0]
  p₊i(t) [defaults to 1.0]
⋮
Parameters (1):
  C [defaults to 0.6]

julia> target_var.v
C₊v(t)

The full code is here

I have Julia 1.7.3 and packages:

 pkg> st
      Status `D:\Projects\playground\julia\grid\Project.toml`
  [336ed68f] CSV v0.10.4
  [a93c6f00] DataFrames v1.3.4
  [41bf760c] DiffEqSensitivity v6.79.0
  [0c46a032] DifferentialEquations v7.2.0
  [f6369f11] ForwardDiff v0.10.30
  [961ee093] ModelingToolkit v8.17.0
  [16a59e39] ModelingToolkitStandardLibrary v1.4.0
  [7f7a1694] Optimization v3.7.1
  [36348300] OptimizationOptimJL v0.1.1
  [42dfb2eb] OptimizationOptimisers v0.1.0
  [91a5bcdd] Plots v1.31.3
  [295af30f] Revise v3.3.3
  [e88e6eb3] Zygote v0.6.41

[ChrisRackauckas]

https://mtk.sciml.ai/dev/basics/FAQ/#Getting-the-index-for-a-symbol

Just grab the index and use it in the loss function.

[mzhen77]

@ChrisRackauckas 2 points:

1. Out of loss function I can do

indexof(sym,syms) = findfirst(isequal(sym),syms)
var_ind = indexof(target_var.v, model_state)

sol = solve(prob, Tsit5(), p=p_init, saveat=t_save)
q1 = sol[target_var.v]
q2 = sol[model_state[var_ind]]

But if I do it in the loss function:

function loss(p::Vector{Float64})
    pred_sol = solve(prob, Tsit5(), p=p, saveat=t_save)
    pred = Array(pred_sol[model_state[var_ind]])
    mse = mean((pred - obs_vals) .^ 2)
    return mse
end

I get the error:

ERROR: ArgumentError: invalid index: C₊v(t) of type Term{Real, Base.ImmutableDict{DataType, Any}}

2. What to do in case if target variable doesn't exist in the state. Foe example: the state:

julia> model_state
1-element Vector{Term{Real, Base.ImmutableDict{DataType, Any}}}:
 C₊v(t)

But I'm interested in the R:

julia> target_var
Model R with 4 (6) equations
States (6):
  v(t) [defaults to 0.0]
  i(t) [defaults to 0.0]
  p₊v(t) [defaults to 1.0]
  p₊i(t) [defaults to 1.0]
⋮
Parameters (1):
  R [defaults to 0.4]

julia> target_var.v
R₊v(t)

Outside of the loss function I can use the same way: sol[target_var.v].

[mzhen77]

I created the minimum working example to reproduce it:

using ModelingToolkit
using ModelingToolkitStandardLibrary.Electrical
using ModelingToolkitStandardLibrary.Blocks: Constant

using DifferentialEquations, Optimization, Zygote, DiffEqSensitivity, OptimizationOptimisers

@variables t
@named R = Resistor(; R=1.0)
@named C = Capacitor(; C=1.0)
@named V = Voltage()
@named ground = Ground()
@named source = Constant(; k=1.0)

rc_eqs = [
    connect(V.p, R.p)
    connect(R.n, C.p)
    connect(C.n, V.n, ground.g)
    connect(V.V, source.output)
]

@named _rc_model = ODESystem(rc_eqs, t)
@named rc_model = compose(_rc_model, [R, C, V, source, ground])
model = structural_simplify(rc_model)

t_save = collect(0:0.1:2)
tspan = (t_save[1], last(t_save))
obs = rand(length(t_save))

function loss(p::Vector{Float64})
    pred_sol = solve(prob, Tsit5(); p=p, saveat=t_save)
    mse = sum((pred_sol[C.v] - obs) .^ 2)  # Error!
    return mse
end

prob = ODAEProblem(model, [0., 0.], tspan)

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, [0.5, 0.2, 0.1])
result = Optimization.solve(optprob, ADAM(); maxiters=2)

[ChrisRackauckas]

Yeah I think the issue is that we need a way to support differentiation with observed functions. The non-observed case can always use indices so that's fine. Observed functions are a different story. This is missing a derivative rule.

[lamorton]

Right, it boils down to #1014 then?