fixup! wip: better parameter indexing

Author: AayushSabharwal

src/parameter_indexing.jl

@@ -42,11 +42,7 @@ function (gpi::GetParameterIndex)(::IsTimeseriesTrait, prob)
     parameter_values(prob, gpi.idx)
 end
 function (gpi::GetParameterIndex)(::Timeseries, prob, args)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is a parameter timeseries object) with parameter index $(gpi.idx) (which is not\
-        a `ParameterTimeseriesIndex`) at index $(args) in the timeseries.
-    """)
+    throw(ParameterTimeseriesValueIndexMismatchError{Timeseries}(prob, gpi, args))
 end
 function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(::Timeseries, prob)
     parameter_values.(
@@ -68,10 +64,7 @@ function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(ts::Timeseries, pr
     gpi.((ts,), (prob,), i)
 end
 function (gpi::GetParameterIndex{<:ParameterTimeseriesIndex})(::NotTimeseries, prob)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is not a parameter timeseries object) using index $(gpi.idx).
-          """)
+    throw(ParameterTimeseriesValueIndexMismatchError{NotTimeseries}(prob, gpi))
 end
 
 function _getp(sys, ::NotSymbolic, ::NotSymbolic, p)
@@ -177,26 +170,14 @@ for (indexerTimeseriesType, timeseriesType) in [
 end
 
 function (mpg::MixedTimeseriesIndexMPG)(::Timeseries, prob, args...)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is a parameter timeseries object) with variables having mixed timeseries indexes\
-        $(mpg.timeseries_idx.indexes).
-    """)
+    throw(MixedParameterTimeseriesIndexError(prob, mpg.timeseries_idx.indexes))
 end
 
 function (mpg::MultipleParametersGetter{IndexerNotTimeseries})(::Timeseries, prob, args)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is a parameter timeseries object) with non-timeseries indexer $mpg at index $args\
-        in the timeseries.
-    """)
+    throw(ParameterTimeseriesValueIndexMismatchError{Timeseries}(prob, mpg, args))
 end
 function (mpg::MultipleParametersGetter{IndexerNotTimeseries})(::AbstractArray, ::Timeseries, prob, args)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is a parameter timeseries object) with non-timeseries indexer $mpg at index $args\
-        in the timeseries.
-    """)
+    throw(ParameterTimeseriesValueIndexMismatchError{Timeseries}(prob, mpg, args))
 end
 function (mpg::AtLeastTimeseriesMPG)(ts::Timeseries, prob)
     map(eachindex(parameter_timeseries(prob, indexer_timeseries_index(mpg)))) do i
@@ -242,16 +223,10 @@ function (mpg::AtLeastTimeseriesMPG)(buffer::AbstractArray, ts::Timeseries, prob
     buffer
 end
 function (mpg::MultipleParametersGetter{IndexerTimeseries})(::NotTimeseries, prob)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is not a parameter timeseries object) with timeseries indexer $mpg.
-          """)
+    throw(ParameterTimeseriesValueIndexMismatchError{NotTimeseries}(prob, mpg))
 end
 function (mpg::MultipleParametersGetter{IndexerTimeseries})(::AbstractArray, ::NotTimeseries, prob)
-    error("""
-        Invalid indexing operation: tried to access object of type $(typeof(prob)) (which\
-        is not a parameter timeseries object) with timeseries indexer $mpg.
-          """)
+    throw(ParameterTimeseriesValueIndexMismatchError{NotTimeseries}(prob, mpg))
 end
 
 for (t1, t2) in [

src/state_indexing.jl

@@ -1,11 +1,7 @@
-state_values(arr::AbstractArray) = arr
-state_values(arr, i) = state_values(arr)[i]
-
 function set_state!(sys, val, idx)
     state_values(sys)[idx] = val
 end
 
-current_time(p, i) = current_time(p)[i]
 
 """
     getu(indp, sym)
@@ -38,9 +34,12 @@ end
 function (gsi::GetStateIndex)(::Timeseries, prob)
     getindex.(state_values(prob), (gsi.idx,))
 end
-function (gsi::GetStateIndex)(::Timeseries, prob, i)
+function (gsi::GetStateIndex)(::Timeseries, prob, i::Union{Int, CartesianIndex})
     getindex(state_values(prob, i), gsi.idx)
 end
+function (gsi::GetStateIndex)(::Timeseries, prob, i)
+    getindex.(state_values(prob, i), gsi.idx)
+end
 function (gsi::GetStateIndex)(::NotTimeseries, prob)
     state_values(prob, gsi.idx)
 end
@@ -62,12 +61,34 @@ end
 function (g::GetpAtStateTime)(::Timeseries, ::NotTimeseries, prob, _...)
     g.getter(prob)
 end
-function (g::GetpAtStateTime, ::Timeseries, ::Timeseries, prob)
+function (g::GetpAtStateTime)(ts::Timeseries, p_ts::Timeseries, prob)
+    g(ts, p_ts, is_indexer_timeseries(g.getter), prob)
+end
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::Union{IndexerTimeseries, IndexerBoth}, prob)
     g.getter.((prob,), parameter_timeseries_at_state_time(prob, indexer_timeseries_index(g.getter)))
 end
-function (g::GetpAtStateTime, ::Timeseries, ::Timeseries, prob, i)
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::IndexerNotTimeseries, prob)
+    g.getter(prob)
+end
+function (g::GetpAtStateTime)(ts::Timeseries, p_ts::Timeseries, prob, i)
+    g(ts, p_ts, is_indexer_timeseries(g.getter), prob, i)
+end
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::Union{IndexerTimeseries, IndexerBoth}, prob, i)
     g.getter(prob, parameter_timeseries_at_state_time(prob, indexer_timeseries_index(g.getter), i))
 end
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::IndexerNotTimeseries, prob, ::Union{Int, CartesianIndex})
+    g.getter(prob)
+end
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::IndexerNotTimeseries, prob, ::Colon)
+    map(_ -> g.getter(prob), current_time(prob))
+end
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::IndexerNotTimeseries, prob, i::AbstractArray{Bool})
+    num_ones = sum(i)
+    map(_ -> g.getter(prob), 1:num_ones)
+end
+function (g::GetpAtStateTime)(::Timeseries, ::Timeseries, ::IndexerNotTimeseries, prob, i)
+    map(_ -> g.getter(prob), 1:length(i))
+end
 function (g::GetpAtStateTime)(::NotTimeseries, prob)
     g.getter(prob)
 end
@@ -126,12 +147,22 @@ struct MultipleGetters{G} <: AbstractStateGetIndexer
     getters::G
 end
 
-function (mg::MultipleGetters)(::Timeseries, prob)
-    return broadcast(i -> map(g -> g(prob, i), mg.getters),
-        eachindex(state_values(prob)))
+function (mg::MultipleGetters)(ts::Timeseries, prob)
+    return mg.((ts,), (prob,), eachindex(current_time(prob)))
+end
+function (mg::MultipleGetters)(::Timeseries, prob, i::Union{Int, CartesianIndex})
+    return map(CallWith(prob, i), mg.getters)
+end
+function (mg::MultipleGetters)(ts::Timeseries, prob, ::Colon)
+    return mg(ts, prob)
 end
-function (mg::MultipleGetters)(::Timeseries, prob, i)
-    return map(g -> g(prob, i), mg.getters)
+function (mg::MultipleGetters)(ts::Timeseries, prob, i::AbstractArray{Bool})
+    return map(only(to_indices(current_time(prob), (i,)))) do idx
+        mg(ts, prob, idx)
+    end
+end
+function (mg::MultipleGetters)(ts::Timeseries, prob, i)
+    mg.((ts,), (prob,), i)
 end
 function (mg::MultipleGetters)(::NotTimeseries, prob)
     return map(g -> g(prob), mg.getters)
@@ -144,9 +175,12 @@ end
 function (atw::AsTupleWrapper)(::Timeseries, prob)
     return Tuple.(atw.getter(prob))
 end
-function (atw::AsTupleWrapper)(::Timeseries, prob, i)
+function (atw::AsTupleWrapper)(::Timeseries, prob, i::Union{Int, CartesianIndex})
     return Tuple(atw.getter(prob, i))
 end
+function (atw::AsTupleWrapper)(::Timeseries, prob, i)
+    return Tuple.(atw.getter(prob, i))
+end
 function (atw::AsTupleWrapper)(::NotTimeseries, prob)
     return Tuple(atw.getter(prob))
 end
@@ -158,9 +192,13 @@ for (t1, t2) in [
 ]
     @eval function _getu(sys, ::NotSymbolic, ::$t1, sym::$t2)
         num_observed = count(x -> is_observed(sys, x), sym)
-        if num_observed <= 1
-            getters = getu.((sys,), sym)
-            return MultipleGetters(getters)
+        if num_observed == 0
+            if all(Base.Fix1(is_parameter, sys), sym) && all(!Base.Fix1(is_timeseries_parameter, sys), sym)
+                GetpAtStateTime(getp(sys, sym))
+            else
+                getters = getu.((sys,), sym)
+                return MultipleGetters(getters)
+            end
         else
             obs = observed(sys, sym isa Tuple ? collect(sym) : sym)
             getter = if is_time_dependent(sys)
@@ -181,7 +219,7 @@ function _getu(sys, ::ArraySymbolic, ::SymbolicTypeTrait, sym)
         idx = variable_index(sys, sym)
         return getu(sys, idx)
     elseif is_parameter(sys, sym)
-        return getp(sys, sym)
+        return GetpAtStateTime(getp(sys, sym))
     end
     return getu(sys, collect(sym))
 end

src/symbol_cache.jl

@@ -169,10 +169,20 @@ function observed(sc::SymbolCache, expr::Expr)
         end
     end
 end
-function observed(sc::SymbolCache, exprs::AbstractArray{Expr})
+function observed(sc::SymbolCache, exprs::AbstractArray)
+    for expr in exprs
+        if !(expr isa Union{Symbol, Expr})
+            throw(TypeError(:observed, "SymbolCache", Union{Symbol, Expr}, expr))
+        end
+    end
     return observed(sc, :(reshape([$(exprs...)], $(size(exprs)))))
 end
-function observed(sc::SymbolCache, exprs::Tuple{Vararg{Expr}})
+function observed(sc::SymbolCache, exprs::Tuple)
+    for expr in exprs
+        if !(expr isa Union{Symbol, Expr})
+            throw(TypeError(:observed, "SymbolCache", Union{Symbol, Expr}, expr))
+        end
+    end
     return observed(sc, :(($(exprs...),)))
 end
 

src/value_provider_interface.jl

@@ -73,21 +73,37 @@ function parameter_timeseries end
 Return the index of the timestep in the parameter timeseries at timeseries index `i` which
 occurs just before or at the same time as the state timestep with index `j`. The two-
 argument version of this function returns an iterable of indexes, one for each timestep in
-the state timeseries.
+the state timeseries. If `j` is an object that refers to multiple values in the state
+timeseries (e.g. `Colon`), return an iterable of the indexes in the parameter timeseries
+at the appropriate points.
 
 Both versions of this function have default implementations relying on
-[`current_time`](@ref) and [`parameter_timeseries`](@ref).
+[`current_time`](@ref) and [`parameter_timeseries`](@ref), for the cases where `j` is one
+of: `Int`, `CartesianIndex`, `AbstractArray{Bool}`, `Colon` or an iterable of the
+aforementioned.
 """
 function parameter_timeseries_at_state_time end
 
-function parameter_timeseries_at_state_time(valp, i, j)
+function parameter_timeseries_at_state_time(valp, i, j::Union{Int, CartesianIndex})
     state_time = current_time(valp, j)
     timeseries = parameter_timeseries(valp, i)
     searchsortedlast(timeseries, state_time)
 end
 
+function parameter_timeseries_at_state_time(valp, i, ::Colon)
+    parameter_timeseries_at_state_time(valp, i)
+end
+
+function parameter_timeseries_at_state_time(valp, i, j::AbstractArray{Bool})
+    parameter_timeseries_at_state_time(valp, i, only(to_indices(current_time(valp), (j,))))
+end
+
+function parameter_timeseries_at_state_time(valp, i, j)
+    (parameter_timeseries_at_state_time(valp, i, jj) for jj in j)
+end
+
 function parameter_timeseries_at_state_time(valp, i)
-    (parameter_timeseries_at_state_time(valp, i, j) for j in eachindex(current_time(valp)))
+    parameter_timeseries_at_state_time(valp, i, eachindex(current_time(valp)))
 end
 
 """
@@ -131,14 +147,21 @@ Return an indexable collection containing the values of all states in the value
 each of which contain the state values at the corresponding timestep. In this case, the
 two-argument version of the function can also be implemented to efficiently return
 the state values at timestep `i`. By default, the two-argument method calls
-`state_values(valp)[i]`
+`state_values(valp)[i]`. If `i` consists of multiple indices (for example, `Colon`,
+`AbstractArray{Int}`, `AbstractArray{Bool}`) specialized methods may be defined for
+efficiency. By default, `state_values(valp, ::Colon) = state_values(valp)` to avoid
+copying the timeseries.
 
 If this function is called with an `AbstractArray`, it will return the same array.
 
 See: [`is_timeseries`](@ref)
 """
 function state_values end
 
+state_values(arr::AbstractArray) = arr
+state_values(arr, i) = state_values(arr)[i]
+state_values(arr, ::Colon) = state_values(arr)
+
 """
     set_state!(valp, val, idx)
 
@@ -162,11 +185,21 @@ also be implemented to efficiently return the time at timestep `i`. By default,
 argument method calls `current_time(p)[i]`. It is assumed that the timeseries is sorted
 in increasing order.
 
+If `i` consists of multiple indices (for example, `Colon`, `AbstractArray{Int}`,
+`AbstractArray{Bool}`) specialized methods may be defined for efficiency. By default,
+`current_time(valp, ::Colon) = current_time(valp)` to avoid copying the timeseries.
+
+By default, the single-argument version acts as the identity function if
+`valp isa AbstractVector`.
 
 See: [`is_timeseries`](@ref)
 """
 function current_time end
 
+current_time(arr::AbstractVector) = arr
+current_time(valp, i) = current_time(valp)[i]
+current_time(valp, ::Colon) = current_time(valp)
+
 ###########
 # Utilities
 ###########
@@ -232,3 +265,74 @@ end
 function (cw::CallWith)(arg)
     arg(cw.args...)
 end
+
+###########
+# Errors
+###########
+
+struct ParameterTimeseriesValueIndexMismatchError{P <: IsTimeseriesTrait} <: Exception
+    valp
+    indexer
+    args
+
+    function ParameterTimeseriesValueIndexMismatchError{Timeseries}(valp, indexer, args)
+        if is_parameter_timeseries(valp) != Timeseries()
+            throw(ArgumentError("""
+                This should never happen. Expected parameter timeseries value provider, \
+                got $(valp). Open an issue in SymbolicIndexingInterface.jl with an MWE.
+            """))
+        end
+        if is_indexer_timeseries(indexer) != IndexerNotTimeseries()
+            throw(ArgumentError("""
+                This should never happen. Expected non-timeseries indexer, got \
+                $(indexer). Open an issue in SymbolicIndexingInterface.jl with an MWE.
+            """))
+        end
+        return new{Timeseries}(valp, indexer, args)
+    end
+    function ParameterTimeseriesValueIndexMismatchError{NotTimeseries}(valp, indexer)
+        if is_parameter_timeseries(valp) != NotTimeseries()
+            throw(ArgumentError("""
+                This should never happen. Expected non-parameter timeseries value \
+                provider, got $(valp). Open an issue in SymbolicIndexingInterface.jl \
+                with an MWE.
+            """))
+        end
+        if is_indexer_timeseries(indexer) != IndexerTimeseries()
+            throw(ArgumentError("""
+                This should never happen. Expected timeseries indexer, got $(indexer). \
+                Open an issue in SymbolicIndexingInterface.jl with an MWE.
+            """))
+        end
+        return new{NotTimeseries}(valp, indexer, nothing)
+    end
+end
+
+function Base.showerror(io::IO, err::ParameterTimeseriesValueIndexMismatchError{Timeseries})
+    print(io, """
+        Invalid indexing operation: tried to access object of type $(typeof(err.valp)) \
+        (which is a parameter timeseries object) with non-timeseries indexer \
+        $(err.indexer) at index $(err.args) in the timeseries.
+    """)
+end
+
+function Base.showerror(io::IO, err::ParameterTimeseriesValueIndexMismatchError{NotTimeseries})
+    print(io, """
+        Invalid indexing operation: tried to access object of type $(typeof(err.valp)) \
+        (which is not a parameter timeseries object) using timeseries indexer \
+        $(err.indexer).
+    """)
+end
+
+struct MixedParameterTimeseriesIndexError <: Exception
+    valp
+    ts_idxs
+end
+
+function Base.showerror(io::IO, err::MixedParameterTimeseriesIndexError)
+    print(io, """
+        Invalid indexing operation: tried to access object of type $(typeof(err.valp)) \
+        (which is a parameter timeseries object) with variables having mixed timeseries \
+        indexes $(err.ts_idxs).
+    """)
+end