Cycle-Consistent and Uncertainty-Aware Neural Surrogates for Tokamak Edge Plasmas

Technical Summary: Cycle-Consistent and Uncertainty-Aware Neural Surrogates for Tokamak Edge Plasmas

Problem Statement
Accurate prediction of boundary and divertor plasma parameters (heat fluxes, temperatures, densities, and detachment onset) is critical for the safe operation of current and future tokamaks. However, high-fidelity edge-plasma simulations using the SOLPS-ITER code are computationally prohibitive for tasks requiring extensive parameter scans, optimization, or real-time control. Each SOLPS-ITER run solves a stiff, nonlinear fixed-point problem that can take several hours, exhibits high sensitivity to initialization, and suffers from significant failure rates (divergence or non-convergence) when scanning broad parameter sets or entering detached regimes. While data-driven surrogates have emerged to accelerate these tasks, existing models often lack physical consistency between forward and inverse mappings and provide no built-in mechanism to assess prediction reliability, particularly in undersampled or sharp-gradient regimes where errors are most likely.

Methodology
The authors propose a framework comprising two complementary neural surrogates and a robust database generation strategy:

1. Database Generation (k-d Tree Warm-Start): To address the high failure rate of SOLPS-ITER, the authors employ a branching k-d tree strategy. Instead of "cold-starting" simulations from generic initial conditions, each new simulation is warm-started from the nearest previously converged design point in the parameter space. This approach propagates reliable initial conditions outward from a single seed, increasing the database completion rate from ~70% (in cold-started ensembles) to over 95%. The final dataset comprises 760 converged runs of a deuterium plasma in a DIII-D lower-single-null configuration, sampled via Latin hypercube sampling over five control parameters: total power ($P_{tot}$), core fueling rate ($\Gamma_{core}$), deuterium puff rate ($\Gamma_{D2}$), cross-field particle diffusivity ($D_{\perp}$), and ion heat diffusivity ($\chi_i$).

2. Cycle-Consistent 2D Forward/Inverse Surrogate:

   • Forward Model: A conditional U-Net ($F_\theta$) maps the five scalar control parameters and a spatial binary mask to four 2D plasma-state fields: electron temperature ($T_e$), ion temperature ($T_i$), electron density ($n_e$), and parallel ion velocity ($u_a$). The control parameters are injected via Feature-wise Linear Modulation (FiLM). The network is trained on log-transformed variables to handle wide dynamic ranges.
   • Inverse Model: Rather than training a separate inverse network, the authors use an optimization-based approach. Given a target field, control parameters are recovered by minimizing the discrepancy between the target and the output of the frozen forward model via gradient descent.
   • Cycle Consistency: A cycle-consistency constraint links the forward and inverse processes. Recovered parameters are fed back through the forward model; the resulting reconstructed field must match the original target. This provides a self-supervised quality metric during inference that requires no ground-truth labels.
   • Training Regularization: During training, a learned pseudo-inverse ($G_\psi$) is trained jointly with the forward model under a composite loss that includes a cycle-consistency term. This reshapes the forward surrogate to ensure mutual consistency between the forward and inverse mappings.

3. Uncertainty-Aware 1D Profile Ensemble:

   • A separate "query-by-committee" ensemble of multilayer perceptrons (MLPs) predicts 1D profiles of $T_e$ and $n_e$ at the outboard midplane and divertor targets.
   • The ensemble variance serves as a proxy for predictive uncertainty. Points where committee disagreement exceeds a threshold are flagged as high-uncertainty regions, guiding active learning strategies to prioritize new, informative SOLPS-ITER simulations.

Key Results

• Forward Accuracy: The 2D U-Net achieves normalized root-mean-square errors (NRMSE) below 2.6% and Pearson correlations above 0.95 for all four plasma-state fields on held-out test cases. The model captures steep gradients near the separatrix and divertor targets.
• Cycle Consistency Gains: Adding cycle-consistency regularization ($\lambda_{cyc} = 0.5$) increases the average cyclical $R^2$ from 0.59 (forward-only baseline) to 0.99 without degrading forward prediction accuracy. Crucially, this regularization enables the recovery of the core fueling rate ($\Gamma_{core}$), which was poorly recovered by the forward-only pair.
• Inverse Parameter Recovery: The optimization-based inverse method recovers all five control parameters with Pearson correlations $r \ge 0.97$. The method successfully identifies parameters even in regimes where the forward mapping is locally insensitive, provided the cycle consistency is maintained.
• Database Robustness: The k-d tree warm-start strategy achieved a >95% completion rate for the 760-run database, a significant improvement over the ~30% failure rate observed in comparable cold-started ensembles.
• Inference Speed: The model (approx. 4.3 million parameters) generates full 2D predictions in milliseconds (sub-millisecond on GPU), representing a speedup of five to six orders of magnitude over original SOLPS-ITER runs.
• Uncertainty Quantification: The 1D ensemble accurately reproduces profiles ($R^2 > 0.89$) and correctly identifies high-uncertainty regions. The inner divertor target, which spans both attached and detached states, showed the highest uncertainty (~13% of points flagged), aligning with known physical challenges.

Significance and Claims
The paper claims that this framework provides a practical solution for real-time plasma control, rapid scenario optimization, and digital-twin applications by replacing slow, brittle physics solvers with fast, reliable neural surrogates. Key contributions include:

• Self-Supervised Reliability: The cycle-consistency mechanism offers a built-in quality check for inverse predictions without requiring ground-truth labels at inference time, addressing a major limitation of standard surrogate models.
• Robust Data Generation: The k-d tree warm-start strategy is presented as an essential infrastructure for building large, uniformly sampled training sets for stiff nonlinear systems, directly mitigating the convergence failures that typically limit edge-plasma databases.
• Active Learning Integration: The uncertainty-aware ensemble provides a mechanism to identify "blind spots" in the parameter space (specifically the inner divertor detachment regime), nominating the most informative new simulations to improve the model iteratively.
• Automated Transport Inference: The differentiable inverse procedure offers a route to automatically infer transport coefficients ($D_{\perp}, \chi_i$) and core fueling rates from diagnostic measurements, potentially replacing manual, iterative fitting workflows.

The authors maintain a modest tone regarding limitations, noting that the current model is restricted to a single DIII-D configuration with deuterium-only fueling and uniform transport coefficients. They state that validation against experimental data and extension to other machines and spatially varying transport coefficients are necessary future steps.