如何加速/向量化海洋生物地球化学加权平均的嵌套循环？

问题背景 #

我刚接触Python，编程经验有限，现在需要在大规模海洋生物地球化学数据集上运行一个加权平均计算——该计算需纳入采样日期间气体传输速率（k）的历史数据，权重由大气通风水柱比例（f）决定，越接近采样时间权重越高（采样时刻权重为1，随时间回溯递减）。当前用嵌套循环实现的代码处理近2年的模型数据（k维度为927360，7天窗口内样本数10080）时，单次运行需数小时，且计算需重复执行，急需提速。

原实现代码：

def kw_omega (k, depth, window, samples_day):
    """
        calculate  the scheme weights for gas transfer velocity of oxygen
        over the previous window of time, where the most recent gas transfer velocity
        has a weight of 1, and the weighting decreases going back in time. The rate of decrease
        depends on the wind history and MLD.

        Parameters
        ----------
        k: ndarray
            instantaneous O2 gas transfer velocity

        depth: ndarray
            Water depth

        window: integer
            weighting period in days which equals the residence time of oxygen at sampling day

        samples_day: integer
            number of samples in each day composing window

        Returns
        ---------
        weighted_kw: ndarray

        Notes
        ---------
        n = the weighting period / the time resolution of the wind data
        samples_day = the time resolution of the wind data
        omega = is the weighting coefficient at each time step within the weighting window
        f = the fraction of the water column (mixed layer, photic zone or full water column) ventilated at each time

    """
    Dt = 1./samples_day
    f = (k*Dt)/depth
    f = np.flip(f)
    k = np.flip(k)
    n = window*samples_day
    weighted_kw = np.zeros(len(k))
    for t in np.arange(len(k) - n):
        omega = np.zeros((n))
        omega[0] = 1.
        for i in np.arange(1,len(omega)):
            omega[i] = omega[i-1]*(1-f[t+(i-1)])
        weighted_kw[t] = sum(k[t:t+n]*omega)/sum(omega)
        print(f"t = {t}")
    return np.flip(weighted_kw)

提速方案 #

1. 移除冗余打印操作 #

直接删除代码中的print(f"t = {t}")——百万级循环的打印操作会显著拖慢运行速度，完全可以在计算完成后再做进度验证。

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

嵌套循环是Python低效的核心原因。观察omega的递推逻辑：omega[i] = omega[i-1]*(1-f[t+i-1])，本质是(1-f)序列从t开始的前i项累积乘积。我们可以预计算全局累积乘积，再通过滑动窗口快速生成每个t对应的omega序列：

import numpy as np

def kw_omega_optimized(k, depth, window, samples_day):
    Dt = 1. / samples_day
    f = (k * Dt) / depth
    f = np.flip(f)
    k = np.flip(k)
    n = window * samples_day
    len_k = len(k)
    weighted_kw = np.zeros(len_k)

    # 预计算(1-f)的累积乘积数组，开头插入1以匹配omega[0]=1的初始条件
    one_minus_f = 1 - f
    cum_prod = np.insert(np.cumprod(one_minus_f), 0, 1.0)

    for t in range(len_k - n):
        # 通过累积乘积的比值快速生成omega序列
        omega = cum_prod[t:t+n] / cum_prod[t]
        # 用NumPy的向量化求和替代Python内置sum
        numerator = np.sum(k[t:t+n] * omega)
        denominator = np.sum(omega)
        weighted_kw[t] = numerator / denominator

    return np.flip(weighted_kw)

3. Numba JIT编译加速 #

如果向量化后仍达不到预期速度，用Numba将Python代码编译为机器码，能大幅提升循环效率：

先安装Numba：pip install numba

优化后代码：

import numpy as np
from numba import jit, prange

@jit(nopython=True, parallel=True)  # 启用并行编译
def kw_omega_numba(k, depth, window, samples_day):
    Dt = 1. / samples_day
    len_k = len(k)
    f = np.empty(len_k, dtype=np.float64)
    for i in range(len_k):
        f[i] = (k[i] * Dt) / depth[i]
    
    # 翻转数组
    f = f[::-1]
    k = k[::-1]
    n = window * samples_day
    weighted_kw = np.zeros(len_k, dtype=np.float64)

    # 并行遍历每个时间步
    for t in prange(len_k - n):
        omega = np.zeros(n, dtype=np.float64)
        omega[0] = 1.0
        for i in range(1, n):
            omega[i] = omega[i-1] * (1 - f[t + i - 1])
        # 手动求和减少数组操作开销
        numerator = 0.0
        denominator = 0.0
        for i in range(n):
            numerator += k[t + i] * omega[i]
            denominator += omega[i]
        weighted_kw[t] = numerator / denominator
    
    return weighted_kw[::-1]

4. 内存优化（可选） #

如果计算精度允许，将k和depth从float64转为float32，减少内存占用和带宽压力：

k = k.astype(np.float32)
depth = depth.astype(np.float32)

5. 分块并行处理（超大规模数据） #

若数据集超出内存承载，用Dask将数组分块，并行处理每个块的计算，适合TB级数据场景。

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