Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PWACG: Partial Wave Analysis Code Generator supporting Newton-conjugate gradient method

Xiang Dong, Yu-Chang Sun, Chu-Cheng Pan, Ao-Yan Cheng, Ao-Bo Wang, Hao Cai, Kai Zhu

TL;DR

PWACG tackles the computational bottlenecks of partial wave analysis by automatically generating GPU-accelerated code from user configurations. It leverages JAX for automatic differentiation and the Newton-conjugate gradient method, implemented via template-based code generation with Jinja2 to assemble optimized PWA computations. The approach demonstrates improved efficiency in converging to global optima and scales across multiple GPUs, enabling larger and more complex fits. This framework offers a practical pathway to faster, scalable PWA analyses in high-energy physics with reduced reliance on traditional optimization toolkits.

Abstract

This paper introduces a novel Partial Wave Analysis Code Generator (PWACG) that automatically generates high-performance partial wave analysis codes. This is achieved by leveraging the JAX automatic differentiation library and the jinja2 template engine. The resulting code is constructed using the high-performance API of JAX, and includes support for the Newton's Conjugate Gradient optimization method, as well as the full utilization of parallel computing capabilities offered by GPUs. By harnessing these advanced computing techniques, PWACG demonstrates a significant advantage in efficiently identifying global optimal points compared to conventional partial wave analysis software packages.

PWACG: Partial Wave Analysis Code Generator supporting Newton-conjugate gradient method

TL;DR

PWACG tackles the computational bottlenecks of partial wave analysis by automatically generating GPU-accelerated code from user configurations. It leverages JAX for automatic differentiation and the Newton-conjugate gradient method, implemented via template-based code generation with Jinja2 to assemble optimized PWA computations. The approach demonstrates improved efficiency in converging to global optima and scales across multiple GPUs, enabling larger and more complex fits. This framework offers a practical pathway to faster, scalable PWA analyses in high-energy physics with reduced reliance on traditional optimization toolkits.

Abstract

This paper introduces a novel Partial Wave Analysis Code Generator (PWACG) that automatically generates high-performance partial wave analysis codes. This is achieved by leveraging the JAX automatic differentiation library and the jinja2 template engine. The resulting code is constructed using the high-performance API of JAX, and includes support for the Newton's Conjugate Gradient optimization method, as well as the full utilization of parallel computing capabilities offered by GPUs. By harnessing these advanced computing techniques, PWACG demonstrates a significant advantage in efficiently identifying global optimal points compared to conventional partial wave analysis software packages.
Paper Structure (7 sections, 7 equations, 7 figures, 1 table)

This paper contains 7 sections, 7 equations, 7 figures, 1 table.

Table of Contents

  1. Introduction
  2. Partial Wave Analysis Code Generator
  3. Architecture
  4. Implementation and Performance Optimization of Computational Templates
  5. Multi-GPU Support
  6. Performance Test
  7. Summary

Figures (7)

  • Figure 1: Workflow diagram illustrating the process for generating Partial Wave Analysis (PWA) code optimized for GPU execution.
  • Figure 2: The Dalitz plot from the MC sample.
  • Figure 3: The $K^+ K^-$ invariant mass spectrum from the MC sample.
  • Figure 4: The results of 300 fits for three different optimization methods. The convergence precision of FCN for method iMinuit+FDM was set to $10^{-2}$ , while the methods iMinuit+AutoDiff and NTCG+AutoDiff was set to $10^{-5}$.
  • Figure 5: Performance comparison of NTCG+AutoDiff at $xtol=10^{-8}$ and $xtol=10^{-5}$
  • ...and 2 more figures