awkward 配列の形状操作#
本章では、Awkward Array における形状（階層構造）の操作方法を、flatten や ravel、規則軸（Regular）と非規則軸（Irregular）の変換、unflatten による層追加、配列結合、構造体配列の操作などを例示しながら解説します。
import numpy as np
import awkward as ak
flatten と ravel#
まず、多重にネストされた Awkward Array を用いて、flatten() がどの軸をどのように潰すかを確認します。
a = ak.Array([[[1, 2], None, [3]], [[None, 4], [5]], None, [[6, None]]])
a
[[[1, 2], None, [3]], [[None, 4], [5]], None, [[6, None]]] --------------------- backend: cpu nbytes: 272 B type: 4 * option[var * option[var * ?int64]]
Awkward Array の axis=0 での flatten は、最上位の None を取り除くだけです。
ak.flatten(a, axis=0)
[[[1, 2], None, [3]], [[None, 4], [5]], [[6, None]]] --------------------- backend: cpu nbytes: 256 B type: 3 * var * option[var * ?int64]
次に、axis=1 や axis=2 を指定すると、対応する層が削除され、要素が一段浅い層へ移動します。
ak.flatten(a, axis=1)
[[1, 2], None, [3], [None, 4], [5], [6, None]] ----------- backend: cpu nbytes: 208 B type: 6 * option[var * ?int64]
ak.flatten(a, axis=2)
[[1, 2, 3], [None, 4, 5], None, [6, None]] -------------- backend: cpu nbytes: 176 B type: 4 * option[var * ?int64]
複数階層を削除したい場合は、flatten を複数回適用します。
ak.flatten(ak.flatten(a, axis=2), axis=1)
[1, 2, 3, None, 4, 5, 6, None] ------ backend: cpu nbytes: 112 B type: 8 * ?int64
axis=None を指定すると、最終軸のすべての値が一次元に取り出され、None は除去されます。
ak.flatten(a, axis=None)
[1, 2, 3, 4, 5, 6] --- backend: cpu nbytes: 48 B type: 6 * int64
一方、ravel() は flatten(None) と似ていますが、None を保持します。
ak.ravel(a)
[1, 2, 3, None, 4, 5, 6, None] ------ backend: cpu nbytes: 112 B type: 8 * ?int64
規則軸と非規則軸#
次に、Numpy 配列から生成した Awkward Array の規則軸（Regular）と、リストから生成した非規則軸（Irregular）の違いを確認します。
np.random.seed(42)
arr = np.random.randint(0, 10, (4, 6))
a = ak.Array(arr)
a
[[6, 3, 7, 4, 6, 9], [2, 6, 7, 4, 3, 7], [7, 2, 5, 4, 1, 7], [5, 1, 4, 0, 9, 5]] -------------------- backend: cpu nbytes: 96 B type: 4 * 6 * int32
b = ak.Array(arr.tolist())
b
[[6, 3, 7, 4, 6, 9], [2, 6, 7, 4, 3, 7], [7, 2, 5, 4, 1, 7], [5, 1, 4, 0, 9, 5]] -------------------- backend: cpu nbytes: 232 B type: 4 * var * int64
固定長の軸を可変長に変換します。
ak.from_regular(a, axis=1)
[[6, 3, 7, 4, 6, 9], [2, 6, 7, 4, 3, 7], [7, 2, 5, 4, 1, 7], [5, 1, 4, 0, 9, 5]] -------------------- backend: cpu nbytes: 136 B type: 4 * var * int32
可変長軸を固定長軸へ変換します。
ak.to_regular(b, axis=1)
規則軸では Numpy と同様のブールインデックスが可能です。
arr[arr >= 4]
array([6, 7, 4, 6, 9, 6, 7, 4, 7, 7, 5, 4, 7, 5, 4, 9, 5], dtype=int32)
a[a >= 4]
[6, 7, 4, 6, 9, 6, 7, 4, 7, 7, 5, 4, 7, 5, 4, 9, 5] --- backend: cpu nbytes: 68 B type: 17 * int32
非規則軸では、ネスト構造を維持したままフィルタされます。
c = b[b >= 4]
c
[[6, 7, 4, 6, 9], [6, 7, 4, 7], [7, 5, 4, 7], [5, 4, 9, 5]] ----------------- backend: cpu nbytes: 176 B type: 4 * var * int64
非規則軸のため to_regular() するとエラーになります。
try:
ak.to_regular(c)
except ValueError as ex:
print(ex)
cannot convert to RegularArray because subarray lengths are not regular (in compiled code: https://github.com/scikit-hep/awkward/blob/awkward-cpp-50/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_toRegularArray.cpp#L22)
長さを揃えてから規則化#
次のコードは、非規則な長さを持つ各部分配列の長さを揃えるための処理です。まず ak.num(c) で各部分配列の長さを取得し、その中の最大値を ak.max(ak.num(c)) で求めます。その最大長を target_length として ak.pad_none(c, target_length) を適用することで、長さが足りない部分配列には末尾に None を追加してすべての部分配列の長さを統一し、後続で ak.to_regular() による規則軸への変換が可能な状態にします。
d = ak.pad_none(c, ak.max(ak.num(c)))
d
[[6, 7, 4, 6, 9], [6, 7, 4, 7, None], [7, 5, 4, 7, None], [5, 4, 9, 5, None]] -------------------- backend: cpu nbytes: 336 B type: 4 * var * ?int64
ak.to_regular(d)
[[6, 7, 4, 6, 9], [6, 7, 4, 7, None], [7, 5, 4, 7, None], [5, 4, 9, 5, None]] -------------------- backend: cpu nbytes: 296 B type: 4 * 5 * ?int64
左側に None を入れたい場合：
ak.pad_none(c[:, ::-1], ak.max(ak.num(c)))[:, ::-1]
[[6, 7, 4, 6, 9], [None, 6, 7, 4, 7], [None, 7, 5, 4, 7], [None, 5, 4, 9, 5]] -------------------- backend: cpu nbytes: 336 B type: 4 * var * ?int64
clip引数をTrueにすると、各部分配列の長さが指定した長さを超えている場合に、余分な要素を末尾から切り捨てて指定した長さに揃えます。次のコードでは、各部分配列の長さを c の最小長に合わせ、長すぎる部分配列は余分な要素を切り捨てることで、すべての部分配列の長さを同じにします。クリップされた軸は規則軸になります。
ak.pad_none(c, int(ak.min(ak.num(c))), clip=True)
[[6, 7, 4, 6], [6, 7, 4, 7], [7, 5, 4, 7], [5, 4, 9, 5]] -------------- backend: cpu nbytes: 264 B type: 4 * 4 * ?int64
新しい規則軸を追加#
NumPyと同じように、None を使って新しい軸を追加できます。この場合、追加された軸の長さは 1 で、規則軸になります。
b[..., None]
[[[6], [3], [7], [4], [6], [9]], [[2], [6], [7], [4], [3], [7]], [[7], [2], [5], [4], [1], [7]], [[5], [1], [4], [0], [9], [5]]] -------------------------------- backend: cpu nbytes: 232 B type: 4 * var * 1 * int64
b[:, None, :]
[[[6, 3, 7, 4, 6, 9]], [[2, 6, 7, 4, 3, 7]], [[7, 2, 5, 4, 1, 7]], [[5, 1, 4, 0, 9, 5]]] ---------------------- backend: cpu nbytes: 232 B type: 4 * 1 * var * int64
singletons() を使うと、指定軸の後ろに長さ 1 の非規則軸が追加されます。
ak.singletons(a, axis=1)
[[[6], [3], [7], [4], [6], [9]], [[2], [6], [7], [4], [3], [7]], [[7], [2], [5], [4], [1], [7]], [[5], [1], [4], [0], [9], [5]]] -------------------------------- backend: cpu nbytes: 296 B type: 4 * 6 * var * int32
ak.singletons(a, axis=0)
[[[6, 3, 7, 4, 6, 9]], [[2, 6, 7, 4, 3, 7]], [[7, 2, 5, 4, 1, 7]], [[5, 1, 4, 0, 9, 5]]] ---------------------- backend: cpu nbytes: 136 B type: 4 * var * 6 * int32
unflatten で層追加#
ak.unflatten() は、一次元配列を指定した長さごとに分割して階層化されたリスト（部分配列）の形に変換する関数です。例えば次のコードは、最初の 3 要素 [1, 2, 3] が第一の部分配列、次の 2 要素 [4, 5] が第二の部分配列、長さ 0 の空配列が第三、最後の 3 要素 [6, 7, 8] が第四の部分配列として、それぞれネストされた配列として返されます。このように unflatten を使うことで、一次元配列を任意の長さごとに階層化して新しい階層構造を作ることができ、後続の処理で部分配列単位での操作が容易になります。
a = ak.Array([1, 2, 3, 4, 5, 6, 7, 8])
b = ak.unflatten(a, [3, 2, 0, 3])
b
[[1, 2, 3], [4, 5], [], [6, 7, 8]] ----------- backend: cpu nbytes: 104 B type: 4 * var * int64
既存の階層に対しても、axis を指定することで新しい層を挿入できます。ak.unflatten() の第二引数は、分割したい各部分配列の長さを順番に指定する一次元配列です。この場合、a は既に二次元の配列（各行が部分配列）ですが、axis=1 を指定することで、各行の中の要素をさらに指定した長さごとに分割して新しい階層を作ることができます。次のコードでは、第二引数 [4, 2, 1, 2, 2, 1, 3] が「各部分配列を分割する長さ」を一次元配列として順番に指定しており、最初の行 [1,2,3,4,5,6] は [1,2,3,4] と [5,6] に、次の行 [1,2,3,4,5] は [1] と [2,3]、最後の行 [1,2,3,4] は [2,1,3] のように、それぞれ指定された長さで部分配列に分割されます。こうすることで、既存の階層の内部に任意の長さの新しい階層（部分配列）を挿入することが可能になります。
a = ak.Array([[1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5], [1, 2, 3, 4]])
ak.unflatten(a, [4, 2, 1, 2, 2, 1, 3], axis=1)
[[[1, 2, 3, 4], [5, 6]], [[1], [2, 3], [4, 5]], [[1], [2, 3, 4]]] ------------------------ backend: cpu nbytes: 216 B type: 3 * var * var * int64
0 軸に新しい層を追加する例：
ak.unflatten(a, [1, 2], axis=0)
[[[1, 2, 3, 4, 5, 6]], [[1, 2, 3, 4, 5], [1, 2, 3, 4]]] --------------------------------- backend: cpu nbytes: 176 B type: 2 * var * var * int64
配列の結合#
concatenate() を使うと、軸方向に配列を結合できます。
デフォルト（axis=0）で結合：
a = ak.Array([1, 2, 3, 4])
b = ak.Array([2, 3, 4, 5])
ak.concatenate([a, b])
[1, 2, 3, 4, 2, 3, 4, 5] --- backend: cpu nbytes: 64 B type: 8 * int64
1 軸方向に結合：
ak.concatenate([a[:, None], b[:, None]], axis=1)
[[1, 2], [2, 3], [3, 4], [4, 5]] -------- backend: cpu nbytes: 64 B type: 4 * 2 * int64
データ型が異なる場合は union 型として結合されます。
c = ak.Array([[1, 2, 3], [3, 4]])
ak.concatenate([a, c])
[1,
2,
3,
4,
[1, 2, 3],
[3, 4]]
-----------
backend: cpu
nbytes: 150 B
type: 6 * union[
int64,
var * int64
]規則軸と非規則軸をもつ配列を結合すると、非規則軸になります。
ak.concatenate([a[:, None], c])
[[1], [2], [3], [4], [1, 2, 3], [3, 4]] ----------- backend: cpu nbytes: 128 B type: 6 * var * int64
ak.concatenate([c, a[:2, None]], axis=1)
[[1, 2, 3, 1], [3, 4, 2]] -------------- backend: cpu nbytes: 80 B type: 2 * var * int64
構造体配列#
辞書のリストも、リストの辞書も、どちらも同様に構造体配列として扱われます。
a = ak.Array({"x": [1, 2, 3, 4], "y": [2, 3, 4, 5]})
b = ak.Array(
[
{"vx": 10, "vy": 20},
{"vx": 11, "vy": 21},
{"vx": 12, "vy": 22},
{"vx": 13, "vy": 23},
]
)
a
[{x: 1, y: 2},
{x: 2, y: 3},
{x: 3, y: 4},
{x: 4, y: 5}]
--------------
backend: cpu
nbytes: 64 B
type: 4 * {
x: int64,
y: int64
}b
フィールドの追加・削除#
辞書のように新しいフィールドを追加できます。
a['z'] = 0
a
[{x: 1, y: 2, z: 0},
{x: 2, y: 3, z: 0},
{x: 3, y: 4, z: 0},
{x: 4, y: 5, z: 0}]
--------------------
backend: cpu
nbytes: 96 B
type: 4 * {
x: int64,
y: int64,
z: int64
}a['z'] = [3, 4, 5, 6]
a
[{x: 1, y: 2, z: 3},
{x: 2, y: 3, z: 4},
{x: 3, y: 4, z: 5},
{x: 4, y: 5, z: 6}]
--------------------
backend: cpu
nbytes: 96 B
type: 4 * {
x: int64,
y: int64,
z: int64
}フィールド削除は del a['z'] ですが、非破壊的に除外するには：
ak.without_field(a, 'z')
[{x: 1, y: 2},
{x: 2, y: 3},
{x: 3, y: 4},
{x: 4, y: 5}]
--------------
backend: cpu
nbytes: 64 B
type: 4 * {
x: int64,
y: int64
}フィールド追加の非破壊版：
ak.with_field(a, np.sqrt(a.x**2 + a.y**2), 'r')
[{x: 1, y: 2, z: 3, r: 2.24},
{x: 2, y: 3, z: 4, r: 3.61},
{x: 3, y: 4, z: 5, r: 5},
{x: 4, y: 5, z: 6, r: 6.4}]
-----------------------------
backend: cpu
nbytes: 128 B
type: 4 * {
x: int64,
y: int64,
z: int64,
r: float64
}unzip でフィールドを展開#
タプルとして出力：
ak.unzip(a)
(<Array [1, 2, 3, 4] type='4 * int64'>,
<Array [2, 3, 4, 5] type='4 * int64'>,
<Array [3, 4, 5, 6] type='4 * int64'>)
辞書として出力：
ak.unzip(a, how=dict)
{'x': <Array [1, 2, 3, 4] type='4 * int64'>,
'y': <Array [2, 3, 4, 5] type='4 * int64'>,
'z': <Array [3, 4, 5, 6] type='4 * int64'>}
unzip + zip で構造体配列をマージ#
ak.zip(ak.unzip(a, how=dict) | ak.unzip(b, how=dict))
[{x: 1, y: 2, z: 3, vx: 10, vy: 20},
{x: 2, y: 3, z: 4, vx: 11, vy: 21},
{x: 3, y: 4, z: 5, vx: 12, vy: 22},
{x: 4, y: 5, z: 6, vx: 13, vy: 23}]
------------------------------------
backend: cpu
nbytes: 160 B
type: 4 * {
x: int64,
y: int64,
z: int64,
vx: int64,
vy: int64
}構造体配列を2次元配列に変換#
構造体配列を 2 次元配列へ展開する方法を示します。
array = ak.Array(
[
{"x": 11, "y": 12, "z": 13},
{"x": 21, "y": 22, "z": 23},
{"x": 31, "y": 32, "z": 33},
{"x": 41, "y": 42, "z": 43},
{"x": 51, "y": 52, "z": 53},
]
)
array
[{x: 11, y: 12, z: 13},
{x: 21, y: 22, z: 23},
{x: 31, y: 32, z: 33},
{x: 41, y: 42, z: 43},
{x: 51, y: 52, z: 53}]
-----------------------
backend: cpu
nbytes: 120 B
type: 5 * {
x: int64,
y: int64,
z: int64
}一次元のレコードを二次元化：
array[:, None]
[[{x: 11, y: 12, z: 13}],
[{x: 21, y: 22, z: 23}],
[{x: 31, y: 32, z: 33}],
[{x: 41, y: 42, z: 43}],
[{x: 51, y: 52, z: 53}]]
-------------------------
backend: cpu
nbytes: 160 B
type: 5 * 1 * {
x: int64,
y: int64,
z: int64
}各フィールドを取り出す：
ak.unzip(array[:, None])
(<Array [[11], [21], [31], [41], [51]] type='5 * 1 * int64'>,
<Array [[12], [22], [32], [42], [52]] type='5 * 1 * int64'>,
<Array [[13], [23], [33], [43], [53]] type='5 * 1 * int64'>)
取り出した値を結合して行ベクトルに：
ak.concatenate(ak.unzip(array[:, None]), axis=-1)
[[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43], [51, 52, 53]] -------------- backend: cpu nbytes: 120 B type: 5 * 3 * int64
singletons() でも同様の処理が可能：
ak.concatenate(ak.unzip(ak.singletons(array)), axis=-1)
[[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43], [51, 52, 53]] -------------- backend: cpu nbytes: 168 B type: 5 * var * int64