问numba函数中的这个循环可以优化为运行得更快吗？

EN

import numba as nb
import numpy as np
import time 
#######create some synthetic data for illustration purpose##################
size=5000
eps = 0.2
theta_c = 0.4
temp = np.ones(size)
neighbour = np.random.randint(size, size=(size, 3)) 
coschi = np.random.random_sample((size))
theta = np.random.random_sample((size))*np.pi/2
pwr = np.cos(theta)
###################end of dummy data##########################

###################-----main loop------###############
@nb.jit(fastmath=True)
def func(theta, pwr, neighbour, coschi, temp):
    for k in range(np.argmax(pwr), 5000*(pwr.size)): 
        n = k%pwr.size
        if (np.abs(theta[n]-np.pi/2.)<np.abs(theta_c)):
                adj = neighbour[n,1]
        else:
                adj = neighbour[n,0]
        psi_diff = np.abs(np.arccos(coschi[adj])-np.arccos(coschi[n]))
        temp5 = temp[adj]**5;
        e_temp = 1.- np.exp(-temp5*psi_diff/np.abs(eps))
        temp[n] = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)
    return temp

#check time
time1 = time.time()
temp = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, " seconds.")