如何使用numba加快以下代码的速度？
Posted
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了如何使用numba加快以下代码的速度？相关的知识，希望对你有一定的参考价值。
我正在做分子动力学模拟。它由数值积分组成，许多积分用于循环，操纵大型NumPy数组。我试图尽可能地使用NumPy函数和数组。但是代码仍然太慢。我想到了使用numba jit作为加速器。但是它总是抛出错误消息。
这里是代码。
# -*- coding: utf-8 -*-
"""
Created on Sat Mar 28 12:10:42 2020
@author: Sandipan
"""
import numpy as np
import matplotlib.pyplot as plt
from numba import jit
import os
import sys
# Setting up the simulation
NSteps =100 # Number of steps
deltat = 0.005 # Time step in reduced time units
temp = 0.851# #Reduced temperature
DumpFreq = 100 # Save the position to file every DumpFreq steps
epsilon = 1.0 # LJ parameter for the energy between particles
DIM =3
N =500
density =0.776
Rcutoff =3
#----------------------Function Definitions---------------------
#------------------Initialise Configuration--------
@jit(nopython=True)
def initialise_config(N,DIM,density):
velocity = (np.random.randn(N,DIM)-0.5)
# Set initial momentum to zero
COM_V = np.sum(velocity)/N #Center of mass velocity
velocity = velocity - COM_V # Fix any center-of-mass drift
# Calculate initial kinetic energy
k_energy=0
for i in range (N):
k_energy+=np.dot(velocity[i],velocity[i])
vscale=np.sqrt(DIM*temp/k_energy)
velocity*=vscale
#Initialize with zeroes
coords = np.zeros([N,DIM]);
# Get the cooresponding box size
L = (N/density)**(1.0/DIM)
""" Find the lowest perfect cube greater than or equal to the number of
particles"""
nCube = 2
while (nCube**3 < N):
nCube = nCube + 1
# Assign particle positions
ip=-1
x=0
y=0
z=0
for i in range(0,nCube):
for j in range(0,nCube):
for k in range(0,nCube):
if(ip<N):
x=(i+0.5)*(L/nCube)
y=(j+0.5)*(L/nCube)
z=(k+0.5)*(L/nCube)
coords[ip]=np.array([x,y,z])
ip=ip+1
else:
break
return coords,velocity,L
@jit(nopython=True)
def wrap(pos,L):
'''Apply perodic boundary conditions.'''
for i in range (len(pos)):
for k in range(DIM):
if (pos[i][k]>0.5):
pos[i][k]=pos[i][k]-1
if (pos[i][k]<-0.5):
pos[i][k]=pos[i][k]+1
return (pos)
@jit(nopython=True)
def LJ_Forces(pos,acc,epsilon,L,DIM,N):
# Compute forces on positions using the Lennard-Jones potential
# Uses double nested loop which is slow O(N^2) time unsuitable for large systems
Sij = np.zeros(DIM) # Box scaled units
Rij = np.zeros(DIM) # Real space units
#Set all variables to zero
ene_pot = np.zeros(N)
acc = acc*0
virial=0.0
# Loop over all pairs of particles
for i in range(N-1):
for j in range(i+1,N): #i+1 to N ensures we do not double count
Sij = pos[i]-pos[j] # Distance in box scaled units
for l in range(DIM): # Periodic interactions
if (np.abs(Sij[l])>0.5):
Sij[l] = Sij[l] - np.copysign(1.0,Sij[l]) # If distance is greater than 0.5 (scaled units) then subtract 0.5 to find periodic interaction distance.
Rij = L*Sij # Scale the box to the real units in this case reduced LJ units
Rsqij = np.dot(Rij,Rij) # Calculate the square of the distance
if(Rsqij < Rcutoff**2):
# Calculate LJ potential inside cutoff
# We calculate parts of the LJ potential at a time to improve the efficieny of the computation (most important for compiled code)
rm2 = 1.0/Rsqij # 1/r^2
rm6 = rm2**3
forcefact=(rm2**4)*(rm6-0.5) # 1/r^6
phi =4*(rm6**2-rm6)
ene_pot[i] = ene_pot[i]+0.5*phi # Accumulate energy
ene_pot[j] = ene_pot[j]+0.5*phi # Accumulate energy
virial = virial-forcefact*Rsqij # Virial is needed to calculate the pressure
acc[i] = acc[i]+forcefact*Sij # Accumulate forces
acc[j] = acc[j]-forcefact*Sij # (Fji=-Fij)
return 48*acc, np.sum(ene_pot)/N, -virial/DIM # return the acceleration vector, potential energy and virial coefficient
@jit(nopython=True)
def Calculate_Temperature(vel,L,DIM,N):
ene_kin = 0.0
for i in range(N):
real_vel = L*vel[i]
ene_kin = ene_kin + 0.5*np.dot(real_vel,real_vel)
ene_kin_aver = 1.0*ene_kin/N
temperature = 2.0*ene_kin_aver/DIM
return ene_kin_aver,temperature
# Main MD loop
@jit(nopython=True)
def main():
# Vectors to store parameter values at each step
ene_kin_aver = np.ones(NSteps)
ene_pot_aver = np.ones(NSteps)
temperature = np.ones(NSteps)
virial = np.ones(NSteps)
pressure = np.ones(NSteps)
pos,vel,L = initialise_config(N,DIM,density)
acc = (np.random.randn(N,DIM)-0.5)
volume=L**3
# Open file which we will save the outputs to
if os.path.exists('energy2'):
os.remove('energy2')
f = open('traj.xyz', 'w')
for k in range(NSteps):
# Refold positions according to periodic boundary conditions
pos=wrap(pos,L)
# r(t+dt) modify positions according to velocity and acceleration
pos = pos + deltat*vel + 0.5*(deltat**2.0)*acc # Step 1
# Calculate temperature
ene_kin_aver[k],temperature[k] = Calculate_Temperature(vel,L,DIM,N)
# Rescale velocities and take half step
chi = np.sqrt(temp/temperature[k])
vel = chi*vel + 0.5*deltat*acc # v(t+dt/2) Step 2
# Compute forces a(t+dt),ene_pot,virial
acc, ene_pot_aver[k], virial[k] = LJ_Forces(pos,acc,epsilon,L,DIM,N) # Step 3
# Complete the velocity step
vel = vel + 0.5*deltat*acc # v(t+dt/2) Step 4
# Calculate temperature
ene_kin_aver[k],temperature[k] = Calculate_Temperature(vel,L,DIM,N)
# Calculate pressure
pressure[k]= density*temperature[k] + virial[k]/volume
# Print output to file every DumpFreq number of steps
if(k%DumpFreq==0): # The % symbol is the modulus so if the Step is a whole multiple of DumpFreq then print the values
f.write("%s
" %(N)) # Write the number of particles to file
# Write all of the quantities at this step to the file
f.write("Energy %s, Temperature %.5f
" %(ene_kin_aver[k]+ene_pot_aver[k],temperature[k]))
for n in range(N): # Write the positions to file
f.write("X"+" ")
for l in range(DIM):
f.write(str(pos[n][l]*L)+" ")
f.write("
")
if (k%5==0):
# print("
Step: {0} KE: {1} PE: {2} Energy: {3}".format(k, ene_kin_aver[k], ene_pot_aver[k],ene_kin_aver[k]+ene_pot_aver[k]))
sys.stdout.write("
Step: {0} KE: {1} PE: {2} Energy: {3}".format(k, ene_kin_aver[k], ene_pot_aver[k],ene_kin_aver[k]+ene_pot_aver[k]))
sys.stdout.flush()
return ene_kin_aver, ene_pot_aver, temperature, pressure, pos
#------------------------------------------------------
ene_kin_aver, ene_pot_aver, temperature, pressure, pos = main()
# Plot all of the quantities
def plot():
plt.figure(figsize=[7,12])
plt.rc('xtick', labelsize=15)
plt.rc('ytick', labelsize=15)
plt.subplot(4, 1, 1)
plt.plot(ene_kin_aver,'k-')
plt.ylabel(r"$E_{K}", fontsize=20)
plt.subplot(4, 1, 2)
plt.plot(ene_pot_aver,'k-')
plt.ylabel(r"$E_{P}$", fontsize=20)
plt.subplot(4, 1, 3)
plt.plot(temperature,'k-')
plt.ylabel(r"$T$", fontsize=20)
plt.subplot(4, 1, 4)
plt.plot(pressure,'k-')
plt.ylabel(r"$P$", fontsize=20)
plt.show()
plot()
我得到的错误是：
runfile('E:/Project/LJMD4.py', wdir='E:/Project')
Traceback (most recent call last):
File "<ipython-input-8-aeebce887079>", line 1, in <module>
runfile('E:/Project/LJMD4.py', wdir='E:/Project')
File "C:UsersSandipanAnaconda3libsite-packagesspyder_kernelscustomizespydercustomize.py", line 827, in runfile
execfile(filename, namespace)
File "C:UsersSandipanAnaconda3libsite-packagesspyder_kernelscustomizespydercustomize.py", line 110, in execfile
exec(compile(f.read(), filename, 'exec'), namespace)
File "E:/Project/LJMD4.py", line 226, in <module>
ene_kin_aver, ene_pot_aver, temperature, pressure, pos = main()
File "C:UsersSandipanAnaconda3libsite-packages
umbadispatcher.py", line 351, in _compile_for_args
error_rewrite(e, 'typing')
File "C:UsersSandipanAnaconda3libsite-packages
umbadispatcher.py", line 318, in error_rewrite
reraise(type(e), e, None)
File "C:UsersSandipanAnaconda3libsite-packages
umbasix.py", line 658, in reraise
raise value.with_traceback(tb)
TypingError: cannot determine Numba type of <class 'builtin_function_or_method'>
[当我在互联网上搜索时，我发现numba可能不支持我正在使用的某些功能。但是我没有使用任何熊猫或其他数据框。我只是使用纯python循环或NumPy，据numba文档建议，它得到了很好的支持。我试图从某些函数中删除numba并使nopython = 0，但是它仍然发送不同的错误消息。我不知道这是怎么回事。没有numba的代码将无法实际使用。提速的任何其他技巧将有很大帮助。预先谢谢你。
答案
一些常见错误
使用不受支持的功能
文件操作，许多字符串操作。这些可以在objmode block中。在此示例中，我将这些内容注释掉。
初始化数组的错误方法
仅支持元组，不支持列表（Numpy接受两者，但文档中仅提及元组）
检查是否被零除并抛出异常
这是Python的标准行为，但不是Numpy。如果要避免这种开销（如果在每个分区上均是/否则），请打开Numpy默认行为（error_model =“ numpy”）。
使用全局变量
全局变量被硬编码到已编译的代码中（就像您将它们直接写入代码中一样）。不重新编译就无法更改它们。
Numpy数组的错误索引
pos[i][k]而不是pos[i,k]。 Numba可能会对此进行优化，但这在Pure Python代码中具有相当明显的负面影响。
工作版本
# -*- coding: utf-8 -*-
"""
Created on Sat Mar 28 12:10:42 2020
@author: Sandipan
"""
import numpy as np
import matplotlib.pyplot as plt
from numba import jit
import os
import sys
# All globals are compile time constants
# recompilation needed if you change this values
# Better way: hand a tuple of all needed vars to the functions
# params=(NSteps,deltat,temp,DumpFreq,epsilon,DIM,N,density,Rcutoff)
# Setting up the simulation
NSteps =100 # Number of steps
deltat = 0.005 # Time step in reduced time units
temp = 0.851# #Reduced temperature
DumpFreq = 100 # Save the position to file every DumpFreq steps
epsilon = 1.0 # LJ parameter for the energy between particles
DIM =3
N =500
density =0.776
Rcutoff =3
params=(NSteps,deltat,temp,DumpFreq,epsilon,DIM,N,density,Rcutoff)
#----------------------Function Definitions---------------------
#------------------Initialise Configuration--------
#error_model=True
#Do you really want to search for division by zeros (additional cost)?
@jit(nopython=True,error_model="numpy")
def initialise_config(N,DIM,density):
velocity = (np.random.randn(N,DIM)-0.5)
# Set initial momentum to zero
COM_V = np.sum(velocity)/N #Center of mass velocity
velocity = velocity - COM_V # Fix any center-of-mass drift
# Calculate initial kinetic energy
k_energy=0
for i in range (N):
k_energy+=np.dot(velocity[i],velocity[i])
vscale=np.sqrt(DIM*temp/k_energy)
velocity*=vscale
#wrong array initialization (use tuple)
#Initialize with zeroes
coords = np.zeros((N,DIM))
# Get the cooresponding box size
L = (N/density)**(1.0/DIM)
""" Find the lowest perfect cube greater than or equal to the number of
particles"""
nCube = 2
while (nCube**3 < N):
nCube = nCube + 1
# Assign particle positions
ip=-1
x=0
y=0
z=0
for i in range(0,nCube):
for j in range(0,nCube):
for k in range(0,nCube):
if(ip<N):
x=(i+0.5)*(L/nCube)
y=(j+0.5)*(L/nCube)
z=(k+0.5)*(L/nCube)
coords[ip]=np.array([x,y,z])
ip=ip+1
else:
break
return coords,velocity,L
@jit(nopython=True)
def wrap(pos,L):
'''Apply perodic boundary conditions.'''
#correct array indexing
for i in range (len(pos)):
for k in range(DIM):
if (pos[i,k]>0.5):
pos[i,k]=pos[i,k]-1
if (pos[i,k]<-0.5):
pos[i,k]=pos[i,k]+1
return (pos)
@jit(nopython=True,error_model="numpy")
def LJ_Forces(pos,acc,epsilon,L,DIM,N):
# Compute forces on positions using the Lennard-Jones potential
# Uses double nested loop which is slow O(N^2) time unsuitable for large systems
Sij = np.zeros(DIM) # Box scaled units
Rij = np.zeros(DIM) # Real space units
#Set all variables to zero
ene_pot = np.zeros(N)
acc = acc*0
virial=0.0
# Loop over all pairs of particles
for i in range(N-1):
for j in range(i+1,N): #i+1 to N ensures we do not double count
Sij = pos[i]-pos[j] # Distance in box scaled units
for l in range(DIM): # Periodic interactions
if (np.abs(Sij[l])>0.5):
Sij[l] = Sij[l] - np.copysign(1.0,Sij[l]) # If distance is greater than 0.5 (scaled units) then subtract 0.5 to find periodic interaction distance.
Rij = L*Sij # Scale the box to the real units in this case reduced LJ units
Rsqij = np.dot(Rij,Rij) # Calculate the square of the distance
if(Rsqij < Rcutoff**2):
# Calculate LJ potential inside cutoff
# We calculate parts of the LJ potential at a time to improve the efficieny of the computation (most important for compiled code)
rm2 = 1.0/Rsqij # 1/r^2
rm6 = rm2**3
forcefact=(rm2**4)*(rm6-0.5) # 1/r^6
phi =4*(rm6**2-rm6)
ene_pot[i] = ene_pot[i]+0.5*phi # Accumulate energy
ene_pot[j] = ene_pot[j]+0.5*phi # Accumulate energy
virial = virial-forcefact*Rsqij # Virial is needed to calculate the pressure
acc[i] = acc[i]+forcefact*Sij # Accumulate forces
acc[j] = acc[j]-forcefact*Sij # (Fji=-Fij)
#If you want to get get the best performance, sum directly in the loop intead of
#summing at the end np.sum(ene_pot)
return 48*acc, np.sum(ene_pot)/N, -virial/DIM # return the acceleration vector, potential energy and virial coefficient
@jit(nopython=True,error_model="numpy")
def Calculate_Temperature(vel,L,DIM,N):
ene_kin = 0.0
for i in range(N):
real_vel = L*vel[i]
ene_kin = ene_kin + 0.5*np.dot(real_vel,real_vel)
ene_kin_aver = 1.0*ene_kin/N
temperature = 2.0*ene_kin_aver/DIM
return ene_kin_aver,temperature
# Main MD loop
@jit(nopython=True,error_model="numpy")
def main(params):
NSteps,deltat,temp,DumpFreq,epsilon,DIM,N,density,Rcutoff=params
# Vectors to store parameter values at each step
ene_kin_aver = np.ones(NSteps)
ene_pot_aver = np.ones(NSteps)
temperature = np.ones(NSteps)
virial = np.ones(NSteps)
pressure = np.ones(NSteps)
pos,vel,L = initialise_config(N,DIM,density)
acc = (np.random.randn(N,DIM)-0.5)
volume=L**3
# Open file which we will save the outputs to
# Unsupported operations have to be in an objectmode block
# or simply write the outputs at the end in a pure Python function
"""
if os.path.exists('energy2'):
os.remove('energy2')
f = open('traj.xyz', 'w')
"""
for k in range(NSteps):
# Refold positions according to periodic boundary conditions
pos=wrap(pos,L)
# r(t+dt) modify positions according to velocity and acceleration
pos = pos + deltat*vel + 0.5*(deltat**2.0)*acc # Step 1
# Calculate temperature
ene_kin_aver[k],temperature[k] = Calculate_Temperature(vel,L,DIM,N)
# Rescale velocities and take half step
chi = np.sqrt(temp/temperature[k])
vel = chi*vel + 0.5*deltat*acc # v(t+dt/2) Step 2
# Compute forces a(t+dt),ene_pot,virial
acc, ene_pot_aver[k], virial[k] = LJ_Forces(pos,acc,epsilon,L,DIM,N) # Step 3
# Complete the velocity step
vel = vel + 0.5*deltat*acc # v(t+dt/2) Step 4
# Calculate temperature
ene_kin_aver[k],temperature[k] = Calculate_Temperature(vel,L,DIM,N)
# Calculate pressure
pressure[k]= density*temperature[k] + virial[k]/volume
# Print output to file every DumpFreq number of steps
"""
if(k%DumpFreq==0): # The % symbol is the modulus so if the Step is a whole multiple of DumpFreq then print the values
f.write("%s
" %(N)) # Write the number of particles to file
# Write all of the quantities at this step to the file
f.write("Energy %s, Temperature %.5f
" %(ene_kin_aver[k]+ene_pot_aver[k],temperature[k]))
for n in range(N): # Write the positions to file
以上是关于如何使用numba加快以下代码的速度？的主要内容，如果未能解决你的问题，请参考以下文章