Idea: Adding confidence intervals to TORAX's transport model

[nytdevansh]

Hey team — I’ve been thinking about uncertainty in TORAX transport modeling and wanted to float an idea.

Observation:
QLKNN produces point estimates devoid of out-of-distribution signal and epistemic uncertainty, but it offers quick and precise flux predictions within its QuaLiKiz training envelope.

Why this matters:
We are unable to determine quantitatively whether a prediction is extrapolating or lies within the training manifold when investigating novel operating regimes, such as when scenario optimisation for SPARC or when using TORAX as an RL training environment.

This creates two risks:
• Overconfidence in scenario optimization (design margins unclear)
• RL agents exploiting surrogate model bias

Potential direction:
Augment (not replace) QLKNN with an uncertainty-aware surrogate in shadow mode — potentially using neural-operator-style architectures or ensemble-based methods — to provide:
• Calibrated predictive intervals
• OOD detection scores
• Regime confidence indicators

Goal:
Know when to trust fast surrogate predictions and when to trigger higher-fidelity validation (e.g., gyrokinetics).

Realistic scope:
Start with a single transport channel, one tokamak dataset, and evaluate:
• Calibration quality
• OOD detection performance
• Long-time solver stability

3-month prototype to test feasibility.

Worth exploring?

[nytdevansh]

Shadow mode could be toggled via runtime_params.enable_transport_shadow for validation runs.

class HybridTransportModel:
"""Run QLKNN (primary) + Neural Operator (shadow) in parallel."""

  def __init__(self, primary_model, shadow_model):
      self.primary = primary_model  # QLKNN_7_11
      self.shadow = shadow_model    # FNO/DeepONet
      
  def predict_transport(self, core_profiles):
      # Primary: point-wise local predictions (current TORAX)
      chi_primary = self.primary(local_params_at_each_radius)
      
      # Shadow: global profile-to-profile prediction
      input_profiles = jnp.stack([Te_profile, ne_profile, Ti_profile])
      chi_shadow, uncertainty = self.shadow(input_profiles)
      
      # Detect disagreement
      disagreement = jnp.abs(chi_primary - chi_shadow)
      if jnp.any(disagreement > 3 * uncertainty):
          log.warning("Non-local effects detected at r/a = ...")
      
      return chi_primary  # Still use QLKNN for actual simulation

[jcitrin]

Thanks for the ideas and motivation.

We and others already have some UQ frameworks that can explore TORAX uncertainty. There's an argument to be made that the API to trigger this should be outside of the TORAX library itself, and more in specific application frameworks.

In any case we are unfortunately very limited in the moment in review capacity and would not be able to work together on this due to other priorities.

[nytdevansh]

Makes sense. Keeping the UQ triggering outside of TORAX avoids unnecessary bloat and lets the application frameworks handle the heavy lifting.
No worries at all regarding the review capacity—I know how it goes. We’ll focus on other areas for now and can touch base if the priority shifts in the future. Thanks for the feedback on the approach!