A conditional latent autoregressive recurrent model for generation and forecasting of beam dynamics in particle accelerators

Abstract

Particle accelerators are complex systems that focus, guide, and accelerate intense charged particle beams to high energy. Beam diagnostics present a challenging problem due to limited non-destructive measurements, computationally demanding simulations, and inherent uncertainties in the system. We propose a two-step unsupervised deep learning framework named as Conditional Latent Autoregressive Recurrent Model (CLARM) for learning the spatiotemporal dynamics of charged particles in accelerators. CLARM consists of a Conditional Variational Autoencoder (CVAE) transforming six-dimensional phase space into a lower-dimensional latent distribution and a Long Short-Term Memory (LSTM) network capturing temporal dynamics in an autoregressive manner. The CLARM can generate projections at various accelerator modules by sampling and decoding the latent space representation. The model also forecasts future states (downstream locations) of charged particles from past states (upstream locations). The results demonstrate that the generative and forecasting ability of the proposed approach is promising when tested against a variety of evaluation metrics.

Figures (8)

• Figure 1: (A) A layout of the LANSCE linear accelerator and the related scientific areas is shown together with the 2D $(z,E)$ projection of the beam's charge density $\rho$ at various times along the accelerator. Modules M$_1$–M$_{48}$ are the resonant structures which accelerate the charged particle beam up to a final energy of 800 MeV. (B) Simplified representation of the beam where $\mathbf{x}_m$ represents all 15 of the 2D projections of the beam's 6D phase space when passing through M$_{m}$. $\mathbf{P}_m$ represents the amplitude and phase set-point of the electromagnetic field in M$_{m}$. (C) The inputs to the CVAE are $(\mathbf{x}_m,m)$, all 15 projections and also we condition the encoder on module number $m$ which is equivalent to conditioning on the amount of time the beam has spent moving through the accelerator. The CVAE allows us to study the beam dynamics based on a very low-dimensional latent representation of the beam's phase space. The beam's phase space $\mathbf{x}_{m-1}$ is embedded into a latent vector $\mathbf{z}_{m-1}$, which is then directly mapped to a latent vector representing the beam's subsequent state $\mathbf{z}_{m}$ by the LSTM, from which the decoder of the CVAE generates an estimate $\hat{\mathbf{x}}_{m}$ of the beam's phase space (only the $(x,y)$, $(p_x,p_y)$ and $(z,E)$ projections shown here).
• Figure 2: Three out of fifteen projections ($x-y$, $E-\phi$, $x'-y'$) of the 6D phase space of charged particle beams at different modules. The plots are shown on a logarithmic scale for better visualization. The plots illustrates how the projections evolve as they experience RF electromagnetic fields while moving through different modules.
• Figure 3: CLARM is a two step DL framework which combines a CVAE and an LSTM acting on the latent embedding of the CVAE. The CVAE transforms the phase space projections ($\mathbf{x}$) into the latent space ($\mathbf{z}$) and learns a probabilistic low dimensional distribution of ($\mathbf{x}$). The LSTM learns the temporal dynamics in the latent space. The CVAE and LSTM are coupled together in an autoregressive loop for forecasting phase space projections through the system.
• Figure 4: Conditional reconstruction ability of CLARM across three different modules: Left columns (3x5) shows original projections ($X$) from the test dataset whereas right set of columns (3x5) shows the reconstructions ($\hat{X}$) using CVAE. Bottom most plots show average reconstruction MSE and SSIM for different projections across 48 different modules and entire training and test set.
• Figure 5: Latent space visualization of CVAE: The top row shows 8D latent space using parallel coordinates plots where different colors corresponds to different modules and every curve is a point in the 8D latent space. The middle row shows various 2D projections of the 8D latent space ($Z_1 - Z_2$, $Z_1 - Z_4$, $Z_1 - Z_6$, $Z_1 - Z_8$). The bottom row shows the 2D PCA, t-SNE and UMAP of 8D latent space. The color maps for all of the rows are the same, corresponding to module number.