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

Alexander Scheinker; Alexander Plastun; Peter Ostroumov

physics.acc-phApr 4, 2026

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

Authors: Alexander Scheinker, Alexander Plastun, Peter Ostroumov

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

Abstract

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.

Figures & tables

Figure 2: Examples of conditional latent diffusion-generated latent embeddings and all 6 unique 2D projections associated with the 4D phase space densities that those latent images are decoded to by the VAE’s decoder.

Figure 3: 4D VAE architecture. Top: the encoder compresses the $128^{4}$ phase space tensor to a compact latent code $\mathbf{z}$ . Bottom: the decoder reconstructs $\hat{X}$ , from which all 2D marginal projections can be computed analytically and compared to the ground-truth projections during training.

Figure 4: The conditional latent diffusion architecture is a standard U-Net approach with 3 residual blocks at each resolution, GroupNorm, and attention and 100 denoising steps.

Figure 5: The conditional latent diffusion generative process is shown at various diffusion steps for different beam conditions. In this image, we show only the first three channels, as RGB values, of each $16\times 16\times 4$ latent images.

Figure 6: Overview of the the adaptive feedback setup.

Figure 7: Top: Error statistics are shown for test and training data, where Gaussians have been fit to the 1D $(x,x^{\prime},y,y^{\prime})$ projections that were created from the generated 4D tensors. Bottom: One detailed example of all of the generated 2D projections relative to the true projections for a single test data point at the end of the beamline, where the test data input was only the conditional vector into the latent diffusion model.

Figure 8: Top: The true $(x,y)$ projection of the beam is shown (blue) relative to the generated prediction (red) where tracking is performed by adaptive tuning of the diffusion conditional vector. Bottom: Although only the $(x,y)$ projection is available and used for tracking, the entire 4D phase space density is accurately tracked as seen by the true vs predicted $(\sigma_{x},\sigma_{x^{\prime}},\sigma_{y},\sigma_{y^{\prime}})$ fits of all 1D projections $(x,x^{\prime},y,y^{\prime})$ of the beam.

Figure 9: ES-based charge state tracking based on the $\pi_{xy}X(t)$ measurement (left) and error with and without tracking (right).

Figure 10: True and tracked projections of the beam are shown for various charge states during the ES-based tracking procedure.

Figure 11: True and generated examples from test set.

Figure 12: True and generated examples from test set.

Figure 13: True and generated examples from test set.

Figure 14: True and generated examples from test set.

Figure 15: True and generated examples from test set.

Figure 16: True and generated examples from test set.

Figure 17: True and generated examples from test set.

Figure 18: True and generated examples from test set.

Figure 19: True and generated examples from test set.

Figure 20: True and generated examples from test set.

Explore similar work

Apr 2, 2026

Generative models on phase space

Deep generative models such as diffusion and flow matching are powerful machine learning tools capable of learning and sampling from high-dimensional distributions. They are particularly useful when the training data appears to be concentrated on a submanifold of the data embedding space. For high-energy physics data, consisting of collections of relativistic energy-momentum 4-vectors, this submanifold can enforce extremely strong physically-motivated priors, such as energy and momentum conservation. If these constraints are learned only approximately, rather than exactly, this can inhibit the interpretability and reliability of such generative models. To remedy this deficiency, we introduce generative models which are, by construction, confined at every step of their sampling trajectory to the manifold of massless N-particle Lorentz-invariant phase space in the center-of-momentum frame. In the case of diffusion models, the "pure noise" forward process endpoint corresponds to the uniform distribution on phase space, which provides a clear starting point from which to identify how correlations among the particles emerge during the reverse (de-noising) process. We demonstrate that our models are able to learn both few-particle and many-particle distributions with various singularity structures, paving the way for future interpretability studies using generative models trained on simulated jet data.

Zachary Bogorad, Ibrahim Elsharkawy, Yonatan Kahn +2

Theory Division, Fermi National Accelerator Laboratory, Batavia, IL, USA · Department of Physics, University of Toronto and Vector Institute, Toronto, ON, Canada and NERSC, Lawrence Berkeley National Laboratory, Berkeley, California, USA · Department of Physics, University of Toronto and Vector Institute, Toronto, ON, Canada +2

Generative Modeling Constrained Generative Modeling Physics-Informed Generative Modeling Diffusion Models High-Energy Physics

Ask PDF

Oct 13, 2026

Dual-Domain Sparse-View 4D CBCT Reconstruction with Spatial-Temporal Consistency Learning

Four-dimensional cone-beam computed tomography (4D CBCT) provides phase-resolved volumetric images for image-guided radiotherapy. Yet good reconstruction quality usually needs many projections. This increases scan time and radiation dose. We propose STCNet, a sparse-view 4D CBCT reconstruction framework that uses spatial-temporal consistency learning. STCNet reconstructs high-quality images from only 16 projections in each respiratory phase. STCNet integrates physics-driven positional encoding to enforce geometric consistency in the projection domain and introduces temporal encoding in latent space to capture respiratory motion dynamics across phases to capture dynamic respiratory motion across phases. A dual-domain collaborative strategy jointly optimizes both projection and image domains, enhancing structural fidelity and temporal coherence. We tested STCNet on the TCIA-LUNG dataset and a CIRS thorax phantom dataset. STCNet performed better than conventional methods and recent deep learning baselines in image quality, artifact reduction, and structural preservation. On TCIA-LUNG, STCNet reached PSNR values of 33.311, 32.695, and 32.658 dB at phases 2, 5, and 8. The corresponding SSIM values were 0.845, 0.823, and 0.825. On CIRS, STCNet achieved PSNR values up to 36.497 dB and SSIM values up to 0.877. These results support STCNet as a promising framework for low-dose, phase-resolved 4D CBCT reconstruction in image-guided radiotherapy.

Qi Lai, YuTong He

Image Reconstruction Cone-Beam CT Sparse-View CT Reconstruction

Ask PDF

Feb 26, 2026

Advancing accelerator virtual beam diagnostics through latent evolution modeling: an integrated solution to forward, inverse, tuning, and UQ problems

Virtual beam diagnostics relies on computationally intensive beam dynamics simulations where high-dimensional charged particle beams evolve through the accelerator. We propose Latent Evolution Model (LEM), a hybrid machine learning framework with an autoencoder that projects high-dimensional phase spaces into lower-dimensional representations, coupled with transformers to learn temporal dynamics in the latent space. This approach provides a common foundational framework addressing multiple interconnected challenges in beam diagnostics. For forward modeling, a Conditional Variational Autoencoder (CVAE) encodes 15 unique projections of the 6D phase space into a latent representation, while a transformer predicts downstream latent states from upstream inputs. For inverse problems, we address two distinct challenges: (a) predicting upstream phase spaces from downstream observations by utilizing the same CVAE architecture with transformers trained on reversed temporal sequences along with aleatoric uncertainty quantification, and (b) estimating RF settings from the latent space of the trained LEM using a dedicated dense neural network that maps latent representations to RF parameters. For tuning problems, we leverage the trained LEM and RF estimator within a Bayesian optimization framework to determine optimal RF settings that minimize beam loss. This paper summarizes our recent efforts and demonstrates how this unified approach effectively addresses these traditionally separate challenges.

Mahindra Rautela, Alexander Scheinker

Los Alamos National Laboratory, Los Alamos, NM, US

Latent Dynamics Modeling Inverse Problems

Ask PDF