Наиболее эффективный способ применения функции к массиву NumPy

В более новой версии (я использую 1.13) библиотеки NumPy вы можете просто вызвать функцию, передав в неё массив NumPy, который вы написали для скалярных типов. NumPy автоматически применит вызов функции к каждому элементу массива и вернёт вам новый массив.

Например:

>>> import numpy as np
>>> squarer = lambda t: t ** 2
>>> x = np.array([1, 2, 3, 4, 5])
>>> squarer(x)
array([ 1,  4,  9, 16, 25])

Таким образом, ваша функция squarer будет применена ко всем элементам массива x, и результат будет представлен в виде нового массива NumPy.

Все вышеуказанные ответы хорошо сопоставимы, но если вам нужно использовать пользовательскую функцию для отображения, и у вас есть numpy.ndarray, и вам необходимо сохранить форму массива, то вот полезное сравнение.

Я протестировал только два подхода, но оба сохраняют форму ndarray. Я использовал массив с 1 миллионом элементов для сравнения. В данном примере используется функция возведения в квадрат, которая также встроена в NumPy и обеспечивает отличную производительность. Если вам понадобится что-то другое, вы можете использовать функцию по вашему выбору.

import numpy, time

def timeit():
    y = numpy.arange(1000000)
    now = time.time()
    
    # Вариант с использованием спискового включения
    numpy.array([x * x for x in y.reshape(-1)]).reshape(y.shape)        
    print(time.time() - now)
    
    now = time.time()
    
    # Вариант с использованием numpy.fromiter
    numpy.fromiter((x * x for x in y.reshape(-1)), y.dtype).reshape(y.shape)
    print(time.time() - now)
    
    now = time.time()
    
    # Вариант с использованием встроенной функции
    numpy.square(y)  
    print(time.time() - now)

Вывод

>>> timeit()
1.162431240081787    # списковое включение и дальнейшее создание numpy массива
1.0775556564331055   # использование numpy.fromiter
0.002948284149169922 # использование встроенной функции

Как видно из вывода, numpy.fromiter показывает хорошую производительность по сравнению с простым подходом, но если вы можете воспользоваться встроенной функцией, то всегда лучше использовать её.

Используйте функцию numpy.fromfunction(), чтобы создать массив, заполняя его значениями, генерируемыми заданной функцией. Аргументы этой функции включают function, которая определяет, как будут вычисляться значения массива, shape, которая задает форму создаваемого массива, и дополнительные параметры **kwargs, которые могут быть переданы в функцию.

Пример использования:

import numpy as np

def func(i, j):
    return i + j

# Создаем массив 3x3, где значения равны сумме индексов
array = np.fromfunction(func, (3, 3), dtype=int)
print(array)

В этом примере func принимает два индекса (i и j), и создает 2D массив размером 3x3, где каждое значение равно сумме индексов. Полученный массив будет выглядеть так:

[[0 1 2]
 [1 2 3]
 [2 3 4]]

Вы можете найти больше информации о numpy.fromfunction в документации NumPy.

Подход, который ещё не так распространен, но легко реализуется и работает быстро, — это Zig и Ziglang.

Установите пакет с помощью команды:

pip install ziglang

Создайте файл zinptest.zig

export fn npprod(inarray: usize, outarray: usize, lenarray: usize) void {
    const inarraydata: [*]u64 = @ptrFromInt(inarray);
    var outarraydata: [*]u64 = @ptrFromInt(outarray);
    for (0..lenarray) |i| {
        outarraydata[i] = inarraydata[i] * inarraydata[i];
    }
}

export fn npprod2(inarray: usize, outarray: usize, lenarray: usize) void {
    const inarraydata: [*]u64 = @ptrFromInt(inarray);
    var outarraydata: [*]u64 = @ptrFromInt(outarray);
    for (0..lenarray) |i| {
        outarraydata[i] = inarraydata[i] + 2 * inarraydata[i] * inarraydata[i] + 4 * inarraydata[i] * inarraydata[i] * inarraydata[i];
    }
}

Напишите обертку

import subprocess
import os
import sys
import ctypes
import numpy as np
# pip install ziglang

def compile_dll():
    subprocess.run(
        [
            sys.executable,
            "-m",
            "ziglang",
            "build-lib",
            "zinptest.zig",
            "-dynamic",
            "-O",
            "ReleaseFast",
        ],
        shell=True,
        env=os.environ,
        cwd=this_folder,
    )

def zigproduct(a):
    out = np.empty_like(a)
    inaddress = a.ctypes._arr.__array_interface__["data"][0] # сырой адрес памяти, будьте внимательны! (данные должны быть правильно выровнены)
    outaddress = out.ctypes._arr.__array_interface__["data"][0]
    lena = np.prod(a.shape)
    zigprod(inaddress, outaddress, lena)
    return out

def zigproduct2(a):
    out = np.empty_like(a)
    inaddress = a.ctypes._arr.__array_interface__["data"][0]
    outaddress = out.ctypes._arr.__array_interface__["data"][0]
    lena = np.prod(a.shape)
    zigprod2(inaddress, outaddress, lena)
    return out

def f(x):
    return x + 2 * x * x + 4 * x * x * x

this_folder = os.path.dirname(__file__)
zigdll = os.path.normpath(os.path.join(this_folder, "zinptest.dll"))
if not os.path.exists(zigdll):
    compile_dll()

# Первая функция Zig (**2)
zigdllloaded = ctypes.cdll.LoadLibrary(zigdll)
zigprod = zigdllloaded.npprod
zigprod.argtypes = [ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64]
zigprod.restype = None

# Вторая функция Zig (x + 2 * x * x + 4 * x * x * x)
zigprod2 = zigdllloaded.npprod2
zigprod2.argtypes = [ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64]
zigprod2.restype = None

a = np.arange(100000000, dtype=np.uint64)

Наслаждайтесь результатами

# Встроенные функции Numpy очень оптимизированы, здесь не так много, что можно улучшить
# In [3]: %timeit a**2
# 213 ms ± 1.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# In [4]: %timeit zigproduct(a)
# 215 ms ± 1.01 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Теперь используем функцию от @ead  return x+2*x*x+4*x*x*x

# In [1]: %timeit zigproduct2(a)
# 206 ms ± 606 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# In [2]: %timeit f(a)
# 1.66 s ± 17.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# zigproduct2 работает в 8 раз быстрее, даже быстрее, чем zigproduct (Бог знает, почему?!?)

# Результаты совпадают
# In[3]: np.all(zigproduct2(a) == f(a))
# Out[3]: True

Заключение

Язык с великолепным будущим, если вы выполняете действия с меньшим количеством операций записи в память (например, np.where / np.argwhere), Zig будет ещё быстрее!

Дополнительные преимущества

Компилятор Zig также компилирует C-код без изменения команды, которую вы видели выше: просто создайте файл с именем "zinptest.c" и используйте его в качестве замены (216 мс для npprod2). Работать непосредственно с C-кодом, на мой взгляд, проще, чем расшифровывать странные сообщения об ошибках Numba. Если вам нужно работать с размерами, используйте шаги или атрибут a.shape и операции / и %.

void npprod(unsigned long long inarray, unsigned long long outarray,
            unsigned long long lenarray) {
  unsigned long long *inarraydata = (unsigned long long *)inarray;
  unsigned long long *outarraydata = (unsigned long long *)outarray;
  for (int i = 0; i < lenarray; i++) {
    outarraydata[i] = inarraydata[i] * inarraydata[i];
  }
}

void npprod2(unsigned long long inarray, unsigned long long outarray,
             unsigned long long lenarray) {
  unsigned long long *inarraydata = (unsigned long long *)inarray;
  unsigned long long *outarraydata = (unsigned long long *)outarray;
  for (int i = 0; i < lenarray; i++) {
    outarraydata[i] = inarraydata[i] + 2 * inarraydata[i] * inarraydata[i] +
                      4 * inarraydata[i] * inarraydata[i] * inarraydata[i];
  }
}

Интересно, что наиболее очевидное решение — использование C, если вы хотите добиться скорости C, не было упомянуто последние 8 лет. C не является магией, и чаще всего, особенно при работе с массивами NumPy (100% C массивы), это проще и быстрее. Компилировать C-код стало проще, чем когда-либо, с использованием Zig.