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

TL;DR

This paper presents a novel deep learning framework combining CVAE and LSTM to perform inverse spatiotemporal beam dynamics, enabling efficient upstream phase space estimation from downstream measurements with uncertainty quantification.

Contribution

It introduces a self-supervised latent evolution reversal model that estimates upstream charged particle beam states from downstream data, incorporating uncertainty awareness.

Findings

Accurately predicts 6D phase space upstream from downstream measurements.

Captures and propagates measurement uncertainty through the model.

Demonstrates robustness against in-distribution data variations.

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.