Machine learning surrogate for charged particle beam dynamics with space charge based on a recurrent neural network with aleatoric uncertainty

Abstract

In this work, we develop a machine learning (ML) model with aleatoric uncertainty for the low energy beam transport (LEBT) region of the LANSCE linear accelerator in which we model the transport of a space-charge-dominated 750 keV proton beam through a lattice of 22 quadrupole magnets. Our ML model is developed based on data generated by a Kapchinsky–Vladimirsky (KV) envelope model of beam transport. We show that a recurrent neural network can be used as a dynamical surrogate model for fast prediction of the LEBT beam envelope. Furthermore, we endow the model with the prediction of aleatoric uncertainty and compare three different approaches. We demonstrate that the ML-based uncertainty quantification models are well calibrated and produce good estimates of the regions where the model is less certain about its predictions. This ML framework is a necessary step in the development of a real-time virtual diagnostic tool with uncertainty quantification that can be integrated into more complex downstream tasks (e.g., adaptive control or learning flexible control policies via reinforcement learning) for improved efficiency in beam operations. In future work, we plan to expand on this preliminary study by considering more realistic envelope models that include longitudinal momentum spread and dispersive effects in bending magnets, as well as particle tracking codes with 3D space charge (such as hpsim and gpt).