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

    pip install jVMC

To install jVMC with GPU support, use

    pip install jVMC[gpu]

This is assuming, that jax[cuda12] will work for you (see the JAX documentation).

Option 2: Clone and pip-install for development

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

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

2. pip-install the package

    pip install -e ".[dev]"
   

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

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 SciPost Phys. Codebases 2 (2022)

    @Article{jVMC,
            title={{jVMC: Versatile and performant variational Monte Carlo leveraging automated differentiation and GPU acceleration}},
            author={Markus Schmitt and Moritz Reh},
            journal={SciPost Phys. Codebases},
            pages={2},
            year={2022},
            publisher={SciPost},
            doi={10.21468/SciPostPhysCodeb.2},
            url={https://scipost.org/10.21468/SciPostPhysCodeb.2},
    }