为啥`arr.take(idx)`比`arr[idx]`快

Posted 2023-02-23

【中文标题】为啥`arr.take(idx)`比`arr[idx]`快

【英文标题】：Why is `arr.take(idx)` faster than `arr[idx]`为什么`arr.take(idx)`比`arr[idx]`快

【发布时间】：2019-08-03 05:56:32

【问题描述】：

似乎普遍认为使用np.take 比使用数组索引要快得多。例如http://wesmckinney.com/blog/numpy-indexing-peculiarities/、Fast numpy fancy indexing 和Fast(er) numpy fancy indexing and reduction?。还有人建议np.ix_在某些情况下更好。

我已经进行了一些分析，在大多数情况下似乎确实如此，尽管随着数组变大，差异会减小。 性能受数组大小、索引长度（对于行）和列数的影响。行数似乎影响最大，即使索引为 1D，数组中的列数也有影响。更改索引的大小似乎对方法之间的影响不大。

所以，问题有两个方面： 1. 为什么方法之间的性能差异如此之大？ 2. 什么时候使用一种方法而不是另一种方法才有意义？是否有一些数组类型、排序或形状总是会更好用？

有很多事情可能会影响性能，因此我在下面展示了其中的一些，并包含了用于尝试使其可重现的代码。

编辑我更新了图表上的 y 轴以显示完整的值范围。它更清楚地表明差异比一维数据显示的要小。

一维索引

查看运行时间与行数的比较表明，索引非常一致，有轻微的上升趋势。随着行数的增加，take 始终变慢。

随着列数的增加，两者都变得更慢，但take 的增加更高（这仍然是一维索引）。

二维索引

与 2D 数据结果相似。还显示了使用ix_，它的整体性能似乎最差。

数字代码

from pylab import *
import timeit


def get_test(M, T, C):
    """
    Returns an array and random sorted index into rows
    M : number of rows
    T : rows to take
    C : number of columns
    """
    arr = randn(M, C)
    idx = sort(randint(0, M, T))
    return arr, idx


def draw_time(call, N=10, V='M', T=1000, M=5000, C=300, **kwargs):
    """
    call : function to do indexing, accepts (arr, idx)
    N : number of times to run timeit
    V : string indicating to evaluate number of rows (M) or rows taken (T), or columns created(C)
    ** kwargs : passed to plot
    """
    pts = 
        'M': [10, 20, 50, 100, 500, 1000, 2000, 5000, 10000, 20000, 50000, 100000, 200000, 500000, ],
        'T': [10, 50, 100, 500, 1000, 5000, 10000, 50000],
        'C': [5, 10, 20, 50, 100, 200, 500, 1000],
    
    res = []

    kw = dict(T=T, M=M, C=C) ## Default values
    for v in pts[V]:
        kw[V] = v
        try:
            arr, idx = get_test(**kw)
        except CallerError:
            res.append(None)
        else:
            res.append(timeit.timeit(lambda :call(arr, idx), number=N))

    plot(pts[V], res, marker='x', **kwargs)
    xscale('log')
    ylabel('runtime [s]')

    if V == 'M':
        xlabel('size of array [rows]')
    elif V == 'T':
        xlabel('number of rows taken')
    elif V == 'C':
        xlabel('number of columns created')

funcs1D = 
    'fancy':lambda arr, idx: arr[idx],
    'take':lambda arr, idx: arr.take(idx, axis=0),


cidx = r_[1, 3, 7, 15, 29]
funcs2D = 
    'fancy2D':lambda arr, idx: arr[idx.reshape(-1, 1), cidx],
    'take2D':lambda arr, idx: arr.take(idx.reshape(-1, 1)*arr.shape[1] + cidx),
    'ix_':lambda arr, idx: arr[ix_(idx, cidx)],


def test(funcs, N=100, **kwargs):
    for descr, f in funcs.items():
        draw_time(f, label="".format(descr), N=100, **kwargs)
    legend()

figure()
title('1D index, 30 columns in data')
test(funcs1D, V='M')
ylim(0, 0.25)
# savefig('perf_1D_arraysize', C=30)

figure()
title('1D index, 5000 rows in data')
test(funcs1D, V='C', M=5000)
ylim(0, 0.07)
# savefig('perf_1D_numbercolumns')

figure()
title('2D index, 300 columns in data')
test(funcs2D, V='M')
ylim(0, 0.01)
# savefig('perf_2D_arraysize')

figure()
title('2D index, 30 columns in data')
test(funcs2D, V='M')
ylim(0, 0.01)
# savefig('perf_2D_arraysize_C30', C=30)

【问题讨论】：

take 可能更快，但从我在 SO 问题（和numpy 函数）中看到的情况来看，它的使用频率不如“普通”索引。跨度>

【参考方案1】：

答案是非常低级别，并且与 C 编译器和 CPU 缓存优化有关。请在numpy issue 上查看与 Sebastian Berg 和 Max Bolingbroke（都是 numpy 的贡献者）的积极讨论。

花式索引试图“聪明”地了解如何读取和写入内存（C 顺序与 F 顺序），而 .take 将始终保持 C 顺序。这意味着花式索引对于 F 有序数组通常会更快，并且在任何情况下对于大型数组都应该更快。现在，numpy 决定什么是“智能”方式，而不考虑数组的大小，或者它运行的特定硬件。因此，对于较小的数组，由于更好地使用 CPU 缓存中的读取，选择“错误”的内存顺序实际上可能会获得更好的性能。

【讨论】：

感谢您的链接，我将不得不花一些时间阅读该链接。从那里的链接来看，take 的想法似乎更好地源于对 pandas 的优化，这让我感觉更好，因为它是一种不总是正确的民间智慧。