JuTrack: A Julia package for auto-differentiable accelerator modeling and particle tracking

Abstract

Efficient accelerator modeling and particle tracking are key for the design and configuration of modern particle accelerators. In this work, we present JuTrack, a nested accelerator modeling package developed in the Julia programming language and enhanced with compiler-level automatic differentiation (AD). With the aid of AD, JuTrack enables rapid derivative calculations in accelerator modeling, facilitating sensitivity analyses and optimization tasks. We demonstrate the effectiveness of AD-derived derivatives through several practical applications, including sensitivity analysis of space-charge-induced emittance growth, nonlinear beam dynamics analysis for a synchrotron light source, and lattice parameter tuning of the future Electron-Ion Collider (EIC). Through the incorporation of automatic differentiation, this package opens up new possibilities for accelerator physicists in beam physics studies and accelerator design optimization.

Program summary

Program Title: JuTrack

CPC Library link to program files: https://doi.org/10.17632/r2g5zkwp7s.1

Developer's repository link: https://github.com/MSU-Beam-Dynamics/JuTrack.jl.git

Licensing provisions: MIT

Programming language: Julia

Nature of problem: Derivatives of the physics parameters calculated in accelerator modeling are critical for sensitivity analysis and optimization of the whole system. Traditional numerical approaches often rely on finite differences for derivative computations, which can prone to numerical inaccuracies. In highly nonlinear accelerator systems, like those encountered in synchrotrons and colliders, accurate sensitivity analysis and optimization require a large number of derivative evaluations. Thus, there is a need for more efficient methods to compute these derivatives accurately, especially when optimizing complex accelerator lattices or studying complicated collective effects, such as space-charge effects, wakefield effects, and beam-beam interaction.

Solution method: JuTrack addresses this problem by integrating compiler-level automatic differentiation (AD) into accelerator modeling routines, offering a powerful toolset for rapid derivative computation. Developed in the Julia programming language, JuTrack uses the Enzyme AD package to perform gradient-based analyses with minimal computational overhead. The package provides an efficient way to compute derivatives by directly differentiating through the model code, thus avoiding approximation errors associated with finite difference methods. It is designed to handle complex beam dynamics simulations, including complicated collective effects, such as space-charge effects, wakefield effects, beam-beam interaction, and combination of Truncated Power Series Algebra (TPSA) with AD. It can be applied to lattice optimization and beam dynamics analysis for future accelerators like the Electron-Ion Collider (EIC). Users can easily apply the package to their models, enabling robust optimization and sensitivity analysis in their accelerator studies.

Additional comments including restrictions and unusual features: JuTrack is particularly well-suited for scenarios requiring frequent derivative calculations, such as during beam dynamics optimization, sensitivity analysis, and accelerator tuning. Its integration with the Julia programming language provides excellent performance due to Julia's just-in-time (JIT) compilation capabilities. The modular nature of JuTrack and Julia's easy-to-understand syntax allows for future extensions and custom modifications, making it adaptable to a variety of accelerator configurations.