优化从各种数组到数组的赋值 - NumPy

Question

我有四个方形矩阵,尺寸为3Nx3N,称为A,B,C和D.

我想将它们组合在一个矩阵中.带有for循环的代码如下:

import numpy
N = 3
A = numpy.random.random((3*N, 3*N))
B = numpy.random.random((3*N, 3*N))
C = numpy.random.random((3*N, 3*N))
D = numpy.random.random((3*N, 3*N))

final = numpy.zeros((6*N, 6*N))

for i in range(N):
    for j in range(N):
        for k in range(3):
            for l in range(3):
                final[6*i + k][6*j + l] = A[3*i+k][3*j+l]
                final[6*i + k + 3][6*j + l + 3] = B[3*i+k][3*j+l]
                final[6*i + k + 3][6*j + l] = C[3*i+k][3*j+l]
                final[6*i + k][6*j + l + 3] = D[3*i+k][3*j+l]

是否有可能以numpythonic方式编写前面的for循环？

Answer 1

array-slicing练习多维张量/数组的好问题！

我们将输出数组初始化为多维6D数组，然后简单地对其进行切片并将要重塑的四个数组分配为4D数组。目的是避免任何堆叠/连接，因为这些会很昂贵，特别是在处理大型数组时，而不是对输入数组进行整形，而这只是视图。

这是实现 -

out = np.zeros((N,2,3,N,2,3),dtype=A.dtype)
out[:,0,:,:,0,:] = A.reshape(N,3,N,3)
out[:,0,:,:,1,:] = D.reshape(N,3,N,3)
out[:,1,:,:,0,:] = C.reshape(N,3,N,3)
out[:,1,:,:,1,:] = B.reshape(N,3,N,3)
out.shape = (6*N,6*N)

为了多解释一下，我们有：

            |------------------------ Axes for selecting A, B, C, D
np.zeros((N,2,3,N,2,3),dtype=A.dtype)
                  |------------------------- Axes for selecting A, B, C, D

因此，这两个长度轴（第二轴和第五轴）(2x2) = 4用于在四个输入之间进行选择。

运行时测试

方法 -

def original_app(A, B, C, D):
    final = np.zeros((6*N,6*N),dtype=A.dtype)
    for i in range(N):
        for j in range(N):
            for k in range(3):
                for l in range(3):
                    final[6*i + k][6*j + l] = A[3*i+k][3*j+l]
                    final[6*i + k + 3][6*j + l + 3] = B[3*i+k][3*j+l]
                    final[6*i + k + 3][6*j + l] = C[3*i+k][3*j+l]
                    final[6*i + k][6*j + l + 3] = D[3*i+k][3*j+l]
    return final

def slicing_app(A, B, C, D):
    out = np.zeros((N,2,3,N,2,3),dtype=A.dtype)
    out[:,0,:,:,0,:] = A.reshape(N,3,N,3)
    out[:,0,:,:,1,:] = D.reshape(N,3,N,3)
    out[:,1,:,:,0,:] = C.reshape(N,3,N,3)
    out[:,1,:,:,1,:] = B.reshape(N,3,N,3)
    return out.reshape(6*N,6*N)

时间安排和验证 -

In [147]: # Setup input arrays
     ...: N = 200
     ...: A = np.random.randint(11,99,(3*N,3*N))
     ...: B = np.random.randint(11,99,(3*N,3*N))
     ...: C = np.random.randint(11,99,(3*N,3*N))
     ...: D = np.random.randint(11,99,(3*N,3*N))
     ...: 

In [148]: np.allclose(slicing_app(A, B, C, D), original_app(A, B, C, D))
Out[148]: True

In [149]: %timeit original_app(A, B, C, D)
1 loops, best of 3: 1.63 s per loop

In [150]: %timeit slicing_app(A, B, C, D)
100 loops, best of 3: 9.26 ms per loop