3d Matrix to 2d Matrix matlab

Question

I am using Matlab R2014a.

I have a 3-dimensional M x N x M matrix A. I would like a vectorized way to extract a 2 dimensional matrix B from it, such that for each i,j I have

B(i,j)=A(i,j,g(i,j))

where g is a 2-dimensional index matrix of size M x N, ie with integral values in {1,2,...,M}.

The context is that I am representing a function A(k,z,k') as a 3-dimensional matrix, the function g(k,z) as a 2-dimensional matrix, and I would like to compute the function

h(k,z)=f(k,z,g(k,z))

This seems like a simple and common thing to try to do but I really can't find anything online. Thank you so much to whoever can help!

My first thought was to try something like B = A(:,:,g) or B=A(g) but neither of these works, unsurprisingly. Is there something similar?

Answer 1

You can employ the best tool for vectorization, bsxfun here -

B = A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1))

Explanation: Breaking it down to two steps

Step #1: Calculate the indices corresponding to the first two dimensions (rows and columns) of A -

bsxfun(@plus,[1:M]',M*(0:N-1))

Step #2: Add the offset needed to include the dim-3 indices being supplied by g and index into A with those indices to get our desired output -

A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1))

---

Benchmarking

Here's a quick benchmark test to compare this bsxfun based approach against the ndgrid + sub2ind based solution as presented in Luis's solution with M and N as 100 .

The benchmarking code using tic-toc would look something like this -

M = 100;
N = 100;
A = rand(M,N,M);
g = randi(M,M,N);

num_runs = 5000; %// Number of iterations to run each approach

%// Warm up tic/toc.
for k = 1:50000
    tic(); elapsed = toc();
end

disp('-------------------- With BSXFUN')
tic
for iter = 1:num_runs
    B1 = A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1));  %//'
end
toc, clear B1

disp('-------------------- With NDGRID + SUB2IND')
tic
for iter = 1:num_runs
    [ii, jj] = ndgrid(1:M, 1:N);
    B2 = A(sub2ind([M N M], ii, jj, g));
end
toc

Here's the runtime results -

-------------------- With BSXFUN
Elapsed time is 2.090230 seconds.
-------------------- With NDGRID + SUB2IND
Elapsed time is 4.133219 seconds.

Conclusions

As you can see bsxfun based approach works really well, both as a vectorized approach and good with performance too.

Why is bsxfun better here -

• bsxfun does replication of offsetted elements and adding them, both on-the-fly .

• In the other solution, ndgrid internally makes two function calls to repmat , thus incurring the function call overheads. At the next step, sub2ind spends time in adding the offsets to get the linear indices, bringing in another function call overhead.

Answer 2

Try using sub2ind . This assumes g is defined as an M x N matrix with possible values 1 , ..., M :

[ii, jj] = ndgrid(1:M, 1:N);
B = A(sub2ind([M N M], ii, jj, g));