refactor: refactor interpolation with indexed clocks

Author: AayushSabharwal

src/solutions/ode_solutions.jl

@@ -180,22 +180,45 @@ function SymbolicIndexingInterface.is_parameter_timeseries(::Type{S}) where {
     Timeseries()
 end
 
+function _hold_discrete(disc_u, disc_t, t::Number)
+    idx = searchsortedlast(disc_t, t)
+    if idx == firstindex(disc_t) - 1
+        error("Cannot access discrete variable at time $t before initial save $(first(disc_t))")
+    end
+    return disc_u[idx]
+end
+
+function hold_discrete(disc_u, disc_t, t::Number)
+    val = _hold_discrete(disc_u, disc_t, t)
+    return DiffEqArray([val], [t])
+end
+
+function hold_discrete(disc_u, disc_t, t::AbstractVector{<:Number})
+    return DiffEqArray(_hold_discrete.((disc_u,), (disc_t,), t), t)
+end
+
 function get_interpolated_discretes(sol::AbstractODESolution, t, deriv, continuity)
     is_parameter_timeseries(sol) == Timeseries() || return nothing
 
     discs::ParameterTimeseriesCollection = RecursiveArrayTools.get_discretes(sol)
     interp_discs = map(discs) do partition
-        ConstantInterpolation(partition.t, partition.u)(t, nothing, deriv, nothing, continuity)
+        hold_discrete(partition.u, partition.t, t)
     end
     return ParameterTimeseriesCollection(interp_discs, parameter_values(discs))
 end
 
 function (sol::AbstractODESolution)(t, ::Type{deriv} = Val{0}; idxs = nothing,
         continuity = :left) where {deriv}
+    if t isa IndexedClock
+        t = canonicalize_indexed_clock(t, sol)
+    end
     sol(t, deriv, idxs, continuity)
 end
 function (sol::AbstractODESolution)(v, t, ::Type{deriv} = Val{0}; idxs = nothing,
         continuity = :left) where {deriv}
+    if t isa IndexedClock
+        t = canonicalize_indexed_clock(t, sol)
+    end
     sol.interp(v, t, idxs, deriv, sol.prob.p, continuity)
 end
 
@@ -247,15 +270,13 @@ function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs,
     if is_parameter(sol, idxs) && !is_timeseries_parameter(sol, idxs)
         return getp(sol, idxs)(ps)
     end
-    # NOTE: This is basically SII.parameter_values_at_time but that isn't public API
-    # and once we move interpolation to SII, there's no reason for it to be
     if is_parameter_timeseries(sol) == Timeseries()
         discs::ParameterTimeseriesCollection = RecursiveArrayTools.get_discretes(sol)
         ps = parameter_values(discs)
         for ts_idx in eachindex(discs)
             partition = discs[ts_idx]
             interp_val = ConstantInterpolation(partition.t, partition.u)(t, nothing, deriv, nothing, continuity)
-            ps = with_updated_parameter_timeseries_values(ps, ts_idx => interp_val)
+            ps = with_updated_parameter_timeseries_values(sol, ps, ts_idx => interp_val)
         end
     end
     state = ProblemState(; u = sol.interp(t, nothing, deriv, ps, continuity), p = ps, t)
@@ -270,15 +291,13 @@ function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs::AbstractVect
     end
     error_if_observed_derivative(sol, idxs, deriv)
     ps = parameter_values(sol)
-    # NOTE: This is basically SII.parameter_values_at_time but that isn't public API
-    # and once we move interpolation to SII, there's no reason for it to be
     if is_parameter_timeseries(sol) == Timeseries()
         discs::ParameterTimeseriesCollection = RecursiveArrayTools.get_discretes(sol)
         ps = parameter_values(discs)
         for ts_idx in eachindex(discs)
             partition = discs[ts_idx]
             interp_val = ConstantInterpolation(partition.t, partition.u)(t, nothing, deriv, nothing, continuity)
-            ps = with_updated_parameter_timeseries_values(ps, ts_idx => interp_val)
+            ps = with_updated_parameter_timeseries_values(sol, ps, ts_idx => interp_val)
         end
     end
     state = ProblemState(; u = sol.interp(t, nothing, deriv, ps, continuity), p = ps, t)
@@ -290,9 +309,21 @@ function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv},
     symbolic_type(idxs) == NotSymbolic() && error("Incorrect specification of `idxs`")
     error_if_observed_derivative(sol, idxs, deriv)
     p = hasproperty(sol.prob, :p) ? sol.prob.p : nothing
+    getter = getu(sol, idxs)
+    if is_parameter_timeseries(sol) == NotTimeseries()
+        interp_sol = augment(sol.interp(t, nothing, deriv, p, continuity), sol)
+        return DiffEqArray(getter(interp_sol), t, p, sol)
+    end
     discretes = get_interpolated_discretes(sol, t, deriv, continuity)
-    interp_sol = augment(sol.interp(t, nothing, deriv, p, continuity), sol; discretes)
-    return DiffEqArray(getu(interp_sol, idxs)(interp_sol), t, p, sol; discretes)
+    interp_sol = sol.interp(t, nothing, deriv, p, continuity)
+    u = map(eachindex(t)) do ti
+        ps = parameter_values(discretes)
+        for i in eachindex(discretes)
+            ps = with_updated_parameter_timeseries_values(sol, ps, i => discretes[i, ti])
+        end
+        return getter(ProblemState(; u = interp_sol.u[ti], p = ps, t = t[ti]))
+    end
+    return DiffEqArray(u, t, p, sol; discretes)
 end
 
 function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv},
@@ -301,34 +332,51 @@ function (sol::AbstractODESolution)(t::AbstractVector{<:Number}, ::Type{deriv},
         error("Incorrect specification of `idxs`")
     error_if_observed_derivative(sol, idxs, deriv)
     p = hasproperty(sol.prob, :p) ? sol.prob.p : nothing
+    getter = getu(sol, idxs)
+    if is_parameter_timeseries(sol) == NotTimeseries()
+        interp_sol = augment(sol.interp(t, nothing, deriv, p, continuity), sol)
+        return DiffEqArray(getter(interp_sol), t, p, sol)
+    end
     discretes = get_interpolated_discretes(sol, t, deriv, continuity)
-    interp_sol = augment(sol.interp(t, nothing, deriv, p, continuity), sol; discretes)
-    return DiffEqArray(
-        getu(interp_sol, idxs)(interp_sol), t, p, sol; discretes)
+    interp_sol = sol.interp(t, nothing, deriv, p, continuity)
+    u = map(eachindex(t)) do ti
+        ps = parameter_values(discretes)
+        for i in eachindex(discretes)
+            ps = with_updated_parameter_timeseries_values(sol, ps, i => discretes[i, ti])
+        end
+        return getter(ProblemState(; u = interp_sol.u[ti], p = ps, t = t[ti]))
+    end
+    return DiffEqArray(u, t, p, sol; discretes)
 end
 
 # public API, used by MTK
 """
-    create_parameter_timeseries_collection(sys, ps)
+    create_parameter_timeseries_collection(sys, ps, tspan)
 
 Create a `SymbolicIndexingInterface.ParameterTimeseriesCollection` for the given system
 `sys` and parameter object `ps`. Return `nothing` if there are no timeseries parameters.
-Defaults to `nothing`.
+Defaults to `nothing`. Falls back on the basis of `symbolic_container`.
 """
 function create_parameter_timeseries_collection(sys, ps, tspan)
-    return nothing
+    if hasmethod(symbolic_container, Tuple{typeof(sys)})
+        return create_parameter_timeseries_collection(symbolic_container(sys), ps, tspan)
+    else
+        return nothing
+    end
 end
 
 const PeriodicDiffEqArray = DiffEqArray{T, N, A, B} where {T, N, A, B <: AbstractRange}
 
 # public API, used by MTK
 """
-    get_saveable_values(ps, timeseries_idx)
+    get_saveable_values(sys, ps, timeseries_idx)
 """
-function get_saveable_values end
+function get_saveable_values(sys, ps, timeseries_idx)
+    return get_saveable_values(symbolic_container(sys), ps, timeseries_idx)
+end
 
 function save_discretes!(integ::DEIntegrator, timeseries_idx)
-    save_discretes!(integ.sol, current_time(integ), get_saveable_values(parameter_values(integ), timeseries_idx), timeseries_idx)
+    save_discretes!(integ.sol, current_time(integ), get_saveable_values(integ, parameter_values(integ), timeseries_idx), timeseries_idx)
 end
 
 save_discretes!(args...) = nothing
@@ -346,9 +394,8 @@ function _save_discretes_internal!(A::AbstractDiffEqArray, t, vals)
 end
 
 function _save_discretes_internal!(A::PeriodicDiffEqArray, t, vals)
-    # This is O(1) because A.t is a range
-    idx = searchsortedlast(A.t, t)
-    if idx == firstindex(A.t) - 1 || A.t[idx] ≉ t
+    idx = length(A.u) + 1
+    if A.t[idx] ≉ t
         error("Tried to save periodic discrete value with timeseries $(A.t) at time $t")
     end
     push!(A.u, vals)
@@ -393,7 +440,11 @@ function build_solution(prob::Union{AbstractODEProblem, AbstractDDEProblem},
     end
 
     ps = parameter_values(prob)
-    discretes = create_parameter_timeseries_collection(prob.f.sys, ps, prob.tspan)
+    if has_sys(prob.f)
+        discretes = create_parameter_timeseries_collection(prob.f.sys, ps, prob.tspan)
+    else
+        discretes = nothing
+    end
     if has_analytic(f)
         u_analytic = Vector{typeof(prob.u0)}()
         errors = Dict{Symbol, real(eltype(prob.u0))}()