QMCTorch: Molecular Wavefunctions with Neural Components for Energy and Force Calculations

TL;DR

QMCTorch is a GPU-native, PyTorch-based framework for quantum Monte Carlo simulations that integrates neural network components into wavefunctions, enabling efficient optimization and analysis of molecular energies and forces.

Contribution

It introduces a modular, extendable platform combining machine learning components with quantum chemistry for improved wavefunction optimization.

Findings

Accurately optimized wavefunctions for four molecules.

Computed dissociation energy curves matching baseline calculations.

Demonstrated the platform's flexibility for prototyping new wavefunctions.

Abstract

In this paper, we present results obtained using QMCTorch, a modular framework for real-space Quantum Monte Carlo (QMC) simulations of small molecular systems. Built on the popular deep learning library PyTorch, QMCTorch is GPU-native and enables the integration of machine learning-inspired components into the wave function ansatz, such as neural network backflow transformations and Jastrow factors, while leveraging efficient optimization algorithms. QMCTorch interfaces with two widely used quantum chemistry packages - PySCF and ADF - which provide initial values for the atomic orbital exponents and molecular orbital coefficients. In this study, we present wavefunction optimizations for four molecules: $H_2$, $LiH$, $Li_2$, and $CO$, using various wavefunction ans\"atze. We also compute their dissociation energy curves and the corresponding interatomic forces along these curves. Our results show good agreement with baseline calculations, recovering a significant portion of the correlation energy. QMCTorch provides a modular and extendable platform for rapidly prototyping new wavefunction ans\"atze, evaluating their performance, and analyzing optimization outcomes.