Most efficient way to map function over numpy array

Какой наиболее эффективный способ сопоставления функции с массивом numpy? В настоящее время я занимаюсь:

import numpy as np 

x = np.array([1, 2, 3, 4, 5])

# Obtain array of square of each element in x
squarer = lambda t: t ** 2
squares = np.array([squarer(xi) for xi in x])

Однако это, вероятно, очень неэффективно, поскольку я использую понимание списка для построения нового массива в виде списка Python, прежде чем преобразовывать его обратно в массив numpy. Можем ли мы сделать лучше?

Я протестировал все предложенные методы плюс np.array(list(map(f, x))) с perfplot (мой небольшой проект).

Сообщение № 1: Если вы можете использовать собственные функции numpy, сделайте это.

Если функция, которую вы пытаетесь векторизовать, уже векторизована (как в x**2 примере в исходном сообщении), использовать это намного быстрее, чем что-либо другое (обратите внимание на логарифмический масштаб):

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

Код для воспроизведения графиков:

import numpy as np
import perfplot
import math


def f(x):
    # return math.sqrt(x)
    return np.sqrt(x)


vf = np.vectorize(f)


def array_for(x):
    return np.array([f(xi) for xi in x])


def array_map(x):
    return np.array(list(map(f, x)))


def fromiter(x):
    return np.fromiter((f(xi) for xi in x), x.dtype)


def vectorize(x):
    return np.vectorize(f)(x)


def vectorize_without_init(x):
    return vf(x)


b = perfplot.bench(
    setup=np.random.rand,
    n_range=[2 ** k for k in range(20)],
    kernels=[
        f,
        array_for,
        array_map,
        fromiter,
        vectorize,
        vectorize_without_init,
    ],
    xlabel="len(x)",
)
b.save("out1.svg")
b.show()

Использование numpy.vectorize:

import numpy as np
x = np.array([1, 2, 3, 4, 5])
squarer = lambda t: t ** 2
vfunc = np.vectorize(squarer)
vfunc(x)

# Output: array([ 1,  4,  9, 16, 25])

TL; DR

Как отметил @User, "прямой" метод применения функции всегда является самым быстрым и простым способом сопоставления функции с массивами Numpy:

import numpy as np
x = np.array([1, 2, 3, 4, 5])
f = lambda x: x ** 2
squares = f(x)

Обычно избегайте np.vectorize, поскольку он работает недостаточно хорошо и имеет (или имел) ряд проблем. Если вы работаете с другими типами данных, возможно, вам захочется изучить другие методы, показанные ниже.

Сравнение методов

Вот несколько простых тестов для сравнения трех методов сопоставления функции, в этом примере используются Python 3.6 и NumPy 1.15.4. Во-первых, функции настройки для тестирования:

import timeit
import numpy as np

f = lambda x: x ** 2
vf = np.vectorize(f)

def test_array(x, n):
    t = timeit.timeit(
        'np.array([f(xi) for xi in x])',
        'from __main__ import np, x, f', number=n)
    print('array: {0:.3f}'.format(t))

def test_fromiter(x, n):
    t = timeit.timeit(
        'np.fromiter((f(xi) for xi in x), x.dtype, count=len(x))',
        'from __main__ import np, x, f', number=n)
    print('fromiter: {0:.3f}'.format(t))

def test_direct(x, n):
    t = timeit.timeit(
        'f(x)',
        'from __main__ import x, f', number=n)
    print('direct: {0:.3f}'.format(t))

def test_vectorized(x, n):
    t = timeit.timeit(
        'vf(x)',
        'from __main__ import x, vf', number=n)
    print('vectorized: {0:.3f}'.format(t))

Тестирование с использованием пяти элементов (отсортированных от самых быстрых к самым медленным):

x = np.array([1, 2, 3, 4, 5])
n = 100000
test_direct(x, n)      # 0.265
test_fromiter(x, n)    # 0.479
test_array(x, n)       # 0.865
test_vectorized(x, n)  # 2.906

Со 100 элементами:

x = np.arange(100)
n = 10000
test_direct(x, n)      # 0.030
test_array(x, n)       # 0.501
test_vectorized(x, n)  # 0.670
test_fromiter(x, n)    # 0.883

И с 1000 элементами массива или более:

x = np.arange(1000)
n = 1000
test_direct(x, n)      # 0.007
test_fromiter(x, n)    # 0.479
test_array(x, n)       # 0.516
test_vectorized(x, n)  # 0.945

Разные версии Python / NumPy и оптимизация компилятора будут иметь разные результаты, поэтому проведите аналогичный тест для вашей среды.

Существуют numexpr, numba и cython, цель этого ответа - учесть эти возможности.

Но сначала давайте констатируем очевидное: независимо от того, как вы сопоставляете Python-функцию с numpy-массивом, она остается функцией Python, что означает для каждой оценки:

• элемент numpy-массива должен быть преобразован в Python-объект (например, a Float).


• все вычисления выполняются с помощью Python-объектов, что означает накладные расходы на интерпретатор, динамическую отправку и неизменяемые объекты.

Итак, какой механизм используется для фактического перебора массива, не играет большой роли из-за упомянутых выше накладных расходов - он работает намного медленнее, чем при использовании встроенной функциональности numpy.

Давайте посмотрим на следующий пример:

# numpy-functionality
def f(x):
    return x+2*x*x+4*x*x*x

# python-function as ufunc
import numpy as np
vf=np.vectorize(f)
vf.__name__="vf"

np.vectorize выбран как представитель подходов класса функций на чистом python. Используя perfplot (см. Код в приложении к этому ответу), мы получаем следующее время выполнения:

Мы можем видеть, что numpy-подход в 10-100 раз быстрее, чем чистая версия python. Снижение производительности при больших размерах массива, вероятно, связано с тем, что данные больше не помещаются в кэш.

Стоит также упомянуть, что vectorize также использует много памяти, поэтому часто использование памяти является бутылочным горлышком (см. Связанный SO-вопрос). Также обратите внимание, что в документации numpy по np.vectorize говорится, что это "предоставляется в первую очередь для удобства, а не для производительности".

Когда требуется производительность, следует использовать другие инструменты, помимо написания расширения C с нуля, существуют следующие возможности:

Часто приходится слышать, что производительность numpy настолько высока, насколько это возможно, потому что под капотом чистый C. Тем не менее, есть много возможностей для улучшения!

Векторизованная версия numpy использует много дополнительной памяти и операций доступа к памяти. Библиотека Numexp пытается разбить numpy-массивы на плитки и, таким образом, улучшить использование кэша:

# less cache misses than numpy-functionality
import numexpr as ne
def ne_f(x):
    return ne.evaluate("x+2*x*x+4*x*x*x")

Приводит к следующему сравнению:

Я не могу объяснить все на приведенном выше графике: мы можем видеть большие накладные расходы для библиотеки numexpr в начале, но поскольку она лучше использует кэш, она примерно в 10 раз быстрее для больших массивов!

Другой подход заключается в jit-компиляции функции и, таким образом, получении реальной UFunc на чистом C. Это подход numba.:

# runtime generated C-function as ufunc
import numba as nb
@nb.vectorize(target="cpu")
def nb_vf(x):
    return x+2*x*x+4*x*x*x

Это в 10 раз быстрее, чем исходный numpy-подход:

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

@nb.njit(parallel=True)
def nb_par_jitf(x):
    y=np.empty(x.shape)
    for i in nb.prange(len(x)):
        y[i]=x[i]+2*x[i]*x[i]+4*x[i]*x[i]*x[i]
    return y

Как и ожидалось, параллельная функция работает медленнее при меньших входных данных, но быстрее (почти в 2 раза) при больших размерах:

В то время как numba специализируется на оптимизации операций с массивами numpy, Cython является более общим инструментом. Сложнее добиться той же производительности, что и с numba - часто это зависит от llvm (numba) или локального компилятора (gcc / MSVC):

%%cython -c=/openmp -a
import numpy as np
import cython

#single core:
@cython.boundscheck(False) 
@cython.wraparound(False) 
def cy_f(double[::1] x):
    y_out=np.empty(len(x))
    cdef Py_ssize_t i
    cdef double[::1] y=y_out
    for i in range(len(x)):
        y[i] = x[i]+2*x[i]*x[i]+4*x[i]*x[i]*x[i]
    return y_out

#parallel:
from cython.parallel import prange
@cython.boundscheck(False) 
@cython.wraparound(False)  
def cy_par_f(double[::1] x):
    y_out=np.empty(len(x))
    cdef double[::1] y=y_out
    cdef Py_ssize_t i
    cdef Py_ssize_t n = len(x)
    for i in prange(n, nogil=True):
        y[i] = x[i]+2*x[i]*x[i]+4*x[i]*x[i]*x[i]
    return y_out

Cython приводит к несколько более медленным функциям:

Заключение

Очевидно, что тестирование только для одной функции ничего не доказывает. Также следует иметь в виду, что для выбранной функции - например, пропускная способность памяти была бутылочным горлышком для размеров, превышающих 10 ^ 5 элементов - таким образом, у нас была одинаковая производительность для numba, numexpr и cython в этом регионе.

В конце концов, окончательный ответ зависит от типа функции, аппаратного обеспечения, дистрибутива Python и других факторов. Например, Anaconda-distribution использует Intel VML для функций numpy и, таким образом, легко превосходит numba (если только он не использует SVML, см. Этот SO-пост) для трансцендентных функций, таких как exp, sin, cos и подобных - см., например, Следующий SO-пост .

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

Построение графика времени выполнения с помощью perfplot-package:

import perfplot
perfplot.show(
    setup=lambda n: np.random.rand(n),
    n_range=[2**k for k in range(0,24)],
    kernels=[
        f, 
        vf,
        ne_f, 
        nb_vf, nb_par_jitf,
        cy_f, cy_par_f,
        ],
    logx=True,
    logy=True,
    xlabel='len(x)'
    )