Python中优化嵌套循环数值计算的Numba加速指南

from numba import njit
import time

@njit
def find_combinations_jit():
    target = 0.3048
    tolerance = 1e-06
    found_count = 0
    for a in range(-100, 101):
        for b in range(-100, 101):
            for c in range(-100, 101):
                for d in range(-100, 101):
                    # 使用浮点数进行幂运算，避免整数溢出或精度问题
                    n = (2.0**a) * (3.0**b) * (5.0**c) * (7.0**d)
                    v = n - target
                    if abs(v) <= tolerance:
                        # 在Numba编译函数中，print语句的性能可能不如纯数值计算
                        # 但为了演示，此处保留
                        print(
                            "a=", a, ", b=", b, ", c=", c, ", d=", d,
                            ", the number=", n, ", error=", abs(v)
                        )
                        found_count += 1
    return found_count

print("--- Numba JIT 模式 ---")
start_time = time.time()
count = find_combinations_jit()
end_time = time.time()
print(f"找到 {count} 组合。执行时间: {end_time - start_time:.2f} 秒")

from numba import njit, prange
import time

@njit(parallel=True)
def find_combinations_parallel():
    target = 0.3048
    tolerance = 1e-06
    found_count = 0
    # 将外层循环设为prange，允许Numba并行化
    for a in prange(-100, 101):
        i_a = 2.0**a  # 预计算2的a次幂
        for b in prange(-100, 101): # 同样可以并行化
            i_b = i_a * (3.0**b) # 预计算2的a次幂乘以3的b次幂
            for c in prange(-100, 101): # 同样可以并行化
                i_c = i_b * (5.0**c) # 预计算2的a次幂乘以3的b次幂乘以5的c次幂
                for d in prange(-100, 101):
                    n = i_c * (7.0**d) # 最终结果
                    v = n - target
                    if abs(v) <= tolerance:
                        # 在并行模式下，print语句可能导致竞争条件或性能下降
                        # 实际应用中，通常会收集结果到列表中，然后在函数外部打印
                        # 为了演示，此处保留
                        print(
                            "a=", a, ", b=", b, ", c=", c, ", d=", d,
                            ", the number=", n, ", error=", abs(v)
                        )
                        # 注意：在并行循环中直接修改共享变量 (如found_count) 需要原子操作，
                        # Numba的prange默认不提供，可能导致计数不准确。
                        # 更安全的做法是每个线程计算自己的部分，然后合并。
                        # 此处仅为演示并行化效果，实际计数应谨慎处理。
                        # found_count += 1 # 暂不推荐在prange内直接计数
    return found_count # 返回的found_count可能不准确，需要外部聚合

print("\n--- Numba 并行模式 (prange) ---")
start_time = time.time()
# 注意：并行模式下的print输出顺序可能不确定
# 并且，如果`print`语句在并行执行的循环内部，其开销可能会抵消部分并行化带来的优势。
# 在实际生产代码中，更推荐将结果收集到数组或列表中，然后在并行循环结束后统一处理。
count = find_combinations_parallel()
end_time = time.time()
print(f"执行时间: {end_time - start_time:.2f} 秒 (请注意并行模式下print和计数的局限性)")