fixup! wip: better parameter indexing

Author: AayushSabharwal

src/index_provider_interface.jl

@@ -62,7 +62,7 @@ function is_timeseries_parameter(indp, sym)
     if hasmethod(symbolic_container, Tuple{typeof(indp)})
         is_timeseries_parameter(symbolic_container(indp), sym)
     else
-        false
+        return false
     end
 end
 
@@ -94,7 +94,7 @@ function timeseries_parameter_index(indp, sym)
     if hasmethod(symbolic_container, Tuple{typeof(indp)})
         timeseries_parameter_index(symbolic_container(indp), sym)
     else
-        nothing
+        return nothing
     end
 end
 

src/parameter_indexing.jl

@@ -15,13 +15,24 @@ Requires that the value provider implement [`parameter_values`](@ref). This func
 may not always need to be implemented, and has a default implementation for collections
 that implement `getindex`.
 
-If the returned function is used on a timeseries object which saves parameter timeseries, it
-can be used to index said timeseries. The timeseries object must implement
+If the returned function is used on a timeseries object which saves parameter timeseries,
+it can be used to index said timeseries. The timeseries object must implement
 [`parameter_timeseries`](@ref), [`parameter_values_at_state_time`](@ref),
-[`parameter_timeseries_at_state_time`](@ref) and [`is_parameter_timeseries`](@ref). The
-function returned from `getp` can be passed `Colon()` (`:`) as the last argument to return
-the entire parameter timeseries for `p`, or any index into the parameter timeseries for a
-subset of values.
+[`parameter_timeseries_at_state_time`](@ref) and [`is_parameter_timeseries`](@ref).
+
+If `sym` is a timeseries parameter, the function will return the timeseries of the
+parameter if the value provider is a parameter timeseries object. An additional argument
+can be provided to the function indicating the specific indexes in the timeseries at
+which to access the values. If `sym` is an array of parameters, the following cases
+apply:
+
+- All parameters are non-timeseries parameters: The function returns the value of each
+  parameter.
+- All parameters are timeseries parameters: All the parameters must belong to the same
+  timeseries (otherwise `getp` will error). The function returns the timeseries of all
+  parameter values, and can be accessed at specific indices in the timeseries.
+- A mix of timeseries and non-timeseries parameters: The function can _only_ be used on
+  non-timeseries objects and will return the value of each parameter at in the object.
 """
 function getp(sys, p)
     symtype = symbolic_type(p)
@@ -51,7 +62,7 @@ function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(buffer::AbstractAr
     for (buf_idx, ts_idx) in zip(eachindex(buffer), eachindex(parameter_timeseries(prob, indexer_timeseries_index(gpi))))
         buffer[buf_idx] = gpi(ts, prob, ts_idx)
     end
-    buffer
+    return buffer
 end
 function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(::Timeseries, prob, i::Union{Int, CartesianIndex})
     parameter_values(prob, gpi.idx, i)
@@ -71,7 +82,7 @@ function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(buffer::AbstractAr
     for (buf_idx, ts_idx) in zip(eachindex(buffer), only(to_indices(parameter_timeseries(prob, indexer_timeseries_index(gpi)), (i,))))
         buffer[buf_idx] = gpi(ts, prob, ts_idx)
     end
-    buffer
+    return buffer
 end
 function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(ts::Timeseries, prob, i)
     gpi.((ts,), (prob,), i)
@@ -80,7 +91,7 @@ function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(buffer::AbstractAr
     for (buf_idx, subidx) in zip(eachindex(buffer), i)
         buffer[buf_idx] = gpi(ts, prob, subidx)
     end
-    buffer
+    return buffer
 end
 function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(::NotTimeseries, prob)
     throw(ParameterTimeseriesValueIndexMismatchError{NotTimeseries}(prob, gpi))
@@ -190,7 +201,7 @@ for (indexerTimeseriesType, timeseriesType) in [
         for (buf_idx, getter) in zip(eachindex(buffer), mpg.getters)
             buffer[buf_idx] = getter(prob)
         end
-        buffer
+        return buffer
     end
 end
 
@@ -227,13 +238,13 @@ function (mpg::AtLeastTimeseriesMPG)(buffer::AbstractArray, ts::Timeseries, prob
     for (buf_idx, ts_idx) in zip(eachindex(buffer), eachindex(parameter_timeseries(prob, indexer_timeseries_index(mpg))))
         mpg(buffer[buf_idx], ts, prob, ts_idx)
     end
-    buffer
+    return buffer
 end
 function (mpg::AtLeastTimeseriesMPG)(buffer::AbstractArray, ::Timeseries, prob, i::Union{Int, CartesianIndex})
     for (buf_idx, getter) in zip(eachindex(buffer), mpg.getters)
         buffer[buf_idx] = getter(prob, i)
     end
-    buffer
+    return buffer
 end
 function (mpg::AtLeastTimeseriesMPG)(buffer::AbstractArray, ts::Timeseries, prob, ::Colon)
     mpg(buffer, ts, prob)
@@ -245,7 +256,7 @@ function (mpg::AtLeastTimeseriesMPG)(buffer::AbstractArray, ts::Timeseries, prob
     for (buf_idx, ts_idx) in zip(eachindex(buffer), i)
         mpg(buffer[buf_idx], ts, prob, ts_idx)
     end
-    buffer
+    return buffer
 end
 function (mpg::MultipleParametersGetter{IndexerTimeseries})(::NotTimeseries, prob)
     throw(ParameterTimeseriesValueIndexMismatchError{NotTimeseries}(prob, mpg))
@@ -285,7 +296,7 @@ end
 function (phw::ParameterHookWrapper)(prob, args...)
     res = phw.setter(prob, args...)
     finalize_parameters_hook!(prob, phw.original_index)
-    res
+    return res
 end
 
 """

src/parameter_timeseries_collection.jl

@@ -13,9 +13,9 @@ The collection is expected to implement `Base.eachindex`, `Base.iterate` and
 returned by calling [`timeseries_parameter_index`](@ref) on the corresponding index
 provider.
 
-This type forwards `eachindex` and `iterate` to the contained `collection`. It implements
-`Base.parent` to allow access to the contained `collection`, and has the following
-`getindex` methods:
+This type forwards `eachindex`, `iterate` and `length` to the contained `collection`. It
+implements `Base.parent` to allow access to the contained `collection`, and has the
+following `getindex` methods:
 
 - `getindex(ptc::ParameterTimeseriesCollection, idx) = ptc.collection[idx]`.
 - `getindex(::ParameterTimeseriesCollection, idx::ParameterTimeseriesIndex)` returns the
@@ -48,6 +48,8 @@ Base.eachindex(ptc::ParameterTimeseriesCollection) = eachindex(ptc.collection)
 
 Base.iterate(ptc::ParameterTimeseriesCollection, args...) = iterate(ptc.collection, args...)
 
+Base.length(ptc::ParameterTimeseriesCollection) = length(ptc.collection)
+
 Base.parent(ptc::ParameterTimeseriesCollection) = ptc.collection
 
 Base.getindex(ptc::ParameterTimeseriesCollection, idx) = ptc.collection[idx]

src/state_indexing.jl

@@ -21,6 +21,11 @@ support symbolic expressions, the value provider must implement [`observed`](@re
 
 This function typically does not need to be implemented, and has a default implementation
 relying on the above functions.
+
+If the value provider is a parameter timeseries object, the same rules apply as
+[`getp`](@ref). The difference here is that `sym` may also contain non-parameter symbols,
+and the values are always returned corresponding to the state timeseries. This utilizes
+[`parameter_values_at_state_time`](@ref) and [`parameter_timeseries_at_state_time`](@ref).
 """
 function getu(sys, sym)
     symtype = symbolic_type(sym)
@@ -102,16 +107,41 @@ struct TimeDependentObservedFunction{F} <: AbstractStateGetIndexer
     obsfn::F
 end
 
-function (o::TimeDependentObservedFunction)(::Timeseries, prob)
+function (o::TimeDependentObservedFunction)(ts::Timeseries, prob)
+    return o(ts, is_parameter_timeseries(prob), prob)
+end
+function (o::TimeDependentObservedFunction)(::Timeseries, ::Timeseries, prob)
     o.obsfn.(state_values(prob),
         parameter_values_at_state_time(prob),
         current_time(prob))
 end
-function (o::TimeDependentObservedFunction)(::Timeseries, prob, i)
+function (o::TimeDependentObservedFunction)(::Timeseries, ::NotTimeseries, prob)
+    o.obsfn.(state_values(prob),
+        (parameter_values(prob),),
+        current_time(prob))
+end
+function (o::TimeDependentObservedFunction)(ts::Timeseries, prob, i)
+    return o(ts, is_parameter_timeseries(prob), prob, i)
+end
+function (o::TimeDependentObservedFunction)(::Timeseries, ::Timeseries, prob, i::Union{Int, CartesianIndex})
     return o.obsfn(state_values(prob, i),
         parameter_values_at_state_time(prob, i),
         current_time(prob, i))
 end
+function (o::TimeDependentObservedFunction)(ts::Timeseries, p_ts::IsTimeseriesTrait, prob, ::Colon)
+    return o(ts, p_ts, prob)
+end
+function (o::TimeDependentObservedFunction)(ts::Timeseries, p_ts::IsTimeseriesTrait, prob, i::AbstractArray{Bool})
+    map(only(to_indices(current_time(prob), (i,)))) do idx
+        o(ts, p_ts, prob, idx)
+    end
+end
+function (o::TimeDependentObservedFunction)(ts::Timeseries, p_ts::IsTimeseriesTrait, prob, i)
+    o.((ts,), (p_ts,), (prob,), i)
+end
+function (o::TimeDependentObservedFunction)(::Timeseries, ::NotTimeseries, prob, i::Union{Int, CartesianIndex})
+    o.obsfn(state_values(prob, i), parameter_values(prob), current_time(prob, i))
+end
 function (o::TimeDependentObservedFunction)(::NotTimeseries, prob)
     return o.obsfn(state_values(prob), parameter_values(prob), current_time(prob))
 end
@@ -168,21 +198,29 @@ function (mg::MultipleGetters)(::NotTimeseries, prob)
     return map(g -> g(prob), mg.getters)
 end
 
-struct AsTupleWrapper{G} <: AbstractStateGetIndexer
+struct AsTupleWrapper{N, G} <: AbstractStateGetIndexer
     getter::G
 end
 
-function (atw::AsTupleWrapper)(::Timeseries, prob)
-    return Tuple.(atw.getter(prob))
+AsTupleWrapper{N}(getter::G) where {N, G} = AsTupleWrapper{N, G}(getter)
+
+wrap_tuple(::AsTupleWrapper{N}, val) where {N} = ntuple(i -> val[i], Val(N))
+
+function (atw::(AsTupleWrapper{N} where {N}))(::Timeseries, prob)
+    return wrap_tuple.((atw,), atw.getter(prob))
+    # return Tuple.(atw.getter(prob))
 end
 function (atw::AsTupleWrapper)(::Timeseries, prob, i::Union{Int, CartesianIndex})
-    return Tuple(atw.getter(prob, i))
+    return wrap_tuple(atw, atw.getter(prob, i))
+    # return Tuple(atw.getter(prob, i))
 end
 function (atw::AsTupleWrapper)(::Timeseries, prob, i)
-    return Tuple.(atw.getter(prob, i))
+    return wrap_tuple.((atw,), atw.getter(prob, i))
+    # return Tuple.(atw.getter(prob, i))
 end
 function (atw::AsTupleWrapper)(::NotTimeseries, prob)
-    return Tuple(atw.getter(prob))
+    wrap_tuple(atw, atw.getter(prob))
+    # return Tuple(atw.getter(prob))
 end
 
 for (t1, t2) in [
@@ -207,7 +245,7 @@ for (t1, t2) in [
                 TimeIndependentObservedFunction(obs)
             end
             if sym isa Tuple
-                getter = AsTupleWrapper(getter)
+                getter = AsTupleWrapper{length(sym)}(getter)
             end
             return getter
         end

src/symbol_cache.jl

@@ -46,6 +46,9 @@ function SymbolCache(vars = nothing, params = nothing, indepvars = nothing;
     params = to_dict_or_nothing(params)
     timeseries_parameters = to_dict_or_nothing(timeseries_parameters)
     if timeseries_parameters !== nothing
+        if indepvars === nothing
+            throw(ArgumentError("Independent variable is required for timeseries parameters to exist"))
+        end
         for (k, v) in timeseries_parameters
             if !haskey(params, k)
                 throw(ArgumentError("Timeseries parameter $k must also be present in parameters."))
@@ -75,7 +78,7 @@ function variable_symbols(sc::SymbolCache, i = nothing)
     for (k, v) in sc.variables
         buffer[v] = k
     end
-    buffer
+    return buffer
 end
 function is_parameter(sc::SymbolCache, sym)
     sc.parameters !== nothing && haskey(sc.parameters, sym)

src/value_provider_interface.jl

@@ -215,15 +215,8 @@ abstract type AbstractSetIndexer <: AbstractIndexer end
 (ai::AbstractStateGetIndexer)(prob, i) = ai(is_timeseries(prob), prob, i)
 (ai::AbstractParameterGetIndexer)(prob) = ai(is_parameter_timeseries(prob), prob)
 (ai::AbstractParameterGetIndexer)(prob, i) = ai(is_parameter_timeseries(prob), prob, i)
-# unfortunately, this is ambiguous
-function (ai::AbstractParameterGetIndexer)(arg1::AbstractArray, arg2)
-    if hasmethod(parameter_values, Tuple{typeof(arg2)})
-        # arg1 is buffer
-        ai(arg1, is_parameter_timeseries(arg2), arg2)
-    else
-        # arg1 is value provider
-        ai(is_parameter_timeseries(arg1), arg1, arg2)
-    end
+function (ai::AbstractParameterGetIndexer)(buffer::AbstractArray, prob)
+    ai(buffer, is_parameter_timeseries(prob), prob)
 end
 function (ai::AbstractParameterGetIndexer)(buffer::AbstractArray, prob, i)
     ai(buffer, is_parameter_timeseries(prob), prob, i)

test/example_test.jl

@@ -71,6 +71,8 @@ sys = SystemMockup(true, [:x, :y, :z], [:a, :b, :c], :t)
 @test all(.!is_parameter.((sys,), [:x, :y, :z, :t, :p, :q, :r]))
 @test all(parameter_index.((sys,), [:c, :a, :b]) .== [3, 1, 2])
 @test all(parameter_index.((sys,), [:x, :y, :z, :t, :p, :q, :r]) .=== nothing)
+@test all(.!is_timeseries_parameter.((sys,), [:x, :y, :z, :t, :p, :q, :r])) # fallback even if not implemented
+@test all(timeseries_parameter_index.((sys,), [:x, :y, :z, :t, :p, :q, :r]) .=== nothing) # fallback
 @test is_independent_variable(sys, :t)
 @test all(.!is_independent_variable.((sys,), [:x, :y, :z, :a, :b, :c, :p, :q, :r]))
 @test all(is_observed.((sys,), [:x, :y, :z, :a, :b, :c, :t]))
@@ -88,6 +90,7 @@ sys = SystemMockup(true, [:x, :y, :z], [:a, :b, :c], :t)
 @test independent_variable_symbols(sys) == [:t]
 @test all_variable_symbols(sys) == [:x, :y, :z]
 @test sort(all_symbols(sys)) == [:a, :b, :c, :t, :x, :y, :z]
+@test default_values(sys) == Dict() # fallback even if not implemented
 
 sys = SystemMockup(true, [:x, :y, :z], [:a, :b, :c], nothing)
 

test/parameter_indexing_test.jl

@@ -5,6 +5,12 @@ using SymbolicIndexingInterface: IndexerTimeseries, IndexerNotTimeseries, Indexe
                                  MixedParameterTimeseriesIndexError
 using Test
 
+arr = [1.0, 2.0, 3.0]
+@test parameter_values(arr) == arr
+@test current_time(arr) == arr
+tp = (1.0, 2.0, 3.0)
+@test parameter_values(tp) == tp
+
 struct FakeIntegrator{S, P}
     sys::S
     p::P
@@ -29,6 +35,9 @@ for sys in [
         p = [1.0, 2.0, 3.0, 4.0]
         fi = FakeIntegrator(sys, pType(copy(p)), Ref(0))
         new_p = [4.0, 5.0, 6.0, 7.0]
+        for i in [7, CartesianIndex(5)]
+            @test parameter_values_at_state_time(fi, i) == parameter_values(fi)
+        end
         for (sym, oldval, newval, check_inference) in [
             (:a, p[1], new_p[1], true),
             (1, p[1], new_p[1], true),
@@ -308,6 +317,9 @@ for (sym, val_is_timeseries, val, check_inference) in [
         ((:x, :b, :c), true, tuple.(xval, bval_state, cval_state), true),
         ([:a, :b, :x], true, vcat.(aval, bval_state, xval), false),
         ((:a, :b, :x), true, tuple.(aval, bval_state, xval), true),
+        (:(2b), true, 2 .* bval_state, true),
+        ([:x, :(2b), :(3c)], true, vcat.(xval, 2 .* bval_state, 3 .* cval_state), true),
+        ((:x, :(2b), :(3c)), true, tuple.(xval, 2 .* bval_state, 3 .* cval_state), true),
     ]
     getter = getu(sys, sym)
     if val isa DataType

test/parameter_timeseries_collection_test.jl

@@ -0,0 +1,47 @@
+using SymbolicIndexingInterface
+using Test
+
+struct MyDiffEqArray
+    t::Vector{Float64}
+    u::Vector{Vector{Float64}}
+end
+SymbolicIndexingInterface.current_time(mda::MyDiffEqArray) = mda.t
+SymbolicIndexingInterface.state_values(mda::MyDiffEqArray) = mda.u
+SymbolicIndexingInterface.is_timeseries(::Type{MyDiffEqArray}) = Timeseries()
+
+@test_throws ArgumentError ParameterTimeseriesCollection((ones(3), 2ones(3)))
+
+a_timeseries = MyDiffEqArray(collect(0:0.1:0.9), [[2.5i, sin(0.2i)] for i in 1:10])
+b_timeseries = MyDiffEqArray(collect(0:0.25:0.9), [[3.5i, log(1.3i)] for i in 1:4])
+c_timeseries = MyDiffEqArray(collect(0:0.17:0.90), [[4.3i] for i in 1:5])
+collection = (a_timeseries, b_timeseries, c_timeseries)
+ptc = ParameterTimeseriesCollection(collection)
+
+@test collect(eachindex(ptc)) == [1, 2, 3]
+@test [x for x in ptc] == [a_timeseries, b_timeseries, c_timeseries]
+@test length(ptc) == 3
+@test parent(ptc) === collection
+
+for i in 1:3
+    @test ptc[i] === collection[i]
+    @test parameter_timeseries(ptc, i) == collection[i].t
+    for j in eachindex(collection[i].u[1])
+        pti = ParameterTimeseriesIndex(i, j)
+        @test ptc[pti] == getindex.(collection[i].u, j)
+        for k in eachindex(collection[i].u)
+            rhs = collection[i].u[k][j]
+            @test ptc[pti, CartesianIndex(k)] == rhs
+            @test ptc[pti, k] == rhs
+            @test ptc[i, k] == collection[i].u[k]
+            @test ptc[i, k, j] == rhs
+            @test parameter_values(ptc, pti, k) == rhs
+        end
+        allidxs = eachindex(collection[i].u)
+        for subidx in [:, rand(allidxs, 3), rand(Bool, length(allidxs))]
+            rhs = getindex.(collection[i].u[subidx], j)
+            @test ptc[pti, subidx] == rhs
+            @test ptc[i, subidx, j] == rhs
+            @test parameter_values(ptc, pti, subidx) == rhs
+        end
+    end
+end

test/runtests.jl

@@ -27,6 +27,9 @@ if GROUP == "All" || GROUP == "Core"
     @safetestset "Fallback test" begin
         @time include("fallback_test.jl")
     end
+    @safetestset "ParameterTimeseriesCollection test" begin
+        @time include("parameter_timeseries_collection_test.jl")
+    end
     @safetestset "Parameter indexing test" begin
         @time include("parameter_indexing_test.jl")
     end