📖 Full documentation: https://atomgptlab.github.io/alignn/

Table of Contents

• Introduction
• Installation
• Examples — train every model type
• Reproducing a JARVIS-Leaderboard contribution
• Colab notebooks
• Pre-trained models
• JARVIS-ALIGNN webapp
• ALIGNN-FF & ASE Calculator
• Inverse design — generate structures
• Peformances on a few datasets
• Useful notes
• References
• How to contribute
• Correspondence
• Funding support

ALIGNN & ALIGNN-FF (Introduction)

The Atomistic Line Graph Neural Network (paper) introduces a graph convolution layer that explicitly models both two- and three-body interactions in atomistic systems. The ALIGNN-FF variant (paper) extends this to a force-field for structurally and chemically diverse systems across 89 elements.

⚡ Pure PyTorch — DGL is no longer required. ALIGNN now runs fully in native PyTorch. Neighbor lists, line graphs, and batched readout are all built with plain torch tensor/scatter ops via alignn/torch_graph_builder.py, so you can train and run inference without installing DGL. To use the pure path, set the model name to the *_pure variant (e.g. alignn_atomwise_pure) and neighbor_strategy to "pure_torch" in your config. The example configs and tests in this repository already default to this pure-PyTorch path.

Installation

See docs/installation.md for conda, GitHub, and pip installation methods.

Examples — train every model type

All training recipes live on this page. Each one ships a self-contained, runnable example under alignn/examples/recipes/ with a make_toy_dataset.py (generates a tiny synthetic id_prop.json), a config_example.json, and its own detailed README.md. Every recipe below runs in ~1–2 minutes on CPU.

⚠️ The toy datasets are smoke tests, not real models. They are 40 rattled Si cells with synthetic labels, meant only to prove the pipeline runs. For a usable model, replace the structures/labels with real DFT data (thousands → millions of entries), raise epochs to 100–300 and batch_size to 32–64, and expect to use a GPU. See each recipe's README.

Every recipe reads an id_prop.json: a JSON list where each entry has a jid, an inline jarvis Atoms dict, and the target(s). See Dataset format for the full spec.

cd alignn/examples/recipes/knn
python make_toy_dataset.py                    # -> id_prop.json (40 toy entries)
train_alignn.py --root_dir . --config_name config_example.json \
    --output_dir toy_out --target_key target --id_key jid

cd alignn/examples/recipes/radius
python make_toy_dataset.py
train_alignn.py --root_dir . --config_name config_example.json \
    --output_dir toy_out --target_key target --id_key jid

cd alignn/examples/recipes/tensor
python make_toy_dataset.py
train_alignn.py --root_dir . --config_name config_example.json \
    --output_dir toy_out --target_key target --id_key jid

cd alignn/examples/recipes/spectra
python make_toy_dataset.py
train_alignn.py --root_dir . --config_name config_example.json \
    --output_dir toy_out --target_key target --id_key jid

cd alignn/examples/recipes/forcefield
python make_toy_dataset.py
train_alignn.py --root_dir . --config_name config_example.json \
    --output_dir toy_out --target_key energy_per_atom --force_key forces --id_key jid

cd alignn/examples/recipes/atomwise
python make_toy_dataset.py
train_alignn.py --root_dir . --config_name config_example.json \
    --output_dir toy_out --target_key target --id_key jid --atomwise_key charges

For the historical per-topic docs see also docs/training/ (dataset format, classification, multi-GPU).

Reproducing a JARVIS-Leaderboard contribution

Every ALIGNN entry on the JARVIS-Leaderboard ships the exact config, split, and run.sh used to produce it, so any result can be reproduced end to end:

# 1) install ALIGNN (pure-PyTorch, no DGL needed)
pip install alignn
#    or from source:
git clone https://github.com/atomgptlab/alignn.git
cd alignn && pip install -e . && cd ..

# 2) get the leaderboard (holds every contribution's config + data split + run.sh)
git clone https://github.com/atomgptlab/jarvis_leaderboard.git
cd jarvis_leaderboard
pip install -e .

# 3) pick a contribution and re-run it
#    contributions live under jarvis_leaderboard/contributions/<name>/
ls jarvis_leaderboard/contributions/alignn_model/
#    each folder has: the benchmark CSV, metadata.json, and run.sh
cat jarvis_leaderboard/contributions/alignn_model/run.sh
bash jarvis_leaderboard/contributions/alignn_model/run.sh

run.sh downloads the benchmark's train/val/test split (from the matching jarvis_leaderboard/benchmarks/.../*.json.zip), writes the id_prop/config, and calls train_alignn.py with the same settings that produced the leaderboard number — so you reproduce the published MAE exactly. To submit a new ALIGNN result, copy an existing contribution folder, drop in your predictions CSV + metadata.json, and open a PR (see the leaderboard's CONTRIBUTING).

Colab notebooks

Ready-to-run notebooks covering property prediction, force-field training, and pretrained-model usage. Click a badge to open in Colab.

Notebook	Open in Colab	Description
Regression task (graph-wise prediction)		Single-output regression for 2D-material exfoliation energies.
ML force-field training from scratch		Train an ALIGNN-FF force field for Silicon.
ALIGNN-FF: relaxation, EV curve, phonons, interfaces		Pretrained ALIGNN-FF for relaxation, EV curves, phonons, and interfaces.
Scaling / timing comparison		Scaling/timing analysis of universal MLFFs.
Melt-Quench MD		Generate amorphous structures via molecular dynamics.
Miscellaneous training tasks		Single-output, multi-output (phonon/electron DOS), classification, and pretrained usage.
Superconductor Tc		Train a model for superconductor transition temperature.
Build id_prop.json from VASP runs		Compile vasprun.xml files into id_prop.json for ALIGNN-FF training.
LAMMPS MD with ALIGNN-FF (pair_alignn)		Build LAMMPS with the native pair_alignn style and run NVE / melt-quench MD with the default ALIGNN-FF mps force field.

Using pre-trained models

See docs/pretrained/:

• Property predictor
• ALIGNN-FF

Web-apps

See docs/usage/webapps.md. Direct links: AtomGPT ALIGNN app, ALIGNN-FF app.

ALIGNN-FF ASE Calculator

from ase.build import bulk
from alignn.ff.unified_calculator import (
    AlignnUnifiedCalculator, AlignnUnifiedConfig)

cfg = AlignnUnifiedConfig(
    energy=True, forces=True, stress=True,
    properties=["formation_energy_peratom", "optb88vdw_bandgap"],
)
calc = AlignnUnifiedCalculator(cfg)          # models loaded once, reused

si = bulk("Si", "diamond", a=5.43); si.calc = calc
si.get_potential_energy(); si.get_forces(); si.get_stress()
print(calc.predictions())                    # extra property predictors

A single pydantic config selects the outputs (force-field energy/forces/stress plus any pretrained ALIGNN 2.0 property predictors — scalar, spectra, or D-dim tensor; radius or kNN graph). See docs/usage/ase-calculator.md for more, and the ASE docs page Calculators → ALIGNN.

Inverse design — generate structures

The same line graph also runs backwards. ALIGNN-CSP is a conditional diffusion model over the lattice and the fractional coordinates that takes a composition and a target property and returns a structure, then relaxes and ranks candidates with ALIGNN-FF. It is used like the ASE calculator: build it once, call it repeatedly.

from alignn.inverse.generate import ALIGNNGenerator

gen = ALIGNNGenerator()          # default model, downloaded + cached once

result = gen.generate("NbN", prop=16.0, num_candidates=8)   # Tc = 16 K
print(result.best)        # GeneratedStructure(NNb, 2 atoms, E=-17.91 eV/atom, relaxed=True)
print(result.to_poscar())
print(result.atoms)       # jarvis Atoms -> feeds any other ALIGNN model

For NbN this returns the rocksalt structure: N at (0.436, 0.691, 0.674) and Nb at (0.936, 0.191, 0.174), a (½,½,½) offset.

Composition takes a formula, a counts dict, or an explicit atom list; formula_units asks for a bigger cell; omitting prop leaves the property unconditioned. Every candidate stays on result.candidates, best first.

gen.generate("Nb3Sn", prop=18.0)                             # formula
gen.generate({"Mg": 1, "B": 2}, prop=39.0, formula_units=2)  # 6-atom cell
gen.generate(["Fe", "Fe", "O", "O", "O"])                    # no target value

from alignn.pretrained import list_alignn2_models
list_alignn2_models("generative")

ALIGNNGenerator(model="csp_supercon_jarvis_pt")   # pick a different one

Conditioning is pluggable rather than fixed to one property — a scalar (Tc, band gap), the composition, a 1-D signal such as an XRD pattern, or a 2-D map such as a STEM image. Each is dropped independently during training, so one checkpoint can generate from any subset at sampling time.

Two knobs matter for speed. num_candidates is nearly free, because cost scales with the number of denoising steps rather than the batch size, and it is the strongest lever on quality. num_steps trades fidelity for latency (~24 s at the trained 1000 steps, ~4.6 s at 200). Training your own model and reproducing the AtomBench benchmarks is covered in alignn/inverse/README.md.

That page also carries a full reference for an optional research extension, off by default: making bond angles an explicit denoising channel rather than only an input feature, with a line-graph topology that varies continuously as the coordinates denoise. See Explicit bond-angle diffusion — reference for the mathematics, the per-ablation breakdown, what each design choice takes from the literature, and the bibliography.

Performances

ALIGNN 2.0 benchmarked across single-property, multi-property (spectra / per-atom / tensor), and interatomic-force-field tasks — the latter including universal-potential stability on Matbench-Discovery / WBM (256,963 materials) and the CHIPS-FF material-property suite (lattice, elastic, vacancy, surface). Columns compare ALIGNN 2.0 on the radius and 8 Å kNN graphs against the original ALIGNN and CGCNN; bold marks the row best. Skill is 100 · (1 − MAE / MAD) vs the mean-absolute-deviation baseline. For the live, continually-updated numbers see the JARVIS-Leaderboard.

Useful notes

References

If ALIGNN or ALIGNN-FF contributed to your work, please cite the relevant papers.

How to contribute

See Contribution instructions and docs/contributing.md.

Correspondence

Please report bugs as GitHub issues or email drkamal@jhu.edu.

Funding support

• NIST-MGI
• NIST-CHIPS

Code of conduct

Please see Code of conduct.