diff --git a/src/TimeAxes.jl b/src/TimeAxes.jl index c9b5853..255c919 100644 --- a/src/TimeAxes.jl +++ b/src/TimeAxes.jl @@ -41,7 +41,10 @@ export @reexport using AxisIndices +include("utils.jl") include("timedim.jl") +include("lead.jl") +include("lag.jl") include("timeaxis.jl") include("timestamps.jl") include("fft.jl") diff --git a/src/lag.jl b/src/lag.jl new file mode 100644 index 0000000..5f11232 --- /dev/null +++ b/src/lag.jl @@ -0,0 +1,125 @@ + +""" + lag(A::AbstractArray, n::Integer) + +Shift the elements of `A` along the the axis of the dimension `dim` by `nshift` +elements later. If `dim` is not specified then the dimension returned by +`timedim` is used. If `A` does not have a time dimension then the last dimension +is assumed to be the time dimension. + +## Examples +```jldoctest +julia> using TimeAxes + +julia> using Unitful: s + +julia> A = NamedAxisArray{(:time,)}(collect(1:5), (1:5)s) +5-element NamedAxisArray{Int64,1} + • time - 1 s:1 s:5 s + + 1 s 1 + 2 s 2 + 3 s 3 + 4 s 4 + 5 s 5 + +julia> lag(A, 1) +4-element NamedAxisArray{Int64,1} + • time - 2 s:1 s:5 s + + 2 s 1 + 3 s 2 + 4 s 3 + 5 s 4 + +julia> lag([1 2 3; 4 5 6; 7 8 9], 1, 1) +2×3 Array{Int64,2}: + 1 2 3 + 4 5 6 + +julia> lag([1 2 3; 4 5 6; 7 8 9], 1, 2) +3×2 Array{Int64,2}: + 1 2 + 4 5 + 7 8 + +``` +""" + +@inline function lag(A::AbstractArray{T,N}, nshift::Integer, dim::Integer) where {T,N} + return A[_lag_indices(Val(N), axes(A), nshift, dim)...] +end + +@inline function lag(A::AxisArray{T,N}, nshift::Integer, dim::Int) where {T,N} + indexing_indices = _lag_indices(Val(N), axes(A), nshift, dim) + p = parent(A)[indexing_indices...] + axs = _shift_axes(axes(A), indexing_indices, axes(p), dim) + return unsafe_reconstruct(A, p, axs) +end + +function lag(A::NamedAxisArray, dim::Int, n::Int) + return NamedDimsArray{dimnames(A)}(lag(parent(A), dim, n)) +end + +@inline function lag(A::AbstractArray{T,N}, nshift::Int) where {T,N} + if has_timedim(A) + return lag(A, nshift, timedim(A)) + else + return lag(A, nshift, N) + end +end + +@inline function _lag_indices(::Val{N}, inds::Tuple, nshift::Integer, dim::Integer) where {N} + ntuple(Val(N)) do i + if i === dim + index = getfield(inds, i) + firstindex(index):(lastindex(index) - nshift) + else + Colon() + end + end +end + +#= +lag(A::AbstractArray, n::Int) = _lag(A, timedim(A), n) + +function _lag(A::NamedAxisArray, dim::Int, n::Int) + return NamedDimsArray{dimnames(A)}(_lag(parent(A), dim, n)) +end +function _lag(A::AxisArray{T,N}, dim::Int, n::Int) where {T,N} + axs = axes(A) + p = parent(A)[_lag_indices(axs, dim, n)...] + newaxs = _lag_axes(axs, axes(p), dim, n) + return AxisArray{eltype(p),N,typeof(p),typeof(newaxs)}(p, newaxs) +end + +# _lag_indices +@inline function _lag_indices(axs::NTuple{N,Any}, dim, n) where {N} + if dim == 1 + axis = first(axs) + return (firstindex(axis):lastindex(axis) - n, ntuple(_-> :, Val(N-1))...) + else + return (:, _lag_indices(tail(axs), dim, n)...) + end +end +_lag_indices(::Tuple{}, dim, n) = () + +# _lag_axes +# TODO should probably add support for non AbstractAxis +_lag_axes(::Tuple{}, ::Tuple{}, dim, n) = () +@inline function _lag_axes(axs::Tuple, newinds::Tuple, dim, n) + if dim == 1 + return (_lag_axis(first(axs), first(newinds), n), map(assign_indices, tail(axs), tail(newinds))...) + else + return (assign_indices(first(axs), first(newinds)), _lag_axes(tail(axs), tail(newinds), dim, n)...) + end +end + +@inline function _lag_axis(axis::AbstractAxis, newinds, n) + if is_indices_axis(axis) + return assign_indices(axis, newinds) + else + return to_axis(axis, keys(axis)[(firstindex(axis) + n):lastindex(axis)], newinds) + end +end +=# diff --git a/src/lead.jl b/src/lead.jl new file mode 100644 index 0000000..21d2c44 --- /dev/null +++ b/src/lead.jl @@ -0,0 +1,81 @@ + +""" + lead(A::AbstractArray, nshift::Integer[, dim::Integer]) + +Shift the elements of `A` along the the axis of the dimension `dim` by `nshift` +elements earlier. If `dim` is not specified then the dimension returned by +`timedim` is used. If `A` does not have a time dimension then the last dimension +is assumed to be the time dimension. + +## Examples +```jldoctest +julia> using TimeAxes + +julia> using Unitful: s + +julia> A = NamedAxisArray{(:time,)}(collect(1:5), (1:5)s) +5-element NamedAxisArray{Int64,1} + • time - 1 s:1 s:5 s + + 1 s 1 + 2 s 2 + 3 s 3 + 4 s 4 + 5 s 5 + +julia> lead(A, 1) +4-element NamedAxisArray{Int64,1} + • time - 1 s:1 s:4 s + + 1 s 2 + 2 s 3 + 3 s 4 + 4 s 5 + +julia> lead([1 2 3; 4 5 6; 7 8 9], 1, 1) +2×3 Array{Int64,2}: + 4 5 6 + 7 8 9 + +julia> lead([1 2 3; 4 5 6; 7 8 9], 1, 2) +3×2 Array{Int64,2}: + 2 3 + 5 6 + 8 9 + +``` +""" +@inline function lead(A::AbstractArray{T,N}, nshift::Integer, dim::Integer) where {T,N} + return A[_lead_indices(Val(N), axes(A), nshift, dim)...] +end + +@inline function lead(A::AxisArray{T,N}, nshift::Integer, dim::Int) where {T,N} + indexing_indices = _lead_indices(Val(N), axes(A), nshift, dim) + p = parent(A)[indexing_indices...] + axs = _shift_axes(axes(A), indexing_indices, axes(p), dim) + return unsafe_reconstruct(A, p, axs) +end + +function lead(A::NamedAxisArray, dim::Int, n::Int) + return NamedDimsArray{dimnames(A)}(lead(parent(A), dim, n)) +end + +@inline function lead(A::AbstractArray{T,N}, nshift::Int) where {T,N} + if has_timedim(A) + return lead(A, nshift, timedim(A)) + else + return lead(A, nshift, N) + end +end + +@inline function _lead_indices(::Val{N}, inds::Tuple, nshift::Integer, dim::Integer) where {N} + ntuple(Val(N)) do i + index = getfield(inds, i) + if i === dim + index = getfield(inds, i) + (firstindex(index) + nshift):lastindex(index) + else + Colon() + end + end +end diff --git a/src/timedim.jl b/src/timedim.jl index a9f4d96..6a7a044 100644 --- a/src/timedim.jl +++ b/src/timedim.jl @@ -68,163 +68,3 @@ Throw an error if the `x` has a time dimension that is not the last dimension. return nothing end end - -""" - lead(A::AbstractArray, n::Integer) - -Shift the elements of `A` along the time axis so that they are `n` time units sooner. - -```jldoctest -julia> using TimeAxes - -julia> using Unitful: s - -julia> A = NamedAxisArray{(:time,)}(collect(1:5), (1:5)s) -5-element NamedAxisArray{Int64,1} - • time - 1 s:1 s:5 s - - 1 s 1 - 2 s 2 - 3 s 3 - 4 s 4 - 5 s 5 - -julia> lead(A, 1) -4-element NamedAxisArray{Int64,1} - • time - 1 s:1 s:4 s - - 1 s 2 - 2 s 3 - 3 s 4 - 4 s 5 - -``` -""" -lead(A::AbstractArray, n::Int) = _lead(A, timedim(A), n) - -""" - lag(A::AbstractArray, n::Integer) - -Shift the elements of `A` along the time axis so that they are `n` time units later. - -```jldoctest -julia> using TimeAxes - -julia> using Unitful: s - -julia> A = NamedAxisArray{(:time,)}(collect(1:5), (1:5)s) -5-element NamedAxisArray{Int64,1} - • time - 1 s:1 s:5 s - - 1 s 1 - 2 s 2 - 3 s 3 - 4 s 4 - 5 s 5 - -julia> lag(A, 1) -4-element NamedAxisArray{Int64,1} - • time - 2 s:1 s:5 s - - 2 s 1 - 3 s 2 - 4 s 3 - 5 s 4 - -``` -""" -lag(A::AbstractArray, n::Int) = _lag(A, timedim(A), n) - - -### -### lead -### -function _lead(A::NamedAxisArray, dim::Int, n::Int) - return NamedDimsArray{dimnames(A)}(_lead(parent(A), dim, n)) -end - -function _lead(A::AxisArray{T,N}, dim::Int, n::Int) where {T,N} - axs = axes(A) - p = parent(A)[_lead_indices(axs, dim, n)...] - newaxs = _lead_axes(axs, axes(p), dim, n) - return AxisArray{eltype(p),N,typeof(p),typeof(newaxs)}(p, newaxs) -end - -# _lead_indices -@inline function _lead_indices(axs::NTuple{N,Any}, dim, n) where {N} - if dim == 1 - axis = first(axs) - return ((firstindex(axis) + n):lastindex(axis), ntuple(_-> :, Val(N-1))...) - else - return (:, _lead_indices(tail(axs), dim, n)...) - end -end -_lead_indices(::Tuple{}, dim, n) = () - -# _lead_axes -# TODO should probably add support for non AbstractAxis -_lead_axes(::Tuple{}, ::Tuple{}, dim, n) = () -@inline function _lead_axes(axs::Tuple, newinds::Tuple, dim, n) - if dim == 1 - return ( - _lead_axis(first(axs), first(newinds), n), - map(assign_indices, tail(axs), tail(newinds))... - ) - else - return ( - assign_indices(first(axs), first(newinds)), - _lead_axes(tail(axs), tail(newinds), dim, n)... - ) - end -end - -@inline function _lead_axis(axis::AbstractAxis, newinds, n) - if is_indices_axis(axis) - return AxisIndices.assign_indices(axis, newinds) - else - return to_axis(axis, keys(axis)[firstindex(axis):(lastindex(axis) - n)], newinds) - end -end - -### -### lags -### -function _lag(A::NamedAxisArray, dim::Int, n::Int) - return NamedDimsArray{dimnames(A)}(_lag(parent(A), dim, n)) -end -function _lag(A::AxisArray{T,N}, dim::Int, n::Int) where {T,N} - axs = axes(A) - p = parent(A)[_lag_indices(axs, dim, n)...] - newaxs = _lag_axes(axs, axes(p), dim, n) - return AxisArray{eltype(p),N,typeof(p),typeof(newaxs)}(p, newaxs) -end - -# _lag_indices -@inline function _lag_indices(axs::NTuple{N,Any}, dim, n) where {N} - if dim == 1 - axis = first(axs) - return (firstindex(axis):lastindex(axis) - n, ntuple(_-> :, Val(N-1))...) - else - return (:, _lag_indices(tail(axs), dim, n)...) - end -end -_lag_indices(::Tuple{}, dim, n) = () - -# _lag_axes -# TODO should probably add support for non AbstractAxis -_lag_axes(::Tuple{}, ::Tuple{}, dim, n) = () -@inline function _lag_axes(axs::Tuple, newinds::Tuple, dim, n) - if dim == 1 - return (_lag_axis(first(axs), first(newinds), n), map(assign_indices, tail(axs), tail(newinds))...) - else - return (assign_indices(first(axs), first(newinds)), _lag_axes(tail(axs), tail(newinds), dim, n)...) - end -end - -@inline function _lag_axis(axis::AbstractAxis, newinds, n) - if is_indices_axis(axis) - return assign_indices(axis, newinds) - else - return to_axis(axis, keys(axis)[(firstindex(axis) + n):lastindex(axis)], newinds) - end -end diff --git a/src/utils.jl b/src/utils.jl new file mode 100644 index 0000000..076a252 --- /dev/null +++ b/src/utils.jl @@ -0,0 +1,66 @@ + +@inline function _noshift_axis(axis::A, inds::I) where {A,I} + if is_indices_axis(axis) + return to_axis(axis, nothing, inds, false, AxisIndices.Interface.Staticness(axis)) + else + return to_axis(axis, keys(axis), inds, false, AxisIndices.Interface.Staticness(axis)) + end +end + +@inline function _shift_axis( + axis::AbstractAxis, + indexing_index::AbstractUnitRange, + parent_index::AbstractUnitRange +) + + if is_indices_axis(axis) + return to_axis( + axis, + nothing, + parent_index, + false, + AxisIndices.Interface.Staticness(axis) + ) + else + return to_axis( + axis, + @inbounds(keys(axis)[indexing_index]), + parent_index, + false, + AxisIndices.Interface.Staticness(axis) + ) + end +end + +_shift_axes(::Tuple{}, ::Tuple{}, ::Tuple{}, dim::Int) = () +@inline function _shift_axes( + old_axes::Tuple, + indexing_indices::Tuple, + parent_indices::Tuple, + dim::Int +) + + if dim === 1 + return ( + _shift_axis( + first(old_axes), + first(indexing_indices), + first(parent_indices), + ), + map(_noshift_axis, tail(old_axes), tail(parent_indices))... + ) + else + return ( + _noshift_axis( + first(old_axes), + first(parent_indices) + ), + _shift_axes( + tail(old_axes), + tail(indexing_indices), + tail(parent_indices), + dim - 1 + )... + ) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 66bb28b..56374a9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,6 +5,7 @@ using AxisIndices using Dates using Documenter using Unitful +using Unitful: s nia = NamedAxisArray(reshape(1:6, 2, 3), x = 2:3, time = 3.0:5.0) @test has_timedim(nia) @@ -46,6 +47,55 @@ t2 = @inferred(t[:ts1..:ts2]) @test values(t2) == 1:3 @test keys(t2) == Second(1):Second(1):Second(3) +@testset "lead" begin + A = [1 2 3; 4 5 6; 7 8 9] + A_axes = AxisArray(A,(1:3, (1:3)s)); + A_named_axes = NamedAxisArray{(:time,:_)}(A, (1:3, (1:3)s)) + + @testset "Array" begin + @test @inferred(lead(A, 1, 1)) == [4 5 6; 7 8 9] + @test @inferred(lead(A, 1, 2)) == [2 3; 5 6; 8 9] + end + + @testset "AxisArray" begin + @test @inferred lead(A_axes, 1, 1) == [4 5 6; 7 8 9] + @test @inferred(lead(A_axes, 1, 2)) == [2 3; 5 6; 8 9] + end + + @testset "NamedAxisArray" begin + @test @inferred lead(A_named_axes, 1, 1) == [4 5 6; 7 8 9] + @test @inferred(lead(A_named_axes, 1, 2)) == [2 3; 5 6; 8 9] + end + # this has a mutable axis and therefore it needs to handle this in a type + # stable way when recreating the axes + @test lead(NamedAxisArray{(:time,)}(collect(1:5), (1:5)s), 1) == [2, 3, 4, 5] +end + + +@testset "lag" begin + A = [1 2 3; 4 5 6; 7 8 9] + A_axes = AxisArray(A,(1:3, (1:3)s)); + A_named_axes = NamedAxisArray{(:time,:_)}(A, (1:3, (1:3)s)) + + @testset "Array" begin + @test @inferred(lag(A, 1, 1)) == [1 2 3; 4 5 6] + @test @inferred(lag(A, 1, 2)) == [1 2; 4 5; 7 8] + end + + @testset "AxisArray" begin + @test @inferred(lag(A_axes, 1, 1)) == [1 2 3; 4 5 6] + @test @inferred(lag(A_axes, 1, 2)) == [1 2; 4 5; 7 8] + end + + @testset "NamedAxisArray" begin + @test @inferred(lag(A_named_axes, 1, 1)) == [1 2 3; 4 5 6] + @test @inferred(lag(A_named_axes, 1, 2)) == [1 2; 4 5; 7 8] + end + # this has a mutable axis and therefore it needs to handle this in a type + # stable way when recreating the axes + @test lag(NamedAxisArray{(:time,)}(collect(1:5), (1:5)s), 1) == [1, 2, 3, 4] +end + @testset "fft-tests" begin A = reshape(1:6, 2, 3) A_named_axes = NamedAxisArray(A, x = 2:3, time = 3.0:5.0)