DOCSEARCHbyAlgolia
Submit Docs

FermiNet

repository·main·Indexed 21 days ago

https://github.com/google-deepmind/ferminet

A JAX-based implementation of Fermionic Neural Networks used to learn highly accurate ground state wavefunctions of atoms and molecules via variational Monte Carlo. It supports training via CLI or Python scripts, inference for calculating energy and observables, and the calculation of excited states using NES-VMC or the Ensemble Penalty Method.

Tokens
3K
Snippets
9
Records
12
Agent score
73%

What's inside FermiNet

  1. Calculate Excited States

    main

    Excited state properties can be calculated by setting cfg.system.states = k (where k is the number of states).

    Two methods are supported:

    1. NES-VMC (Natural Excited States for VMC): The default method. It requires no additional parameter tuning.
    2. Ensemble Penalty Method: Enabled by setting cfg.optim.objective = 'vmc_overlap'. This method allows tuning weights for energies and the overlap penalty via cfg.optim.overlap. If weights are not provided, energies are automatically set to $1/k$.

    NES-VMC is generally more accurate, but the ensemble penalty method is provided for completeness.

  2. Install FermiNet

    main

    Install FermiNet and its dependencies using pip install -e .. It is recommended to use a virtual environment. If you have a GPU, you should install JAX with CUDA support. Ensure the jaxlib version matches your CUDA installation.

    To run tests, install the testing extras and use pytest.

    # Standard installation in a virtual environment
virtualenv ~/venv/ferminet
source ~/venv/ferminet/bin/activate
pip install -e .

# GPU installation (example for CUDA 11.0)
pip install --upgrade jax jaxlib==0.1.57+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

# Install for testing
pip install -e '.[testing]'
python -m pytest

  3. Run FermiNet via CLI

    main

    You can run FermiNet training from the command line using the ferminet executable or by calling ferminet/main.py. Configuration is managed via ml_collections.ConfigDict. You can specify a config file and override specific settings using flags with the --config.<key> <value> syntax.

    Example: Training a Li atom with a specific batch size and pretraining iterations.

    # Using the ferminet command
ferminet --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

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

  4. Perform Inference and accumulate observables

    main

    To run inference (calculating energy and observables with fixed parameters to get low-variance estimates), re-run your training command but set the optimizer to 'none' using --config.optim.optimizer 'none'.

    Ensure that cfg.log.save_path matches the original training run, or set cfg.log.restore_path to the original save_path.

    You can also enable additional observables at inference time by adding flags for:

    • --config.observables.s2 (Spin magnitude)
    • --config.observables.dipole (Dipole moments)
    • --config.observables.density (Density matrices)

  5. Extract and analyze excited state data from .npy files

    main

    During excited states calculations, FermiNet saves data to .npy files and logs certain observables to results.csv.

    Data Formats:

    • Density Matrix: Saved to a .npy file (even for ground state calculations).
    • Spin Magnitude and Dipole Moment: Logged directly to results.csv. For excited states calculations, only the total spin magnitude and dipole moments are logged to this file.

    To analyze the results, you can use utility functions to load the accumulated .npy files, compute the mean and variance, and demix the states using the energy matrix eigenvalues.

    import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm

# Example of loading data using the provided utility pattern
energy_mat_and_var = load_data("energy_matrix.npy", tail=1000)
s2_mat_and_var = load_data("s2_matrix.npy", tail=1000)
dipole_mat_and_var = load_data("dipole_matrix.npy", tail=1000)
density_mat_and_var = load_data("density_matrix.npy", tail=1000)

# Demix states and compute observables
energy_est, energy_std, observables, observable_std = get_results(
    energy_mat_and_var, 
    s2_mat_and_var, 
    dipole_mat_and_var, 
    density_mat_and_var
)

  6. Configure FermiNet with a custom Python script

    main

    To run custom training jobs, you can write a Python script that uses base_config.default() to initialize settings and train.train(cfg) to start the process. You can define the molecular system using ferminet.utils.system.Atom objects or by passing a PySCF Molecule object directly into cfg.system.pyscf_mol.

    import sys
from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train

# Optional: Print training progress to STDOUT
logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)

# Define system (e.g., 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)

  7. Understand FermiNet output files

    main

    After training or inference, the results directory contains:

    • train_stats.csv: Contains local energy and MCMC acceptance probability for each iteration.
    • checkpoints/: Directory containing training checkpoints.
    • .npy files: Saved when computing observables for excited states or the ground state density matrix. A single NumPy array is saved for every optimization iteration into the same file.

  8. Compute oscillator strength from transition dipole moments

    main

    To compute oscillator strengths from transition dipole moments, you must account for the transition between states. The oscillator strength $f$ can be calculated using the estimated energies and the transition dipole moments.

    Note: When working with the dipole matrix, you may need to multiply elementwise by its transpose to correctly recover off-diagonal terms for the calculation.

    # f = 2/3 * (E_i - E_j) * sum(|<i|mu|j>|^2)
# where dipole.transpose((1, 2, 0)) * dipole.transpose((2, 1, 0)) recovers the squared magnitude of the transition dipole

f = 2/3 * (energy_est[:, None] - energy_est[None, :]) * np.sum(
    dipole.transpose((1, 2, 0)) * dipole.transpose((2, 1, 0)), 
    axis=-1
)

  9. Load and process accumulated .npy data with load_data()

    main

    The load_data function is used to read arrays appended to a .npy file (common in iterative training/inference) and calculate their mean and variance.

    Arguments:

    • fname (str): Name of the file to load.
    • tail (int): If 0, loads all arrays in the file. If n != 0, loads only the last n arrays.

    Returns:

    • mean_data: The average value of the arrays.
    • var_data: The variance of the arrays.
    def load_data(fname, tail=0):
  """Load the arrays from a file and take the average and variance.

Args:
    fname: Name of file to load.
    tail: If 0, load all arrays in the file. If n !=0, load the last n arrays.

Returns:
    Average value of the arrays in the file and the variance.
  """
  # ... implementation ...
  return mean_data, var_data / n ** 2 - second_mom(mean_data) / n

  10. Demix states and compute observables with get_results()

    main

    The get_results function takes the energy matrix (and its variance) along with other observable matrices (and their variances) to perform state demixing. It uses the eigenvalues of the energy matrix to identify the states and then transforms the observables into the demixed basis.

    Arguments:

    • energy_matrix_and_var: A tuple containing the energy matrix and its variance.
    • *observable_matrices_and_var: Variable number of tuples, each containing an observable matrix and its variance (e.g., spin magnitude, dipole moments, density matrices).

    Returns:

    • energy: The sorted estimated energies.
    • energy_std: Standard deviation of the energies.
    • observables: A list of the demixed observable matrices.
    • observable_std: A list of the standard deviations for each demixed observable.
    def get_results(energy_matrix_and_var, *observable_matrices_and_var):
  """Given energy matrix and observable matrices, demix the states and compute the std deviations."""
  energy_mat, energy_var = energy_matrix_and_var
  energy, demix = np.linalg.eig(energy_mat)
  demix = demix[:, np.argsort(energy)]
  energy = np.sort(energy)
  energy_std = get_demixed_std(energy_matrix_and_var, demix)
  observables = [np.linalg.inv(demix) @ mv[0] @ demix for mv in observable_matrices_and_var]
  observable_std = [get_demixed_std(mv, demix) for mv in observable_matrices_and_var]
  return energy, energy_std, observables, observable_std

  11. Use the --config flag to specify experiment settings

    main

    The ferminet.main module defines a --config flag (via ml_collections.config_flags) which accepts a path to a configuration file. This file contains the hyperparameters and settings required for the FermiNet training workflow.

    --config PATH_TO_CONFIG_FILE

  12. Run FermiNet experiments via CLI

    main

    FermiNet is designed to be executed as a command-line tool. You can run experiments by providing a path to a configuration file using the --config flag. The entrypoint resolves the configuration using base_config.resolve and then initiates the training process via train.train(cfg).

    python -m ferminet.main --config /path/to/your_config.py