diff --git a/src/core/time_arrays.jl b/src/core/time_arrays.jl index 473e05ce9..dede9963d 100644 --- a/src/core/time_arrays.jl +++ b/src/core/time_arrays.jl @@ -39,6 +39,14 @@ function _missing_data_check(data, t) end end +function _missing_data_check(data) + if data === missing + throw(MissingException("Cannot get index; data is missing.")) + else + return data + end +end + # Helper function for getindex; throws an error if the TimestepIndex index is out of range of the TimestepArray function _index_bounds_check(data, dim, t) if size(data, dim) < t @@ -391,12 +399,12 @@ time_labels(obj::TimestepArray{VariableTimestep{TIMES}, T, N, ti}) where {TIMES, function Base.getindex(arr::TimestepArray{FixedTimestep{A_FIRST, A_STEP, A_LAST}, T, N, ti}, idxs::Union{FixedTimestep{T_FIRST, T_STEP, T_LAST}, AnyIndex}...) where {A_FIRST, A_STEP, A_LAST, T_FIRST, T_STEP, T_LAST, T, N, ti} idxs1, ts, idxs2 = idxs[1:ti - 1], idxs[ti], idxs[ti + 1:end] - return arr.data[idxs1..., ts.t, idxs2...] + return _missing_data_check(arr.data[idxs1..., ts.t, idxs2...]) end function Base.getindex(arr::TimestepArray{VariableTimestep{A_TIMES}, T, N, ti}, idxs::Union{VariableTimestep{T_TIMES}, AnyIndex}...) where {A_TIMES, T_TIMES, T, N, ti} idxs1, ts, idxs2 = idxs[1:ti - 1], idxs[ti], idxs[ti + 1:end] - return arr.data[idxs1..., ts.t, idxs2...] + return _missing_data_check(arr.data[idxs1..., ts.t, idxs2...]) end function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T_data, N, ti}, idxs::Union{TimestepValue{T_time}, AnyIndex}...) where {T_data, N, ti, FIRST, STEP, LAST, T_time} @@ -405,7 +413,7 @@ function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T_da t = _get_time_value_position([FIRST:STEP:LAST...], ts) arr.data isa SubArray ? view_offset = arr.data.indices[ti][1] - 1 : view_offset = 0 t = t - view_offset - return arr.data[idxs1..., t, idxs2...] + return _missing_data_check(arr.data[idxs1..., t, idxs2...]) end function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T_data, N, ti}, idxs::Union{TimestepValue{T_time}, AnyIndex}...) where {T_data, N, ti, TIMES, T_time} @@ -414,7 +422,7 @@ function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T_data, N, ti t = _get_time_value_position(TIMES, ts) arr.data isa SubArray ? view_offset = arr.data.indices[ti][1] - 1 : view_offset = 0 t = t - view_offset - return arr.data[idxs1..., t, idxs2...] + return _missing_data_check(arr.data[idxs1..., t, idxs2...]) end function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T, N, ti}, idxs::Union{TimestepIndex, AnyIndex}...) where {T, N, ti, FIRST, STEP, LAST} @@ -422,7 +430,7 @@ function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T, N idxs1, ts, idxs2 = idxs[1:ti - 1], idxs[ti], idxs[ti + 1:end] t = ts.index _index_bounds_check(arr.data, ti, t) - return arr.data[idxs1..., t, idxs2...] + return _missing_data_check(arr.data[idxs1..., t, idxs2...]) end function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T, N, ti}, idxs::Union{TimestepIndex, AnyIndex}...) where {T, N, ti, TIMES} @@ -430,7 +438,7 @@ function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T, N, ti}, id idxs1, ts, idxs2 = idxs[1:ti - 1], idxs[ti], idxs[ti + 1:end] t = ts.index _index_bounds_check(arr.data, ti, t) - return arr.data[idxs1..., t, idxs2...] + return _missing_data_check(arr.data[idxs1..., t, idxs2...]) end function Base.setindex!(arr::TimestepArray{FixedTimestep{A_FIRST, A_STEP, A_LAST}, T, N, ti}, val, idxs::Union{FixedTimestep{T_FIRST, T_STEP, T_LAST}, AnyIndex}...) where {A_FIRST, A_STEP, A_LAST, T_FIRST, T_STEP, T_LAST, T, N, ti} @@ -475,7 +483,7 @@ end function Base.getindex(arr::TimestepArray{TS, T, N, ti}, idxs::Union{Array{TimestepIndex,1}, AnyIndex}...) where {TS, T, N, ti} idxs1, ts_array, idxs2 = idxs[1:ti - 1], idxs[ti], idxs[ti + 1:end] ts_idxs = _get_ts_indices(ts_array) - return arr.data[idxs1..., ts_idxs, idxs2...] + return _missing_data_check(arr.data[idxs1..., ts_idxs, idxs2...]) end function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T_data, N, ti}, idxs::Union{Array{TimestepValue{T_time},1}, AnyIndex}...) where {T_data, N, ti, FIRST, STEP, LAST, T_time} @@ -483,7 +491,7 @@ function Base.getindex(arr::TimestepArray{FixedTimestep{FIRST, STEP, LAST}, T_da ts_idxs = _get_ts_indices(ts_array, [FIRST:STEP:LAST...]) arr.data isa SubArray ? view_offset = arr.data.indices[ti][1] - 1 : view_offset = 0 ts_idxs = ts_idxs .- view_offset - return arr.data[idxs1..., ts_idxs, idxs2...] + return _missing_data_check(arr.data[idxs1..., ts_idxs, idxs2...]) end function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T_data, N, ti}, idxs::Union{Array{TimestepValue{T_time},1}, AnyIndex}...) where {T_data, N, ti, TIMES, T_time} @@ -491,7 +499,7 @@ function Base.getindex(arr::TimestepArray{VariableTimestep{TIMES}, T_data, N, ti ts_idxs = _get_ts_indices(ts_array, TIMES) arr.data isa SubArray ? view_offset = arr.data.indices[ti][1] - 1 : view_offset = 0 ts_idxs = ts_idxs .- view_offset - return arr.data[idxs1..., ts_idxs, idxs2...] + return _missing_data_check(arr.data[idxs1..., ts_idxs, idxs2...]) end function Base.setindex!(arr::TimestepArray{TS, T, N, ti}, vals, idxs::Union{Array{TimestepIndex,1}, AnyIndex}...) where {TS, T, N, ti}