numpyでfor文の計算を高速化したい

Question

1次元配列xを下のような3次元の配列yに変換したいです。
xは0～9215が格納されている配列です。
この配列xを8ずつ抜き出して, 256×6×6の配列yの深さ方向に格納していきたいです。
たとえば, x[0:8]をy[0:8,0,0]に格納, x[9:16]をy[0:8,0,1]格納して,y[0:8,6,6]まで
格納したら, y[8:16,0,0]から y[8:16,6,6]まで入れていき, 保存したいです。
for文でプログラムは作成しましたが, 繰り返しの計算に時間がかかってしまい。
numpyのreshapeなどで高速化することはできるかどうか教えていただきたく質問しました。

for文でのプログラムは下記になります。

import numpy as np

x = np.arange(1152*16)

print(x)

y = np.zeros((256,6,6))
for k in range(32):
    for i in range(6):
        for j in range(6):
            y[8*k:8*(k+1),i,j] = x[k*8*36+(i+j)*8:k*8*36+(i+j+1)*8]
            
print(y)

Comments

[bsdfan]

x[k*8*36+(i+j)*8:k*8*36+(i+j+1)*8] はこれであっています？
(i*6+j)*8ではない？
あっているなら、複数回参照されているところがあるので、reshapeではできないですね。

[coco_bauer]

(256,6,6)のshapeにreshapeするなら、numpy.reshapeが使えるかもしれませんが、8ずつ抜き出してという動作は無理と思いますよ。

[redcatML]

ありがとうございます。
(i*6+j)*8の誤りでした。
8ずつ抜き出すのは難しいのですね。
回答ありがとうございます

Answer 1

いったん4次元にreshapeして、transposeで次元を並べ替えてから、望みの形にreshapeすればできます。

y = x[:256*6*6].reshape((32, 6, 6, 8)).transpose((0, 3, 1, 2)).reshape((256, 6, 6))

元の配列がどう分割されて、どう並べ替えられているかを、うまくイメージしてみてください。

Comments

[redcatML]

回答ありがとうございます。
そのような方法でできるのですね。
toransposeの方法をよく調べてみます。
ありがとうございます。