物理の駅 Physics station by 現役研究者

テクノロジーは共有されてこそ栄える

VScode: 独自の syntax highlighter 拡張機能を作る

シンタックス(構文) ハイライター用の拡張機能(Extension)を作る。拡張機能として公開することは目指さない。

インストールしていない場合は以下をインストールする。

VS Codeの雛形インストール

npm install -g yo generator-code
yo code

色々聞かれる。今回はLISE++ファイル(*.lpp)を読む拡張機能を作るので以下の設定とした。

 New Language Support を選ぶ
URL or file to import, or none for new: → `Enter`
What's the name of your extension? → `LISE-Reader`
What's the identifier of your extension? → `Enter`
What's the description of your extension? → `Enter`
 Language id: → `lise`
Language name: →`LISE`
File extensions: → `.lpp`
Scope names: → `lise-file.lpp`
Initialize a git repository? → Y
Open with code を選択

F5 でデバッグ開始すると、新しいvscode のウィンドウが立ち上がる。右下の言語モードの選択から LISE(lise) を選べるようになる。

lise.tmLanguage.json を見ると

{
    "$schema": "https://raw.githubusercontent.com/martinring/tmlanguage/master/tmlanguage.json",
    "name": "LISE",
    "patterns": [
        {
            "include": "#keywords"
        },
        {
            "include": "#strings"
        }
    ],
    "repository": {
        "keywords": {
            "patterns": [{
                "name": "keyword.control.lise",
                "match": "\\b(if|while|for|return)\\b"
            }]
        },
        "strings": {
            "name": "string.quoted.double.lise",
            "begin": "\"",
            "end": "\"",
            "patterns": [
                {
                    "name": "constant.character.escape.lise",
                    "match": "\\\\."
                }
            ]
        }
    },
    "scopeName": "lise-file.lpp"
}

これが意味するところは、

  • if while for return が完全一致でキーワードに
  • " で囲まれた文字は文字列扱いに。その中で \ 直後の1文字はエスケープ文字

である。実際に試してみよう。

LISE++ファイルの構文用を実装する必要がある。説明がめんどくさいので、GIthub のリポジトリを参考にしてほしい。

github.com

拡張機能として公開したい人はこのブログを参照されたし。

qiita.com

Python+matplotlib: ログスケールでの目盛り、サブ目盛りを設定する

十分な広さのfigsizeにグラフを書くと、軸はそれなりに表示される

import matplotlib.pyplot as plt

fig, axes = plt.subplots(ncols=2, nrows=2)
for ax in axes.flat:
    ax.plot([1,2], [0.00001,1])
    ax.set_yscale("log")
plt.tight_layout()
plt.show()

グラフに使える面積が小さくなると、目盛りやラベルを勝手に省略してしまうことがある。例えば、4x4のsubplotsの場合は目盛りがほとんど表示されなくなる。

import matplotlib.pyplot as plt

fig, axes = plt.subplots(ncols=4, nrows=4)
for ax in axes.flat:
    ax.plot([1,2], [0.00001,1])
    ax.set_yscale("log")
plt.tight_layout()
plt.show()

対処療法として、文字サイズを小さくするか、figsizeを大きくする方法があるものの、根本解決ではない。

1つ目の対応策は、set_xticks() set_yticks() で、目盛り=ticksを明示的に設定し、必要なticksを表示させる。

import matplotlib.pyplot as plt

fig, axes = plt.subplots(ncols=4, nrows=4)
for ax in axes.flat:
    ax.plot([1,2], [0.00001,1])
    ax.set_xticks([1.0,1.5,2.0])
    ax.set_yscale("log")
    ax.set_yticks([1e-5,1e-4,1e-3,1e-2,1e-1,1e-0])
plt.tight_layout()
plt.show()

2つ目の対応策は、locator を使う方法である。set_ticks では設定できない minor ticksの設定が可能になる。また、描画範囲が変わった時でもコードの修正を必要としない。

import matplotlib.pyplot as plt
from matplotlib.ticker import (MultipleLocator, LogLocator)

fig, axes = plt.subplots(ncols=4, nrows=4)
for ax in axes.flat:
    ax.plot([1,2], [0.00001,1])
    ax.set_yscale("log")
    ax.xaxis.set_major_locator(MultipleLocator(0.5)) #major ticksは0.5刻み
    ax.xaxis.set_minor_locator(MultipleLocator(0.1)) #minor ticksは0.1刻み

    ax.yaxis.set_major_locator(LogLocator(base=100)) #基数 100でに描画
    ax.yaxis.set_minor_locator(LogLocator(base=100,subs=[10])) #基数 100、その 10倍 も描画
    ax.yaxis.set_minor_formatter("") #minor ticksのlabelは表示されるので非表示に

plt.tight_layout()
plt.show()

Python awkward: parquet形式で1億イベントのtreeのIO(入出力)速度を圧縮形式で比較する

Python の awkward パッケージに実装されている parquet ファイル形式を入出力するときの処理時間を、圧縮形式ごとに比較してみる。配列の型は awkward.highlevel.Array と呼ばれる形式である。

8バイト(int64_t)✕8メンバー✕1億エントリーなので、圧縮なしの場合は6.4GBのデータサイズとなる。

圧縮形式 ファイルサイズ 出力時間 読み込み時間 理想IO時間
NONE(無圧縮) 6.40 GB 13.8 sec 16.5 sec 12.8 sec
SNAPPY(デフォルト) 3.21 GB 15.5 sec 18.6 sec 6.42 sec
LZ4 3.21 GB 11.0 sec 7.3 sec 6.42 sec
ZSTD 0.82 GB 19.0 sec 13.2 sec 1.64 sec
参考: ROOT LZ4 3.22 GB 37 sec 16 sec 6.44 sec
参考: ROOT ZSTD 0.56 GB 39 sec 25 sec 1.12 sec

理想IO時間は、SSDの読み書き速度を500MB/sとしたときのIO時間である。理想IO時間に、圧縮展開の時間、他の処理の時間を加えると、出力時間、読み込み時間になる。

NONEは、無圧縮なので6.4GBもの容量を使ってしまうものの、理想IO時間との差は当然ながら少なかった。デフォルト SNAPPYLZ4 は 圧縮率は同じだが、IO速度はLZ4の圧勝だった。ZSTDは圧縮率は高いもののLZ4より遅い。LZ4ZSTDは、圧縮率を取るか、IO速度を取るかで選べばよさそうだ。

なお、C++でコンパイルしたROOTのTFileのIO速度と同じアルゴリズムで比べると(圧縮率の差はあるが)Pythonの方が2倍早い。Treeをシンプルな処理で使う分にはPythonで十分であることが分かるだろう。

Python awkward の parquet ファイルのIO速度を検証したコード。ディスクキャッシュを防ぐため、数20GBのディスク容量を使う。

import awkward as ak
import numpy as np
import time
import os

comps = ["NONE", "SNAPPY", "LZ4", "ZSTD"] #"GZIP", "BROTLI"は遅すぎるので除外
for comp in comps:
    for i in range(3):
        start = time.time()

        N = 1 * 10000 * 10000
        tree = ak.Array({})
        tree["x"] = np.arange(N, dtype=np.int64)
        tree["y"] = np.arange(N, dtype=np.int64)
        tree["ax"] = np.arange(N, dtype=np.int64)
        tree["ay"] = np.arange(N, dtype=np.int64)
        tree["ax0"] = np.arange(N, dtype=np.int64)
        tree["ay0"] = np.arange(N, dtype=np.int64)
        tree["ax1"] = np.arange(N, dtype=np.int64)
        tree["ay1"] = np.arange(N, dtype=np.int64)
        
        filename = f"tree{i}_{comp}.parquet"
        ak.to_parquet(tree, filename, compression=comp)

        end = time.time()
        print(f"WRITE {i} {comp} {os.path.getsize(filename)/1e9:.2f} GB {(end-start):.1f} sec")

        del tree

for comp in comps:
    for i in range(3):
        start = time.time()
        filename = f"tree{i}_{comp}.parquet"
        tree = ak.from_parquet(filename, lazy=True)
        print(np.max(tree["x"]),end=" ")
        print(np.max(tree["y"]),end=" ")
        print(np.max(tree["ax"]),end=" ")
        print(np.max(tree["ay"]),end=" ")
        print(np.max(tree["ax0"]),end=" ")
        print(np.max(tree["ay0"]),end=" ")
        print(np.max(tree["ax1"]),end=" ")
        print(np.max(tree["ay1"]),end=" ")
        end = time.time()
        print(f"READ {i} {comp} {(end-start):.1f} sec")
        
        del tree

ROOT: TFile形式で1億イベントのtreeのIO(入出力)速度を圧縮形式で比較する

CERN ROOT の TTreeを TFileで入出力するときの処理時間を圧縮形式ごとに比較してみる。Pythonで同様のことをやった記事も参照。

8バイト(int64_t)✕8メンバー✕1億エントリーなので、圧縮なしの場合は6.4GBのデータサイズとなる。

圧縮形式 ファイルサイズ 出力時間 読み込み時間
無圧縮=0 6.40 GB 36 sec 14 sec
kZLIB=1 1.23 GB 63 sec 22 sec
kLZMA=2 0.24 GB 220 sec 59 sec
=3 1.23 GB 62 sec 23 sec
kLZ4=4 3.22 GB 37 sec 16 sec
kZSTD=5 0.56 GB 39 sec 25 sec

結果を見る限り、デフォルト kZLIB は圧縮率、時間のバランスはよくない。データの質にもよるが、kZSTD を選んでおけばほとんどのデータに対応できるだろう。ただし、ROOT6のZSTDで圧縮したファイルをROOT5で読むことは(たぶん)できない。

ROOTファイル出力のコード。データは圧縮しやすい内容にした。WindowsでVisual Studioでコンパイルして実行した。

void writeTreeToFile(std::string filename, int compressionAlgorithm = 1) {
    TFile* file = new TFile(filename.c_str(), "RECREATE");
    if (compressionAlgorithm == 0) {
        file->SetCompressionLevel(0);
    }
    else {
        file->SetCompressionAlgorithm(compressionAlgorithm);
    }
    TTree* tree = new TTree("tree", "Tree Title");

    int64_t x, y, ax, ay, ax0, ay0, ax1, ay1;

    // 1億本
    int N = 1 * 10000 * 10000;
    tree->Branch("x", &x, "x/L", N);
    tree->Branch("y", &y, "y/L", N);
    tree->Branch("ax", &ax, "ax/L", N);
    tree->Branch("ay", &ay, "ay/L", N);
    tree->Branch("ax0", &ax0, "ax0/L", N);
    tree->Branch("ay0", &ay0, "ay0/L", N);
    tree->Branch("ax1", &ax1, "ax1/L", N);
    tree->Branch("ay1", &ay1, "ay1/L", N);

    for (int i = 0; i < N; i++) {
        x = i;
        y = i;
        ax = i;
        ay = i;
        ax0 = i;
        ay0 = i;
        ax1 = i;
        ay1 = i;
        tree->Fill();
    }

    tree->Write();

    file->Close();
    delete file;
}

ROOTファイル読み込みコード

#include <TFile.h>
#include <TTree.h>
#include <iostream>
void readTreeFromFile(std::string filename) {

    TFile* f = new TFile(filename.c_str());
    TTree* tree = (TTree*)f->Get("tree");
    std::cout << tree->GetMaximum("x") << " ";
    std::cout << tree->GetMaximum("y") << " ";
    std::cout << tree->GetMaximum("ax") << " ";
    std::cout << tree->GetMaximum("ay") << " ";
    std::cout << tree->GetMaximum("ax0") << " ";
    std::cout << tree->GetMaximum("ay0") << " ";
    std::cout << tree->GetMaximum("ax1") << " ";
    std::cout << tree->GetMaximum("ay1") << " ";
    f->Close();
}

圧縮アルゴリズムのROOT上の定義

ROOT::RCompressionSetting::EAlgorithm::kZLIB; //ZLIB デフォルト
ROOT::RCompressionSetting::EAlgorithm::kLZMA; //LZMA Lempel-Ziv-Markov chain-Algorithm
ROOT::RCompressionSetting::EAlgorithm::kOldCompressionAlgo; // 古いの
ROOT::RCompressionSetting::EAlgorithm::kLZ4; //LZ4
ROOT::RCompressionSetting::EAlgorithm::kZSTD; //ZSTD 

1億イベントのtreeデータのグラフ化の速度をPythonとROOTで比較する

データはここのを使う。64バイト(double 8個)×1億イベント(エントリー)である。

phst.hateblo.jp

結果のまとめ

言語 ライブラリ ファイル形式 処理時間 スレッド数
Python awkward parquet 46 sec 1
C++ ROOT TTree ROOT 100 sec 1
C++ ROOT RDataFrame ROOT 22 sec (不安定) 8
Python uproot+awkward ROOT 76 sec 1

まずはPythonコード。Pythonでは標準であろうparquet形式を読み込み、テンポラリーなメンバー da を作成する。さらに、カット条件を適用したテンポラリー tree_cutも作成する。matplotlibで描画する。手元のPCでは約46秒かかった。

import time
import awkward as ak
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter

start = time.time()

tree = ak.from_parquet("tree.parquet", lazy=True)

tree["da"] = ((tree["ax0"]-tree["ax1"])**2+(tree["ay0"]-tree["ay1"])**2)**0.5
tree_cut = tree[tree["da"]<0.4]

fig, axes = plt.subplots(ncols=2,nrows=2)
ax=axes.flat[0]
ax.hist2d(tree["x"],tree["y"],bins=[300,300],range=[[0,300],[0,300]])

ax=axes.flat[1]
ax.hist2d(tree["ax"],tree["ay"],bins=[150,150],range=[[-1.5,1.5],[-1.5,1.5]])

ax=axes.flat[2]
ax.hist(tree["da"], bins=1000,range=[0,2])
ax.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
ax.ticklabel_format(axis="y", style='scientific', scilimits=[3,3])
ax.set_xlim(0,2.0)

ax=axes.flat[3]
ax.hist(tree_cut["da"], bins=1000,range=[0,2])
ax.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
ax.ticklabel_format(axis="y", style='scientific', scilimits=[3,3])
ax.set_xlim(0,2.0)

plt.savefig("plot_python.pdf")

plt.clf()
plt.cla()
plt.close()

end = time.time()
print(f"{(end-start)*1000:.0f} [msec]")

次に、ROOTで書く。普通の TFile TTreeを使った書き方である。手元のPCでは約100秒 (Pythonの約2倍)かかった。

#include <TFile.h>
#include <TTree.h>
#include <TH2D.h>
#include <TH1D.h>
#include <TCanvas.h>

#include <iostream>

void makeGraph1() {

    TFile* f = new TFile("output.root");
    TTree* tree = (TTree*)f->Get("tree");

    TH2D* h1 = new TH2D("h1", "pos. dist.", 300, 0, 300, 300, 0, 300);
    TH2D* h2 = new TH2D("h2", "ang. dist.", 150, -1.5, 1.5, 150, -1.5, 1.5);

    TH1D* h3 = new TH1D("h3", "before cut", 1000, 0, 2);
    TH1D* h4 = new TH1D("h4", "after cut", 1000, 0, 2);

    TCanvas* c1 = new TCanvas();
    c1->Divide(2, 2);
    c1->cd(1); tree->Draw("y:x>>h1", "", "colz");
    c1->cd(2); tree->Draw("ay:ax>>h2", "", "colz");
    c1->cd(3); tree->Draw("sqrt((ax0-ax1)*(ax0-ax1)+(ay0-ay1)*(ay0-ay1))>>h3");
    c1->cd(4); tree->Draw("sqrt((ax0-ax1)*(ax0-ax1)+(ay0-ay1)*(ay0-ay1))>>h4", "(ax0-ax1)*(ax0-ax1)+(ay0-ay1)*(ay0-ay1)<0.4*0.4");
    c1->Print("plot_root1.pdf");
}

RDataFrameを使うと高速化できるらしい。ただし、手元の環境では、コンパイルには成功するものの、実行すると、正常終了したり、コード -1073740940 で終了しました。 というエラーメッセージが出て終了してしまったりと、挙動が安定しない。成功すると、全スレッド使って約22秒で終了する。数スレッド並列にするだけで30秒台になったので、有効なのは間違いないが、Pythonは並列なしで46秒なのでコアあたりの速度はPythonの方がこの程度ではまだ優位である。

環境

  • root_v6.30.04.win64.vc17
  • VS 2022 17.9.1
  • Windows 11 Pro 23H2
#include <ROOT/RDataFrame.hxx>
#include <TCanvas.h>

void makeGraph2() {

    ROOT::EnableImplicitMT();
    ROOT::RDataFrame df("tree", "output.root");

    auto dang_cut = [](double ax0, double ay0, double ax1, double ay1) {return (ax0 - ax1) * (ax0 - ax1) + (ay0 - ay1) * (ay0 - ay1) < 0.4 * 0.4; };
    auto df_cut = df.Filter(dang_cut, { "ax0","ay0","ax1","ay1" });

    auto h1 = df.Histo2D({ "h1","pos. dist.",300,0,300,300,0,300 }, "x", "y");
    auto h2 = df.Histo2D({ "h2","ang. dist.",150,-1.5,1.5,150,-1.5,1.5 }, "ax", "ay");
    auto h3 = df.Define("da", "sqrt((ax0-ax1)*(ax0-ax1)+(ay0-ay1)*(ay0-ay1))").Histo1D({ "h3", "before cut", 1000, 0, 2 }, "da");
    auto h4 = df_cut.Define("da", "sqrt((ax0-ax1)*(ax0-ax1)+(ay0-ay1)*(ay0-ay1))").Histo1D({ "h4", "after cut", 1000, 0, 2 }, "da");

    TCanvas* c1 = new TCanvas();
    c1->Divide(2, 2);
    c1->cd(1); h1->Draw("colz");
    c1->cd(2); h2->Draw("colz");
    c1->cd(3); h3->Draw();
    c1->cd(4); h4->Draw();
    c1->Print("plot_root_MT.pdf");
}

コード -1073740940 はDWORDで C0000374。ヒープ破損の例外が飛んだことを意味する。

最後に、Pythonでrootファイルを読む。uprootというライブラリを使うと良い。読んだ後の形式は awkward=ak とする。awkward から直接 parquetファイルを読んだ場合よりも 30秒余分にかかったので、(少なくともデフォルトの)ROOTファイルのパフォーマンスはparquetより低いことが分かる。

import uproot
tree = uproot.concatenate("output.root:tree",library ="ak")