pythonで配列を任意の軸に沿ってシャッフル

2022年4月20日追記
コメント欄でご指摘して頂いている通り，

rng = np.random.default_rng()
a = rng.random((3, 3))
print(a)
# [[0.79327762 0.07504873 0.25900636]
#  [0.75265331 0.49884496 0.55722056]
#  [0.89652968 0.97670661 0.22210966]]

print(rng.permutation(a, 0))
# [[0.75265331 0.49884496 0.55722056]
#  [0.79327762 0.07504873 0.25900636]
#  [0.89652968 0.97670661 0.22210966]]
print(rng.permutation(a, 1))
# [[0.25900636 0.79327762 0.07504873]
#  [0.55722056 0.75265331 0.49884496]
#  [0.22210966 0.89652968 0.97670661]]

でいけます．

以下は上記の下位互換かつ非推奨らしいので読む必要なし．

0次元目に沿ったシャッフルで良ければ，

import numpy as np
a = np.random.randn(3,3) #配列作成
print(a)
#array([[-1.53792152,  0.79386219,  0.48612455],
#       [-0.54037524, -0.00331165,  0.17748696],
#       [-1.77427981,  1.81990354,  0.9835304 ]])

np.random.shuffle(a)
print(a)
#array([[-0.54037524, -0.00331165,  0.17748696],
#       [-1.77427981,  1.81990354,  0.9835304 ],
#       [-1.53792152,  0.79386219,  0.48612455]])

でOK．

任意の軸でシャッフルしたい

他の次元でシャッフルしたい時は，

def shuffle_along_axis(a, axis):
    idx = np.random.rand(*a.shape).argsort(axis=axis)
    return np.take_along_axis(a,idx,axis=axis)

print(a)
#array([[-0.54037524, -0.00331165,  0.17748696],
#       [-1.77427981,  1.81990354,  0.9835304 ],
#       [-1.53792152,  0.79386219,  0.48612455]])

a_axis1 = shuffle_along_axis(a, 1)
print(a_axis1)
#[[-0.00331165  0.17748696 -0.54037524]
# [ 1.81990354  0.9835304  -1.77427981]
# [ 0.48612455  0.79386219 -1.53792152]]

これでOK.

参考