diff --git a/src/core/time.jl b/src/core/time.jl index 03496c2b6..fe15a34fc 100644 --- a/src/core/time.jl +++ b/src/core/time.jl @@ -125,6 +125,14 @@ function Base.:+(ts::VariableTimestep{TIMES}, val::Int) where {TIMES} new_ts = VariableTimestep{TIMES}(ts.t + val) end +function Base.:-(ts::Timestep, val::Int) + return Timestep(ts.value, ts.offset_in_timesteps - val) +end + +function Base.:+(ts::Timestep, val::Int) + return Timestep(ts.value, ts.offset_in_timesteps + val) +end + # # 2. CLOCK # @@ -195,6 +203,20 @@ function _missing_data_check(data) end end +# Helper functions for Timestep type +function _get_time_value_position_(times::Array, ts::Timestep) + t = findfirst(isequal.(ts.value, times)) + if t === nothing + error("cannot use Timestep with value $(ts.value), value is not in the TimestepArray") + end + + t_offset = t + ts.offset_in_timesteps + if t_offset > length(times) + error("cannot get offset of $(ts.offset_in_timesteps) from $(ts.value), offset is after the end of the TimestepArray") + end + return t_offset +end + # Helper macro used by connector macro allow_missing(expr) let e = gensym("e") @@ -239,6 +261,18 @@ function Base.getindex(v::TimestepVector{VariableTimestep{D_TIMES}, T}, ts::Vari _missing_data_check(data) end +function Base.getindex(v::TimestepVector{FixedTimestep{FIRST, STEP, LAST}, T}, ts::Timestep) where {T, FIRST, STEP, LAST} + t = _get_time_value_position_([FIRST:STEP:LAST...], ts) + data = v.data[t] + _missing_data_check(data) +end + +function Base.getindex(v::TimestepVector{VariableTimestep{TIMES}, T}, ts::Timestep) where {T, TIMES} + t = _get_time_value_position_(TIMES, ts) + data = v.data[t] + _missing_data_check(data) +end + # int indexing version supports old-style components and internal functions, not # part of the public API @@ -333,6 +367,30 @@ function Base.getindex(mat::TimestepMatrix{VariableTimestep{D_TIMES}, T, 2}, idx _missing_data_check(data) end +function Base.getindex(mat::TimestepMatrix{FixedTimestep{FIRST, STEP, LAST}, T, 1}, ts::Timestep, idx::AnyIndex) where {T, FIRST, STEP, LAST} + t = _get_time_value_position_([FIRST:STEP:LAST...], ts) + data = mat.data[t, idx] + _missing_data_check(data) +end + +function Base.getindex(mat::TimestepMatrix{VariableTimestep{TIMES}, T, 1}, ts::Timestep, idx::AnyIndex) where {T, TIMES} + t = _get_time_value_position_(TIMES, ts) + data = mat.data[t, idx] + _missing_data_check(data) +end + +function Base.getindex(mat::TimestepMatrix{FixedTimestep{FIRST, STEP, LAST}, T, 2}, idx::AnyIndex, ts::Timestep) where {T, FIRST, STEP, LAST} + t = _get_time_value_position_([FIRST:STEP:LAST...], ts) + data = mat.data[idx, t] + _missing_data_check(data) +end + +function Base.getindex(mat::TimestepMatrix{VariableTimestep{TIMES}, T, 2}, idx::AnyIndex, ts::Timestep) where {T, TIMES} + t = _get_time_value_position_(TIMES, ts) + data = mat.data[idx, t] + _missing_data_check(data) +end + function Base.setindex!(mat::TimestepMatrix{FixedTimestep{FIRST, STEP}, T, 1}, val, ts::FixedTimestep{FIRST, STEP, LAST}, idx::AnyIndex) where {T, FIRST, STEP, LAST} setindex!(mat.data, val, ts.t, idx) @@ -449,6 +507,18 @@ function Base.getindex(arr::TimestepArray{VariableTimestep{D_TIMES}, T, N, ti}, return arr.data[idxs1..., t, idxs2...] end +function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T, N, ti}, idxs::Union{Timestep, AnyIndex}...) where {T, N, ti, FIRST, STEP, LAST} + idxs1, ts, idxs2 = split_indices(idxs, ti) + t = _get_time_value_position_([FIRST:STEP:LAST...], ts) + return arr.data[idxs1..., t, idxs2...] +end + +function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T, N, ti}, idxs::Union{Timestep, AnyIndex}...) where {T, N, ti, TIMES} + idxs1, ts, idxs2 = split_indices(idxs, ti) + t = _get_time_value_position_(TIMES, ts) + return arr.data[idxs1..., t, idxs2...] +end + function Base.setindex!(arr::TimestepArray{FixedTimestep{FIRST, STEP}, T, N, ti}, val, idxs::Union{FixedTimestep{FIRST, STEP, LAST}, AnyIndex}...) where {T, N, ti, FIRST, STEP, LAST} idxs1, ts, idxs2 = split_indices(idxs, ti) setindex!(arr.data, val, idxs1..., ts.t, idxs2...) diff --git a/src/core/types.jl b/src/core/types.jl index bdbefc1b5..98f54e9ce 100644 --- a/src/core/types.jl +++ b/src/core/types.jl @@ -22,6 +22,15 @@ struct VariableTimestep{TIMES} <: AbstractTimestep end end +struct Timestep <: AbstractTimestep + value::Int + offset_in_timesteps::Int + + function Timestep(v::Int, offset_in_timesteps::Int=0) + return new(v, offset_in_timesteps) + end +end + mutable struct Clock{T <: AbstractTimestep} ts::T