Reputation: 1145
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?
Upvotes: 1
Views: 180
Reputation: 221564
You can employ the best tool for vectorization, bsxfun
here -
B = A(bsxfun(@plus,[1:M]',M*(0:N-1)) + M*N*(g-1))
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))
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.
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.
Upvotes: 2
Reputation: 112679
Try using sub2ind
. This assumes g
is defined as an M
xN
matrix with possible values 1
, ..., M
:
[ii, jj] = ndgrid(1:M, 1:N);
B = A(sub2ind([M N M], ii, jj, g));
Upvotes: 0