AGeDi

AGeDi (Atomistic Generative Diffusion) is a Python package for training and sampling diffusion models for atomistic structures. It is built around PyTorch, PyTorch Geometric, PyTorch Lightning, and ASE.

Documentation

AGeDi pronounced "A Jedi" is a library for Atomistic Generative Diffusion built on PyG, Lightning and ASE and offers customizable diffusion models for periodic atomistic material generation.

• Full docs: https://agedi.readthedocs.io
• CLI entrypoint: agedi
• Primary model backend today: PaiNN (via SchNetPack)

What AGeDi does

AGeDi provides:

• Data conversion from ASE Atoms to graph data (AtomsGraph)
• Training pipeline for diffusion models over positions and atom types
• Sampling pipeline from trained checkpoints (with optional templates)
• CLI and Python functional API for reproducible workflows

Installation

Minimal install:

pip install "agedi @ git+https://github.com/nronne/agedi.git"

This installs the core package only. For the current release, training and sampling require PaiNN via SchNetPack:

pip install "agedi[full] @ git+https://github.com/nronne/agedi.git"

For contributors:

pip install -e ".[test,full]"

Quickstart (CLI)

# Train (example: 3 hours, surface/slab system)
agedi train -t 3 --noisers ConfinedCellPositions --mask MaskFixed --confinement 2 10 PdO_training_data.traj

# Inspect saved hyperparameters
agedi inspect logs/version_0

# Sample structures
agedi sample logs/version_0 -f Pd2O2 --template_path template.traj --confinement 2 10

# Predict energies and forces (requires model trained with --force_field)
agedi predict logs/version_0 structures.traj

Quickstart (Python API)

from ase.io import read
from agedi import train_from_atoms, sample, AtomsGraph

data = read("PdO_training_data.traj", ":")
diffusion, dataset, trainer = train_from_atoms(
    data,
    noisers=("Positions",),
    style="surface",
    mask="MaskFixed",
    confinement=(2.0, 10.0),
    max_time=3,
)

template = AtomsGraph.from_atoms(read("template.traj"), confinement=(2.0, 10.0))
structures = sample(diffusion, n_samples=8, formula="Pd2O2", template=template)

To additionally train a force-field and run predictions:

from ase.io import read, write
from agedi import train_from_atoms, load_diffusion, predict

data = read("PdO_training_data.traj", ":")  # must contain forces and energy
diffusion, _, _ = train_from_atoms(data, noisers=("CellPositions",), force_field=True)

# Later, predict on new structures
diffusion = load_diffusion("logs/version_0")
predicted = predict(diffusion, read("structures.traj", index=":"))
write("predicted.traj", predicted)

Documentation map

The documentation has dedicated pages for:

• System overview and code architecture
• Installation and environment setup
• CLI and Python workflows
• End-to-end PdO tutorial
• Pitfalls and troubleshooting
• Publication references and citation text
• API reference (auto-generated)

References

• N. Rønne, A. Aspuru-Guzik, B. Hammer, Phys. Rev. B 110, 235427 (2024): https://doi.org/10.1103/PhysRevB.110.235427
• AGeDi preprint: https://arxiv.org/abs/2507.18314

Citation

If you use AGeDi in research, please cite the paper above and the AGeDi preprint.