Resumen Técnico: TORAX – Un simulador de transporte de tokamaks rápido y diferenciable en JAX

Planteamiento del problema

La simulación integrada precisa del plasma es crítica para el diseño y la operación de dispositivos de confinamiento magnético toroidal, como los tokamaks, sirviendo a aplicaciones que van desde la interpretación experimental hasta el diseño de controladores. Sin embargo, los códigos de transporte de núcleo existentes (por ejemplo, JETTO, ASTRA, PTRANSP) enfrentan un compromiso entre la fidelidad de la física y la velocidad de simulación; los modelos de alta fidelidad a menudo requieren horas o días para simular unos pocos tiempos de confinamiento de energía. Además, los códigos tradicionales, típicamente escritos en Fortran o MATLAB, carecen de diferenciabilidad nativa, lo que dificulta la optimización basada en gradientes, el análisis de sensibilidad y la integración fluida con modelos sustitutos (surrogates) de aprendizaje automático (ML). Si bien algunos códigos dedicados como RAPTOR ofrecen velocidad, dependen de Jacobianos codificados manualmente, lo que limita la extensibilidad a modelos complejos de ML. Existe una necesidad distintiva de un simulador que combine tiempos de ejecución rápidos, diferenciabilidad automática y soporte nativo para la integración con ML.

Metodología

Los autores presentan TORAX, un simulador de transporte de núcleo de tokamak de código abierto implementado en Python utilizando el marco de trabajo JAX.

Ecuaciones gobernantes y discretización

TORAX resuelve un conjunto acoplado de EDPs (ecuaciones diferenciales parciales) de transporte 1D en coordenadas de flujo toroidal normalizadas ($\hat{\rho}$). El sistema gobierna:

1. Transporte de calor iónico ($T_i$)
2. Transporte de calor electrónico ($T_e$)
3. Transporte de partículas electrónicas ($n_e$)
4. Difusión de corriente (flujo poloidal $\psi$)

Las ecuaciones se derivan de los momentos promediados por superficie de flujo de las ecuaciones cinéticas y la ley de Faraday. La discretización espacial se realiza mediante un Método de Volúmenes Finitos (FVM) en una rejilla uniforme. El esquema emplea un pesaje de ley de potencia para los flujos de convección-difusión para transicionar suavemente entre la diferencia central y el upwinding basándose en el número de Péclet.

Implementación del solver

La integración temporal utiliza el método theta, un promedio ponderado entre los esquemas Euler explícito e implícito. TORAX ofrece tres opciones de solver para manejar los sistemas algebraicos no lineales resultantes:

• Solver Lineal (Predictor-Corrector): Utiliza iteración de punto fijo. Para mejorar la estabilidad con coeficientes rígidos (stiff) (por ejemplo, provenientes de modelos de ML), implementa opcionalmente el método de Pereverzev-Corrigan, añadiendo difusión artificial equilibrada por convección.
• Solver Newton-Raphson: Un método de búsqueda de raíces basado en gradientes. Crucialmente, aprovecha la diferenciación automática de JAX para computar el Jacobiano del vector residual automáticamente, eliminando la necesidad de codificación manual del Jacobiano.
• Solver de Optimización: Recastea el residuo como una función de pérdida minimizada mediante algoritmos de optimización iterativos (usando jaxopt).

Modelos de Física

TORAX cuenta con un marco modular para componentes físicos:

• Geometría: Soporta datos de equilibrio CHEASE y una geometría circular simplificada.
  um
• Transporte: Incluye modelos de constante, Modelo de Gradiente Crítico (CGM) y QLKNN, un sustituto de aprendizaje automático entrenado con el código de la cinética de giros (gyrokinetic) QuaLiKiz. El modelo QLKNN está implementado nativamente en JAX utilizando la biblioteca Flax.
• Fuentes: Maneja el intercambio de calor ion-electrón, potencia de fusión (parametrización de Bosch-Hale), calentamiento óhmico y fuentes de partículas (gas puff, inyección de pellets, NBI). Las fuentes pueden tratarse como implícitas o explícitas.
• Física Neoclásica: Utiliza el modelo de Sauter para la corriente bootstrap y la conductividad neoclásica.

Arquitectura de Software

El código está estructurado en distintos módulos (Geometría, Modelos de Transporte, Fuentes, Steppers, Calculadores de Paso de Tiempo, Estado). Utiliza la compilación Just-In-Time (JIT) de JAX para acelerar los componentes computacionalmente intensivos (modelos de fuente, coeficientes de transporte, álgebra lineal) manteniendo el flujo de control en Python estándar para mayor flexibilidad. Este diseño permite la ejecución en CPUs, GPUs y TPUs.

Resultados Clave

Los autores verificaron TORAX contra el código establecido RAPTOR, el cual ha sido validado a su vez contra otros solvers de transporte.

1. Benchmark de Estado Estacionario (ITER L-mode):

   • Se simularon 10 segundos de un escenario tipo ITER con corriente y calentamiento constantes.
   • Los resultados mostraron una excelente concordancia entre TORAX y RAPTOR para los perfiles de temperatura iónica/electrónica, densidad, flujo poloidal y factor de seguridad.
   • Los valores de Desviación de Raíz Cuadrática Media Normalizada (NRMSD) fueron aproximadamente de 0.7% a 1.1% para perfiles estacionarios.
   • Las desviaciones transitorias de temperatura se limitaron a ~2.5%.

2. Benchmark de Escenario Dinámico:

   • Se simuló un escenario de 300 segundos con corriente de plasma variable en el tiempo (rampa de subida, flattop, rampa de bajada) y potencia de calentamiento modulada.
   • TORAX y RAPTOR siguieron la dinámica de cerca, con valores de NRMSD de ~3.5% a 4.0% para los perfiles de temperatura tanto en las escalas de tiempo de confinamiento de energía como de difusión de corriente.

3. Rendimiento:

   • En un nodo de CPU de 48 núcleos, una simulación de 80 segundos de la rampa de subida de ITER (usando QLKNN y Newton-Raphson) completó 22 segundos de tiempo de ejecución, con un tiempo de compilación de 15.6 segundos.
   • Esto demuestra una capacidad de tiempo real más rápido (faster-than-realtime) para configuraciones específicas, aunque los autores señalan que esto no es una conclusión general para todos los fenómenos transitorios.

Significancia y Reivindicaciones

El artículo posiciona a TORAX como una herramienta versátil diseñada para acelerar la investigación en el modelado de escenarios de tokamaks, diseño de pulsos y control. Su principal significancia radica en tres áreas:

1. Diferenciabilidad: Al aprovechar JAX, TORAX proporciona un residuo de EDP diferenciable. Esto permite la optimización basada en gradientes y el análisis de sensibilidad (vía Análisis de Sensibilidad hacia Adelante) sin la codificación manual del Jacobiano requerida por códigos anteriores como RAPTOR.
2. Integración de ML: El soporte nativo para redes neuronales dentro del ecosistema JAX facilita el acoplamiento fluido de sustitutos de ML (como QLKNN) como modelos de física, permitiendo que sean desarrollados en el mismo lenguaje que el simulador.
3. Velocidad y Flexibilidad: La compilación JIT asegura tiempos de ejecución rápidos, mientras que la arquitectura modular basada en Python simplifica la extensibilidad del código, el mantenimiento y el acoplamiento con flujos de trabajo externos.

Los autores concluyen que TORAX ofrece una "herramienta poderosa y versátil" para la comunidad de fusión, dirigida específicamente a aplicaciones que requieren tanto velocidad como diferenciabilidad, tales como la cuantificación de la incertidumbre, la optimización de escenarios y el diseño de controladores basados en modelos. El trabajo futuro planea expandir la fidelidad de la física (por ejemplo, MHD, radiación, múltiples especies iónicas) y mejorar las características de diferenciabilidad (por ejemplo, cachés de compilación persistentes, cálculos de sensibilidad hacia adelante).