diff --git a/src/SciMLBase.jl b/src/SciMLBase.jl index ee1dc0d79..c8c8d0b66 100644 --- a/src/SciMLBase.jl +++ b/src/SciMLBase.jl @@ -724,6 +724,7 @@ include("problems/problem_interface.jl") include("problems/optimization_problems.jl") include("clock.jl") +include("solutions/save_idxs.jl") include("solutions/basic_solutions.jl") include("solutions/nonlinear_solutions.jl") include("solutions/ode_solutions.jl") diff --git a/src/solutions/ode_solutions.jl b/src/solutions/ode_solutions.jl index bb2a9fabc..c92644da0 100644 --- a/src/solutions/ode_solutions.jl +++ b/src/solutions/ode_solutions.jl @@ -104,9 +104,12 @@ https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/ successfully, whether it terminated early due to a user-defined callback, or whether it exited due to an error. For more details, see [the return code documentation](https://docs.sciml.ai/SciMLBase/stable/interfaces/Solutions/#retcodes). +- `saved_subsystem`: a [`SavedSubsystem`](@ref) representing the subset of variables saved + in this solution, or `nothing` if all variables are saved. Here "variables" refers to + both continuous-time state variables and timeseries parameters. """ struct ODESolution{T, N, uType, uType2, DType, tType, rateType, discType, P, A, IType, S, - AC <: Union{Nothing, Vector{Int}}, R, O} <: + AC <: Union{Nothing, Vector{Int}}, R, O, V} <: AbstractODESolution{T, N, uType} u::uType u_analytic::uType2 @@ -124,6 +127,7 @@ struct ODESolution{T, N, uType, uType2, DType, tType, rateType, discType, P, A, retcode::ReturnCode.T resid::R original::O + saved_subsystem::V end function ConstructionBase.constructorof(::Type{O}) where {T, N, O <: ODESolution{T, N}} @@ -137,7 +141,7 @@ function ConstructionBase.setproperties(sol::ODESolution, patch::NamedTuple) patch = merge(getproperties(sol), patch) return ODESolution{T, N}(patch.u, patch.u_analytic, patch.errors, patch.t, patch.k, patch.discretes, patch.prob, patch.alg, patch.interp, patch.dense, patch.tslocation, patch.stats, - patch.alg_choice, patch.retcode, patch.resid, patch.original) + patch.alg_choice, patch.retcode, patch.resid, patch.original, patch.saved_subsystem) end Base.@propagate_inbounds function Base.getproperty(x::AbstractODESolution, s::Symbol) @@ -154,12 +158,12 @@ end function ODESolution{T, N}( u, u_analytic, errors, t, k, discretes, prob, alg, interp, dense, tslocation, stats, alg_choice, retcode, resid = nothing, - original = nothing) where {T, N} + original = nothing, saved_subsystem = nothing) where {T, N} return ODESolution{T, N, typeof(u), typeof(u_analytic), typeof(errors), typeof(t), typeof(k), typeof(discretes), typeof(prob), typeof(alg), typeof(interp), - typeof(stats), typeof(alg_choice), typeof(resid), - typeof(original)}(u, u_analytic, errors, t, k, discretes, prob, alg, interp, - dense, tslocation, stats, alg_choice, retcode, resid, original) + typeof(stats), typeof(alg_choice), typeof(resid), typeof(original), + typeof(saved_subsystem)}(u, u_analytic, errors, t, k, discretes, prob, alg, interp, + dense, tslocation, stats, alg_choice, retcode, resid, original, saved_subsystem) end error_if_observed_derivative(_, _, ::Type{Val{0}}) = nothing @@ -180,6 +184,53 @@ function SymbolicIndexingInterface.is_parameter_timeseries(::Type{S}) where { Timeseries() end +const SolutionWithSavedSubsystem = ODESolution{T1, + T2, + T3, + T4, + T5, + T6, + T7, + T8, + T9, + T10, + T11, + T12, + T13, + T14, + T15, + T16} where { + T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16 <: SavedSubsystem} + +for method in [is_timeseries_parameter, timeseries_parameter_index, + with_updated_parameter_timeseries_values, get_saveable_values] + fname = nameof(method) + mod = parentmodule(method) + @eval function $(mod).$(fname)(sol::SolutionWithSavedSubsystem, args...) + $(method)(SavedSubsystemWithFallback(sol.saved_subsystem, symbolic_container(sol)), + args...) + end +end + +function SymbolicIndexingInterface.state_values(sol::SolutionWithSavedSubsystem, i) + original = state_values(sol.prob) + saved = sol.u[i] + if !(saved isa AbstractArray) + saved = [saved] + end + ss = sol.saved_subsystem + idxs = similar(saved, eltype(keys(ss.state_map))) + for (k, v) in ss.state_map + idxs[v] = k + end + replaced = remake_buffer(sol, original, idxs, saved) + return replaced +end + +function SymbolicIndexingInterface.state_values(sol::SolutionWithSavedSubsystem) + return map(Base.Fix1(state_values, sol), eachindex(sol.u)) +end + function _hold_discrete(disc_u, disc_t, t::Number) idx = searchsortedlast(disc_t, t) if idx == firstindex(disc_t) - 1 @@ -409,15 +460,25 @@ const PeriodicDiffEqArray = DiffEqArray{T, N, A, B} where {T, N, A, B <: Abstrac # public API, used by MTK """ get_saveable_values(sys, ps, timeseries_idx) + +Return the values to be saved in parameter object `ps` for timeseries index `timeseries_idx`. Called by +`save_discretes!`. If this returns `nothing`, `save_discretes!` will not save anything. """ function get_saveable_values(sys, ps, timeseries_idx) return get_saveable_values(symbolic_container(sys), ps, timeseries_idx) end +""" + save_discretes!(integ::DEIntegrator, timeseries_idx) + +Save the parameter timeseries with index `timeseries_idx`. Calls `get_saveable_values` to +get the values to save. If it returns `nothing`, then the save does not happen. +""" function save_discretes!(integ::DEIntegrator, timeseries_idx) - save_discretes!(integ.sol, current_time(integ), - get_saveable_values(integ, parameter_values(integ), timeseries_idx), - timeseries_idx) + inner_sol = get_sol(integ) + vals = get_saveable_values(inner_sol, parameter_values(integ), timeseries_idx) + vals === nothing && return + save_discretes!(integ.sol, current_time(integ), vals, timeseries_idx) end save_discretes!(args...) = nothing @@ -451,6 +512,7 @@ function build_solution(prob::Union{AbstractODEProblem, AbstractDDEProblem}, interp = LinearInterpolation(t, u), retcode = ReturnCode.Default, destats = missing, stats = nothing, resid = nothing, original = nothing, + saved_subsystem = nothing, kwargs...) T = eltype(eltype(u)) @@ -482,7 +544,12 @@ function build_solution(prob::Union{AbstractODEProblem, AbstractDDEProblem}, ps = parameter_values(prob) if has_sys(prob.f) - discretes = create_parameter_timeseries_collection(prob.f.sys, ps, prob.tspan) + sswf = if saved_subsystem === nothing + prob.f.sys + else + SavedSubsystemWithFallback(saved_subsystem, prob.f.sys) + end + discretes = create_parameter_timeseries_collection(sswf, ps, prob.tspan) else discretes = nothing end @@ -503,7 +570,8 @@ function build_solution(prob::Union{AbstractODEProblem, AbstractDDEProblem}, alg_choice, retcode, resid, - original) + original, + saved_subsystem) if calculate_error calculate_solution_errors!(sol; timeseries_errors = timeseries_errors, dense_errors = dense_errors) @@ -524,7 +592,8 @@ function build_solution(prob::Union{AbstractODEProblem, AbstractDDEProblem}, alg_choice, retcode, resid, - original) + original, + saved_subsystem) end end diff --git a/src/solutions/save_idxs.jl b/src/solutions/save_idxs.jl new file mode 100644 index 000000000..4481c9739 --- /dev/null +++ b/src/solutions/save_idxs.jl @@ -0,0 +1,303 @@ +#= +(Symbolic) save_idxs interface: + +Allows symbolically indexing solutions where a subset of the variables are saved. + +To implement this interface, the solution must store a `SavedSubsystem` if it contains a +subset of all timeseries variables. It must forward `is_timeseries_parameter`, +`timeseries_parameter_index`, `with_updated_parameter_timeseries_values` and +`get_saveable_values` to `SavedSubsystemWithFallback(sol.saved_subsystem, symbolic_container(sol))`. + +Additionally, it must implement `state_values` to always return the full state vector, using +the `SciMLProblem`'s `u0` as a reference, and updating the saved values in it. + +See the implementation for `ODESolution` as a reference. +=# + +struct VectorTemplate + type::DataType + size::Int +end + +struct TupleOfArraysWrapper{T} + x::T +end + +function TupleOfArraysWrapper(vt::Vector{VectorTemplate}) + return TupleOfArraysWrapper(Tuple(map(t -> Vector{t.type}(undef, t.size), vt))) +end + +function Base.getindex(t::TupleOfArraysWrapper, i::Tuple{Int, Int}) + t.x[i[1]][i[2]] +end + +function Base.setindex!(t::TupleOfArraysWrapper, val, i::Tuple{Int, Int}) + t.x[i[1]][i[2]] = val +end + +function as_diffeq_array(vt::Vector{VectorTemplate}, t) + return DiffEqArray(typeof(TupleOfArraysWrapper(vt))[], t, (1, 1)) +end + +""" + $(TYPEDSIGNATURES) + +A representation of the subsystem of a given system which is saved in a solution. Created +by providing an index provider and the indexes of saved variables in the system. The indexes +can also be symbolic variables. All indexes must refer to state variables, or timeseries +parameters. + +The arguments to the constructor are an index provider, the parameter object and the indexes +of variables to save. + +This object is stored in the solution object and used for symbolic indexing of the subsetted +solution. + +In case the provided `saved_idxs` is `nothing` or `isempty`, or if the provided +`saved_idxs` includes all of the variables and timeseries parameters, returns `nothing`. +""" +struct SavedSubsystem{V, T, I, P, Q, C} + """ + `Dict` mapping indexes of saved variables in the parent system to corresponding + indexes in the saved continuous timeseries. + """ + state_map::V + """ + `Dict` mapping indexes of saved timeseries parameters in the parent system to + corresponding `ParameterTimeseriesIndex`es in the save parameter timeseries. + """ + timeseries_params_map::T + """ + `Set` of all timeseries_idxs that are saved as-is. + """ + identity_partitions::I + """ + `Dict` mapping timeseries indexes to a vector of `VectorTemplate`s to use for storing + that subsetted timeseries partition + """ + timeseries_partition_templates::P + """ + `Dict` mapping timeseries indexes to a vector of `ParameterTimeseriesIndex` in that + partition. Only for saved partitions not in `identity_partitions`. + """ + indexes_in_partition::Q + """ + Map of timeseries indexes to the number of saved timeseries parameters in that + partition. + """ + partition_count::C +end + +function SavedSubsystem(indp, pobj, saved_idxs) + # nothing saved + if saved_idxs === nothing || isempty(saved_idxs) + return nothing + end + + # array state symbolics must be scalarized + saved_idxs = collect(Iterators.flatten(map(saved_idxs) do sym + if symbolic_type(sym) == NotSymbolic() + (sym,) + elseif sym isa AbstractArray && is_variable(indp, sym) + collect(sym) + else + (sym,) + end + end)) + + saved_state_idxs = Int[] + ts_idx_to_type_to_param_idx = Dict() + ts_idx_to_count = Dict() + num_ts_params = 0 + TParammapKeys = Union{} + for var in saved_idxs + if (idx = variable_index(indp, var)) !== nothing + push!(saved_state_idxs, idx) + elseif (idx = timeseries_parameter_index(indp, var)) !== nothing + TParammapKeys = Base.promote_typejoin(TParammapKeys, typeof(idx)) + # increment total number of ts params + num_ts_params += 1 + # get dict mapping type to idxs for this timeseries_idx + buf = get!(() -> Dict(), ts_idx_to_type_to_param_idx, idx.timeseries_idx) + # get type of parameter + val = parameter_values(pobj, parameter_index(indp, var)) + T = typeof(val) + # get vector of idxs for this type + buf = get!(() -> [], buf, T) + # push to it + push!(buf, idx) + # update count of variables in this partition + cnt = get(ts_idx_to_count, idx.timeseries_idx, 0) + ts_idx_to_count[idx.timeseries_idx] = cnt + 1 + else + throw(ArgumentError("Can only save variables and timeseries parameters. Got $var.")) + end + end + + # type of timeseries_idxs + Ttsidx = Union{} + for k in keys(ts_idx_to_type_to_param_idx) + Ttsidx = Base.promote_typejoin(Ttsidx, typeof(k)) + end + + # timeseries_idx to timeseries_parameter_index for all params + all_ts_params = Dict() + num_all_ts_params = 0 + for var in parameter_symbols(indp) + if (idx = timeseries_parameter_index(indp, var)) !== nothing + num_all_ts_params += 1 + buf = get!(() -> [], all_ts_params, idx.timeseries_idx) + push!(buf, idx) + end + end + + save_all_states = length(saved_state_idxs) == length(variable_symbols(indp)) + save_all_tsparams = num_ts_params == num_all_ts_params + if save_all_states && save_all_tsparams + # if we're saving everything + return nothing + end + if save_all_states + sort!(saved_state_idxs) + state_map = saved_state_idxs + else + state_map = Dict(saved_state_idxs .=> collect(eachindex(saved_state_idxs))) + end + + if save_all_tsparams + if isempty(ts_idx_to_type_to_param_idx) + identity_partitions = () + else + identity_partitions = Set{Ttsidx}(keys(ts_idx_to_type_to_param_idx)) + end + return SavedSubsystem( + state_map, nothing, identity_partitions, nothing, nothing, nothing) + end + + if num_ts_params == 0 + return SavedSubsystem(state_map, nothing, (), nothing, nothing, nothing) + end + + identitypartitions = Set{Ttsidx}() + parammap = Dict() + timeseries_partition_templates = Dict() + TsavedParamIdx = ParameterTimeseriesIndex{Ttsidx, NTuple{2, Int}} + indexes_in_partition = Dict{Ttsidx, Vector{TParammapKeys}}() + for (tsidx, type_to_idxs) in ts_idx_to_type_to_param_idx + if ts_idx_to_count[tsidx] == length(all_ts_params[tsidx]) + push!(identitypartitions, tsidx) + continue + end + templates = VectorTemplate[] + for (type, idxs) in type_to_idxs + template = VectorTemplate(type, length(idxs)) + push!(templates, template) + for (i, idx) in enumerate(idxs) + pti = ParameterTimeseriesIndex(tsidx, (length(templates), i)) + parammap[idx] = pti + + buf = get!(() -> TsavedParamIdx[], indexes_in_partition, tsidx) + push!(buf, idx) + end + end + timeseries_partition_templates[tsidx] = templates + end + parammap = Dict{TParammapKeys, TsavedParamIdx}(parammap) + timeseries_partition_templates = Dict{Ttsidx, Vector{VectorTemplate}}(timeseries_partition_templates) + ts_idx_to_count = Dict{Ttsidx, Int}(ts_idx_to_count) + return SavedSubsystem(state_map, parammap, identitypartitions, + timeseries_partition_templates, indexes_in_partition, ts_idx_to_count) +end + +""" + $(TYPEDEF) + +A combination of a `SavedSubsystem` and a fallback index provider. The provided fallback +is used as the `symbolic_container` for the `SavedSubsystemWithFallback`. Manually +implements `is_timeseries_parameter` and `timeseries_parameter_index` using the +`SavedSubsystem` to return the appropriate indexes for the subset of saved variables, +and `nothing`/`false` otherwise. + +Also implements `create_parameter_timeseries_collection`, `get_saveable_values` and +`with_updated_parameter_timeseries_values` to appropriately handled subsetted timeseries +parameters. +""" +struct SavedSubsystemWithFallback{S <: SavedSubsystem, T} + saved_subsystem::S + fallback::T +end + +function SymbolicIndexingInterface.symbolic_container(sswf::SavedSubsystemWithFallback) + sswf.fallback +end + +function SymbolicIndexingInterface.is_timeseries_parameter( + sswf::SavedSubsystemWithFallback, sym) + timeseries_parameter_index(sswf, sym) !== nothing +end + +function SymbolicIndexingInterface.timeseries_parameter_index( + sswf::SavedSubsystemWithFallback, sym) + ss = sswf.saved_subsystem + ss.timeseries_params_map === nothing && return nothing + if symbolic_type(sym) == NotSymbolic() + sym isa ParameterTimeseriesIndex || return nothing + sym.timeseries_idx in ss.identity_partitions && return sym + return get(ss.timeseries_params_map, sym, nothing) + end + v = timeseries_parameter_index(sswf.fallback, sym) + return timeseries_parameter_index(sswf, v) +end + +function create_parameter_timeseries_collection(sswf::SavedSubsystemWithFallback, ps, tspan) + original = create_parameter_timeseries_collection(sswf.fallback, ps, tspan) + ss = sswf.saved_subsystem + if original === nothing + return nothing + end + newcollection = map(enumerate(original)) do (i, buffer) + i in ss.identity_partitions && return buffer + ss.partition_count === nothing && return buffer + cnt = get(ss.partition_count, i, 0) + cnt == 0 && return buffer + + return as_diffeq_array(ss.timeseries_partition_templates[i], buffer.t) + end + + return ParameterTimeseriesCollection(newcollection, parameter_values(original)) +end + +function get_saveable_values(sswf::SavedSubsystemWithFallback, ps, tsidx) + ss = sswf.saved_subsystem + original = get_saveable_values(sswf.fallback, ps, tsidx) + tsidx in ss.identity_partitions && return original + ss.partition_count === nothing && return nothing + cnt = get(ss.partition_count, tsidx, 0) + cnt == 0 && return nothing + + toaw = TupleOfArraysWrapper(ss.timeseries_partition_templates[tsidx]) + for idx in ss.indexes_in_partition[tsidx] + toaw[ss.timeseries_params_map[idx].parameter_idx] = original[idx.parameter_idx] + end + return toaw +end + +function SymbolicIndexingInterface.with_updated_parameter_timeseries_values( + sswf::SavedSubsystemWithFallback, ps, args...) + ss = sswf.saved_subsystem + for (tsidx, val) in args + if tsidx in ss.identity_partitions + ps = with_updated_parameter_timeseries_values(sswf.fallback, ps, tsidx => val) + continue + end + ss.partition_count === nothing && continue + cnt = get(ss.partition_count, tsidx, 0) + cnt == 0 && continue + + for idx in ss.indexes_in_partition + set_parameter!(ps, parameter_values(val, ss.timeseries_params_map[idx]), idx) + end + end + + return ps +end diff --git a/test/downstream/solution_interface.jl b/test/downstream/solution_interface.jl index 709901c85..c7985596f 100644 --- a/test/downstream/solution_interface.jl +++ b/test/downstream/solution_interface.jl @@ -1,4 +1,5 @@ using ModelingToolkit, OrdinaryDiffEq, RecursiveArrayTools, StochasticDiffEq, Test +using SymbolicIndexingInterface using ModelingToolkit: t_nounits as t, D_nounits as D using Plots: Plots, plot @@ -170,3 +171,131 @@ sol10 = sol(0.1, idxs = 2) end end end + +@testset "Saved subsystem" begin + @testset "Pure ODE" begin + @variables x(t) y(t) + @parameters p + @mtkbuild sys = ODESystem([D(x) ~ x + p * y, D(y) ~ 2p + x^2], t) + @test length(unknowns(sys)) == 2 + xidx = variable_index(sys, x) + prob = ODEProblem(sys, [x => 1.0, y => 1.0], (0.0, 5.0), [p => 0.5]) + + @test SciMLBase.SavedSubsystem(sys, prob.p, []) === + SciMLBase.SavedSubsystem(sys, prob.p, nothing) === nothing + @test SciMLBase.SavedSubsystem(sys, prob.p, [x, y]) === nothing + @test begin + ss1 = SciMLBase.SavedSubsystem(sys, prob.p, [x]) + ss2 = SciMLBase.SavedSubsystem(sys, prob.p, [xidx]) + ss1.state_map == ss2.state_map + end + + sol = solve(prob, Tsit5(); save_idxs = xidx) + subsys = SciMLBase.SavedSubsystem(sys, prob.p, [xidx]) + xvals = sol[x] + # FIXME: hack for save_idxs + SciMLBase.@reset sol.saved_subsystem = subsys + @test is_variable(sol, x) + @test variable_index(sol, x) == 1 + @test !is_variable(sol, y) + @test variable_index(sol, y) === nothing + @test sol[x] == xvals + @test isequal(only(variable_symbols(sol)), x) + @test is_parameter(sol, p) + @test parameter_index(sol, p) == parameter_index(sys, p) + @test isequal(only(parameter_symbols(sol)), p) + @test is_independent_variable(sol, t) + + tmp = copy(prob.u0) + tmp[xidx] = xvals[2] + @test state_values(sol, 2) == tmp + @test state_values(sol) == [state_values(sol, i) for i in 1:length(sol)] + end + + @testset "ODE with callbacks" begin + @variables x(t) y(t) + @parameters p q(t) r(t) s(t) u(t) + evs = [0.1 => [q ~ q + 1, s ~ s - 1], 0.3 => [r ~ 2r, u ~ u / 2]] + @mtkbuild sys = ODESystem([D(x) ~ x + p * y, D(y) ~ 2p + x], t, [x, y], + [p, q, r, s, u], discrete_events = evs) + @test length(unknowns(sys)) == 2 + @test length(parameters(sys)) == 5 + @test is_timeseries_parameter(sys, q) + @test is_timeseries_parameter(sys, r) + xidx = variable_index(sys, x) + qidx = parameter_index(sys, q) + qpidx = timeseries_parameter_index(sys, q) + ridx = parameter_index(sys, r) + rpidx = timeseries_parameter_index(sys, r) + sidx = parameter_index(sys, s) + uidx = parameter_index(sys, u) + + @test SciMLBase.SavedSubsystem(sys, prob.p, [x, y, q, r, s, u]) === nothing + + prob = ODEProblem(sys, [x => 1.0, y => 1.0], (0.0, 5.0), + [p => 0.5, q => 0.0, r => 1.0, s => 10.0, u => 4096.0]) + sol = solve(prob; save_idxs = xidx) + xvals = sol[x] + subsys = SciMLBase.SavedSubsystem(sys, prob.p, [x, q, r]) + qvals = sol.ps[q] + rvals = sol.ps[r] + # FIXME: hack for save_idxs + SciMLBase.@reset sol.saved_subsystem = subsys + discq = DiffEqArray(qvals, sol.discretes[qpidx.timeseries_idx].t) + discr = DiffEqArray(rvals, sol.discretes[rpidx.timeseries_idx].t) + SciMLBase.@reset sol.discretes.collection[qpidx.timeseries_idx] = discq + SciMLBase.@reset sol.discretes.collection[rpidx.timeseries_idx] = discr + + @test is_variable(sol, x) + @test variable_index(sol, x) == 1 + @test !is_variable(sol, y) + @test sol[x] == xvals + @test variable_index(sol, y) === nothing + + @test all(Base.Fix1(is_parameter, sol), [p, q, r, s, u]) + @test all(Base.Fix1(is_timeseries_parameter, sol), [q, r]) + @test all(!Base.Fix1(is_timeseries_parameter, sol), [s, u]) + @test timeseries_parameter_index(sol, q) == + ParameterTimeseriesIndex(qpidx.timeseries_idx, 1) + @test timeseries_parameter_index(sol, r) == + ParameterTimeseriesIndex(rpidx.timeseries_idx, 1) + @test sol[q] == qvals + @test sol[r] == rvals + end + + @testset "SavedSubsystemWithFallback" begin + @variables x(t) y(t) + @parameters p q(t) r(t) s(t) u(t) + evs = [0.1 => [q ~ q + 1, s ~ s - 1], 0.3 => [r ~ 2r, u ~ u / 2]] + @mtkbuild sys = ODESystem([D(x) ~ x + p * y, D(y) ~ 2p + x^2], t, [x, y], + [p, q, r, s, u], discrete_events = evs) + prob = ODEProblem(sys, [x => 1.0, y => 1.0], (0.0, 5.0), + [p => 0.5, q => 0.0, r => 1.0, s => 10.0, u => 4096.0]) + ss = SciMLBase.SavedSubsystem(sys, prob.p, [x, q, s, r]) + sswf = SciMLBase.SavedSubsystemWithFallback(ss, sys) + xidx = variable_index(sys, x) + qidx = parameter_index(sys, q) + qpidx = timeseries_parameter_index(sys, q) + ridx = parameter_index(sys, r) + rpidx = timeseries_parameter_index(sys, r) + sidx = parameter_index(sys, s) + uidx = parameter_index(sys, u) + @test qpidx.timeseries_idx in ss.identity_partitions + @test !(rpidx.timeseries_idx in ss.identity_partitions) + @test timeseries_parameter_index(sswf, q) == timeseries_parameter_index(sys, q) + @test timeseries_parameter_index(sswf, s) == timeseries_parameter_index(sys, s) + + ptc = SciMLBase.create_parameter_timeseries_collection(sswf, prob.p, prob.tspan) + @test eltype(ptc[qpidx.timeseries_idx].u) <: SciMLBase.TupleOfArraysWrapper + @test eltype(ptc[rpidx.timeseries_idx].u) <: SciMLBase.TupleOfArraysWrapper + vals = SciMLBase.get_saveable_values(sswf, prob.p, qpidx.timeseries_idx) + @test vals isa SciMLBase.TupleOfArraysWrapper + @test vals.x isa Tuple{Vector{Float64}} + @test vals.x[1][1] == prob.ps[q] + + vals[(1, 1)] = 2prob.ps[q] + newp = with_updated_parameter_timeseries_values( + parameter_values(ptc), qpidx.timeseries_idx => vals) + @test newp[qidx] == 2prob.ps[q] + end +end