在python中为依赖于索引的函数向量化嵌套的for循环

Posted 2023-02-23

【中文标题】在python中为依赖于索引的函数向量化嵌套的for循环

【英文标题】：Vectorizing a nested for-loop in python for index-dependent function

【发布时间】：2014-03-17 16:26:01

【问题描述】：

我目前正在使用 Numpy 数组将 C++ 程序移植到 Python。如果可能的话，我正在寻找一种以更 Pythonic 的方式实现以下循环的方法：

for (int j = start_y; j < end_y; j++)
    
        for (int i = start_x; i < end_x; i++)
        
            plasmaFreq[i][j] = plasmaFreq_0*(tanh((i - 50)/10) - tanh((i - (nx - 50))/10))/2.0;
        
    

在上面，plasmaFreq_0 是传递给周围函数的常数，nx 也是如此。显然，将循环边界向量化以对 numpy 数组的特定区域进行操作很容易，但这给我留下了如何在数组中映射上述索引相关函数的问题。

【问题讨论】：

你对plasmaFreq[i][j]的计算目前没有参考j；对吗？

这个循环是正确的。它在程序的其他地方必然是一个二维数组，但在这里，我在 x 方向建立一个在 y 方向上恒定的轮廓。

【参考方案1】：

你需要一个数组i，

i = np.arange(start_x, end_x)
plasmaFreq[start_x:end_x, start_y: end_y] = plasmaFreq_0 *(np.tanh((i - 50)/10) - np.tanh((i - (nx - 50))/10))/2.0

我认为广播应该从那里开始。

---

请注意，您的原始代码效率很低¹...首先，您正在计算每个j 的右侧，但它不依赖于j，所以你只需要计算一次。其次，您的内部循环超出了慢速索引，因此您不会有效地使用缓存。我可能会写成：

for (int i = start_x; i < end_x; i++)

    rhs = plasmaFreq_0*(tanh((i - 50)/10) - tanh((i - (nx - 50))/10))/2.0;
    for (int j = start_y; j < end_y; j++)
    
        plasmaFreq[i][j] = rhs;
    

^{1效率如何取决于编译器在找出循环方面的能力。有一天，也许一些编译器可以从你的和我的生成相同的代码}

【讨论】：

这比直接实现循环更有效吗？

@Scott -- 通常，但在效率方面，通常最好忘记您认为自己知道的内容，而只需 timeit。即使是专业人士也经常对结果感到惊讶。