在numpy中加速数组分析

Question

我有一个Python代码，它导入了一个包含数字的四列文本文件。前三列是x、y、z坐标，第四列是该坐标处的密度。以下是读取代码，转换为ndarray数组，对该字段进行傅里叶变换，计算到原点（k=(0,0,0)）和一个转换坐标的距离，并取平均值，然后绘制它们的图表。由于pandas（用于数据分析的Python库）和Python FFT，加载256^3行并对其进行傅里叶变换非常快，并在几秒钟内完成。但是，将加载的文本转换为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

- Tom

---

请尝试将您的代码简化为更短但仍能展示运行缓慢问题的代码。您的代码长度远远超过了典型成功的SO问题长度。此外，请提供一个通过使用 np.random 生成有效输入的简短程序。 - John Zwinck

如果您能更具体地指出哪些部分较慢，那将非常有帮助。我认为我可以找出您所指的部分，但您应该清晰地指出它们，以确保没有人花费太多时间思考代码中错误的部分。 - Mike

b = a[1:] 应该可以摆脱第一个循环。a[tuple(f[:,:2])]=f[:,3] 可能适用于第二个; np.sum(f[:,3]) 适用于第三个。 - hpaulj

谢谢，我刚刚编辑并更新了最新的代码，并考虑了John和Mike的建议。 - Tom

Answer 1

Numpy通常可以比纯Python快数百倍，而且只需要很少的额外努力（有时甚至会自动使用多个核心，而无需您付出任何工作）。您只需要知道编写代码的正确方法。仅列举我能想到的第一件事：

• 思考您的方法 - 简化它并确保它正确无误，然后再编写代码
• 学习使用numpy的索引
• 尽可能避免复制数据
• 尽量避免使用循环
• 使用numpy的内置函数
• 整体处理数组（或大块的数组），而不是逐个元素处理
• 分析您的代码以找出慢点
• 使用numba，特别是对于您无法避免或numpy无法处理的循环

您的方法

我认为我们所有人都犯过的一个错误是仅仅试图用代码来表达思路，即使在我们仍在努力理解代码应该做什么的时候也是如此——基本上，使用我们并不真正了解的语言，这种语言描述性很差，词汇量极其有限。我发现最好先用简单的英语将我的方法概述为一系列步骤。这让我清楚地看到，用自己的话来说，我想让代码做什么。所以很容易得到一个概览并考虑正在发生的事情，但如果我意识到需要重写某些内容，也很容易进行修改。

实际上，我将步骤作为注释写在代码文件的不同行中。一旦我对方法感到满意，我就开始在这些注释之间编写代码本身。这留给我一个结构良好且适当注释的代码。

在原始代码中的特定示例中，可能并不需要整个f数组存在。如果您只是用简单的英语将过程写出来，我认为您会看到这一点，并能够轻松地在编写代码之前用简单的英语更改方法。目前，代码如此依赖于f，以至于需要完全重新编写才能更改它。

索引

纯Python通常处理索引时非常缓慢。例如：

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

这让我非常怀疑。当你说“将加载的文本转换为numpy ndarray”需要很长时间时，是指这个吗？这并不令人惊讶，因为每次Python看到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)

这是一个极其快速的操作，只需要不到一毫秒的时间。可能形状不正确，您可能需要更改参数的顺序或进行转置，但基本思想肯定存在。您可以在文档中了解更多信息。

更具体地说，如果我正确理解代码，那么前三个代码部分可以完全替换为：

a = np.random.rand(ngridx,ngridy,ngridz)

我的理解是，据我所见，f 实际上可以不必存在。

复制数据是不好的

你不需要复制全部内容，当你需要复制数组（或部分数组）时，应该让numpy为你完成。例如，不要使用你的 Firstdel 函数，而是使用 a[1:]：

def Firstdel(a):
    return a[1:]

这并不会复制数据；它返回的东西只是无视了它之前存储在RAM中的另一个数字。或者，如果你真的需要复制数据（仅仅为了绘图），可以使用以下方法：

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

通常情况下，您可以只使用numpy数组的“切片”，而不是复制它们。请在这里阅读有关此内容的信息。

但是，请注意，如果您只使用了一个切片（没有发生复制），那么这意味着如果您对该切片进行写入，则原始数组也将被写入。

循环也是时间浪费的罪魁祸首——尤其是当这些循环在python代码中而不是C代码中时。实际上，我认为numpy的整个重点是用C语言编写的循环替换您在python中编写的循环。

首先，while在Python中不常用于简单循环。因此，代替：

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

你应该使用

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

但这只是样式和可读性的问题。为了提高速度，您需要尽可能减少循环——特别是嵌套循环。对于OP参数，为了设置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] = np.arange(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函数的任何内容；它们只需完成工作即可。）

针对整个数组（或其切片）进行操作

一个常见的错误是使用numpy函数，但仍然自己编写循环。在这段代码中就有一个典型的例子：

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

不要使用大型嵌套循环计算每次循环中的K值，而是将K设置为包含相应值的数组：

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

K 将和 kx、ky 和 kz 的大小和形状相同。然后，您可以使用高级索引来设置 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):
    p_i = p[K==i]  # This will be just the elements of p where K==i
    q[i,0] = i
    q[i,1] = np.sum(np.abs(p_i)**2)
    q[i,2] = np.sum(p_i)
    q[i,3] = len(p_i)
print q

将这些技巧结合起来，我得到了大约70倍的优化，相比您问题中当前的代码。如果它仍然太慢，您可以使用性能分析器找出缓慢的部分。如果你还是无法改善它们，那么尝试使用numba可能会有所帮助，它会根据需要编译代码，运行速度几乎与相当的C代码一样快。

- Mike

---

非常感谢你的帮助，Mike。我正在重新编写和最小化代码，以便问题部分更清晰。基本上，减少循环并使用内置的numpy函数使代码运行速度更快。然而，仍然无法弄清楚最后一个索引部分。 - Tom

我已经更新了您在此答案中提出的所有建议，完善了代码财务部分。非常感谢您的帮助！ - Tom

在这个上下文中，“run off the edge”是什么意思？ - cxrodgers

不客气！我认为你可能错过了我在第一次发布后添加到问题中的其他优化，包括刚刚添加的一个，可以根据ngrid等因素缩短大约一分钟的时间。只要记住：您确实希望尽可能删除许多Python循环，并尽可能使用切片。掌握好切片技巧，这是我的建议。 - Mike

嗯，很奇怪。我想你可能做了一些不同于我所做的事情，但我找不出可能是什么。或者你正在使用旧版本的numpy。我已经清理了我版本的你的代码（我已经在ipython笔记本中），你可以在这里在线查看它（http://nbviewer.ipython.org/gist/moble/820999c8901968bbcd29）。 - Mike

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，您可以使用广播来消除循环：

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

- wwii

---

您不必使用np.floor_divide和np.fmod，您可以像平常一样使用//和%。由于a是一个numpy数组，它知道该怎么做。 - Mike

谢谢，这比我做的好多了。 - Tom

@Mike，使用ufunc有什么缺点吗？ - wwii

我认为它们完全相同 - 一个调用另一个。在我看来，这只是一种符号上的美观，使代码更易于阅读。 - Mike

哇！我觉得使用广播+newaxis非常聪明。非常感谢帮助我的每个人！ - Tom