Fast Graph-based Higher-Order Clustering Statistics on the GPU

技術要約：GPUを用いた高速グラフベース高次クラスタリング統計量

問題提起
銀河の空間分布は、原始的な密度場および重力進化に関する決定的な情報をエンコードしている。2点相関関数（2pCF）は、この分布を解析するための標準的なツールとして長らく用いられてきたが、後期非線形重力進化や銀河バイアスによって生成される非ガウス構造を捉えるには不十分である。高次統計量、特に3点相関関数（3pCF）および4点相関関数（4pCF）は、バイアスパラメータの縮退を解いたり、原始的な非ガウス性やパリティ対称性の破れを探索したりするための補完的な情報を提供する。しかし、大規模な宇宙論的カタログに対してこれらの統計量を計算することは、計算量的に極めて困難である。総当たり法は $O(N^k)$ のスケールを持つが、ツリーベースの手法を用いても $O(N m^{k-1})$ （$m$ は平均近傍数）に削減されるものの、現代のサーベイ規模（例：DESI）における $k=3$ および $k=4$ のコストは依然として高い。さらに、既存のコードの多くは、パリティ分解を伴う連結4pCFの直接的な実空間実装や、デバイスメモリを超えるグラフを処理可能な効率的なGPU加速機能を欠いている。

手法
本論文では、k-d木に基づいたグラフデータベース・アプローチを利用したアルゴリズムである GRAMSCI (GRAph Made Statistics for Cosmological Information) のバージョン2を提示する。コアとなる手法は、初期座標からコンパクトでソートされた隣接構造を構築し、生データを再訪問することなく $N$ 個のタプルをカウントするものである。

このアップデートでは、5つの主要な手法上の進展を導入している：

1. マージ・ウォーク（Merge-Walk）アルゴリズム： 元のバイナリサーチによる内側ループは、ハブ・スポークのペアごとに $O(m \log m)$ のコストを発生させていたが、これをマージ・ウォークアルゴリズムに置き換えた。隣接リストのソート順を利用することで、アルゴリズムは集合の積集合を見つけるための線形パスを実行し、計算量を $O(m)$ に低減させる。
2. パリティ分解4pCF： 新しいエスティメータは、4pCFの信号をパリティ偶成分（$\zeta_+$）とパリティ奇成分（$\zeta_-$）のチャネルに分離する。これは銀河四面体の符号付き体積を計算することによって達成される。浮動小数点ノイズによるパリティ奇成分への汚染を避けるため、退化した体積（実質的にゼロ）を持つ四面体には、決定論的にゼロの重みが割り当てられる。
3. 連結4pCFの推定： コードは、同一のグラフおよびビニング上で、非連結な寄与（2pCFの積）を内部的に計算する。これにより、全4pCFから非連結項を直接差し引いて、連結成分（トリスペクトラム）の信号を孤立させることが可能となり、これは原始的な非ガウス性の研究に不可ントである。
4. Pythonインターフェース： 軽量なラッパーにより、FortranエンジンをNumPy配列から直接呼び出すことができ、ファイルI/Oの処理および構造化された結果オブジェクトの返却が可能である。
5. アウト・オブ・コア（Out-of-Core）サポートを備えたGPU移植： 完全なクエリエンジン（3pCF、4pCF、およびパリティ4pCF）が、OpenACCを用いてNVIDIA GPUに移植された。実装には圧縮疎行列（CSR）データレイアウトを使用し、「ギャング・アンド・ベクター（gang-and-vector）」マッピング戦略を採用している。これは、1つのスレッドブロックが1つのハブを処理し、ベクターレーンが内側の近傍ループを処理する手法である。デバイスメモリを超えるグラフ（例：BAOスケールのサーベイ）を扱うために、アウト・オブ・コア・タイリングスキームを用いてエッジ配列をウィンドウごとに分割し、正確なカウントを維持しながらグラフをチャンクごとに処理する。

主な結果
著者らは、新しいコードを元のCPU実装および解析的な注入テストと比較してベンチマークを行った：

• アルゴリズムの高速化： マージ・ウォークカーネルは、64スレッドのCPUノードにおいて、バイナリサーチ方式と比較して3pCFで $1.5\times$ から $2.0\times$、4pCFで $2.8\times$ から $3.8\times$ の高速化を実現した。この利得は、最大スケール（$r_{max}$）および近傍数が増大するにつれて大きくなる。
• GPU性能： 単体のコンシューマー向けGPU（NVIDIA RTX 3090 Ti）において、コードは64スレッドのCPUノードに対し、3pCFで $2.6\times$、4pCFで $9\times$ の高速化を達成した。GPUの性能は、浮動小数点スループットではなくメモリ帯域幅に依存してスケールする。
• スケーラビリティ： アウト・オブ・コア・タイリングスキームにより、デバイスメモリを大幅に超えるグラフの処理が可能となった。具体的なテストでは、24 GBのカード上で $9 \times 10^9$ エッジの3pCFを、仮に全データがメモリ上に存在する場合と比較してわずか約20%のオーバーヘッドで処理した。
• 正確性： GPUカーネルは、ビニングされたカウントにおいてCPUのリファレンス結果を正確に再現した。パリティ対称性の破れに対する解析的な信号注入テスト、およびパリティ対称なシミュレーション（Quijote）に対するヌルテストによる検証により、エスティメータが偏りなく、かつ退化した幾何学形状を正しく処理していることが確認された。
• 応用： コードはDESI DR1 LRGサンプル（300万個の銀河）に適用された。測定された3pCFおよび連結4pCFは、EZmockアンサンブルの予測値と一致しており、3pCFは期待通りのBAO特徴を示し、連結4pCFは非ガウス信号を孤立させている。

意義
本論文は、これらのアップデートが宇宙論における高次統計量の実用性を大幅に拡大すると主張している。4pCFの計算コストをコンシューマー向けハードウェア上でほぼ1桁削減することにより、本コードはDESIのような大規模なカタログに対して、全BAO範囲にわたる連結4pCFを評価することを可能にする。これは、共分散推定のための数千のモック実現系が必要であったという従来の制約を打破するものである。パリティ分解エスティメータの導入は、球面調和関数解析とは独立した、パリティ対称性の破れ探索のための直接的な実空間クロスチェックを提供する。最後に、外部入力なしに連結4pCFを内部的に計算できる能力は、ビニングおよびペアセットの自己整合性を保証し、これは非ガウス的なトリスペクトラム信号を孤立させる上で極めて重要である。著者らは、コードが公開されており、現在および将来の大規模構造サーベイへの適用準備が整っていると結論付けている。