Advanced plasma-based particle accelerators and colliders promise more powerful, compact, and cheaper machines than today’s technologies. Called laser-plasma accelerators, or LPAs, they could herald a new era in fundamental and applied science, promising breakthroughs in high-energy physics and fusion research, and lead to better medical diagnosis and treatments.
However, extending the capabilities of existing accelerators and developing next-generation technologies like LPAs requires novel computational tools capable of capturing the full complexity of particle beam acceleration and beam dynamics.
For example, simulating the intricate, ultrafast processes that accelerate and guide the beams in LPAs—which can often act over picosecond (trillionths of a second) and micrometer scales—can be computationally intensive and costly. Moreover, reaching the relativistic energies needed in accelerators may require coupling hundreds of LPA stages (one after the other), further increasing the computational time and cost.
Researchers from the Accelerator Technology & Applied Physics Division at Berkeley Lab are developing a powerful new simulation tool that combines traditional accelerator simulation techniques with high-performance computing and machine learning (ML) algorithms. This so-called “ML-surrogate model” can simulate hybrid accelerator beamlines (comprising conventional and plasma-based accelerator elements) and promises a faster path toward start-to-end (whole-machine) modeling of plasma-based acceleration and transportation in and between coupled LPA stages.
While traditional hybrid beamline modeling usually involves alternating between simulation codes or reducing the physics sufficiently to fit it into one code, which can significantly increase the modeling time, their ML-surrogate model promises to dramatically improve the speed of modeling the plasma stages, making it faster to optimize beam transport between the acceleration stages (known as “transport gaps”).
The work is a significant step toward improving the performance of current accelerators and modeling advanced accelerator designs. It supports the development of next-generation accelerator technologies and their applications, including plasma-based colliders and X-ray free-electron lasers.
Hybrid beamlines
Rapidly and accurately modeling hybrid accelerator beamlines “is an important next step toward advanced accelerator modeling,” says Axel Huebl, a research scientist in ATAP’s Advanced Modeling Program (AMP) and the principal investigator for the research.
“Current approaches to modeling complex, multi-stage plasma accelerator designs and particle beams with realistic energy distributions spend most of their time modeling the intricate plasma-based acceleration elements. So, we focused on a new approach to optimizing transport between a chain of plasma acceleration stages to maintain beam quality as it passes from one stage to the next.”
He says, however, that accurate three-dimensional (3D) modeling of the plasma-based acceleration elements can require “highly detailed simulations involving thousands of hours of computational time on graphic processing units [GPUs] and can sometimes fill whole supercomputers, which is costly.” GPUs are powerful processors containing hundreds of arithmetic units that accelerate computationally intensive simulations.
Schematic of laser-plasma acceleration stages (grey) and interstage transport with focusing optics (blue). An electron beam is accelerated left to right through the accelerating and focusing elements. The modeling of each laser-plasma stage is replaced with surrogate models built from neural networks. (Credit: Berkeley Lab)
Although cheaper to model, these interstage transport sections contain intricate and sophisticated focusing plasma lenses, beam drift sections, and chromatic correction elements that require hundreds or thousands of simulations to fine-tune accurately.
The challenge, says Huebl, is “running the thousands of iterations needed to design and optimize the beamline without conducting thousands of costly simulations of these LPA elements.” However, he adds, building models of future or planned colliders that include complex plasma-based elements requires identifying “self-consistent parameters” to enable whole-machine modeling.
He says such models require “increasing the speed (thereby reducing the time) of modeling the complex plasma-based elements so they match the fast modeling of conventional beam transport elements.”
Combining WarpX and ImpactX
To build a model that can rapidly and accurately capture a hybrid beamline’s acceleration and beam dynamics as it passes from one acceleration stage to another, the researchers combined two state-of-the-art accelerator simulation codes—WarpX and ImpactX—into one framework. (This framework also has broader utility for modeling interaction point physics in conventional particle colliders.)
“Our approach, for the first time, combines these codes, which are tailored to address specific problems in accelerator modeling and cover the widely different time and length scales between conventional and plasma-based accelerator elements,” says Ryan Sandberg, a postdoctoral fellow in AMP, who performed the work to build and train the ML-surrogate model.
Particle-In-Cell (PIC) codes such as WarpX and ImpactX simulate a particle beam bunch using “numerical particles” that evolve self-consistently as they pass through the accelerator elements. The codes are part of the Beam, Plasma, & Accelerator Simulation Toolkit (BLAST), a unique and important open-source suite of several interoperable accelerator simulation codes. They are used for modeling acceleration and beam dynamics in conventional and plasma-based accelerators and colliders over various time and length scales—often spanning nanoseconds to seconds and micrometers to kilometers. (AMP leads and manages BLAST in collaboration with global accelerator modeling community partners.)
Developed under the U.S. Department of Energy’s Exascale Computing Project, the award-winning WarpX is a time-based code that simulates the detailed plasma-based accelerator elements, including the laser wakefields and plasma lenses. ImpactX is a specialized accelerator code for simulating the particle bunch relative to a reference beam trajectory. The two codes traditionally communicate via input/output particle data files using the open standard for particle-mesh data (openPMD).
“As a first key step toward a workflow that can handle the self-consistent modeling of the chain of plasma-based acceleration stages, we were able to show how WarpX and ImpactX can be combined to leverage and create more value by using a data-driven approach to modeling accelerator elements that each code cannot achieve individually,” explains Sandberg.
He says this combination of GPU-accelerated particle accelerator codes is unique to BLAST and meets the need for training “data-hungry” neural networks, which are the ML algorithms used in the tool. “We relied on the high-quality data from a single 3D WarpX simulation to train these neural networks to learn nearly arbitrary particle trajectories for our surrogate model for ImpactX.”
Combining traditional physics-based simulations on supercomputers with this data-driven workflow provides a surrogate model that can approximate computationally costly segments of a hybrid accelerator beamline and is significantly faster in evaluating simulations.
Sandberg says a detailed, surrogate-embedded simulation hybrid took only ten GPU seconds to evaluate (and could be tailored for sub-second runtime), compared with a well-resolved PIC simulation, which can take over a thousand GPU hours to run.
Moreover, he adds, their model showed close agreement (within a few percent) with the predicted particle beam emittance from conventional high-fidelity accelerator simulations after incorporating the optimized transport parameters into a full 3D WarpX simulation for final verification.
The work provides a powerful, data-driven tool for particle accelerator simulations and complements existing full-scale and reduced-physics computational tools. It is an important step toward improving the performance and energy reach of existing accelerators and modeling advanced accelerator designs.
Next steps
The researchers plan to advance the application of ML models to accelerator modeling by increasing the complexity of their surrogate models to capture the collective effects of multiple particles, including how they interact with each other and the accelerating structure and how these interactions impact the acceleration and beam dynamics of hybrid accelerator beamlines.
“Extending the new workflow to train surrogate models that can incorporate these collective effects is not straightforward and is a research topic being addressed by the accelerator community as part of a series of upcoming workshops,” says Jean-Luc Vay, a senior scientist and head of AMP. “The team, however, is excited to tackle this challenge with its newly built GPU-accelerated codes that can tightly integrate GPU-accelerated ML models.”
Under the leadership of Vay, the AMP team has “positioned itself at the forefront of advanced accelerator modeling,” said ATAP Division Director Cameron Geddes.
“Their application of artificial intelligence and machine learning techniques combined with supercomputing capabilities has created increasingly powerful models. Their work is helping to extend the capabilities and performance of current accelerators and advance the development of next-generation accelerator technologies for breakthrough science.”
The research described here won the best paper award at PASC’24, held from June 3-5, 2024, in Zurich, Switzerland. It was funded by Berkeley Lab’s Laboratory Directed Research and Development Program and the U.S. Department of Energy, Office of High Energy Physics, and used the resources of the National Energy Research Scientific Computing Center (NERSC).
To learn more …
- Sandberg, R. T., Lehe, R., Mitchell, C. E., Garten, M., Myers, A., Qiang, J., Vay, J.-L., and Huebl, A. “Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines,” Proc. of Platform for Advanced Scientific Computing (PASC’24), PASC24 best paper, 2024, https://doi.org/10.1145/3659914.3659937
- Huebl, A., Lehe, R., Zoni, E., Shapoval, O., Sandberg, R. T., Garten, M., Formenti, A., Jambunathan, R., Kumar, P., Gott, K., Myers, A., Zhang, W., Almgren, A., Mitchell, C. E., Qiang, J., Sinn, A., Diederichs, S., Thevenet, M., Grote, D., Fedeli, L., Clark, T., Zaim, N., Vincenti, H., and Vay, J.-L. “From Compact Plasma Particle Sources to Advanced Accelerators with Modeling at Exascale,” Proceedings of the 20th Advanced Accelerator Concepts Workshop (AAC’22), in print, 2023, https://arxiv.org/abs/2303.12873
- Huebl. A, Lehe, R., Mitchell, C. E., Qiang, J., Ryne, R. D., Sandberg, R. T., and Vay, J.-L. “Next Generation Computational Tools for the Modeling and Design of Particle Accelerators at Exascale,” 2022 North American Particle Accelerator Conference (NAPAC’22), TUYE2, pp. 302-306, 2022, DOI:10.18429/JACoW-NAPAC2022-TUYE2, https://impactx.readthedocs.io and https://warpx.readthedocs.io
- Sandberg, R. T., Lehe, R., Mitchell, C. E., Garten, M., Qiang, J., Vay, J.-L., and Huebl, A. “Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code,” 14th International Particle Accelerator Conference (IPAC’23), WEPA101, 2023, DOI:10.18429/JACoW-IPAC2023-WEPA101
- Myers, A., Zhang, W., Almgren, A., Antoun, T., Bell, J., Huebl, A., and Sinn, A. “AMReX and pyAMReX: Looking Beyond ECP,” submitted, arXiv:2403.12179
- Huebl, A., Ananthan, S., Grote, D. P., Sandberg, R. T., Zoni, E., Jambunathan, R., Lehe, R., Myers, A., and Zhang, W. “pyAMReX: GPU-Enabled, Zero-Copy AMReX Python Bindings including AI/ML,” software, 2023. DOI:10.5281/zenodo.8408733; github.com/AMReX-Codes/pyamrex
For more information on ATAP News articles, contact caw@lbl.gov.
