Python jupyter lab + awkward arrayで2列データの片方にゲートをかけながら動的にグラフを作る

この記事の Awkward のバージョンは1.X.Xです

基本的な例

phst.hateblo.jp

Tree形式のデータがあり、あるパラメータにカット条件をかけながら(ゲートをかけながら)、他のグラフをインタラクティブに確認したいことがある

次の例では、dataXとdataYの2つの変数をもつTree形式のデータ (ここでは Awkward Arrayを使う )があり、その片方 dataXにゲートをかけながら、他のグラフを描画する。

%matplotlib inline
import matplotlib.pyplot as plt
import ipywidgets
import numpy as np
import awkward as ak

plt.rcParams['figure.figsize'] = [5, 4]

seed=0
np.random.seed(seed)
tree = ak.Array({"dataX":np.random.randn(1000),"dataY":np.random.randn(1000)})

def f(xmax):
    fig, axes = plt.subplots(1, 3, tight_layout=True,figsize=(12,4))
    tree2 = tree[tree["dataX"]<xmax] #ゲート
    axes[0].hist(tree2["dataX"],range=[-5,5],bins=50)
    axes[1].hist(tree2["dataY"],range=[-5,5],bins=50)
    axes[0].set_ylim(0,110)
    axes[1].set_ylim(0,110)
    axes[2].hist2d(tree2["dataX"],tree2["dataY"],range=[[-5,5],[-5,5]],bins=[50,50])
    plt.show()

ipywidgets.interact(f, xmax=(-5,5,0.2))