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

Question

我正在分析我的一些代码,发现一个让我感到惊讶的结果np.where().我想where()在我的数组切片上使用(知道2D数组的很大一部分与我的搜索无关),并发现它是我的代码中的瓶颈.作为测试,我创建了一个新的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结果:

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

由于两种方法都返回一个新数组,我不清楚如何预先获取切片的完整副本能够加速np.where()到这个程度.我还做了几个健全检查确认:

1. 在这种情况下,它们都返回相同的结果.
2. 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)

我很好奇在这种情况下增加的开销来自何处？

Answer 1

以下是一些至少部分解释观察的源片段.我没有调查,where因为差异似乎是之前创建的.相反,我ufuncs一般都在看.

ufuncs的基本功能

暂时忽略一些特殊的套管ufuncs由外部环路内部可能优化的最内层1D环路计算,该环路覆盖其他维度.

外循环比较昂贵,它使用numpy nditer,所以必须设置它,并为每个迭代调用iternext,这是一个函数指针,所以没有内联.

相比之下,内环是一个简单的C循环.

跨步ufunc评估具有显着的开销

来自numpy/core/src/private/lowlevel_strided_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具有连续参数的a可以通过对快速内循环的单次调用来评估,完全绕过外循环.

要了解复杂性和开销的差异,请查看numpy/core/src/umath/ufunc_object.c 中的函数trivial_two/three_operand_loop和iterator_loopnumpy/core/src/multiarray/nditer_templ.c中的所有npyiter_iternext_*函数.

跨步的ufunc eval比跨步副本更昂贵

来自自动生成的numpy/core/src/multiarray/lowlevel_strided_loops.c

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

这个文件几乎是25万行.

相比之下,提供ufunc最内层循环的自动生成文件numpy/core/src/umath/loops.c是一个可疑的15'000行.

这本身就表明复制可能比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))

这是用来通过功能raw_array_assign_array在numpy的/核心/ SRC /多阵列/ array_assign_array.c来完成实际的复制为Python的ndarray.copy方法.

我们可以看到,与ufuncs使用的"完全迭代"相比,"原始迭代"的开销相当适中.