本文旨在提供一种优化Python中嵌套循环计算效率的方法，特别是针对计算密集型任务。通过使用Numba库的即时编译（JIT）技术，可以显著提升代码的执行速度，避免传统嵌套循环带来的
性能瓶颈
。文章将展示如何使用Numba加速原始代码，并提供并行化的优化方案，以及性能对比。
在Python中，嵌套循环是常见的编程结构，但当循环次数较多时，其执行效率会显著下降。对于需要处理大量数据的科学计算或数据分析任务，优化嵌套循环至关重要。Numba是一个开源的即时编译器，它可以将Python代码转换为优化的机器码，从而显著提高程序的运行速度。本文将介绍如何使用Numba来优化包含嵌套循环的Python函数。
使用Numba加速计算
首先，我们来看一个包含嵌套循环的示例函数 U_p_law，该函数计算两个概率密度函数之间的关系：
为了使用Numba加速这个函数，我们需要导入 numba 库，并使用 @njit 装饰器修饰函数。同时，为了获得更好的性能，我们也可以对 probability_of_loss 函数进行加速。
立即学习
“
Python免费学习笔记（深入）
”；
注意:
@njit 装饰器告诉Numba将该函数编译为机器码。
为了保证Numba能够成功编译，我们需要确保函数中使用的所有操作和数据类型都受Numba支持。例如，如果使用NumPy数组，需要确保数组的数据类型是Numba支持的类型。
Numba 推荐使用 np.float64 作为浮点数类型，以获得更好的性能。
使用并行化进一步加速
对于计算量更大的任务，我们可以利用多核CPU的优势，使用Numba的并行化功能。为了实现并行化，我们需要使用 parallel=True 参数修饰 @njit 装饰器，并将外层循环替换为 prange。
Python 3.14.3
微软官方的 Python 扩展，是 VS Code 安装量最高的扩展（209M+）。集成 IntelliSense（通过 Pylance）、调试（通过 Python Debugger）、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。
下载
注意:
prange 是 Numba 提供的并行循环，它会将循环迭代分配到多个线程上执行。
并行化可以显著提高程序的运行速度，但也会带来一些额外的开销，例如线程创建和同步。因此，只有当计算量足够大时，并行化才能带来明显的性能提升。
性能测试
为了验证Numba的加速效果，我们可以使用 timeit 模块来测试不同版本的函数的运行时间。
在我的机器上（AMD 5700x），运行结果如下：
从结果可以看出，使用 Numba 可以显著提高程序的运行速度。在这个例子中，使用 Numba JIT 可以提速约 170 倍，而使用多线程 Numba JIT 可以提速约 640 倍。
总结
Numba 是一个强大的工具，可以用来优化 Python 代码的性能，特别是对于包含嵌套循环的计算密集型任务。通过使用 Numba 的即时编译和并行化功能，我们可以显著提高程序的运行速度，从而更快地完成计算任务。在使用 Numba 时，需要注意确保函数中使用的所有操作和数据类型都受 Numba 支持，并根据实际情况选择合适的优化策略。
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_pfrom 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_pfrom 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_pfrom 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)10 calls using vanilla Python : 2.4276352748274803
10 calls using Numba : 0.013957140035927296
10 calls using Numba (+ parallel) : 0.003793451003730297