numba，一个强大的 python 库

news2024/10/5 14:52:17

今天给大家分享一个强大的 python 库，numba。

https://github.com/numba/numba

Pandas 是 Python 中流行的数据分析库。然而，随着数据集大小的增长，原生 Python 代码对于滚动窗口计算等关键操作可能会变得很慢。这就是 Numba 用武之地。

Numba 是 Python 的即时编译器，可在运行时将 Python 代码转换为机器代码。

Numba 可用于加速 Python 函数，而无需编写任何 C 或 C++ 代码。

在本文中，我们将探讨如何利用 Numba 加速常见的 Pandas 工作流程，例如滚动统计。‍

你将了解启用 Numba 编译所需的简单注释以及如何使用并行执行等选项来调整性能。

Numba 的安装

我们可以直接使用pip 来进行安装。

pip install numba

如何对 Pandas 使用 Numba？

将 Numba 与 Pandas 结合使用有两种方法：

在 pandas 方法中指定 engine=“numba” 关键字。
定义你自己的用 @jit 修饰的 Python 函数，并将 DataFrame 或 Series 的底层 NumPy 数组（使用 to_numpy()）传递到函数中。

使用 engine=“numba”

你可以通过指定 engine=“numba” 来增强所选 Pandas 方法的执行。

这指示该方法利用 Numba 来加速性能。支持该关键字的方法主要涉及窗口操作，例如：

滚动平均值、中位数、最大值、最小值、总和和标准差等。
groupby 滚动平均值、中位数、最大值、最小值、总和和标准差等。
expanding 平均值、中位数、最大值、最小值、总和和标准差等。
groupby expanding 平均值、中位数、最大值、最小值、总和和标准差等。

# create a sample DataFrame  
data = np.random.rand(int(1e5), 4)  
df = pd.DataFrame(data)  
  
window_size = 10  
# rolling sum with Numba   
rolling_sum = df.rolling(window_size).sum(engine='numba')

上面的代码片段生成一个随机 DataFrame 并使用 Numba 引擎对其执行滚动求和操作，从而优化计算以提高效率。

定义你自己的函数

你还可以定义用 @jit 或 @njit 装饰的 Python 函数，并将 DataFrame 或 Series 的基础 NumPy 数组（使用 to_numpy()）传递到函数中。此方法比使用 engine=“numba” 更灵活，因为你可以定义自己的自定义函数。

from numba import njit  
  
# create a sample DataFrame  
data = np.random.rand(int(1e5), 4)  
df = pd.DataFrame(data)  
  
# Define the custom function  
@njit  
def sum_of_squares(x):  
    return np.sum(x**2, axis=0)  
  
result = sum_of_squares(df.to_numpy())

此代码片段使用 Numba 的 @njit 装饰器（即 @jit(nopython=True) 的别名）来定义名为 sum_of_squares 的自定义函数。

它采用 NumPy 数组作为输入并计算沿列的平方和。

该代码生成一个随机 DataFrame，使用 to_numpy() 将其转换为 NumPy 数组，并应用 sum_of_squares 函数计算每列的平方和。结果存储在 result 变量中。

性能比较

我们比较执行时间并评估通过利用 Numba 的即时编译功能获得的潜在加速。

为此，我们在如下三种场景下测量代码的执行时间：有和没有 Numba 优化，以及并行执行。

下面提供了执行这些基准测试的代码：

import time  
import matplotlib.pyplot as plt  
import numba  
import numpy as np  
import pandas as pd  
  
plt.style.use('ggplot')  
# Define the custom rolling apply function  
def root_mean_square(x):  
    return np.sqrt(np.mean(x**2))  
  
def bench(df: pd.DataFrame, use_numba: bool, use_parallel: bool = False, n_times: int = 10):  
    engine = 'numba' if use_numba else None  
    engine_kwargs = {'parallel': True, 'nopython': True} if use_parallel else None  
      
    elapsed_time_list = []  
    for _ in range(n_times):  
        start_time = time.time()  
        df.rolling(window_size).apply(root_mean_square, raw=True, engine=engine, engine_kwargs=engine_kwargs)  
        elapsed_time_list.append(time.time() - start_time)  
    return np.mean(elapsed_time_list)  
  
  
# Generate sample data  
np.random.seed(42)  
data = np.random.rand(int(1e5), 4)  
df = pd.DataFrame(data)  
  
window_sizes = range(10, 1001, 100)  
mean_time_without_numba = []  
mean_time_with_numba = []  
mean_time_with_numba_parallel = []  
# Benchmark the performance for different window sizes  
for window_size in window_sizes:  
    # Without Numba  
    mean_time = bench(df, use_numba=False)  
    mean_time_without_numba.append(mean_time)  
      
    mean_time = bench(df, use_numba=True)  
    mean_time_with_numba.append(mean_time)  
  
    mean_time = bench(df, use_numba=True, use_parallel=True)  
    mean_time_with_numba_parallel.append(mean_time)  
  
mean_time_without_numba = np.array(mean_time_without_numba)  
mean_time_with_numba = np.array(mean_time_with_numba)  
mean_time_with_numba_parallel = np.array(mean_time_with_numba_parallel)  
print("mean_time_without_numba", mean_time_without_numba)  
print("mean_time_with_numba", mean_time_with_numba)  
print("mean_time_with_numba_parallel", mean_time_with_numba_parallel)  
  
fig, ax = plt.subplots(figsize=(8, 6))  
ax.plot(window_sizes, mean_time_without_numba / mean_time_with_numba, label='Numba')  
ax.plot(window_sizes, mean_time_without_numba / mean_time_with_numba_parallel, label='Numba + Parallel')  
  
ax.set_xlabel('Window Size')  
ax.set_ylabel('Speedup')  
ax.set_title('Speedup compared to without Numba: Pandas Rolling Apply')  
ax.legend()  
fig.savefig("numba_speedup_pandas_rolling_apply.png")