Python三维物理空间三重嵌套循环单核心性能优化求助

问题背景 #

代码性能瓶颈集中在遍历三维物理空间(r, theta, phi)的三重嵌套循环中，无法通过重构逻辑避免该循环（必须处理每个数据点）。该循环会被调用8825次以上，且仅允许使用单核心。当前循环耗时14-15秒，目标是降至8-10秒。

已尝试的优化手段：

• 预计算np.sin/np.cos等结果并存储到数组，避免重复计算
• 使用O(n*log n)的二分查找函数在有序数组中查找

附最小可复现代码：

import numpy as np
from time import perf_counter

Nb = 300
l_max = 50
Z = 1.0

def Coulomb(r):
    return (-Z / r)

def find_idx_closest_to_input(array, target_value):
    # it assumes array is sorted and is 1 dimensional
    n = array.shape[0]
    if (target_value < array[0]):
        return -1
    elif (target_value > array[n-1]):
        return n

    jl = 0 # Initialize lower
    ju = n-1 # and upper limits.

    while (ju-jl > 1): # If we are not yet done,
        jm = (ju+jl) >> 1 # compute a midpoint with a bitshift
        if (target_value >= array[jm]):
            jl = jm # and replace either the lower limit
        else:
            ju = jm # or the upper limit, as appropriate.
        # Repeat until the test condition is satisfied.

    if (target_value == array[0]): # edge cases at bottom
        return 0
    elif (target_value == array[n-1]): # and top
        return n-1
    else:
        return jl

coulomb_array = np.random.rand( Nb )

ris = np.random.rand( Nb )
sin_thetas = np.random.rand( l_max+1 )
cos_thetas = np.random.rand( l_max+1 )
sin_phis = np.random.rand( 2*l_max+2 )
cos_phis = np.random.rand( 2*l_max+2 )
r_vec_cartesian = np.zeros ( (3) )
b_vec = np.random.rand ( 3 )
b_dot_vec = np.random.rand ( 3 )

k_vec_cartesian = np.random.rand ( 3 )
k_wave = 2.0
conhyp_results = np.random.rand( 15, 10000, 2 ) + 1j * np.random.rand( 15, 10000, 2 ) 

Psi_scattering_prelim = np.zeros ( ( Nb, (l_max+1), (2*l_max+2) ) ) + 1j * np.zeros ( ( Nb, (l_max+1), (2*l_max+2) ) )
tensor_of_exp_of_minusI_bdot_times_r = np.zeros ( (Nb, (l_max+1), (2*l_max+2)) ) + 1j * np.zeros ( (Nb, (l_max+1), (2*l_max+2)) )
difference_in_potential_values = np.zeros( (Nb, (l_max+1), (2*l_max+2)) )

k_idx = 10

prefactor = 10.0

t_coordinates_start = perf_counter()
for a in range(Nb):
    for b in range(l_max+1):
        for c in range(2*l_max+2):
            r_vec_cartesian[0] = ris[a] * sin_thetas[b] * cos_phis[c]
            r_vec_cartesian[1] = ris[a] * sin_thetas[b] * sin_phis[c]
            r_vec_cartesian[2] = ris[a] * cos_thetas[b]                 
            r_minus_b_vec_cartesian = r_vec_cartesian - b_vec
            r_minus_b_vec_cartesian_modulus = np.sqrt( r_minus_b_vec_cartesian.dot(r_minus_b_vec_cartesian) )

            arg3_imag_part = -( k_wave*r_minus_b_vec_cartesian_modulus + k_vec_cartesian.dot(r_minus_b_vec_cartesian) ) # a np.float64

            index_from_conhyp_results = find_idx_closest_to_input( np.imag(conhyp_results[k_idx, :, 0]), arg3_imag_part )

            result_cpp = conhyp_results[k_idx, index_from_conhyp_results, 1]
            
            factor = np.exp( 1j * k_vec_cartesian.dot(r_minus_b_vec_cartesian) )
            Psi_scattering_prelim[a, b, c] = prefactor * factor * result_cpp
            b_dot_vec_times_r_vec = b_dot_vec.dot(r_vec_cartesian)
            tensor_of_exp_of_minusI_bdot_times_r[a, b, c] = np.exp(-1j * b_dot_vec_times_r_vec)

            difference_in_potential_values[a, b, c] = coulomb_array[a] - Coulomb(r_minus_b_vec_cartesian_modulus)
# END for loop through coordinates

t_coordinates_end = perf_counter()
print("We have gone through the coordinates")
print("To go through coordinates it took: " + str(t_coordinates_end - t_coordinates_start) + " seconds")

优化方案 #

1. 用numpy内置二分查找替代自定义函数 #

自定义查找函数可以直接替换为np.searchsorted——这是底层优化的C实现，速度远快于Python循环。提前把目标数组提取到循环外，再在循环内调用：

# 循环外预提取目标数组（确保数组有序）
target_array = np.imag(conhyp_results[k_idx, :, 0])

# 循环内替换查找逻辑
index_from_conhyp_results = np.searchsorted(target_array, arg3_imag_part, side='right') - 1
# 处理边界情况
if index_from_conhyp_results < 0:
    index_from_conhyp_results = 0
elif index_from_conhyp_results >= len(target_array):
    index_from_conhyp_results = len(target_array)-1

2. 向量化计算笛卡尔坐标与点积 #

利用numpy广播特性，把内层循环的逐元素计算转为批量操作，避免Python循环的开销：

# 预先生成三维网格索引
a_grid, b_grid, c_grid = np.meshgrid(np.arange(Nb), np.arange(l_max+1), np.arange(2*l_max+2), indexing='ij')

# 向量化计算笛卡尔坐标
r_x = ris[a_grid] * sin_thetas[b_grid] * cos_phis[c_grid]
r_y = ris[a_grid] * sin_thetas[b_grid] * sin_phis[c_grid]
r_z = ris[a_grid] * cos_thetas[b_grid]

# 计算r - b的笛卡尔分量
r_minus_b_x = r_x - b_vec[0]
r_minus_b_y = r_y - b_vec[1]
r_minus_b_z = r_z - b_vec[2]

# 批量计算模长
r_minus_b_modulus = np.sqrt(r_minus_b_x**2 + r_minus_b_y**2 + r_minus_b_z**2)

# 批量计算点积
k_dot_r_minus_b = k_vec_cartesian[0]*r_minus_b_x + k_vec_cartesian[1]*r_minus_b_y + k_vec_cartesian[2]*r_minus_b_z
b_dot_r = b_dot_vec[0]*r_x + b_dot_vec[1]*r_y + b_dot_vec[2]*r_z

3. 批量计算Coulomb函数与势能差 #

把Coulomb函数的逐次调用转为数组批量运算：

# 批量计算Coulomb值
coulomb_values = -Z / r_minus_b_modulus
# 批量计算势能差
difference_in_potential_values = coulomb_array[a_grid] - coulomb_values

4. 减少循环内的函数查找 #

循环内重复调用的numpy函数（如np.exp）可以提前赋值给局部变量，减少属性查找开销：

# 循环外预赋值
np_exp = np.exp
# 循环内使用
factor = np_exp(1j * k_dot_r_minus_b[a,b,c])

5. 内存与类型优化 #

确保输出数组的 dtype 与输入匹配（比如如果精度允许，用complex64替代complex128），减少内存占用和计算开销：

Psi_scattering_prelim = np.zeros((Nb, l_max+1, 2*l_max+2), dtype=np.complex128)
tensor_of_exp_of_minusI_bdot_times_r = np.zeros_like(Psi_scattering_prelim)

内容的提问来源于stack exchange，提问作者velenos14