Resumo Técnico: TORAX – Um Simulador de Transporte de Tokamak Rápido e Diferenciável em JAX

Problema

A simulação integrada precisa do plasma é crítica para o projeto e operação de dispositivos de confinamento magnético toroidal, como os tokamaks, servindo aplicações que variam desde a interpretação experimental até o design de controladores. No entanto, os códigos de transporte de núcleo existentes (ex: JETTO, ASTRA, PTRANSP) enfrentam um compromisso entre fidelidade física e velocidade de simulação; modelos de alta fidelidade frequentemente exigem horas ou dias para simular alguns tempos de confinamento de energia. Além disso, códigos tradicionais, tipicamente escritos em Fortran ou MATLAB, carecem de diferenciabilidade nativa, dificultando a otimização baseada em gradiente, análise de sensibilidade e integração contínua com modelos substitutos (surrogates) de aprendizado de máquina (ML). Embora alguns códigos dedicados como o RAPTOR ofereçam velocidade, eles dependem de Jacobianos codificados manualmente, o que limita a extensibilidade para modelos complexos de ML. Existe uma necessidade distinta de um simulador que combine tempos de execução rápidos, diferenciabilidade automática e suporte nativo para integração com ML.

Metodologia

Os autores apresentam o TORAX, um simulador de transporte de núcleo de tokamak de código aberto implementado em Python utilizando o framework JAX.

Equações Governantes e Discretização

O TORAX resolve um conjunto acoplado de Equações Diferenciais Parciais (PDEs) de transporte 1D em coordenadas de fluxo toroidal normalizadas ($\hat{\rho}$). O sistema governa:

1. Transporte de calor iônico ($T_i$)
2. Transporte de calor eletrônico ($T_e$)
3. Transporte de partículas eletrônicas ($n_e$)
4. Difusão de corrente (fluxo poloidal $\psi$)

As equações são derivadas de momentos de superfícies de fluxo das equações cinéticas e da lei de Faraday. A discretização espacial é realizada através de um Método de Volumes Finitos (FVM) em uma grade uniforme. O esquema emprega uma ponderação por lei de potência para fluxos de convecção-difusão para transicionar suavemente entre diferenciação central e upwinding com base no número de Péclet.

Implementação do Solver

A integração temporal utiliza o método theta, uma média ponderada entre os esquemas Euler explícito e implícito. O TORAX oferece três opções de solver para lidar com os sistemas algébricos não lineares resultantes:

• Solver Linear (Preditor-Corretor): Utiliza iteração de ponto fixo. Para aumentar a estabilidade com coeficientes rígidos (stiff) (ex: de modelos de ML), ele implementa opcionalmente o método de Pereverzev-Corrigan, adicionando difusão artificial balanceada por convecção.
• Solver Newton-Raphson: Um método de busca de raiz baseado em gradiente. Crucialmente, ele aproveita a diferenciação automática do JAX para computar o Jacobiano do vetor resíduo automaticamente, eliminando a necessidade de codificação manual de Jacobianos.
• Solver de Otimização: Recodifica o resíduo como uma função de perda minimizada via algoritmos de otimização iterativos (usando jaxopt).

Modelos de Física

O TORAX apresenta um framework modular para componentes físicos:

• Geometria: Suporta dados de equilíbrio CHEASE e uma geometria circular simplificada.
• Transporte: Inclui constante, Modelo de Gradiente Crítico (CGM) e QLKNN, um substituto de aprendizado de máquina treinado no código criogênico QuaLiKiz. O modelo QLKNN é implementado nativamente em JAX usando a biblioteca Flax.
• Fontes: Lida com troca de calor íon-elétron, potência de fusão (parametrização Bosch-Hale), aquecimento Ôhmico e fontes de partículas (injeção de gás, injeção de pellets, NBI). As fontes podem ser tratadas como implícitas ou explícitas.
• Física Neoclássica: Utiliza o modelo de Sauter para a corrente de bootstrap e condutividade neoclássica.

Arquitetura de Software

A base de código é estruturada em módulos distintos (Geometria, Modelos de Transporte, Fontes, Steppers, Calculadores de Passo de Tempo, Estado). Ela utiliza a compilação Just-In-Time (JIT) do JAX para acelerar componentes computacionalmente intensivos (modelos de fonte, coeficientes de transporte, álgebra linear) enquanto mantém o fluxo de controle em Python padrão para flexibilidade. Este design permite a execução em CPUs, GPUs e TPUs.

Principais Resultados

Os autores verificaram o TORAX contra o código estabelecido RAPTOR, que por sua vez foi validado contra outros solvers de transporte.

1. Benchmark de Estado Estacionário (ITER L-mode):

   • Simulou 10 segundos de um cenário tipo ITER com corrente e aquecimento constantes.
   • Os resultados mostraram excelente concordância entre o TORAX e o RAPTOR para perfis de temperatura iônica/eletrônica, densidade, fluxo poloidal e fatores de segurança.
   • Valores de Desvio Médio Quadrático Normalizado (NRMSD) foram de aproximadamente 0,7% a 1,1% para perfis estacionários.
   • Desvios transientes de temperatura foram limitados a ~2,5%.

2. Benchmark de Cenário Dinâmico:

   • Simulou um cenário de 300 segundos com corrente de plasma variante no tempo (rampa de subida, flattop, rampa de descida) e potência de aquecimento modulada.
   • O TORAX e o RAPTOR acompanharam a dinâmica de perto, com valores de NRMSD de ~3,5% a 4,0% para perfis de temperatura sobre tanto as escalas de tempo de confinamento de energia quanto de difusão de corrente.

3. Desempenho:

   • Em um nó de CPU de 48 núcleos, uma simulação de 80 segundos de rampa de subida do ITER (usando QLKNN e Newton-Raphson) completou em 22 segundos de tempo de execução, com um tempo de compilação de 15,6 segundos.
   • Isso demonstra capacidade de tempo real mais rápido que o real para configurações específicas, embora os autores notem que esta não é uma conclusão geral para todos os fenômenos transientes.

Significância e Alegações

O artigo posiciona o TORAX como uma ferramenta versátil projetada para acelerar a pesquisa em modelagem de cenários de tokamak, design de pulsos e controle. Sua principal significância reside em três áreas:

1. Diferenciabilidade: Ao aproveitar o JAX, o TORAX fornece um resíduo de PDE diferenciável. Isso permite otimização baseada em gradiente e análise de sensibilidade (via Análise de Sensibilidade Direta) sem a codificação manual de Jacobiano exigida por códigos diferenciais anteriores como o RAPTOR.
2. Integração de ML: O suporte nativo para redes neurais dentro do ecossistema JAX facilita o acoplamento contínuo de substitutos de ML (como o QLKNN) como modelos físicos, permitindo que sejam desenvolvidos na mesma linguagem do simulador.
3. Velocidade e Flexibilidade: A compilação JIT garante tempos de execução rápidos, enquanto a arquitetura modular baseada em Python simplifica a extensibilidade do código, manutenção e acoplamento com fluxos de trabalho externos.

Os autores concluem que o TORAX oferece uma "ferramenta poderosa e versátil" para a comunidade de fusão, visando especificamente aplicações que exigem tanto velocidade quanto diferenciabilidade, como quantificação de incerteza, otimização de cenários e design de controlador baseado em modelo. Trabalhos futuros estão planejados para expandir a fidelidade física (ex: MHD, radiação, múltiplas espécies iônicas) e aprimorar recursos de diferenciabilidade (ex: caches de compilação persistentes, cálculos de sensibilidade direta).