加快numpy阵列的分析

Question

我有一个python代码,它导入4列txt文件,数字前三列是x,y,z坐标,第四列是该坐标的密度.

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

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

我认为问题是最后的部分,但我无法弄清楚优化它的方法.

我有一个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)

这是一个非常快速的操作,只需要几分之一毫秒.形状可能是错误的,所以你可能不得不改变参数的顺序,用转座做一些事情,但基本的想法肯定存在.您可以在文档中阅读它.

复制数据很糟糕

您不需要复制所有内容,当您需要复制数组(或数组的一部分)时,您应该让numpy为您执行此操作.例如,Firstdel只需使用,而不是您的功能a[1:].或者,如果您确实需要复制数据(不仅仅是用于绘图),请使用以下命令:

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

但一般来说,你可以只使用numpy数组的"切片",而不是复制它们.在这里阅读它.

循环

循环也是臭名昭着的浪费时间.首先,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存在的原因是将这些缓慢的python循环转换为快速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,等等.然后,您可以使用高级索引来设置值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倍.