Эффективная итерация по трехмерному массиву?

Question

Я использую Python и Numpy для анализа данных.

У меня есть большая 3D-матрица (NxNxN), где каждая ячейка снова является матрицей, на этот раз матрицей 3x3. Называя матрицу data, она выглядит так:

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

Мне нужно найти собственные значения всех матриц 3х3, и для этого я использую подпрограмму 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

Answer 2

Я уверен, что в NumPy есть хороший способ сделать это, но в целом itertools.product быстрее, чем вложенные циклы в диапазонах.

from itertools import product

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

Answer 3

, поскольку все вычисления независимы, вы можете использовать многопроцессорный модуль для ускорения расчета, если у вас многоядерный процессор.