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

Question

我目前正在使用 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 数组的特定区域进行操作很容易，但这给我留下了如何在数组中映射上述索引相关函数的问题。

Answer 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效率如何取决于编译器在找出循环方面的能力。有一天，也许一些编译器可以从你的和我的生成相同的代码}