Как я могу упростить вложенный цикл в операции с тензором горелки? - PullRequest
Новая
       6

Как я могу упростить вложенный цикл в операции с тензором горелки?

0 голосов
/ 15 апреля 2019

Я пытаюсь преобразовать некоторый код, который я написал в numpy, который содержит вложенный цикл, в тензорные операции, найденные в PyTorch. Тем не менее, после попытки реализовать мою собственную версию я не получаю то же значение на выходе. Мне удалось сделать то же самое с помощью одного цикла, поэтому я не совсем уверен, что я делаю неправильно. 

    #(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

Массивы / тензоры k_values, kp_values, wavef и wavefp имеют размеры (1000,1). Значения self.hbar, self.mu и step являются скалярами. Переменная self.MTV представляет собой матрицу размера (1000, 1000).

Я ожидаю, что оба метода дадут одинаковый результат, но это не так. Код для расчета кинетической энергии (как в Numpy, так и в PyTorch) дает одинаковое значение. Тем не менее, расчеты потенциальной энергии отличаются, и я не совсем уверен, почему.

Большое спасибо заранее!

1 Ответ

1 голос
/ 16 апреля 2019

Проблема в формах. У вас есть kp_values и wavef в (1000, 1), которые необходимо преобразовать в (1000,) перед умножением. Результат (kp_values.pow(2)).mul(wavef).mul(MTV[i,:]) - это матрица, но вы предположили, что это вектор.

Итак, должно работать следующее. 

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

И решение Numpy and PyTorch без петель будет: 

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())

выход 

12660.407492918514
12660.407492918514
...