yancc: A GPU-accelerated, differentiable solver for neoclassical transport in tokamaks and stellarators

Resumen Técnico de yancc: A GPU-accelerated, differentiable solver for neoclassical transport in tokamaks and stellarators

Planteamiento del Problema
La teoría del transporte neoclásico describe el transporte colisional de partículas, calor y momento en dispositivos de confinamiento magnético (tokamaksaks y stellarators). Resolver la ecuación cinética de deriva (DKE) es esencial para determinar los campos eléctricos radiales, las corrientes bootstrap y los coeficientes de transporte. Sin embargo, la solución numérica de la DKE es notoriamente difícil debido al fuerte dominio de la advección, flujos recirculantes, capas límite internas (particularmente entre las poblaciones de partículas atrapadas y libres) una severa anisotropía y alta dimensionalidad (cuatro dimensiones: dos espaciales, dos de velocidad).

Los solvers existentes enfrentan un compromiso entre velocidad y fidelidad física. Los códigos monoenergéticos (p. ej., DKES, MONKES) sacrifican el acoplamiento entre especies y la dispersión de energía en favor del rendimiento. Los solvers de 4D completos (p. ej., SFINCS, NEO) ofrecen alta precisión pero sufren de altas huellas de memoria y una escalabilidad deficiente hacia el hardware de GPU debido a factorizaciones de matrices dispersas con acceso irregular a la memoria. Además, la falta de diferenciabilidad automática en los códigos existentes impide la optimización basada en gradientes y el análisis de sensibilidad adjunta, que son cada vez más necesarios para el diseño de stellarators y la cuantificación de la incertidumbre.

Metodología
Los autores presentan yancc, un nuevo solver construido sobre el framework JAX para aprovechar la aceleración por GPU y la diferenciación automática. El código resuelve tanto la DKE de 4D completa (manteniendo la dependencia de la velocidad de las colisiones, la dispersión de energía y el acoplamiento completo entre especies) como la forma monoenergética reducida.

• Discretización:

  • Velocidad ($x$): Una rejilla de colocación utilizando polinomios de Maxwell ponderados (ortogonales con respecto a $e^{-x^2}$), que requiere solo de 5 a 10 puntos.
  • Ángulos de Pitch ($\alpha$), Poloidal ($\theta$) y Toroidal ($\zeta$): Discretización por diferencias finitas.
  • Modificación de la Esténcil: Para mejorar la eficiencia de multigrid, los autores emplean un esténcil de flujo ascendente (upwind) modificado (p. ej., usando puntos $\{0, 1, 4\}$ para 2º orden y $\{-2, 0, 1, 3, 4\}$ para 4º orden) en lugar de los esténciles de flujo ascendente estándar. Esto aumenta la dominancia diagonal de las matrices de los sistemas lineales resultantes sin aumentar significativamente la constante de error de orden principal.
  • Operador de Colisión: El proceso de dispersión campo-partícula (potenciales de Rosenbluth) se maneja analíticamente mediante un enfoque de función de Green, resolviendo las ecuaciones de Poisson asociadas en el espacio de velocidades utilizando expansiones de Legendre y funciones gamma incompletas, evitando la necesidad de resolver los potenciales como incógnitas adicionales.

• Solver Lineal:

  • La DKE discretizada forma un sistema lineal disperso de gran tamaño $Af = s$.
  • En lugar de solvers directos dispersos (que fallan en escalar en GPUs), yancc utiliza un método de subespacio de Krylov (GCROT con GMRES interno) precondicionado por un esquema de multigrid geométrico.
  • Estrategia de Multigrid:
    • Jerarquía de Rejillas: Se utiliza el semicoarsening (refinamiento parcial); la coordenada de velocidad ($x$) se mantiene en la resolución más fina en todos los niveles, mientras que $\alpha$, $\theta$ y $\zeta$ se coarsen por un factor de 2.
    • Suavizado (Smoothing): Se utiliza una aproximación de diagonal de bloque del operador, ciclando a través de permutaciones que mantienen el acoplamiento completo a lo largo de una coordenada a la vez. Los esténciles de flujo ascendente modificados mejoran la efectividad de este suavizador de bloque diagonal.
    • Corrección de Rejilla Gruesa: Se prefiere la rediscretización directa sobre la proyección de Galerkin para mantener un bajo uso de memoria y evitar el alto costo de las operaciones de rejilla gruesa sin matriz (matrix-free).
    • Ciclo: Se utiliza un ciclo V (índice 1) para la DKE completa, mientras que un ciclo W o superior (índice 3) resulta eficiente para el caso monoenergético.

Contribuciones Clave

1. Solver de 4D Completo Acelerado por GPU: yancc es el primer solver que combina la solución de una DKE de 4D completa (incluyendo el acoplamiento completo entre especies y la dispersión de energía) con una ejecución eficiente en GPU, superando las limitaciones de memoria y escalabilidad de los códigos basados en matrices dispersas como SFINCS.
2. Diferenciabilidad: Construido en JAX, el solver es totalmente diferenciable, lo que permite la optimización de la geometría magnética basada en gradientes y el análisis de sensibilidad adjunta sin la derivación manual de los coeficientes de transporte.
3. Innovaciones Algorítmicas: La introducción de esténciles de flujo ascendente ensanchados para mejorar la dominancia diagonal y el uso de semicoarsening con suavizadores de diagonal de bloque adaptados a la anisotropía específica de la DKE.
4. Potenciales de Rosenbluth Analíticos: Una implementación eficiente del operador de colisión campo-partícula utilizando funciones de Green precomputadas y contracciones de tensores, evitando la sobrecarga computacional de resolver ecuaciones de Poisson auxiliares en la rejilla.

Resultados y Benchmarks

• Precisión: Las comparaciones contra MONKES (monoenergético) y SFINCS (4D completo) muestran una concordancia dentro del 1% a través de un amplio rango de colisionalidades ($\nu^* \in [10^{-4}, 10^2]$), geometrías magnéticas (W7-X, tipo NCSX) y configuraciones de plasma (una y múltiples especies).
• Rendimiento vs. SFINCS:
  • Velocidad: yancc logra una aceleración de un orden de magnitud sobre SFINCS en una base por escaneo. Mientras que el tiempo de ejecución de SFINCS aumenta significativamente a alta colisionalidad (debido a la degradación del precondicionador), el tiempo de ejecución de yancc permanece casi plano.
  • Memoria: yancc utiliza un orden de magnitud menos de memoria (p. ej., ~6 GB frente a >50 GB para SFINCS en los escenarios probados), lo que lo hace factible para flujos de trabajo basados en GPU.
• Rendimiento vs. MONKES: En benchmarks monoenergéticos, yancc es de 2 a 4 veces más rápido que un algoritmo MONKES reimplementado en JAX en la mayor parte del rango de colisionalidad, aunque MONKES puede ser más rápido para geometrías muy simples que requieren baja resolución.

Significancia y Reivindicaciones
El artículo afirma que la combinación de velocidad, baja huella de memoria y diferenciabilidad hace que yancc sea único para su integración en los flujos de trabajo modernos de optimización de stellarators, la cuantificación de la incertidumbre (vía muestreo de Monte Carlo) y la predicción de perfiles integrados. A diferencia de los enfoques actuales que dependen de bases de datos precomputadas o funciones de ajuste, yancc puede computar coeficientes neoclásicos on-the-fly dentro de los solvers de transporte. La capacidad de calcular gradientes de coeficientes de transporte con respecto a los parámetros de entrada permite una optimización eficiente sobre el espacio de configuración completo. Los autores posicionan a yancc como una herramienta para cerrar la brecha entre la física de alta fidelidad y las demandas computacionales de diseño y análisis, señalando que el trabajo futuro incluirá cálculos de campo eléctrico radial autoconstante y refinamiento de malla adaptativo.