Learning efﬁcient hybrid particle-continuum representations of non-equilibrium N-body systems
Abstract
An important class of multi-scale, non-equilibrium, N-body physical systems deals with an interplay between particle and continuum phenomena.
These include hypersonic ﬂow and plasma dynamics, materials science, and astrophysics.
Hybrid solvers that combine particle and continuum representations could provide an efﬁcient framework to model these systems.
However, the coupling between these two representations has been a key challenge, which is often limited to inaccurate or incomplete prescriptions.
In this work, we introduce a method for Learning
Hybrid Particle-Continuum (LHPC) models from the data of ﬁrst-principles particle simulations.
LHPC analyzes the local velocity-space particle distribution function and separates it into near-equilibrium (thermal) and far-from-equilibrium (nonthermal) components.
The most computationally-intensive particle solver is used to advance the non-thermal particles, whereas a neural network solver is used to efﬁciently advance the thermal component using a continuum representation.
Most importantly , an additional neural network learns the particle-continuum coupling: the dynamical exchange of mass, momentum, and energy between the particle and continuum representations.
Training of the different neural network components is done in an integrated manner to ensure global consistency and stability of the
LHPC model.
We demonstrate our method in an intense laser-plasma interaction problem involving highly nonlinear, far-from-equilibrium dynamics associated with
[Figure]
[Table]
/sD6/
AFAQ5MJ</latexit
[Table]
<
[Table] [Table]
[Table] [Table]
[Table]
u
[Table]
u
[Table]
u
[Table]
/ud1UxNe+xmXSWpQssVHYSqIicnsfDLgCpkRE0soU9xmJWxEFWXGllSyJXjLJ6+S1kXVq1Uv72uV+l1eRxFO4BTOwYMrqMMtNKAJDCQ8wyu8Odp5cd6dj8Vowcl3juEPnM8f/aqRKQ==</latexit> u <latexit sha1_base64="o4uDT5G0zTHJLZrwITDN6TpIXjU=">AAAB8XicbVDLSgMxFL1TX7W+qi7dBIvgqsxIRZcFN+Kqgn1gO5RMeqcNzWSGJCOUoX/hxoUibv0bd/6NaTsLbT0QOJxzLzn3BIng2rjut1NYW9/Y3Cpul3Z29/YPyodHLR2nimGTxSJWnYBqFFxi03AjsJMopFEgsB2Mb2Z++wmV5rF8MJME/YgOJQ85o8ZKj72ImlEQZum0X664VXcOskq8nFQgR6Nf/uoNYpZGKA0TVOuu5ybGz6gynAmclnqpxoSyMR1i11JJI9R+Nk88JWdWGZAwVvZJQ+bq742MRlpPosBOzhLqZW8m/ud1UxNe+xmXSWpQssVHYSqIicnsfDLgCpkRE0soU9xmJWxEFWXGllSyJXjLJ6+S1kXVq1Uv72uV+l1eRxFO4BTOwYMrqMMtNKAJDCQ8wyu8Odp5cd6dj8Vowcl3juEPnM8f/aqRKQ==</latexit> our LHPC architecture.
It consists of three components: (1) a neural network g ﬂuid for evolution of thermal sub-population with ﬂuid representation; (2) a solver g ∗ for the evolution of the non-thermal energetic sub-population with particle representation; (3) neural networks that model the injection of particles from ﬂuid state ( g inject ).
(b) and (c) illustrate the architectures of the components (1) and (3).
the coupling between electromagnetic ﬁelds and multiple particle species. More efﬁcient modeling of these interactions is critical for the design and optimization of compact accelerators for material science and medical applications. Our method achieves an important balance between accuracy and speed: LHPC is 8 × faster than a classical particle solver and achieves up to 6.8-fold reduction of long-term prediction error for key quantities of interest compared to deep-learning baselines using uniform representations  Introduction
The dynamics of physical systems is often nonlinear and involves the competition of different processes across a wide range of spatial and temporal scales. This gives rise to local non-equilibrium conditions (in the thermodynamic sense) that results in the failure of common numerical approaches.
While continuum models ( e.g. , based on ﬂuid equations) can provide accurate and computationally efﬁcient descriptions of near-equilibrium systems at large scales, they break down when signiﬁcant departures from local equilibrium are encountered (often at small scales) and give rise to important N-body phenomena. Kinetic ( e.g. , particle-based) numerical methods can accurately describe these non-equilibrium phenomena but are very computationally intensive, limiting their practical application to small scales. Over the last decades, this has motivated efforts to develop hybrid algorithms that can more efﬁciently couple continuum and particle representations in a variety of ﬁelds, including hypersonic gas dynamics (Schwartzentruber & Boyd, 2006), high-energy-density physics (Fiuza et al., 2011), and plasma physics (Bai et al., 2015).
Plasmas — hot ionized gases of charged particles — are a particularly challenging class of complex physics systems, where long-range electromagnetic interactions inevitably drive multi-scale and far-from-equilibrium dynamics.
Indeed, plasma research associated with controlled nuclear fusion (Zylstra et al., 2022), advanced laser-plasma particle accelerators (Haberberger et al., 2012), and most space and high-energy astrophysical environments (Drake et al., 2006) have long stimulated the development of hybrid particle-continuum representations to accurately and efﬁciently capture the nonlinear dynamics of these systems (Winske, 1985; Fiuza et al., 2011; Bai et al., 2015; Liu et al., 2022).
In all these cases, small-scale kinetic processes can accelerate a very small group of particles to energies signiﬁcantly above the mean (thermal) energy, driving the system out of equilibrium.
Importantly, this small group (few %) of non-thermal particles can carry away a large fraction (up to 50%) of the system energy and thus impact its global evolution.
This motivated the recent development of hybrid representations that use a ﬂuid solver to model the near-equilibrium (thermal) part of the particle distribution and a particle-based kinetic solver to model the evolution of high-energy particles (Kowal et al., 2011; Bai et al., 2015; Guidoni et al., 2016).
However, the coupling of the two representations has been based on over-simpliﬁed phenomenological prescriptions that can limit their validity and applicability.
In this work, we introduce a method for Learning Hybrid Particle-Continuum (LHPC) models to address the key challenge of efﬁciently and accurately coupling ﬂuid and kinetic representations of farfrom-equilibrium N-body systems.
LHPC combines a classical particle-in-cell (PIC) solver to advance the non-thermal particle distribution with a neural network surrogate model that efﬁciently advances the thermal component using a continuum representation.
Most importantly, our key contribution is the use of an additional neural network to learn the self-consistent particle-continuum coupling: the dynamical exchange of mass, momentum, and energy between the particle and continuum representations.
This coupling is learned from the data of ﬁrst-principles PIC simulations, providing an accurate and physics-informed description that addresses the main limitation of previous hybrid methods.
We review related works in Appendix A.
While the combination of classical numerical solvers and deep learning has been explored before (Um et al., 2020; Vlachas et al., 2022), these have been primarily based on uniform representations.
The use of machine learning techniques to learn efﬁcient and accurate coupling between continuum and particle representations introduced in this work is a promising and important route for addressing the multi-scale and multi-physics challenge of modeling N-body non-equilibrium systems.
We demonstrate our method on a challenging non-linear, far-from-equilibrium N-body system: the interaction of an intense laser with a solid-density plasma and resulting particle acceleration.
Presentday high-power lasers reach intensities in excess of 10 20 W / cm 2 , which nearly instantaneously vaporize and ionize solid-state matter upon interaction, resulting in high-energy-density plasma.
These interactions give rise to nonlinear and kinetic processes that establish strong electric ﬁelds in the plasma and accelerate particles to high energies over very compact distances.
In the last two decades, there has been great interest in exploiting these laser-plasma based accelerators for a wide range of applications, including material science (Patel et al., 2003), imaging (Rygg et al., 2008), and medical therapy (Kroll et al., 2022).
Due to the need to capture detailed kinetic physics associated with these interactions, numerical modeling has relied primarily on fully kinetic PIC simulations, which are too computationally intensive – in fact, one-to-one modeling of experimental systems is computationally prohibitive even on the largest supercomputers.
Effective hybrid particle-continuum representations have not been previously established for this problem.
This is thus a prime example for testing our model and evaluating the ability of LHPC to learn efﬁcient and accurate hybrid representations that can have a transformative impact in the modeling of these complex laser-plasma interactions.
Our results show that LHPC achieves an important balance between accuracy and speed: LHPC is 8 times faster than a classical particle solver, and achieves up to 6.8-fold reduction of long-term prediction error for key quantities of interest when compared to deep-learning baselines using uniform representations.
This opens up a novel route to develop more efﬁcient and accurate particle-continuum descriptions across different domains involving non-equilibrium N-body systems, and represents a key step for data-driven multi-scale algorithms for nonlinear physical systems.
Problem setup and preliminaries
Task setup .
In general, simulations of the time-evolution of a dynamical system can be described as follows.
The state of the system at time t is s t , and there exists a ground-truth evolution g ∗ that evolves this state:
[Equation]
The ground-truth evolution g ∗ can either be the physical world which is challenging to predict, or a ﬁrst-principles numerical solver which can be slow and expensive for large-scale systems.
Assume that we have a pre-deﬁned mapping h which maps the state s t to a representation S t :
S t = h ( s t ) , t = 0 , 1 , 2 , . .
, such that S t captures the essence of the system.
For example, consider an original state s t represented by particles, and S t a ﬂuid representation which describes the system with statistics of the particles within each cell (see below).
Alternatively, h may be an identity mapping resulting in S t = s t .
A hybrid h mapping will apply both of these representations to different groups of particles.
Given the states { s t } , t = 1 , 2 , . . .
, T , the task is to design a proper representation S t = h ( s t ) , and learn a surrogate model g θ , typically in terms of a neural network, [Figure] [Figure]
un
i t )
[Figure]
2x [c/
[Figure]
[Figure]
[Figure]
u x [ m e c ]
[Figure]
[Figure]
[Figure]
10 0 # e l e c t r o n ( a r b .
un
i t ) n [ n c ]
[Figure]
[Figure]
{+< n v x > [ n c c ] 
+}
[Figure]
< n v x v y > [ n c c 2 ]
[Figure]
2x [c
[Figure]
which approximates the ground truth evolution of S t :
[Equation]
Here ˆ S t +1 is the model g θ ’s prediction of S t +1 .
Our goal is to have the long-term autoregressive evolution of the learned surrogate model ˆ S t + k
= g θ ◦ ... ◦
g θ ( S t ) (composing k times) to closely match the ground-truth S t + k .
Physical problem overview .
In this work, we consider the interactions of intense laser pulses with plasmas, where the laser ﬁeld accelerates the plasma electrons near the front surface, which in turn sets up a charge separation, generating electric (E-) and magnetic (B-) ﬁelds that accelerate and deﬂect the ions (Fig. 2; see Appendix B for details).
These interactions of the charged particles with E- and B-ﬁelds (collectively EM-ﬁeld) occur on ultra-fast time scales ( ∼ 10 − 15 s), and can produce electrons and ions with energies exceeding 10 6 electronvolt ( ≡ 1 MeV) – with the electrons moving at close to the speed of light c .
These are prime examples of non-equilibrium systems, for which computationally expensive ﬁrst-principles simulations are necessary to accurately capture the associated non-linear dynamics.
Particle and continuum representations .
The two dominant representations of the plasma system considered in our work are particle (kinetic) and continuum (ﬂuid) representations.
A particle representation { ( x ti , u ti , Q i ) } describes the position x ti , momentum vector u ti , and other static properties
[-i t ) 4 3 12x [c/ ω 0 ]   0 -1 u x [ m e c ] 310101010 
-]
Q i ( e.g. , mass or charge) of each particle i .
The particles interact with each other, and change their position and momentum accordingly.
In contrast, a ﬂuid representation M t is deﬁned on a ﬁxed grid or mesh.
The values on each vertex represent the moments (statistical averages) of the momenta of the particles residing around the vertex.
These statistical quantities summarize the states of a large number of particles, permitting a more compact representation.
Their physical meaning is well established.
For example, the ﬁrst three moments correspond to the ﬂuid observables mass density, ﬂuid velocity, and pressure.
We note that in general an exact representation of the full particle distribution requires an inﬁnite number of ﬂuid moments (see Appendix B.1).
However, if the particle distribution is near local thermodynamic equilibrium it can be described by a compact set of ﬂuid moments.
A hybrid particle-continuum model aims to provide an efﬁcient description of a more general distribution of particles by coupling a particle representation for far-from-equilibrium particles with a continuum representation for particles near equilibrium.
Thermal and non-thermal population .
In this context, we refer to the group of particles in a near equilibrium state as thermal , for which the momentum distribution can be well described by a
Gaussian ( i.e. requiring only the mean and variance).
On the other hand, we refer to the particles that are far from equilibrium as non-thermal .
These are the outliers from this Gaussian distribution.
In terms of the total particle distribution function, f , this allows us to decompose it into a thermal , f T , and non-thermal , f NT components, with f = f T + f NT .
See Fig.
3 for an illustration of this separation in physical problem of intense laser-plasma interactions considered in this work; details of the separation algorithm are presented in Appendix C.
Method 3.1 Method overview
Here we introduce Learning Hybrid Particle-Continuum Models (LHPC) to address the key challenges in hybrid representations of non-equilibrium dynamical systems.
Speciﬁcally LHPC learns two key components, for the evolution of the ﬂuid, and coupling between the particle and ﬂuid representations.
In Figs. 1, 3 and the following sections, we detail the data generation and preparation procedure, the main architecture, and its application to the study of intense laser-plasma interactions.
3.1.1 Data generation and preparation from ﬁrst-principles simulations
Although computationally expensive, full PIC simulations provide ground truth data from ﬁrstprinciples particle-based evolution of plasma dynamics.
We run full PIC simulations of laser-plasma interactions in a one spatial dimension, three velocity dimension (1D3V) setting, where the laser intensity is varied to generate multiple trajectories.
A trajectory is composed of a simulation with 2000 time steps of evolution, involving 1.5 million (numerical) particles evolving on a (cid:39) 10k-cell grid, and is 200 GB in size (see Appendix B.3 for details).
To train the two learned components of LHPC, we transform the ground truth particle data into the ﬂuid and particle representations at each time step.
This is achieved by ﬁrst applying an algorithm that separates the particles into thermal and non-thermal populations ( e.g. Fig.
3; for details see
Appendix C).
The labeled particle data is then transformed, with the non-thermal population keeping the original particle representation, and using the velocity moments (up to 2nd order) to represent the thermal ﬂuid.
As the system evolves, a small fraction of the particles from the thermal population will gain a signiﬁcant amount of energy to get detached from the thermal distribution, becoming non-thermal.
They are said to be “injected” into the non-thermal population.
The transformed ground truth data therefore provides the information of the evolution of both representations, and the coupling (injection) between them.
3.1.2 Architecture
LHPC evolves the system state S t at time t to state S t +1 at time t + 1 with three components (Fig. 1):
(1) A neural network g ﬂuid ,θ with
parameters
θ which evolves the (thermal) ﬂuid
Fig. 1b).
(2) A ground-truth solver g ∗ which evolves the non-thermal particles.
(3) A neural network g inject ,ϕ with parameters ϕ which models the interplay of the ﬂuid and particle populations (Fig. 1c).
In particular, it models the injection of new non-thermal particles from the ﬂuid, and updates the existing ﬂuid and particle populations accordingly.
At inference, these processes are iterated autoregressively to predict the state of the system into the long-term future.
3.2 LHPC for laser-plasma interactions
Here, we employ LHPC to model laser-plasma simulations.
The space in x is divided into C consecutive cells: c ∈ [ c left , c left +1 , . .
, c right ] in the x direction.
The state S t consists of:
[Equation]
where the velocity moments of different populations M and the non-thermal particles ( x, u ) interact through the ﬁelds ( E, B ).
The full pipeline in Fig.
7 in Appendix D further visualizes the dependencies amongst these components and documents what each quantity represents.
The following sequence of operations is used to evolve the system from time t to time t + 1 :
[Equation]
Advance ﬂuid moments according to electromagnetic (EM-) ﬁeld on the grid.
g M,θ is an instantiation of the g ﬂuid ,θ introduced in Sec. 3.1.2.
[Equation]
Compute moments of new non-thermal particles to be injected g M NT ,ϕ is an instantiation of the g inject ,ϕ introduced in Sec. 3.1.2.
[Equation]
Inject new non-thermal particles by sampling from distribution with M t +1 NT
[Equation]
Update ﬂuid moments to conserve mass, momentum, and energy
[Equation]
Advance existing particles and append newly injected non-thermal particles
[Equation]
Deposition of electric current J from particles and ﬂuid to grid
[Equation]
Advance EM-ﬁeld on the grid using Maxwell’s equations
[Equation]
The operators g M,θ and g M NT ,ϕ are instantiated from g ﬂuid ,θ and g inject ,ϕ and are deep neural network models learned from data.
The operator G and the ﬁeld solver advances the non-thermal particles with the classical PIC solver.
A more detailed description of each component a – g can be found in
Appendix D.1, and detailed network architecture is in Appendix F.
3.2.1 Baselines
We compare with state-of-the-art deep learning-based surrogate model of Fourier Neural Operator (FNO; Li et al. 2021), and a baseline CNN that has the same architecture as our ﬂuid evolution model without hybrid representation.
Both architectures are evaluated on two representations: (1) an all ﬂuid representation which models all particles as ﬂuid deﬁned on each cell; (2) a bi-Gaussian ﬂuid representation, that models the more static thermal particles and the more dynamic non-thermal particles each as a Gaussian.
We compare the hybrid pipeline with naive baseline (1), which treats the entire system as a ﬂuid and advances the system by learning g as:
[Equation]
We also compare with an improved “bi-ﬂuid" baseline that models the thermal population and the non-thermal particles each as a Gaussian, and advances the system by learning g as:
[Equation]
We care especially about the error on the moments of the non-thermal subpopulation M PIC to understand how the respective methods model the most dynamic subpopulation of particles.
Experiments
In this section, we aim to answer two questions: (1) does our LHPC model with hybrid representation achieve better accuracy than learning the evolution with a pure ﬂuid representation?
(2) Does LHPC offer speedup, compared to the full-PIC solver that evolves the full system in ﬁrst principles with particle representation?
We evaluate LHPC and baselines in the N-body laser-plasma interactions system described above.
We evaluate both single-step prediction error and auto-regressive rollout error over 50 steps, where the state of the system changes signiﬁcantly.
We use relative L2 as the error metric – the ratio between L2 norm of error over L2 norm of ground-truth values.
4.1 Data and Experiment
Both the baseline and main pipeline are trained over 9 trajectories of PIC simulation.
10 trajectories were generated that vary over the normalized laser intensity a 0 from [2 , 4 , . . .
, 20] and trajectory( a 0 = 10 ) is held-out for testing (see Appendix B for details).
In this way, we test the generalization of the models to novel initial conditions and environment setup.
Each dataset consists of T =time steps and uses the separation strategy (details in Appendix C).
For the nine datasets, we train with single-step loss on the time-range [0,1400], and validate on [1400, 1500].
We iterate on hyperparameters to obtain the best individual settings for both the baseline and hybrid pipelines respectively (detailed in Appendix).
Then, the best models are ﬁnetuned using the multi-step loss with data from timeranges [980 , 1000] ∪ [1080 , 1100] ∪ [1180 , 1200] ∪ [1280 , 1300] ∪ [1380 , 1400] , and validated on [1400, 1500].
We report the conﬁdence intervals of the 1-, 20- and 50-step relative
L2 errors on the held-out dataset and report the results in Table 1, which shows the error for the
EM-ﬁeld, the non-thermal particles ( M PIC ) and the full ﬂuid, respectively.
Runtime for ground-truth (GT) solver and the compared models are also provided in Table 1.
We provide additional results in
Appendix E with an ablation study on trajectory 10 with the most intense laser ( a 0 = 20 ).
4.2 Results
Error comparison .
From Table 1, we see that our LHPC typically outperforms the baselines in terms of error by a wide margin.
Among the three aspects we evaluate, the EM-ﬁeld error and the non-thermal particles ( M PIC ) are the most important, since the EM-ﬁeld is most sensitive to the local particle dynamics, and the non-thermal particles are typically the most important observable/deliverable of laser-plasma accelerators.
An all-ﬂuid description can be very efﬁcient and accurate at describing the thermal ( M ) population, as is common for traditional ﬂuid solvers.
However, it fails to capture non-thermal particle acceleration, manifested in large errors for the ﬁeld evolution (Field) and non-thermal population ( M PIC ).
Speciﬁcally, for Field, our LHPC achieves an error reduction of 91.0%, 91.3% and 85.4% (6.8-fold) for 1-, 20- and 50-step autoregressive rollout, compared to the best models among the baselines.
Similarly, in the evaluation of M PIC , our LHPC achieves an error reduction of 63.5%, 36.2%, 25.8% for 1-, 20- and 50-step predictions, respectively, demonstrating the ability of our LHPC to model accurately the dynamics of those energetic non-thermal particles.
The Baseline CNN, although having the same architecture as our LHPC, cannot model accurately the non-thermal particles via either all-ﬂuid or Bi-Gaussian representations, thus achieving much larger error.
Speciﬁcally, while the use of a separate ﬂuid description for the non-thermal population (Bi-Gaussian model) can help capture its impact on the system evolution, it still leads to signiﬁcantly larger errors in the ﬁeld evolution when compared to the LHPC model and cannot be used to fully describe the complex distribution of the non-thermal particles (see also
Sec. 4.3).
Comparing baseline and state-of-the-art FNO, we ﬁnd that FNO performs comparable or slightly worse than Baseline CNN, likely due to the simpler architecture of CNN being more appropriate to capture the local ﬂuid evolution.
On predicting the ﬂuid moments ( M ), our LHPC model performs comparably to Baseline that has the same CNN architecture, showing both models can model reasonably the dynamics of the ﬂuid.
Runtime comparison .
While achieving signiﬁcant error reduction on predicting the EM-ﬁeld and the non-thermal particles, our LHPC also achieves signiﬁcant speedup compared to the Ground-truth (GT) full PIC solver, with a reduction of runtime by 8.0-fold.
We also see that the runtime of
LHPC is not much higher than full deep learning based baselines, showing that although our LHPC involves a GT solver to evolve the non-thermal particles, the increase in runtime is negligible since the non-thermal particles only constitute a small fraction of the system. [Table]
4.3 Visualization of results
In this section we visualize the predictions of LHPC.
An accurate evolution of the system dynamics relies on the correct EM-ﬁeld (Eq. 1g), which is dictated by the electric current J (Eq. 1f).
Figureshows that our model can predict the evolution of J contribution from the non-thermal population accurately over 50 time steps of rollout.
This is further illustrated in Fig.
5, where for the different contributions of J , good agreement between the predictions of our model and the ground truth is observed.
Similarly, our model predictions components of E- ﬂuid density match well the ground truth after 20 steps of rollout.
These results demonstrate faithful evolution of the important quantities that describe the dynamics of the ﬂuid, the non-thermal particles, and the coupling between them.
Importantly, our model is able to accurately capture the critical aspect of laser-plasma interactions — non-thermal particle acceleration (Fig.
Note that the highly non-trivial energy distribution of the accelerated (non-thermal) particles cannot be properly described with a few ﬂuid moments and must rely on the accurate evolution of the particle representation.
Discussion and conclusion
This work helps address a key challenge in modeling multi-scale, non-equilibrium, N-body systems by providing a novel hybrid particle-continuum model, whose coupling is learned from the data of ﬁrst-principles simulations.
LHPC outperforms the classical ﬁrst-principles solver by an order of magnitude in speed, and the baseline method that models the entire system as a ﬂuid by up to an order of magnitude in accuracy.
Further speed-ups are expected in multiple dimensions as the non-thermal population constitutes a much smaller fraction of the system representation (see Appendix B.4).
Our method can be applied to other domains where efﬁcient computational descriptions of nonequilibrium N-body dynamics are required, and represents a key step towards the development of data-driven algorithms for multi-scale modeling of nonlinear physical systems.
A Related Works
Deep learning has recently emerged as a powerful tool to complement Um et al.
or serve as a surrogate for classical solvers (Sanchez-Gonzalez et al., 2020).
They accelerate simulation of physical systems through larger spatial resolution (Um et al., 2020; Kochkov et al., 2021) or temporal intervals (Li et al., 2021), explicit forward method (Tang et al., 2020; Wu et al., 2022b), or via reduced representations (Sanchez-Gonzalez et al., 2020; Wu et al., 2022a).
However, almost all works employ a uniform representation, either all particles or all ﬂuid, without considering the multi-scale characteristics present in many physical systems.
One important exception is LED (Vlachas et al., 2022), which performs hybrid evolution in the temporal dimension.
Speciﬁcally, they employ a solver to simulate the full system in some interval of time and use latent evolution to evolve in other parts of time, with a pre-deﬁned alternative schedule.
In comparison, our work addresses the multi-scale, multi-physics modeling of dynamical physical systems from a different perspective: it aims to learn an accurate and efﬁcient coupling between different representations involving different physical processes and scales, namely a hybrid particle-continuum representation which is critical to describe non-equilibrium N-body systems.
To the best of our knowledge, machine learning methods have not yet been used to address this problem.
B First-principles simulations of laser-plasma interactions
A plasma is a collection of unbound, moving ions and electrons, interacting through the electromagnetic (EM-) ﬁelds.
It can be produced through for example the irradiation of a high-intensity ( I > 10 18 W/cm 2 ) laser pulse onto a solid material such as a metal foil.
Through the nonlinear interplay between the electrons, ions, and the EM-ﬁelds, these laser-plasma interactions can generate high-energy ( > MeV) ions for applications in material science (Patel et al., 2003), imaging (Rygg et al., 2008), and medical therapy (Kroll et al., 2022).
To understand these nonlinear processes, it is necessary to use kinetic treatment of the plasma.
B.1 Particle and continuum descriptions of plasmas
One example of kinetic treatment of plasma is the particle-in-cell (PIC) method (Dawson, 1983;
Birdsall & Langdon, 1991).
In this method we solve the Klimontovich equation (Klimontovich 1967) for ﬁnite size particles, coupled to Maxwell’s equations for the EM-ﬁelds.
The numerical procedure consists of solving Maxwell’s equations on a spatial grid using the current and charge densities that are obtained by weighting the discrete plasma particles onto the grid.
The particles are then advanced via the Lorentz force associated with the EM-ﬁelds.
This particle-based simulation technique captures the kinetic microphysics of plasmas, and to the extent that quantum mechanical effects can be neglected, provides a ﬁrst-principles description of plasma dynamics.
A plasma can also be described exactly as a continuum, by ﬂuid equations – the evolution of the velocity moments of the distribution function f of the plasma particles ( M N ( x ) :
= (cid:82) d N vf ( x , v ) v N is the N -th order velocity moment).
The ﬂuid equations can be derived from the kinetic equations and form an inﬁnite hierarchy of exact coupled conservation equations for each ﬂuid moment (with the n -th moment depending explicitly on the ( n + 1) -th).
In practice, this inﬁnite hierarchy needs to be truncated after the ﬁrst few moments, through the so-called closure relation – a relation that expresses the evolution of the highest-order moment considered in terms of the lower-order moments.
For a plasma near local thermodynamic equilibrium, such as the thermal population of the electrons and ions studied in our work, these closures allow us to describe the plasma self-consistently at larger spatial and temporal scales (relative to the PIC scales), using a ﬁnite number of moments.
Indeed, we use the ﬁrst three moments – the mass, momentum, and energy densities of the particles – to describe the dynamics of the thermal ﬂuids, and learn their coupling with the kinetic ( non-thermal ) particles and the EM-ﬁelds.
B.2 Laser-plasma interactions
The generation of high-energy ion beams from intense laser-solid interactions has been an active area of research due to the potential of producing high-energy, high-charge, high-current ion beams in much more compact systems than solid-state based accelerators ( e.g. linear accelerators, cyclotrons).
The enormous accelerating gradients (TV/m; ∼ 5 orders of magnitude larger than solid-state based accelerators) can be sustained in a plasma (Wilks et al., 2001), accelerating ions to MeV energies within millimeters.
For these laser intensities it is useful to deﬁne the normalized vector potential a(cid:39) 0 .
85 (cid:113)
I
[ W/cm 2 ]( λ 0 [ µ m ]) 2 / 10 18 , where λ 0 is the laser wavelength.
The laser electric ﬁeld can accelerate electrons to relativistic speeds in one cycle if a 0 > 1 .
For ion acceleration, a solid-density target is typically used, for which the laser E-ﬁeld accelerates the electrons in the small, skin depth, layer near the front surface, producing very energetic electrons.
In addition to the acceleration near the front surface, these energetic electrons cross the dense target and escape into the vacuum on the rear side, setting up a strong charge separation E-ﬁeld that accelerates the ions from the back surface.
As a result, the ions will be laminar and possess small divergence angle, and exhibit an energy spectrum typically characterized by an exponentially decreasing distribution (Snavely et al., 2000; Wilks et al., 2001).
B.3 Simulation parameters
In our PIC simulations we consider an intense laser interacting with a planar, solid-density target.
In our 1D simulations a laser with frequency ω 0 is launched along the x direction from the left boundary and irradiates an electron-proton plasma ( i.e. m i = m p = 1836 m e ).
The laser pulse duration is (cid:39) 15 fs, with intensity 2 (cid:54) a 0 (cid:54) 20 .
The plasma density follows a step-like proﬁle with thickness (cid:39) 5 µ m and electron number density 10 23 cm − 3 (considering a laser wavelength of 1 µ m), corresponding to a solid-density target such as a metal foil.
The target is simulated with 1000 particles per cell per species, and the total simulation domain of (cid:39) 80 µ m is resolved with a spatial resolution (cell size) of 0.03 c/ω 0 .
The time step is chosen according to the Courant–Friedrichs–Lewy condition, and the system is evolved over 2000 time steps (or 8500 for a 0 = 20 ).
Periodic boundary conditions for both particles and ﬁelds are used.
B.4 Computational challenges and opportunities
Studying these laser-plasma interactions using the particle-in-cell (PIC) method requires resolving the fastest and smallest oscillations of the electrons in the plasmas.
As a consequence, modeling these interactions with ﬁrst-principles simulations is very computational demanding.
For example, a 3D one-to-one simulation of laser-plasma interactions typically involves evolving the dynamics of billions of (numerical) particles on a billion-cell grid, requiring millions of CPU hours to compute.
In this work, we consider only one spatial dimension x , while retaining the three dimensions in momentum.
This is identical to simulating a particle distribution uniform in the y and z directions.
See Appendix B for details of laser-plasma interactions and the PIC method.
It is important to note that in more realistic 2D and 3D geometries a signiﬁcant enhancement of speed-up can be achieved, due to the much smaller proportion of the non-thermal particles.
This is a result of the ﬁnite spot size of the laser, which will only interact with a ﬁnite volume of the plasma, and accelerate a small fraction of the particles to non-thermal.
In a 3D simulation of these laser-plasma interactions, the non-thermal particles compose typically < 0 .
1% of the total number of particles while encompassing > 50% of the system energy.
Compared with the 1D simulations presented in this work ( (cid:39) 20% ), one expects a speed-up of > 1000 x in 3D.
C Details for separation method
Here we describe the details of the separation method for preparing the labeled (thermal and nonthermal) data.
The algorithm separates the particles (from the original simulations) into thermal and non-thermal populations.
This is done by computing locally the moments of the particle velocity distribution f ( v ) and considering non-thermal particles those with velocities v that exceed a threshold αv th – a given multiple ( α ) of the thermal velocity v th ( = the velocity spread σ v for non-relativistic velocities; Cohen et al. 2010;
Fiuza et al. 2011).
At each time step, we study the velocity distribution function of the particles in the local neighborhood in space N i of each particle P i .
Without knowing which particles belong to the thermal population a priori , we begin by identifying a population that is likely thermal at iteration 1.
The velocity spread σ v 1 computed from this population provides a ﬁrst estimate of the threshold ασ v , allowing us remove particles with v > ασ v from the distribution.
The procedure is repeated (on the updated distribution) until a convergence of the value of the threshold is reached at iteration N .
The remaining particles now constitute the thermal population, characterized by the ﬂuid velocity v ﬂ = (cid:104) v j (cid:105) ∀ j ∈ N and thermal velocity v th = σ v N ( (cid:104)· · · (cid:105) denotes the average value).
These are the mean and sigma of the Gaussian.
The corresponding threshold αv th is then used to determine the population P i belongs to.
Note that only P i has been given a label, the other particles in N i were only used to calculate v ﬂ and v th .
# v v abs h abs # v
X h abs (v abs ,h abs ) vh
# v Y h(v abs ,h abs ) (v 0 ,h 0 ) v + v - # v (v abs
-v
fl ,h abs ) (v
-v
fl ,h 0 )
# v (v abs
-v fl ,h abs ) (v
-v
fl ,h 0 )
σ v ασ v [Figure] (e) [Figure]
Figure 6 illustrates an iteration in detail.
From the example velocity distribution function shown in Fig. 6(a), we recognize an absolute peak located at v abs with height h abs .
Among the peaks ( v µ , h µ ) ∀ h µ (cid:54) Xh abs for a constant X , pick the one with v µ closest to zero as the ﬁrst guess of thermal population and denote the peak location as ( v 0 , h 0 ) (Fig. 6b).
Search down and outwards from ( v 0 , h 0 ) ( i.e. in both directions of v with v + > v 0 > v − ) until the heights h
+ ≡ h ( v = v + ) =
Y hand h − ≡ h ( v = v − ) =
Y h 0 (Fig. 6c).
The population enclosed by { v | v + > v > v − } is regarded as the population to estimate the threshold.
Transform all the particles in N into the local ﬂuid frame according to v ﬂ (Fig.
6d; for non-relativistic velocities this amounts to v i → v i − v ﬂ ).
Calculate σ v to obtain the threshold (Fig. 6e).
The size of N i , values of X , Y , and α are tunable parameters determined by the physics and ﬁnetuned empirically.
We have found that N i (cid:39) 60 cells, X = 0 .
8 , Y = 0 .
1 , and α = 3 and 5 for the water and laser-plasma systems lead to the most physically accurate separation.
The accuracy of the algorithm is evaluated by the pressure tensor P ij of the thermal population, where P ij = (cid:104) v i v j (cid:105) − (cid:104) v i (cid:105)(cid:104) v j (cid:105) characterizes the second-order velocity moments and a is directly related to how well the distribution is described by a Gaussian.
Figure 6(f) exempliﬁes that indeed, for the identiﬁed thermal population of the ions for the laser-plasma system, the diagonal terms are all comparable and the off-diagonal terms are negligibly small.
Namely, the velocity distribution can be described by a Gaussian.
D Full Pipeline D.1 Component details D.1.1 Components (a) and (b)
For both M or M NT , there’s 20 moment directional components that needs to be predictedfor ions, 10 for electrons - of the zero ( 1 ), ﬁrst ( 3 ), and second ( 6 ) moment orders.
Thus, for both components, we instantiate 20 models with the same hyperparameter conﬁguration in f M or f M NT .
We also train Component (1) with the push-forward trick and multi-step loss, shown in Brandstetter et al.
(2022) to greatly enhance capability of generalization.
D.1.2 Component (c)
In order to parameterize the distribution to sample, a three-dimensional Gaussian is taken to have mean
E [ u ] ≈ the ﬁrst order moments in ˆ M t +1 NT , and covariance matrix E
[( u − E [ u ])( u − E [ u ]) T ] ∈ R 3 ×informed by the second order moments in ˆ M t +1 NT .
To be able to sample, we require a valid (positive deﬁnite) covariance matrix.
In order to enforce this condition, we obtain the ground-truth L using the Cholesky decomposition E [( u − E
[ u ])( u − E [ u ]) T ] = LL T and train f M NT to predict ˆ L .
During rollout, we approximate E [( u − E [ u ])( u − E [ u ]) T ]
≈ ˆ L ˆ L T , and use the valid covariance matrix to sample a set of particles which is then “injected".
The cumulative set of injected particles constitute the PIC loop.
D.1.3 Component (d)
After obtaining the sample of injected particles, we subtract their moments from ˆ M t +1 to obtain the thermal population’s moments ˆ M (cid:48) t +1 for the next time step.
D.1.4 Components (e) (f) (g)
Each of these components by default use the respective components from the imitation solver.
The imitation solver implements the equations used in the OSIRIS PIC simulation, and is tested to match the simulation to machine precision error on any single timestep.
These components run efﬁciently because we now only advance the much smaller sub-population of non-thermal particles ( x, u ) t +1 .
Figure 7 illustrates the full pipeline of our method, concretized into the laser-plasma simulation.
[Figure]
[Figure]
Here we provide additional results for our experiments, as shown in Fig.
8 of energy spectrum, and
Tables 2 and 3.
Table 2 uses the same dataset split as Table 1, but focus on rollout starting from t = 1200 , where there is most injection and dynamical activity.
From the table, we see that the conclusion is similar as Table 1 in main text: LHPC achieves signiﬁcant error reduction for key quantities of interest: EM ﬁeld and MPIC (energetic particles).
To ablate different aspect of our model, we perform an additional experiment with a different train: test split, where the test dataset is in the future of the training (Table 3).
We take the 10 th trajectory with the highest laser intensity that consists of T = 8480 time steps, where the training time-range is [1000, 1900] and the testing (rollout) time-range is 1900 onwards.
We explore an ablation study on various strategies to use the additional data [1800, 1900].
Namely, we compare ﬁnetuning with simply retraining the model on [1000, 1900].
We also use the multi-step loss with the push-forward trick, reported by Brandstetter et al.
and other works to enhance the generalization of the model, and because it can accurately mimic the input distribution during rollout.
[Figure]
[Table]
plasma interactions for time range t = 1200 –1250.
Trajectory 5 ( a 0 = 10 ) is held out for testing, and the model is trained on the other 9 datasets.
We report performance on
Dataset 5 and rollout at t = 1200 , where there is the most injection and dynamical activity.
As shown, the model remains stable within this chaotic, unseen time range, showing generalization to an unseen trajectory.
F Additional details for architecture and training
We use a convolutional network (CNN) as the base architecture for modeling ﬂuid network g ﬂuid ,θ (Fig. 9) and injection network g inject ,ϕ .
For both networks, we use the same 6 layers, with the same kernel sizes of 1,1,3,7,3,1, with feature sizes of 64, allowing feature extraction and local information exchange among neighboring cells.
For each feature of the moments, we use a different feature head as shown in Fig.
We use Adam (Kingma & Ba, 2014) optimizer, with starting learning
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
We compare with the All-ﬂuid baseline, and report result from the best hyperparameter conﬁguration.
Each cell reports the relative L2 error for the EM-ﬁeld [ E t , B t ] , for M t .
rate of 10 − 3 .
We ﬁrst train with predicting single-step, and then train with predicting 3 steps into the future with push-forward trick (Brandstetter et al. 2022; which improves long-term prediction performance).
For single-step and multi-step training, we both train 500 epochs, with and cosine learning rate scheduling (Loshchilov & Hutter, 2016).