Time-inversion of spatiotemporal beam dynamics using uncertainty-aware latent evolution reversal

TL;DR

This work tackles the inverse problem of estimating upstream $6D$ beam phase space from downstream measurements in a high-dimensional accelerator setting. It introduces the reverse latent evolution model (rLEM), a two-step framework that first uses a conditional variational autoencoder (CVAE) to map phase-space projections into a latent space and then employs an autoregressive LSTM to learn reverse temporal dynamics in that space. The approach captures aleatoric uncertainty in the latent representation and propagates it to upstream predictions, achieving high accuracy (training/test MSE around $5\times10^{-7}$ and $1\times10^{-6}$, SSIM around $0.998$ and $0.976$) while delivering a ~600x speedup over full physics simulations. Applied to the LANSCE accelerator with 48 modules and 15 projected phase-space views, the method provides fast, uncertainty-aware upstream reconstructions, with potential for real-time diagnostics and control, and it sets the stage for incorporating epistemic uncertainty in future work.

Abstract

Charged particle dynamics under the influence of electromagnetic fields is a challenging spatiotemporal problem. Many high performance physics-based simulators for predicting behavior in a charged particle beam are computationally expensive, limiting their utility for solving inverse problems online. The problem of estimating upstream six-dimensional phase space given downstream measurements of charged particles in an accelerator is an inverse problem of growing importance. This paper introduces a reverse Latent Evolution Model (rLEM) designed for temporal inversion of forward beam dynamics. In this two-step self-supervised deep learning framework, we utilize a Conditional Variational Autoencoder (CVAE) to project 6D phase space projections of a charged particle beam into a lower-dimensional latent distribution. Subsequently, we autoregressively learn the inverse temporal dynamics in the latent space using a Long Short-Term Memory (LSTM) network. The coupled CVAE-LSTM framework can predict 6D phase space projections across all upstream accelerating sections based on single or multiple downstream phase space measurements as inputs. The proposed model also captures the aleatoric uncertainty of the high-dimensional input data within the latent space. This uncertainty, which reflects potential uncertain measurements at a given module, is propagated through the LSTM to estimate uncertainty bounds for all upstream predictions, demonstrating the robustness of the LSTM against in-distribution variations in the input data.

Figures (21)

• Figure 1: Three 2D projections ($x-y$,$E-\phi$,$x'-y'$) out of 15 projections of 6d phase space of charged particle beam in the LANSCE linear accelerator. Accelerating modules - 1 to 4 are 201 MHz drift tube linac (DTL) and 5 to 48 are 805 MHz coupled cavity linac (CCL). The beam serves various scientific areas like isotope production facility (IPF), ultra-cold neutrons (UCN), proton radiography (PRAD), weapons neutron research (WNR), proton storage ring (PSR).
• Figure 2: Reverse latent evolution model (RLEM): CVAE captures a lower-dimensional latent distribution of a higher-dimensional 6D phase space followed by an autoregressive LSTM learning to reverse temporal dynamics in the latent space.
• Figure 3: Latent space visualization: Top row shows 7 2d projections ($Z_1 - Z_{2:7}$) of 8d latent space. The bottom row shows 2D PCA and 2D t-SNE of the 8D latent space. The first two figures in the row shows train-test points and the last two figures shows combined train-test colored based on 48 modules.
• Figure 4: $x-y$ projection
• Figure 5: $E-\phi$ projection