augment-atoms 0.1.0

Add your description here

augment-atoms

augment-atoms is a tool for augmenting datasets of atomic configurations via a model-driven, GPU-accelerated, rattle-relax-repeat procedure.

For each structure in the starting dataset, augment-atoms uses the provided potential energy surface (PES) model to generate a "family tree" of new structures. In the beginning, the tree consists of the single starting structure. To generate a new "child" structure, augment-atoms:

1. selects a "parent" structure from the tree,
2. rattles the atomic positions and unit cell,
3. relaxes using the PES model to get a new structure,
4. labels the child structure with the PES model, and
5. inserts the child structure into the tree.

For precise details of each of these steps, see the Details section below.

Installation

pip install augment-atoms

Usage

augment-atoms config.yaml

where config.yaml is a YAML file containing the following:

data:
  # an ase-readable file containing the starting structures
  input: input.xyz 

  # an ase-writeable path to append the new structures to
  output: output.xyz

config:
  # number of augmentations per starting structure
  n_per_structure: 10
  
  # the temperature
  T: 300  # units are Kelvin

  # the explore-vs-exploit trade-off (see below)
  beta: 0.5

  # the range of values from which to sample a 
  # standard deviation to rattle with at each step
  sigma_range: [0.01, 0.1]  # units are Å

  # the random seed to use (for reproducibility)
  seed: 42

  # the standard deviation of the cell perturbation
  # if null, no cell perturbation is applied
  cell_sigma: null  # units are Å
  
  # the units of the energies generated by the PES model
  units: eV

  # the maximum force magnitude to relax to
  max_force: 30  # units are (energy / Å)

  # the minimum separation between atoms to consider
  min_separation: 0.5  # units are Å

  # the maximum number of relaxations to perform per iteration
  max_relax_steps: 20

model:
  # the calculator to use to generate the PES model
  calculator: +lennard_jones()

In-built options for the calculator are:

• a Lennard-Jones calculator:

model:
  calculator: +lennard_jones()

• any model from the graph-pes package. If a GPU is available, it will be used to accelerate the PES model.

model:
  calculator:
    +graph_pes_calculator:
      path: path/to/model.pt

Alternatively, you are free to point to any instance of an ase.Calculator object. If you have my_function in my_file.py that returns an ase.Calculator object, you can use it as follows:

model:
  calculator: +my_file.my_function()

Details

1. Selecting a parent structure

To choose a new parent structure, we randomly sample from all structures in the tree, such that atom $i$ in structure $i$ has a probability of being picked given by

$$\mathbb{P}_i = \beta \cdot \frac{e^{-E_i / kT}}{\sum_j e^{-E_j / kT}} + (1-\beta) \cdot \frac{G_i}{\sum_j G_j}$$

where $E_i$ is the energy of structure $i$ and $G_i \in \mathbb{Z}^+$ is the `generation' of the structure, $k$ is the Boltzmann constant, $T$ is the temperature and $\beta \in [0, 1]$. Small values of $\beta$ favour the sampling of "younger" structures in the family tree, and hence a greater degree of exploration. Large values of $\beta$ favour the sampling of lower energy structures, and hence a denser sampling of the PES around energy minima.

2. Rattling the atomic positions and unit cell

To create a "child" from this parent structure, we perform the following transformation:

$$\begin{aligned} R^\prime &\leftarrow [(A + I) \times R] + B \ C^\prime &\leftarrow (A + I) \times C_0 \end{aligned}$$

where

• $R$ are the atomic positions
• $C_0$ is the unit cell of the original seed structure
• $A \in \mathbb{R}^{3\times 3}$ has entries sampled from $\mathcal{N}(0, \sigma_{A})$ where $\sigma_{A} \in \rm{sigma \_ range}$
• $B \in \mathbb{R}^{N \times 3}$ has entries sampled from $\mathcal{N}(0, \sigma_{B})$ where $\sigma_{B} \in [0, \rm{cell \_ sigma}]$

In the case of isolated structures, we only rattle the positions (i.e. $A = 0^{3 \times 3}$).

3. Relaxing the rattled child structure

To relax the rattled child structure, we use energies and forces generated by the PES model using a scheme inspired by the Robbins--Monro algorithm.

Step $x$ of this relaxation involves updating the atomic positions according to:

$$R^\prime \leftarrow R + \frac{\sigma_B}{x} \cdot \frac{F}{||F||}$$

where $F/||F||$ are the normalised unit vectors corresponding to the direction of each atomic force. We perform up to $M$ relaxations steps, but stop early with probability $\min(0.25, e^{-\Delta E / kT})$ providing the maximum force magnitude is less than config.max_force and where $\Delta E$ is the energy difference between the relaxed child and its starting parent structure. We reject all final structures that have any pair of atoms closer than config.min_separation Å.