An implementation of the algorithm and experiments defined in "Ab-Initio Solution of the Many-Electron Schroedinger Equation with Deep Neural Networks", David Pfau, James S. Spencer, Alex G de G Matthews and W.M.C. Foulkes, Phys. Rev. Research 2, 033429 (2020). FermiNet is a neural network for learning the ground state wavefunctions of atoms and molecules using a variational Monte Carlo approach.
pip install -e . will install all required dependencies. This is best done
inside a virtual environment.
virtualenv -p python3.7 ~/venv/ferminet
source ~/venv/ferminet/bin/activate
pip install -e .
If you have a GPU available (highly recommended for fast training), then use
pip install -e '.[tensorflow-gpu]' to install TensorFlow with GPU support.
We use python 3.7 (or earlier) because there is TensorFlow 1.15 wheel available for it. TensorFlow 2 is not currently supported.
The tests are easiest run using pytest:
pip install pytest
python -m pytest
ferminet --batch_size 1024 --pretrain_iterations 100
will train FermiNet to find the ground-state wavefunction of the LiH molecule with a bond-length of 1.63999 angstroms using a batch size of 1024 MCMC configurations ("walkers" in variational Monte Carlo language), and 100 iterations of pretraining (the default of 1000 is overkill for such a small system). The system and hyperparameters can be controlled by flags. Run
ferminet --help
to see the available options. Several systems used in the FermiNet paper are
included by default. Other systems can easily be set up, by setting the
appropriate system flags to ferminet, modifying ferminet.utils.system or
writing a custom training script. For example, to run on the H2 molecule:
import sys
from absl import logging
from ferminet.utils import system
from ferminet import train
# Optional, for also printing training progress to STDOUT
logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)
# Define H2 molecule
molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]
train.train(
molecule=molecule,
spins=(1, 1),
batch_size=256,
pretrain_config=train.PretrainConfig(iterations=100),
logging_config=train.LoggingConfig(result_path='H2'),
)
train.train is controlled by a several lightweight config objects. Only
non-default settings need to be explicitly supplied. Please see the docstrings
for train.train and associated *Config classes for details.
Note: to train on larger atoms and molecules with large batch sizes, multi-GPU
parallelisation is essential. This is supported via TensorFlow's
MirroredStrategy
and the --multi_gpu flag.
The results directory contains pretrain_stats.csv, which contains the
pretraining loss for each iteration, train_stats.csv which contains the local
energy and MCMC acceptance probability for each iteration, and the checkpoints
directory, which contains the checkpoints generated during training. If
requested, there is also an HDF5 file, data.h5, which contains the walker
configuration, per-walker local energies and per-walker wavefunction values for
each iteration. Warning: this quickly becomes very large!
If you use this code in your work, please cite the associated paper:
@article{ferminet,
title={Ab-Initio Solution of the Many-Electron Schr{\"o}dinger Equation with Deep Neural Networks},
author={D. Pfau and J.S. Spencer and A.G. de G. Matthews and W.M.C. Foulkes},
journal={Phys. Rev. Research},
year={2020},
volume={2},
issue = {3},
pages={033429},
doi = {10.1103/PhysRevResearch.2.033429},
url = {https://link.aps.org/doi/10.1103/PhysRevResearch.2.033429}
}
This is not an official Google product.