使3d numpy数组的每N个切片连续

Question

声明

假设我们有一个三维的numpy数组，A，它的形状是(X, Y, Z)。我想创建一个新的数组B，它的形状也是(X, Y, Z)。你知道吗

我们希望沿第0轴的B的前n个片（:n）对应于沿第0轴的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的第0个切片，我们会发现：

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

如果我明白了，也许你可以用^{}：

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

---

代替tile，沿轴0使用np.repeat，并使用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)