如何使用Numpy对两层嵌套for循环进行向量化优化以缩短运行时长

核心优化点 #

• 移除所有i、j层显式循环，全部用numpy广播和矩阵运算实现，速度提升至少2个数量级
• 预计算所有和hexind绑定的波矢分量，避免循环内重复计算
• 利用delta项仅依赖k、i的特性，不需要在j维度重复计算
• 特征向量内积直接用矩阵乘法批量计算，替代逐元素dot

完整实现代码 #

预计算全局依赖（仅需执行1次） #

import numpy as np
epsilon = 0.15
im = complex(0,1)

def delta(x):
    # 洛伦兹型δ函数
    return (1/np.pi)*(epsilon)/(epsilon**2+x**2)

# 按你给出的最小示例定义基础变量
G1 = np.random.rand(1,2)
G2 = np.random.rand(1,2)
index = np.array([[ 0,  0],
       [ 1,  1],
       [ 1,  0],
       [ 0, -1],
       [-1, -1],
       [-1,  0],
       [ 0,  1]])
hexind2 = np.repeat(index,2, axis=0)
hexind = np.tile(hexind2,(2,1))
n = len(hexind)

# 预计算每个i对应的波矢分量，仅计算一次，后续不用循环重复算
G_mat = np.vstack([G2, G1])  # shape (2,2)
k_vecs = hexind @ G_mat  # shape (n, 2)
kx = k_vecs[:, 0].reshape(-1,1,1,1)  # 调整形状方便后续广播运算
ky = k_vecs[:, 1].reshape(-1,1,1,1)

优化后的主函数 #

def quantity_abs_squared_vec(Vecs, Egrid, Eigen, Rx, Ry):
    kdots = Vecs.shape[0]
    # 1. 预计算指数相位矩阵，shape (n,n,M,M)，M为Rx/Ry的网格边长
    dkx = kx - kx.reshape(1,n,1,1)
    dky = ky - ky.reshape(1,n,1,1)
    phase = im * (dkx * Rx + dky * Ry)
    exp_mat = np.exp(phase)
    
    # 2. 预计算delta项，仅和k、i有关，shape (kdots, n, 1, 1)
    delta_vals = delta(Egrid - Eigen).reshape(kdots, n, 1, 1)
    
    holder = 0
    # 仅保留k层循环，1575次循环在numpy矩阵运算加持下速度极快
    for k in range(kdots):
        # 单次矩阵乘法直接得到所有i、j组合的特征向量内积，shape (n,n)
        vec_k = Vecs[k]
        coeff_mat = vec_k.T @ vec_k.conj()
        # 广播乘指数矩阵后求和i、j维度
        term = (coeff_mat[..., np.newaxis, np.newaxis] * exp_mat).sum(axis=(0,1))
        # 乘delta项后求和i维度，加入总结果
        holder += (delta_vals[k] * term).sum(axis=0)
    
    return 4 * holder / kdots

调用示例 #

kdots = 100
n = 28
# 生成测试用特征值、特征向量
E = np.random.rand(kdots, n)
vec = np.random.rand(kdots,n,n) + 1j*np.random.rand(kdots,n,n)

# 生成网格
rx = np.linspace(-20, 20)
ry = np.linspace(-20, 20)
Rx,Ry = np.meshgrid(rx,ry)

# 调用优化后的函数
Eval = 300
result = quantity_abs_squared_vec(vec, Eval, E, Rx,Ry)

性能说明 #

原代码n=244时，每个k点要执行244*244≈6万次python层循环，1575个k点累计近1亿次循环，耗时自然很长。优化后所有i、j维度的运算都交给numpy底层C实现，原10分钟的计算任务优化后仅需几秒到几十秒即可完成。如果内存不足可以将exp_mat的计算拆到k循环内，或者用numba对k层循环做JIT编译进一步提速。

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