numpy np.apply_along_axis функция ускоряется?

Question

numpy np.apply_along_axis функция ускоряется?

Функция np.apply_along_axis() кажется очень медленной (без вывода через 15 минут). Есть ли быстрый способ выполнить эту функцию на длинном массиве без распараллеливания операции? Я специально говорю о массивах с миллионами элементов.

Вот пример того, что я пытаюсь сделать. Пожалуйста, игнорируйте упрощенное определение my_func, цель состоит не в том, чтобы умножить массив на 55 (что, конечно, можно сделать на месте), а в качестве иллюстрации. На практике my_func немного сложнее, принимает дополнительные аргументы, и в результате каждый элемент a изменяется по-разному, то есть не просто умножается на 55.

>>> def my_func(a):
...     return a[0]*55
>>> a = np.ones((200000000,1))
>>> np.apply_along_axis(my_func, 1, a)

Редактировать:

a = np.ones((20,1))

def my_func(a, i,j):
...     b = np.zeros((2,2))
...     b[0,0] = a[i]
...     b[1,0] = a[i]
...     b[0,1] = a[i]
...     b[1,1] = a[j]
...     return  linalg.eigh(b)


>>> my_func(a,1,1)
(array([ 0.,  2.]), array([[-0.70710678,  0.70710678],
   [ 0.70710678,  0.70710678]]))

22

python numpy

Источник

user890135 24 май '14 в 19:58

2 ответа

Другие вопросы по тегам python numpy

user1763356 24 май '14 в 20:15 2014-05-24 20:15 · Answer 1 · 2014-05-24 20:15

np.apply_along_axis не для скорости.

Невозможно применить чистую функцию Python к каждому элементу массива Numpy, не вызывая его много раз, если не считать переписывания AST...

К счастью, есть решения:

векторизации
Хотя это часто сложно, обычно это простое решение. Найдите способ выразить свои вычисления так, чтобы они обобщали элементы, чтобы вы могли работать сразу со всей матрицей. Это приведет к тому, что циклы будут выведены из Python и в сильно оптимизированные подпрограммы C и Fortran.
JITing: Numba и Parakeet, и в меньшей степени PyPy с NumPyPy
Numba и Parakeet имеют дело с циклами JITing над структурами данных Numpy, поэтому, если вы встроите цикл в функцию (это может быть функцией-оберткой), вы можете получить огромное повышение скорости практически бесплатно. Это зависит от используемых структур данных.
Символические оценщики, такие как Theano и Numberxpr
Это позволяет вам использовать встроенные языки для выражения вычислений, которые могут оказаться намного быстрее, чем даже векторизованные версии.
Расширения Cython и C
Если все остальное потеряно, вы всегда можете вырыть вручную до C. Cython скрывает много сложности и имеет много прекрасной магии, так что это не всегда так плохо (хотя это помогает узнать, что вы делаете).

Ну вот.

Это моя тестовая "среда" (вам действительно следовало предоставить такую информацию:P):

import itertools
import numpy

a = numpy.arange(200).reshape((200,1)) ** 2

def my_func(a, i,j):
    b = numpy.zeros((2,2))
    b[0,0] = a[i]
    b[1,0] = a[i]
    b[0,1] = a[i]
    b[1,1] = a[j]
    return  numpy.linalg.eigh(b)

eigvals = {}
eigvecs = {}

for i, j in itertools.combinations(range(a.size), 2):
    eigvals[i, j], eigvecs[i, j] = my_func(a,i,j)

Теперь гораздо проще получить все перестановки вместо комбинаций, потому что вы можете просто сделать это:

# All *permutations*, not combinations
indexes = numpy.mgrid[:a.size, :a.size]

Это может показаться расточительным, но перестановок всего в два раза больше, так что это не имеет большого значения.

Поэтому мы хотим использовать эти индексы, чтобы получить соответствующие элементы:

# Remove the extra dimension; it's not wanted here!
subs = a[:,0][indexes]

и тогда мы можем сделать наши матрицы:

target = numpy.array([
    [subs[0], subs[0]],
    [subs[0], subs[1]]
])

Нам нужно, чтобы матрицы были в двух последних измерениях:

target.shape
#>>> (2, 2, 200, 200)

target = numpy.swapaxes(target, 0, 2)
target = numpy.swapaxes(target, 1, 3)

target.shape
#>>> (200, 200, 2, 2)

И мы можем проверить, что это работает:

target[10, 20]
#>>> array([[100, 100],
#>>>        [100, 400]])

Ура!

Итак, мы просто бежим numpy.linalg.eigh:

values, vectors = numpy.linalg.eigh(target)

И посмотрите, это работает!

values[10, 20]
#>>> array([  69.72243623,  430.27756377])

eigvals[10, 20]
#>>> array([  69.72243623,  430.27756377])

Итак, я думаю, что вы можете объединить эти:

numpy.concatenate([values[row, row+1:] for row in range(len(values))])
#>>> array([[  0.00000000e+00,   1.00000000e+00],
#>>>        [  0.00000000e+00,   4.00000000e+00],
#>>>        [  0.00000000e+00,   9.00000000e+00],
#>>>        ..., 
#>>>        [  1.96997462e+02,   7.78160025e+04],
#>>>        [  3.93979696e+02,   7.80160203e+04],
#>>>        [  1.97997475e+02,   7.86070025e+04]])

numpy.concatenate([vectors[row, row+1:] for row in range(len(vectors))])
#>>> array([[[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        [[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        [[ 1.        ,  0.        ],
#>>>         [ 0.        ,  1.        ]],
#>>> 
#>>>        ..., 
#>>>        [[-0.70890372,  0.70530527],
#>>>         [ 0.70530527,  0.70890372]],
#>>> 
#>>>        [[-0.71070503,  0.70349013],
#>>>         [ 0.70349013,  0.71070503]],
#>>> 
#>>>        [[-0.70889463,  0.7053144 ],
#>>>         [ 0.7053144 ,  0.70889463]]])

Также возможно сделать этот конкатенационный цикл сразу после numpy.mgrid вдвое сократить объем работы:

# All *permutations*, not combinations
indexes = numpy.mgrid[:a.size, :a.size]

# Convert to all *combinations* and reduce the dimensionality
indexes = numpy.concatenate([indexes[:, row, row+1:] for row in range(indexes.shape[1])], axis=1)

# Remove the extra dimension; it's not wanted here!
subs = a[:,0][indexes]

target = numpy.array([
    [subs[0], subs[0]],
    [subs[0], subs[1]]
])

target = numpy.rollaxis(target, 2)

values, vectors = numpy.linalg.eigh(target)

Да, последний образец - это все, что тебе нужно.

user1038318 24 май '14 в 20:25 2014-05-24 20:25 · Answer 2 · 2014-05-24 20:25

и в результате каждый элемент изменяется по-разному

Если между элементами массива нет связи, ответ Veedracs суммирует типичные стратегии. Однако чаще всего можно найти схему векторизации, которая значительно ускоряет вычисления. Если вы предоставите соответствующие фрагменты кода самой функции, мы можем дать вам гораздо более полезный ответ.

Редактировать:

Следующий код иллюстрирует, как можно векторизовать вашу примерную функцию. Хотя он не является полным (блочная матрица и поиск собственных значений), он должен дать вам некоторые основные идеи, как можно это сделать. Посмотрите на каждую матрицу и подматрицу в функции, чтобы увидеть, как можно настроить такой расчет. Кроме того, я использовал плотные матрицы, которые, скорее всего, не поместятся в память при использовании миллионов элементов в a и большое количество индексных пар. Но большинство матриц при расчете редки. Таким образом, вы всегда можете преобразовать код для использования разреженных матриц. Функция теперь берет вектор a и вектор индекса pairs,

import numpy as np

def my_func(a,pairs):
    #define mask matrix
    g=np.zeros((4,2))
    g[:3,0]=1
    g[3,1]=1

    # k is the number of index pairs which need calculation
    k=pairs.shape[0]
    h=np.kron(np.eye(k),g)
    b=np.dot(h,a[pairs.ravel()[:2*k]]) # this matrix product generates your matrix b
    b.shape=-1,2

    out = np.zeros((2*k,2*k)) # pre allocate memory of the block diagonal matrix

    # make block diagonal matrix
    for i in xrange(k):
        out[i*2:(i+1)*2, i*2:(i+1)*2] = b[i*2:(i+1)*2,:]

    res = np.linalg.eigh(out) # the eigenvalues of each 2by2 matrix are the same as the ones of one large block diagonal matrix
    # unfortunately eigh sorts the eigenvalues
    # to retrieve the correct pairs of eigenvalues
    # for each submatrix b, one has to inspect res[1] and pick
    # corresponding eigenvalues 
    # I leave that for you, remember out=res[1] diag(res[0]) res[1].T
    return res

#vektor a
a=np.arange(20)
#define index pairs for calculation
pairs=np.asarray([[1,3],[2,7],[1,7],[2,3]])
print my_func(a,pairs)