Python 列联表

作为我正在编写的项目的一部分，我正在生成很多很多列联表。

工作流程是：

• 获取具有连续（浮点）行的大型数据数组，并通过分箱将其转换为离散整数值（例如，结果行的值为 0-9）
• 将两行切片为向量 X 和 Y 并生成列联表 https://en.wikipedia.org/wiki/Contingency_table来自他们，这样我就有了二维频率分布
• 例如，我有一个 10 x 10 数组，计算出现的 (xi, yi) 数量
• 使用列联表进行一些信息论数学计算

最初，我将其写为：

def make_table(x, y, num_bins):
    ctable = np.zeros((num_bins, num_bins), dtype=np.dtype(int))
    for xn, yn in zip(x, y):
        ctable[xn, yn] += 1
    return ctable

这工作正常，但速度太慢，以至于占用了整个项目 90% 的运行时间。

我能想到的最快的纯Python优化是这样的：

def make_table(x, y, num_bins):
    ctable = np.zeros(num_bins ** 2, dtype=np.dtype(int))
    reindex = np.dot(np.stack((x, y)).transpose(), 
                     np.array([num_bins, 1]))
    idx, count = np.unique(reindex, return_counts=True)
    for i, c in zip(idx, count):
        ctable[i] = c
    return ctable.reshape((num_bins, num_bins))

这（不知何故）快了很多，但对于看起来不应该成为瓶颈的东西来说，它仍然相当昂贵。有没有我没有看到的有效方法可以做到这一点，或者我应该放弃并在 cython 中执行此操作？

另外，这里还有一个基准测试函数。

def timetable(func):
    size = 5000
    bins = 10
    repeat = 1000
    start = time.time()
    for i in range(repeat):
        x = np.random.randint(0, bins, size=size)
        y = np.random.randint(0, bins, size=size)
        func(x, y, bins)
    end = time.time()
    print("Func {na}: {ti} Ms".format(na=func.__name__, ti=(end - start)))

表示元素的巧妙技巧np.stack((x, y))因为整数可以变得更快：

In [92]: %timeit np.dot(np.stack((x, y)).transpose(), np.array([bins, 1]))
109 µs ± 6.55 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [94]: %timeit bins*x + y
12.1 µs ± 260 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

此外，第二个解决方案的最后部分可以稍微简化，只需考虑

np.unique(bins * x + y, return_counts=True)[1].reshape((bins, bins))

更重要的是，由于我们正在处理等距的非负整数，np.bincount https://docs.scipy.org/doc/numpy/reference/generated/numpy.bincount.html将会超越np.unique;这样，上面的内容可以归结为

np.bincount(bins * x + y).reshape((bins, bins))

总而言之，这比您当前正在做的事情提供了相当多的性能：

In [78]: %timeit make_table(x, y, bins)  # Your first solution
3.86 ms ± 159 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [79]: %timeit make_table2(x, y, bins)  # Your second solution
443 µs ± 23.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [101]: %timeit np.unique(bins * x + y, return_counts=True)[1].reshape((bins, bins))
307 µs ± 25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [118]: %timeit np.bincount(bins * x + y).reshape((10, 10))
30.3 µs ± 3.44 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

您可能还想知道np.histogramdd https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogramdd.html#numpy.histogramdd它同时处理舍入和合并，尽管它可能会比舍入和使用慢np.bincount.