二维数组中基于一列的数据分块，并使用cython估计每个分格中的平均值

Question

为了优化代码的速度，这对MCMC的速度至关重要，我想用cython替换python代码中的一些瓶颈。由于我使用的是一个巨大的二维数组，我需要基于2D数组的一列来存储数据，然后根据第一列中的binning查找所有其他列中每个bin的平均值，所以我使用以下python代码：

   import numpy as np
   d = np.random.random((10**5, 3))
  #binning data again based on first column 
   bins = np.linspace(ndata[0,0], ndata[-1,0], 10)
   #compute the mean in each bin for different input parameters
   digitized = np.digitize(ndata[:,0], bins)
   r= np.array([ndata[digitized == i,0].mean() for i in range(1, len(bins))])
   p= np.array([ndata[digitized == i,1].mean() for i in range(1, len(bins))])
   q= np.array([ndata[digitized == i,2].mean() for i in range(1, len(bins))])

如何使用cython代码而不是numpy.digitize来加速至少两个数量级的速度？在

Answer 1

我不认为你需要cython来做这个，我想你在找numpy.bincount。下面是一个例子：

import numpy as np
d = np.random.random(10**5)
numbins = 10

bins = np.linspace(d.min(), d.max(), numbins+1)
# This line is not necessary, but without it the smallest bin only has 1 value.
bins = bins[1:]
digitized = bins.searchsorted(d)

bin_means = (np.bincount(digitized, weights=d, minlength=numbins) /
             np.bincount(digitized, minlength=numbins))

更新

让我们花点时间讨论一下为什么上面的代码比您问题中的代码快，以及为什么cython（可能）在这种情况下不会有太大帮助。在您的代码中，当您执行[digitized == i] for i in range(numbins)]操作时，您正在对digitized数组执行numbins传递。如果你熟悉big O notation，那就是O（n*m）。另一方面，bincount的做法有些不同。Bincount或多或少相当于：

^{pr2}$

它在digitized上有1次通过（如果你计算最大值的话是2次通过），所以它的复杂度为O（n）。此外，bincount已经用C编写并编译，因此它的开销非常小，而且非常快。Cython在代码有大量解释器和类型检查开销时非常有用，这样声明类型和编译代码可以消除这些开销。希望对你有帮助。在