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

Résumé technique de yancc : Un solveur accéléré par GPU et différentiable pour le transport néoclassique dans les tokamaks et les stellarators

Énoncé du problème
La théorie du transport néoclassique décrit le transport collisionnel des particules, de la chaleur et de la quantité de mouvement dans les dispositifs de confinement magnétique (tokamaks et stellarators). La résolution de l'équation de Boltzmann cinétique dérivée (DKE) est essentielle pour déterminer les champs électriques radiaux, les courants bootstrap et les coefficients de transport. Cependant, la résolution numérique de la DKE est notoirement difficile en raison d'une forte dominance de l'advection, de flux recirculants, de couches limites internes (particulièrement entre les populations de particules piégées et passantes), d'une anisotropie sévère et d'une haute dimensionnalité (quatre dimensions : deux spatiales, deux de vitesse).

Les solveurs existants font face à un compromis entre vitesse et fidélité physique. Les codes monoénergétiques (ex. : DKES, MONKES) sacrifient le couplage entre espèces et la diffusion d'énergie au profit de la performance. Les solveurs 4D complets (ex. : SFINCS, NEO) offrent une grande précision mais souffrent d'empreintes mémoire élevées et d'une faible extensibilité sur le matériel GPU en raison de factorisations de matrices creuses avec des accès mémoire irréguliers. De plus, l'absence de différentiabilité automatique dans les codes existants entrave l'optimisation basée sur le gradient et l'analyse de sensibilité par adjoint, qui sont de plus en plus requises pour la conception de stellarators et la quantification de l'incertitude.

Méthodologie
Les auteurs présentent yancc, un nouveau solveur construit sur le framework JAX pour tirer parti de l'accélération GPU et de la différenciation automatique. Le code résout à la fois la DKE 4D complète (conservant la dépendance de la collision vis-à-vis de la vitesse, la diffusion d'énergie et le plein couplage entre espèces) et la forme monoénergétique réduite.

• Discrétisation :

  • Vitesse ($x$) : Une grille de collocation utilisant des polynômes de Maxwell pondérés (orthogonaux par rapport à $e^{-x^2}$), ne nécessitant que 5 à 10 points.
  • Angles de pitch ($\alpha$), poloidal ($\theta$) et toroïdal ($\zeta$) : Discrétisation par différences finies.
  • Modification du stencil : Pour améliorer l'efficacité du multigrid, les auteurs utilisent un stencil amont modifié (par exemple, en utilisant les points $\{0, 1, 4\}$ pour l'ordre 2 et $\{-2, 0, 1, 3, 4\}$ pour l'ordre 4) plutôt que des stencils amont standards. Cela augmente la dominance diagonale des matrices du système linéaire résultant sans augmenter significativement la constante d'erreur de l'ordre principal.
  • Opérateur de collision : La diffusion champ-particule (potentiels de Rosenbluth) est traitée analytiquement via une approche par fonction de Green, résolvant les équations de Poisson associées dans l'espace des vitesses à l'aide d'expansions de Legendre et de fonctions gamma incomplètes, évitant ainsi de devoir résoudre les potentiels comme des inconnues supplémentaires.

• Solveur linéaire :

  • La DKE discrétisée forme un grand système linéaire creux $Af = s$.
  • Au lieu de solveurs directs creux (qui ne passent pas à l'échelle sur GPU), yancc utilise une méthode de sous-espace de Krylov (GCROT avec GMRES interne) préconditionnée par un schéma de multigrid géométrique.
  • Stratégie Multigrid :
    • Hiérarchie de grilles : Un sous-maillage (semi-coarsening) est utilisé ; la coordonnée de vitesse ($x$) est maintenue à sa résolution la plus fine sur tous les niveaux, tandis que $\alpha$, $\theta$ et $\zeta$ sont affinés par un facteur de 2.
    • Lissage (Smoothing) : Une approximation bloc-diagonale de l'opérateur est utilisée, effectuant des cycles de permutations qui maintiennent un couplage complet le long d'une coordonnée à la fois. Les stencils amont modifiés améliorent l'efficacité de ce lisseur bloc-diagonal.
    • Correction de grille grossière : La re-discrétisation directe est préférée à la projection de Galerkin pour maintenir une faible utilisation de la mémoire et éviter le coût élevé des opérations de grille grossière sans matrice (matrix-free).
    • Cycle : Un cycle V (indice 1) est utilisé pour la DKE complète, tandis qu'un cycle W ou supérieur (indice 3) s'avère efficace pour le cas monoénergétique.

Contributions clés

1. Solveur 4D complet accéléré par GPU : yancc est le premier solveur combinant la solution d'une DKE 4D complète (incluant le plein couplage entre espèces et la diffusion d'énergie) avec une exécution GPU efficace, surmontant les limitations de mémoire et de mise à l'échelle des codes basés sur des matrices creuses comme SFINCS.
2. Différentiabilité : Construit en JAX, le solveur est entièrement différentiable, permettant l'optimisation par gradient de la géométrie magnétique et l'analyse de sensibilité par adjoint sans dérivation manuelle des coefficients de transport.
3. Innovations algorithmiques : L'introduction de stencils amont élargis pour renforcer la dominance diagonale et l'utilisation du sous-maillage avec des lisseurs bloc-diagonaux adaptés à l'anisotropie spécifique de la DKE.
4. Potentiels de Rosenbluth analytiques : Une implémentation efficace de l'opérateur de collision champ-particule utilisant des fonctions de Green précalculées et des contractions de tenseurs, évitant la surcharge de calcul liée à la résolution d'équations de Poisson auxiliaires sur la grille.

Résultats et Benchmarks

• Précision : Les tests de référence par rapport à MONKES (monoénergétique) et SFINCS (4D complet) montrent un accord à moins de 1 % sur une large gamme de collisionalités ($\nu^* \in [10^{-4}, 10^2]$), de géométries magnétiques (W7-X, type NCSX) et de configurations de plasma (simple et multi-espèces).
• Performance vs SFINCS :
  • Vitesse : yancc atteint une accélération d'un ordre de grandeur par rapport à SFINCS sur une base de scan. Alors que le temps d'exécution de SFINCS augmente significativement à haute collisionalité (en raison de la dégradation du préconditionneur), le temps d'exécution de yancc reste presque plat.
  • Mémoire : yancc utilise un ordre de grandeur de mémoire en moins (ex. : ~6 Go contre >50 Go pour SFINCS dans les scénarios testés), ce qui le rend compatible avec les flux de travail basés sur GPU.
• Performance vs MONKES : Dans les benchmarks monoénergétiques, yancc est 2 à 4 fois plus rapide qu'un algorithme MONKES réimplémenté en JAX sur la majeure partie de la plage de collisionalité, bien que MONKES puisse être plus rapide pour des géométries très simples nécessitant une faible résolution.

Signification et Revendications
L'article affirme que la combinaison de la vitesse, de la faible empreinte mémoire et de la différentiabilité rend yancc idéalement adapté à l'intégration dans les flux de travail modernes d'optimisation de stellarators, la quantification de l'incertitude (via l'échantillonnage de Monte Carlo) et la prédiction de profils intégrés. Contrairement aux approches actuelles qui reposent sur des bases de données précalculées ou des fonctions d'ajustement, yancc peut calculer les coefficients néoclassiques à la volée au sein de solveurs de transport. La capacité de calculer les gradients des coefficients de transport par rapport aux paramètres d'entrée permet une optimisation efficace sur l'espace complet des configurations. Les auteurs positionnent yancc comme un outil pour combler le fossé entre la physique de haute fidélité et les exigences de calcul de la conception et de l'analyse, notant que les travaux futurs incluront le calcul auto-cohérent du champ électrique radial et le raffinement de maillage adaptatif.