Physics-Informed Latent Space Dynamics Identification for Time-Dependent NLTE Atomic Kinetics

This paper introduces a physics-informed Latent Space Dynamics Identification (pLaSDI) framework that overcomes the computational bottlenecks of time-dependent non-local thermodynamic equilibrium (NLTE) atomic kinetics by learning explicit reduced governing equations, achieving massive speedups while ensuring physical stability and accuracy in predicting plasma charge-state evolution for EUV lithography applications.

The Gist

The Big Problem: The "Slow Motion" of Atoms

Imagine you are trying to predict the weather for a city. To do this accurately, you need to track millions of tiny air molecules, how they bump into each other, and how they react to heat. This is incredibly hard and slow.

In the world of physics, specifically when designing things like EUV lithography (the technology used to print computer chips), scientists face a similar problem. They need to simulate how plasmas (super-hot, ionized gas) behave. Inside this plasma, atoms are constantly changing their "charge" (losing or gaining electrons) in a chaotic dance called Non-Local Thermodynamic Equilibrium (NLTE).

To get the answer right, traditional computer codes have to solve millions of complex equations for every tiny fraction of a second. It's like trying to calculate the path of every single raindrop in a storm to predict if you'll get wet. It takes hours or even days to run one simulation, which makes it impossible to design new chips or fusion reactors quickly.

The Old Solution: The "Guess-and-Check" Machine

Scientists tried to use Artificial Intelligence (AI) to speed this up. They treated the problem like a simple "input-output" machine:

• Input: Give the AI the temperature and density.
• Output: The AI guesses the state of the atoms.

The Flaw: This works well for the specific scenarios the AI was trained on, but it's a "black box." If you ask the AI to predict something slightly different (like a temperature it hasn't seen before), it might hallucinate. Worse, if you let the AI run the simulation for a long time, it might start to drift, predicting that atoms have negative mass or that the plasma explodes. It lacks common sense and physics rules.

The New Solution: The "Physics-Guided GPS" (pLaSDI)

The authors of this paper created a new AI framework called pLaSDI (Physics-Informed Latent Space Dynamics Identification). Think of this not as a guesser, but as a GPS with a built-in rulebook.

Here is how it works, broken down into three simple steps:

1. The Compression (The Suitcase)

The original simulation has 1,583 different "variables" (like 1,583 different types of atoms and their energy levels). That's too heavy to carry.

• The Analogy: Imagine trying to pack a whole library into a single suitcase. You can't fit every book, so you write a summary of the story.
• The Tech: The AI uses a "neural network" to squash those 1,583 variables down into just 3 numbers (a "latent space"). It's like summarizing the entire library into three key themes.

2. The Map (The Rules of the Road)

Instead of just guessing the next step, the AI learns the rules of the road for these 3 numbers.

• The Analogy: A normal GPS might just tell you "turn left." This new GPS knows the laws of physics. It knows that if you drive too fast, you might crash, so it has built-in brakes.
• The Tech: The team taught the AI to follow a specific mathematical equation (an ODE) that describes how the plasma changes over time. Crucially, they added three "safety rules" (loss functions) to the training:
  1. Stability: "Don't go off the road." (Mathematically, this ensures the numbers don't explode to infinity).
  2. Conservation: "Don't create or destroy matter." (The total number of atoms must stay the same).
  3. Steady State: "If the road is flat, stop." (If the temperature stops changing, the plasma should settle into a calm, predictable state, not keep wobbling).

3. The Expansion (Unpacking the Suitcase)

Once the AI runs the simulation using just those 3 simple numbers and the safety rules, it "unpacks" the suitcase back into the full 1,583 variables to show the scientists the detailed result.

Why This is a Game-Changer

The results are impressive:

• Speed: The new model is 50,000 to 100,000 times faster than the old method. A calculation that took hours now takes a fraction of a second.
• Accuracy: It predicts the behavior of the plasma with less than 2% error.
• Reliability: Because of the "safety rules," if you ask the AI to simulate a scenario it has never seen before (extrapolation), it doesn't crash or hallucinate. It stays stable and converges to the correct physical answer.

The Takeaway

Think of the old method as a student who memorized the answers to a specific test. If you ask a different question, they fail.

This new pLaSDI method is like a student who memorized the laws of physics and learned how to apply them. They can solve the specific test, but they can also solve any test you throw at them, even ones they've never seen, because they understand the underlying rules.

This breakthrough means scientists can now run thousands of simulations in the time it used to take to run one, accelerating the design of better computer chips and bringing us closer to clean fusion energy.

Technical Summary

1. Problem Statement

Non-Local Thermodynamic Equilibrium (NLTE) atomic kinetics is a critical component in radiation-hydrodynamics simulations, particularly for applications like Extreme Ultraviolet (EUV) lithography and fusion energy. However, it presents a severe computational bottleneck:

• Stiffness: The governing rate equations involve a wide range of transition rates (from $10^{-50}$ to $1$), creating a stiff Ordinary Differential Equation (ODE) system that requires expensive implicit integration.
• Computational Cost: In hydrodynamic simulations, NLTE calculations must be performed at every time step for every spatial cell, often dominating the total simulation cost.
• Limitations of Existing ML Surrogates: Most current machine learning (ML) approaches treat NLTE as a static input-output mapping or a "black-box" regression problem. These methods often fail to:
  • Maintain physical consistency (e.g., population conservation, correct steady states) during long-time integration.
  • Extrapolate robustly outside the training regime.
  • Ensure dynamical stability (e.g., preventing unphysical divergence).

The authors aim to develop a fast, stable, and physically reliable surrogate model that captures the time-dependent evolution of atomic populations, rather than just predicting static snapshots.

2. Methodology: Physics-Informed Latent Space Dynamics Identification (pLaSDI)

The authors propose a framework called pLaSDI, which combines dimensionality reduction with explicit physics-constrained dynamics identification.

A. Data Representation and Preprocessing

• Input: High-dimensional atomic population vectors ($\mathbf{n} \in \mathbb{R}^N$, where $N \approx 1,583$ for the tin super-configuration model).
• Scaling: To handle the extreme dynamic range of populations ($10^{-50}$ to $1$), a **$w$-scaled representation** is used. This involves a logarithmic transformation followed by min-max normalization to map values to $[0, 1]$, ensuring small populations contribute meaningfully to the loss function.
• Control Variables: Time-dependent electron temperature ($T$) and density ($\rho$) are treated as control inputs.

B. Architecture

The framework consists of three main components:

1. Nonlinear Autoencoder: Compresses the high-dimensional population vector $\mathbf{n}$ into a low-dimensional latent space $\mathbf{z} \in \mathbb{R}^{N_z}$ (where $N_z \ll N$).
   • Encoder: $\mathcal{G}_{enc}: \mathbf{n} \to \mathbf{z}$
   • Decoder: $\mathcal{G}_{dec}: \mathbf{z} \to \mathbf{n}$
2. Latent Dynamics Identification (DMDc): Instead of learning a black-box map, the temporal evolution in the latent space is modeled by an explicit linear ODE with control inputs:
   $$\frac{d\mathbf{z}}{dt} = \mathbf{a} + \mathbf{A}\mathbf{z} + \mathbf{B}\mathbf{u}(t)$$
   where $\mathbf{u}(t)$ represents the scaled $T$ and $\rho$. The coefficients $\{\mathbf{a}, \mathbf{A}, \mathbf{B}\}$ are identified globally across all training trajectories using Dynamic Mode Decomposition with Control (DMDc).
3. Physics-Informed Constraints: To ensure physical reliability, specific loss terms are added to the training objective:
   • Hurwitz Stability Loss: Enforces that the eigenvalues of matrix $\mathbf{A}$ have strictly negative real parts. This guarantees that the latent dynamics converge to a steady state rather than diverging during long-time integration.
   • Steady-State Consistency Loss: Explicitly constrains the analytically computed equilibrium of the latent ODE ($\mathbf{z}^* = -\mathbf{A}^{-1}(\mathbf{a} + \mathbf{B}\mathbf{u}^*)$) to match the reference NLTE steady-state populations under fixed conditions.
   • Macroscopic Regularization: Additional penalties enforce population conservation, charge-state distribution (CSD) consistency, and mean charge state ($\bar{q}$) accuracy.

C. Training Objective

The total loss function combines reconstruction error, DMDc fitting error, and the physics-informed constraints:
$$\mathcal{L}_{total} = \mathcal{L}_{AE} + \mathcal{L}_{DMDc} + \mathcal{L}_{stability} + \mathcal{L}_{steady}$$

3. Key Contributions

1. Explicit Latent Dynamics: Unlike standard neural ODEs or black-box surrogates, pLaSDI identifies an explicit, interpretable linear operator in the latent space, allowing for analytical stability analysis.
2. Physics-Informed Design: The introduction of Hurwitz stability and steady-state consistency constraints directly addresses the failure modes of data-driven models (divergence and incorrect asymptotic behavior) in stiff kinetic systems.
3. Control-Driven Formulation: The use of DMDc allows the model to handle time-varying plasma conditions ($T(t), \rho(t)$) as continuous control inputs, rather than requiring interpolation of fixed-parameter coefficients.
4. Robust Extrapolation: The model demonstrates stability and physical correctness even when extrapolating beyond the training time horizon or to fixed plasma conditions not explicitly seen during training.

4. Results

The model was tested on Tin (Sn) plasma data relevant to EUV lithography, generated via a two-stage workflow (FLASH for hydrodynamics $\to$ SCFLY for NLTE kinetics).

• Dimensionality Reduction: Compressed a 1,583-dimensional state space into 3 latent variables (a reduction factor of ~521).
• Accuracy:
  • Mean Charge State ($\bar{q}$): Achieved relative errors of ~1.92% (training) and ~1.98% (validation).
  • Population Reconstruction: Root-mean-square error (RMSE) of $2.9 \times 10^{-3}$ for fractional populations and maintained accuracy in the low-population regime via $w$-scaling.
  • Steady State: The model correctly converged to the reference steady state across a wide range of temperatures (1–50 eV) and densities, with an RMSE of $5.7 \times 10^{-4}$ for populations.
• Stability:
  • Without the Hurwitz constraint, the model diverged rapidly even within the training window.
  • With the constraint, the model remained bounded and converged to the correct equilibrium.
• Computational Efficiency:
  • The surrogate evaluates a full 4 ns time-dependent trajectory in < 0.038 seconds.
  • This represents a speedup of $5 \times 10^4$ to $10^5$ compared to the conventional SCFLY solver (which takes tens of minutes to hours).

5. Significance and Conclusion

This work establishes a new paradigm for constructing surrogates for stiff kinetic systems. The authors demonstrate that data fitting alone is insufficient for long-time reliability; instead, the structural design of the latent dynamics must be constrained by physical principles (stability and asymptotic behavior).

• Impact: The pLaSDI framework enables scalable, real-time NLTE modeling, making large-scale ensemble simulations, uncertainty quantification, and design optimization for EUV lithography and fusion energy feasible.
• Future Work: The authors plan to extend the framework to include radiation transport effects, more detailed atomic structure models (beyond super-configurations), and integration into full radiation-hydrodynamic simulation codes.

In summary, pLaSDI successfully bridges the gap between high-fidelity physics and machine learning speed by embedding physical laws directly into the mathematical structure of the learned dynamics.