jQMC is an ab initio quantum Monte Carlo (QMC) simulation package developed entirely from scratch using Python and JAX. Originally designed for molecular systems --with future extensions planned for periodic systems-- jQMC implements two well-established QMC algorithms: Variational Monte Carlo (VMC) and a robust and efficient variant of Diffusion Monte Carlo algorithm known as Lattice Regularized Diffusion Monte Carlo (LRDMC). By leveraging JAX just-in-time (jit) compilation and vectorized mapping (vmap) functionalities, jQMC achieves high-performance computations especially on GPUs while remaining portable across CPUs and GPUs. See here for the details of JAX. The jQMC users and developers manual is available from GitHub Pages.
What sets jQMC apart:
- It employs not only the standard Jatrow Slater determinant (JSD) wavefunction, but also the resonating valence bond (RVB)-type wave function, as known as Jastrow Antisymmetrized Geminal (JAGP) wavefunction, which captures correlation effects beyond the conventional JSD wave function used in many other QMC codes.
- Neural network Quantum States (NQSs) are implmented via
flaxmodule (a neural network library in theJAXecosystem). Curretnly, Jastrow Neural Network (JNN) is implemented. - It features a state-of-the-art optimization algorithm, stochastic reconfiguration, that enables stable optimization of both the amplitudes and nodal surfaces of many-body wave functions at the variational level.
- It implements the LRDMC method, providing a numerically stable approach to diffusion Monte Carlo calculations.
- The use of adjoint algorithmic differentiation in
JAXallows for efficient differentiation of many-body wave functions, facilitating the computation of atomic forces analytically. - Written in
Python, jQMC is designed to be user-friendly for executing simulations and easily extensible for developers implementing and testing new QMC methods. - By leveraging
JAXjust-in-time (jit) compilation and vectorized mapping (vmap) functionalities, the code achieves high-performance computations especially on GPUs while remaining portable across CPUs, GPUs, and TPUs. - MPI support enables the execution of large-scale computations on HPC facilities.
- Automated workflows: The
jqmc-workflowmodule automates the entire simulation pipeline — from pilot runs and step-count estimation through production runs and convergence monitoring — allowing users to obtain publication-quality results with minimal manual intervention. - To minimize bugs, the code is written in a loosely coupled manner and includes comprehensive unit tests and regression tests (managed by
pytest).
This combination of features makes jQMC a versatile and powerful tool for both users and developers in the field of quantum Monte Carlo simulations.
- On CPUs,
jQMCis slower than other QMC packages written in compiled languages (e.g., C++ or Fortran). On GPUs, however,jQMCachieves performance comparable to (or even faster than) compiled-language QMC codes, thanks toJAX's just-in-time compilation and hardware-level optimizations. Please use GPUs with a large number of walkers to fully exploit the performance. - Periodic boundary condition calculations are not supoorted yet. It will be implemented in the future as
JAXsupportscomplex128. Work in progress.
- Implementing periodic boundary conditions (PBC), at least at the Gamma point (i.e., using a real-space wave function).
Kosuke Nakano (National Institute for Materials Science (NIMS), Japan)
The release version of jQMC can be installed from PyPI via pip.
% pip install jqmcThe latest version of jQMC can be installed via pip from the cloned GitHub repository.
% git clone https://github.com/kousuke-nakano/jQMC
% cd jQMC
% pip install .Examples are in examples directory.
jQMC can prepare a trial (guiding) wavefunction from a TREX-IO file. Below is the list of HF/DFT packages that adopt TREX-IO for writing wave functions:
See the TREX-IO website for the detail.
jQMC user documentation is written using python sphinx. The source files are
stored in doc directory. Please see how to write the documentation at
doc/README.md.
main: main branch.devel*: development branches.rc: the latest stable version ready for deployment of the package.rc-gh-pages: the latest stable version ready for deployment of the documentation.
Every time a change is pushed to the main or devel* branch, the GitHub workflow launches the implemented unit and integration tests (jqmc-run-short-pytest.yml and jqmc-run-full-pytest.yml for the main and devel* branches, respectively).
Once the main branch is merged into the rc branch, the GitHub workflow launches the implemented unit and integration tests (jqmc-run-full-pytest.yml) and test a deployment using test-PyPI. Then, once a tag is attached to (the latest) commit in the rc branch, the GitHub workflow checks the tag format (PEP 440 with the starting v, e.g., v0.1.0b4, v0.1.1, v1.0) and deploy the package to PyPI.
Once the main branch is merged into the rc-gh-pages branch, the GitHub workflow launches the implemented documentaion building process (jqmc-deploy-gh-pages.yml) and deploy the compiled documentaiton to GitHub Pages.
Please see CONTRIBUTING.md for contribution guidelines.
Formatting rules are written in pyproject.toml.
Pre-commit (https://pre-commit.com/) is mainly used for applying the formatting rules automatically. Therefore, it is strongly encouraged to use it at or before git-commit. Pre-commit is set-up and used in the following way:
- Installed by
pip install pre-commit,conda install pre_commitor see https://pre-commit.com/#install. - pre-commit hook is installed by
pre-commit install. - pre-commit hook is run by
pre-commit run --all-files.
Unless running pre-commit, pre-commit.ci may push the fix at PR by github action. In this case, the fix should be merged by the contributor's repository.
-
Not strictly, but VSCode's
settings.jsonmay be written like below"ruff.lint.args": [ "--config=${workspaceFolder}/pyproject.toml", ], "[python]": { "editor.defaultFormatter": "charliermarsh.ruff", "editor.codeActionsOnSave": { "source.organizeImports": "explicit" } },
Tests are written using pytest. To run tests, pytest has to be installed. The tests can be run by
% pytest -s -v # with jax-jit
% pytest -s -v --disable-jit # without jax jit
% pytest -s -v --skip-heavy # skip heavy (slow) testsTo mark a test as heavy (skipped when --skip-heavy is passed), use the activate_if_skip_heavy marker:
@pytest.mark.activate_if_skip_heavy
def test_something_heavy():
...If you used jQMC in your reseach project, please cite the following articles. This indeed helps the jQMC project to continue:
-
"jQMC: JAX-based ab initio Quantum Monte Carlo package",
Kousuke Nakano and Michele Casula, in preparation (2025)
@article{jqmc, author = {Nakano, Kousuke and Casula, Michele}, title = {jQMC: JAX-based ab initio Quantum Monte Carlo package}, journal = {in preparation}, %volume = {}, %number = {}, %pages = {}, year = {2025}, %doi = {} } -
"Load-Balanced Diffusion Monte Carlo Method with Lattice Regularization",
K. Nakano, S. Sorella, and M. Casula, J. Chem. Phys. 163, 194117 (2025)
@article{10.1063/5.0296986, author = {Nakano, Kousuke and Sorella, Sandro and Casula, Michele}, title = {Load-balanced diffusion Monte Carlo method with lattice regularization}, journal = {J. Chem. Phys.}, volume = {163}, number = {19}, pages = {194117}, year = {2025}, month = {11}, doi = {10.1063/5.0296986} }