augment-atoms

November 21, 2025 · View on GitHub

Test PyPI GitHub last commit License

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

This will install the augment-atoms command line tool (see pyproject.toml for the dependencies, requires Python 3.9+). Using uv is recommended, and will install augment-atoms with the correct dependencies in under 20 seconds starting from scratch.

There are no specific hardware requirements for augment-atoms. If a GPU is available, and the PES model supports it, the GPU will be used to accelerate structure generation. augment-atoms has been tested on both Linux and macOS.

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

  # the threshold for considering a structure too similar to the existing pool
  similarity_threshold: 0.1  # units are Å

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 ii in structure ii has a probability of being picked given by

Pi=β⋅e−Ei/kT∑je−Ej/kT+(1−β)⋅Gi∑jGj\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 EiE_i is the energy of structure ii and Gi∈Z+G_i \in \mathbb{Z}^+ is the `generation' of the structure, kk is the Boltzmann constant, TT is the temperature and β∈[0,1]\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:

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](https://en.wikipedia.org/wiki/Stochastic_approximation). 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` Å. ## Demo > This demo uses structures and a model taken from this repo's sister repository, found [here](https://github.com/dft-dutoit/synthetic-distillation). We include a stand-alone demo usage in the `demo` directory. This takes 3 water structures as input and uses a PaiNN model to generate and label 27 new structures, for a total of 30 structures. The `demo` directory has the following files: - `input.xyz` contains 3 starting water structures - `config.yaml` contains the configuration for the demo - `model.pt` is a PaiNN model trained on water structures from ... - `output.xyz` is the augmented dataset output. To run this demo yourself: ```bash # clone the repository git clone https://github.com/jla-gardner/augment-atoms.git cd augment-atoms/demo # remove the output file if it exists rm -rf output.xyz # run the demo augment-atoms config.yaml ``` This entire script took under 10 seconds on my M1 MacBook Pro. ## Citation If you use `augment-atoms` in your research, please cite the following pre-print: ```bibtex @misc{Gardner-25-06, title = {Distillation of Atomistic Foundation Models across Architectures and Chemical Domains}, author = {Gardner, John L. A. and du Toit, Daniel F. Thomas and Mahmoud, Chiheb Ben and Beaulieu, Zo{\'e} Faure and Juraskova, Veronika and Pa{\c s}ca, Laura-Bianca and Rosset, Louise A. M. and Duarte, Fernanda and Martelli, Fausto and Pickard, Chris J. and Deringer, Volker L.}, year = {2025}, number = {arXiv:2506.10956}, doi = {10.48550/arXiv.2506.10956}, } ```