高效的数组整形和分块操作

Question

我想以更高效的方式进行以下一组操作：

import numpy as np
from time import time

n1  = 5
n2  = 800
n3  = 1000  

a    = np.random.rand( n1,n2,n3 )

ts = time()

a_new  = np.array( [ np.repeat( a[np.newaxis,ll,:,:], n1, axis=0  ) for ll in range(n1)] )

print(time()-ts)
ts = time()
d      = n1**2
a_new  = np.reshape( a_new, ( d, n2, n3  ) )    
a_new  = np.repeat( a_new[np.newaxis,:,:,:], n1, axis=0 )

print(time()-ts)
ts = time()
d2     = d*n1      
a_new  = np.reshape( a_new, ( d2, n2, n3 ))
a_new  = np.reshape( a_new, ( d2*n2 , n3 ) )

print(time()-ts)

当 n1、n2、n3 变得非常大时，这变得有点低效。最好的方法是什么？

Answer 1

您的代码，跟踪结果的大小：

In [22]: n1 = 5
    ...: n2 = 800
    ...: n3 = 1000
    ...: 
    ...: a = np.random.rand(n1, n2, n3)
In [23]: a.shape
Out[23]: (5, 800, 1000)
In [24]: a_new = np.array([np.repeat(a[np.newaxis, ll, :, :], n1, axis=0) for ll in range(n1)])
    ...: 
In [25]: a_new.shape
Out[25]: (5, 5, 800, 1000)
In [26]: d = n1**2
    ...: a1 = np.reshape(a_new, (d, n2, n3))
    ...: a1 = np.repeat(a1[np.newaxis, :, :, :], n1, axis=0)
    ...: 
In [27]: a1.shape
Out[27]: (5, 25, 800, 1000)
In [28]: d2 = d * n1
    ...: a2 = np.reshape(a1, (d2, n2, n3))
    ...: a2 = np.reshape(a2, (d2 * n2, n3))
    ...: 
    ...: 
In [29]: a2.shape
Out[29]: (100000, 1000)

一个小问题，但是 a2 中的双 reshape 不是必需的。您可以直接转到最后一个形状。在任何情况下，重塑都很快（除非它必须强制复制，如转置之后）。

有时

In [34]: %timeit a_new = np.array([np.repeat(a[np.newaxis, ll, :, :], n1, axis=0) for ll in range(n1)])
263 ms ± 11.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [36]: a1  = np.reshape( a_new, ( d, n2, n3  ) )
In [37]: %timeit np.repeat( a_new[np.newaxis,:,:,:], n1, axis=0 )
912 ms ± 134 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

所以是的，重复确实花费了最多的时间。但它将大型阵列的大小增加了 5 倍。 repeat 相对较快 - 对于构成大数组的东西。

您的 a_new 列表理解不是必需的：

In [41]: np.repeat(a[:, None, :, :], n1, axis=1).shape
Out[41]: (5, 5, 800, 1000)
In [42]: np.allclose(np.repeat(a[:, None, :, :], n1, axis=1), a_new)
Out[42]: True
In [43]: %timeit np.repeat(a[:, None, :, :], n1, axis=1)
210 ms ± 23.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

节省的时间并不多，因为 n1 循环只有 5 倍。

请注意，以下 repeat 大约需要 5 倍。 a_new 是 a 的 5 倍，a1 是 a 的 25 倍。

我们可以在一行中完成所有工作

In [44]: a4 = np.repeat(np.repeat(a[:,None,:,:], n1, axis=1)[None],n1, axis=0).reshape(-1,n3)
In [45]: a4.shape
Out[45]: (100000, 1000)
In [46]: np.allclose(a2, a4)
Out[46]: True

时间变化不大。有趣的是 [46] allclose 花费了最明显的时间。

双重重复可以写成tile:

a5 = np.tile(a[None,None,:,:],(5,5,1,1)).reshape(-1,n3)

它不会节省时间，因为 tile 在每个维度上重复使用 repeat。

甚至

a6 = np.repeat(a[None,:,:], n1*n1,axis=0).reshape(-1,n3)

我不打算计时或测试它，因为创建多个这种大小的数组会超出我的内存限制。

高效的数组整形和分块操作

Efficient array reshaping and tile operations

python

numpy

numpy-ndarray