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

Technische Samenvatting van yancc: Een GPU-versnelde, differentieerbare solver voor neoclassicistische transport in tokamaks en stellarators

Probleemstelling
Neoclassicistische transporttheorie beschrijft het collisionele transport van deeltjes, warmte en impuls in magnetische opsluitingsapparaten voor kernfusie (tokamaks en stellarators). Het oplossen van de drift-kinetische vergelijking (DKE) is essentieel voor het bepalen van radiale elektrische velden, bootstrap-stromen en transportcoëfficiënten. Het numeriek oplossen van de DKE is echter berucht moeilijk vanwege sterke advectie-dominantie, recirculerende stromingen, interne grenslagen (met name tussen gevangen en passerende deeltjespopulaties), extreme anisotropie en hoge dimensionaliteit (vier dimensies: twee ruimtelijk, twee snelheid).

Bestaande solvers staan voor een afweging tussen snelheid en fysieke getrouwheid. Monoenergetische codes (bijv. DKES, MONKES) offeren interspecies-koppeling en energieverstrooiing op voor prestaties. Full 4D-solvers (bijv. SFINCS, NEO) bieden een hoge nauwkeurigheid, maar lijden onder een hoge geheugenvoetafdruk en slechte schaalbaarheid naar GPU-hardware door ijle matrixfactorisaties met onregelmatige geheugentoegang. Bovendien belemmert het gebrek aan automatische differentieerbaarheid in bestaande codes gradiëntgebaseerde optimalisatie en adjoint-gevoeligheidsanalyse, wat steeds vaker vereist is voor stellaratorontwerp en onzekerheidskwantificering.

Methodologie
De auteurs presenteren yancc, een nieuwe solver gebouwd op het JAX-framework om gebruik te maken van GPU-versnelling en automatische differentiatie. De code lost zowel de volledige 4D DKE op (waarbij snelheid-afhankelijke botsingen, energieverstrooiing en volledige interspecies-koppeling behouden blijven) als de gereduceerde monoenergetische vorm.

• Discretisatie:

  • Snelheid ($x$): Een collocatiegrid met gewogen Maxwell-polynomen (orthogonaal met betrekking tot $e^{-x^2}$), die slechts 5–10 punten vereist.
  • Pitch-hoek ($\alpha$), poloidale ($\theta$) en toroidale ($\zeta$) hoeken: Eindige-verschillen-discretisatie.
  • Stencil-modificatie: Om de multigrid-efficiëntie te verbeteren, gebruiken de auteurs een gemodificeerde upwind-stencil (bijv. gebruikmakend van punten $\{0, 1, 4\}$ voor 2e orde en $\{-2, 0, 1, 3, 4\}$ voor 4e orde) in plaats van standaard upwind-stencils. Dit vergroot de diagonaaldominantie van de resulterende lineaire systeemmatrices zonder de leidende-orde foutconstante significant te verhogen.
  • Botsingsoperator: De veld-deeltje verstrooiing (Rosenbluth-potentialen) wordt analytisch afgehandeld via een Green-functie benadering, waarbij de bijbehorende Poisson-vergelijkingen in de snelheidsruimte worden opgelost met behulp van Legendre-expansies en incomplete gammafuncties, waardoor het niet nodig is om voor de potentialen als extra onbekenden op te lossen.

• Lineaire Solver:

  • De gediscretiseerde DKE vormt een groot, ijl lineair systeem $Af = s$.
  • In plaats van ijle directe solvers (die slecht schalen op GPU's), gebruikt yancc een Krylov-subruimte methode (GCROT met innerlijke GMRES) gepreconditioneerd door een geometrisch multigrid-schema.
  • Multigrid Strategie:
    • Grid Hiërarchie: Semi-coarsening wordt gebruikt; de snelheidscoördinaat ($x$) wordt op de fijnste resolutie gehouden over alle niveaus, terwijl $\alpha$, $\theta$ en $\zeta$ met een factor 2 worden vergroefd (coarsened).
    • Smoothing: Een blokdiagonaal benadering van de operator wordt gebruikt, waarbij door permutaties wordt gecyceld die volledige koppeling langs één coördinaat tegelijkertijd in stand houden. De gemodificeerde upwind-stencils verbeteren de effectiviteit van deze blokdiagonaal smoother.
    • Coarse Grid Correctie: Directe her-discretisatie wordt verkozen boven Galerkin-projectie om het geheugengebruik laag te houden en de hoge kosten van matrix-vrije coarse-grid operaties te vermijden.
    • Cycle: Een V-cycle (index 1) wordt gebruikt voor de volledige DKE, terwijl een W-cycle of hogere (index 3) efficiënt wordt bevonden voor het monoenergetische geval.

Belangrijkste Bijdragen

1. GPU-versnelde Full 4D Solver: yancc is de eerste solver die een volledige 4D DKE-oplossing (inclusclusief volledige interspecies-koppeling en energieverstrooiing) combineert met efficiënte GPU-executie, waarmee de geheugen- en schaalbaarheidsbeperkingen van sparse-matrix-gebaseerde codes zoals SFINCS worden overwonnen.
2. Differentieerbaarheid: Gebouwd in JAX, is de solver volledig differentieerbaar, wat gradiëntgebaseerde optimalisatie van magnetische geometrie en adjoint-gevoeligheidsanalyse mogelijk maakt zonder handmatige afleiding van transportcoëfficiënten.
3. Algoritmische Innovaties: De introductie van verbreedde upwind-stencils om de diagonaaldominantie te verhogen en het gebruik van semi-coarsening met blokdiagonaal smoothers die zijn afgestemd op de specifieke anisotropie van de DKE.
4. Analytische Rosenbluth-potentialen: Een efficiënte implementatie van de veld-deeltje botsingsoperator met behulp van vooraf berekende Green-functies en tensorcontracties, wat de computationele overhead van het oplossen van hulp-Poisson-vergelijkingen op het grid vermijdt.

Resultaten en Benchmarks

• Nauwkeurigheid: Benchmarks tegenover MONKES (monoenergetisch) en SFINCS (full 4D) tonen overeenstemming binnen 1% over een breed bereik van collisionaliteiten ($\nu^* \in [10^{-4}, 10^2]$), magnetische geometrieën (W7-X, NCSX-achtig) en plasma-configuraties (enkelvoudig en meervoudig deeltjestype).
• Prestaties vs. SFINCS:
  • Snelheid: yancc bereikt een orde van grootte sneller dan SFINCS op basis van een per-scan analyse. Terwijl de runtime van SFINCS significant toeneemt bij hoge collisionaliteit (door degradatie van de preconditioner), blijft de runtime van yancc bijna vlak.
  • Geheugen: yancc gebruikt een orde van grootte minder geheugen (bijv. ~6 GB vs. >50 GB voor SFINCS in geteste scenario's), wat het geschikt maakt voor GPU-gebaseerde workflows.
• Prestaties vs. MONKES: In monoenergetische benchmarks is yancc 2–4 keer sneller dan een in JAX geïmplementeerd MONKES-algoritme over het grootste deel van het collisionaliteitsbereik, hoewel MONKES sneller kan zijn voor zeer eenvoudige geometrieën die een lage resolutie vereisen.

Betekenis en Claims
Het artikel beweert dat de combinatie van snelheid, lage geheugenvoetafdruk en differentieerbaarheid yancc uniek geschikt maakt voor integratie in moderne stellarator-optimalisatieworkflows, onzekerheidskwantificering (via Monte Carlo-sampling) en geïntegreerde profielvoorspelling. In tegen tegenstelling tot huidige benaderingen die vertrouwen op vooraf berekende databases of fitfuncties, kan yancc neoclassicistische coëfficiënten on-the-fly berekenen binnen transport-solvers. Het vermogen om gradiënten van transportcoëfficiënten te berekenen ten opzichte van inputparameters maakt efficiënte optimalisatie over de volledige configuratieruimte mogelijk. De auteurs positioneren yancc als een instrument om de kloof te overbruggen tussen hoge-getrouwheid fysica en de computationele eisen van ontwerp en analyse, waarbij zij opmerken dat toekomstig werk zelfconsistente radiale elektrische veldberekeningen en adaptieve mesh-verfijning zal omvatten.