python - 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

如帝斯曼所说，数量级更小。甚至单独对数组进行切片的工作量也更大:

>>> 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