Robust Confinement State Classification with Uncertainty Quantification through Ensembled Data-Driven Methods

TL;DR

The paper tackles automatic confinement-state labeling in tokamaks with quantified uncertainty and robustness to missing signals. It introduces a hierarchical ensemble that combines dynamic FNOLSTM (Fourier Neural Operator + LSTM) and static GBDT models across diverse feature sets, optimized on 302 labeled TCV discharges. The approach yields high prediction accuracy with calibrated uncertainties (Cohen's κ ≈ 0.93–0.95; ECE ≈ 0.02–0.03) and demonstrates robustness to signal loss and out-of-distribution regimes, while highlighting moderate challenges in pinpointing exact transition times. The work provides a public dataset and a practical, scalable framework suitable for large-scale confinement-state analyses and potential real-time control implementations.

Abstract

Maximizing fusion performance in tokamaks relies on high energy confinement, often achieved through distinct operating regimes. The automated labeling of these confinement states is crucial to enable large-scale analyses or for real-time control applications. While this task becomes difficult to automate near state transitions or in marginal scenarios, much success has been achieved with data-driven models. However, these methods generally provide predictions as point estimates, and cannot adequately deal with missing and/or broken input signals. To enable wide-range applicability, we develop methods for confinement state classification with uncertainty quantification and model robustness. We focus on off-line analysis for TCV discharges, distinguishing L-mode, H-mode, and an in-between dithering phase (D). We propose ensembling data-driven methods on two axes: model formulations and feature sets. The former considers a dynamic formulation based on a recurrent Fourier Neural Operator-architecture and a static formulation based on gradient-boosted decision trees. These models are trained using multiple feature groupings categorized by diagnostic system or physical quantity. A dataset of 302 TCV discharges is fully labeled, and will be publicly released. We evaluate our method quantitatively using Cohen's kappa coefficient for predictive performance and the Expected Calibration Error for the uncertainty calibration. Furthermore, we discuss performance using a variety of common and alternative scenarios, the performance of individual components, out-of-distribution performance, cases of broken or missing signals, and evaluate conditionally-averaged behavior around different state transitions. Overall, the proposed method can distinguish L, D and H-mode with high performance, can cope with missing or broken signals, and provides meaningful uncertainty estimates.

Figures (20)

• Figure 1: Overview of the total time and transitions present in the dataset. The top shows the cumulative time spent in the different states, whereas the bottom depicts the total number of state transitions. The left depicts the total number of transitions, whereas the right excludes unstable or transient transitions, which we define as those where the plasma changes state within 10ms before or after a transition.
• Figure 2: The distribution of the dates of the discharges present in the dataset. It includes TCV plasmas from several decades, with the majority from the last 10 years.
• Figure 3: Distributions of key plasma parameters in the dataset, stacked for L, D and H-mode plasmas. We plot average values for phases of 20ms--around the TCV energy confinement time--to exclude fast transient measurements.
• Figure 4: An overview of the ensemble structure and prediction procedure. The structure is defined in a top-down fashion. First, we ensemble over different models and feature sets (level 1). Each (model + feature set) combination consists of a small ensemble of models fit on different folds of the data (level 2). At the bottom, we consider individual models (level 3). Prediction is done in a bottom-up fashion. Individual models map the input signals to a prediction of being in L, D or H (level 3). These predictions are combined and re-scaled on the level of a (model + feature set) (level 2). The resulting predictions are combined weighted by their confidences and a constant that is computed using a (model + feature set)'s average classification performance, giving the full ensemble prediction (level 1).
• Figure 5: A simplified illustration of the two modeling approaches, a static formulation operating on timeslices of signals (top) and a dynamic formulation operating on sequences of signals (bottom). Static (GBDT): The static formulation is implemented with gradient boosted decision trees. A collection of small decision trees is fit in a sequential manner: each tree is fit using the residual errors of the previously fit trees. These trees operate on a vector of signal data corresponding to a single timeslice of the discharge. The final prediction is computed using a weighted combination of all individual trees. Dynamic (FNOLSTM): The dynamic formulation is implemented through a neural network extracting features on both small and large timescales. The input consists of a matrix corresponding to a small time window of signal data. This input is transformed with the FNO, a convolution-like operator acting primarily in the frequency domain. This computed abstract representation is processed in a recurrent manner over the entire duration of the shot with an LSTM. The final prediction is then computed using this local and global representation of the input signals.