Vectorising 5 nested FOR loops

Question

I am writing a program in MATLAB as a part of my project based on DFT.

Let the N x N data matrix be X and the corresponding DFT matrix be Y, then the DFT coefficients can be expressed as

 Y(k1,k2) = ∑(n1=0:N-1)∑(n2=0:N-1)[X(n1,n2)*(WN^(n1k1+n2k2))]               (1)

            0≤k1,k2≤N-1
            Where WN^k=e^((-j2πk)/N)

Since the twiddle factor WN is periodic, (1) can be expressed as

Y(k1,k2)=∑(n1=0:N-1)∑(n1=0:N-1)[X(n1,n2)*(WN^([(n1k1+n2k2)mod N) ]          (2)

The exponent ((n1k1 +n2k2)) N = p is satisfied by a set of (n1,n2) for a given (k1,k2). Hence, by grouping such data and applying the property that WN^(p+N /2) = -(WN^P),

(2) can be expressed as

 Y(k1,k2)= ∑(p=0:M-1)[Y(k1,k2,p)*(WN^p)]                                    (3)

Where

 Y(k1,k2,p)= ∑(∀(n1,n2)|z=p)X(n1,n2) - ∑(∀(n1,n2)|z=p+M)X(n1,n2)           (4)

 z=[(n1k1+n2k2)mod N]                                                       (5)

I am coding a program to find Y(k1,k2,p).ie I need to create slices of 2d matrices(ie a 3D matrix in which each slice is a 2D matrix )from a given 2D square matrix (which is the matrix X)..Dimensions of X can be upto 512.

Based on the above equations,I have written a code as follows.I need to vectorise it.

N=size(X,1);
M=N/2;
Y(1:N,1:N,1:M)=0;
for k1 = 1:N
for k2 = 1:N
   for p= 1:M
       for n1=1:N
           for n2=1:N
               N1=n1-1; N2=n2-1; P=p-1; K1=k1-1; K2=k2-1;
             z=mod((N1*K1+N2*K2),N); 
             if (z==P)                                       
               Y(k1,k2,p)= Y(k1,k2,p)+ X(n1,n2);
           elsif (z==(P+M))
               Y(k1,k2,p)= Y(k1,k2,p)- X(n1,n2);
            end
           end
       end
    end
   end

As there is 5 FOR loops, the execution time is very large for large dimensions of N. Hence please provide me a solution for eliminating the FOR loops and vectorising the code..I need to make the code execute in maximum speed...Thanks Again..

Answer 1

Here is a first hint to vectorize the most inner loop.

From your code, we can notice that n1, N1, P, K1 and K2 are constant in this loop. So we can rewrite z as a mask vector as follows:

z = mod(N1*K1+K2*(0:N-1));

Then your if-statement is equivalent to adding the sum of all elements in X so that z==P minus the sum of all elements in X so that z==P+M. Rewriting this is straightforward:

Y(k1,k2,p)= Y(k1,k2,p)+sum(X(n1,z==P))-sum(X(n1,z==P+M));

So your program can be first written as follows:

N=size(X,1);
M=N/2;
Y(1:N,1:N,1:M)=0;
for k1 = 1:N
  for k2 = 1:N
    for p= 1:M
      for n1=1:N
        N1=n1-1; P=p-1; K1=k1-1; K2=k2-1;
        z=mod(N1*K1+K2*(0:N-1),N); 
        Y(k1,k2,p) = sum(X(n1,z==P))-sum(X(n1,z==P+M));
      end
    end
  end
end

Then you can do the same thing with n1; for that, you need to construct a 2D array for z, such as:

z = mod(K1*repmat(0:N-1,N,1)+K2*repmat((0:N-1).',1,N));

Notice that size(z)==size(X).Then the 2D sum for Y becomes:

Y(k1,k2,p) = Y(k1,k2,p)+sum(X(z==P))-sum(X(z==P+M));

The += operation is here no longer needed, since you access only once to each element of Y:

Y(k1,k2,p)= sum(X(n1,z==P))-sum(X(n1,z==P+M));

And so we discard one more loop:

N=size(X,1);
M=N/2;
Y(1:N,1:N,1:M)=0;
for k1 = 1:N
  for k2 = 1:N
    for p= 1:M
      P=p-1; K1=k1-1; K2=k2-1;
      z = mod(K1*repmat(0:N-1,N,1)+K2*repmat((0:N-1).',1,N)); 
      Y(k1,k2,p) = sum(X(z==P))-sum(X(z==P+M));
    end
  end
end

Concerning the other loops, I don't think it worths it to vectorize them, as you have to build a 5D array, which could be very huge in memory. My advise is to keep z as a 2D array, as it is of the size of X. If it does not fit well in memory, just vectorize the most inner loop.