使 3-d numpy 数组的每个第 n 个切片连续

Question

声明

假设我们有一些 3 维 numpy 数组 A，其形状为 (X, Y, Z)。我想创建一个新数组B，它的形状也将是(X, Y, Z)。

我们希望B 沿第零轴的前 n 个切片 (:n) 对应于 A 沿第零轴的每个第 m 切片 (::m)。

我们还希望B 的切片n:2*n 对应于A 的每个m+1 切片（1::m）。数组的其余部分依此类推。

使用矢量化 numpy 计算实现此目的的最佳方法是什么？

示例

上面的陈述最好通过一个例子来理解。所以让我们开始设置一些示例数组A：

import numpy as np

# Create array A with shape (15, 3, 3)
n = 3; m = 5
a = np.array([i * np.eye(3) for i in range(1, 1+m)])
A = np.tile(a, (n, 1, 1))

如果我们看一下A 的一些零切片，我们会发现：

print(A[0])
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
print(A[1])
[[2. 0. 0.]
 [0. 2. 0.]
 [0. 0. 2.]]

...

print(A[4])
[[5. 0. 0.]
 [0. 5. 0.]
 [0. 0. 5.]]
print(A[5])
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]

等等。

A 中的值并不重要，但应该有助于说明原始语句。

我想知道我们是否可以仅使用 numpy 函数创建矩阵 B。数组B 应该有切片：

print(B[0])
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
print(B[1])
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
print(B[2])
[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
print(B[3])
[[2. 0. 0.]
 [0. 2. 0.]
 [0. 0. 2.]]

等等。

有没有办法用纯 numpy 解决方案从 A 生成 B？

我尝试过的

以下内容根据需要提供B，但随着m 变大，它变得乏味：

# vstack solution
B = np.vstack((A[::m], A[1::m], A[2::m], A[3::m], A[4::m]))

使用列表推导也可以，但我想避免使用循环：

# List comprehension solution
B = np.vstack([A[i::m] for i in range(m)])

Answer 1

我认为这是你想要的：

import numpy as np

# Create array A with shape (15, 3, 3)
a = np.array([i * np.eye(3) for i in range(1, 6)])
A = np.tile(a, (3, 1, 1))

B = np.swapaxes(A.reshape(3, 5, 3, 3), 0, 1)
B = B.reshape(-1, 3, 3)
print(B)
# [[[1. 0. 0.]
#   [0. 1. 0.]
#   [0. 0. 1.]]
#
#  [[1. 0. 0.]
#   [0. 1. 0.]
#   [0. 0. 1.]]
#
#  [[1. 0. 0.]
#   [0. 1. 0.]
#   [0. 0. 1.]]
#
#  [[2. 0. 0.]
#   [0. 2. 0.]
#   [0. 0. 2.]]
# ...

Answer 2

如果我明白，也许你可以使用numpy.sort：

B = np.sort(A, axis = 0)

Answer 3

设置

n = 3; m = 5
a = np.array([i * np.eye(n) for i in range(1, 1+m)])

---

沿轴 0 使用 np.repeat 代替 tile，并使用 Fortran 样式排序进行整形。

np.repeat(a, n, 0).reshape(m*n, n, n, order='F')

array([[[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]],

       [[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]],

       [[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]],

       [[2., 0., 0.],
        [0., 2., 0.],
        [0., 0., 2.]],    
       ...    
       [[5., 0., 0.],
        [0., 5., 0.],
        [0., 0., 5.]]])

验证

# your approach
A = np.tile(a, (n, 1, 1))
B = np.vstack((A[::m], A[1::m], A[2::m], A[3::m], A[4::m]))

# my approach
usr_B = np.repeat(a, n, 0).reshape(m*n, n, n, order='F')

>>> np.array_equal(B, usr_B)
True

时间

%%timeit
A = np.tile(a, (n, 1, 1))
B = np.vstack((A[::m], A[1::m], A[2::m], A[3::m], A[4::m]))
19 µs ± 57.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%%timeit
A = np.tile(a, (3, 1, 1))
B = np.swapaxes(A.reshape(3, 5, 3, 3), 0, 1)
B = B.reshape(-1, 3, 3)
11 µs ± 74.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit np.repeat(a, n, 0).reshape(m*n, n, n, order='F')
2.68 µs ± 21.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)