尝试将 for 循环转换为 numpy 矩阵运算

Question

这是一个有点奇怪的问题，但我试图通过丢失 for 循环并依赖于 numpy 操作来提高我的代码的运行时间。我仍然是处理 numpy 矩阵运算的初学者，我似乎无法正确翻译它。

for i in range (Ma):

    for j in range (Na):
    
        for h in range(Mt):
        
            for k in range(Nt):

                dx = xrfa[i] - xrft[h]
                dy = yrfa[j] - yrft[k] 

                Wat[i,j,h,k] = at * np.exp(- ((np.square(dx) + np.square(dy))/ (2  * np.square(sigat)

任何帮助将不胜感激。我特别挣扎于 dx 和 dy。

Answer 1

依靠广播的力量！

设置：

import numpy as np

Ma, Na, Mt, Nt = 2, 3, 4, 5

# Assuming these are all just 1D arrays
xrfa = np.arange(Ma)
yrfa = np.arange(Na)
xrft = np.arange(Mt)
yrft = np.arange(Nt)

# and these scalars:
at = 2
sigat = 3

你可以（ab）使用

np.ix_

来引入单例维度，但我会在这里明确地做到这一点：

xrfa = xrfa[:, None, None, None]
yrfa = yrfa[None, :, None, None]
xrft = xrft[None, None, :, None]
yrft = yrft[None, None, None, :]

现在你的数组全部广播到

(Ma, Na, Mt, Nt)

的形状，此时你可以使用一个衬里来进行所有计算（希望我正确地修复了你的括号）：

Wat = at * np.exp(-(np.square(xrfa - xrft) + np.square(yrfa - yrft)) / (2  * np.square(sigat)))