Physics-integrated Latent Space Dynamics Learning for Stiff Collisional-radiative Models

Xuping Xie² , Qi Tang³, Xianzhu Tang¹

​

​

1. T-5 Applied Mathematics and Plasma Physics, Los Alamos National Laboratory

2. Department of Mathematics and Statistics, Old Dominion University

​

NERSC User Group (NUG) Annual Meeting

October 8, 205

​

LA-UR-24-26289

​

​

3. School of Computation Science and Engineering, Georgia Tech

​

1

10/6/25

Acknowledgement

​

• The work is supported by the U.S. Department of Energy (DOE) Office of Fusion Energy Science (FES) SciML/AI program.

​

• This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility using NERSC award FES-ERCAP002815

​

2

10/6/25

Collisional radiative (CR) models

​

​

• Collisional-radiative (CR) models describe the atomic

processes in a plasma by tracking the population dens-

ity in ground and excited states for each charge state

of the atom/ion.

• Disruptions can occur when the plasma becomes

unstable.

• These instabilities can lead to a sudden loss of confinement, causing the plasma to touch the reactor walls
• To prevent severe damage, we need to accurately predict the ion population distribution and radiative power cooling rate before a disruption occurs

Tokamak operated by ITER

​

3

10/6/25

Collisional radiative (CR) models

​

• Autoionization. It occurs when an electron in a high-energy excited state undergoes a decay that leads to the emission of a photon and potentially ionizes another electron in an excited state (ES).

​

​

​

​

• Electron capture

4

10/6/25

Collisional radiative (CR) models

​

• The high-fidelity CR model is a high-dimensional, stiff ODE system where dim(N) is O(10⁶) or larger:

​

​

• A high-fidelity model is computationally intractable in an integrated plasma simulation, e.g., a high-fidelity CR model needs to be solved on a fast time scale at each grid point in an MHD simulation at every time step.

​

• A low-fidelity corona model exists, but its fidelity is not enough for the need of the high-fidelity fusion simulations.

5

10/6/25

CR dynamics and ML models

• Consider lithium as an example with dim(N) = 94, and there exist four charge states:

​

​

• It is stiff and requires time steps following an exponential growth.
• The scales in both time steps and densities span multiple orders

of magnitudes.

​

Motivations of ML surrogate models:

​

• A low-fidelity CR model can be used to enhance data driven ML

models.

• Coupling the CR surrogate model to dynamical plasma simulation.
• Predicting long-term dynamics and significantly accelerate conventional computational models.

​

​

Figure: a typical CR dynamics

6

10/6/25

Grey-box Latent Space

​

White space:

• Refers to the part of the latent space composed of physical quantities that are directly interpretable and have well-established meanings within the CR model.

​

• CR Coronal Equilibrium Model, i.e., Charge state distribution

​

• By incorporating these known quantities, we ensure that essential physical properties are retained in the reduced model

​

​

​

7

10/6/25

Grey-box Latent Space

​

Black space:

• It represents the portion of the latent space learned by the autoencoder without direct physical interpretation.
• The low-dimension representations capture complex patterns and relationships not easily described by known physical variables.
• It functions as a black-box.

​

Grey space: refers to the combined latent space that integrates both white

space (interpretable physical quantities) and black space (learned abstract features).

​

​

​

8

10/6/25

Physics-integrated Deep Reduced Physics Model (DeepRPM)

​

Full-order model approximation consists of two steps:

1. Discover a latent space 2. Discover a latent dynamics

​

​

White space

Autoencoder

Black space

Physics

9

10/6/25

DeepRPM training

​

1. Autoencoder loss (the radiative loss rate is �incorporated as a weak constraint):

​

​

​

• Flow map loss (the total mass conservation is �incorporated as a weak constraint):

​

​

Figure: physics-assisted autoencoder architecture

10

10/6/25

DeepRPM long-term dynamics prediction

After trained, the DeepRPM model is used to predict long-term CR dynamics:

Input – the initial state vector and parameters

Output – the full state vector and RL at any time

11

10/6/25

Key component 1: resampling in time

​

​

​

​

​

​

​

Two advantages for ML training:

• reduce the total number of training data, 14.8 million pairs.
• move the target flow map away from a near-identity map for small Δt

12

10/6/25

Key component 2: proper transformations for both time steps and densities

• Temporal transformation (convert Δt to a range of [0, 1]):

​

​

​

​

​

• Density transformation (convert N to a range of [0, 1]):

13

10/6/25

Key component 3: neural network architecture search

​

We conduct a grid search for layers of 2 to 7 and hidden units of 16 to 512.

​

Our team is actively working on replacing this simple search with a ML-assisted advanced search tool of DeepHyper for hyperparameter optimization (effort led by Romit Maulik).

14

10/6/25

Dynamics prediction with different parameters

​

15

10/6/25

Dynamics prediction with different parameters

​

An excellent agreement is observed for a learned model predicting dynamics of different parameters!

16

10/6/25

Radiative loss rate prediction with different parameters

A physical quantity of interest can be predicted accurately, thanks to the built-in physics-assisted architecture.

17

10/6/25

Impact of training data size

• Fix testing Te values, increase the training dataset with more Te values.

Testing Te = [15,45,75,95].

18

10/6/25

Impact of training data size

Dynamics prediction using the model trained with Te=[5,25,35,55,65,85] (dark dots)

19

10/6/25

Impact of training data size

Dynamics prediction using the model trained with full Te values (blue dots)

20

10/6/25

Conclusion and future work

• A physics-integrated ML-based surrogate model is proposed for the collisional radiative model for plasma simulation.
• Several key components for a successful data-driven training have been identified
• We will explore other dynamics prediction architectures.

​

Reference:�[1] X. Xie, Q. Tang, and X.-Z. Tang. Physics-assisted latent space dynamics learning for stiff collisional-radiative models, https://arxiv.org/abs/2409.05893

[2] N. Garland, R. Maulik, Q. Tang, X. Tang, and P. Balaprakash. Efficient data acquisition and training of collisional-radiative model artificial neural network surrogates through adaptive parameter space sampling.

[3] N. Garland, H. Chung, C. Fontes, M. Zammit, J. Colgan, X. Tang. Impact of a minority relativistic electron tail interacting with a thermal plasma containing high-atomic-number impurities. Physics of Plasmas, 27(4), 2020.

​

​

21

10/6/25