L EARNING EFFICIENT HYBRID PARTICLE - CONTINUUM REPRESENTATIONS OF NON - EQUILIBRIUM N- BODY SYSTEMS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 (non-thermal) 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 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.
I NTRODUCTION
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 fu-
[Figure]
[Table]
/sD6/
AFAQ5MJ</latexit>
[Table]
<latexit sha1_base64="bgHI6uDYE3O8vHJUirzY1snKUV4=">AAAB+HicbVDLSsNAFJ3UV62PRl26CRbBVUlE0WXBjbiqYB/QhjCZ3rRjJ5MwcyPW0i9x40IRt36KO//GaZuFth4YOJxzL3PPCVPBNbrut1VYWV1b3yhulra2d3bL9t5+UyeZYtBgiUhUO6QaBJfQQI4C2qkCGocCWuHwauq3HkBpnsg7HKXgx7QvecQZRSMFdrkfdBEecczlPTCcBHbFrbozOMvEy0mF5KgH9le3l7AsBolMUK07npuiP6YKORMwKXUzDSllQ9qHjqGSxqD98ezwiXNslJ4TJco8ic5M/b0xprHWozg0kzHFgV70puJ/XifD6NI3mdIMQbL5R1EmHEycaQtOjyuTVowMoUxxc6vDBlRRhqarkinBW4y8TJqnVe+sen57Vqnd5HUUySE5IifEIxekRq5JnTQIIxl5Jq/kzXqyXqx362M+WrDynQPyB9bnD4H1k60=</latexit>
[Table]
<latexit sha1_base64="s17+usCfIGl/vd8nXq5dUrL39zo=">AAAB9XicbVBNS8NAEJ3Ur1q/qh69BIvgqSRS0WPBi3iqYD+gjWWz2bRLN5uwO1FL6P/w4kERr/4Xb/4bt20O2vpg4PHeDDPz/ERwjY7zbRVWVtfWN4qbpa3tnd298v5BS8epoqxJYxGrjk80E1yyJnIUrJMoRiJfsLY/upr67QemNI/lHY4T5kVkIHnIKUEj3Q/6PWRPmIUi5cGkX644VWcGe5m4OalAjka//NULYppGTCIVROuu6yToZUQhp4JNSr1Us4TQERmwrqGSREx72ezqiX1ilMAOY2VKoj1Tf09kJNJ6HPmmMyI41IveVPzP66YYXnoZl0mKTNL5ojAVNsb2NAI74IpRFGNDCFXc3GrTIVGEogmqZEJwF19eJq2zqlurnt/WKvWbPI4iHMExnIILF1CHa2hAEygoeIZXeLMerRfr3fqYtxasfOYQ
[Table]
u
[Table]
u
[Table]
u
[Table]
> sion (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 far-from-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.
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.
Present-day 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 elec- tric ﬁ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 particlecontinuum 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.
R ELATED W ORKS
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 coarser spatial resolution (Um et al., 2020; Kochkov et al., 2021) or temporal intervals (Li et al., 2021; Vlachas et al., 2022), explicit forward method (Tang et al., 2020; Wu et al., 2022b), or via reduced representations (Sanchez-Gonzalez et al., 2020; Wu et al., 2022a).
Two important examples are LED (Vlachas et al., 2022) and solver-in-the-loop (Um et al., 2020).
The former optimizes the temporal evolution by employing 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. The latter uses a ﬂuid solver on a coarser spatial grid and a neural network to correct for the ﬂuid solver’s error.
However, almost all works employ a uniform representation, based on either all particles or all ﬂuid. This is not appropriate when trying to optimize far-fromequilibrium, N-body systems where a particle description is too computationally expensive for the near-thermal population (where most particles reside) and a ﬂuid description cannot accurately describe the complex distributions associated with the far-from-equilibrium population. Our method aims to address this long-standing challenge by learning an accurate and efﬁcient coupling between different representations involving different physical processes and scales, namely a hybrid particlecontinuum 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.
P ROBLEM 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, which approximates the ground truth evolution of S t :
[Figure] [Figure]
e l e c t r o n ( a r b .
un
i t )
[Figure]
2x [c/ ω 0 ]
[Figure]
1 u x [ m e c ]
[Figure]
[Figure]
# 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] [Figure]
[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 A 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
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 par- ticle distribution requires an inﬁnite number of ﬂuid moments (see Appendix A.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 B.
M ETHOD 4. M ETHOD 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.
4.1. D ATA GENERATION AND PREPARATION FROM FIRST - 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 A.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 B).
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.
4.1. A RCHITECTURE
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.
4. 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. 8 in Appendix C 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. 4.
[Equation]
Compute moments of new non-thermal particles to be injected g M NT ,ϕ is an instantiation of the g inject ,ϕ introduced in Sec.
[Equation]
1b)
[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]
1d)
[Equation]
(cid:2)
[Equation]
(( x, u [Equation] t
[Equation]
( E, B
[Equation] [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 C.1, and detailed architecture, learning procedure and objective is in Appendix E.
4.2. B ASELINES
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 thermal particles and non-thermal particles separately but each as a Gaussian.
Baseline (1) treats the entire system as a ﬂuid and advances the system by learning g as: [Equation] t + [Equation]
E, B
[Equation]
t +
[Equation]
g ( [Equation] E, B ) t , M t ) . The improved “bi-Gaussian” baseline (2) models the thermal population and the non-thermal particles each as a Gaussian, and advances the system by learning g as:
[Equation]
To evaluate the importance of the coupling network g M NT ,ϕ , we also compare our hybrid model with an ablation where we remove g M NT ,ϕ , while keeping all other parts the same (“LHPC (nocoupling)”).
For all models, we care especially about the error on the moments of the non-thermal subpopulation M PIC to understand the ability of the respective methods to correctly capture the most dynamic subpopulation of particles. This is important because in many of the relevant applications of far-from-equilibrium, N-body systems we are interested in the spectrum of accelerated particles.
E XPERIMENTS
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?
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 L as the error metric – the ratio between L2 norm of error over L2 norm of ground-truth values.
5. D ATA AND E XPERIMENT
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 A for details).
In this way, we test the generalization of the models to novel initial conditions and environment setup.
Each dataset consists of T = 1500 time steps and uses the separation strategy (details in Appendix B).
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 D with an ablation study on trajectory 10 with the most intense laser ( a 0 = 20 ).
5. R ESULTS
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. 5.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.
The LHPC (no-coupling) ablation shows that the coupling component is essential in ensuring good long-term accuracy, especially for Field and M PIC .
On [Table]
to the best-performing baseline.
Similarly, for M PIC LHPC achieves error reduction of 63.5%, 36.2%, 25.8% for 1-, 20- and 50-step predictions. Note that the performance is N/A for all-ﬂuid representation models for M PIC as they cannot model the non-thermal particles separately. For M prediction, our LHPC has comparable performance with the baselines, since both use similar architectures to model the thermal part of the distribution as ﬂuid, but the baselines also receive M as an input whereas LHPC does not.
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 Groundtruth (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.
5. V ISUALIZATION 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).
Figure shows 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 for components of the E-ﬁeld and the thermal ﬂ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. 6).
Note that the highly nontrivial 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.
[Figure]
D ISCUSSION 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 A.4).
Our method is generalizable towards ﬁnding more efﬁcient and accurate particle-continuum descriptions across different domains involving non-equilibrium N-body systems, and represent a key step towards the development of advanced data-driven algorithms for multi-scale modeling of nonlinear physical systems. Appendix A F IRST - 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.
A. P ARTICLE 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.
A. L ASER - 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
A. S IMULATION 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.
A. C OMPUTATIONAL 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 A 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.
B D ETAILS 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 .
Figure 7 illustrates an iteration in detail.
From the example velocity distribution function shown in Fig. 7(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.
7b).
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 h 
+}
and h − ≡ h ( v = v − ) =
Y h 0 (Fig. 7c).
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
# 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 (a) (b) (c) (d) (e) (f) ﬂuid frame according to v ﬂ (Fig.
7d; for non-relativistic velocities this amounts to v i
→ v i − v ﬂ ).
Calculate σ v to obtain the threshold (Fig. 7e).
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 7(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.
C F ULL P IPELINE C. C OMPONENT DETAILS C.1. C
OMPONENTS ( A )
AND ( B )
For both M or M NT , there’s 20 moment directional components that needs to be predicted - 10 for 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.
C.1. C OMPONENT ( 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 × 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.
C.1. C OMPONENT ( 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.
C.1. C OMPONENTS ( 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 8 illustrates the full pipeline of our method, concretized into the laser-plasma simulation.
[Figure]
D A DDITIONAL RESULTS FOR MULTI - SCALE LASER - PLASMA INTERACTION
Here we provide additional results for our experiments, as shown in 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 Tablein 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.
E A DDITIONAL 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
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
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. 9.
Training .
We use Adam (Kingma & Ba, 2014) optimizer, with starting learning rate of 10 − 3 .
The training consists of two stages.
In the ﬁrst stage, we train both g M,θ and g M NT ,ϕ with single-step loss:
[Equation]
In the second stage, we ﬁne-tune with predicting N = 4 steps into the future with push-forward trick (Brandstetter et al. 2022 that stops the gradient on the input.
The loss is given by:
[Equation]
For both losses, we have
[Equation]
Notice the hat notation ˆ · in the input arguments, denoting that they are autoregressive prediction of the LHPC at the previous time step, but the gradient is stopped (the “sg” notation).
Essentially, we rollout the pipeline of LHPC for multiple steps, and use it to provide a realistic input that contains the rollout error, requiring the model to not only predict well, but able to adapt to noise due to the rollout error.
We ﬁnd that this signiﬁcantly improve the long-term prediction performance.
For the coefﬁcient of the multi-step loss in Eq.
4, we set ( α 1 , α 2 , α 3 , α 4 ) =
(1 , 0 . 5 , 0 . 1 , 0 . 1) with decreasing weight for longer time steps.
This put more emphasis on the single-step loss to make the training more stable, and also have weight on longer-term future to improve long-term prediction.
For the loss function (cid:96) ( · , · ) in Eq. 3 and 4, we use
[Equation]
We ﬁnd that this achieves a better performance than the alternative of MSE loss (with exponent of 2) and MAE loss (with exponent of 1).
MSE will give very small gradient if the loss is small, and not able to encourage that the prediction to be exactly 0 in the vacuum.
On the other hand, MAE is harder to train and do not penalize more for larger errors.
Our choice of loss function (cid:96) ( · , · ) strikes a good balance between the two and enjoys the beneﬁt of both loss functions.
For single-step and multi-step training, we both train 500 epochs, with and cosine learning rate scheduling (Loshchilov & Hutter, 2016).