Indexing an array of arbitrary dimension

New to Julia
1

I’m trying to write a function which returns a random index for each dimension of an array.
It’s easy to do this if the dimensions of the array are known, but I want to abstract my function to work on arrays of arbitrary dimension.

The 1D case is very simple:

function rand_index_1D(a::AbstractVector)
    return rand(1:length(a))
end

Similarly, the 2D case:

function rand_index_2D(a::AbstractMatrix)
    i = rand(1:length(a[1,:]))
    j = rand(1:length(a[:,1]))
    return (i,j)
end

The trouble I’m having is writing a single function which will work for an array of any dimension. I’ve been trying to rewrite the 2D case into a loop, but I’m not sure how to slice the array correctly. What I have so far:

function rand_index_nD(a::AbstractArray)
    dims = size(a)
    nout = length(dims)
    v = Vector{Int64}(undef, nout)
    for n = 1:nout
    #    assign v[n] somehow
    end
    return v
end

I suspect that Julia has a built-in function (eachslice, selectdim, etc) which would really help here, but I haven’t quite figured it out.

Any help would be appreciated,
Thanks in advance.

2

I’d do: rand.(axes(A)) or perhaps even CartesianIndex(rand.(axes(A)).

4 Likes
3

axes() is exactly what I needed. The following works for a (rectangular) array of any dimension:

function uniform_index(a::AbstractArray)
    foo = axes(a)
    return [rand(foo[i]) for i=1:length(foo)]
end

Unfortunately it does allocate seem to allocate, but that may just be an artifact of benchmarking in global scope. I’ll get back to you on that one.

Thanks again,

4

Yes, that will allocate because you’re explicitly creating an array with that comprehension. It really can be as simple as what I wrote:

uniform_index(a::AbstractArray) = rand.(axes(a))
5 Likes
5

If you want to avoid allocations, you could use an NTuple:

function uniform_index(a::AbstractArray{T, N}) where {T, N}
    return ntuple(Val(N)) do i
        rand(axes(a, i))  # safer than 1:size(a, i)
    end
end

using BenchmarkTools
A = rand(2, 3, 4)
@btime uniform_index($A)
#  16.016 ns (0 allocations: 0 bytes)
# (2, 3, 2)

A more elegant and equally performant (or perhaps a bit faster) approach is to use CartesianIndices:

@btime rand(CartesianIndices($A))
#   14.000 ns (0 allocations: 0 bytes)
# CartesianIndex(2, 3, 3)

Edit: A third option is

@btime map(rand, axes($A))
#  11.912 ns (0 allocations: 0 bytes)
# (1, 3, 1)

Edit2: Well, apparently I didn’t understand how broadcast works, since mbauman’s rand.(axes(A)) was already non-allocating :slight_smile:

@btime rand.(axes($A))
#  11.912 ns (0 allocations: 0 bytes)
# (1, 2, 1)
4 Likes