Skip to content

Commit

Permalink
Update docs/tests for lead/lag
Browse files Browse the repository at this point in the history
- More examples for lead/lag docstrings
- can choose what dimension to apply lead/lag
- Automatically chooses last dimension if no timedim
  • Loading branch information
Tokazama committed Jul 8, 2020
1 parent 33d8660 commit 78211b4
Show file tree
Hide file tree
Showing 6 changed files with 325 additions and 160 deletions.
3 changes: 3 additions & 0 deletions src/TimeAxes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
125 changes: 125 additions & 0 deletions src/lag.jl
Original file line number Diff line number Diff line change
@@ -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
=#
81 changes: 81 additions & 0 deletions src/lead.jl
Original file line number Diff line number Diff line change
@@ -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
160 changes: 0 additions & 160 deletions src/timedim.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading

2 comments on commit 78211b4

@Tokazama
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

Release notes:

  • More examples for lead/lag docstrings
  • can choose what dimension to apply lead/lag
  • Automatically chooses last dimension if no timedim
  • Fixed related type inference issues

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/17656

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.3.1 -m "<description of version>" 78211b40769fd1b0c3e3ab97955cb98608950dda
git push origin v0.3.1

Please sign in to comment.