Accessing multi-dimensional arrays in 1 go

General Usage
,
1

I am trying to figure out if there is an easy (and quick) way of accessing elements of a multi-dimensional array according to multi-dimensional array of indices.

For example:

A = reshape(1:8,2,2,2);
indxs = repeat([1],1,2,2);

I would like it to output

Mat = zeros(size(indxs,1),size(A,2),size(A,3));
for i=1:size(A,2)
     for j=1:size(A,3)
          Mat[:,i,j] = A[indxs[:,i,j],i,j];
     end
end

whose output is

[:, :, 1] =
 1.0  3.0

[:, :, 2] =
 5.0  7.0

all in one go, so I try

A[indxs,[1;2],[1;2]]

but, what I thought would work is off (it adds dimensions by doing all pair-wise combinations):

[:, :, 1, 1, 1] =
 1  1
[:, :, 2, 1, 1] =
 1  1
[:, :, 1, 2, 1] =
 3  3
[:, :, 2, 2, 1] =
 3  3
[:, :, 1, 1, 2] =
 5  5
[:, :, 2, 1, 2] =
 5  5
[:, :, 1, 2, 2] =
 7  7
[:, :, 2, 2, 2] =
 7  7

Any ideas if what I want exists?

2

Nothing wrong with writing loops, although I’d write a 3rd one, since here A[indxs[:,1,2],1,2] isa Array which you need not create. I also have a package you may like, which does exactly that:

@tullio Mat2[c,i,j] := A[indxs[c,i,j],i,j]
Mat2 == Mat
2 Likes
3

I was hoping there was a way to do it without resorting to loops, just a matter of indexation but haven’t found a way yet.

Thank you for your suggestion, however, in this context

appears to be slower than the loop.

4

You may like What to do about nonscalar indexing? · Issue #30845 · JuliaLang/julia · GitHub for discussion of whether to make things like A.[indxs, :, :] mean something. But as you’ve observed, right now A[indxs, :, :] adds dimensions.

Maybe there’s a clever way of broadcasting getindex. I guess you can write a comprehension:

[A[indxs[I], I.I[2:end]...] for I in CartesianIndices(indxs)]

Are you sure? I get it’s about 3x faster at this size, and with everything 100x100x100.

1 Like
5

I tried this:

M = 50; T = 20; N = M*T*20;
A = reshape(1:N,Int(N/(M*T)),M,T);
indxs = rand(1:Int(N/(M*T)),Int(N/(M*T)),M,T);


@time begin
Mat = zeros(size(indxs,1),size(A,2),size(A,3)); 
for i=1:size(A,2)
     for j=1:size(A,3)
          Mat[:,i,j] = A[indxs[:,i,j],i,j];
     end
end
end

@time @tullio Mat2[c,i,j] := A[indxs[c,i,j],i,j];

and the times were 0.001876 seconds vs. 0.363211 seconds. But I might have done something wrong.

6

OK, I think that is mostly the time of generating the code, which should happen once. Using GitHub - JuliaCI/BenchmarkTools.jl: A benchmarking framework for the Julia language and being careful about globals, at these sizes, I get:

julia> @btime @tullio Mat2[c,i,j] := $A[$indxs[c,i,j],i,j];
  27.247 μs (3 allocations: 156.38 KiB)

julia> @btime [$A[$indxs[I], I.I[2:end]...] for I in CartesianIndices($indxs)];
  84.033 μs (2 allocations: 156.33 KiB)

julia> @btime f1($A, $indxs); # as in original post
  153.832 μs (2002 allocations: 625.08 KiB)

The last two can be made quicker with @inbounds, but perhaps with a check that issubset(extrema(indxs), axes(A,1)) first.

1 Like