如何将嵌套循环简化为割炬张量操作？

Question

I'm trying to convert some code I have written in numpy which contains a nested-loop into tensor operations found in PyTorch. However, after trying to implement my own version I'm not getting the same value on the output. I have managed to do the same with a single loop, so I'm not entirely sure what I'm doing wrong.

    #(Numpy Version)
    #calculate Kinetic Energy
    summation = 0.0
    for i in range(0,len(k_values)-1):
        summation += (k_values[i]**2.0)*wavefp[i]*(((self.hbar*kp_values[i])**2.0)/(2.0*self.mu))*wavef[i]
    Ek = step*(4.0*np.pi)*summation

    #(Numpy Version)
    #calculate Potential Energy
    summation = 0.0
    for i in range(0,len(k_values)-1):
        for j in range(0,len(kp_values)-1):
            summation+= (k_values[i]**2.0)*wavefp[i]*(kp_values[j]**2.0)*wavef[j]*self.MTV[i,j]
    Ep = (step**2.0)*(4.0*np.pi)*(2.0/np.pi)*summation

            #####################################################

    #(PyTorch Version)
    #calcualte Kinetic Energy
    Ek = step*(4.0*np.pi)*torch.sum( k_values.pow(2)*wavefp.mul(wavef)*((kp_values.mul(self.hbar)).pow(2)/(2.0*self.mu)) )

    #(PyTorch Version)
    #calculate Potential Energy
    summation = 0.0
    for i in range(0,len(k_values)-1):
        summation += ((k_values[i].pow(2)).mul(wavefp[i]))*torch.sum( (kp_values.pow(2)).mul(wavef).mul(self.MTV[i,:]) )
    Ep = (step**2.0)*(4.0*np.pi)*(2.0/np.pi)*summation

The arrays/tensors k_values, kp_values, wavef, and wavefp have dimensions of (1000,1). The values self.hbar, and self.mu, and step are scalars. The variable self.MTV is a matrix of size (1000,1000).

I would expect that both methods would give the same output but they don't. The code for calculating the Kinetic Energy (in both Numpy and PyTorch) give the same value. However, the potential energy calculation differ, and I'm not entirely sure why.

Many Thanks in advance!

Answer 1

The problem is in the shapes. You have kp_values and wavef in (1000, 1) which needs to be converted to (1000, ) before the multiplications. The outcome of (kp_values.pow(2)).mul(wavef).mul(MTV[i,:]) is a matrix but you asummed it is a vector.

So, the following should work.

summation += ((k_values[i].pow(2)).mul(wavefp[i]))*torch.sum((kp_values.squeeze(1)
    .pow(2)).mul(wavef.squeeze(1)).mul(MTV[i,:]))

---

And a loop-free Numpy and PyTorch solution would be:

step = 1.0
k_values = np.random.randint(0, 100, size=(1000, 1)).astype("float") / 100
kp_values = np.random.randint(0, 100, size=(1000, 1)).astype("float") / 100
wavef = np.random.randint(0, 100, size=(1000, 1)).astype("float") / 100
wavefp = np.random.randint(0, 100, size=(1000, 1)).astype("float") / 100
MTV = np.random.randint(0, 100, size=(1000, 1000)).astype("float") / 100

# Numpy solution
term1 = k_values**2.0 * wavefp # 1000 x 1 
temp = kp_values**2.0 * wavef # 1000 x 1
term2 = np.matmul(temp.transpose(1, 0), MTV).transpose(1, 0) # 1000 x 1000
summation = np.sum(term1 * term2)
print(summation)

# PyTorch solution
term1 = k_values.pow(2).mul(wavefp) # 1000 x 1
term2 = kp_values.pow(2).mul(wavef).transpose(0, 1).matmul(MTV) # 1000 x 1000
summation = torch.sum(term2.transpose(0, 1).mul(term1)) # 1000 x 1000
print(summation.item())

Output

12660.407492918514
12660.407492918514