python numpy多次元のsortについて

Question

３次元配列の３次元目を，ある２次元配列の持つ値の大きさでソートしたいです．

その時

a_3d = np.random.rand(100,10,20)
a_2d = np.random.rand(20,100)
sort_col_index = np.argsort(a_2d, axis=0)
a_sorted = a_3d[0,:,sort_col_index[:,0]]

のように行うと

a_sorted.shape
# (20, 10)

となり，a_3dの形が崩れてしまいます．
なぜなのでしょうか．

また，実際にはa_3dの1次元目（要素数は100）でforループにかけたく思います．
一体どのようなコードを書けばよろしいのでしょうか．

どうぞよろしくお願いいたします．

Answer 1

a_sorted.shape
# (20, 10)

となり，a_3dの形が崩れてしまいます．
なぜなのでしょうか．

最後の添字表記で0番目の要素を取り出しているからです。

# 元のコード
# a_sorted = a_3d[0,:,sort_col_index[:,0]]

# 改善
# a_sorted = a_3d[:,:,sort_col_index[:,0]] # どうして:の代わりに0って書いたんだろう？
# => a_sorted.shape : (100, 10, 20)

また，実際にはa_3dの1次元目（要素数は100）でforループにかけたく思います．
一体どのようなコードを書けばよろしいのでしょうか．

こういうことでしょうか？

for x in a_3d:
   ... 

Comments

[Renkon]

実際には，100というのが試行回数となっていて
a_3dの各試行ごとに列ベクトルをソートしたものを取得したく思います．

ご回答いただいた
＞# 改善
＞# a_sorted = a_3d[:,:,sort_col_index[:,0]] # どうして:の代わりに0って書いたんだろう？
＞# => a_sorted.shape : (100, 10, 20)
こちらのコードでa_3dを各試行（1次元目）ごとに列ベクトルをソートした配列が得られるのでしょうか？
理解が浅く申し訳ありません．

[hayataka2049]

ちょっと理解が追いついていないので整理させてください。
＞a_3dの各試行ごとに列ベクトルをソートしたもの
ということですが、その際に1回目の試行a_3d[0]に対して、sort_col_indexを使って書くとすると、
a_2dのshapeは(20, 100)ですが100が試行回数に相当し20がa_3dのaxis=2と相当するので
a_sorted = a_3d[0,:,sort_col_index[:,0]]
で0番目の試行の結果が得られるということですね。それなら確かに納得行きます。
でも実際には、配列全体のソートを一括で行いたいのですね。その場合は、私の回答のコードは間違いです（単にshapeを揃えただけという）。

[Renkon]

返信が遅くなり申し訳ありません．

＞ということですが、その際に1回目の試行a_3d[0]に対して、sort_col_indexを使って書くとすると、
＞a_2dのshapeは(20, 100)ですが100が試行回数に相当し20がa_3dのaxis=2と相当するので
＞a_sorted = a_3d[0,:,sort_col_index[:,0]]
＞で0番目の試行の結果が得られるということですね。それなら確かに納得行きます。
おっしゃる通りです．
(10, 20)の配列を，axis=1つまり20個の列ベクトルをsort_col_indexでソートしたいです．

これを全試行で行いたいので，質問に乗せたコードで良いかと思ったのですがいかがでしょうか．
要領を得ない質問をしてしまい恐縮です．

[hayataka2049]

質門のコードはあくまでも0つめの試行の結果を得ているだけなので、全試行で行いたいのであれば間違っています。
たとえば
np.array([a_3d[i,:,sort_col_index[:,i]] for i in range(100)])
のようにして全試行のソートされた結果を得られるはずです。

[Renkon]

forで回す前に確認と思って0を代入するとa_sortedは（20, 10）の配列となっており，
期待していた（10, 20）の配列ではありませんでした．

＞np.array([a_3d[i,:,sort_col_index[:,i]] for i in range(100)])
こちらのコードでも同様に(100, 20, 10)の配列になるのですが
なぜこうなってしまうのでしょうか？

[Renkon]

np.array([a_3d[i][:,sort_col_index[:,i]] for i in range(100)])
上記のようにindexの指定を切り分けることで(100, 10, 20)の配列を得ることができました．
ありがとうございました．