加快对numpy数组的分析

Question

我有一个python代码，它导入4列txt文件，其中第 前三列是x，y，z坐标，第四列是该坐标下的密度。

下面是读取代码，转换为ndarray，傅立叶变换该字段，计算原点距离（k =（0,0,0））和变换后的坐标，并取平均值并绘制它们。 感谢大熊猫（用于数据分析的python库）和python FFT，加载256^3行并对它们进行傅里叶变换非常快，并在几秒钟内完成。

但是，将加载的txt转换为numpy ndarray，计算平均密度（每个坐标的平均值）以及计算与原点的距离（k =（0,0,0））需要很长时间。

我认为问题是在最后np.around部分，但我不知道如何优化它。

我有一个32核心机器的资源。

有人可以教我如何加速，使其成为一个多进程代码或类似的东西，以便这些可以很快完成？谢谢。

（如果你是宇宙学家和以往任何时候都需要这个代码，你可以，如果你可以的。谢谢使用它，但请与我联系，）

from __future__ import division 
import numpy as np 

ngridx = 128 
ngridy = 128  
ngridz = 128 

maxK = max(ngridx,ngridy,ngridz) 

#making input file 
f = np.zeros((ngridx*ngridy*ngridz,4)) 

i = 0 
for i in np.arange(len(f)): 
    f[i][0] = int(i/(ngridy*ngridz)) 
    f[i][1] = int((i/ngridz))%ngridy 
    f[i][2] = int(i%ngridz) 
    f[i][3] = np.random.rand(1) 
    if i%1000000 ==0: 
     print i 
#This takes forever 
#end making input file 

#Thanks to Mike, 
a = f[:,3].reshape(ngridx,ngridy,ngridz) 

avg =np.sum(f[:,3])/len(f) 
a /= avg 
p = np.fft.fftn(a) 
#This part is much much faster than before (Original Post). 

#Keeping track of corresponding wavenumbers (k_x, k_y,k_z) for each element in p 
#This is just a convension on fourier transformation so you can ignore this part 
kValx = np.fft.fftfreq(ngridx , (1.0/ngridx)) 
kValy = np.fft.fftfreq(ngridy , (1.0/ngridy)) 
kValz = np.fft.fftfreq(ngridz , (1.0/ngridz)) 
kx = np.zeros((ngridx,ngridy,ngridz)) 
ky = np.zeros((ngridx,ngridy,ngridz)) 
kz = np.zeros((ngridx,ngridy,ngridz)) 
rangecolx = np.arange(ngridx) 
rangecoly = np.arange(ngridy) 
rangecolz = np.arange(ngridz) 
for row in np.arange(ngridx): 
    for column in np.arange(ngridy): 
     for height in np.arange(ngridz): 
      kx[row][column][height] = (kValx[row]) 
      ky[row][column][height] = (kValy[column]) 
      kz[row][column][height] = (kValz[height]) 
    if row%10 == 0: 
     print row 
print 'wavenumber generate complete!' 

#Calculating the average powerspectrum in terms of fixed K (Distance from origin to a point in fourier space) 
#by taking the spherical shell of thickness 1 and averaging out the values inside it. 
#I am sure that this process can be optimised somehow, but I gave up. 

qlen = maxK/2 #Nyquist frequency 
q = np.zeros(((qlen),4),dtype=complex) 
#q is a four column array with length maxK/2. 
#q[:,0] is integer wavenumber (K, which is the distance from the origin = sqrt(kx^2+ky^2+kz^2)) 
#q[:,1] is the sum of square of the fourier transformed value 
#q[:,2] is the sum of the fourier transformed value, 
#and q[:,3] is the total number of samples with K=q[:,0] 

for i in np.arange(len(q)): 
    q[i][0] = i 
i = 0 
for i in np.arange(len(p)): 
    for r in np.arange(len(p[0])): 
     for s in np.arange(len(p[0,0])): 
      K = np.around(np.sqrt(kx[i,r,s]**2+ky[i,r,s]**2+kz[i,r,s]**2)) 
      if K < qlen: 
       q[K][1]=q[K][1]+np.abs(p[i,r,s])**2 
       q[K][2]=q[K][2]+p[i,r,s] 
       q[K][3]=q[K][3]+1 
    if i%10 ==0: 
     print 'i = ',i,' !' 
print q 

Answer 1

numpy的通常可以做的事情数百时间比普通的Python快，只需很少的额外努力。你只需要知道编写代码的正确方法。仅举的第一件事情我想：

索引

普通的Python是经常的事情，电脑应该是很大的很慢。一个例子是索引，所以像

a[f[i,0]][f[i,1]][f[i,2]]=f[i,3] 

让我非常怀疑。当你说“将加载的txt转换为numpy ndarray”需要很长时间时，这就是你指的那个吗？这并不会让我感到意外，因为每次python看到a[f[i,0]]时，它必须首先索引f，确保i是一个整数，并且您还没有跑出f的边缘;那么它必须确保f[i,0]是一个整数，并且您还没有跑出a的边缘。然后在它知道你要设置哪个元素之前，必须重复这两次。

一个改进是使用a[f[i,0],f[i,1],f[i,2]]，因为在这种索引下numpy速度更快。

但我假设你的数据实际上是某种顺序。例如，f[i,2]是否从0循环到256，然后f[i,1]增加1，并且f [i，2]从0开始？如果是这样，你真正需要做的是这样说

a = f[:,3].reshape(ngridx,ngridy,ngridz) 

这是一个可笑的快速操作，以一毫秒的一小部分。形状可能是错误的，所以你可能不得不改变参数的顺序，对转置做些事情，但基本的想法肯定是存在的。您可以在the documentation中阅读。

复制数据是很糟糕

你并不需要复制的一切，当你需要复制（数组或部分）的数组，你应该让numpy的为你做它。例如，您可以使用a[1:]而不是Firstdel函数。或者，如果你真的需要使数据的副本（你只是为了绘制不）使用：

def Firstdel(a): 
    return numpy.copy(a[1:]) 

但在一般情况下，你可以只使用numpy的阵列的“切片”，而不是复制它们。阅读关于它here。

循环

循环也臭名昭着的时间浪费。首先，while在Python的简单循环中并不常见。因此，而不是

while i < len(f): 
    # do stuff 
    i = i+1 

你应该使用

for i in range(len(f)): 
    # do stuff 

摆脱尽可能多的循环，你可以。要设置kx，ky，并kz，该代码比嵌套循环快约10倍，但秤如N代替N-^3（其中N = ngridx ngridy ngridz）：

for row in range(ngridx): 
    kx[row,:,:] = kValx[row] 
for column in range(ngridy): 
    ky[:,column,:] = kValy[column] 
for height in range(ngridz): 
    kz[:,:,height] = kValz[height] 

切片可以对于设置值也很有用，因为循环进入了numpy。取而代之的是代码

i = 0 
while i < len(q): 
    q[i][0] = i 
    i = i + 1 

只使用

q[:,0] = range(len(q)) 

在这里，我只是设置的q“切片”等于另一个数组。

该循环之后的嵌套循环也可以加速，但它们会更复杂一些。

但是你也希望尽可能避免循环。这给我们带来...

使用内置numpy的功能

numpy的存在的原因是把这些蟒蛇缓慢环路和成等快速C代码（我们并不需要知道其存在）。所以有很多功能可以完成你想要做的事情，而这些功能已经内置在numpy中。

而不是

while i < len(f): 
    masstot = masstot + f[i,3] 
    i = i+1 

你应该使用类似

masstot = np.sum(f[:,3]) 

所以，很简单阅读，但也很可能是方式更快，因为numpy的必须在该数据的直接访问计算机的内存，并可以使用快速的C函数来查找总和，而不是使用慢python函数。 （再一次，你不需要知道任何关于C函数的东西，他们只会完成这项工作。）

代替K表明大嵌套循环计算值通过循环每一次，只是使K用适当的值的数组：

K = np.around(np.sqrt(kx**2+ky**2+kz**2)) 

K将是大小相同kx等。然后，您可以使用advanced indexing来设置q的值。这是我会怎么做，最后一节：

# Again, we get rid of nested loops, to get a large improvement in speed and scaling 
K = np.around(np.sqrt(kx**2+ky**2+kz**2)).astype(int) 
for i in range(qlen): 
    indices = (K==i) # This will be an array of True/False values, 
        # which will be used for "advanced indexing" of p 
    q[i,0] = i 
    q[i,1] = sum(np.abs(p[indices])**2) 
    q[i,2] = sum(p[indices]) 
    q[i,3] = sum(indices) 
print q 

把这些都在一起，我得到的35提高的因素在代码目前在你的问题。

Answer 2

还不如加快输入文件创建以及：

size = ngridx*ngridy*ngridz 
f = np.zeros((size,4)) 
a = np.arange(size) 
f[:, 0] = np.floor_divide(a, ngridy*ngridz) 
f[:, 1] = np.fmod(np.floor_divide(a, ngridz), ngridy) 
f[:, 2] = np.fmod(a, ngridz) 
f[:, 3] = np.random.rand(size) 

为了kx，ky，并kz，你可以摆脱循环使用broadcasting：

kx += kValx[:, np.newaxis, np.newaxis] 
ky += kValy[np.newaxis, :, np.newaxis] 
kz += kValz[np.newaxis, np.newaxis, :]