U-Net-based surrogate modeling for attosecond X-ray free-electron lasers

TL;DR

Attosecond XFEL pulse generation demands precise control of the longitudinal phase space (LPS), but conventional optimization is computationally intensive and invasive diagnostics impede operations. The authors develop a modified U-Net surrogate that predicts two-dimensional LPS distributions from three-phase accelerator settings, trained on start-to-end simulations for AttoSHINE with structure-aware and multi-scale losses. The surrogate achieves high fidelity to simulations, with a mean normalized mean absolute error of $13.83\%$ and an $R^2 = 0.9866$ on a 200-case test set, and delivers inference times around $150$ ms on CPU after GPU training. This forward model enables fast, noninvasive diagnostics and real-time, data-efficient tuning for LPS-sensitive attosecond XFEL operation at high repetition rates.

Abstract

Attosecond X-ray pulse generation in modern X-ray free-electron lasers relies on strongly compressed, precisely tailored electron bunches, making accurate diagnostics and control of the longitudinal phase space (LPS) essential. In the self-chirping scheme, collective effects in the linac generate a strong energy chirp that is converted into high peak current through pre-undulator compression, enabling isolated attosecond pulse generation. Reliable operation of this scheme depends on precise LPS control and fast diagnostics. In this work, we present a U-Net-based neural network surrogate that predicts two-dimensional LPS distributions directly from accelerator settings. The model exhibits excellent agreement with start-to-end simulation results. These results demonstrate the potential of neural network surrogates to facilitate real-time tuning and control in attosecond X-ray pulse generation.

Figures (5)

• Figure 1: Schematic of the modified U-Net architecture. Feature maps are annotated as $C\times H\times W$. The encoder uses $2\times2$ max pooling followed by residual blocks. The bottleneck includes an ASPP module. The decoder uses bilinear upsampling and attention gated skip fusion. Encoder features are concatenated with decoder features at matching scales, and padding is applied only when size alignment is required. A final $1\times1$ convolution projects decoder features to the output density map.
• Figure 2: Schematic layout of the self-chirping process in AttoSHINE and the corresponding evolution of the electron-beam LPS. Insets show representative LPS distributions (colored density) together with the current profile (gray curve) at three locations along the beamline. (a) After nonlinear compression downstream of BC2. (b) At the end of L4. (c) After further compression in LTU1. The arrows labeled a--c indicate the corresponding observation points. The left is the head of the electron beam.
• Figure 3: Comparison of simulated and U-Net-predicted current profiles derived from the LPS at the end of L4 (observation point (b) in Fig. \ref{['example1']}) for the 200-case test set. (a) Representative reconstructed current profiles (including examples with the highest and lowest peak currents in the selected set). (b) Shot-by-shot comparison of peak current with shots sorted by the simulated peak current. (c) Shot-by-shot comparison of the full width at half maximum for the same shot ordering as in (b), showing the corresponding FWHM values for the peak-current-sorted cases. The close overlap between prediction and simulation indicates good agreement.
• Figure 4: Histogram of the NMAE and $R^2$ for 200 test set cases.
• Figure 5: Three representative test-set examples. For each case, the left panel shows the simulated longitudinal phase space, the middle panel shows the U-Net prediction, and the right panel compares the corresponding current profiles from simulation and prediction. From top to bottom, the NMAE values in phase space are 12.38%, 13.44%, and 18.21%, respectively.