lagrange2art / vmc_jax

Impementation of Variational Monte Carlo (VMC) for quantum many-body dynamics using JAX.

jVMC

This is an impementation of Variational Monte Carlo (VMC) for quantum many-body dynamics using the JAX library (and Flax on top) to exploit the blessings of automatic differentiation for easy model composition and just-in-time compilation for execution on accelerators.

1. Documentation
2. Installation
3. Online example
4. Important gotchas
5. Citing jVMC

Please report bugs as well as other issues or suggestions on our issues page.

Documentation

Documentation is available here.

Installation

Option 1: pip-install

1. We recommend you create a new conda environment to work with jVMC:

    conda create -n jvmc python=3.8
    conda activate jvmc
   

2. pip-install the package

    pip install jVMC
   

Option 2: Clone and pip-install

1. Clone the jVMC repository and check out the development branch:

    git clone https://github.com/markusschmitt/vmc_jax.git
    cd vmc_jax
   

2. We recommend you create a new conda environment to work with jVMC:

    conda create -n jvmc python=3.8
    conda activate jvmc
   

3. pip-install the package

    pip install .  
   

   Alternatively, for development:

    pip install -e .[dev]
   

Test that everything worked, e.g. run 'python -c "import jVMC"' from a place different than vmc_jax.

Compiling JAX

How to compile JAX on a supercomputing cluster

Online example

Click on the badge above to open a notebook that implements an exemplary ground state search in Google Colab.

Important gotchas

Out-of-memory issues and batching

Memory requirements grow with increasing network sizes. To avoid out-of-memory issues, the batchSize parameter of the NQS class has to be adjusted. The batchSize indicates on how many input configurations the network is evaluated concurrently. Out-of-memory issues are usually resolved by reducing this number. The numChains parameter of the Sampler class for Markov Chain Monte Carlo sampling plays a similar role, but its optimal values in terms of computational speed are typically not memory critical.

Citing jVMC

If you use the jVMC package for your research, please cite our reference paper arXiv:2108.03409.