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

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

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

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

• векторизации

  Хотя это часто бывает трудно, обычно это простое решение. Найдите способ выразить свои вычисления таким образом, чтобы они обобщали элементы, поэтому вы можете работать сразу со всей матрицей. Это приведет к тому, что петли будут выведены из Python и будут использоваться в сильно оптимизированных подпрограммах C и Fortran.

• JITing : Numba и Parakeet , и в меньшей степени PyPy с NumPyPy

  Numba и Parakeet имеют дело с JITing-контурами над структурами данных Numpy, поэтому, если вы встроите цикл в функцию (это может быть функция-обертка), вы можете получить огромное ускорение скорости для почти бесплатного. Это зависит от используемых структур данных.

• Символические оценщики, такие как Theano и numexpr

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

• Расширения 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) 

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

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

Если между элементами массива нет связи, тогда ответ 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)