从大型数据集中高效创建二维直方图

时间:2012-01-10 15:18:57

标签: python numpy matplotlib

我想在存储在HDF5文件中的大型数据集(100000+样本)中创建python中的2d直方图。我想出了以下代码:

import sys
import h5py
import numpy as np
import matplotlib as mpl
import matplotlib.pylab

f = h5py.File(sys.argv[1], 'r')

A = f['A']
T = f['T']

at_hist, xedges, yedges = np.histogram2d(T, A, bins=500)
extent = [yedges[0], yedges[-1], xedges[0], xedges[-1]]

fig = mpl.pylab.figure()
at_plot = fig.add_subplot(111)

at_plot.imshow(at_hist, extent=extent, origin='lower', aspect='auto')

mpl.pylab.show()

f.close()

执行大约需要15秒(100000个数据点)。然而,CERN的Root(使用自己的树数据结构而不是HDF5)可以在不到1秒内完成此操作。你知道如何加速代码吗?如果它有用,我也可以改变HDF5数据的结构。

3 个答案:

答案 0 :(得分:14)

我会尝试一些不同的东西。

  1. 从hdf文件加载数据,而不是传入有效的内存映射数组。
  2. 如果这不能解决问题,您可以利用scipy.sparse.coo_matrix来制作2D直方图。对于numpy的旧版本,digitize(所有各种histogram*函数在内部使用)在某些情况下可能会使用过多的内存。但是,最近(> 1.5 ??)版本的numpy不再是这种情况了。
  3. 作为第一个建议的一个例子,你可以这样做:

    f = h5py.File(sys.argv[1], 'r')
    A = np.empty(f['A'].shape, f['A'].dtype)
    T = np.empty(f['T'].shape, f['T'].dtype)
    f['A'].read_direct(A)
    f['T'].read_direct(T)
    

    这里的不同之处在于整个数组将被读入到内存,而不是h5py的类似数组的对象,这些对象基本上是高效的内存映射数组在磁盘上

    至于第二个建议,除非第一个建议没有帮助你解决问题,否则不要尝试。

    它可能不会明显更快(并且对于小型数组可能更慢),并且对于numpy的最新版本,它只是稍微更高的内存效率。我确实有一段代码,我故意这样做,但我不推荐它。这是一个非常hackish的解决方案。在非常精选的情况下(许多点和许多箱子),它可以比histogram2d更好地进行预制。

    除了所有这些警告之外,这里是:

    import numpy as np
    import scipy.sparse
    import timeit
    
    def generate_data(num):
        x = np.random.random(num)
        y = np.random.random(num)
        return x, y
    
    def crazy_histogram2d(x, y, bins=10):
        try:
            nx, ny = bins
        except TypeError:
            nx = ny = bins
        xmin, xmax = x.min(), x.max()
        ymin, ymax = y.min(), y.max()
        dx = (xmax - xmin) / (nx - 1.0)
        dy = (ymax - ymin) / (ny - 1.0)
    
        weights = np.ones(x.size)
    
        # Basically, this is just doing what np.digitize does with one less copy
        xyi = np.vstack((x,y)).T
        xyi -= [xmin, ymin]
        xyi /= [dx, dy]
        xyi = np.floor(xyi, xyi).T
    
        # Now, we'll exploit a sparse coo_matrix to build the 2D histogram...
        grid = scipy.sparse.coo_matrix((weights, xyi), shape=(nx, ny)).toarray()
    
        return grid, np.linspace(xmin, xmax, nx), np.linspace(ymin, ymax, ny)
    
    if __name__ == '__main__':
        num=1e6
        numruns = 1
        x, y = generate_data(num)
        t1 = timeit.timeit('crazy_histogram2d(x, y, bins=500)',
                setup='from __main__ import crazy_histogram2d, x, y',
                number=numruns)
        t2 = timeit.timeit('np.histogram2d(x, y, bins=500)',
                setup='from __main__ import np, x, y',
                number=numruns)
        print 'Average of %i runs, using %.1e points' % (numruns, num)
        print 'Crazy histogram', t1 / numruns, 'sec'
        print 'numpy.histogram2d', t2 / numruns, 'sec'
    

    在我的系统上,这会产生:

    Average of 10 runs, using 1.0e+06 points
    Crazy histogram 0.104092288017 sec
    numpy.histogram2d 0.686891794205 sec
    

答案 1 :(得分:3)

您需要确定瓶颈是在数据加载中还是在histogram2d中。尝试在代码中插入一些时间测量。

是A和T数组还是生成器对象?如果是后者,则需要更加谨慎地区分瓶颈所在;您可能必须先将它们解压缩到numpy数组才能进行测试。

答案 2 :(得分:2)

整个事情是在15s运行还是只调用histogram2d?导入pylab系列可能需要很长时间。如果我没记错的话,numpy histogram2d函数应该在C中实现,所以我怀疑那里存在性能问题。您可以通过使用优化标志--OO

调用脚本来加速python
python -OO script.py 

另请考虑使用Psycho来提高效果。