Question

我正在使用np.roll（）进行最接近邻居的平均，但我觉得有更快的方法。这是一个简化的example，但想象3维和更复杂的平均“模板”。例如，请参阅此paper的第6部分。

以下是该简化示例中的几行：

for j in range(nper):
    phi2 = 0.25*(np.roll(phi,  1, axis=0) +
                 np.roll(phi, -1, axis=0) +
                 np.roll(phi,  1, axis=1) +
                 np.roll(phi, -1, axis=1) )
    phi[do_me] = phi2[do_me]

那么我应该寻找返回视图而不是数组的东西（因为它似乎roll返回数组）？在这种情况下，每次调用时都会初始化一个新数组？我注意到小阵列的开销很大。

事实上，对于我的笔记本电脑上[100,100]到[300,300]的数组，效率最高。可能缓存上面的问题。

scipy.ndimage.interpolation.shift()执行方式here会更好吗？如果是，fixed?是Zimbabwean dollars circa 2009在上面的链接示例中，无论如何我都会扔掉被包裹的部分，但可能并不总是。

注意：在这个问题中，我只关注NumPy / SciPy中的可用内容。当然有很多好方法可以加速Python甚至NumPy，但这不是我在这里寻找的，因为我真的想要更好地理解NumPy。谢谢！

Answer 1

np.roll每次都必须创建数组的副本，这就是（相对）缓慢的原因。与scipy.ndimage.filters.convolve()之类的卷积会更快一些，但仍可能创建副本（取决于实现方式）。

在这种情况下，我们可以避免使用numpy views完全复制并在开头填充原始数组。

import numpy as np


def no_copy_roll(nx, ny):
    phi_padded = np.zeros((ny+2, nx+2))
    # these are views into different sub-arrays of phi_padded
    # if two sub-array overlap, they share memory
    phi_north = phi_padded[:-2, 1:-1]
    phi_east = phi_padded[1:-1, 2:]
    phi_south = phi_padded[2:, 1:-1]
    phi_west = phi_padded[1:-1, :-2]
    phi = phi_padded[1:-1, 1:-1]

    do_me = np.zeros_like(phi, dtype='bool')
    do_me[1:-1, 1:-1] = True

    x0, y0, r0 = 40, 65, 12
    x = np.arange(nx, dtype='float')[None, :]
    y = np.arange(ny, dtype='float')[:, None]
    rsq = (x-x0)**2 + (y-y0)**2
    circle = rsq <= r0**2
    phi[circle] = 1.0
    do_me[circle] = False

    n, nper = 100, 100
    phi_hold = np.zeros((n+1, ny, nx))
    phi_hold[0] = phi
    for i in range(n):
        for j in range(nper):
            phi2 = 0.25*(phi_south +
                         phi_north +
                         phi_east +
                         phi_west)

            phi[do_me] = phi2[do_me]

        phi_hold[i+1] = phi

    return phi_hold

对于像这样的简单基准测试，这将节省大约35％的时间。

from original import original_roll
from mwe import no_copy_roll
import numpy as np

nx, ny = (301, 301)
arr1 = original_roll(nx, ny)
arr2 = no_copy_roll(nx, ny)

assert np.allclose(arr1, arr2)

这是我的分析结果

37.685 <module>  timing.py:1
├─ 22.413 original_roll  original.py:4
│  ├─ 15.056 [self]
│  └─ 7.357 roll  <__array_function__ internals>:2
│     └─ 7.243 roll  numpy\core\numeric.py:1110
│           [10 frames hidden]  numpy
├─ 14.709 no_copy_roll  mwe.py:4
└─ 0.393 allclose  <__array_function__ internals>:2
   └─ 0.393 allclose  numpy\core\numeric.py:2091
         [2 frames hidden]  numpy
            0.391 isclose  <__array_function__ internals>:2
            └─ 0.387 isclose  numpy\core\numeric.py:2167
                  [4 frames hidden]  numpy

对于更精细的模具，此方法仍然有效，但是可能有点笨拙。在这种情况下，您可以查看skimage.util.view_as_windows，它使用此技巧的一种变体（numpy stride tricks）返回一个数组，该数组使您可以廉价地访问每个元素周围的窗口。但是，您将必须自己执行填充操作，并且需要注意不要创建结果数组的副本，这会很快变得昂贵。

Answer 2

目前我能得到的最快实现是基于您已经提到的 scipy.ndimage.interpolation.shift 的低级实现：

from scipy.ndimage.interpolation import _nd_image, _ni_support

cval = 0.0  # unused for mode `wrap`
mode = _ni_support._extend_mode_to_code('wrap')
_nd_image.zoom_shift(data, None, shift, data, 0, mode, cval)  # in-place update

预计算 mode、cval 和 shift 以直接调用低级 zoom_shift 方法让我了解了 x5 加速比。改为调用 shift，并且 x10 加速 wrt np.roll。

如何更快地np.roll（）？

2 个答案: