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

Question

我想根据存储在 HDF5 文件中的大型数据集(超过 100000 个样本)在 Python 中创建二维直方图。我想出了以下代码:

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

执行大约需要 15s(100000 个数据点)。然而，CERN 的 Root(使用其自己的树数据结构而不是 HDF5)可以在不到 1 秒的时间内完成此操作。你知道我如何加速代码吗？如果有帮助，我还可以更改 HDF5 数据的结构。

Answer 1

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

1. 从 hdf 文件加载数据，而不是传入有效的内存映射数组。
2. 如果这不能解决问题，您可以利用 scipy.sparse.coo_matrix 制作二维直方图。对于旧版本的 numpy，digitize(所有各种 histogram* 函数都在内部使用)在某些情况下可能会使用过多的内存。不过，最近 (>1.5??) 版本的 numpy 已不再是这种情况。

作为第一个建议的示例，您可以执行以下操作:

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，它只是稍微提高了内存效率。我确实有一段代码是我故意这样做的，但我一般不会推荐它。这是一个非常骇人听闻的解决方案。不过，在非常特定的情况下(很多点和很多 bin)，它可以比 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