Fast solvers for Tokamak fluid models with PETSC

Este artículo presenta un nuevo solucionador geométrico multigrid de semi-afinamiento implementado en PETSc para el código de tokamak M3D-C1, que aborda las limitaciones de convergencia del precondicionador de Jacobi por bloques existente aprovechando la estructura de la red toroidal para lograr una robustez y un rendimiento superiores en modelos complejos de magnetohidrodinámica.

Lo esencial

Imagina que intentas predecir cómo se comporta una sopa supercaliente y turbulenta de partículas cargadas (plasma) dentro de una máquina gigante con forma de dona llamada tokamak. Esta máquina está diseñada para crear energía de fusión, como la del Sol. Sin embargo, esta "sopa" es increíblemente caótica. Si intentas calcular su movimiento paso a paso usando una computadora, las matemáticas se vuelven tan complicadas que la computadora se atasca, o tarda tanto que la respuesta es inútil para cuando llega.

Este artículo trata sobre construir una calculadora más inteligente y rápida para este tipo específico de problema.

Aquí tienes el desglose de lo que hicieron los autores, usando analogías sencillas:

1. El Problema: El "Atasco de Tráfico" de las Matemáticas

El código informático que utilizan (llamado M3D-C1) intenta resolver ecuaciones que describen cómo se mueve el plasma. Para hacerlo, tiene que resolver un rompecabezas masivo millones de veces.

• La Vieja Forma (Jacobi por Bloques): Imagina que tienes un mapa enorme de una ciudad con atascos de tráfico. El método antiguo consistía en pedirle a una persona diferente que arreglara el tráfico en una sola calle a la vez, ignorando el resto de la ciudad. Si la ciudad es pequeña, esto funciona. Pero a medida que la ciudad se hace más grande (más "planos" o rebanadas de la forma de dona), las personas que arreglan las calles no pueden hablar entre sí lo suficientemente rápido. Los atascos de tráfico empeoran y la solución se ralentiza o deja de funcionar por completo.
• El Desafío Específico: El plasma en estas máquinas es "anisotrópico". Piénsalo como una pila de papel. Es muy fácil deslizar una hoja a lo largo de la superficie (dirección fácil), pero es muy difícil empujarla a través de la pila (dirección difícil). El antiguo solucionador matemático no entendía esta estructura de "pila de papel", así que intentó resolver la dirección difícil y la dirección fácil con el mismo método torpe.

2. La Solución: El Ascensor "Multigrilla"

Los autores construyeron un nuevo solucionador utilizando un método llamado Multigrilla (MG).

• La Analogía: Imagina que intentas encontrar un juguete perdido en una mansión masiva de varios pisos.
  • La Vieja Forma: Revisas cada habitación, cada cajón y cada rincón en la planta baja antes de subir. Esto toma una eternidad.
  • La Forma Multigrilla: Primero miras un modelo en miniatura de toda la mansión desde una vista de pájaro. Rápidamente localizas el área general donde falta el juguete (la "cuadrícula" gruesa). Luego, haces zoom en un mapa de tamaño medio para estrechar la búsqueda. Finalmente, vas a la habitación real (la "cuadrícula" fina) para recoger el juguete.
  • Al resolver el problema primero en los niveles de "gran perspectiva", el solucionador sabe exactamente dónde buscar cuando llega a los detalles minúsculos. Esto lo hace increíblemente rápido.

3. El "Ingrediente Secreto": Semi-Afinamiento

Los autores se dieron cuenta de que el tokamak es como una pila de rebanadas 2D (planos poloidales) retorcidas en una dona 3D.

• Aplicaron su "Ascensor Multigrilla" específicamente a la dirección de apilamiento (la dirección toroidal).
• En lugar de intentar simplificar todo el desorden 3D de una vez, mantuvieron las rebanadas 2D detalladas (porque la forma de la pared del tokamak es compleja), pero hicieron que la pila de rebanadas fuera más simple a medida que subían de nivel.
• Esto es como tomar un libro grueso y solo reducir el número de páginas mientras mantienes el texto de cada página claro. Es un ajuste perfecto para la forma de la máquina.

4. Los Resultados: Velocidad y Fiabilidad

El equipo probó este nuevo solucionador en dos escenarios muy difíciles:

• Escenario A: El "Electrón Desbocado" (SPARC): Esto simula un evento peligroso donde las partículas se aceleran de forma incontrolable.
  • Resultado: El nuevo solucionador fue competitivo con el antiguo en configuraciones más pequeñas y mucho más rápido en las configuraciones más grandes y complejas. Resolvió el problema en menos pasos, ahorrando tiempo.
• Escenario B: El "Estelarator" (Una máquina diferente, más retorcida): Esta geometría es aún más retorcida e irregular que una dona estándar.
  • Resultado: El antiguo solucionador falló por completo y no pudo encontrar una respuesta. El nuevo solucionador Multigrilla tuvo éxito. Fue lo suficientemente robusto para manejar la geometría retorcida que rompió la antigua herramienta.

5. El Hardware: Usando Supercomputadoras

Realizaron estas pruebas en Perlmutter, una de las supercomputadoras más rápidas del mundo, que utiliza tanto CPUs potentes como GPUs (tarjetas gráficas).

• Descubrieron que, aunque la "configuración" (construir los modelos en miniatura) era costosa, la resolución real era increíblemente rápida en las GPUs.
• Descubrieron que para los problemas más difíciles, necesitaban usar un suavizador "de trabajo pesado" (un truco matemático específico) para evitar que el solucionador se atascara, lo cual requería un poco más de potencia de cálculo, pero valió la pena en velocidad.

Resumen

El artículo afirma que, al comprender la forma específica de "pila de papel" de las máquinas de plasma de fusión, crearon una nueva herramienta matemática (Multigrilla) que:

1. Resuelve problemas más rápido que el método estándar actual en simulaciones grandes y complejas.
2. No se bloquea en formas retorcidas y complejas donde el método antiguo falla.
3. Es un primer paso crucial para hacer que las simulaciones de energía de fusión sean prácticas y lo suficientemente rápidas para ayudar a diseñar centrales eléctricas reales.

No afirmaron que esto resuelva la energía de fusión en sí, sino que proporciona la calculadora rápida y fiable necesaria para simular la física que eventualmente conducirá a la energía de fusión.

Resumen técnico

Resumen Técnico: Solucionadores Rápidos para Modelos de Fluidos de Tokamak con PETSc

Planteamiento del Problema
El código M3D-C1 es una herramienta principal para simular plasmas magnetohidrodinámicos (MHD) extendidos en la investigación de energía de fusión, abordando específicamente inestabilidades a gran escala y disrupciones en tokamaks y estelares. El código emplea un integrador temporal semi-implícito que requiere resolver numerosos sistemas lineales por paso de tiempo. Entre estos, la avance implícito de la ecuación de momento se identifica como el componente más desafiante y menos robusto.

El solucionador de producción actual utiliza un método de Krylov (FGMRES/GMRES) precondicionado por un método de Jacobi por Bloques (BJ), donde los bloques agrupan los grados de libertad en planos de coordenada toroidal constante. Aunque efectivo para configuraciones más pequeñas, la convergencia del solucionador BJ se degrada significativamente a medida que aumenta el número de planos toroidales debido a las propiedades espectrales de la matriz precondicionada. Además, el solucionador BJ falla en converger por completo en configuraciones de estelares no axisimétricos, limitando la aplicabilidad del código a dispositivos de fusión complejos.

Metodología
Este trabajo desarrolla un solucionador de multigrid geométrico (MG) adaptado a la discretización específica y la estructura de la malla de M3D-C1. La metodología aprovecha las siguientes características estructurales:

• Estructura de la Malla: M3D-C1 utiliza mallas 2D no estructuradas en el plano poloidal (extruidas alrededor del toro) combinadas con una malla 1D regular en la dirección toroidal.
• Coarsening Semi-Geométrico: El solucionador emplea un semi-coarsening geométrico específicamente en la dirección toroidal. Este enfoque se elige porque la malla toroidal está estructurada, lo que permite técnicas estándar de MG geométrico (coarsening 2:1) mientras se deja el plano poloidal complejo sin coarsening para trabajos futuros.
• Interfaz Algebraica: Aunque la aplicación calcula jacobianos en mallas finas, los operadores de malla gruesa se construyen algebraicamente utilizando un proceso de Galerkin ($A_{i+1} = P^T A_i P$), donde $P$ es un operador de prolongación de plantilla de 3 puntos. Esto crea un solucionador híbrido geométrico/algebraico con una interfaz puramente algebraica compatible con PETSc.
• Suavizadores: Se investigan dos suavizadores:
  1. Jacobi por Bloques Puntual (PBJ): Calcula inversas densas explícitas de los bloques diagonales (tamaño 36 por vértice). Es rápido y eficiente en memoria, adecuado para aceleración por GPU.
  2. Jacobi por Bloques (BJ): El solucionador de producción existente se reutiliza como suavizador. Implica factorizaciones LU completas en planos poloidales. Aunque costoso, se utiliza como suavizador posterior en un ciclo $V(0,1)$ para los problemas más difíciles donde PBJ se estanca.
• Implementación: El solucionador se implementa dentro de la biblioteca PETSc, utilizando sus solucionadores de Krylov e interfaces abstractas para back-ends de álgebra lineal (incluyendo MUMPS y SuperLU para resoluciones directas).

Resultados Clave
El solucionador se evaluó en dos casos de prueba utilizando el superordenador Perlmutter (nodos CPU y GPU): un modelo de electrones desbocados de una disrupción de tokamak SPARC y un modelo de estelar cuasi-axisimétrico.

1. Tokamak SPARC (Escalabilidad y Velocidad):

   • En configuraciones con 4 a 16 planos toroidales, el solucionador MG (usando suavizado PBJ) convergió en una sola iteración cuando el precondicionador se actualizó, demostrando una excelente escalabilidad algorítmica.
   • En la configuración más grande (64 planos), la dificultad del problema aumentó, requiriendo una tolerancia del solucionador más estricta ($rtol = 10^{-7}$) y un ciclo $V(6,6)$ más agresivo con actualizaciones frecuentes.
   • Rendimiento: El solucionador MG logró un rendimiento competitivo en comparación con el solucionador BJ. En el caso de 64 planos, el solucionador MG redujo el número total de iteraciones de Krylov en un factor de aproximadamente 16x en comparación con el solucionador BJ. Aunque el tiempo total de resolución fue aproximadamente comparable en GPUs, el solucionador MG mostró tiempos de resolución significativamente más rápidos en CPUs.
   • Complejidad: El análisis teórico sugiere que el solucionador MG tiene una complejidad de trabajo asintótica de aproximadamente 2x la del solucionador BJ debido a la suma geométrica del trabajo a través de los niveles de malla. El rendimiento observado indica que las resoluciones de malla gruesa están funcionando actualmente de manera deficiente, lo que sugiere margen para optimización.

2. Estelar (Robustez):

   • En un caso de prueba de estelar no axisimétrico, el solucionador estándar de Jacobi por Bloques falló en converger por completo.
   • Por el contrario, el solucionador MG (usando BJ como suavizador) convergió con éxito. Esto demuestra una ventaja crítica de robustez del enfoque MG para geometrías no axisimétricas donde las propiedades espectrales de la matriz derrotan al precondicionador BJ.

Significado y Afirmaciones
El artículo afirma presentar el primer paso en la aplicación de métodos de multigrid a los complejos modelos MHD de tokamak relevantes para la ciencia dentro del código M3D-C1. Las contribuciones principales son:

• Demostración de Viabilidad: Probar que los métodos de multigrid pueden proporcionar un camino hacia solucionadores más rápidos, escalables y robustos para modelos de plasma magnetizado altamente anisotrópicos y mal condicionados.
• Semi-Coarsening Toroidal: Establecer que el semi-coarsening en la dirección toroidal es un primer paso útil y efectivo para códigos comunes de tokamak con estructuras de malla extruidas.
• Robustez: Mostrar que los solucionadores MG pueden resolver problemas (específicamente disrupciones de estelares) donde el solucionador de producción actual de Jacobi por Bloques falla.

Trabajo Futuro y Limitaciones
Los autores permanecen modestos respecto al estado actual del solucionador, señalando varias áreas de mejora:

• Rendimiento de Malla Gruesa: El rendimiento actual está limitado por resoluciones de malla gruesa ineficientes, en parte debido a requisitos de paralelismo redundantes y fallos de MPI en solucionadores directos (MUMPS/SuperLU) en mallas gruesas.
• Multigrid 3D Completo: La optimización futura más impactante es la adición de multigrid 2D en el plano poloidal para crear un solucionador MG 3D completo. Esto requeriría manejar mallas geométricas gruesas no estructuradas.
• Suavizadores: El trabajo futuro incluye optimizar los tamaños de bloque para los suavizadores PBJ para ajustarse mejor a los grupos de hilos de GPU e investigar precondicionadores FieldSplit que traten las tres ondas MHD desacopladas por separado.
• Estrategias de Actualización: El modelo de actualización actual (factorización LU periódica) es simple; se necesitan algoritmos más inteligentes para manejar estados cuasi-estacionarios de manera más eficiente.

En resumen, este trabajo establece un marco de multigrid viable y robusto para M3D-C1 que supera al solucionador existente en términos de recuentos de iteraciones y robustez en geometrías difíciles, mientras identifica obstáculos algorítmicos y de implementación específicos que deben abordarse para lograr una eficiencia de multigrid de libro de texto completa.