如何加速Python中执行复杂公式的代码运行速度

你的遗传算法模拟中，AF函数是核心性能瓶颈——它在个体适应度计算、均匀阵列特性计算、阵列综合特性计算中被频繁调用，且当前嵌套Python循环的实现效率极低。以下是针对性的加速优化方案：

1. 预计算固定参数 #

将不随theta和个体基因变化的常量参数提前计算，避免每次调用AF时重复计算：

# 在模块级别预计算每个环的固定参数
RING_PARAMS = []
for m in range(RING_NUM):
    ring_size = RING_ELEMENT_NUM[m]
    # 预计算每个环元素的角度项：PHI - 2π*(i-1)/ring_size
    angles = PHI - 2 * np.pi * (np.arange(ring_size) - 1) / ring_size
    # 预计算K*(m*WAVELENGTH/2)项
    k_term = K * (m * WAVELENGTH / 2)
    RING_PARAMS.append((k_term, angles))

2. 向量化计算替代嵌套循环 #

利用Numpy的向量化操作替代Python原生嵌套循环，大幅提升计算效率：

def AF(self, cur, theta):
    cur_np = np.array(cur, dtype=np.float64)
    sin_theta = np.sin(theta)
    af_total = 1.0 + 0.0j  # 初始化为复数类型
    start_idx = 0
    
    for m in range(RING_NUM):
        ring_size = RING_ELEMENT_NUM[m]
        k_term, angles = RING_PARAMS[m]
        
        # 提取当前环的基因片段
        ring_genes = cur_np[start_idx:start_idx+ring_size]
        # 向量化计算指数项
        exp_terms = np.exp(1j * k_term * sin_theta * np.cos(angles))
        # 向量化求和
        ring_sum = np.sum(ring_genes * exp_terms)
        
        af_total += ring_sum
        start_idx += ring_size
    
    return af_total

3. 使用Numba即时编译加速 #

通过Numba的JIT编译将AF函数编译为机器码，进一步提升数值计算速度：

1. 先安装Numba：pip install numba
2. 修改AF函数为Numba优化版本：

from numba import jit, complex128, float64

# 将预计算的参数转为Numba可处理的数组
K_TERMS = np.array([p[0] for p in RING_PARAMS], dtype=float64)
ANGLES_ARRAY = np.array([p[1] for p in RING_PARAMS], dtype=float64)

@jit(complex128(float64[:], float64), nopython=True)
def af_numba(cur, theta):
    sin_theta = np.sin(theta)
    af_total = 1.0 + 0.0j
    start_idx = 0
    
    for m in range(RING_NUM):
        ring_size = RING_ELEMENT_NUM[m]
        k_term = K_TERMS[m]
        angles = ANGLES_ARRAY[m]
        
        ring_genes = cur[start_idx:start_idx+ring_size]
        exp_terms = np.exp(1j * k_term * sin_theta * np.cos(angles))
        ring_sum = np.sum(ring_genes * exp_terms)
        
        af_total += ring_sum
        start_idx += ring_size
    
    return af_total

# 在Individual类的AF方法中调用该函数
def AF(self, cur, theta):
    return af_numba(np.array(cur, dtype=np.float64), theta)

4. 减少重复计算 #

fn_uniform函数的计算结果不依赖个体基因，可缓存其结果避免重复计算：

from functools import lru_cache

class Individual:
    # 将fn_uniform改为静态方法，并用lru_cache缓存结果
    @staticmethod
    @lru_cache(maxsize=None)
    def fn_uniform():
        AF_uni = np.zeros(1801, dtype=np.complex128)
        cur = np.ones(CHROMOSOME_LENGTH, dtype=np.float64)
        for theta_idx in range(-900, 901):
            theta = theta_idx / 10
            AF_uni[theta_idx + 900] = af_numba(cur, theta)
        
        null1 = 0
        for i in range(899):
            if np.abs(AF_uni[i + 901]) > np.abs(AF_uni[i + 900]):
                null1 = i / 10
                break
        
        null2 = 0
        for i in range(899):
            if np.abs(AF_uni[-i + 900 - 1]) > np.abs(AF_uni[i + 900]):
                null2 = i / 10
                break
        
        fn = np.abs(null2 - null1)
        return [fn]

额外优化点 #

• 将个体的genes属性改为Numpy数组，避免每次调用AF时重复转换类型
• 优化fn_com中的循环，同样使用向量化方式批量计算所有theta对应的AF值，而非逐个循环调用

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