google-deepmind/ferminet

google-deepmind/ferminet

Description: An implementation of the Fermionic Neural Network for ab-initio electronic structure calculations

Language: Python

License: Apache-2.0

Stars: 845

Forks: 165

Open issues: 3

Created: 2020-10-06T12:21:06Z

Pushed: 2026-05-21T12:30:41Z

Default branch: main

Fork: no

Archived: no

README:

FermiNet: Fermionic Neural Networks

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.

Installation

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.html

Note 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 pytest

Usage

ferminet 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 100

or

python3 ferminet/main.py --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

will 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 cfg

and 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.

Inference

After training, it is useful to run calculations of the energy and other observables over many time steps with the parameters fixed to accumulate low-variance estimates of physical quantities. To do this, just re-run the same command used for training with the flag --config.optim.optimizer 'none'. Make sure that either the value of cfg.log.save_path is the same, or that the value of cfg.log.restore_path is set to the value of cfg.log.save_path from the original training run.

It can also be useful to accumulate statistics about observables at inference time which were not included in the original training run. Spin magnitude, dipole moments and density matrices can be tracked by adding --config.observables.s2, --config.observables.dipole and --config.observables.density to the command line if they are not set to true in the config file.

Excited States

Excited state properties of systems can be calculated using either the Natural Excited States for VMC (NES-VMC) algorithm or an ensemble penalty method.…