本文旨在提供一种使用 Numba 库优化 Python 中嵌套循环的有效方法。通过 JIT 编译和并行计算，可以显著提高计算密集型任务的执行速度，特别是在处理概率密度函数等数值计算时。我们将通过实例展示如何应用 Numba 来加速包含嵌套循环的函数，并对比优化前后的性能差异。
在 Python 中，嵌套循环经常出现在需要对多个维度的数据进行迭代和计算的场景中。然而，由于 Python 的解释型特性，当循环次数较多时，嵌套循环的执行效率往往较低。对于计算密集型的任务，这会成为性能瓶颈。为了解决这个问题，可以使用 Numba 库对 Python 代码进行优化。
Numba 简介
Numba 是一个开源的 JIT（Just-In-Time）编译器，它可以将 Python 代码编译成机器码，从而显著提高程序的执行速度。Numba 特别适用于数值计算，它可以自动识别并优化循环、数组操作等常见的计算模式。
优化嵌套循环的步骤
以下步骤展示了如何使用 Numba 优化包含嵌套循环的 Python 函数。我们将以计算概率密度函数的例子为例进行说明。
立即学习“Python免费学习笔记（深入）”；
1. 安装 Numba
首先，需要安装 Numba 库。可以使用 pip 命令进行安装：
pip install numba
2. 原始代码
假设有以下 Python 函数，其中包含嵌套循环：
import numpy as np
def probability_of_loss(x):
return 1 / (1 + np.exp(x / 67))
def U_p_law(W, L, L_P, L_Q):
omega = np.arange(0, 3501, 10)
U_p = np.zeros_like(omega, dtype=float)
for p_idx, p in enumerate(omega):
for q_idx, q in enumerate(omega):
U_p[p_idx] += (
probability_of_loss(q - p) ** W
* probability_of_loss(p - q) ** L
* L_Q[q_idx]
* L_P[p_idx]
)
normalization_factor = np.sum(U_p)
U_p /= normalization_factor
return omega, U_p3. 使用 Numba 进行 JIT 编译
要使用 Numba 优化上述函数，只需添加 @njit 装饰器即可。@njit 装饰器告诉 Numba 编译该函数，以便在运行时获得更高的性能。
from numba import njit
@njit
def probability_of_loss_numba(x):
return 1 / (1 + np.exp(x / 67))
@njit
def U_p_law_numba(W, L, L_P, L_Q):
omega = np.arange(0, 3501, 10, dtype=np.float64)
U_p = np.zeros_like(omega)
for p_idx, p in enumerate(omega):
for q_idx, q in enumerate(omega):
U_p[p_idx] += (
probability_of_loss_numba(q - p) ** W
* probability_of_loss_numba(p - q) ** L
* L_Q[q_idx]
* L_P[p_idx]
)
normalization_factor = np.sum(U_p)
U_p /= normalization_factor
return omega, U_p注意：
- 为了保证 Numba 的最佳性能，建议在 Numba 编译的函数中使用 NumPy 数组，并指定数组的数据类型。
- probability_of_loss 也需要使用 @njit 装饰器编译。
4. 使用并行计算
对于某些计算密集型的嵌套循环，还可以利用 Numba 的并行计算能力进一步提高性能。要启用并行计算，可以使用 parallel=True 参数。
from numba import njit, prange
@njit(parallel=True)
def U_p_law_numba_parallel(W, L, L_P, L_Q):
omega = np.arange(0, 3501, 10, dtype=np.float64)
U_p = np.zeros_like(omega)
for p_idx in prange(len(omega)):
p = omega[p_idx]
for q_idx in prange(len(omega)):
q = omega[q_idx]
U_p[p_idx] += (
probability_of_loss_numba(q - p) ** W
* probability_of_loss_numba(p - q) ** L
* L_Q[q_idx]
* L_P[p_idx]
)
normalization_factor = np.sum(U_p)
U_p /= normalization_factor
return omega, U_p注意：
- 使用 parallel=True 参数需要使用 prange 代替 range。prange 是 Numba 提供的并行循环，它可以将循环任务分配给多个线程并行执行。
5. 性能测试
为了验证 Numba 的优化效果，可以使用 timeit 模块对原始函数和优化后的函数进行性能测试。
from timeit import timeit
P_mean = 1500
P_std = 100
Q_mean = 1500
Q_std = 100
W = 1 # Number of matches won by P
L = 0 # Number of matches lost by P
L_P = np.exp(-0.5 * ((np.arange(0, 3501, 10) - P_mean) / P_std) ** 2) / (
P_std * np.sqrt(2 * np.pi)
)
L_Q = np.exp(-0.5 * ((np.arange(0, 3501, 10) - Q_mean) / Q_std) ** 2) / (
Q_std * np.sqrt(2 * np.pi)
)
omega_1, U_p_1 = U_p_law(W, L, L_P, L_Q)
omega_2, U_p_2 = U_p_law_numba(W, L, L_P, L_Q)
omega_3, U_p_3 = U_p_law_numba_parallel(W, L, L_P, L_Q)
assert np.allclose(omega_1, omega_2)
assert np.allclose(omega_1, omega_3)
assert np.allclose(U_p_1, U_p_2)
assert np.allclose(U_p_1, U_p_3)
t1 = timeit("U_p_law(W, L, L_P, L_Q)", number=10, globals=globals())
t2 = timeit("U_p_law_numba(W, L, L_P, L_Q)", number=10, globals=globals())
t3 = timeit("U_p_law_numba_parallel(W, L, L_P, L_Q)", number=10, globals=globals())
print("10 calls using vanilla Python :", t1)
print("10 calls using Numba :", t2)
print("10 calls using Numba (+ parallel) :", t3)在我的机器上（AMD 5700x），测试结果如下：
10 calls using vanilla Python : 2.4276352748274803 10 calls using Numba : 0.013957140035927296 10 calls using Numba (+ parallel) : 0.003793451003730297
可以看到，使用 Numba 进行 JIT 编译后，函数的执行速度提高了约 170 倍。如果启用并行计算，函数的执行速度可以提高约 640 倍。
总结
通过使用 Numba 库，可以显著提高 Python 中嵌套循环的执行效率。Numba 的 JIT 编译和并行计算能力可以有效地优化计算密集型任务，从而提高程序的整体性能。在实际应用中，可以根据具体情况选择合适的优化策略，例如，对于简单的循环，可以使用 @njit 装饰器进行 JIT 编译；对于计算密集型的循环，可以启用并行计算。
