如何使用numpy将两个嵌套for循环矢量化？

Question

我想在下面对我的部分代码进行矢量化。特别是从range(0,n)迭代的两个for循环。如果有什么问题的话，我可以将j循环矢量化，因为当我在操作中包含它时，运行时间超过10分钟

我在想，对于称为wavea和waveb的指数部分，用np.reshape做一些事情，然后减去它们，这样我就可以在指数函数上输入它们，但这会给我一个矩阵，我不确定如何处理它

epsilon = 0.15
def delta(x):
    
    # Definining the delta dirac function as a Lorentzian distribution 
    # with a width epsilon
    
    return (1/np.pi)*(epsilon)/(epsilon**2+x**2)

# Doing it inefficiently at first

im = complex(0,1)
n = 244      
kdots = 1575      


def quantity_abs_squared(Vecs, Evalue, Eigen, rx,ry):
    holder = 0
    
    for k in range(0,kdots):
        
        for i in range(0,n):

                ind_G2a = G2*hexind[i,0] #G1, G2 are fixed (1,2) vectors
                ind_G1a = G1*hexind[i,1] #hexind is a list of (1,2) tuples of length n

                vectora = Vecs[k][:,i]    #Vecs has shape (1575, 244, 244) 

                wavea = np.exp(im*(ind_G1a[0] + ind_G2a[0])*rx + im*(ind_G1a[1] + ind_G2a[1])*ry)


            for j in range(0,n):

                ind_G2b = G2*hexind[j,0]
                ind_G1b = G1*hexind[j,1]

                # Extracting the eigenvector for the eigenvalues

                vectorb = Vecs[k][:,j]

                waveb = np.exp(im*(ind_G1b[0] + ind_G2b[0])*rx + im*(ind_G1b[1] + ind_G2b[1])*ry)

                coeffs = np.dot(vectora, np.conjugate(vectorb))

                exps = np.dot(wavea, np.conjugate(waveb))

                # For every kdots, there are n number of Eigenvalues
                # For every kdots, there are n number of Eigenvectors of length n

                holder = holder + coeffs*exps*delta((Egrid - Eigen[k][i]))

    quant = holder
    
    return (2*2)*quant/kdots

编辑：

我有一个想法，但我不知道如何实现它。最重要的是要去掉数组hexind的组合，它是（n，2）的形状，并且保存了我需要乘以G1和G2向量的数字。如果我做这样的事情会怎么样：

differences = hexind.reshape([n, 1, 2])-hexind.reshape([1,n,2])

differences现在是一个（n，n，2）矩阵，它保存了G1和G2部分所需的索引，但也提供了可以在指数内部输入的差异。同样，我不知道如何在没有for循环的情况下循环这个矩阵

编辑#2：

输入一个最小的可重复的例子。假设索引是以下数组：

G1 = np.random.rand((1,2))
G2 = np.random.rand((1,2))

index = np. array([[ 0,  0],
       [ 1,  1],
       [ 1,  0],
       [ 0, -1],
       [-1, -1],
       [-1,  0],
       [ 0,  1]])

然后，我按以下方式更改此索引：

hexind2 = np.repeat(index,2, axis=0)   # repeat two times for every row
hexind = np.tile(hexind2,(2,1))

然后我们可以为特征值和特征向量创建一些随机numpy数组：

kdots = 100
n = 28

E = np.random.rand(100, n)
vec = np.random.rand(100,n,n) + 1j*np.random.rand(100,n,n)

最后，我调用函数quantity_abs_squared，如下所示：

Eval = 300

rx = np.linspace(-20, 20)
ry = np.linspace(-20, 20)

Rx,Ry = np.meshgrid(rx,ry)

result = np.zeros((len(Rx), len(Ry)))


result = quantity_abs_squared(vec, Eval, E, Rx,Ry)