Manipulating Matrices in Matlab

Question

Suppose I have a matrix $A$ of size

$n_1 \times n_2 \times n_3$

Now, I have another matrix $B$ of size $n_1 \times n_2 \times n_3 \times N$ where $N<n_3$

I'd like to create the following matrix $C = [A(m,l,B(m,l,k,:))]_{1 \le m \le n_1, 1 \le l \le n_2, 1 \le k \le n_3}$. What is the most efficient way of doing this without having to create loops?

B essentially tells me which elements to pick out of the 3-rd dimension given an m,l and k. So a loop would go along the lines of:

for m=1:n1; 
    for l=1:n2; 
       for k=1:n3 
          C(m,l,:,k)=A(m,l,B(m,l,k,:)) 
       end 
    end 
end

Thanks!!!

Answer 1

If I got the question right, the problem is given $a_{ijk}$ and $b_{ijkl}$ to form
$$ c_{ijkl} = a_{ijb_{ijlk}} $$ or in matlab notation

C(i,j,k,l) = A(i,j,B(i,j,l,k));

4 nested loops

Without colon notation this would require 4 nested loops over variables i,j,k,l:

C(n1,n2,N,n3) = 0.0;          % equivalent to C = zeros(n1,n2,N,n3);
for i = 1:n1,
  for j = 1:n2,
    for k = 1:N,
      for l = 1:n3,
        C(i,j,k,l) = A(i,j,B(i,j,l,k));
      end,
    end,
  end,
end,

Note the preallocation of array C: having variable sized arrays inside a for loop is normally the main reason for inefficiency. In fact matlab is pretty good (compared to python) in evaluating for loops, provided that the l.h.s. does not change size during the iterations.

2 nested loops

Now let's go for the easy vectorization, eliminating the loops over k and l. The main problem here is that k,l are swapped in C and B so we have to permute the array dimensions:

C(n1,n2,n3,N) = 0.0;
for i = 1:n1,
  for j = 1:n2,
    C(i,j,:) = A(i,j,B(i,j,:));
  end,
end,
C = permute(C, [1 2 4 3]);

(Remember that a trailing colon in an array squeezes all remaining dimension into a single one: size(C(i,j,:)) == [1 1 n3*N] ; correct stride is preserved, so that C(i,j,:) is OK as a l.h.s.)

no loops

Edited: corrected error and inserted final solution.

Now for the hard part of the question: how to eliminate the loops over i and j. The idea is simple: find an integer array indx such that C == A(indx): in fact A(indx) is an array with the same shape of indx and elements taken from A regarded as a single column.

In order to construct this array, we have to remember that matlab follows fortran convention:

size(A) == [n1 n2 n3]
A(i,j,k) == A(i + (j-1)*n1 + (k-1)*n1*n2)

So here is how to construct this matrix with 2 nested loops:

indx(n1,n2,n3,N) = 0.0;
for j=1:n2,
  for i=1:n1,
    indx(i,j,:) = i+(j-1)*n1+(B(i,j,:)-1)*n1*n2;
  end,
end,

or with two instructions

indx = (B-1)*n1*n2;
indx(:) = indx(:) + repmat(1:n1*n2, n3*N)'

The no loops code reads

indx = (B-1)*n1*n2;
indx(:) = indx(:) + repmat(1:n1*n2, 1, n3*N)';
indx = permute(indx, [1 2 4 3]);
C = A(indx);

conclusions

In my opinion the most efficient way of computing C should be the 2 nested loops version. In the no loops version we put a great effort in computing indx instead of letting matlab do his own pointer arithmetic, so I would discourage it, unless the same array B and therefore the same vector indx can be used multiple times with different A arrays.

If this is the case I would still recommend to compute indx by two nested loops, instead of the no loops version. (Here Godric Seer is right: it is almost impossible to understand the logic behind.)

Let me stress out an important point: if array C is pre-allocated, two nested loops will be very fast in matlab, and there will be no substantial advantage in going to a no loops version, unless we have a fast method of computing the indx vector. (E.g. by using integer arithmetic.)

Hope this lengthy reply will be useful.

Edited previous wrong no loops version:

C(n1,n2,n3,N) = 0.0;
indx =  ... % to be defined later
C(:) = A(indx);
C = permute(C, [1 2 4 3]);

Here we have to exploit the fact that C(:) is the multi-dimensional array viewed as a single column, and A(indx) is a single column of elements picked from A. We have to know how pointer arithmetic works in matlab

size(C) == [n1 n2 n3 N]
C(i,j,k,l) == C(i + (j-1)*n1 + (k-1)*n1*n2 + (l-1)*n1*n2*n3 )

Now

C(i,j,k,l) = A(i,j,B(i,j,k,l))

can be written as

C(i,j,k,l) = A(i + (j-1)*n1 + (B(i,j,k,l)-1)*n1*n2 );

or even

C(i,j,:) = A(i + (j-1)*n1 + (B(i,j,:)-1)*n1*n2);

therefore we can compute indx as

indx = [];
for j=1:n2,
  for i=1:n1,
    indx = [ indx, reshape( i+(j-1)*n1+(B(i,j,:)-1)*n1*n2, 1, n3*N ) ];
  end,
end,

With some effort it should be possible to compute indx without loops: however I have no time to derive it now. By the way I fear that it would be neither elegant nor efficient, but who knows for sure!

Answer 2

I am assuming that you want to vectorize the operation for efficiency. While a vectorized assignment may be possible (I personally have tried several things and have failed up to this point) I am going to have to argue that you would be better off not using a vectorized matlab for this application.

The reason for my argument is that the clarity of your code will go out the window. Even with your single line assignment

C(m,l,:,k)=A(m,l,B(m,l,k,:))

takes a moment to wrap your head around, and the complex vector notation required to perform the three loops in a single line would completely obscure the assignment process to anyone else reading the code (this includes you a month from now).

If you must have the efficiency boost you would see from the vectorization, you could likely get better performance and more clarity by utilizing a short fortran or C function. The above loops could be implemented with minimal effort, and both languages would likely beat the vectorized matlab in performance while being called from within your matlab script.

I know this does not directly answer your question, but hopefully it helps in the case that no one here finds an acceptable vectorization for your loops.

Answer 3

Thanks for all your response. I actually worked it out over the weekend but haven't had the chance to check this forum (To be frank, I'm quite surprised by how helpful the people are in this community. I actually just assumed that my questioned would be shrugged off given that no answers were given that night.) In any case, I used the following command, which I think is the same as the first answer, but I haven't had the chance to thoroughly go over the generously long response. Again, MANY thanks to the two responders for your time and insights. I'll read over your suggestions and incorporate them in my code.

A(repmat(reshape(1:n1*n2,n1,n2),[1,1,n3,N])+n1*n2*(B-1);