使用 numpy 构建一个数组，其中的行从另一个二维数组中提取为 2x2 块

Question

假设我有以下二维数组：

x = np.array([[10,20,30,40], [50,60,70,80],[90,100,110,120]])  
print(x)

array([[ 10,  20,  30,  40],
       [ 50,  60,  70,  80],
       [ 90, 100, 110, 120]])

我想构建一个新数组，y，其中每一行都有来自 x 的 2x2 块的值，按顺时针顺序：

print(y)
array([[ 10,  20,  60,  50],
       [ 20,  30,  70,  60],
       [ 30,  40,  80,  70],
       [ 50,  60,  100, 90],
       [ 60,  70,  110, 100],
       [ 70,  80,  120, 110]])

我可以使用 Python for 循环实现，如下所示：

n_rows, n_cols = x.shape
y = []
for i in range(n_rows-1): 
     for j in range(n_cols-1): 
         row = [x[i,j],x[i,j+1],x[i+1, j+1],x[i+1,j]] 
         y.append(row) 
y = np.array(y)

我想知道是否有更快的方法利用 Numpy 函数并避免使用 Python 循环。

Answer 1

你可以缓存你的代码，因为循环主要是一次又一次地迭代同一个矩阵（如果你想在循环中保留相同的代码）。我已经为你的代码做了缓存前后的速度对比。

# Before caching
def loop_before_cache():
    n_rows, n_cols = x.shape
    y = []
    for i in range(n_rows-1): 
        for j in range(n_cols-1): 
            row = [x[i,j],x[i,j+1],x[i+1, j+1],x[i+1,j]] 
            y.append(row) 
    return np.array(y)


%timeit loop_before_cache()
11.6 µs ± 318 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

现在有了缓存

# After caching
from functools import lru_cache

@lru_cache()
def loop_after_cache():
    n_rows, n_cols = x.shape
    y = []
    for i in range(n_rows-1): 
        for j in range(n_cols-1): 
            row = [x[i,j],x[i,j+1],x[i+1, j+1],x[i+1,j]] 
            y.append(row) 
    return np.array(y)

%timeit loop_after_cache()
83.6 ns ± 2.42 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

额外

我使用 range 添加了带有 (1000,5000) 数组的模拟数据，以显示缓存的效率。

x = np.array([i for i in range(1,5000001)])
x = np.reshape(x, (1000,5000))

# Before caching
def loop_before_cache():
    n_rows, n_cols = x.shape
    y = []
    for i in range(n_rows-1): 
        for j in range(n_cols-1): 
            row = [x[i,j],x[i,j+1],x[i+1, j+1],x[i+1,j]] 
            y.append(row) 
    return np.array(y)

%timeit loop_before_cache()
8.58 s ± 113 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# After caching
@lru_cache(maxsize = 256)
def loop_after_cache():
    n_rows, n_cols = x.shape
    y = []
    for i in range(n_rows-1): 
        for j in range(n_cols-1): 
            row = [x[i,j],x[i,j+1],x[i+1, j+1],x[i+1,j]] 
            y.append(row) 
    return np.array(y)

%timeit loop_after_cache()
82.2 ns ± 5.58 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

Answer 2

首先，在 x 中创建一个 sliding_window_view，其中包含您要查看的 2x2 框：

b = np.lib.stride_tricks.sliding_window_view(x, (2, 2))

最里面的每个 2x2 数组都包含您想要的未拆解版本，但数组的第二部分颠倒了。到目前为止，我们没有复制任何数据。现在通过拼凑最后一个维度来制作副本。由于 b 高度 non-contiguous:

，因此 reshape 将始终在此处制作副本

c = b.reshape(*b.shape[:2], 4)

交换最后两列：

c[..., 2:] = c[..., -1:1:-1]

现在阐明主要维度：

y = c.reshape(-1, c.shape[-1])

如果您的 numpy 版本低于 1.20，您可以将 b 的定义替换为

b = np.lib.stride_tricks.as_strided(x, shape=(x.shape[0] - 1, x.shape[1] - 1, 2, 2), strides=x.strides * 2)

使用 numpy 构建一个数组，其中的行从另一个二维数组中提取为 2x2 块

Using numpy to construct an array with rows extracted from another 2D array as 2x2 blocks

python

numpy

sliding-window

额外