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, i.e. 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

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

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

Answer 2

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.