feat: implement new SII and SciMLBase interfaces for discrete saving/indexing

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -475,6 +475,46 @@ function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym::Sym
     return nothing
 end
 
+function SymbolicIndexingInterface.is_timeseries_parameter(sys::AbstractSystem, sym)
+    has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing || return false
+    is_timeseries_parameter(ic, sym)
+end
+
+function SymbolicIndexingInterface.timeseries_parameter_index(sys::AbstractSystem, sym)
+    has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing || return false
+    timeseries_parameter_index(ic, sym)
+end
+
+function SymbolicIndexingInterface.parameter_observed(sys::AbstractSystem, sym)
+    if has_index_cache(sys) && (ic = get_index_cache(sys)) !== nothing
+        allvars = vars(sym; op = Symbolics.Operator)
+        ts_idxs = Set{Int}()
+        for var in allvars
+            var = unwrap(var)
+            # FIXME: Shouldn't have to shift systems
+            if istree(var) && (op = operation(var)) isa Shift && op.steps == 1
+                var = only(arguments(var))
+            end
+            ts_idx = check_index_map(ic.discrete_clocks, unwrap(var))
+            ts_idx === nothing && continue
+            push!(ts_idxs, ts_idx)
+        end
+        if length(ts_idxs) == 1
+            ts_idx = only(ts_idxs)
+        else
+            ts_idx = nothing
+        end
+        obsfn = let raw_obs_fn = build_explicit_observed_function(
+                sys, sym; param_only = true)
+            f(p::MTKParameters, t) = raw_obs_fn(p..., t)
+        end
+    else
+        ts_idx = nothing
+        obsfn = build_explicit_observed_function(sys, sym; param_only = true)
+    end
+    return ParameterObservedFunction(ts_idx, obsfn)
+end
+
 function SymbolicIndexingInterface.parameter_symbols(sys::AbstractSystem)
     return full_parameters(sys)
 end

src/systems/parameter_buffer.jl

@@ -420,6 +420,54 @@ function SymbolicIndexingInterface.remake_buffer(sys, oldbuf::MTKParameters, val
     return newbuf
 end
 
+struct NestedGetIndex{T}
+    x::T
+end
+
+function Base.getindex(ngi::NestedGetIndex, idx::Tuple)
+    i, j, k... = idx
+    return ngi.x[i][j][k...]
+end
+
+# Required for DiffEqArray constructor to work during interpolation
+Base.size(::NestedGetIndex) = ()
+
+function SymbolicIndexingInterface.with_updated_parameter_timeseries_values(
+        ps::MTKParameters, args::Pair{A, B}...) where {A, B <: NestedGetIndex}
+    for (i, val) in args
+        ps.discrete[i] = val.x
+    end
+    return ps
+end
+
+function SciMLBase.create_parameter_timeseries_collection(
+        sys::AbstractSystem, ps::MTKParameters, tspan)
+    ic = get_index_cache(sys) # this exists because the parameters are `MTKParameters`
+    has_discrete_subsystems(sys) || return nothing
+    (dss = get_discrete_subsystems(sys)) === nothing && return nothing
+    _, _, _, id_to_clock = dss
+    buffers = []
+
+    for (i, partition) in enumerate(ps.discrete)
+        clock = id_to_clock[i + 1]
+        if clock isa Clock
+            ts = tspan[1]:(clock.dt):tspan[2]
+            push!(buffers, DiffEqArray(NestedGetIndex{typeof(partition)}[], ts, (1, 1)))
+        elseif clock isa SolverStepClock
+            push!(buffers,
+                DiffEqArray(NestedGetIndex{typeof(partition)}[], eltype(tspan)[], (1, 1)))
+        else
+            error("Unhandled clock $clock")
+        end
+    end
+
+    return ParameterTimeseriesCollection(Tuple(buffers), copy(ps))
+end
+
+function SciMLBase.get_saveable_values(ps::MTKParameters, timeseries_idx)
+    return NestedGetIndex(deepcopy(ps.discrete[timeseries_idx]))
+end
+
 function DiffEqBase.anyeltypedual(
         p::MTKParameters, ::Type{Val{counter}} = Val{0}) where {counter}
     DiffEqBase.anyeltypedual(p.tunable)