为什么np.compress比布尔索引更快？

Question

np.compress内部的操作比布尔索引更快，原因是什么？

在这个例子中，compress大约快20%，但是时间节省取决于数组a和布尔数组b中的True值数量及其大小，但在我的机器上，compress总是更快。

import numpy as np

a = np.random.rand(1000000,4)
b = (a[:,0]>0.5)

%timeit a[b]
#>>> 10 loops, best of 3: 24.7 ms per loop
%timeit a.compress(b, axis=0)
#>>> 10 loops, best of 3: 20 ms per loop

布尔索引的文档中提到：

返回的是数据的副本，而不像切片一样返回视图。

相比之下，压缩文档则表示：

沿给定轴返回数组的选定切片。

但是，使用此方法来确定两个数组是否共享相同的数据缓冲区，结果显示这两种方法都不与其父级a共享数据，我认为也就意味着这两种方法都没有返回实际切片。

def get_data_base(arr):
    base = arr
    while isinstance(base.base, np.ndarray):
        base = base.base
    return base

def arrays_share_data(x, y):
    return get_data_base(x) is get_data_base(y) 

arrays_share_data(a, a.compress(b, axis=0))
#>>> False
arrays_share_data(a, a[b])
#>>> False

我很好奇，因为我在工作中经常执行这些操作。 我运行的是Python 3.5.2，Numpy版本为1.11.1，通过Anaconda安装。

- saintsfan342000

---

您还可以使用 np.where 将布尔值转换为索引数组，使用 np.take 选择具有该数组的行。所有操作都会生成所选值的副本。时间可能更多地涉及调用开销和普遍性。 - hpaulj

Answer 1

当我在 numpy 的 github 上追踪 a.compress 函数通过多层函数调用时，我找到了以下内容：

/numpy/core/src/multiarray/item_selection.c
PyArray_Compress(PyArrayObject *self, PyObject *condition, int axis,
             PyArrayObject *out)
    # various checks
    res = PyArray_Nonzero(cond);
    ret = PyArray_TakeFrom(self, PyTuple_GET_ITEM(res, 0), axis,
                       out, NPY_RAISE);

使用您的示例数组，compress与使用where获取索引数组并使用take相同：

In [135]: a.shape
Out[135]: (1000000, 4)
In [136]: b.shape
Out[136]: (1000000,)
In [137]: a.compress(b, axis=0).shape
Out[137]: (499780, 4)
In [138]: a.take(np.nonzero(b)[0], axis=0).shape
Out[138]: (499780, 4)
In [139]: timeit a.compress(b, axis=0).shape
100 loops, best of 3: 14.3 ms per loop
In [140]: timeit a.take(np.nonzero(b)[0], axis=0).shape
100 loops, best of 3: 14.3 ms per loop

事实上，如果我在[]索引中使用此索引数组，我将获得可比的时间：

In [141]: idx=np.where(b)[0]
In [142]: idx.shape
Out[142]: (499780,)
In [143]: timeit a[idx,:].shape
100 loops, best of 3: 14.6 ms per loop
In [144]: timeit np.take(a,idx, axis=0).shape
100 loops, best of 3: 9.9 ms per loop

np.take 代码更加复杂，因为它包括 clip 和 wrap 模式。

[] 索引被转换为 __getitem__ 调用，并通过各种层级。我没有追踪过那段代码，但我认为可以说 compress （或者更确切地说是 take）只是采取了更直接的方法来完成任务，因此速度有所提升。30-50% 的速度差异表明编译代码细节上存在差异，而不是像 views vs copies 或解释 vs 编译这样的重大问题。

- hpaulj

Answer 2

当沿一个轴选择的索引由布尔掩码向量指定时，函数compress是一种替代fancy indexing的方法。明显的速度增益是由于轴选择被预先指定，而fancy indexing可以用于对数组进行任意选择，因此会产生性能损失以使这成为可能。这也是您经历的可变速度增益的原因。

i = np.random.random_sample(n) < .5

b1 = a[i]
b2 = np.compress(i, a, axis=0)

%timeit a[i]
10 loops, best of 3: 59.8 ms per loop
%timeit np.compress(i, a, axis=0)
10 loops, best of 3: 24.1 ms per loop

- SerialDev

---

因此，为了实现这一点，会产生性能损失。这应该是一个恒定的惩罚，不随数组大小而变化。 - Eric

但是为什么numpy不能像从compress的轴参数中一样，从布尔索引语法中同样很好地推断出轴选择呢？ - saintsfan342000