fix: use new trait in parameter timeseries indexing

Author: AayushSabharwal

src/parameter_indexing.jl

@@ -1,21 +1,3 @@
-parameter_values(arr::AbstractArray) = arr
-parameter_values(arr::Tuple) = arr
-parameter_values(arr::AbstractArray, i) = arr[i]
-parameter_values(arr::Tuple, i) = arr[i]
-parameter_values(prob, i) = parameter_values(parameter_values(prob), i)
-
-parameter_values_at_time(p, i) = parameter_values(p)
-
-parameter_values_at_state_time(p, i) = parameter_values(p)
-
-parameter_timeseries(_) = [0]
-
-# Tuple only included for the error message
-function set_parameter!(sys::Union{AbstractArray, Tuple}, val, idx)
-    sys[idx] = val
-end
-set_parameter!(sys, val, idx) = set_parameter!(parameter_values(sys), val, idx)
-
 """
     getp(indp, sym)
 
@@ -35,18 +17,19 @@ 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
-[`parameter_timeseries`](@ref), [`parameter_values_at_time`](@ref) and
-[`parameter_values_at_state_time`](@ref). The function returned from `getp` will 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`](@ref), [`parameter_values_at_time`](@ref),
+[`parameter_values_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.
 """
 function getp(sys, p)
     symtype = symbolic_type(p)
     elsymtype = symbolic_type(eltype(p))
     _getp(sys, symtype, elsymtype, p)
 end
 
-struct GetParameterIndex{I} <: AbstractGetIndexer
+struct GetParameterIndex{I} <: AbstractParameterGetIndexer
     idx::I
 end
 
@@ -59,10 +42,16 @@ function (gpi::GetParameterIndex)(::Timeseries, prob, i::Union{Int, CartesianInd
             prob, only(to_indices(parameter_timeseries(prob), (i,)))),
         gpi.idx)
 end
-function (gpi::GetParameterIndex)(::Timeseries, prob, i::Union{AbstractArray{Bool}, Colon})
+function (gpi::GetParameterIndex)(::Timeseries, prob, i::Colon)
     parameter_values.(
         parameter_values_at_time.((prob,),
-            (j for j in only(to_indices(parameter_timeseries(prob), (i,))))),
+            only(to_indices(parameter_timeseries(prob), (i,)))),
+        (gpi.idx,))
+end
+function (gpi::GetParameterIndex)(::Timeseries, prob, i::AbstractArray{Bool})
+    parameter_values.(
+        map(Base.Fix1(parameter_values_at_time, prob),
+            only(to_indices(parameter_timeseries(prob), (i,)))),
         (gpi.idx,))
 end
 function (gpi::GetParameterIndex)(::Timeseries, prob, i)
@@ -84,14 +73,19 @@ struct MultipleParameterGetters{G} <: AbstractGetIndexer
 end
 
 function (mpg::MultipleParameterGetters)(::IsTimeseriesTrait, prob)
-    map(g -> g(prob), mpg.getters)
+    map(CallWith(prob), mpg.getters)
 end
 function (mpg::MultipleParameterGetters)(::Timeseries, prob, i::Union{Int, CartesianIndex})
-    map(g -> g(prob, i), mpg.getters)
+    map(CallWith(prob, i), mpg.getters)
 end
 function (mpg::MultipleParameterGetters)(::Timeseries, prob, i)
-    [map(g -> g(prob, j), mpg.getters)
-     for j in only(to_indices(parameter_timeseries(prob), (i,)))]
+    map.(CallWith.((prob,), only(to_indices(parameter_timeseries(prob), (i,)))),
+        (mpg.getters,))
+end
+function (mpg::MultipleParameterGetters)(::Timeseries, prob, i::AbstractArray{Bool})
+    callers = map(
+        Base.Fix1(CallWith, prob), only(to_indices(parameter_timeseries(prob), (i,))))
+    map.(callers, (mpg.getters,))
 end
 function (mpg::MultipleParameterGetters)(buffer::AbstractArray, ::Timeseries, prob)
     for (g, bufi) in zip(mpg.getters, eachindex(buffer))
@@ -123,10 +117,10 @@ function (mpg::MultipleParameterGetters)(buffer::AbstractArray, ::NotTimeseries,
 end
 
 function (mpg::MultipleParameterGetters)(buffer::AbstractArray, prob, i...)
-    mpg(buffer, is_timeseries(prob), prob, i...)
+    mpg(buffer, is_parameter_timeseries(prob), prob, i...)
 end
 function (mpg::MultipleParameterGetters)(prob, i...)
-    mpg(is_timeseries(prob), prob, i...)
+    mpg(is_parameter_timeseries(prob), prob, i...)
 end
 
 for (t1, t2) in [

src/state_indexing.jl

@@ -32,7 +32,7 @@ function getu(sys, sym)
     _getu(sys, symtype, elsymtype, sym)
 end
 
-struct GetStateIndex{I} <: AbstractGetIndexer
+struct GetStateIndex{I} <: AbstractStateGetIndexer
     idx::I
 end
 function (gsi::GetStateIndex)(::Timeseries, prob)
@@ -49,13 +49,12 @@ function _getu(sys, ::NotSymbolic, ::NotSymbolic, sym)
     return GetStateIndex(sym)
 end
 
-struct GetpAtStateTime{G} <: AbstractGetIndexer
+struct GetpAtStateTime{G} <: AbstractStateGetIndexer
     getter::G
 end
 
 function (g::GetpAtStateTime)(::Timeseries, prob)
-    [g.getter(parameter_values_at_state_time(prob, i))
-     for i in eachindex(current_time(prob))]
+    g.getter.(parameter_values_at_state_time(prob))
 end
 function (g::GetpAtStateTime)(::Timeseries, prob, i)
     g.getter(parameter_values_at_state_time(prob, i))
@@ -64,20 +63,19 @@ function (g::GetpAtStateTime)(::NotTimeseries, prob)
     g.getter(prob)
 end
 
-struct GetIndepvar <: AbstractGetIndexer end
+struct GetIndepvar <: AbstractStateGetIndexer end
 
 (::GetIndepvar)(::IsTimeseriesTrait, prob) = current_time(prob)
 (::GetIndepvar)(::Timeseries, prob, i) = current_time(prob, i)
 
-struct TimeDependentObservedFunction{F} <: AbstractGetIndexer
+struct TimeDependentObservedFunction{F} <: AbstractStateGetIndexer
     obsfn::F
 end
 
 function (o::TimeDependentObservedFunction)(::Timeseries, prob)
-    curtime = current_time(prob)
-    return o.obsfn.(state_values(prob),
-        (parameter_values_at_state_time(prob, i) for i in eachindex(curtime)),
-        curtime)
+    o.obsfn.(state_values(prob),
+        parameter_values_at_state_time(prob),
+        current_time(prob))
 end
 function (o::TimeDependentObservedFunction)(::Timeseries, prob, i)
     return o.obsfn(state_values(prob, i),
@@ -88,7 +86,7 @@ function (o::TimeDependentObservedFunction)(::NotTimeseries, prob)
     return o.obsfn(state_values(prob), parameter_values(prob), current_time(prob))
 end
 
-struct TimeIndependentObservedFunction{F} <: AbstractGetIndexer
+struct TimeIndependentObservedFunction{F} <: AbstractStateGetIndexer
     obsfn::F
 end
 
@@ -115,7 +113,7 @@ function _getu(sys, ::ScalarSymbolic, ::SymbolicTypeTrait, sym)
     error("Invalid symbol $sym for `getu`")
 end
 
-struct MultipleGetters{G} <: AbstractGetIndexer
+struct MultipleGetters{G} <: AbstractStateGetIndexer
     getters::G
 end
 
@@ -130,7 +128,7 @@ function (mg::MultipleGetters)(::NotTimeseries, prob)
     return map(g -> g(prob), mg.getters)
 end
 
-struct AsTupleWrapper{G} <: AbstractGetIndexer
+struct AsTupleWrapper{G} <: AbstractStateGetIndexer
     getter::G
 end
 

src/value_provider_interface.jl

@@ -9,15 +9,23 @@
 Return an indexable collection containing the value of each parameter in `valp`. The two-
 argument version of this function returns the parameter value at index `i`. The
 two-argument version of this function will default to returning
-`parameter_values(valp)[i]`.
+`parameter_values(valp)[i]`. For a parameter timeseries object, this should return the
+parameter values at the initial time.
 
 If this function is called with an `AbstractArray` or `Tuple`, it will return the same
 array/tuple.
 """
 function parameter_values end
 
+parameter_values(arr::AbstractArray) = arr
+parameter_values(arr::Tuple) = arr
+parameter_values(arr::AbstractArray, i) = arr[i]
+parameter_values(arr::Tuple, i) = arr[i]
+parameter_values(prob, i) = parameter_values(parameter_values(prob), i)
+
 """
     parameter_values_at_time(valp, i)
+    parameter_values_at_time(valp)
 
 Return an indexable collection containing the value of all parameters in `valp` at time
 index `i`. This is useful when parameter values change during the simulation (such as
@@ -28,11 +36,21 @@ By default, this function returns `parameter_values(valp)` regardless of `i`, an
 to be specialized for timeseries objects where parameter values are not constant at all
 times. The resultant object should be indexable using [`parameter_values`](@ref).
 
-If this function is implemented, [`parameter_values_at_state_time`](@ref) must be 
-implemented for [`getu`](@ref) to work correctly.
+If this function is implemented for a type, [`parameter_values_at_state_time`](@ref) must
+be  implemented for [`getu`](@ref) to work correctly. Additionally,
+[`is_parameter_timeseries`](@ref) must be [`Timeseries`](@ref) for the type.
+
+The single-argument version of this function is a shorthand to return parameter values
+at each point in the parameter timeseries. This has a default implementation relying on
+[`parameter_timeseries`](@ref) and the two-argument version of this method.
 """
 function parameter_values_at_time end
 
+parameter_values_at_time(valp, i) = parameter_values(valp)
+function parameter_values_at_time(valp)
+    parameter_values_at_time.((valp,), eachindex(parameter_timeseries(valp)))
+end
+
 """
     parameter_values_at_state_time(valp, i)
 
@@ -48,20 +66,31 @@ all times. The resultant object should be indexable using [`parameter_values`](@
 
 If this function is implemented, [`parameter_values_at_time`](@ref) must be implemented for
 [`getp`](@ref) to work correctly.
+
+The single-argument version of this function is a shorthand to return parameter values
+at each point in the state timeseries. This has a default implementation relying on
+[`current_time`](@ref) and the two-argument version of this method.
 """
 function parameter_values_at_state_time end
 
+parameter_values_at_state_time(p, i) = parameter_values(p)
+function parameter_values_at_state_time(p)
+    parameter_values_at_state_time.((p,), eachindex(current_time(p)))
+end
+
 """
     parameter_timeseries(valp)
 
 Return an iterable of time steps at which the parameter values are saved. This is only
-required for objects where `is_timeseries(valp) === Timeseries()` and the parameter values
-change during the simulation (such as through callbacks). By default, this returns `[0]`.
+required for objects where `is_parameter_timeseries(valp) === Timeseries()`. By default,
+this returns `[0]`.
 
 See also: [`parameter_values_at_time`](@ref).
 """
 function parameter_timeseries end
 
+parameter_timeseries(_) = [0]
+
 """
     set_parameter!(valp, val, idx)
 
@@ -74,6 +103,12 @@ See: [`parameter_values`](@ref)
 """
 function set_parameter! end
 
+# Tuple only included for the error message
+function set_parameter!(sys::Union{AbstractArray, Tuple}, val, idx)
+    sys[idx] = val
+end
+set_parameter!(sys, val, idx) = set_parameter!(parameter_values(sys), val, idx)
+
 """
     finalize_parameters_hook!(valp, sym)
 
@@ -139,7 +174,21 @@ function current_time end
 abstract type AbstractIndexer end
 
 abstract type AbstractGetIndexer <: AbstractIndexer end
+abstract type AbstractStateGetIndexer <: AbstractGetIndexer end
+abstract type AbstractParameterGetIndexer <: AbstractGetIndexer end
 abstract type AbstractSetIndexer <: AbstractIndexer end
 
-(ai::AbstractGetIndexer)(prob) = ai(is_timeseries(prob), prob)
-(ai::AbstractGetIndexer)(prob, i) = ai(is_timeseries(prob), prob, i)
+(ai::AbstractStateGetIndexer)(prob) = ai(is_timeseries(prob), prob)
+(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)
+
+struct CallWith{A}
+    args::A
+
+    CallWith(args...) = new{typeof(args)}(args)
+end
+
+function (cw::CallWith)(arg)
+    arg(cw.args...)
+end

test/parameter_indexing_test.jl

@@ -124,6 +124,7 @@ function SymbolicIndexingInterface.parameter_values_at_state_time(fs::FakeSoluti
 end
 SymbolicIndexingInterface.parameter_timeseries(fs::FakeSolution) = fs.pt
 SymbolicIndexingInterface.is_timeseries(::Type{FakeSolution}) = Timeseries()
+SymbolicIndexingInterface.is_parameter_timeseries(::Type{FakeSolution}) = Timeseries()
 sys = SymbolCache([:x, :y, :z], [:a, :b, :c], :t)
 fs = FakeSolution(
     sys,