向量化5个嵌套的FOR循环

Question

我正在用MATLAB编写一个程序，作为我基于DFT的项目的一部分。

设N x N数据矩阵为X，相应的DFT矩阵为Y，则DFT系数可表示为

 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)

由于旋转因子WN是周期性的，因此(1)可以表示为

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

对于给定的(k1,k2)，指数((n1k1 +n2k2)) N = p由一组(n1,n2)满足。因此，通过对这些数据进行分组并应用WN^(p+N /2) = -(WN^P)的属性，

(2)可以表示为

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

哪里

 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)

我正在编写一个程序来寻找Y(k1,k2,p)，我需要从给定的2d矩阵(这是矩阵X)创建2D矩阵的切片(即，每个切片是2D矩阵的3D矩阵) X的..Dimensions可以高达512。

基于上面的等式，我写了一个代码，如下所示。我需要对它进行矢量化。

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

由于有5个FOR循环，对于N的大维度，执行时间非常长。因此，请为我提供一个消除for循环和向量化代码的解决方案。我需要使代码再次以最大speed...Thanks执行。

Answer 1

这里是向量化最内层循环的第一个提示。

从您的代码中，我们可以注意到n1、N1、P、K1和K2在此循环中是常量。因此我们可以将z重写为掩码向量，如下所示：

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

那么if语句等同于将X中所有元素的和相加，使得z==P减去X中所有元素的和，因此z==P+M。重写这段代码很简单：

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

因此，您的程序可以首先按如下方式编写：

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

然后，您可以使用n1执行相同的操作；为此，您需要为z构造一个2D数组，例如：

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

请注意，size(z)==size(X).Then Y的二维总和变成：

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

这里不再需要+=操作，因为您只访问Y的每个元素一次：

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

所以我们又丢弃了一个循环：

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

至于其他循环，我不认为将它们矢量化是值得的，因为你必须构建一个5D数组，这可能会占用非常大的内存。我的建议是保持z为2D数组，因为它的大小是X。如果它在内存中不能很好地适应，只需向量化最内部的循环。