Machine learning prediction of plasma behavior from discharge configurations on WEST

This study presents a transformer-based machine learning model trained on 550 WEST tokamak discharges that rapidly and accurately predicts key global plasma parameters from pre-discharge configurations, offering a computationally efficient alternative to physics-based codes for discharge planning and real-time control.

The Gist

Imagine you are trying to bake the perfect cake. In the world of fusion energy, the "cake" is a super-hot ball of plasma (a state of matter like the sun) held inside a giant magnetic donut called a tokamak.

For decades, scientists have tried to predict how this plasma will behave before they even turn the machine on. They used to rely on physics-based simulations. Think of these like trying to bake a cake by solving complex math equations for every single grain of sugar and every molecule of flour. It's incredibly accurate, but it takes hours (or even days) to run the calculation. By the time you get the answer, the oven is already cold, and you can't use it to adjust the recipe in real-time.

This paper introduces a new, faster way to bake: The "AI Chef."

The Problem: Too Slow for Real-Time

The researchers at the WEST tokamak (a high-tech fusion experiment in France) needed a way to predict the plasma's behavior instantly. They needed to know: If I set the magnetic coils to this position and turn on this much heating power, will the plasma stay stable? Will it hold enough energy?

The old physics models were too slow for this. They were like trying to calculate the flight path of a rocket using a slide rule while the rocket is already launching.

The Solution: A Machine Learning "Cheat Sheet"

The team built a Machine Learning model (specifically, a "Transformer," the same type of AI architecture that powers tools like ChatGPT). Instead of solving physics equations from scratch every time, this AI learned by looking at 550 past "baking sessions" (discharges) from the WEST machine.

Think of it like this:

• The Old Way: Calculating the weather by simulating every air molecule in the atmosphere.
• The New Way: Looking at a massive database of past weather patterns and saying, "Hey, when the wind blows from the north and the humidity is high, it usually rains. I'll bet it rains today."

How the AI Chef Works

The model takes in 19 "ingredients" (inputs) that are decided before the experiment starts:

• How much magnetic power to use (like setting the oven temperature).
• How much heating power to apply (like turning on the stove).
• How much fuel (plasma density) to inject.

It then predicts 6 key outcomes (the "taste" of the cake):

1. Stability: Will the plasma hold together or explode?
2. Efficiency: How much energy is stored inside?
3. Safety: Are we pushing the limits too hard?

The Results: Fast and Accurate

The results are impressive:

• Speed: The AI makes a prediction in 0.1 seconds. That's faster than you can blink. The old physics models took minutes or hours.
• Accuracy: It got the prediction right about 94% of the time (a score of 0.94).
• Reliability: It works so well that scientists can now use it to run thousands of "what-if" scenarios instantly to find the perfect settings for a new experiment.

The "Glitch" in the Recipe

The paper admits the AI isn't perfect at predicting two specific things: the exact shape of the magnetic field in the very center and the very edge of the plasma (called $q_0$ and $q_{95}$).

Why? Imagine trying to guess the exact texture of a cake's center just by looking at the ingredients list. You know the flour and eggs, but you don't know exactly how they mixed inside the batter. Similarly, the AI knows the inputs, but it can't "see" the invisible internal currents of the plasma. Because of this, it sometimes guesses a little low on these specific numbers.

Why This Matters

This isn't just a cool trick; it's a game-changer for the future of clean energy.

• Scenario Planning: Scientists can now test thousands of different settings in the time it used to take to test one.
• Real-Time Control: In the future, this AI could act like an autopilot for the fusion reactor, adjusting the knobs instantly to keep the plasma stable, preventing it from crashing.

In short: The researchers built a super-fast "plasma oracle." Instead of doing heavy math to guess what will happen, the AI looks at what happened before and tells us exactly what to expect, helping us get closer to building a star on Earth that provides limitless clean energy.

Technical Summary

1. Problem Statement

Accurate prediction of plasma behavior based on pre-discharge configurations is critical for the safe and efficient operation of tokamak experiments.

• Limitations of Physics-Based Models: Traditional "Integrated Modeling" codes (e.g., TRANSP, CRONOS, RAPTOR) rely on first-principles physics but suffer from high computational costs, making them unsuitable for fast scenario design, real-time control optimization, or rapid parameter sweeps. They also rely on empirical coefficients and reduced physics models that may lack shot-by-shot fidelity.
• Limitations of Existing Data-Driven Models: Previous machine learning approaches often rely on signals that evolve during the discharge (e.g., gas puffing data, real-time plasma state), which limits their ability to predict outcomes based solely on the initial discharge schedule. Furthermore, many existing models struggle to generalize across different operational regimes.

Goal: Develop a fast, data-driven surrogate model that predicts key global plasma parameters using only signals defined before the discharge begins (inputs from the Discharge Control System), enabling rapid scenario evaluation and planning.

2. Methodology

Dataset

• Source: 550 stable, reproducible discharges from the WEST tokamak (campaigns #57381 to #60286).
• Preprocessing:
  • Signals were resampled at 1 kHz.
  • Discharges shorter than 2 seconds (ramp-up phase) were excluded.
  • A Simple Moving Average (SMA) filter (window size 11) was applied to suppress noise without introducing temporal lag.
  • Data was segmented using a sliding window (length 1024, step 512) to augment the dataset, resulting in ~24,000 sliced windows. Overlapping segments were averaged during inference to improve stability.
• Split: Randomly divided into Training (60%), Validation (20%), and Test (20%) sets.

Input and Output Signals

• Inputs (19 signals): Defined strictly by the discharge schedule or feedback-regulated parameters.
  • Lower Hybrid Wave (LHW) power and phase (4 channels).
  • Ion Cyclotron Resonance Heating (ICRH) power (3 channels).
  • Reference plasma current ($I_p$).
  • Real-time line-averaged electron density ($n_e$) as a proxy for fueling/wall conditions.
  • Poloidal Field (PF) coil currents (10 channels).
• Outputs (6 global parameters):
  • Normalized beta ($\beta_n$), Toroidal beta ($\beta_t$), Poloidal beta ($\beta_p$).
  • Plasma stored energy ($W_{mhd}$).
  • Safety factor at the magnetic axis ($q_0$) and at the 95% flux surface ($q_{95}$).

Model Architecture

• Core Architecture: A Transformer-based model was selected after benchmarking against Multilayer Perceptrons (MLP), Long Short-Term Memory (LSTM) networks, and standard Transformer Encoder/Decoder models.
• Rationale: Transformers utilize self-attention mechanisms to capture long-range temporal dependencies and complex cross-channel interactions inherent in multi-signal, long-duration discharge sequences.
• Hyperparameters: Optimized via Bayesian optimization. Key settings included 8 attention heads, 2 layers, a learning rate of $10^{-3}$, dropout of 0.1, and batch size of 16.
• Training Environment: PyTorch on Red Hat Enterprise Linux 8, utilizing four NVIDIA A100 GPUs.

3. Key Contributions

1. Pre-Discharge Prediction Framework: The model successfully predicts plasma responses using only pre-discharge inputs, bypassing the need for real-time state measurements (like gas puffing data) that are unavailable during the planning phase.
2. Transformer Application in Fusion: Demonstrates the superiority of Transformer architectures over traditional time-series models (LSTM) and MLPs for tokamak discharge prediction, achieving the lowest Mean Squared Error (MSE) among tested models.
3. Statistical Validation: Conducted rigorous Granger causality and Pearson correlation analyses to validate that the selected 19 input signals have statistically significant causal and linear relationships with the target outputs.
4. High-Speed Inference: Achieved inference times of approximately 0.1 seconds, enabling real-time or near-real-time scenario sweeps and optimization tasks that are computationally prohibitive with physics-based codes.

4. Results

• Overall Performance:
  • Average MSE: 0.026.
  • Average $R^2$: 0.94.
  • Inference Time: ~0.1 seconds.
• Parameter-Specific Performance:
  • High Accuracy: The model performed exceptionally well for $\beta_n$, $\beta_t$, $\beta_p$, and $W_{mhd}$, with $R^2$ values generally exceeding 0.89.
  • Lower Accuracy for Safety Factors: Predictions for $q_0$ and $q_{95}$ showed lower accuracy ($R^2 \approx 0.65 - 0.78$).
    • Reasoning: These parameters depend heavily on the internal radial distribution of plasma current density ($j(r)$), which is not uniquely determined by the pre-discharge inputs alone. The model struggles to reconstruct these internal profiles without kinetic measurements.
• Case Studies:
  • The model successfully reproduced discharge pulses (ramp-up, flat-top, ramp-down) for Ohmic, LHW-heated, and ICRH-heated scenarios.
  • Failure Modes: Poor performance was observed in rare cases involving significant, unintended fluctuations in magnetic coil currents (e.g., discharge #57821) or extremely rare dual-heating scenarios not well-represented in the training data.

5. Significance and Future Directions

• Operational Impact: This work provides a fast surrogate model for discharge planning and scenario evaluation. It allows operators to rapidly test thousands of configuration variations to optimize for stability and performance before running actual experiments.
• Complementary Role: The model is not intended to replace physics-based simulators but to complement them. While physics codes answer "what should happen based on theory," this data-driven model answers "what actually happens" based on historical machine behavior.
• Future Work:
  • Generalization: Addressing the challenge of transferring the model to different tokamaks or after major hardware modifications.
  • Physics-Informed ML: Integrating physical constraints or knowledge into the training process to improve robustness and extrapolation capabilities.
  • Data Expansion: Collecting more diverse data, including rare operational scenarios and data from multiple devices, to improve generalization.
  • Profile Prediction: Extending the model to predict 1D profiles (density, temperature) and additional parameters like loop voltage ($V_{loop}$) and internal inductance ($l_i$).

In conclusion, this study establishes a robust, high-speed, transformer-based framework for predicting tokamak plasma behavior from discharge configurations, significantly enhancing the efficiency of scenario design and control optimization for the WEST tokamak and future fusion devices.