Accessing Different Indexes in Arrays


I have the current function in Julia but I realize it is quite repetitive and bad. Is there any way of generalizing this function? I want to create a function that takes in an array and range and returns the array indexed from the last dimension.

    function lastdimrange(arr::Array{T,N}, r::R) where {T,N,R}
        if N == 1
            return arr[r]
        elseif N > 1
            new_dims = tuple((size(arr, i) for i in 1:N-1)..., length(r))
            result = similar(arr[:, :, 1], new_dims)
            if N == 2
                result .= arr[:, r]
            elseif N == 3
                result .= arr[:, :, r]
            elseif N == 4
                result .= arr[:, :, :, r]
            elseif N == 5
                result .= arr[:, :, :, :, r]
            elseif N == 6
                result .= arr[:, :, :, :, :, r]
            elseif N == 7
                result .= arr[:, :, :, :, :, :, r]
                error("This function does not support arrays with more than 7 dimensions.")
            return result
            error("Invalid array dimensions.")

I think you may want some variation of the following.

    eachslice(arr; dims=N)[r]

This looks promising! But when r is a range from, say, 1:2, is there any way of obtaining an array object? I get a strange two element vector.

Does stack(eachslice(A, dims=3)[1:2]) do what you want?

Another approach more similar to your original one might be

idx = ntuple(i->i == N ? r : Colon(), N)
colons(::V) where {V<:Val} = ntuple((i -> (:)), V())
lastdimrange(a::AbstractArray{T, n}, r) where {T, n} = a[colons(Val{n - 1}())..., r]
lastdimrange(a::AbstractArray{<:Any, 0}, ::Any) = error("nonsensical")

Thank you all! I found that selectdim() actually was what I was looking for.


naturally the solution is the one already indicated


but only to add a different way to the first proposals

1 Like