对3D阵列进行高效迭代？

Question

我正在使用Python和Numpy进行一些数据分析.

我有一个大的3D矩阵(NxNxN),其中每个单元格也是一个矩阵,这次是3x3矩阵.调用矩阵data,它看起来像这样:

data[N,N,N,3,3]  

我需要找到所有3x3矩阵的特征值,为此我使用Numpy的eigvals例程,但需要很长时间才能完成.现在我几乎这样做:

for i in range(N):
    for j in range(N):
        for k in range(N):
            a = np.linalg.eigvals(data[i,j,k,:,:])

对于N = 256,这需要大约一个小时.关于如何提高效率的任何想法？

非常感谢任何建议!

Answer 1

itertools.product从美学角度看,它比嵌套循环更好.但我不认为它会使你的代码更快.我的测试表明,迭代不是你的瓶颈.

>>> bigdata = numpy.arange(256 * 256 * 256 * 3 * 3).reshape(256, 256, 256, 3, 3)
>>> %timeit numpy.linalg.eigvals(bigdata[100, 100, 100, :, :])
10000 loops, best of 3: 52.6 us per loop

所以低估了:

>>> .000052 * 256 * 256 * 256 / 60
14.540253866666665

这在我的电脑上至少需要14分钟,这是非常新的.让我们看看循环需要多长时间......

>>> def just_loops(N):
...     for i in xrange(N):
...         for j in xrange(N):
...             for k in xrange(N):
...                 pass
... 
>>> %timeit just_loops(256)
1 loops, best of 3: 350 ms per loop

DSM表示,数量级更小.即使是单独切割阵列的工作也更为重要:

>>> def slice_loops(N, data):
...     for i in xrange(N):
...         for j in xrange(N):
...             for k in xrange(N):
...                 data[i, j, k, :, :]
... 
>>> %timeit slice_loops(256, bigdata)
1 loops, best of 3: 33.5 s per loop