https://github.com/bytedance/LapNet
Efficient and Accurate Neural-Network Ansatz for Quantum Monte Carlo
https://github.com/bytedance/LapNet
Last synced: 5 days ago
JSON representation
Efficient and Accurate Neural-Network Ansatz for Quantum Monte Carlo
- Host: GitHub
- URL: https://github.com/bytedance/LapNet
- Owner: bytedance
- License: apache-2.0
- Created: 2023-11-13T08:19:53.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2024-12-04T05:37:32.000Z (4 months ago)
- Last Synced: 2025-04-10T02:42:30.146Z (6 days ago)
- Language: Python
- Size: 126 KB
- Stars: 62
- Watchers: 2
- Forks: 11
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
Awesome Lists containing this project
- best-of-atomistic-machine-learning - GitHub
README
# LapNet
A JAX implementation of the algorithm and calculations described in [A Computational Framework for Neural Network-based Variational Monte Carlo with Forward Laplacian](https://www.nature.com/articles/s42256-024-00794-x).
## Installation
To install LapNet together with all the dependencies, you need to have JAX and LapJAX preinstalled.
JAX can be installed by `pip install jax`.
To install JAX with CUDA support, use e.g.:
```shell
pip3 install --upgrade jax[cuda]==0.3.24 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```
Note that the jaxlib version must be compatible to the CUDA version
you use. Please see the
[JAX documentation](https://github.com/google/jax#installation) for more
details.
[LapJAX](https://github.com/YWolfeee/lapjax) is a efficient laplacian computation package described in our paper. It can be installed via:
```
git clone https://github.com/YWolfeee/lapjax.git
pip install ./lapjax
```
Once you have install both packages, go to the `lapnet` directory and run
```shell
pip install -e .
```
## Usage
LapNet uses the `ConfigDict` from
[ml_collections](https://github.com/google/ml_collections) to configure the
system. A few example scripts are included under `lapnet/configs/`. To estimate the ground state of atom C, simply run
```shell
python3 main.py --config lapnet/configs/atom.py --config.system.atom C
```
The system and hyperparameters can be controlled by modifying the config file or
(better, for one-off changes) using flags. Details of all available config settings are
in `lapnet/base_config.py`.
For example, to reproduce the result for the benzene dimer, use the following script:
```
python3 main.py --config lapnet/configs/benzene_dimer/benzene_dimer.py:4.95 --config.pretrain.iterations 50000 --config.pretrain.basis augccpvdz
```
There are some configuration files that contain multiple configurations. For example, `ferminet_systems_configs.py` contains multiple configurations used in the original ferminet paper. To choose a specific configuration, you need to specify the configuration through `config.system.molecule_name`. For example, to choose the configuration of CH4, one should use:
```
python3 main.py --config lapnet/configs/ferminet_systems_configs.py --config.system.molecule_name CH4
```
For the other config files with multiple configurations, one can check in that configuration file about how to specify different configurations in a command.
The Forward Laplacian option defaults to True. To check the result from original Laplacian calculation method, one can run script:
```
python3 main.py --config lapnet/configs/benzene_dimer/benzene_dimer.py:4.95 --config.pretrain.iterations 50000 --config.pretrain.basis augccpvd --config.optim.forward_laplacian=False
```
See the
[ml_collections](https://github.com/google/ml_collections)' documentation for
further details on the flag syntax.
Note: to train on large atoms and molecules with large batch sizes, multi-GPU
parallelisation is essential. This is supported via JAX's
[pmap](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap).
Multiple GPUs will be automatically detected and used if available.
## Output
The results directory contains the checkpoints generated during training and `train_stats.csv` which contains the local energy and MCMC acceptance probability for each iteration.
## Citation
If you find this repo useful, please cite our paper:
```
@article{li_computational_2024,
title = {A computational framework for neural network-based variational {Monte} {Carlo} with {Forward} {Laplacian}},
volume = {6},
issn = {2522-5839},
url = {https://doi.org/10.1038/s42256-024-00794-x},
doi = {10.1038/s42256-024-00794-x},
number = {2},
journal = {Nature Machine Intelligence},
author = {Li, Ruichen and Ye, Haotian and Jiang, Du and Wen, Xuelan and Wang, Chuwei and Li, Zhe and Li, Xiang and He, Di and Chen, Ji and Ren, Weiluo and Wang, Liwei},
month = feb,
year = {2024},
pages = {209--219},
}
```