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

Resumo Técnico de yancc: Um solver acelerado por GPU e diferenciável para transporte neoclássico em tokamaks e stellarators

Problema
A teoria do transporte neoclássico descreve o transporte colisional de partículas, calor e momento em dispositivos de confinamento magnético por fusão (tokamaks e stellarators). Resolver a equação cinética de deriva (DKE) é essencial para determinar campos elétricos radiais, correntes bootstrap e coeficientes de transporte. No entanto, a solução numérica da DKE é notoriamente difícil devido à forte dominância de advecção, fluxos recirculantes, camadas de contorno internas (particularmente entre populações de partículas presas e de passagem), anisotropia severa e alta dimensionalidade (quatro dimensões: duas espaciais, duas de velocidade).

Os solvers existentes enfrentam um compromisso entre velocidade e fidelidade física. Códigos monoenergéticos (ex: DKES, MONKES) sacrificam o acoplamento entre espécies e o espalhamento de energia em prol do desempenho. Solvers 4D completos (ex: SFINCS, NEO) oferecem alta precisão, mas sofrem com altas pegadas de memória e baixa escalabilidade para hardware GPU devido a fatorações de matrizes esparsas com acesso irregular à memória. Além disso, a falta de diferenciabilidade automática nos códigos existentes impede a otimização baseada em gradiente e a análise de sensibilidade adjunta, que são cada vez mais necessárias para o design de stellarators e quantificação de incertezas.

Metodologia
Os autores apresentam o yancc, um novo solver construído sobre o framework JAX para aproveitar a aceleração por GPU e a diferenciação automática. O código resolve tanto a DKE 4D completa (mantendo colisões dependentes da velocidade, espalhamento de energia e acoplamento total entre espécies) quanto a forma monoenergética reduzida.

• Discretização:

  • Velocidade ($x$): Uma grade de colocação usando polinômios de Maxwell ponderados (ortogonais em relação a $e^{-x^2}$), exigindo apenas 5–10 pontos.
  • Ângulos de Pitch ($\alpha$), Poloidal ($\theta$) e Toroidal ($\zeta$): Discretização por diferenças finitas.
  • Modificação do Stencil: Para melhorar a eficiência do multigrid, os autores empregam um stencil de upwind modificado (por exemplo, usando pontos $\{0, 1, 4\}$ para 2ª ordem e $\{-2, 0, 1, 3, 4\}$ para 4ª ordem) em vez de stencils de upwind padrão. Isso aumenta a dominância diagonal das matrizes do sistema linear resultante sem aumentar significativamente a constante de erro de ordem principal.
  • Operador de Colisão: O espalhamento campo-partícula (potenciais de Rosenbluth) é tratado analiticamente via uma abordagem de função de Green, resolvendo as equações de Poisson associadas no espaço de velocidade usando expansões de Legendre e funções gama incompletas, evitando a necessidade de resolver pelos potenciais como incógnitas adicionais.

• Solver Linear:

  • A DKE discretizada forma um sistema linear grande e esparso $Af = s$.
    imposto de que, em vez de solvers diretos esparsos (que falham em escalar em GPUs), o yancc utiliza um método de subespaço de Krylov (GCROT com GMRES interno) pré-condicionado por um esquema de multigrid geométrico.
  • Estratégia de Multigrid:
    • Hierarquia de Grades: O refinamento semicoarsening é usado; a coordenada de velocidade ($x$) é mantida na resolução mais fina em todos os níveis, enquanto $\alpha$, $\theta$ e $\zeta$ são refinados por um fator de 2.
    • Suavização (Smoothing): Uma aproximação de bloco-diagonal do operador é usada, ciclando através de permutações que mantêm o acoplamento total ao longo de uma coordenada por vez. Os stencils de upwind modificados aumentam a eficácia deste suavizador de bloco-diagonal.
    • Correção de Grade Grossa: A re-discretização direta é preferida em vez da projeção de Galerkin para manter o baixo uso de memória e evitar o alto custo de operações de grade grossa sem matriz (matrix-free).
    • Ciclo: Um ciclo V (índice 1) é usado para a DKE completa, enquanto um ciclo W ou superior (índice 3) mostra-se eficiente para o caso monoenergético.

Principais Contribuições

1. Solver 4D Completo Acelerado por GPU: O yancc é o primeiro solver que combina a solução da DKE 4D completa (incluindo acoplamento total entre espécies e espalhamento de energia) com execução eficiente em GPU, superando as limitações de memória e escalabilidade de códigos baseados em matrizes esparsas como o SFINCS.
2. Diferenciabilidade: Construído em JAX, o solver é totalmente diferenciável, permitindo a otimização de geometria magnética baseada em gradiente e análise de sensibilidade adjunta sem a necessidade de derivação manual de coeficientes de transporte.
3. Inovações Algorítmicas: A introdução de stencils de upwind alargados para aumentar a dominância diagonal e o uso de semi-coarsening com suavizadores de bloco-diagonal adaptados à especificidade da anisotropia da DKE.
4. Potenciais de Rosenbluth Analíticos: Uma implementação eficiente do operador de colisão campo-partícula usando funções de Green pré-computadas e contrações de tensores, evitando o overhead computacional de resolver equações de Poisson auxiliares na grade.

Resultados e Benchmarks

• Precisão: Testes de benchmark contra MONKES (monoenergético) e SFINCS (4D completo) mostram concordância dentro de 1% em uma ampla gama de colisionalidades ($\nu^* \in [10^{-4}, 10^2]$), geometrias magnéticas (W7-X, similares a NCSX) e configurações de plasma (espécie única e múltiplas espécies).
• Desempenho vs. SFINCS:
  • Velocidade: O yancc alcança uma aceleração de uma ordem de magnitude sobre o SFINCS por varredura (per-scan). Enquanto o tempo de execução do SFINCS aumenta significativamente em altas colisionalidades (devido à degradação do pré-condicionador), o tempo de execução do yancc permanece quase constante.
  • Memória: O yancc utiliza uma ordem de magnitude menos memória (ex: ~6 GB vs. >50 GB para o SFINCS nos cenários testados), tornando-o viável para fluxos de trabalho baseados em GPU.
• Desempenho vs. MONKES: Em benchmarks monoenergéticos, o yancc é de 2 a 4 vezes mais rápido que um algoritmo MONKES reimplementado em JAX na maior parte do intervalo de colisionalidade, embora o MONKES possa ser mais rápido para geometrias muito simples que exigem baixa resolução.

Significado e Alegações
O artigo afirma que a combinação de velocidade, baixa pegada de memória e diferenciabilidade torna o yancc unicamente adequado para integração em fluxos de trabalho modernos de otimização de stellarators, quantificação de incertezas (via amostragem de Monte Carlo) e predição de perfis integrados. Ao contrário das abordagens atuais que dependem de bancos de dados pré-computados ou funções de ajuste, o yancc pode computar coeficientes neoclássicos "on-the-fly" dentro de solvers de transporte. A capacidade de computar gradientes de coeficientes de transporte em relação aos parâmetros de entrada permite a otimização eficiente sobre todo o espaço de configuração. Os autores posicionam o yancc como uma ferramenta para preencher a lacuna entre a física de alta fidelidade e as demandas computacionais de design e análise, observando que trabalhos futuros incluirão cálculos de campo elétrico radial autossustentáveis e refinamento de malha adaptativo.