HubLensTrendingbytedance/jaqmc
bytedance

jaqmc

AIJAXQuantum Monte CarloNeural NetworksPhysicsScientific Computing
View on GitHub
102

// summary

JaQMC is a modular framework built on JAX designed for performing neural network quantum Monte Carlo simulations. It utilizes deep neural networks as variational wavefunctions to solve the electronic Schrödinger equation with high accuracy. The platform supports diverse quantum systems including molecules, solids, and fractional quantum Hall states through swappable, independent components.

// technical analysis

JaQMC is a JAX-based framework designed for neural network quantum Monte Carlo simulations, enabling the solution of the electronic Schrödinger equation without reliance on traditional basis sets or density functionals. Its architecture is strictly modular, decoupling wavefunctions, samplers, estimators, and optimizers to allow for flexible experimentation with different physical models and loss functions. By leveraging JAX, the project provides high-performance features like automatic differentiation, JIT compilation, and multi-device parallelism, making it a robust tool for modeling complex quantum systems like molecules, solids, and fractional quantum Hall states.

// key highlights

01
Modular design allows users to independently swap wavefunctions, samplers, and optimizers without rewriting the core codebase.
02
Built-in support for diverse quantum systems including molecules, solid-state materials, and fractional quantum Hall effects.
03
Includes pre-implemented advanced architectures like FermiNet and PsiFormer for immediate research application.
04
Leverages JAX to provide automatic differentiation, JIT compilation, and seamless multi-device parallelism for high-performance computing.
05
Offers extensive configuration options via CLI or code, supporting various optimizers such as KFAC, SR, and Adam.
06
Provides specialized installation paths for GPU acceleration, including options for bundled CUDA libraries or local system installations.

// use cases

01
Solving the electronic Schrödinger equation for atoms, molecules, and crystals
02
Implementing advanced neural network architectures like FermiNet and PsiFormer
03
Conducting high-performance quantum simulations with JAX-accelerated optimization

// getting started

To begin, clone the repository and install the dependencies using uv or pip within a Python 3.12 environment. Once installed, you can execute simulations directly from the command line, such as running a hydrogen atom training task with 'jaqmc hydrogen-atom train'. Users can further customize these runs by modifying parameters via the CLI or by exploring the provided documentation for molecular and solid-state workflows.