Learn the Dynamics, Enforce the Physics: Generative Latent Operators for High-Dimensional Beam Evolution

Haohan Zou; Alexander Scheinker; Yuanyuan Shi

physics.acc-phOct 10, 2026

Learn the Dynamics, Enforce the Physics: Generative Latent Operators for High-Dimensional Beam Evolution

Authors: Haohan Zou, Alexander Scheinker, Yuanyuan Shi

Organizations: University of California San Diego · Los Alamos National Laboratory

Abstract

Particle-in-cell (PIC) simulation is indispensable for modeling beam dynamics in particle accelerators, but its computational cost limits many-query workflows, such as configuration exploration, virtual diagnostics, and controller design. We introduce GLOBE, a physics-constrained Generative Latent Operator for Beam Evolution, modeling the nonlinear, control-dependent, multiscale evolution of 3D beam--field dynamics covering $10^7$ spatial cells. GLOBE compresses beam--field states into a compact latent space and factors their evolution into two components. A control-conditioned neural operator learns the dominant dynamics, while an amplitude-normalized generative model captures the residual distribution, recovering fine-scale structures and conditional variations. GLOBE enforces field constraints by regime-specific, parameter-free operators: a Poisson solve reconstructs the electric field in the electrostatic regime, and a spectral projection imposes Gauss's law and a divergence-free magnetic field in the electromagnetic regime. On two high-resolution 3D particle-accelerator benchmarks, bunched-beam transport and plasma-wakefield acceleration, GLOBE lowers relative- $L_2$ error by 9% and 13%, low-frequency spectral error by 69% and 51%, and mid-frequency spectral error by 54% and 20% relative to the strongest baseline (DPOT), satisfies field constraints, and reduces hours of PIC simulation to a minute per rollout.

Figures & tables

Figure 1: Framework Overview: The encoder $\mathcal{E}_{\psi}$ compresses initial beam–field states $\mathbf{u}_{1:k}$ into latent $\mathbf{z}_{1:k}$ . At each rollout step $t$ , the latest $k$ latents $\mathbf{z}_{t-k+1:t}$ and machine operation controls $\mathbf{c}$ are given to the latent dynamics models. The neural operator $\mathcal{G}_{\theta}$ predicts the dominant evolution $\tilde{\mathbf{z}}^{\text{det}}_{t+1:t+h}$ and the flow-matching model $\mathcal{R}_{\xi}$ samples the residual $\tilde{\mathbf{r}}_{t+1:t+h}$ . Their sum $\tilde{\mathbf{z}}_{t+1:t+h}$ is appended to the history and the prediction advances. Finally, $\mathcal{D}_{\psi}$ reconstructs the mesh-resolved states and $\mathcal{P}$ enforces regime-specific field constraints.

Figure 2: Compressive Autoencoder : For each 3D input, every $2\times 2\times 2$ block is split by the low-/high-pass pair $(\mathbf{h},\mathbf{g})$ into $8$ wavelet subbands, then the encoder $\mathcal{E}_{\psi}$ compresses it to a grid latent $\mathbf{z}$ . The decoder $\mathcal{D}_{\psi}$ mirrors the hierarchy, closing with the inverse Haar transform. Around every down-/up-sampling module, a parameter-free space–channel shortcut ( $\oplus$ ) provides references across resolution by channel averaging and duplication with reshaping, so the corresponding convolutional branch fits only a residual correction.

	Bunched-Beam Transport				PWFA
Method	Params	Rel- $L_{2}$ $\downarrow$	$L_{\mathcal{F}}$ (low / mid / high) $\downarrow$	Time $\downarrow$	Params	Rel- $L_{2}$ $\downarrow$	$L_{\mathcal{F}}$ (low / mid / high) $\downarrow$	Time $\downarrow$
DPOT	91.4M	0.177	0.112 / 0.986 / 0.954	7.55 s	109.4M	0.224	0.156 / 0.377 / 0.634	3.04 s
UPT	92.6M	0.261	0.147 / 0.999 / 1.000	1384s	107.7M	0.378	0.251 / 0.640 / 2.585	114.5s
MPP	90.0M	0.347	0.384 / 0.991 / 0.978	943s	110.9M	0.325	0.443 / 1.279 / 3.212	73.5s
P3D-D	91.4M	0.329	0.197 / 0.961 / 2.160	91.3s	110.9M	0.520	0.314 / 0.789 / 19.3	17.9s
P3D-G	91.4M	0.490	0.437 / 0.323 / 1.115	1490s	111.0M	1.219	0.794 / 151 / 6002	309s

Table 1: Quantitative comparison of GLOBE and baselines on bunched-beam transport and PWFA.

Figure 3: Bunched-beam transport. Left: Transverse charge projection $\int\rho\,\mathrm{d}z$ , top: ground truth, bottom: one prediction sample from GLOBE-C. Right: Evolution of $\sigma_{x}$ , $\sigma_{y}$ , $\sigma_{z}$ , and $\sigma_{x}/\sigma_{y}$ . Due to rapid oscillation, the main curves of transverse quantities show period-wise extrema and the insets show frame-wise evolution. Prediction curve is the mean of eight GLOBE-C samples, and shading indicates the sample standard deviation.

Figure 4: PWFA. Left: longitudinal electric field $E_{z}$ in the $x$ – $z$ mid-plane, overlaid with pseudo-particles of driver, witness, and plasma electrons sampled from charge density fields of WarpX (top) and GLOBE-C (bottom). First column is the last input frame; the rest show the forecast. Right: evolution of $E_{z,\text{wit}}$ and $r_{\text{bubble}}$ .

Haar	Shortcut	$\mathcal{L}_{\mathrm{rel}}\!\downarrow$	$\mathcal{L}_{\mathrm{spec}}\!\downarrow$	$\mathcal{L}_{\mathrm{fd}}\!\downarrow$	Runtime (s) $\downarrow$
✗	✗	0.0891	0.1070	0.697	6.04
✓	✗	0.0547	0.0740	0.566	0.97
✓	✓	0.0457	0.0685	0.489	1.00

Table 2: Ablation of two designs in our autoencoder: Haar transform (Haar) and multiscale resolution shortcut (Shortcut). The last row combines both designs and corresponds to the proposed autoencoder.

	Rel- $L_{2}$ $\downarrow$	$\nabla\!\cdot\!\mathbf{E}-\rho/\varepsilon_{0}$ $\downarrow$	$\nabla\!\cdot\!\mathbf{B}$ $\downarrow$
w/o $\mathcal{P}$	0.191	$2.72\times 10^{-2}$	$6.85\times 10^{-3}$
w/ $\mathcal{P}$	0.194	$\mathbf{2.77\times 10^{-7}}$	$\mathbf{2.78\times 10^{-7}}$

Table 3: Ablation of the physics-constrained operator on PWFA. Constraint residuals are normalized RMS values averaged over frames and trajectories.

	Bunched-Beam Transport		PWFA
Method	Rel- $L_{2}$ $\downarrow$	$L_{\mathcal{F}}$ $\downarrow$	Rel- $L_{2}$ $\downarrow$	$L_{\mathcal{F}}$ $\downarrow$
HFS	0.163	0.090 / 0.995 / 1.004	0.835	2.04 / 6.44 / 258
GLOBE-D (Ours)	0.144	0.039 / 0.994 / 1.003	0.174	0.083 / 0.350 / 0.592
HFS+DM	0.173	0.088 / 0.682 / 1.034	3.258	46.8 / 1301 / 2446
GLOBE-C (Ours)	0.161	0.035 / 0.457 / 0.973	0.194	0.077 / 0.300 / 0.618

Table 4: Ablation of latent Transformer backbone vs. HFS.

Appendix figures & tables32 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Visualization of the Transformer backbone used by deterministic neural operator $\mathcal{G}_{\theta}$ and generative residual model $\mathcal{R}_{\xi}$ .

	Bunched-Beam Transport		PWFA
Method	Rel- $L_{2}$ $\downarrow$	$L_{\mathcal{F}}$ (low / mid / high) $\downarrow$	Rel- $L_{2}$ $\downarrow$	$L_{\mathcal{F}}$ (low / mid / high) $\downarrow$
GLOBE-D (no residual)	0.144	0.039 / 0.994 / 1.003	0.174	0.083 / 0.350 / 0.592
GLOBE-C ( $\bm{\sigma}_{r}$ only)	0.147	0.039 / 0.894 / 0.986	0.266	0.231 / 4.32 / 6.66
GLOBE-C (Ours)	0.161	0.035 / 0.457 / 0.973	0.194	0.077 / 0.300 / 0.618

Appendix

Table 5: Ablation of the amplitude normalization for the generative residual. GLOBE-C uses the factorized normalizer

s_{\omega}=\bm{\sigma}_{r}\odot s_{\mathrm{win}}\odot s_{\mathrm{tok}}

; the variant keeps only the global per-channel scale

\bm{\sigma}_{r}

. GLOBE-D (deterministic anchor, no residual) is shown for reference.

Figure 6: Bunched-beam Transport Run 392: transverse beam evolution. Ground-truth and predicted density projections are shown together with the evolution of $\sigma_{x}$ , $\sigma_{y}$ , $\sigma_{z}$ , and $\sigma_{x}/\sigma_{y}$ .

Figure 7: Bunched-beam Transport Run 392: longitudinal beam evolution. Ground-truth and predicted longitudinal projections are compared over the rollout, with relative density errors shown below.

Figure 8: Bunched-beam Transport Run 392: fine-scale structure in the $y=0$ longitudinal slice. The deterministic prediction captures the bulk density evolution, while generative refinement recovers the realization-dependent mid-band structure observed in WarpX.

Figure 9: Bunched-beam Transport Run 244: transverse beam evolution. Ground-truth and predicted density projections are shown together with the evolution of $\sigma_{x}$ , $\sigma_{y}$ , $\sigma_{z}$ , and $\sigma_{x}/\sigma_{y}$ .

Figure 10: Bunched-beam Transport Run 244: longitudinal beam evolution. Ground-truth and predicted longitudinal projections are compared over the rollout, with relative density errors shown below.

Figure 11: Bunched-beam Transport Run 244: fine-scale structure in the $y=0$ longitudinal slice. The deterministic prediction captures the bulk density evolution, while generative refinement recovers the realization-dependent mid-band structure observed in WarpX.

Figure 12: Bunched-beam Transport Run 389: transverse beam evolution. Ground-truth and predicted density projections are shown together with the evolution of $\sigma_{x}$ , $\sigma_{y}$ , $\sigma_{z}$ , and $\sigma_{x}/\sigma_{y}$ .

Figure 13: Bunched-beam Transport Run 389: longitudinal beam evolution. Ground-truth and predicted longitudinal projections are compared over the rollout, with relative density errors shown below.

Figure 14: Bunched-beam Transport Run 389: fine-scale structure in the $y=0$ longitudinal slice. The deterministic prediction captures the bulk density evolution, while generative refinement recovers the realization-dependent mid-band structure observed in WarpX.

Figure 15: Bunched-beam Transport Run 87: transverse beam evolution. Ground-truth and predicted density projections are shown together with the evolution of $\sigma_{x}$ , $\sigma_{y}$ , $\sigma_{z}$ , and $\sigma_{x}/\sigma_{y}$ .

Figure 16: Bunched-beam Transport Run 87: longitudinal beam evolution. Ground-truth and predicted longitudinal projections are compared over the rollout, with relative density errors shown below.

Figure 17: Bunched-beam Transport Run 87: fine-scale structure in the $y=0$ longitudinal slice. The deterministic prediction captures the bulk density evolution, while generative refinement recovers the realization-dependent mid-band structure observed in WarpX.

Figure 18: Bunched-beam Transport Run 343: transverse beam evolution. Ground-truth and predicted density projections are shown together with the evolution of $\sigma_{x}$ , $\sigma_{y}$ , $\sigma_{z}$ , and $\sigma_{x}/\sigma_{y}$ .

Figure 19: Bunched-beam Transport Run 343: longitudinal beam evolution. Ground-truth and predicted longitudinal projections are compared over the rollout, with relative density errors shown below.

Figure 20: Bunched-beam Transport Run 343: fine-scale structure in the $y=0$ longitudinal slice. The deterministic prediction captures the bulk density evolution, while generative refinement recovers the realization-dependent mid-band structure observed in WarpX.

Figure 21: PWFA Run 312: wake evolution in the $x$ – $z$ plane together with trajectory-level observables. The color map shows $E_{z}$ ; particles are overlaid by species. Right: predicted and reference observables with the ensemble sample deviation.

Figure 22: PWFA Run 312: detailed comparison at $t=25$ for representative field quantities. Left to right: ground truth, prediction, and error.

Figure 23: PWFA Run 312: relative-error maps for representative channels at selected times. Colors show $(\mathrm{prediction}-\mathrm{truth})/\max|\mathrm{truth}|$ .

Figure 24: PWFA Run 355: wake evolution in the $x$ – $z$ plane together with trajectory-level observables. The color map shows $E_{z}$ ; particles are overlaid by species. Right: predicted and reference observables with the ensemble sample deviation.

Figure 25: PWFA Run 355: detailed comparison at $t=25$ for representative field quantities. Left to right: ground truth, prediction, and error.

Figure 26: PWFA Run 355: relative-error maps for representative channels at selected times. Colors show $(\mathrm{prediction}-\mathrm{truth})/\max|\mathrm{truth}|$ .

Figure 27: PWFA Run 271: wake evolution in the $x$ – $z$ plane together with trajectory-level observables. The color map shows $E_{z}$ ; particles are overlaid by species. Right: predicted and reference observables with the ensemble sample deviation.

Figure 28: PWFA Run 271: detailed comparison at $t=25$ for representative field quantities. Left to right: ground truth, prediction, and error.

Figure 29: PWFA Run 271: relative-error maps for representative channels at selected times. Colors show $(\mathrm{prediction}-\mathrm{truth})/\max|\mathrm{truth}|$ .

Figure 30: PWFA Run 27: wake evolution in the $x$ – $z$ plane together with trajectory-level observables. The color map shows $E_{z}$ ; particles are overlaid by species. Right: predicted and reference observables with the ensemble sample deviation.

Figure 31: PWFA Run 27: detailed comparison at $t=25$ for representative field quantities. Left to right: ground truth, prediction, and error.

Figure 32: PWFA Run 27: relative-error maps for representative channels at selected times. Colors show $(\mathrm{prediction}-\mathrm{truth})/\max|\mathrm{truth}|$ .

Figure 33: PWFA Run 118: wake evolution in the $x$ – $z$ plane together with trajectory-level observables. The color map shows $E_{z}$ ; particles are overlaid by species. Right: predicted and reference observables with the ensemble sample deviation.

Figure 34: PWFA Run 118: detailed comparison at $t=25$ for representative field quantities. Left to right: ground truth, prediction, and error.

Figure 35: PWFA Run 118: relative-error maps for representative channels at selected times. Colors show $(\mathrm{prediction}-\mathrm{truth})/\max|\mathrm{truth}|$ .

Latent Dynamics Modeling

Latent dynamics modeling learns compact hidden-state representations and transition rules to predict, explain, and control evolving systems.

See all 251 papers

PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion
Apr 4, 2026Alexander Scheinker, Alexander Plastun, Peter Ostroumov4D Reconstruction Physics-Informed Diffusion Models
Los Alamos National Laboratory Los Alamos 87545, NM, USA · Facility for Rare Isotopes Beams Michigan State University, East Lansing, MI 48824, USA
Classifier-pruned Bayesian optimization for particle accelerator tuning: Exploring temporally structured manifold of 6D beam phase space
Dec 2, 2024Mahindra Rautela, Alan Williams, Alexander ScheinkerBayesian Optimization
Los Alamos National Laboratory NM, United States
Learning a Latent Pulse Shape Interface for Photoinjector Laser Systems
Feb 19, 2026Alexander Klemps, Denis Ilia, Pradeep Kr. Banerjee +3Representation Learning Generative Modeling
Institute for Data Science Foundations, Hamburg University of Technology, Germany · Deutsches Elektron-Synchrotron (DESY), Hamburg, Germany

Apr 4, 2026J/K move · Enter open · S save

PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion

Alexander Scheinker, Alexander Plastun, Peter Ostroumov

Los Alamos National Laboratory Los Alamos 87545, NM, USA · Facility for Rare Isotopes Beams Michigan State University, East Lansing, MI 48824, USA

4D Reconstruction Physics-Informed Diffusion Models Latent Diffusion Models High-Energy Physics

Open paper AskPDF arXiv

We address the problem of recovering a time-varying 4D distribution from a sparse sequence of 2D projections - analogous to novel-view synthesis from sparse cameras, but applied to the 4D transverse phase space density

ρ(x,p_x,y,p_y)

of charged particle beams. Direct single shot measurement of this high-dimensional distribution is physically impossible in real particle accelerator systems; only limited 1D or 2D projections are accessible. We propose PhaseFlow4D, a feedback-guided latent diffusion model that reconstructs and tracks the full 4D phase space from incomplete 2D observations alone, with built-in hard physics constraints. Our core technical contribution is a 4D VAE whose decoder generates the full 4D phase space tensor, from which 2D projections are analytically computed and compared against 2D beam measurements. This projection-consistency constraint guarantees physical correctness by construction - not as a soft penalty, but as an architectural prior. An adaptive feedback loop then continuously tunes the conditioning vector of the latent diffusion model to track time-varying distributions online without retraining. We validate on multi-particle simulations of heavy-ion beams at the Facility for Rare Isotope Beams (FRIB), where full physics simulations require

\sim

6 hours on a 100-core HPC system. PhaseFlow4D achieves accurate 4D reconstructions 11000

\times

faster while faithfully tracking distribution shifts under time-varying source conditions - demonstrating that principled generative reconstruction under incomplete observations transfers robustly beyond visual domains.

Abstract

Latent Dynamics Modeling

Explore similar work

PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion

Classifier-pruned Bayesian optimization for particle accelerator tuning: Exploring temporally structured manifold of 6D beam phase space

Learning a Latent Pulse Shape Interface for Photoinjector Laser Systems