Fast solvers for Tokamak fluid models with PETSC

Este artigo apresenta um novo solucionador geométrico multigrid de semi-refinamento implementado no PETSc para o código de tokamak M3D-C1, que aborda as limitações de convergência do pré-condicionador Jacobi por blocos existente ao explorar a estrutura da grade toroidal para alcançar robustez e desempenho superiores em modelos complexos de magnetohidrodinâmica.

O essencial

Imagine que você está tentando prever como uma sopa superquente e turbulenta de partículas carregadas (plasma) se comporta dentro de uma máquina gigante em forma de donut chamada tokamak. Esta máquina é projetada para criar energia de fusão, como a potência do Sol. No entanto, essa "sopa" é incrivelmente caótica. Se você tentar calcular seu movimento passo a passo usando um computador, a matemática fica tão complicada que o computador trava, ou leva tanto tempo que a resposta se torna inútil no momento em que chega.

Este artigo trata de construir uma calculadora mais inteligente e rápida para este tipo específico de problema.

Aqui está a explicação do que os autores fizeram, usando analogias simples:

1. O Problema: O "engarrafamento" da Matemática

O código de computador que eles usam (chamado M3D-C1) tenta resolver equações que descrevem como o plasma se move. Para fazer isso, ele precisa resolver um quebra-cabeça massivo milhões de vezes.

• O Jeito Antigo (Block Jacobi): Imagine que você tem um mapa enorme de uma cidade com engarrafamentos. O método antigo era como pedir para uma pessoa diferente consertar o tráfego em apenas uma rua de cada vez, ignorando o resto da cidade. Se a cidade for pequena, isso funciona. Mas, à medida que a cidade fica maior (mais "planos" ou fatias da forma de donut), as pessoas que consertam as ruas não conseguem conversar umas com as outras rápido o suficiente. Os engarrafamentos pioram e a solução fica lenta ou para de funcionar completamente.
• O Desafio Específico: O plasma nessas máquinas é "anisotrópico". Pense nisso como uma pilha de papel. É muito fácil deslizar uma folha de papel ao longo da superfície (direção fácil), mas é muito difícil empurrá-la através da pilha (direção difícil). O antigo solucionador matemático não entendia essa estrutura de "pilha de papel", então tentou resolver a direção difícil e a direção fácil com o mesmo método desajeitado.

2. A Solução: O Elevador "Multigrid"

Os autores criaram um novo solucionador usando um método chamado Multigrid (MG).

• A Analogia: Imagine que você está tentando encontrar um brinquedo perdido em uma mansão enorme de vários andares.
  • O Jeito Antigo: Você verifica cada quarto, cada gaveta e cada canto no térreo antes de subir. Isso leva uma eternidade.
  • O Jeito Multigrid: Primeiro, você olha para um modelo em miniatura de toda a mansão de uma visão de pássaro. Você identifica rapidamente a área geral onde o brinquedo está faltando (a grade "grossa"). Depois, você dá zoom para um mapa de tamanho médio para estreitar a busca. Finalmente, você vai para o quarto real (a grade "fina") para pegar o brinquedo.
  • Ao resolver o problema nos níveis de "grande visão" primeiro, o solucionador sabe exatamente onde procurar quando chega aos detalhes minúsculos. Isso o torna incrivelmente rápido.

3. O "Segredo": Semi-Afinamento

Os autores perceberam que o tokamak é como uma pilha de fatias 2D (planos poloidais) torcidas em um donut 3D.

• Eles aplicaram seu "Elevador Multigrid" especificamente na direção de empilhamento (a direção toroidal).
• Em vez de tentar simplificar toda a bagunça 3D de uma vez, eles mantiveram as fatias 2D detalhadas (porque a forma da parede do tokamak é complexa), mas tornaram a pilha de fatias mais simples à medida que subiam nos níveis.
• Isso é como pegar um livro grosso e reduzir apenas o número de páginas, mantendo o texto em cada página claro. É um ajuste perfeito para a forma da máquina.

4. Os Resultados: Velocidade e Confiabilidade

A equipe testou esse novo solucionador em dois cenários muito difíceis:

• Cenário A: O "Elétron Descontrolado" (SPARC): Isso simula um evento perigoso onde partículas aceleram de forma incontrolável.
  • Resultado: O novo solucionador foi competitivo com o antigo em configurações menores e muito mais rápido nas configurações maiores e mais complexas. Ele resolveu o problema em menos etapas, economizando tempo.
• Cenário B: O "Estelarator" (uma máquina diferente e mais torcida): Esta geometria é ainda mais torcida e irregular do que um donut padrão.
  • Resultado: O antigo solucionador falhou completamente e não conseguiu encontrar uma resposta. O novo solucionador Multigrid teve sucesso. Foi robusto o suficiente para lidar com a geometria torcida que quebrou a ferramenta antiga.

5. O Hardware: Usando Supercomputadores

Eles executaram esses testes no Perlmutter, um dos supercomputadores mais rápidos do mundo, que usa tanto CPUs poderosas quanto GPUs (placas de vídeo).

• Eles descobriram que, embora a "configuração" (construir os modelos em miniatura) fosse cara, a resolução real foi incrivelmente rápida nas GPUs.
• Eles descobriram que, para os problemas mais difíceis, era necessário usar um "suavizador" de "alta capacidade" (um truque matemático específico) para evitar que o solucionador travasse, o que exigiu um pouco mais de poder de computação, mas valeu a pena em velocidade.

Resumo

O artigo afirma que, ao entender a forma específica de "pilha de papel" das máquinas de plasma de fusão, eles criaram uma nova ferramenta matemática (Multigrid) que:

1. Resolve problemas mais rápido do que o método padrão atual em simulações grandes e complexas.
2. Não trava em formas torcidas e complexas onde o método antigo falha.
3. É um primeiro passo crucial para tornar as simulações de energia de fusão práticas e rápidas o suficiente para ajudar a projetar usinas de energia reais.

Eles não afirmaram que isso resolve a energia de fusão em si, mas sim que fornece a calculadora rápida e confiável necessária para simular a física que eventualmente levará à energia de fusão.

Resumo técnico

Resumo Técnico: Solucionadores Rápidos para Modelos de Fluidos de Tokamak com PETSc

Declaração do Problema
O código M3D-C1 é uma ferramenta primária para simular plasmas magnetohidrodinâmicos (MHD) estendidos em pesquisas de energia de fusão, abordando especificamente instabilidades e disrupções em grande escala em tokamaks e estelares. O código emprega um integrador de tempo semi-implícito que requer a resolução de inúmeros sistemas lineares por passo de tempo. Dentre estes, o avanço implícito da equação de momento é identificado como o componente mais desafiador e menos robusto.

O solucionador de produção atual utiliza um método de Krylov (FGMRES/GMRES) pré-condicionado por um método de Jacobi de Bloco (BJ), onde os blocos agrupam graus de liberdade em planos de coordenada toroidal constante. Embora eficaz para configurações menores, a convergência do solucionador BJ degrada-se significativamente à medida que o número de planos toroidais aumenta, devido às propriedades espectrais da matriz pré-condicionada. Além disso, o solucionador BJ falha em convergir totalmente em configurações de estelar não-eixo-simétricas, limitando a aplicabilidade do código a dispositivos de fusão complexos.

Metodologia
Este trabalho desenvolve um solucionador de multigrid geométrico (MG) adaptado à discretização específica e à estrutura de grade do M3D-C1. A metodologia aproveita as seguintes características estruturais:

• Estrutura de Grade: O M3D-C1 utiliza grades 2D não estruturadas no plano poloidal (extrudadas ao redor do toro) combinadas com uma grade 1D regular na direção toroidal.
• Semi-Afinamento: O solucionador emprega semi-afinamento geométrico especificamente na direção toroidal. Esta abordagem é escolhida porque a grade toroidal é estruturada, permitindo técnicas de MG geométrico padrão (afinamento 2:1), enquanto deixa o plano poloidal complexo sem afinamento para trabalhos futuros.
• Interface Algébrica: Embora a aplicação calcule Jacobianos em grades finas, os operadores de grade grossa são construídos algebricamente usando um processo de Galerkin ($A_{i+1} = P^T A_i P$), onde $P$ é um operador de prolongamento de stencil de 3 pontos. Isso cria um solucionador híbrido geométrico/algebrico com uma interface puramente algebrica compatível com o PETSc.
• Suavizadores: Dois suavizadores foram investigados:
  1. Jacobi de Bloco Pontual (PBJ): Calcula inversas densas explícitas de blocos diagonais (tamanho 36 por vértice). É rápido e eficiente em memória, adequado para aceleração por GPU.
  2. Jacobi de Bloco (BJ): O solucionador de produção existente é reaproveitado como suavizador. Envolve fatorações LU completas em planos poloidais. Embora caro, é usado como suavizador pós-em um ciclo $V(0,1)$ para os problemas mais difíceis onde o PBJ estagna.
• Implementação: O solucionador é implementado dentro da biblioteca PETSc, utilizando seus solucionadores de Krylov e interfaces abstratas para back-ends de álgebra linear (incluindo MUMPS e SuperLU para resoluções diretas).

Principais Resultados
O solucionador foi avaliado em dois casos de teste usando o supercomputador Perlmutter (nós de CPU e GPU): um modelo de elétron runaway de uma disrupção de tokamak SPARC e um modelo de estelar quase-eixo-simétrico.

1. Tokamak SPARC (Escalabilidade e Velocidade):

   • Em configurações com 4 a 16 planos toroidais, o solucionador MG (usando suavização PBJ) convergiu em uma única iteração quando o pré-condicionador foi atualizado, demonstrando excelente escalabilidade algorítmica.
   • Na maior configuração (64 planos), a dificuldade do problema aumentou, exigindo uma tolerância de solucionador mais rigorosa ($rtol = 10^{-7}$) e um ciclo $V(6,6)$ mais agressivo com atualizações frequentes.
   • Desempenho: O solucionador MG alcançou desempenho competitivo em comparação com o solucionador BJ. No caso de 64 planos, o solucionador MG reduziu o número total de iterações de Krylov por um fator de aproximadamente 16x em comparação com o solucionador BJ. Embora o tempo total de resolução fosse aproximadamente comparável em GPUs, o solucionador MG mostrou tempos de resolução significativamente mais rápidos em CPUs.
   • Complexidade: A análise teórica sugere que o solucionador MG tem uma complexidade de trabalho assintótica de aproximadamente 2x a do solucionador BJ devido à soma geométrica do trabalho através dos níveis de grade. O desempenho observado indica que as resoluções de grade grossa estão atualmente performando mal, sugerindo espaço para otimização.

2. Estelar (Robustez):

   • Em um caso de teste de estelar não-eixo-simétrico, o solucionador padrão Jacobi de Bloco falhou em convergir completamente.
   • Em contraste, o solucionador MG (usando BJ como suavizador) convergiu com sucesso. Isso demonstra uma vantagem crítica de robustez da abordagem MG para geometrias não-eixo-simétricas onde as propriedades espectrais da matriz derrotam o pré-condicionador BJ.

Significado e Alegações
O artigo alega apresentar o primeiro passo na aplicação de métodos de multigrid aos complexos modelos de MHD relevantes para a ciência dentro do código M3D-C1. As principais contribuições são:

• Demonstração de Viabilidade: Provar que métodos de multigrid podem fornecer um caminho para solucionadores mais rápidos, escaláveis e robustos para modelos de plasma magnetizado altamente anisotrópicos e mal condicionados.
• Semi-Afinamento Toroidal: Estabelecer que o semi-afinamento na direção toroidal é um primeiro passo útil e eficaz para códigos comuns de tokamak com estruturas de grade extrudadas.
• Robustez: Mostrar que solucionadores MG podem resolver problemas (especificamente disrupções de estelares) onde o atual solucionador de produção Jacobi de Bloco falha.

Trabalho Futuro e Limitações
Os autores permanecem modestos quanto ao estado atual do solucionador, notando várias áreas para melhoria:

• Desempenho de Grade Grossa: O desempenho atual é limitado por resoluções de grade grossa ineficientes, parcialmente devido a requisitos redundantes de paralelismo e falhas de MPI em solucionadores diretos (MUMPS/SuperLU) em grades grossas.
• Multigrid 3D Completo: A otimização futura mais impactante é a adição de multigrid 2D no plano poloidal para criar um solucionador MG 3D completo. Isso exigiria o manuseio de grades geométricas grossas não estruturadas.
• Suavizadores: Trabalhos futuros incluem a otimização dos tamanhos de bloco para suavizadores PBJ para se ajustarem melhor a grupos de threads de GPU e a investigação de pré-condicionadores FieldSplit que tratam as três ondas MHD desacopladas separadamente.
• Estratégias de Atualização: O modelo de atualização atual (fatoração LU periódica) é simples; algoritmos mais inteligentes são necessários para lidar com estados quase-estacionários de forma mais eficiente.

Em resumo, este trabalho estabelece uma estrutura de multigrid viável e robusta para o M3D-C1 que supera o solucionador existente em termos de contagens de iteração e robustez em geometrias difíceis, ao mesmo tempo em que identifica obstáculos algorítmicos e de implementação específicos a serem abordados para alcançar a eficiência completa de multigrid de livro-texto.