🔬 physics

jQMC: A JAX-based ab initio quantum Monte Carlo package designed for GPU-accelerated computing

O artigo apresenta o jQMC, um pacote de Monte Carlo Quântico ab initio baseado em Python e habilitado para JAX que aproveita a aceleração por GPU e a diferenciação automática para alcançar ganhos de desempenho significativos sobre as implementações tradicionais em Fortran tanto para simulações de Monte Carlo de Difusão Variacional quanto de Difusão Regularizada por Rede.

Autores originais: Kousuke Nakano, Michele Casula

Publicado 2026-07-16

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Kousuke Nakano, Michele Casula

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que o universo é construído de minúsculas partículas dançantes chamadas elétrons, zunindo ao redor de núcleos atômicos como abelhas hiperativas em torno de uma colmeia. Para entender como os materiais se comportam, como os remédios funcionam ou por que algumas coisas conduzem eletricidade enquanto outras não, os cientistas precisam resolver um problema matemático massivo e emaranhado chamado equação de Schrödinger. É como tentar prever a trajetória exata de cada abelha individual em um enxame enquanto todas elas colidem entre si e com as paredes da colmeia. Por décadas, as melhores ferramentas que tínhamos eram ou lentas demais para lidar com grandes enxames ou precisavam fazer muitos atalhos, levando a suposições que às vezes erravam o alvo.

Entra o Quantum Monte Carlo (QMC). Pense nisso não como um único cálculo, mas como um jogo de simulação massivo onde milhares de "caminhadores" (fantasmas digitais dos elétrons) correm por um mundo virtual, dando passos aleatórios para explorar cada arranjo possível. Ao tirar a média de onde esses caminhadores terminam, os cientistas podem obter uma imagem muito precisa da energia e do comportamento do sistema. É um pouco como tentar adivinhar a altura média de todas as pessoas em um estádio fazendo com que alguns milhares de pessoas circulem pelas arquibancadas e relatem o que viram, em vez de medir cada pessoa individualmente. O problema? Esse jogo é incrivelmente caro em termos computacionais. Requer supercomputadores para rodar tempo suficiente para obter uma resposta confiável e, por muito tempo, o software que executava esses jogos foi construído para tipos de computadores mais antigos e lentos.

Isso nos traz ao jQMC, um novo software introduzido por Kousuke Nakano e Michele Casula. O artigo apresenta o jQMC como um motor moderno, baseado em Python, projetado especificamente para rodar nas placas de vídeo (GPUs) mais recentes e rápidas encontradas nos supercomputadores mais poderosos de hoje. Os autores construíram essa ferramenta para ver se conseguiam fazer essas simulações de elétrons rodarem significativamente mais rápido e de forma mais eficiente do que o software que é o padrão ouro atual, que foi escrito em uma linguagem de programação mais antiga (Fortran) e não foi originalmente projetado para esses novos chips.

A equipe descobriu que o jQMC é, de fato, um demônio da velocidade quando combinado com GPUs modernas. Para a parte de "Monte Carlo Variacional" (VMC) da simulação — que é como a rodada de prática onde os caminhadores aprendem o básico — o jQMC em uma GPU foi até dez vezes mais rápido que o software tradicional rodando em uma CPU, e ainda mais rápido que o software tradicional tentando rodar em uma GPU. Para a parte mais complexa de "Monte Carlo de Difusão" (LRDMC), que é a rodada final de alto risco, o jQMC ainda foi cerca de três vezes mais rápido em uma GPU. Os pesquisadores também testaram o quão bem o jQMC consegue calcular as "forças" que empurram e puxam os átomos (essencial para entender como as moléculas se movem e reagem). Eles descobriram que, para uma molécula chamada etanol, os cálculos de força do jQMC foram incrivelmente precisos, aproximando-se da teoria "coupled-cluster" (cluster acoplado) em cerca de 1,5 kcal/mol/Å. No entanto, para uma molécula diferente, a malonaldeído, os resultados foram bastante diferentes do padrão ouro, sugerindo que, mesmo com esta nova e poderosa ferramenta, alguns quebra-cabeças químicos complexos ainda guardam surpresas.

Em resumo, o jQMC é um experimento bem-sucedido de modernização da química quântica. Ele prova que, ao reescrever as regras do jogo para se adequar ao hardware do futuro (GPUs), os cientistas podem rodar essas simulações complexas muito mais rápido, abrindo as portas para o estudo de sistemas maiores e mais complicados do que nunca. Embora não resolva todos os mistérios da química, ele oferece uma maneira muito mais rápida e flexível de enfrentar aqueles que já conhecemos.

Resumo Técnico de jQMC: Um pacote de Monte Carlo Quântico ab initio baseado em JAX projetado para computação acelerada por GPU

Definição do Problema
Resolver a equação de Schrödinger de muitos corpos continua sendo um desafio central na química e física computacionais. Embora a Teoria do Funcional da Densidade (DFT) seja amplamente utilizada, sua dependência de funcionais de troca-correlação aproximados limita sua precisão em sistemas fortemente correlacionados. Métodos de Monte Carlo Quântico (QMC) ab initio, como o Monte Carlo Variacional (VMC) e o Monte Carlo de Difusão (DMC), oferecem uma alternativa estocástica com escalonamento favorável, mas são computacionalmente exigentes. A mudança das instalações de Computação de Alto Desempenho (HPC) para arquiteturas dominadas por GPUs apresenta um obstáculo significativo para os códigos legados de QMC, que são frequentemente escritos em Fortran ou C e dependem de designs centrados em CPU. Refatorar esses códigos para explorar o paralelismo de GPU é difícil, e os códigos de QMC acelerados por GPU existentes muitas vezes carecem da flexibilidade de frameworks modernos ou das características algorítmicas específicas necessárias para simulações ab initio de alta precisão. Além disso, há uma necessidade de formatos de função de onda padronizados e técnicas de otimização eficientes que possam aproveitar a diferenciação automática.

Metodologia
Os autores apresentam o jQMC, um pacote computacional baseado em Python para simulações de QMC ab initio construído inteiramente sobre a biblioteca JAX. O JAX permite a compilação Just-In-Time (JIT), diferenciação automática (AD) e portabilidade contínua entre backends de CPU e GPU.

Algoritmos: O jQMC implementa dois algoritmos de QMC de estado fundamental:
1. Monte Carlo Variacional (VMC): Utiliza o algoritmo de Metropolis-Hastings para amostrar configurações eletrônicas. Emprega o esquema de reponderação de Attacali e Sorella (AS) para lidar com singularidades próximas às superfícies nodais.
2. Monte Carlo de Difusão Regularizado em Rede (LRDMC): Uma variante do Monte Carlo de Função Verde (GFMC) que discretiza o Hamiltoniano em uma rede para melhorar a estabilidade numérica. O jQMC implementa esquemas de ramificação (branching) convencional e de equilíbrio de carga para mitigar o desequilíbrio de carga em cálculos paralelos.
Ansatz da Função de Onda: A função de onda de tentativa é construída como um produto de um fator Jastrow e uma parte antissimétrica. O componente antissimétrico suporta:
- Potencial Geminal Antissimetrizado (AGP): Uma função de onda do tipo ligação de valência ressonante (RVB) que inclui pareamentos de singlete e triplete de spin.
- AGP com orbitais moleculares restritos (AGPn): Uma variante que restringe o posto da matriz de pareamento.
- Determinante de Slater Único (SD): Como um caso especial de AGP.
  A função de onda pode ser inicializada a partir de cálculos externos de Hartree-Fock (HF) ou Teoria do Funcional da Densidade (DFT) via biblioteca TREX-IO.
Otimização: O pacote implementa a Reconfiguração Estocástica (SR) e o Método Linear (LM) para otimizar parâmetros variacionais. Esses métodos utilizam a métrica do espaço de parâmetros para atualizar amplitudes da função de onda e superfícies nodais simultaneamente.
Diferenciação e Forças: Aproveitando a diferenciação automática do JAX, o jQMC computa eficientemente as derivadas de energia. Ele implementa cálculos de forças atômicas tanto para VMC (usando Transformação de Coordenadas de Espaço-Warp) quanto para LRDMC (usando a aproximação de Reynolds e regularização de Pathak-Wagner).
Arquitetura: O código separa estruturas de dados (implementadas como dataclasses compatíveis com JAX) de algoritmos. Utiliza vmap para vetorização sobre os walkers (o paralelismo de alto nível em QMC) e suporta escalonamento multi-GPU via MPI (mpi4py) e particionamento (sharding) do JAX. Também suporta aritmética de precisão mista (FP32 para avaliações de orbitais atômicos, FP64 para determinantes/geminais) para acelerar o cálculo sem sacrificar a precisão.

Principais Contribuições

Implementação Baseada em JAX: O jQMC é o primeiro pacote de QMC a aproveitar totalmente o JAX tanto para diferenciação algorítmica quanto para aceleração por GPU, permitindo código compilado via JIT e portátil entre arquiteturas de CPU e GPU.
Suporte Avançado de Função de Onda: Fornece uma implementação robusta do ansatz Jastrow-Antisymmetrized Geminal Power (JAGP), que generaliza o determinante de Slater único e permite a otimização de superfícies nodais.
Vetorização Eficiente: Ao vetorizar sobre os walkers no nível superior, o jQMC alcança alta eficiência intra-GPU, saturando os recursos da GPU de forma mais eficaz do que os códigos tradicionais paralelizados em CPU.
Padronização: A integração com o TREX-IO garante a compatibilidade com outros pacotes de estrutura eletrônica (ex: PySCF, CP2K).
Avaliação de Forças: A implementação de forças atômicas usando diferenciação automática e técnicas de regularização específicas (Pathak-Wagner) para LRDMC permite a otimização de geometria e cálculos de fônon dentro do framework QMC.

Resultados

Benchmarks de Desempenho:
- GPU vs. CPU: Em GPUs NVIDIA (A100 e H100), o jQMC demonstra acelerações de mais de duas ordens de magnitude em comparação com sua própria implementação em CPU.
- Comparação com TurboRVB: Quando comparado ao TurboRVB (um código Fortran90 que implementa os mesmos algoritmos e funções de onda), o jQMC em GPUs é significativamente mais rápido. Para VMC, os ganhos de tempo de execução (wall-time) atingem uma ordem de magnitude. Para LRDMC, o ganho é mais moderado, mas ainda substancial (aprox. 3x mais rápido para os tamanhos de sistema testados).
- Escalabilidade: Benchmarks de escalonamento fraco nos supercomputadores Miyabi (Universidade de Tóquio) e Leonardo (CINECA) mostram que o jQMC mantém uma eficiência quase ideal (eficiência $\approx$ 1.0) até 1024 GPUs para VMC e até 512 GPUs para LRDMC (usando o algoritmo de equilíbrio de carga).
Validação de Precisão:
- Reprodutibilidade: Os resultados do jQMC para energias totais e energias de ligação de metano, água e o dímero metano-água concordam com outros códigos principais de QMC (CASINO, QMCPACK, CMQMC, TurboRVB) dentro das margens de erro estatístico, validando a implementação dos algoritmos DLA (Determinant Locality Approximation) e LRDMC.
- Forças Atômicas: O artigo avalia a precisão das forças atômicas para etanol e malonaldeia contra cálculos CCSD(T) all-electron.
  - Para o etanol, as forças de LRDMC alcançam um Erro Médio Absoluto (MAE) de $\sim$ 1.5 kcal/mol/Å em relação ao CCSD(T), consistente com a literatura anterior.
  - Para a malonaldeia, observa-se uma discrepância significativa, onde as forças de LRDMC e CCSD(T) diferem em $\sim$ 5 kcal/mol/Å (MAE). Os autores observam que essa discrepância persiste mesmo com superfícies nodais otimizadas, sugerindo que o problema pode não ser apenas devido à aproximação da superfície nodal, exigindo investigação adicional.

Significância
O artigo posiciona o jQMC como uma plataforma moderna e versátil que preenche a lacuna entre métodos QMC ab initio de alta precisão e as modernas arquiteturas de HPC aceleradas por GPU. Ao utilizar o JAX, o pacote oferece um ambiente flexível para prototipagem rápida de novos algoritmos, ao mesmo tempo em que entrega desempenho de nível de produção. Os autores afirmam que o jQMC aborda com sucesso os desafios de portar códigos legados de QMC para GPUs, fornecendo uma ferramenta que não é apenas mais rápida que as implementações baseadas em Fortran existentes em hardware GPU, mas também mais fácil de estender e manter. A validação bem-sucedida de forças e energias contra outros códigos e métodos de química quântica de alto nível demonstra que o jQMC está pronto para simulações de produção de sistemas moleculares e estendidos, particularmente em regimes onde efeitos de correlação forte são críticos.

Mais como este