Вы можете построить совершенно общее решение с причудливой индексацией, используя широковещательное сложение и упорядочение.
Возьмем одномерный случай:
arr = np.random.randint(10, size=973)
S = arr.shape[0]
N = 64
D = 128
# how many D-sized chunks?
nd = np.ceil(S / D)
# how many indices to chop from the end? I.e., which part of the last chunk doesn't fit in S?
nn = N - S + (nd - 1) * D
index = (np.arange(N) + D * np.arange(nd)[:, None]).ravel()[:-nn]
result = arr[index]
В 2D это будет выглядеть как
arr = np.random.randint(10, size=(1024, 768))
S = np.array(arr.shape)
N = 64
D = 128
nd = np.ceil(S / D)
nn = N - S + (nd - 1) * D
r_index = (np.arange(N) + D * np.arange(nd[0])[:, None]).ravel()[:-nn[0]]
c_index = (np.arange(N) + D * np.arange(nd[1])[:, None]).ravel()[:-nn[1]]
result = arr[np.ix_(r_index, c_index)]
Вы можете расширить это до N измерений с помощью небольшого трюка с широковещательной передачей и понимания небольшого списка:
arr = np.random.randint(10, size=(128, 200, 64))
S = np.array(arr.shape)
N = 64 # Could be array with different value for each dimension
D = 128 # Same with this
nd = np.ceil(S / D)
nn = N - S + (nd - 1) * D
Скорее всего, вы получите рваный массив индексов все, поэтому было бы разумно переключиться на список:
index = [(np.arange(N) + D * np.arange(ndx)[:, None]).ravel()[:-nnx] for ndx, nnx in zip(nd, nn)]
result = arr[np.ix_(*index)]