Функция 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]]))

ОТВЕТЫ

Ответ 1

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

Yay!

Итак, мы просто запускаем 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 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)

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

Ответ 2

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

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

Edit:

Следующий код иллюстрирует, как можно векторизовать вашу функцию образца. Хотя он не является полным (блок-матрица и поиск собственных значений), он должен предоставить вам некоторые основные идеи, как это можно сделать. Посмотрите на каждую матрицу и подматрицу в функции, чтобы увидеть, как можно настроить такой расчет. Кроме того, я использовал плотные матрицы, которые, скорее всего, не будут вписываться в память при использовании миллионов элементов в 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)