Efficient 6D Beam Distribution Reconstruction with Machine Learning

6D Phase Space Reconstruction with Generative Machine Learning Ryan Roussel rroussel@slac.stanford.edu LPAW 2025

Modern Challenges in Beam Diagnostics Beam shaping for high efficiency PWFA Modern accelerator applications require precise understanding of high dimensional beam structure ? ? = ?,??,?,??,?,?? RMS quantities Exact beam structure C. A. Lindstr m, et al. Phys. Rev. Lett. 126, 014801 Laser Wakefield Acceleration ?? ?? Vay, J-L., et al. NIMA 909 (2018): 476-479. 2

Machine Learning-Based Virtual Diagnostics Use ML to represent complex relationships between accelerator parameters and simulated / experimental measurements non-destructive measurements A.L. Edelen et al., IPAC18, THYGBE2 C. Emma et al, Instruments (2021) A. Scheinker Sci. Reports 14.1 (2024): 29303. Requires detailed knowledge of system dynamics and/or large data sets to accurately predict measurements as a function of accelerator parameters. 3

Directly Measuring 6D Beam Distributions Multiple scanning slits with dipole and bunch shape monitor @ SNS (ORNL) 5D tomography with polarizable X-band TDC @ ARES (DESY) Jaster-Merz, S., et al Proc. IPAC 22 (2022): 279-283. Cathey, Brandon, et al. PRL 121.6 (2018): 064804. 960 measurements ~ 28 hrs Tomographic reconstruction ~ 24 hrs 5 million (!) measurements ~ 32 hrs 4

Difficulties Scaling Conventional Tomography Analysis to 6D Information compression Costs of detailed beam representations 2D image 6D distribution Histogramming scales poorly with number of dimensions, ????? ?? {< ? >,??} Reasonable resolution, ? = 100 For a 6D distribution, ?????= 1012! Scalar values 1D projection(s) Often required by analysis constraints (analytical tractability, optimization simplicity, etc.) 5

Inferring Beam Distributions Using Optimization Solve an optimization problem: 1. Parameterize an initial beam distribution 2. Simulate beam dynamics to predict measurements 3. Adjust initial beam distribution to match predictions to measurements 6

Enabling Efficient 6D Reconstructions (2) Implement differentiable beam dynamics simulations to enable learning (1) Use generative machine learning to represent complex 6D beam distributions. 7

Generative Beam Representations Use a generative machine learning model to create arbitrary beam distributions Example: arrange particles to make the SLAC logo O(~1000) parameters of the neural network control the distribution of particles in 6D phase space Learn the beam distribution by optimizing the NN parameters 8

Enabling Efficient 6D Reconstructions (2) Implement differentiable beam dynamics simulations to enable learning O(~1000) parameters 9

Differentiable Beam Dynamics in Cheetah Differentiable beam dynamics simulation (written in PyTorch) Gonzalez-Aguilera, J., et al. Proc. IPAC 23 (2023): 1725-1729. Cheetah collaboration between DESY, SLAC, LBNL, UChicago, KIT https://github.com/desy-ml/cheetah J. Kaiser, et.al. Phys. Rev. Accel. Beams 27, 054601 (2024) 10

Generative Phase Space Reconstruction (GPSR) Now we can use gradient descent to optimize (train) beam parameters necessary for learning O(~1000) parameters Roussel, Ryan, et al. PRL (2023) 11

Reconstructing 4D Beam Distributions from Quadrupole Scans Reconstructed distribution Magnetized beam scan Proof of concept reconstruction from quadrupole scan @ AWA Roussel, R, et al. PRL (2023) FACET-II Quadrupole Scan (Online GPSR analysis < 1 min) Kim, S., et al. PRAB (2024) 12

Extension to 6-Dimensional Reconstructions Add longitudinal diagnostics to basic quadrupole scan: Training data: Roussel, R., et al. PRAB (2024) 13

Reconstruction Predictions Evaluate reconstruction accuracy: Compare reconstruction predictions with test data Predicted Measured Roussel, R., et al. PRAB (2024) 14

6-Dimensional Reconstruction of Beams at AWA Reconstructed 6D distribution Measurement: ~8 mins Training: ~16 mins 30 hrs 24 mins: 75x faster than previous measurements ML Reconstructions from scratch, no pretraining or prior data required! Roussel, R., et al. PRAB (2024) 15

6D Phase Space Reconstructions at PAL-XFEL 16

Conclusion GPSR enables 6D reconstructions using small datasets (< 20 measurements), no prior model training required! GPSR enables substantial improvements in analysis capabilities that can benefit our understanding of beam dynamics + high dimensional structure of beam distributions Available for use right now: https://github.com/roussel-ryan/gpsr pip install gpsr 17

Questions? SLAC - - - - - - - U. Chicago - Juan Pablo Gonzalez-Aguilera Argonne Wakefield Accelerator - Seongyeol Kim - John Power - Eric Wisniewski - Wanming Liu - Alexander Ody Auralee Edelen Chris Mayes Daniel Ratner Doug Storey Nathan Majernik Brendan O Shea Claudio Emma Thanks to the team! 18

Backup 19

Addressing Beam Jitter Heuristic: Center beam on screen for each shot using centroid measurement Pros: - Straightforward to implement Cons: - Cannot reconstruct centroid offset of 6d dist. - Assumes no shot-to-shot variance in distribution Currently used for analysis Formal: Variational inference (VI) / Hamiltonian Monte Carlo (HMC) Pros: - Isolates distribution of centroid jitter Cons: - Requires additional R&D - Assumes no shot-to-shot variance in distribution (so far) 20

Uncertainty Quantification via Ensembling FACET-II injector reconstructions using ensembling heuristic 21

Reconstruction Results (Gaussian Beam) 22

Reconstruction Results (EEX Beam) 23