Towards differentiable beam dynamics modeling in BLAST/ImpactX

Differentiable simulations are in demand in accelerator physics, demonstrating order-of-magnitude improvements for complex tasks such as many-parameter optimization for accelerator working points and reconstruction of hard-to-measure quantities. At its core, a differentiable simulation does not only solve a forward problem, but additionally provides gradients of output parameters (e.g. beam parameters) with respect to input parameters (e.g. beamline or source parameters). How to effectively program large dynamic simulations differentiably is still an open question, but there is general consensus that a “single-source” approach aided by automatic differentiation (AD) is desirable. Addressing this, there are a) emerging domain-specific languages in machine learning that are intrinsically differentiable, and b) highly-performing & scalable, general-purpose languages like ISO C++ of existing codes. The challenge of approach a) is syntax specialization, which can limit ease of implementation & performance for physics algorithms, while b) requires additional work for AD. Performance is important for modeling high-order beam dynamics and collective effects in accelerators. We compare the fast, modern codes ImpactX (C++/Python) and Cheetah (PyTorch) using traditional, gradient-free modeling. We then show progress in introducing single-source differentiability in ImpactX using modern compiler techniques, producing performant executables for gradient-based and gradient-free modeling.