Replacing tunables of dependent parameter vector fails

[BenChung]

Describe the bug 🐞

I'm working on implementing a successive convexification algorithm and as part of that need to do structurally-aware sensitivity analysis for the subproblems of a multiple shooting system. I've already successfully implemented this treating all of the subproblems as a single EnsembleProblem, but this ends up (with a FOH discretization) doing a lot of extra work.

Towards this end, I've broken up the FOH parameterization into individual parameters that can be set as tunable or not tunable. However, I then need to pass this into the mechanics for SampledData, which requires that I then construct an array out of it. When I try to do this using parameter_dependencies it causes the replacement of tunables in the forward-mode AD to fail.

The code I'm using (apologies for how non-minimal it is) is:

using LinearAlgebra
using OrdinaryDiffEq
using ModelingToolkit, Symbolics
using SymbolicIndexingInterface, SciMLStructures
using ComponentArrays, RuntimeGeneratedFunctions

using ModelingToolkit: t_nounits as t, D_nounits as D
using ModelingToolkitStandardLibrary, ModelingToolkitStandardLibrary.Blocks

Symbolics.@register_symbolic get_sampled_data_internal(t::Float64, buffer::Vector{Float64}, dt, circular_buffer)
function get_sampled_data_internal(a, b, c, d)
    @show b
    ModelingToolkitStandardLibrary.Blocks.get_sampled_data(a, collect(b), convert(eltype(b), c), d)
end
@component function first_order_hold(; name, N, dt)
    syms = [Symbol("val$i") for i=1:N]
    params = [first(@parameters $sym = 0.0) for sym in syms]
    @parameters vals[1:N]=zeros(N) [tunable=false]
    systems = @named begin
        output = RealOutput()
    end
    eqs = [
        output.u ~ ifelse(t < dt*N, get_sampled_data_internal(t, vals, dt, false), params[end])
    ]
    return ODESystem(eqs, t, [], params; name, systems, continuous_events = [t % dt ~ 0], parameter_dependencies = [vals => params])
end

@mtkmodel DblInt begin
    @parameters begin
        m, [tunable = false]
        τ = 1.0, [tunable = false]
    end
    @variables begin
        f(t)
        x(t)
        v(t)
    end
    @equations begin
        D(v) ~ τ * f / m
        D(x) ~ τ * v
    end
end

function build_example_problem()
    @named dblint = DblInt()
    @named input = first_order_hold(N = 20, dt=0.05)
        
    @named model = ODESystem([dblint.f ~ input.output.u], t,
        systems = [dblint, input])
    return model # structural_simplify(model)
end

sys = build_example_problem()
RuntimeGeneratedFunctions.init(@__MODULE__)


tsys = structural_simplify(sys)
params = ModelingToolkit.MTKParameters(tsys, [tsys.dblint.m => 0.25])
tunable, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), params)
SciMLStructures.replace(SciMLStructures.Tunable(), params, tunable)

This currently fails with the exception

ERROR: BoundsError: attempt to access 1-element Vector{Vector{Float64}} at index [1:20]
Stacktrace:
  [1] throw_boundserror(A::Vector{Vector{Float64}}, I::Tuple{UnitRange{Int64}})
    @ Base .\abstractarray.jl:737
  [2] checkbounds
    @ .\abstractarray.jl:702 [inlined]
  [3] getindex
    @ .\array.jl:973 [inlined]
  [4] _split_helper(::Type, buf_v::Vector{Float64}, ::Val{-1}, raw::Vector{Vector{Float64}}, idx::Base.RefValue{Int64})
    @ ModelingToolkit C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:239
  [5] #145
    @ C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:226 [inlined]
  [6] iterate
    @ .\generator.jl:47 [inlined]
  [7] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
    @ Base .\array.jl:854
  [8] collect_similar(cont::Vector{Vector{…}}, itr::Base.Generator{Vector{…}, ModelingToolkit.var"#145#146"{…}})
    @ Base .\array.jl:763
  [9] map(f::Function, A::Vector{Vector{Float64}})
    @ Base .\abstractarray.jl:3285
 [10] _split_helper
    @ C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:226 [inlined]
 [11] _split_helper
    @ C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:222 [inlined]
 [12] #149
    @ C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:246 [inlined]
 [13] ntuple
    @ .\ntuple.jl:48 [inlined]
 [14] split_into_buffers
    @ C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:246 [inlined]
 [15] macro expansion
    @ C:\Users\benchung\.julia\packages\Setfield\PdKfV\src\sugar.jl:197 [inlined]
 [16] replace(::SciMLStructures.Tunable, p::ModelingToolkit.MTKParameters{…}, newvals::Vector{…})
    @ ModelingToolkit C:\Users\benchung\.julia\packages\ModelingToolkit\H1TZ3\src\systems\parameter_buffer.jl:296
 [17] top-level scope
    @ c:\Users\benchung\Documents\work\SCP.jl\src\SCP.jl:265
Some type information was truncated. Use `show(err)` to see complete types.

@AayushSabharwal