是否可以使用 numba 输出 2 列并将其用作 df 2 新列

Question

我正在尝试使用 numba 来减少 250k 行的运行时间，这可以使用 df.loc 完成，因为当我使用 df.loc 时，运行会花费很多时间递归输入和输出。

这是我的输入 df

input
1
2
3
4

和所需的输出 'a' 和 'b'

inputa  a    b
1       0    0
2       2    3
3       5    7
4       9    12

基本上，'a'和'b'的初始值都是0。

a = previous value of 'a' + current value of inputa

同时 b = previous value of inputa + a current values.

我现在的代码是这样的。

@jit(nopython=True)
def foo(inputa):
    a = np.zeros(inputa.shape)
    b = np.zeros(inputa.shape)
    a[0] = 0
    b[0] = 0
    for i in range(1, a.shape[0]) :
        a[i] = a[i-1] + inputa[i]
        b[i] = inputa[i-1] + a[i]
    return a, b
 
df[['a','b']] = foo(df['inputa'].values)     
print(df)

但是，我遇到了这个错误

TypingError: cannot determine Numba type of <class 'method'>

如果我正在为 numba 使用 1 列输出，代码工作正常。请参阅下面的代码。

@jit(nopython=True)
def foo(inputa):
    a = np.zeros(inputa.shape)
    #b = np.zeros(inputa.shape)
    a[0] = 0
    #b[0] = 0
    print(a[x])
    for i in range(1, a.shape[0]) :
        a[i] = a[i-1] + inputa[i]
        #b[i] = inputa[i-1] + a[i]
    return a #, b
 
#df[['a','b']] = foo(df['inputa'].values)
df['a'] = foo(df['inputa'].values)

print(df)

但是，如果我尝试使用 2 列 outputs/results，则会遇到问题。基本上，我想做多个递归列 w/o 使用 df.loc.

请指教

非常感谢。

Answer 1

你的方法很好，除了 b[i] = input[i-1] + a[i] 中有错别字，因为 @Jérôme Richard 正确提到的 input 是 python 中的内置方法，

import numba as nb
import numpy as np
import pandas as pd

# DataFrame
df = pd.DataFrame(np.arange(1,101),columns=['data'])
df_vals = df.values.ravel() # .ravel to unfold a 2d array into 1d array


# Regular Python Function
def foo(arr):
    arrlen = arr.shape[0]
    a = np.zeros(arrlen)
    b = np.zeros(arrlen)
    for i in range(1, a.shape[0]) :
        a[i] = a[i-1] + arr[i]
        b[i] = arr[i-1] + a[i]
    return a, b

# Jitted (nopython) Function
foo_nb = nb.njit()(foo)


# Numba Warmup
_ = foo_nb(df_vals)


a_vals, b_vals = foo_nb(df_vals)

df['a'] = a_vals
df['b'] = b_vals

# Performance Benchmarks *NOTE : To be used in Jupyter Notebook.
print('Non Numba Code Performance : ')
foo_timeit = %timeit -o -n 1000 foo(df_vals)
print('\nNumba Code Performance : ')
foo_nb_timeit = %timeit -o -n 1000 foo_nb(df_vals)

输出：

Non Numba Code Performance : 
289 µs ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Numba Code Performance : 
1.22 µs ± 174 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

是否可以使用 numba 输出 2 列并将其用作 df 2 新列

Is it possible to have 2 columns output using numba and use it as df 2 new columns

python

numpy

function

pandas

numba