🔬 physics

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

El artículo presenta jQMC, un paquete de Monte Carlo cuántica ab initio basado en Python y habilitado para JAX que aprovecha la aceleración por GPU y la diferenciación automática para lograr mejoras significativas de rendimiento sobre las implementaciones tradicionales en Fortran tanto para simulaciones de Monte Carlo de difusión variacional como de red regularizada.

Autores originales: Kousuke Nakano, Michele Casula

Publicado 2026-07-16

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Kousuke Nakano, Michele Casula

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que el universo está construido de diminutas partículas danzantes llamadas electrones, que zumban alrededor de núcleos atómicos como abejas hiperactivas alrededor de una colmena. Para entender cómo se comportan los materiales, cómo funcionan las medicinas o por qué algunas cosas conducen la electricidad mientras otras no, los científicos necesitan resolver un problema matemático masivo y enredado llamado la ecuación de Schrödinger. Es como intentar predecir la trayectoria exacta de cada una de las abejas en un enjambre mientras todas chocan entre sí y contra las paredes de la colmena. Durante décadas, las mejores herramientas que teníamos eran demasiado lentas para manejar enjambres grandes o tenían que tomar demasiados atajos, lo que llevaba a conjeturas que a veces no daban en el clavo.

Entra el Quantum Monte Carlo (QMC). Piensa en esto no como un único cálculo, sino como un masivo juego de simulación donde miles de "caminantes" (fantasmas digitales de los electrones) corren por un mundo virtual, dando pasos aleatorios para explorar cada disposición posible. Al promediar dónde terminan estos caminantes, los científicos pueden obtener una imagen muy precisa de la energía y el comportamiento del sistema. Es un poco como intentar adivinar la altura promedio de todos en un estadio haciendo que unos pocos miles de personas deambulen por las gradas e informen de vuelta, en lugar de medir a cada una de las personas. ¿El problema? Este juego es increíblemente costoso computacionalmente. Requiere supercomputadoras para ejecutarse el tiempo suficiente para obtener una respuesta fiable, y durante mucho tiempo, el software que ejecutaba estos juegos fue construido para tipos de computadoras más antiguos y lentos.

Esto nos lleva a jQMC, una nueva pieza de software introducida por Kousuke Nakano y Michele Casula. El artículo presenta a jQMC como un motor moderno basado en Python, diseñado específicamente para ejecutarse en las últimas tarjetas gráficas (GPU) que se encuentran en las supercomputadoras más potentes de hoy en día. Los autores construyeron esta herramienta para ver si podían hacer que estas simulaciones de electrones corrieran significativamente más rápido y de manera más eficiente que el software actual que es el estándar de oro, el cual fue escrito en un lenguaje de programación más antiguo (Fortran) y no fue diseñado originalmente para estos nuevos chips.

El equipo descubrió que jQMC es, de hecho, un demonio de la velocidad cuando se combina con las GPU modernas. Para la parte de "Monte Carlo Variacional" (VMC) de la simulación —que es como la ronda de práctica donde los caminantes aprenden los rudimentos—, jQMC en una GPU fue hasta diez veces más rápido que el software tradicional ejecutándose en una CPU, e incluso más rápido que el software tradicional intentando ejecutarse en una GPU. Para la parte de "Monte Carlo de Difusión" (LRDMC), que es la ronda final de alto riesgo, jQMC seguía siendo aproximadamente tres veces más rápido en una GPU. Los investigadores también probaron qué tan bien jQMC podía calcular las "fuerzas" que empujan y tiran de los átomos (esencial para entender cómo se mueven y reaccionan las moléculas). Descubrieron que, para una molécula llamada etanol, los cálculos de fuerza de jQMC fueron increíblemente precisos, coincidiendo con la teoría de "clúster acoplado" (coupled-cluster) dentro de aproximadamente 1.5 kcal/mol/Å. Sin embargo, para una molécula diferente, la malonaldehído, los resultados fueron muy distintos al estándar de oro, lo que sugiere que incluso con esta nueva y poderosa herramienta, algunos rompecabezas químicos complejos todavía guardan sorpresas.

En resumen, jQMC es un experimento exitoso en la modernización de la química cuántica. Demuestra que, al reescribir las reglas del juego para adaptarse al hardware del futuro (GPUs), los científicos pueden ejecutar estas complejas simulaciones mucho más rápido, abriendo la puerta al estudio de sistemas más grandes y complicados que nunca. Si bien no resuelve todos los misterios de la química, proporciona una forma mucho más rápida y flexible de abordar los que ya conocemos.

Resumen Técnico de jQMC: Un paquete de Monte Carlo cuántica ab initio basado en JAX diseñado para la computación acelerada por GPU

Planteamiento del Problema
Resolver la ecuación de Schrödinger de muchos cuerpos sigue siendo un desafío central en la química y la física computacionales. Aunque la Teoría del Funcional de la Densidad (DFT) es ampliamente utilizada, su dependencia de funcionales de intercambio-correlación aproximados limita su precisión en sistemas fuertemente correlacionados. Los métodos de Monte Carlo cuántica (QMC) ab initio, como el Monte Carlo Variacional (VMC) y el Monte Carlo de Difusión (DMC), ofrecen una alternativa estocástica con un escalamiento favorable, pero son computacionalmente demandantes. El cambio de las instalaciones de Computación de Alto Rendimiento (HPC) hacia arquitecturas dominadas por GPU presenta un obstáculo significativo para los códigos QMC heredados, que a menudo están escritos en Fortran o C y dependen de diseños centrados en la CPU. Refactorizar estos códigos para explotar el paralelismo de GPU es difícil, y los códigos QMC acelerados por GPU existentes suelen carecer de la flexibilidad de los marcos de trabajo modernos o de las características algorítmicas específicas requeridas para simulaciones ab initio de alta precisión. Además, existe la necesidad de formatos de función de onda estandarizados y técnicas de optimización eficientes que puedan aprovechar la diferenciación automática.

Metodología
Los autores presentan jQMC, un paquete computacional basado en Python para simulaciones QMC ab initio construido enteramente sobre la biblioteca JAX. JAX permite la compilación Just-In-Time (JIT), la diferenciación automática (AD) y una portabilidad fluida entre los backends de CPU y GPU.

Algoritmos: jQMC implementa dos algoritmos de QMC de estado fundamental:
1. Monte Carlo Variacional (VMC): Utiliza el algoritmo de Metropolis-Hastings para muestrear configuraciones electrónicas. Emplea el esquema de reponderación de Attaccalite y Sorella (AS) para manejar las singularidades cerca de las superficies nodales.
2. Monte Carlo de Difusión con Regularización de Red (LRDMC): Una variante de Monte Carlo de Función de Green (GFMC) que discretiza el Hamiltoniano en una red para mejorar la estabilidad numérica. jQMC implementa esquemas de ramificación (branching) tanto convencionales como de balance de carga para mitigar el desequilibrio de carga en cálculos paralelos.
Ansatz de la Función de Onda: La función de onda de prueba se construye como un producto de un factor Jastrow y una parte antisimétrica. El componente antisimétrico soporta:
- Potencia Geminal Antisimetrizada (AGP): Una función de onda de tipo enlace de valencia resonante (RVB) que incluye apareamientos de singlete y triplete de espín.
- AGP con orbitales moleculares restringidos (AGPn): Una variante que restringe el rango de la matriz de apareamiento.
- Determinante de Slater Único (SD): Como un caso especial de AGP.
  La función de onda puede inicializarse desde cálculos externos de Hartree-Fock (HF) o de la Teoría del Funcional de la Densidad (DFT) a través de la biblioteca TREX-IO.
Optimación: El paquete implementa la Reconfiguración Estocástica (SR) y el Método Lineal (LM) para optimizar los parámetros variacionales. Estos métodos utilizan la métrica del espacio de parámetros para actualizar las amplitudes de la función de onda y las superficies nodales simultáneamente.
Diferenciación y Fuerzas: Aprovechando la diferenciación automática de JAX, jQMC calcula eficientemente las derivadas de energía. Implementa el cálculo de fuerzas atómicas tanto para VMC (usando la Transformación de Coordenadas de Space-Warp) como para LRDMC (usando la aproximación de Reynolds y la regularización de Pathak-Wagner).
Arquitectura: El código separa las estructuras de datos (implementadas como dataclasses compatibles con JAX) de los algoritmos. Utiliza vmap para la vectorización sobre los caminantes (walkers, el paralelismo de alto nivel en QMC) y soporta el escalado multi-GPU mediante MPI (mpi4py) y el particionamiento (sharding) de JAX. También soporta aritmética de precisión mixta (FP32 para evaluaciones de orbitales atómicos, FP64 para determinantes/geminales) para acelerar el cómputo sin sacrificar la precisión.

Contribuciones Clave

Implementación basada en JAX: jQMC es el primer paquete de QMC que aprovecha plenamente JAX tanto para la diferenciación algorítmica como para la aceleración por GPU, permitiendo un código compilado con JIT y portable entre arquitecturas de CPU y GPU.
Soporte Avanzado de Funciones de Onda: Proporciona una implementación robusta del ansatz Jastrow-Potencia Geminal Antisimetrizada (JAGP), que generaliza el determinante de Slater único y permite la optimización de las superficies nodales.
Vectorización Eficiente: Al vectorizar sobre los caminantes al nivel superior, jQMC logra una alta eficiencia intra-GPU, saturando los recursos de la GPU de manera más efectiva que los códigos tradicionales paralelizados en CPU.
Estandarización: La integración con TREX-IO asegura la compatibilidad con otros paquetes de estructura electrónica (por ejemplo, PySCF, CP2K).
Evaluación de Fuerzas: La implementación de fuerzas atómicas utilizando diferenciación automática y técnicas de regularización específicas (Pathak-Wagner) para LRDMC permite la optimización de la geometría y los cálculos de fonones dentro del marco de QMC.

Resultados

Pruebas de Rendimiento (Benchmarks):
- GPU vs. CPU: En GPUs NVIDIA (A100 y H100), jQMC demuestra aceleraciones de más de dos órdenes de magnitud en comparación con su propia implementación en CPU.
- Comparación con TurboRVB: Al compararse con TurboRVB (un código Fortran90 que implementa los mismos algoritmos y funciones de onda), jQMC en GPUs es significativamente más rápido. Para VMC, las ganancias en tiempo de pared (wall-time) alcanzan un orden de magnitud. Para LRDMC, la ganancia es más moderada pero sigue siendo sustancial (aprox. 3 veces más rápido para los tamaños de sistema probados).
- Escalabilidad: Las pruebas de escalamiento débil en los supercomputadores Miyabi (Universidad de Tokio) y Leonardo (CINECA) muestran que jQMC mantiene una eficiencia casi ideal (eficiencia $\approx$ 1.0) hasta 1024 GPUs para VMC y hasta 512 GPUs para LRDMC (usando el algoritmo de balance de carga).
Validación de Precisión:
- Reproducibilidad: Los resultados de jQMC para las energías totales y las energías de enlace de metano, agua y el dímero metano-agua concuerdan con otros códigos principales de QMC (CASINO, QMCPACK, CMQMC, TurboRVB) dentro de los márgenes de error estadístico, validando la implementación de los algoritmos DLA (Aproximación de Localidad del Determinante) y LRDMC.
- Fuerzas Atómicas: El artículo evalúa la precisión de las fuerzas atómicas para etanol y malonaldehído frente a cálculos de todo el electrón CCSD(T).
  - Para el etanol, las fuerzas de LRDMC logran un Error Absoluto Medio (MAE) de $\sim$ 1.5 kcal/mol/Å respecto a CCSD(T), consistente con la literatura previa.
  - Para el malonaldehído, se observa una discrepancia significativa, donde las fuerzas de LRDMC y CCSD(T) difieren en $\sim$ 5 kcal/mol/Å (MAE). Los autores señalan que esta discrepancia persiste incluso con superficies nodales optimizadas, lo que sugiere que el problema puede no deberse únicamente a la aproximación de la superficie nodal, sino que requiere mayor investigación.

Significado
El artículo posiciona a jQMC como una plataforma moderna y versátil que cierra la breancia entre los métodos QMC ab initio de alta precisión y las modernas arquitecturas de HPC aceleradas por GPU. Al utilizar JAX, el paquete ofrece un entorno flexible para el prototipado rápido de nuevos algoritmos, proporcionando al mismo tiempo un rendimiento de nivel de producción. Los autores afirman que jQcia logra abordar con éxito los desafíos de portar los códigos QMC heredados a las GPU, proporcionando una herramienta que no solo es más rápida que las implementaciones existentes basadas en Fortran en hardware de GPU, sino también más fácil de extender y mantener. La validación exitosa de fuerzas y energías frente a otros códigos y métodos de química cuántica de alto nivel demuestra que jQMC está listo para simulaciones de producción de sistemas moleculares y extendidos, particularmente en regímenes donde los efectos de correlación fuerte son críticos.

Más como este