numpy 有没有一种方法可以加速或向量化嵌套的for循环？

4si2a6ki 于 2022-12-29 发布在其他

关注(0)|答案(1)|浏览(123)

我对编程，尤其是Python，还是一个新手，我正在尝试将加权平均方案应用到大数据集，目前这需要几个小时才能完成，我希望加快这个过程，因为这需要重复几次。
加权平均值代表海洋生物地球化学中使用的一种方法，包括气体转移速度的历史（k）采样日期之间，其中k根据水柱分数加权（f）作为k的历史的函数由大气通风，并且向更接近采样时间的值分配更大的重要性（因此，在采样时间step = 1处的权重，然后其随时间远离而减小）：
Weight average equation extracted from (https://doi.org/10.1029/2017GB005874) pp. 1168
在我的尝试中，我使用了一个嵌套的for循环，在每个时间步长t，我计算了加权平均值：

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)

这应用于设置为运行近2年的模型模拟数据，其中模型时间步长设置为60秒，采样间隔为7天。因此k具有形状（927360）和n，表示7天内分钟数的形状为（10080）。目前它需要几个小时来运行。有什么方法可以使这个计算更快吗？

numpy

来源：https://stackoverflow.com/questions/74903553/is-there-a-way-to-speed-up-or-vectorize-a-nested-for-loop

1条答案

按热度按时间

gopyfrb31#

我建议使用numba包来加快计算速度。

import numpy as np
from numba import njit
from numpy.lib.stride_tricks import sliding_window_view

@njit
def k_omega(k_win, f_win):
   delta_t = len(k_win)
   omega_sum = omega = 1.0
   k_omega_sum = k_win[0]
   for t in range(1, delta_t):
       omega *= (1 - f_win[t])
       omega_sum += omega
       k_omega_sum = k_win[t] * omega
   return k_omega_sum / omega_sum

@njit
def windows_k_omega(k_wins, f_wins):
   size = len(k_wins)
   result = np.empty(size)
   for i in range(size):
       result[i] = k_omega(k_wins[i], f_wins[i])
   return result

def kw_omega(k, depth, window, samples_day):
   n = window * samples_day  # delta_t
   f = k / depth / samples_day
   k_wins = sliding_window_view(k, n)
   f_wins = sliding_window_view(f, n)
   k_omegas = windows_k_omega(k_wins, f_wins)
   weighted_kw = np.pad(weighted_kw, (len(k)-len(k_omegas), 0))
   return weighted_kw

在这里，我将函数分为三个部分，以使其更易于理解。函数k_omega基本上是将加权平均函数应用于k和f窗口。函数windows_k_omega只是为了加快循环速度，以便将函数元素应用于窗口。最后，外部函数kw_omega实现了你原来的函数接口，它使用numpy函数sliding_window_view来创建移动窗口（请注意，这是一个巧妙的numpy索引，所以这不是创建原始数组的副本），并使用helper函数执行计算，处理结果数组的填充（初始零）。
对你的原始函数做了一个简短的测试，显示了一些不同的结果，这可能是因为你的np.flip调用颠倒了索引数组。我刚刚实现了你提供的公式，没有深入检查索引，所以我把这个任务留给你。你可能应该用一些你可以手动检查的虚拟输入来调用它。
作为代码的附加注解：如果你想在index上循环，你应该使用range中的构建，而不是使用np.arange。在内部，python使用range的生成器，而不是先创建索引数组，然后逐个迭代每个索引。此外，你应该尝试减少需要创建的数组的数量，而是重用它们。例如，omega = np.zeros(n)可以在外部for循环外部使用omega = np.empty(n)创建，并且只在每次新的迭代omega[:] = 0.0时在内部初始化。注意，除了通过索引访问数组元素之外，所有类型的内存管理通常都会影响速度，您需要自己处理numpy，因为没有编译器可以帮助你，所以我推荐使用numba，它可以编译你的python代码，并在很多方面帮助你使你的数字运算更快。

赞(0）回复(0）举报 2022-12-29

我来回答

numpy 有没有一种方法可以加速或向量化嵌套的for循环？

1条答案

相关问题

热门标签

最新问答