Two-stage Convolutional Neural Network for six-dimensional phase space reconstruction

This paper presents a two-stage convolutional neural network that efficiently reconstructs the six-dimensional beam phase space at a cathode surface from just sixteen transverse screen images, offering a faster and less computationally demanding alternative to conventional tomography techniques for particle accelerator diagnostics.

The Gist

Imagine you are trying to figure out what a mysterious, invisible 3D object looks like, but you can only see its shadow on a wall. Now, imagine that shadow changes shape depending on how you tilt the object or the light source. If you only look at one shadow, you might guess it's a ball, but it could actually be a cube or a pyramid. To know the truth, you need to see many shadows from different angles and use your brain to piece them together.

This is exactly the challenge physicists face with particle accelerators. They need to know the full "shape" of a beam of electrons (which exists in six dimensions: position, direction, and time/energy). But they can't see the beam directly; they can only take pictures of it hitting a screen.

Here is a simple breakdown of how this paper solves that puzzle using a new kind of "digital detective."

The Problem: The 6D Puzzle

Think of an electron beam not as a simple stick, but as a complex, squishy cloud of particles moving at near light speed. To understand how well this cloud is performing, scientists need to map its 6D Phase Space.

• The Hard Part: Traditional methods are like trying to solve a 3D puzzle by looking at only one 2D picture at a time. They are slow, often require destroying the beam to measure it, and can't see the full picture (like missing the "time" or "energy" parts of the cloud).
• The Old Way (Tomography): Imagine trying to figure out the shape of a hidden object by rotating it 360 degrees and taking a photo at every single degree. This takes forever and requires very precise, expensive machinery.

The Solution: The "Two-Stage" AI Detective

The authors created a new Artificial Intelligence (AI) model, specifically a Convolutional Neural Network (CNN), which acts like a super-smart detective that can solve this puzzle in seconds.

Think of this AI as a two-step process:

Stage 1: Learning the "Language" of Shadows

First, the AI needs to learn how the beam behaves.

• The Analogy: Imagine a student learning to draw. They practice by looking at a single photo of a cat and trying to guess what the cat looks like from the side. They do this thousands of times with different cats.
• In the Paper: The AI is trained on millions of computer simulations. It learns: "If I see a beam image like this on the screen, and the magnets are set to this strength, the original beam must have looked like that."
• The Trick: They didn't just use random shapes; they taught the AI using mathematical patterns (Fourier series) so it could understand complex, wavy shapes it had never seen before.

Stage 2: Putting the Pieces Together

Once the AI knows the rules, it moves to the real puzzle.

• The Analogy: Now, the student is given 16 different photos of the same cat, taken from slightly different angles (by moving the light or the cat). Instead of guessing based on just one photo, the student looks at all 16 at once and says, "Aha! When I combine these 16 views, I can see the whole 3D cat perfectly."
• In the Paper: The AI takes 16 real photos of the electron beam taken at the KEK-ATF facility. These photos were taken by slightly changing the magnetic fields and the timing of the electron gun (like rotating the object). The AI combines all 16 images to reconstruct the full 6D shape of the beam at the very beginning (the cathode).

Why This is a Big Deal

1. Speed: The old way took hours or days. This AI does it in less than a minute. It's like switching from developing film in a darkroom to taking a digital photo and seeing it instantly.
2. No Special Hardware: You don't need a fancy, dedicated machine to rotate the beam 360 degrees. You just need a standard accelerator and a screen.
3. Non-Destructive: It doesn't need to break the beam to measure it.
4. Affordable: It runs on a standard, relatively cheap computer chip (a GPU), not a supercomputer.

The Results

The team tested this AI on real data from the KEK-ATF accelerator in Japan.

• They fed it 16 photos of the beam.
• The AI reconstructed the full 6D shape of the electron beam.
• The results matched what they expected: the beam was about 3 millimeters wide and lasted for about 13 picoseconds (a trillionth of a second).
• It successfully identified complex features, like "tails" on the beam (like a comet), even though it had never seen those specific shapes during its training.

The Bottom Line

This paper introduces a "smart camera" for particle physics. Instead of spending hours manually rotating magnets and taking measurements, scientists can now take a quick snapshot of 16 beam images, hit "run," and have a complete, 3D (well, 6D!) map of their particle beam instantly. This allows them to fix problems in real-time and build better, faster, and more efficient particle accelerators for everything from medical treatments to discovering new physics.

Technical Summary

1. Problem Statement

In modern particle accelerators (e.g., linear colliders, free-electron lasers, and synchrotron light sources), beam quality is critical and is often limited by the full six-dimensional (6D) phase space distribution (transverse $x, y$, longitudinal $t$, and their conjugate momenta $x', y', p_z$).

• Limitations of Conventional Methods: Traditional diagnostics (pepper pots, wire scanners) typically measure only 2D transverse projections. Tomographic reconstruction of 4D or 6D phase space requires dense angular coverage (0°–360°), precise magnet settings, and computationally intensive back-projection, making it time-consuming and often impractical for routine operation.
• Limitations of Existing ML Approaches: Recent Generative Phase Space Reconstruction (GPSR) methods can reconstruct 6D space but rely on backward-differentiable particle tracking, which is not available in standard simulation packages (requiring code modification). Furthermore, GPSR is computationally expensive, requiring high-end GPUs (e.g., A100) and long optimization times to maximize entropy.

The authors aim to develop a fast, practical, and computationally efficient method to reconstruct the full 6D phase space from a limited set of 2D real-space images without requiring differentiable simulations.

2. Methodology

The authors propose a Two-Stage Convolutional Neural Network (CNN) trained on synthetic data generated by the ASTRA particle tracking code. The method solves the inverse problem: mapping 16 transverse ($x-y$) screen images (taken at a dispersive chicane) to the 6D phase space distribution at the photocathode.

A. Data Acquisition Strategy

Instead of rotating the beam through 360°, the method utilizes phase-space rotations induced by varying two machine parameters:

1. Solenoid Field: Rotates the transverse phase space ($x-x'$ and $y-y'$).
2. RF Gun Phase: Rotates the longitudinal phase space ($t-p_z$) and induces energy spread.
   By scanning 4 RF phases and 4 solenoid fields, the system acquires 16 distinct 2D beam images at a dispersive chicane. The dispersion allows longitudinal information (energy spread) to manifest as transverse spatial spread.

B. Network Architecture

The model consists of an Encoder, a Transformer, and a Decoder, trained in two distinct stages:

• Stage 1 (Single-View Pre-training):
  • The model learns the inverse mapping from a single 2D image (plus its specific RF/Solenoid settings) to the 6D phase space.
  • The Transformer is frozen during this stage.
  • Goal: To teach the Encoder-Decoder how beam distribution changes affect the projected image under fixed beamline conditions.
• Stage 2 (Multi-View Fusion):
  • The model takes all 16 images and their corresponding settings simultaneously as input.
  • The Transformer is unfrozen and learns to combine information from all 16 views to resolve ambiguities inherent in single-image inversion.
  • Goal: To fuse complementary constraints from different rotation angles into a single, consistent 6D distribution.

C. Training Data & Loss Function

• Data Generation: 1,700 unique cathode distributions were generated using Fourier series functions (sin/cos combinations) to create diverse beam shapes (peaks, tails). This ensures the model can extrapolate to unseen beam shapes.
• Loss Function: A composite loss function is used to ensure accuracy in both intensity and shape:
  • Poisson Negative Log-Likelihood: Respects particle count statistics.
  • Mean Absolute Error (MAE): Penalizes pixel-wise intensity differences.
  • Cosine Similarity: Ensures the global shape and pattern of the distribution match the truth.

3. Key Contributions

1. Two-Stage Learning Strategy: The novel separation of learning single-view physics (Stage 1) and multi-view fusion (Stage 2) allows the model to handle the high complexity of 6D reconstruction without overfitting or generating average estimates.
2. Independence from Differentiable Simulations: Unlike GPSR, this method uses standard forward simulation codes (ASTRA) for training, making it compatible with existing accelerator software ecosystems.
3. Computational Efficiency:
   • Training: Can be performed on a modest, affordable GPU (NVIDIA RTX A400).
   • Inference: Reconstruction of the full 6D phase space takes less than one minute (experimentally ~5 minutes for data acquisition + <1 min for processing), enabling potential online beam diagnostics.
4. Robust Extrapolation: By training on Fourier-based synthetic shapes, the model successfully generalized to "comet-like" distributions (Gaussian + tail) that were explicitly excluded from the training set.

4. Results

The method was validated using both synthetic test data and experimental data from the KEK-ATF (Accelerator Test Facility) injector.

• Synthetic Validation:

  • The model successfully reconstructed 6D distributions with "comet" tails and complex correlations.
  • Quantitative Metric: The reduced chi-square ($\chi^2_{red}$) between truth and prediction for 15 phase-space planes averaged between 1.0 and 2.0, indicating strong agreement.
  • The model accurately captured correlations between transverse and longitudinal degrees of freedom.

• Experimental Demonstration (KEK-ATF):

  • The model reconstructed the 6D phase space from 16 real images taken at the chicane.
  • Transverse Beam Size: Reconstructed values ($2.6 - 3.0$ mm) matched UV laser spot size measurements ($2.0 - 4.0$ mm) within a few millimeters.
  • Bunch Length: Reconstructed temporal width ($13.0 - 13.2$ ps) was consistent with streak-camera references ($\approx 10.0$ ps).
  • Limitations: The model showed some convergence to specific momentum values ($x', y', p_z$) across different datasets, suggesting a need for broader momentum-space coverage in future training.

5. Significance

This work represents a significant step toward practical, real-time 6D beam diagnostics for accelerator facilities.

• Operational Feasibility: By reducing measurement time from hours (traditional tomography) to minutes and removing the need for specialized differentiable simulation codes, this method makes 6D phase space monitoring accessible for routine operation.
• Diagnostic Capability: It enables operators to visualize complex beam structures (correlations, tails, multiple peaks) that are invisible to standard emittance monitors, providing deeper insight into beam quality degradation mechanisms.
• Scalability: The framework is adaptable; while trained on RF/Solenoid scans, it can be retrained for other tuning parameters (e.g., quadrupoles, deflecting cavities) in different accelerator facilities.

In conclusion, the proposed two-stage CNN offers a computationally efficient, non-invasive, and highly accurate alternative to existing tomographic and generative methods, paving the way for advanced beam control in next-generation accelerators.