Resumen Técnico: TokaGrad – Simulador de Tokamak Diferenciable de Extremo a Extremo

1. Planteamiento del Problema

A medida que la investigación en fusión magnética transiciona de la viabilidad científica a la comercialización de reactores a escala, el diseño de nuevos conceptos de reactores y la optimización de escenarios de alto rendimiento se han vuelto críticos. Los enfoques tradicionales para el desarrollo de escenarios dependen en gran medida de la prueba y error costosa, el ajuste de formas de onda dirigido por expertos y los escaneos de parámetros de baja dimensión. Aunque las herramientas de modelado integrado (por ejemplo, ASTRA, TRANSP) han mejorado la experimentación virtual, el diseño de escenarios sigue siendo un proceso de "humano en el bucle".

Las limitaciones principales de los métodos actuales son:

• Optimización de Caja Negra: Muchos enfoques de aprendizaje automático (optimización bayesiana, aprendizaje por refuerzo) tratan al simulador como una caja negra, requiriendo evaluaciones directas repetidas e inferencia indirecta de sensibilidad, lo cual se vuelve computacionalmente prohibitivo para escenarios de descarga completa, de alta dimensión y múltiples actuadores.
• Falta de Diferenciabilidad: Las simulaciones numéricas convencionales se ejecutan como procedimientos paso a paso donde los gradientes no pueden propagarse a través de todo el flujo de trabajo. Esto impide el control y la optimización directa basados en gradientes mediante sensibilidades internas.
• Modelado de Escenario Completo Incompleto: Los resolvedores de transporte diferenciables existentes (por ejemplo, TORAX) se han centrado principalmente en el transporte del núcleo y escenarios prescritos. Carecen de un modelado autoconsistente de escenarios dinámicos de descarga completa que involucren transiciones de L a H, formación de pedestal de borde y el acoplamiento de la evolución de la forma del plasma con la reconstrucción del equilibrio.

2. Metodología

Los autores presentan TokaGrad, un simulador de transporte de tokamak diferenciable de extremo a extremo construido sobre el framework JAX. El sistema representa todo el flujo de trabajo de simulación —desde las condiciones iniciales hasta el rendimiento final del plasma— como un grafo computacional conectado y diferenciable. Esto permite que la diferenciación automática (AD) compute los Jacobianos a través de toda la simulación, vinculando los parámetros de la máquina, las formas de onda de los actuadores y las respuestas del plasma.

Arquitectura Central

TokaGrad integra los siguientes módulos diferenciables dentro de un único grafo:

1. Marco de Transporte: El núcleo consiste en ecuaciones de transporte unidimensionales acopladas para temperatura electrónica/iónica ($T_{e,i}$), densidad electrónica ($n_e$) y flujo magnético poloidal ($\psi$) en una coordenada de flujo radial normalizada ($\rho$).

   • Esquema Numérico: Se utiliza un método de volumen finito semi-implícito para la estabilidad, resuelto mediante un resolvedor de matriz tridiagonal diferenciable (algoritmo de Thomas) compatible con la compilación Just-In-Time (JIT) de JAX.
   • Ecuaciones: Incluye transporte de partículas, transporte de energía conservativa (para $T_{e,i}$) y evolución del flujo inductivo acoplado a fuentes de bootstrap y de conducción de corriente.

2. Geometría y Equilibrio:

   • Modo EQDSK: Soporta equilibrios generados externamente (por ejemplo, desde EFIT) para comparativas realistas, aunque los gradientes respecto a la geometría son limitados a menos que se parametericen.
   • Modo Miller/Moment: Un modelo de frontera fija reducido utilizando parametrización analítica de superficies de flujo (radio mayor, elongación, triangularidad, desplazamiento de Shafranov). Este modo es totalmente diferenciable, permitiendo la optimización de los parámetros de forma de la máquina ($R_0, a, \kappa$).

3. Cierres de Transporte:

   • Transporte Turbulento: Utiliza sustitutos de Redes Neuronales (NN) (TGLF-NN, QLKNN) para reemplazar los costosos resolvedores de la teoría de la cinética de deriva (gyrokinetic). Estos sustitutos proporcionan mapeos diferenciables de los perfiles de plasma a los coeficientes de transporte.
   • Transporte Neoclásico: Soporta fórmulas analíticas (Sauter-Angioni) y sustitutos de NN (NEO-NN).
   • Modelos Empíricos: Incluye un modelo Bohm-gyroBohm (BgB) para líneas base rápidas.

4. Fuentes y Sumideros:

   • Calentamiento: Calentamiento óhmico, calentamiento alfa (usando el ajuste de Bosch-Hale), intercambio colisional y pérdidas por radiación. El calentamiento auxiliar se modela mediante perfiles gaussianos suaves y normalizados.
     el control de retroalimentación de densidad (rastreando la fracción de Greenwald) o modelos de forma prescrita para evitar cambios bruscos no diferenciables.

5. Transición L-H y Pedestal:

   • Transición: Utiliza una función sigmoide suave para controlar la transición entre los modos L y H basada en la potencia que cruza la separatriz ($P_{sep}$) relativa a un umbral empírico ($P_{LH}$, por ejemplo, escalamientos de Martin o Delabie). Esto evita interruptores binarios no diferenciables.
   • Pedestal: Construye objetivos de pedestal utilizando un modelo reducido de alfa-crítico o un sustituto EPED1-NN. El perfil del pedestal se mezcla suavemente con la solución de transporte del núcleo usando una función smoothmax, asegurando que los gradientes se propaguen a través de la transición.

3. Contribuciones Clave

• Primer Simulador de Descarga Completa Autoconsistente y Diferenciable: TokaGrad es el primer marco capaz de modelar escenarios de descarga completa dinámicos donde los actuadores, el equilibrio del plasma, la formación del pedestal y las transiciones de régimen de confinamiento (L a H) evolucionan conjuntamente dentro de un grafo diferenciable.
• Diferenciabilidad de Extremo a Extremo: Al convertir los cierres de física de alta fidelidad (TGLF, EPED1, NEO) en sustitutos de NN compatibles con JAX, el sistema permite que los gradientes se propaguen desde las métricas de rendimiento final hacia atrás, a través del transporte, la turbulencia y los modelos de equilibrio, hasta los inputs de diseño y control.
• Simulación Interactiva en Tiempo Real: El framework soporta simulaciones más rápidas que el tiempo real (demostrado como un aumento de velocidad >10x respecto a TORAX para pulsos largos tras la compilación JIT), permitiendo la exploración interactiva de las relaciones causa-efecto.
• Optimización Basada en Gradientes: El sistema facilita la optimización directa de las formas de onda de los actuadores, las trayectorias de los escenarios y los parámetros de diseño del reactor sin requerir escaneos de diferencia finita.

4. Resultados y Demostraciones

El artículo valida TokaGrad mediante comparativas y estudios de caso de optimización:

• Comparativa (Benchmarking):

  • Contra TORAX: En un escenario de modo L simplificado con geometría circular, TokaGrad reprodujo los perfiles de transporte del núcleo (temperatura, densidad, corriente) con un error de aproximadamente el 5% respecto al simulador TORAX basado en JAX.
  • Rendimiento: TokaGrad demostró una aceleración de tiempo de ejecución de más de un orden de magnitud en comparación con TORAX para simulaciones de pulso largo una vez amortizado el overhead de la compilación JIT.
  • Simulación de Modo H: Las simulaciones del escenario base de ITER utilizando TGLF-NN/EPED1-NN y modelos QLKNN/alfa-crítico lograron una ganancia de fusión $Q \sim 11$ y una potencia de fusión $\sim 550$ MW en ~7 segundos (más rápido que el tiempo real).

• Simulación de Descarga Completa L-a-H:

  • Una simulación de 150 segundos del escenario base de ITER modeló con éxito el rampa de subida, la fase de modo L y la transición al modo H.
  • La simulación capturó la histéresis de la transición L-H: un sobreimpulso de calentamiento transitorio (73 MW) desencadenó la transición, tras lo cual el calentamiento alfa sostuvo el modo H incluso cuando la potencia auxiliar se redujo a 50 MW.

• Aplicaciones de Optimización:

  • Control de Actuadores: El control basado en gradientes (optimizador Adam) ajustó con éxito la corriente del plasma ($I_p$), el calentamiento auxiliar ($P_{aux}$) y la fracción de Greenwald ($f_{GW}$) para maximizar las métricas de fusión ($Q + 0.01 P_{fus}$), respetando las restricciones de estabilidad ($q_{95} > 3, \beta_N < 3$).
  • Optimización de Forma de Onda de Escenarios: Partiendo de una rampa de subida ingenua que fallaba en acceder al modo H, el optimizador descubrió automáticamente una estrategia de sobreimpulso de calentamiento transitorio. Esto permitió que el plasma cruzara el umbral L-H y mantuviera un modo H de alto rendimiento bajo las mismas condiciones finales de los actuadores.
  • Optimización del Diseño del Reactor: El framework optimizó los parámetros de diseño de la máquina ($R_0, a, B_T, I_p$) para maximizar el rendimiento de la fusión. El optimizador identificó correctamente tendencias como el aumento del radio mayor y el campo magnético, mientras ajustaba los parámetros para satisfacer las restricciones de estabilidad.

5. Significación y Reivindicaciones

Los autores afirman que TokaGrad proporciona una vía práctica hacia el diseño automatizado de escenarios de tokamak. Al reemplazar la optimización de caja negra con métodos basados en gradientes habilitados por la física diferenciable, el framework permite:

• Análisis de Sensibilidad Directa: Computar cómo la ganancia de fusión depende de las trayectorias de calentamiento o cómo la altura del pedestal depende de la geometría sin perturbaciones externas.
• Optimización Dependiente de la Trayectoria: Optimizar no solo el estado final, sino toda la trayectoria dinámica (por ejemplo, las formas de onda de rampa de subida) necesaria para alcanzar regímenes de alto rendimiento como el modo H.
• Escalabilidad: La capacidad de realizar evaluaciones repetidas hacia adelante y hacia atrás en hardware de escala de portátil (Apple M1 Pro) hace que sea factible optimizar escenarios de reactores complejos y multiobjetivo que anteriormente eran computacionalmente impracticables.

El artículo concluye que, si bien las demostraciones actuales se centran en condiciones similares al escenario base de ITER, el framework es general y ofrece una nueva vía para la optimación automatizada y basada en la física de escenarios de plasma en combustión y conceptos de reactores de próxima generación.