jQMC

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 flax module (a neural network library in the JAX ecosystem). 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 JAX allows 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 JAX just-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-workflow module 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.

Known issues

• On CPUs, jQMC is slower than other QMC packages written in compiled languages (e.g., C++ or Fortran). On GPUs, however, jQMC achieves performance comparable to (or even faster than) compiled-language QMC codes, thanks to JAX'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 JAX supports complex128. Work in progress.

Ongoing works

• Implementing periodic boundary conditions (PBC), at least at the Gamma point (i.e., using a real-space wave function).

Developer(s)

Kosuke Nakano (National Institute for Materials Science (NIMS), Japan)

How to install jQMC

The release version of jQMC can be installed from PyPI via pip.

% pip install jqmc

The 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

Examples are in examples directory.

Supporting HF/DFT packages

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:

• Quantum Package
• PySCF
• FHI-aims
• CP2K
• Dirac
• pymolpro

See the TREX-IO website for the detail.

Documentation

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.

Branches

• 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).

How to deploy the package

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.

How to deploy the documentation

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.

Contribution

Please see CONTRIBUTING.md for contribution guidelines.

Formatting

Formatting rules are written in pyproject.toml.

Pre-commit

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_commit or 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.

VSCode setting

• Not strictly, but VSCode's settings.json may be written like below

  "ruff.lint.args": [
      "--config=${workspaceFolder}/pyproject.toml",
  ],
  "[python]": {
      "editor.defaultFormatter": "charliermarsh.ruff",
      "editor.codeActionsOnSave": {
          "source.organizeImports": "explicit"
      }
  },
  

How to run tests

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) tests

To 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():
    ...

Citation of jQMC

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}
  }