Created
December 9, 2021 21:53
-
-
Save cossio/a04523e90036208ec2f46af8475ac62d to your computer and use it in GitHub Desktop.
iterator over columns of high-dimensional array
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
columns(A) | |
Returns an array over the columns of `A` (as views). Similar to `eachcol` but | |
for higher-dimensional arrays. In general column (i,j,k,...) is defined as | |
`A[:,i,j,k,...]`. | |
""" | |
function columns(A::AbstractArray) | |
[A[:,I] for I in CartesianIndices(Base.tail(axes(A)))] | |
end | |
@testset "columns" begin | |
A = randn(4,5) | |
@test columns(A) == collect(eachcol(A)) | |
@inferred columns(A) | |
A = randn(5,5,3) | |
@test vec(columns(A)) == collect(eachcol(reshape(A,5,15))) | |
@inferred columns(A) | |
end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment