Setup of OEDLayer with discrete measurements fails

[chplate]

Describe the bug 🐞

Setup of discrete OEDLayer fails if measurement may be taken at final time, i.e., if the end of the time horizon is in the measurement points

Minimal Reproducible Example

using Corleone, CorleoneOED
using LuxCore
using Random
using OrdinaryDiffEq


function compartmental(du, u, p, t)
    θ₁, θ₂, θ₃ = p
    du[1] = θ₃ * (-θ₁ * exp(-θ₁ * t) + θ₂ * exp(-θ₂ * t))
    return
end

tspan = (0.0, 48.0)
u0 = [0.0]
p0 = [0.05884, 4.298, 21.8]
prob = ODEProblem(compartmental, u0, tspan, p0)


measurement_points = [0.166, 0.333, 0.5, 0.666, 1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 6.0, 8.0, 10.0, 12.0, 24.0, 30.0, 48.0]

oed = OEDLayer{true}(
    prob, Tsit5(),
    params = [1, 2, 3],
    bounds_p = (p0, p0),
    measurements = [
        ControlParameter(measurement_points, controls =  0.5 * ones(length(measurement_points)), bounds = (0.0, 1.0)),
    ],
    observed = (u, p, t) -> u[1:1],
)

ps, st = LuxCore.setup(Random.default_rng(), oed)

Error & Stacktrace ⚠️

ps, st = LuxCore.setup(Random.default_rng(), oed)
ERROR: BoundsError: attempt to access 17-element Vector{Int64} at index [18]
Stacktrace:
  [1] getindex(A::Vector{Int64}, i1::Int64)
    @ Base ./essentials.jl:13
  [2] (::CorleoneOED.var"#79#84"{Int64, Vector{Vector{Int64}}})(j::Int64)
    @ CorleoneOED ~Corleone.jl/lib/CorleoneOED/src/oed.jl:169
  [3] iterate
    @ ./generator.jl:47 [inlined]
  [4] _collect(c::Base.OneTo{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
    @ Base ./array.jl:854
  [5] collect_similar
    @ ./array.jl:763 [inlined]
  [6] map
    @ ./abstractarray.jl:3286 [inlined]
  [7] #78
    @ ~Corleone.jl/lib/CorleoneOED/src/oed.jl:166 [inlined]
  [8] iterate
    @ ./generator.jl:47 [inlined]
  [9] collect_to!(dest::Vector{…}, itr::Base.Generator{…}, offs::Int64, st::Int64)
    @ Base ./array.jl:892
 [10] collect_to_with_first!(dest::Vector{…}, v1::Vector{…}, itr::Base.Generator{…}, st::Int64)
    @ Base ./array.jl:870
 [11] _collect(c::Base.OneTo{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
    @ Base ./array.jl:864
 [12] collect_similar(cont::Base.OneTo{Int64}, itr::Base.Generator{Base.OneTo{Int64}, CorleoneOED.var"#78#83"{Vector{…}}})
    @ Base ./array.jl:763
 [13] map(f::Function, A::Base.OneTo{Int64})
    @ Base ./abstractarray.jl:3286
 [14] initialstates(rng::TaskLocalRNG, oed::OEDLayer{true, true, true, SingleShootingLayer{…}, @NamedTuple{…}})
    @ CorleoneOED ~Corleone.jl/lib/CorleoneOED/src/oed.jl:165
 [15] setup(rng::TaskLocalRNG, l::OEDLayer{true, true, true, SingleShootingLayer{…}, @NamedTuple{…}})
    @ LuxCore ~/.julia/packages/LuxCore/kQC9S/src/LuxCore.jl:128