Amorphous Flow Matching (AFM)

Flow-based models for amorphous structures for paper: URL. Also includes diffusion-based variants for modelling amorphous structures.

Note: This project is based on Jasper Butcher's MChem thesis work carried out in 2024, Department of Chemistry, Oxford University.

Description

Amorphous Flow Matching (AFM) is an open-source generative model for amorphous structures. AFM is based on conditional flow matching (CFM) and uses the MACE architecture to learn how to generate samples from data.

Installation / Getting started

pip install -e .

To download the dataset:

python setup.py download_dataset

Note: since the optimal-transport algorithm uses the lapjv package, this repository does not suport MacOS. See here for more details.

If you want to log detailed metrics during your trajectories:

python setup.py precompute_metrics

and you should supply --VERBOSE_LOGGING to afm.core; this will log ring statistics, KL-divergences, defect frequencies, clash frequencies, maximum SOAP similarity. this is very much recommended for analysis, it will give you an indication of the quality of generated structures (especially to WandB) - however, note it may slow generation, or crash at, or above 10,000-atom runs.

Training models: To run a simple training run use:

python -m afm.train

This will train a model (small enough for the CPU) and save the model to results/run/model_DEFAULT.pt. You can sample from the model using:

python -m afm.core --model_weight_path='results/run/model_DEFAULT.pt'

On clusters you can run:

qsub main.sh

Just make sure to add your environment to the script as appropriate. To train a decent model you'll need around 90 epochs, as well as using NODEPLUS for a 4-interaction model. Also consider using optimized paths via solve_method='full' (or the faster 'rot') or supplying --AMORTISE_TRAINING to store optimized paths for future training runs (will take a while). NB: Optimized paths do not generate more training data, they simply find a better fit between $X_0$ and $X_1$. They generate examples of more $X_0$ (uniform noise) pairings to $X_1$ (actual data), which can be expensive if you want to make sure the pairing is optimal.

Pre-Trained models: Pre-trained model weights are also supplied in this repository:

tar -xzvf examples/models/production-model.tar.gz -C examples/models/

Where you now have to supply --model_config_key='NODEPLUS' to afm.core:

python -m afm.core --model_config_key='NODEPLUS' --model_weight_path='examples/models/production-model.pt'

Repository details

Everything diffusion and flow-based modelling for large material systems with periodic symmetry. This repository contains code for two primary methods: diffusion- and conditional flow based models. We found our flow-based appraoch performs generally better than our own diffusion model.

Annealed Restarts

One key component of Amorphous Flow Matching is the use of restart sampling to generate a diverse range of structures;

You can use restarts in the code by suppling --n_restarts=50 to the generation script (python -m afm.core). To control the graphitic character of generated structures you can use: --restart_sigma=0.3, higher values will generate more graphitic structures.

---

Repository Authors and References

If you use this codebase for your work, please consider citing:

@article{
    ...
}

If you run into any issues feel free to contact any of the developers of this repository directly, you can reach us via:

• Jasper Butcher (jbutch@uw.edu)
• Zakariya El-Machachi (zakariya.el.machachi@fu-berlin.de)

---