Как сделать функцию на основе последовательности быстрее в NumPy?

Question

С учетом функции ниже:

import numpy as np

a = np.ones(16).reshape(4,4)

def fn(a):
    b = np.array(a)
    for i in range(b.shape[0]):
        for j in range(b.shape[1] - 1):
            b[i][j+1] += b[i][j]
    return b

print(fn(a))

То есть, для общей функции, которая вычисляет t+1 на основе t в массиве, могу ли я сделать это быстрее? Я знаю, что np.vectorize, но не подходит для этого случая.

Answer 1

То, что вы ищете, называется накопить вот пример:

import numpy as np
from itertools import accumulate

def fn(a):
    acc = accumulate(a, lambda prev, row: prev + row)
    return np.array(list(acc))

a = np.arange(16).reshape(4, 4)
print(fn(a))
# [[ 0  1  2  3]
#  [ 4  6  8 10]
#  [12 15 18 21]
#  [24 28 32 36]]

В numpy нет оптимизированной функции накопления, потому что на самом деле невозможно записать накопление так, чтобы оно было как быстродействующим, так и общим. Реализация Python является общей, но она будет работать во многом как код lok, написанный вручную.

Чтобы получить оптимальную производительность, вам, вероятно, понадобится найти или написать низкоуровневую реализацию конкретной функции накопления, которая вам нужна. Вы уже упоминали о Numba, и вы также можете заглянуть в Cython.

Answer 2

Можно уменьшить два цикла for до одного цикла for, добавив при этом немного дополнительных затрат на копирование.

In [86]: a 
Out[86]: 
array([[1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.]])

In [87]: b = a.copy() 

In [88]: for col in range(b.shape[1]-1): 
    ...:     b[:, col+1] = np.sum(a[:, :col+2], axis=1) 

In [89]: b
Out[89]: 
array([[1., 2., 3., 4.],
       [1., 2., 3., 4.],
       [1., 2., 3., 4.],
       [1., 2., 3., 4.]])

Чтобы это работало для универсальной функции , вы можете найти эквивалентную функцию в numpy или реализовать ее с помощью numpy операций (векторизованной).В приведенном вами примере я просто использовал numpy.sum(), который делает эту работу за нас.

С точки зрения производительности этот подход был бы намного лучше, чем работа с использованием двух циклов for на уровне индексов, особеннодля больших массивов.В подходе, который я использовал выше, мы работаем со срезами столбцов.

---

Вот моменты времени, которые предполагают более чем 3X ускорение по сравнению с собственной реализацией Python.

---

Родной Python:

def fn(a):
    b = np.array(a)
    for i in range(b.shape[0]):
        for j in range(b.shape[1] - 1):
            b[i][j+1] += b[i][j]
    return b

---

Немного векторизовано:

In [104]: def slightly_vectorized(b): 
     ...:     for col in range(b.shape[1]-1): 
     ...:         b[:, col+1] = np.sum(a[:, :col+2], axis=1) 
     ...:     return b 

---

In [100]: a = np.ones(625).reshape(25, 25) 

In [101]: %timeit fn(a) 
303 µs ± 2.05 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [102]: b = a.copy() 

In [103]: %timeit slightly_vectorized(b) 
99.8 µs ± 501 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 3

Вы можете использовать cumsum Я думаю, это было бы полезно.

import numpy as np
import pandas as pd
a = np.ones(16).reshape(4,4)
df =pd.DataFrame(a)
df.cumsum(axis=1)

Или вы можете использовать np.cumsum():

np.cumsum(a,axis=1)