python - 为什么np.where()在数组切片的副本上比在原始数组上的 View 更快？

我分析了我的一些代码，发现了一个让我吃惊的结果np.where()。我想在数组的一个片段上使用where()（知道2d数组的很大一部分与我的搜索无关），发现它是我代码中的一个瓶颈。作为测试，我创建了一个新的二维数组作为该切片的副本，并测试了where()的速度。结果它跑得快得多。在我的实际案例中，速度的提高是非常显著的，但是我认为这个测试代码仍然展示了我的发现：

import numpy as np

def where_on_view(arr):
    new_arr = np.where(arr[:, 25:75] == 5, arr[:, 25:75], np.NaN)

def where_on_copy(arr):
    copied_arr = arr[:, 25:75].copy()
    new_arr = np.where(copied_arr == 5, copied_arr, np.NaN)

arr = np.random.choice(np.arange(10), 1000000).reshape(1000, 1000)

结果是：

%timeit where_on_view(arr)
398 µs ± 2.82 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit where_on_copy(arr)
295 µs ± 6.07 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

由于这两种方法都返回一个新数组，我不清楚预先获取切片的完整副本如何能够将timeit加速到这种程度。我也做了几次精神检查，确认：
在这种情况下，它们都返回相同的结果。
np.where()搜索实际上仅限于切片，不检查整个数组，然后过滤输出。
在这里：

# Sanity check that they do give the same output

test_arr = np.random.choice(np.arange(3), 25).reshape(5, 5)
test_arr_copy = test_arr[:, 1:3].copy()

print("No copy")
print(np.where(test_arr[:, 1:3] == 2, test_arr[:, 1:3], np.NaN))
print("With copy")
print(np.where(test_arr_copy == 2, test_arr_copy, np.NaN))

# Sanity check that it doesn't search the whole array

def where_on_full_array(arr):
    new_arr = np.where(arr == 5, arr, np.NaN)

#%timeit where_on_full_array(arr)
#7.54 ms ± 47.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

我很好奇在这种情况下，额外的开销是从哪里来的？

最佳答案

以下是一些源代码片段，至少部分地解释了观察结果。我不研究where，因为差异似乎是以前造成的。相反，我一般都在看ufuncs。
ufuncs的基本功能
暂时忽略一些特殊的外壳UFRESs由一个潜在优化的最内层1D环在一个覆盖其他维度的外环中计算。
外部循环比较昂贵，它使用一个numpynditer，因此必须设置它，对于每个迭代调用iternext，它是一个函数指针，所以没有内联。
通过比较，内环是一个简单的C环。
跨跃性UFUNC评价具有显著的开销
来自numpy/core/src/private/lowlevel_-spired_loops.h，它包含在numpy/core/src/umath/ufunc_object.c中

/*
 *            TRIVIAL ITERATION
 *
 * In some cases when the iteration order isn't important, iteration over
 * arrays is trivial.  This is the case when:
 *   * The array has 0 or 1 dimensions.
 *   * The array is C or Fortran contiguous.
 * Use of an iterator can be skipped when this occurs.  These macros assist
 * in detecting and taking advantage of the situation.  Note that it may
 * be worthwhile to further check if the stride is a contiguous stride
 * and take advantage of that.

因此，我们可以看到，一个具有连续参数的ufunc可以通过对快速内环的单个调用来评估，完全绕过外部循环。
为了理解复杂性和开销的差异，在NUMPY/CARC/SRC/UMySuthObj.c中查看函数trivial_two/three_operand_loop vsiterator_loop，以及在NUMPY/CARC/SRC/Multuald/ NdTythTunl C中的所有npyiter_iternext_*函数。
步履蹒跚的UFUNC EVE比跨步复制更昂贵
从自动生成的numpy/core/src/multiarray/lowlevel掼u-spired掼loops.c

/*
 * This file contains low-level loops for copying and byte-swapping
 * strided data.
 *

这个文件差不多有25万行。
通过比较，还提供了自动生成的文件NUMPY/CARC/SRC/UMys/Loop.C，它提供了UFUNC最内层的循环，是一条15万条线。
这本身表明，复制可能比UFUNC评价更为优化。
这里的相关位是宏。

/* Start raw iteration */
#define NPY_RAW_ITER_START(idim, ndim, coord, shape) \
        memset((coord), 0, (ndim) * sizeof(coord[0])); \
        do {

[...]

/* Increment to the next n-dimensional coordinate for two raw arrays */
#define NPY_RAW_ITER_TWO_NEXT(idim, ndim, coord, shape, \
                              dataA, stridesA, dataB, stridesB) \
            for ((idim) = 1; (idim) < (ndim); ++(idim)) { \
                if (++(coord)[idim] == (shape)[idim]) { \
                    (coord)[idim] = 0; \
                    (dataA) -= ((shape)[idim] - 1) * (stridesA)[idim]; \
                    (dataB) -= ((shape)[idim] - 1) * (stridesB)[idim]; \
                } \
                else { \
                    (dataA) += (stridesA)[idim]; \
                    (dataB) += (stridesB)[idim]; \
                    break; \
                } \
            } \
        } while ((idim) < (ndim))

由numpy/core/src/multiarray/array_assign_array.c中的函数raw_array_assign_array使用，该函数为pythonndarray.copy方法执行实际复制。
我们可以看到，与ufunc使用的“完全迭代”相比，“原始迭代”的开销相当小。