Numpy Изменение формы / индексация группы

Question

Ситуация такова, что я хотел бы взять следующий Python / NumPy код:

# Procure some data:
z = np.zeros((32,32))

chunks = []
for i in range(0,32,step):
    for j in range(0,32,step):
        chunks.append( z[i:i+step,j:j+step] )
chunks = np.array(chunks)
chunks.shape # (256, 2, 2)

и векторизовать его / удалить циклы for. Это возможно? Я не особо возражаю против упорядочения окончательного массива, например 256,2,2 против 2,2,256, пока пространственная структура остается прежней. То есть блоки 2x2 из исходного массива.

Может, какой-нибудь маг c, использующий :: в дополнение к обычной индексации, сможет это сделать? Здесь есть NumPy мастера?

Answer 1

Вам может понадобиться transpose:

a = np.arange(1024).reshape(32,32)

a.reshape(16,2,16,2).transpose((0,2,1,3)).reshape(-1,2,2)

Вывод:

array([[[   0,    1],
        [  32,   33]],

       [[   2,    3],
        [  34,   35]],

       [[   4,    5],
        [  36,   37]],

       ...,

       [[ 986,  987],
        [1018, 1019]],

       [[ 988,  989],
        [1020, 1021]],

       [[ 990,  991],
        [1022, 1023]]])