FermiNet is a neural network for learning highly accurate ground state wavefunctions of atoms and molecules using a variational Monte Carlo approach.
This repository contains an implementation of the algorithm and experiments first described 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), along with subsequent research and developments.
WARNING: This is a research-level release of a JAX implementation and is under
active development. The original TensorFlow implementation can be found in the
tf branch.
pip install -e . will install all required dependencies. This is best done
inside a virtual environment.
virtualenv ~/venv/ferminet
source ~/venv/ferminet/bin/activate
pip install -e .If you have a GPU available (highly recommended for fast training), then you can install JAX with CUDA support, using e.g.:
pip install --upgrade jax jaxlib==0.1.57+cuda110 -f
https://storage.googleapis.com/jax-releases/jax_releases.htmlNote that the jaxlib version must correspond to the existing CUDA installation you wish to use. Please see the JAX documentation for more details.
The tests are easiest run using pytest:
pip install -e '.[testing]'
python -m pytestferminet uses the ConfigDict from
ml_collections to configure the
system. A few example scripts are included under ferminet/configs/. These are
mostly for testing so may need additional settings for a production-level
calculation.
ferminet --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100will train FermiNet to find the ground-state wavefunction of the Li atom 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
modifying the config file or (better, for one-off changes) using flags. See the
ml_collections' documentation for
further details on the flag syntax. Details of all available config settings are
in ferminet/base_config.py.
Other systems can easily be set up, by creating a new config file or ferminet,
or writing a custom training script. For example, to run on the H2 molecule, you
can create a config file containing:
from ferminet import base_config
from ferminet.utils import system
# Settings in a config files are loaded by executing the the get_config
# function.
def get_config():
# Get default options.
cfg = base_config.default()
# Set up molecule
cfg.system.electrons = (1,1)
cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]
# Set training hyperparameters
cfg.batch_size = 256
cfg.pretrain.iterations = 100
return cfgand then run it using
ferminet --config /path/to/h2_config.py
or equivalently write the following script (or execute it interactively):
import sys
from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train
# Optional, for also printing training progress to STDOUT.
# If running a script, you can also just use the --alsologtostderr flag.
logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)
# Define H2 molecule
cfg = base_config.default()
cfg.system.electrons = (1,1) # (alpha electrons, beta electrons)
cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]
# Set training parameters
cfg.batch_size = 256
cfg.pretrain.iterations = 100
train.train(cfg)Alternatively, you can directly pass in a PySCF 'Molecule'. You can create PySCF Molecules with the following:
from pyscf import gto
mol = gto.Mole()
mol.build(
atom = 'H 0 0 1; H 0 0 -1',
basis = 'sto-3g', unit='bohr')Once you have this molecule, you can pass it directly into the configuration by running
from ferminet import base_config
from ferminet import train
# Add H2 molecule
cfg = base_config.default()
cfg.system.pyscf_mol = mol
# Set training parameters
cfg.batch_size = 256
cfg.pretrain.iterations = 100
train.train(cfg)Note: to train on larger atoms and molecules with large batch sizes, multi-GPU parallelisation is essential. This is supported via JAX's pmap. Multiple GPUs will be automatically detected and used if available.
Excited state properties of systems can be calculated using the Natural Excited
States for VMC (NES-VMC) algorithm.
To enable the calculation of k states of a system, simply set
cfg.system.states=k in the config file.
The results directory contains 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. When
computing observables of excited states or the density matrix for the ground
state, .npy files are also saved to the same folder. A single NumPy array is
saved for every iteration of optimization into the same file. An example Colab
notebook analyzing these outputs is given in notebooks/excited_states_analysis.ipynb.
A collection of pretrained models trained with KFAC can be found on Google Cloud
here.
These are all systems from the original PRResearch paper: carbon and neon atoms,
and nitrogen, ethene, methylamine, ethanol and bicyclobutane molecules. Each
folder contains samples from the wavefunction in walkers.npy, parameters in
parameters.npz and geometries for the molecule in geometry.npz. To load the
models and evaluate the local energy, run:
import numpy as np
import jax
from functools import partial
from ferminet import networks, train
with open('params.npz', 'rb') as f:
params = dict(np.load(f, allow_pickle=True))
params = params['arr_0'].tolist()
with open('walkers.npy', 'rb') as f:
data = np.load(f)
with open('geometry.npz', 'rb') as f:
geometry = dict(np.load(f, allow_pickle=True))
signed_network = partial(networks.fermi_net, envelope_type='isotropic', full_det=False, **geometry)
# networks.fermi_net gives the sign/log of the wavefunction. We only care about the latter.
network = lambda p, x: signed_network(p, x)[1]
batch_network = jax.vmap(network, (None, 0), 0)
loss = train.make_loss(network, batch_network, geometry['atoms'], geometry['charges'], clip_local_energy=5.0)
print(loss(params, data)[0]) # For neon, should give -128.94165If you use this code in your work, please cite the associated papers. The initial paper details the architecture and results on a range of systems:
@article{pfau2020ferminet,
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}
}
and a NeurIPS Workshop Machine Learning and Physics paper describes the JAX implementation:
@misc{spencer2020better,
title={Better, Faster Fermionic Neural Networks},
author={James S. Spencer and David Pfau and Aleksandar Botev and W. M.C. Foulkes},
year={2020},
eprint={2011.07125},
archivePrefix={arXiv},
primaryClass={physics.comp-ph},
url={https://arxiv.org/abs/2011.07125}
}
The PsiFormer architecture is detailed in an ICLR 2023 paper, preprint reference:
@misc{vonglehn2022psiformer,
title={A Self-Attention Ansatz for Ab-initio Quantum Chemistry},
author={Ingrid von Glehn and James S Spencer and David Pfau},
year={2022},
eprint={2211.13672},
archivePrefix={arXiv},
primaryClass={physics.chem-ph},
url={https://arxiv.org/abs/2211.13672},
}
Periodic boundary conditions were originally introduced in a Physical Review Letters article:
@article{cassella2023discovering,
title={Discovering quantum phase transitions with fermionic neural networks},
author={Cassella, Gino and Sutterud, Halvard and Azadi, Sam and Drummond, ND and Pfau, David and Spencer, James S and Foulkes, W Matthew C},
journal={Physical review letters},
volume={130},
number={3},
pages={036401},
year={2023},
publisher={APS}
}
Wasserstein QMC (thanks to Kirill Neklyudov) is described in a NeurIPS 2023 article:
@article{neklyudov2023wasserstein,
title={Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schr{\"o}dinger Equation},
author={Neklyudov, Kirill and Nys, Jannes and Thiede, Luca and Carrasquilla, Juan and Liu, Qiang and Welling, Max and Makhzani, Alireza},
journal={NeurIPS},
year={2023}
}
Natural excited states was introduced in this article, which is also the first paper from our group using pseudopotentials
@article{pfau2023natural,
title={Natural Quantum Monte Carlo Computation of Excited States},
author={Pfau, David and Axelrod, Simon and Sutterud, Halvard and von Glehn, Ingrid and Spencer, James S},
journal={arXiv preprint arXiv:2308.16848},
year={2023}
}
This repository can be cited using:
@software{ferminet_github,
author = {James S. Spencer, David Pfau and FermiNet Contributors},
title = {{FermiNet}},
url = {http://github.com/deepmind/ferminet},
year = {2020},
}
This is not an official Google product.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent