Indexing 3-D matrix with 2-D matrix of indexes

I have a 3-D matrix that represents data for XYZ. I would like to pull the values from data given a 2-D matrix reprenting the Z indices of the values I want to pull. I thought there would be an easy way to index this without using a loop since my data is rather large, but I am stumped.

Below I have an example of the output I want using a loop...

% Sample 3-D Matrix of Data
data = ones(3,3,3);
data(:,:,1) = 0.1*data(:,:,1);
data(:,:,2) = 0.2*data(:,:,2);
data(:,:,3) = 0.3*data(:,:,3);

% Sample 2-D matrix of Z indices 
idx = randi(3,3,3);

% Code for Loop, would like to do this without loop
sz = size(idx);
output = nan(sz);
for i = 1:numel(idx)
   [j,k] = ind2sub(sz,i);
   output(i) = data(j,k,idx(i));
end

[Stephen23]

See also:

https://www.mathworks.com/matlabcentral/answers/320694-how-to-index-3d-array-with-2d-array

[meghannmarie]

So, I am not applying the a 3D logical matrix to all the data. I have a 4D matrix with indexes to the points in 3rd dimension which I want to keep and store in a 3D matrix (obviously this value is not static, it changes from point to point). I am still struggling with correct answer. I wonder if trying to turn 3D matrix into a 4D logical matrix would be easier, but not sure....

To put this in perspective, I have a very large 4 dimensional matrix (time, depth, lat, and lon) and I am trying to pull just the data from the bottom layer which the depth changes. My output would be a 3 dimensional matrix (time, lat, and lon) - all the data is from the bottom.

I could easily do this in a loop but would take lots of time...

Answer 1

[Matt J]

I don't think your example works because idx wasn't 2D as the comments say it should be. Here is what I think you want, though.

[m,n,p]=size(data);

[I,J]=ndgrid(1:m,1:n);

ouput = data(  sub2ind([m,n,p], I,J,idx)   );

[meghannmarie]

That worked! I am try to get it to work for my 4-D dataset (time,depth,lat,lon) and when I plot the output the data is wrong. I need to check my code to make sure I am not doing anything else wrong earlier on. But do you notice anything wrong with this code, I am trying to use your code from above but with one more dimension?

d = permute(data,[1,3,4,2]);
[m,n,p,o]=size(d);
[I,J,K]=ndgrid(1:m,1:n,1:p);
output = data(sub2ind([m,n,p,o],I,J,K,squeeze(bottom_idx)));
output = permute(shiftdim(output,-1),[2 1 3 4]);
 data = cat(2,data,output);

[Matt J]

No, I see nothing untoward in the first 4 lines, at least.

[meghannmarie]

I was super tired yesterday, probably something earlier in code I screwed up...