N次元の配列を２次元配列として分割したい。（配列arr を 軸 i, j により分割する split(arr, i, j) のような関数を制作したい）

Question

N次元の配列から特定の２次元配列を取り出したいです。
例えば2次元の以下の配列があるとして、

A = np.arange(24).reshape(4,6)

#[[ 0  1  2  3  4  5]
# [ 6  7  8  9 10 11]
# [12 13 14 15 16 17]
# [18 19 20 21 22 23]]

以下の関数を適用すると、

np.vsplit(A, 4)

# [array([[0, 1, 2, 3, 4, 5]]),
#  array([[ 6,  7,  8,  9, 10, 11]]),
#  array([[12, 13, 14, 15, 16, 17]]),
#  array([[18, 19, 20, 21, 22, 23]])]

のように簡単にsplitすることができます。

多次元の場合は

A = np.arange(24).reshape(2,4,3)
# array([[[ 0,  1,  2],
#         [ 3,  4,  5],
#         [ 6,  7,  8],
#         [ 9, 10, 11]],
# 
#        [[12, 13, 14],
#         [15, 16, 17],
#         [18, 19, 20],
#         [21, 22, 23]]])

C = [A[i,:,:] for i  in range(A.shape[0])]
# [array([[ 0,  1,  2],
#         [ 3,  4,  5],
#         [ 6,  7,  8],
#         [ 9, 10, 11]]), 
# array([[12, 13, 14],
#         [15, 16, 17],
#         [18, 19, 20],
#         [21, 22, 23]])]

のように分割を行うことができますが、 A[i,:,:]とA.shape(0)の二つの情報の指定が２度手間であると感じられます。うまく対処する関数などありましたら教えていただけると幸いです。

Answer 1

np.vsplitではどうしていけないのでしょうか？

---

分割してNからN-1次元の配列を得たいと勘違いしていました。

1、Nから2次元配列
2、2次元配列の行と列を元のN次元配列のi、j軸
という目標であれば、

def to_2d(A, i, j):
    tmp = A.swapaxes(i, -2)
    tmp = tmp.swapaxes(j, -1)
    tmp = tmp.reshape(-1, tmp.shape[-2], tmp.shape[-1])
    return np.vsplit(tmp, tmp.shape[0])

のような関数を作ればよいはずです。

Comments

[simasima]

np.vsplit では２次元配列を１次元配列の集合として分割することはできますが、３次元以上の配列をそれ以下の次元の配列の集合として分割することができないからです。
解釈が間違っていたらすみません。。。

[mkgrei]

公式ドキュメントでは高次元の例があります。
https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.vsplit.html

numpyのバージョンによって異なるのでしょうか。

[mkgrei]

vsplit/hsplit/dsplit/array_splitのセットのようですね。
おもしろい関数たちです。勉強になりました。

[simasima]

確かに公式ドキュメントで高次元の例がありますね。気づいておりませんでした。ありがとうございます。
ただ、３次元配列を２次元配列に分割する際、理論上３つの軸で分割することが可能ですが、この例だと関数がvsplit と hsplit のみで、２つの軸にしか分割できません。
もう少し柔軟な関数があれば嬉しいのですが。。。

[R.Shigemori]

dsplitが3軸目にあたるのではないでしょうか？

[mkgrei]

おっしゃる通りで、3軸目で分割するのならdsplitです。
axisキーワードに任意の軸iを渡せば、n次元のi軸で分割してくれる関数がarray_splitとなります。

[simasima]

そうなのですね！不勉強ですみません。
大変ありがとうございます！

[simasima]

すみません。この方法だとN次元の配列をN-1次元の配列に分割するのみで、求めている２次元の配列を得ることができません。。。
配列arr を 軸 i, j により分割する split(arr, i, j) のような関数を制作したいと考えております。

[mkgrei]

N次元から2次元を取り出したいのならば、ほしい軸以外を潰すようにreshapeしてからvsplitなどをするとよいと思います。
split(arr, i, j)のようにすると4次元から2次元の取り出しになります。
Nが3か4ならばいたずらに一般化する必要もありませんが。

[simasima]

ありがとうございます！
とてもスマートな方法で感動しました。
今７次元の情報を扱っていますので、この関数があるととても助かります。

Answer 2

今までの議論を元に解決法をまとめます。
例えば以下のような３次元配列がある場合。

import numpy as np

arr = np.arange(24).reshape(3,4,2)
arr

array([[[ 0, 1],
[ 2, 3],
[ 4, 5],
[ 6, 7]],

   [[ 8,  9],
    [10, 11],
    [12, 13],
    [14, 15]],

   [[16, 17],
    [18, 19],
    [20, 21],
    [22, 23]]])

以下のようにarray_splitを利用すると、特定の軸を基準にして二次元配列に分割することができました。

(ax0, ax1,ax2) = (0,1,2)
arr0 = np.array_split(arr, arr.shape[ax0], ax0)
arr1 = np.array_split(arr, arr.shape[ax1], ax1)
arr2 = np.array_split(arr, arr.shape[ax2], ax2)

Comments

[simasima]

すみません。この方法だとN次元の配列をN-1次元の配列に分割するのみで、求めている２次元の配列を得ることができませんでした。
配列arr を 軸 i, j により分割する split(arr, i, j) のような関数を制作したいと考えております。