From 1899515fb2e40f98d23fa3ba58165a770ecce058 Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 12:02:28 -0400 Subject: [PATCH] docs --- docs/api.md | 62 +++++++++++++++ docs/contributing.md | 37 +++++++++ docs/index.md | 49 ++++++++++++ docs/installation.md | 81 +++++++++++++++++++ docs/notes.md | 43 ++++++++++ docs/performance.md | 92 ++++++++++++++++++++++ docs/pretrained/alignn-ff.md | 60 ++++++++++++++ docs/pretrained/index.md | 18 +++++ docs/pretrained/property-predictor.md | 72 +++++++++++++++++ docs/references.md | 44 +++++++++++ docs/requirements.txt | 3 + docs/training/classification.md | 47 +++++++++++ docs/training/dataset-format.md | 56 +++++++++++++ docs/training/force-field.md | 56 +++++++++++++ docs/training/multi-gpu.md | 51 ++++++++++++ docs/training/multi-output-regression.md | 45 +++++++++++ docs/training/single-output-regression.md | 55 +++++++++++++ docs/usage/ase-calculator.md | 95 +++++++++++++++++++++++ docs/usage/webapps.md | 23 ++++++ mkdocs.yml | 90 +++++++++++++++++++++ 20 files changed, 1079 insertions(+) create mode 100644 docs/api.md create mode 100644 docs/contributing.md create mode 100644 docs/index.md create mode 100644 docs/installation.md create mode 100644 docs/notes.md create mode 100644 docs/performance.md create mode 100644 docs/pretrained/alignn-ff.md create mode 100644 docs/pretrained/index.md create mode 100644 docs/pretrained/property-predictor.md create mode 100644 docs/references.md create mode 100644 docs/requirements.txt create mode 100644 docs/training/classification.md create mode 100644 docs/training/dataset-format.md create mode 100644 docs/training/force-field.md create mode 100644 docs/training/multi-gpu.md create mode 100644 docs/training/multi-output-regression.md create mode 100644 docs/training/single-output-regression.md create mode 100644 docs/usage/ase-calculator.md create mode 100644 docs/usage/webapps.md create mode 100644 mkdocs.yml diff --git a/docs/api.md b/docs/api.md new file mode 100644 index 0000000..c2e80da --- /dev/null +++ b/docs/api.md @@ -0,0 +1,62 @@ +# Package Overview + +High-level tour of the `alignn` package. For authoritative details, read the source — +this page is a map, not an API reference. + +## Top-level modules + +| Module | Role | +|---|---| +| `alignn.train_alignn` | Training CLI entry point | +| `alignn.train` | Training loop (called by the CLI) | +| `alignn.config` | Pydantic config schema for training / models | +| `alignn.data` | Dataset loaders — CSV/JSON index → torch datasets | +| `alignn.dataset` | Low-level dataset classes | +| `alignn.graphs` | Crystal graph + line graph construction | +| `alignn.lmdb_dataset` | LMDB-backed dataset for large-scale training | +| `alignn.pretrained` | Load and apply pretrained property predictors | +| `alignn.run_alignn_ff` | ALIGNN-FF CLI entry point | +| `alignn.cli` | Shared CLI argument parsing | +| `alignn.utils` | Misc helpers (logging, config loading, …) | +| `alignn.profiler` | Optional training profiler | + +## Models (`alignn.models`) + +| Module | Model | +|---|---| +| `alignn.models.alignn` | Original ALIGNN property predictor | +| `alignn.models.alignn_atomwise` | ALIGNN with atomwise outputs (forces/charges/mag) | +| `alignn.models.ealignn_atomwise` | Equivariant atomwise variant | +| `alignn.models.utils` | Shared layers and helpers | + +## Force-field (`alignn.ff`) + +| Module | Role | +|---|---| +| `alignn.ff.ff` | `AlignnAtomwiseCalculator`, `default_path`, training utils | +| `alignn.ff.calculators` | ASE calculator implementations | +| `alignn.ff.all_models_alignn.json` | Registry of property-predictor checkpoints | +| `alignn.ff.all_models_alignn_atomwise.json` | Registry of ALIGNN-FF checkpoints | + +Bundled pretrained checkpoints live in sub-directories of `alignn/ff/`, e.g. +`v10.30.2024_dft_3d_307k/`, `v12.2.2024_dft_3d_307k/`, +`v2024.12.12_dft_3d_multi_prop/`, `alignnff_wt01/`. + +## Examples & scripts + +- [`alignn/examples/`](https://github.com/atomgptlab/alignn/tree/main/alignn/examples) — + runnable sample datasets and configs (`sample_data`, `sample_data_ff`, + `sample_data_multi_prop`, …). +- [`alignn/scripts/`](https://github.com/atomgptlab/alignn/tree/main/alignn/scripts) — + high-throughput training scripts that download public datasets and train one model + per target. + +## CLIs at a glance + +```bash +train_alignn.py -h # training (regression, classification, atomwise, multi-output) +pretrained.py -h # apply a pretrained property predictor +run_alignn_ff.py -h # apply a pretrained ALIGNN-FF (optimize, EV curve, phonons, …) +``` + +All three are installed to your environment's `bin/` directory. diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 0000000..7431693 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,37 @@ +# Contributing + +Contributions are welcome — bug fixes, new features, benchmarks, and docs. + +## How to contribute + +1. Fork [atomgptlab/alignn](https://github.com/atomgptlab/alignn) and create a topic + branch off `develop`. +2. Install in development mode: + ```bash + git clone https://github.com//alignn + cd alignn + python -m pip install -e . + ``` +3. Make your change and add or update tests in `alignn/tests/`. +4. Run the test suite: + ```bash + pytest alignn/tests + ``` +5. Open a pull request targeting `develop`. + +Detailed instructions mirror the general JARVIS guide: +[Contribution.rst](https://github.com/atomgptlab/jarvis/blob/master/Contribution.rst). + +## Code of conduct + +See the JARVIS [Code of Conduct](https://github.com/atomgptlab/jarvis/blob/master/CODE_OF_CONDUCT.md). + +## Reporting issues + +Open a GitHub issue at . Please include: + +- The ALIGNN version (`pip show alignn`) +- PyTorch and DGL versions +- Operating system and GPU (if applicable) +- A minimal reproducer (config, a few structures, the command you ran) +- The full error traceback diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..933776b --- /dev/null +++ b/docs/index.md @@ -0,0 +1,49 @@ +# ALIGNN & ALIGNN-FF + +[![PyPI](https://badge.fury.io/py/alignn.svg)](https://badge.fury.io/py/alignn) +[![Build](https://github.com/atomgptlab/alignn/actions/workflows/main.yml/badge.svg)](https://github.com/atomgptlab/alignn/actions/workflows/main.yml) +[![codecov](https://codecov.io/gh/atomgptlab/alignn/branch/main/graph/badge.svg?token=S5X4OYC80V)](https://codecov.io/gh/atomgptlab/alignn) +[![Downloads](https://pepy.tech/badge/alignn)](https://pepy.tech/project/alignn) + +The **Atomistic Line Graph Neural Network (ALIGNN)** introduces a graph convolution layer +that explicitly models both two- and three-body interactions in atomistic systems. It +composes two edge-gated graph convolution layers: one applied to the atomistic line graph +*L(g)* (triplet interactions) and one to the atomistic bond graph *g* (pair interactions). + +**ALIGNN-FF** is a universal force-field built on ALIGNN. It was trained on the JARVIS-DFT +dataset (~75,000 materials and 4M+ energy/force entries) and supports any combination of +89 elements. Pretrained models can be fine-tuned or trained from scratch on new data. + +![ALIGNN layer schematic](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/schematic_lg.jpg?raw=true) + +## Highlights + +- Property prediction (regression & binary classification) +- Multi-output regression (e.g. energy + bandgap + DOS) +- Universal force-field (ALIGNN-FF) with ASE calculator +- Pretrained models on JARVIS, Materials Project, QM9, MOF datasets +- Multi-GPU training via `torchrun` (`DistributedDataParallel`) +- CLI entry points: `train_alignn.py`, `pretrained.py`, `run_alignn_ff.py` + +## Quick links + +- [Install ALIGNN](installation.md) +- [Train your first model](training/single-output-regression.md) +- [Use a pretrained model](pretrained/index.md) +- [Run the ALIGNN-FF ASE calculator](usage/ase-calculator.md) +- [Performance benchmarks](performance.md) + +## Citing + +If you use ALIGNN, please cite the relevant papers listed on the +[References](references.md) page. The primary references are: + +1. Choudhary, K., DeCost, B. *Atomistic Line Graph Neural Network for improved + materials property predictions.* **npj Comput Mater** 7, 185 (2021). +2. Choudhary, K., et al. *Unified graph neural network force-field for the periodic + table.* **Digital Discovery** (2023). + +## Correspondence + +Please open issues at or email +`drkamal@jhu.edu`. diff --git a/docs/installation.md b/docs/installation.md new file mode 100644 index 0000000..cd1d687 --- /dev/null +++ b/docs/installation.md @@ -0,0 +1,81 @@ +# Installation + +ALIGNN supports Linux, macOS, and Windows with Python 3.10. + +## Prerequisites + +Install Miniconda from . Pick the installer matching +your OS: + +```bash +bash Miniconda3-latest-Linux-x86_64.sh # Linux +bash Miniconda3-latest-MacOSX-x86_64.sh # macOS +``` + +On Windows, use the 64-bit Python 3.10 Miniconda installer. + +## Method 1 — conda (recommended) + +```bash +conda create --name my_alignn python=3.10 -y +conda activate my_alignn +conda install dgl=2.1.0 pytorch torchvision torchaudio pytorch-cuda -c pytorch -c nvidia +conda install alignn -y +``` + +## Method 2 — from GitHub (development install) + +```bash +conda create --name my_alignn python=3.10 -y +conda activate my_alignn +conda install dgl=2.1.0 pytorch torchvision torchaudio pytorch-cuda -c pytorch -c nvidia +git clone https://github.com/atomgptlab/alignn +cd alignn +python -m pip install -e . +``` + +## Method 3 — pip + +If you prefer pip, install DGL first from the wheel index that matches your CUDA/PyTorch +version (see ). + +PyTorch 2.1 + CUDA 12.1 (Windows/Linux): + +```bash +pip install -q dgl -f https://data.dgl.ai/wheels/torch-2.1/cu121/repo.html +pip install alignn +``` + +CPU only: + +```bash +pip install -q dgl -f https://data.dgl.ai/wheels/torch-2.1/repo.html +pip install alignn +``` + +## Verifying your install + +```bash +train_alignn.py -h +pretrained.py -h +run_alignn_ff.py -h +``` + +All three are Python executable scripts installed to your environment's `bin/` directory +— you do not need to provide an absolute path. + +## Common issues + +!!! warning "DGL + CUDA mismatches" + The most common install problem is a DGL build that does not match your CUDA version. + If you see import errors about `libtorch_cuda.so` or similar, reinstall DGL from the + wheel index matching your CUDA driver. + +- Use `batch_size` of 32 or 64 for real trainings (the examples ship with `batch_size: 2`). +- Complex `.cif` and `.pdb` files may require `cif2cell==2.0.0a3` and `pytraj` + respectively. +- `pandas >= 1.2.3` is required. +- From March 2024, `pytorch-ignite` is no longer required (removed for conda-forge + compatibility). + +See the [Notes & FAQ](notes.md) for more tips. diff --git a/docs/notes.md b/docs/notes.md new file mode 100644 index 0000000..efff406 --- /dev/null +++ b/docs/notes.md @@ -0,0 +1,43 @@ +# Notes & FAQ + +A grab-bag of tips collected from user issues and questions. + +## GPU / DGL + +- Install a DGL build that matches your CUDA runtime, e.g. `pip install dgl-cu111` for + CUDA 11.1. Mismatched builds are the most common install failure. +- If you see errors involving `libcudart.so` or `libtorch_cuda.so`, reinstall DGL from + the wheel index that matches your PyTorch CUDA version (see + [Installation](installation.md)). + +## Structure file parsing + +- Simple `.cif` and `.pdb` files are handled by `jarvis-tools` directly. +- For more complex CIFs, install `cif2cell==2.0.0a3`. +- For complex PDBs, install `pytraj` via `conda install -c ambermd pytraj`. + +## Training hyperparameters + +- The example `config_example.json` ships with `batch_size: 2` so the test suite runs + fast. **Use `batch_size: 32` or `64` for real trainings** — otherwise training will + be very slow and under-performing. +- `pandas >= 1.2.3` is required. +- Starting March 2024, `pytorch-ignite` is no longer a dependency (removed for + conda-forge build compatibility). + +## CLIs are importable scripts + +`train_alignn.py`, `pretrained.py`, and `run_alignn_ff.py` are installed as executables +in your environment's `bin/` directory. You do not need the absolute path — just run +them. + +## Known dataset issues + +- **QM9** results: see + [issue #54](https://github.com/atomgptlab/alignn/issues/54) for details on a data-split + discrepancy that affects reproducibility. + +## Getting help + +- File a GitHub issue: +- Email: `drkamal@jhu.edu` diff --git a/docs/performance.md b/docs/performance.md new file mode 100644 index 0000000..1db02be --- /dev/null +++ b/docs/performance.md @@ -0,0 +1,92 @@ +# Performance + +ALIGNN has been benchmarked across many public materials datasets. For the most +up-to-date numbers see [JARVIS-Leaderboard](https://pages.nist.gov/jarvis_leaderboard/). + +## JARVIS-DFT 2021 — classification + +| Classifier | Threshold | ALIGNN AUC | +|---|---|---| +| Metal / non-metal (OPT) | 0.01 eV | 0.92 | +| Metal / non-metal (MBJ) | 0.01 eV | 0.92 | +| Magnetic / non-magnetic | 0.05 µB | 0.91 | +| High / low SLME | 10 % | 0.83 | +| High / low spillage | 0.1 | 0.80 | +| Stable / unstable (ehull) | 0.1 eV | 0.94 | +| High / low n-Seebeck | -100 µV K⁻¹ | 0.88 | +| High / low p-Seebeck | 100 µV K⁻¹ | 0.92 | +| High / low n-PF | 1000 µW (mK²)⁻¹ | 0.74 | +| High / low p-PF | 1000 µW (mK²)⁻¹ | 0.74 | + +## JARVIS-DFT 2021 — regression (MAE) + +| Property | Units | MAD | CFID | CGCNN | ALIGNN | MAD:MAE | +|---|---|---|---|---|---|---| +| Formation energy | eV/atom | 0.86 | 0.14 | 0.063 | **0.033** | 26.06 | +| Bandgap (OPT) | eV | 0.99 | 0.30 | 0.20 | **0.14** | 7.07 | +| Total energy | eV/atom | 1.78 | 0.24 | 0.078 | **0.037** | 48.11 | +| Ehull | eV | 1.14 | 0.22 | 0.17 | **0.076** | 15.00 | +| Bandgap (MBJ) | eV | 1.79 | 0.53 | 0.41 | **0.31** | 5.77 | +| Bulk modulus K_v | GPa | 52.80 | 14.12 | 14.47 | **10.40** | 5.08 | +| Shear modulus G_v | GPa | 27.16 | 11.98 | 11.75 | **9.48** | 2.86 | +| Magnetic moment | µB | 1.27 | 0.45 | 0.37 | **0.26** | 4.88 | +| SLME (%) | — | 10.93 | 6.22 | 5.66 | **4.52** | 2.42 | +| Spillage | — | 0.52 | 0.39 | 0.40 | **0.35** | 1.49 | +| ε (DFPT: elec + ionic) | — | 45.81 | 43.71 | 38.78 | **28.15** | 1.63 | +| Max. piezo dij | C N⁻¹ | 24.57 | 36.41 | 34.71 | **20.57** | 1.19 | +| Exfoliation energy | meV/atom | 62.63 | 63.31 | 50.0 | **51.42** | 1.22 | + +(Full table in the repository README — trimmed here for readability.) + +## Materials Project 2018 + +| Property | Unit | MAD | CGCNN | MEGNet | SchNet | ALIGNN | +|---|---|---|---|---|---|---| +| Formation energy | eV/atom | 0.93 | 0.039 | 0.028 | 0.035 | **0.022** | +| Bandgap | eV | 1.35 | 0.388 | 0.33 | — | **0.218** | + +## QM9 — MAE + +| Target | Units | SchNet | MEGNet | DimeNet++ | ALIGNN | +|---|---|---|---|---|---| +| HOMO | eV | 0.041 | 0.043 | 0.0246 | **0.0214** | +| LUMO | eV | 0.034 | 0.044 | 0.0195 | **0.0195** | +| Gap | eV | 0.063 | 0.066 | **0.0326** | 0.0381 | +| µ | Debye | 0.033 | 0.050 | 0.0297 | **0.0146** | +| ZPVE | eV | 0.0017 | 0.00143 | **0.00121** | 0.0031 | + +!!! info "QM9 caveat" + See [issue #54](https://github.com/atomgptlab/alignn/issues/54) for known + discrepancies related to the QM9 split. + +## hMOF — regression + +| Property | Unit | MAE | R² | +|---|---|---|---| +| Gravimetric surface area | m² g⁻¹ | 91.15 | 0.99 | +| Volumetric surface area | m² cm⁻³ | 107.81 | 0.91 | +| Void fraction | — | 0.017 | 0.98 | +| LCD | Å | 0.75 | 0.83 | +| PLD | Å | 0.92 | 0.78 | +| CO₂ adsorption (all) | mol kg⁻¹ | 0.18 | 0.95 | + +## qMOF + +MAE on electronic bandgap: **0.20 eV**. + +## Open Catalyst — IS2RE 10k + +| Model | CGCNN | DimeNet | SchNet | DimeNet++ | ALIGNN | +|---|---|---|---|---|---| +| 10k | 0.988 | 1.0117 | 1.059 | 0.8837 | **0.61** | + +## Coming soon + +OMDB, HOPV, QETB. + +--- + +Claims of *best* performance should be verified on the latest +[JARVIS-Leaderboard](https://pages.nist.gov/jarvis_leaderboard/). Numbers from models +other than ALIGNN are reported as-published by the original authors and are not +necessarily reproduced in-house. diff --git a/docs/pretrained/alignn-ff.md b/docs/pretrained/alignn-ff.md new file mode 100644 index 0000000..680afcb --- /dev/null +++ b/docs/pretrained/alignn-ff.md @@ -0,0 +1,60 @@ +# Pre-trained ALIGNN-FF + +ALIGNN-FF is a universal graph neural network force-field. Several pretrained versions +are shipped with the package under +[`alignn/ff/`](https://github.com/atomgptlab/alignn/tree/main/alignn/ff) and on +[Figshare](https://figshare.com/projects/ALIGNN_models/126478). + +## CLI + +```bash +run_alignn_ff.py -h +``` + +Common tasks: + +```bash +# Single-point energy on the unrelaxed structure +run_alignn_ff.py --file_path POSCAR --task="unrelaxed_energy" + +# Structure optimization +run_alignn_ff.py --file_path POSCAR --task="optimize" + +# Energy vs. volume curve +run_alignn_ff.py --file_path POSCAR --task="ev_curve" +``` + +Additional tasks supported by the CLI include phonon calculations, interface +gamma-surface scans, and simple MD — run `-h` to see the full set. + +## Available checkpoints + +Model folders bundled with the package: + +- `v10.30.2024_dft_3d_307k` — trained on 307k JARVIS-DFT configurations (Oct 2024) +- `v12.2.2024_dft_3d_307k` — updated December 2024 release +- `v2024.12.12_dft_3d_multi_prop` — multi-property (energy + atomwise targets) +- `alignnff_wt01` — weighted loss variant +- `alex_band_gap`, `jv_mbj_bandgap_alignn` — bandgap-oriented models + +Additional models are listed in +[`all_models_alignn_atomwise.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/ff/all_models_alignn_atomwise.json). + +## Default path helper + +```python +from alignn.ff.ff import default_path +print(default_path()) # directory containing the default checkpoint +``` + +## Using it programmatically + +See the [ASE calculator](../usage/ase-calculator.md) page for a full worked example +covering relaxation and an EV curve. + +## See also + +- [Train / fine-tune a force-field](../training/force-field.md) +- [Property predictors](property-predictor.md) +- [Colab: ALIGNN-FF relaxation, EV curve, phonons, interfaces](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) +- [Colab: Melt-quench MD](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) diff --git a/docs/pretrained/index.md b/docs/pretrained/index.md new file mode 100644 index 0000000..e1acde1 --- /dev/null +++ b/docs/pretrained/index.md @@ -0,0 +1,18 @@ +# Pre-trained Models + +ALIGNN ships a large catalog of pre-trained checkpoints for direct use. + +- Models are hosted on [Figshare](https://figshare.com/projects/ALIGNN_models/126478). +- Two CLIs load them: `pretrained.py` (single-property predictors) and + `run_alignn_ff.py` (universal force-field). +- Model names and sources are listed in + [`alignn/ff/all_models_alignn.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/ff/all_models_alignn.json) and + [`alignn/ff/all_models_alignn_atomwise.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/ff/all_models_alignn_atomwise.json). + +| Page | What it covers | +|---|---| +| [Property predictors](property-predictor.md) | Per-property checkpoints (formation energy, bandgap, …) | +| [ALIGNN-FF](alignn-ff.md) | Universal force-field: EV curve, optimization, phonons, MD | +| [ASE calculator](../usage/ase-calculator.md) | Drive ALIGNN-FF from Python via ASE | + +See [Performance](../performance.md) for per-dataset accuracy numbers. diff --git a/docs/pretrained/property-predictor.md b/docs/pretrained/property-predictor.md new file mode 100644 index 0000000..84905ea --- /dev/null +++ b/docs/pretrained/property-predictor.md @@ -0,0 +1,72 @@ +# Property Predictors + +Single-property ALIGNN checkpoints are invoked through +[`pretrained.py`](https://github.com/atomgptlab/alignn/blob/develop/alignn/pretrained.py). + +## CLI help + +```bash +pretrained.py -h +``` + +## Example — formation energy (JARVIS-DFT) + +```bash +pretrained.py \ + --model_name jv_formation_energy_peratom_alignn \ + --file_format poscar \ + --file_path alignn/examples/sample_data/POSCAR-JVASP-10.vasp +``` + +## Available models + +A non-exhaustive list — names match the `--model_name` flag: + +| Target | Model name | +|---|---| +| Formation energy per atom (JARVIS-DFT) | `jv_formation_energy_peratom_alignn` | +| Bandgap — OPT88vdW | `jv_optb88vdw_bandgap_alignn` | +| Bandgap — MBJ | `jv_mbj_bandgap_alignn` | +| Total energy per atom (JARVIS-DFT) | `jv_optb88vdw_total_energy_alignn` | +| Ehull | `jv_ehull_alignn` | +| Bulk modulus (K_v) | `jv_bulk_modulus_kv_alignn` | +| Shear modulus (G_v) | `jv_shear_modulus_gv_alignn` | +| Dielectric (εx, OPT) | `jv_epsx_alignn` | +| Max. piezo dielectric (DFPT) | `jv_dfpt_piezo_max_dielectric_alignn` | +| Spillage | `jv_spillage_alignn` | +| SLME | `jv_slme_alignn` | +| Magnetic moment | `jv_magmom_oszicar_alignn` | +| Raman | `jv_raman_alignn` | +| Superconductor T_c | `jv_supercon_tc_alignn` | +| Interface CBM / VBM | `intermat_cbm`, `intermat_vbm` | +| hMOF CO₂ adsorption | `hmof_co2_absp_alignn` | + +See +[`all_models_alignn.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/ff/all_models_alignn.json) +for the full machine-readable list. + +## File formats + +Pass `--file_format` matching your structure file: + +- `poscar` — VASP POSCAR +- `cif` +- `xyz` +- `pdb` + +## Using from Python + +```python +from alignn.pretrained import get_prediction + +prediction = get_prediction( + model_name="jv_formation_energy_peratom_alignn", + atoms=my_jarvis_atoms, # jarvis.core.atoms.Atoms +) +print(prediction) +``` + +## See also + +- [ALIGNN-FF pretrained models](alignn-ff.md) +- [Performance](../performance.md) diff --git a/docs/references.md b/docs/references.md new file mode 100644 index 0000000..6f27b93 --- /dev/null +++ b/docs/references.md @@ -0,0 +1,44 @@ +# References + +If ALIGNN or ALIGNN-FF contributed to your work, please cite the relevant papers below. + +## Core + +1. Choudhary, K. & DeCost, B. **Atomistic Line Graph Neural Network for improved + materials property predictions.** *npj Computational Materials* 7, 185 (2021). + [Link](https://www.nature.com/articles/s41524-021-00650-1) +2. Choudhary, K., DeCost, B., Major, L., Butler, K., Thiyagalingam, J., Tavazza, F. + **Unified graph neural network force-field for the periodic table.** + *Digital Discovery* (2023). + [Link](https://pubs.rsc.org/en/content/articlehtml/2023/dd/d2dd00096b) + +## Applications + +3. **Prediction of the Electron Density of States for Crystalline Compounds with ALIGNN.** + [Link](https://link.springer.com/article/10.1007/s11837-022-05199-y) +4. **Recent advances and applications of deep learning methods in materials science.** + [Link](https://www.nature.com/articles/s41524-022-00734-6) +5. **Designing High-Tc Superconductors with BCS-inspired Screening, DFT, and + Deep-learning.** [Link](https://arxiv.org/abs/2205.00060) +6. **A Deep-learning Model for Fast Prediction of Vacancy Formation in Diverse Materials.** + [Link](https://arxiv.org/abs/2205.08366) +7. **Graph neural network predictions of MOF CO₂ adsorption properties.** + [Link](https://www.sciencedirect.com/science/article/pii/S092702562200163X) +8. **Rapid Prediction of Phonon Structure and Properties using ALIGNN.** + [Link](https://journals.aps.org/prmaterials/abstract/10.1103/PhysRevMaterials.7.023803) +9. **Large Scale Benchmark of Materials Design Methods.** + [Link](https://www.nature.com/articles/s41524-024-01259-w) +10. **Prediction of Magnetic Properties in van der Waals Magnets using GNNs.** + [Link](https://doi.org/10.1103/PhysRevMaterials.8.114002) +11. **CHIPS-FF: Benchmarking universal force-fields.** + [Link](https://github.com/atomgptlab/chipsff) + +## Full publication list + +A complete list of ALIGNN-related publications is maintained at +[jarvis-tools publications](https://jarvis-tools.readthedocs.io/en/master/publications.html). + +## Funding + +- [NIST-MGI](https://www.nist.gov/mgi) +- [NIST-CHIPS](https://www.nist.gov/chips) diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..dac2c55 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,3 @@ +mkdocs>=1.5,<2.0 +mkdocs-material>=9.5 +pymdown-extensions>=10.0 diff --git a/docs/training/classification.md b/docs/training/classification.md new file mode 100644 index 0000000..be50281 --- /dev/null +++ b/docs/training/classification.md @@ -0,0 +1,47 @@ +# Binary Classification + +Convert a regression dataset into binary labels and train a classifier — for example, +**metal vs. non-metal** based on a bandgap threshold. + +## Command + +```bash +train_alignn.py \ + --root_dir "alignn/examples/sample_data" \ + --classification_threshold 0.01 \ + --config "alignn/examples/sample_data/config_example.json" \ + --output_dir=temp +``` + +Values in `id_prop.csv` are thresholded: `1` if above `--classification_threshold`, `0` +otherwise. + +## Supported tasks + +!!! note "Binary only" + The current training script supports **binary classification only**. For multi-class + problems, open a GitHub issue — it is being considered. + +## Typical classification targets (JARVIS-DFT) + +| Classifier | Threshold | +|---|---| +| Metal / non-metal (OPT / MBJ bandgap) | 0.01 eV | +| Magnetic / non-magnetic | 0.05 µB | +| Stable / unstable (ehull) | 0.1 eV | +| High / low SLME | 10 % | +| High / low spillage | 0.1 | +| High / low Seebeck | ±100 µV K⁻¹ | +| High / low power factor | 1000 µW (mK²)⁻¹ | + +See the [Performance](../performance.md) page for accuracy numbers. + +## Metrics + +The training loop reports AUC and accuracy on the validation set and writes +`prediction_results_test_set.csv` with per-structure predicted probabilities. + +## See also + +- [Single-output regression](single-output-regression.md) +- [Colab: miscellaneous tasks (incl. classification)](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) diff --git a/docs/training/dataset-format.md b/docs/training/dataset-format.md new file mode 100644 index 0000000..6828b89 --- /dev/null +++ b/docs/training/dataset-format.md @@ -0,0 +1,56 @@ +# Dataset Format + +All ALIGNN training entry points read a directory containing **structure files** plus an +index file (`id_prop.csv` for property prediction, `id_prop.json` for force-field +training). + +Supported structure formats: **POSCAR**, **CIF**, **XYZ**, **PDB**. + +## Directory layout + +``` +my_dataset/ +├── POSCAR-1 +├── POSCAR-2 +├── ... +└── id_prop.csv +``` + +## id_prop.csv (property prediction) + +Filename and target value(s) per row, no header: + +```csv +POSCAR-1,1.234 +POSCAR-2,0.567 +``` + +For multi-output regression, append additional columns per target. + +## id_prop.json (force-field training) + +JSON list of entries containing `jid`, `energy` (per atom), `forces`, and `stress`. See +`alignn/examples/sample_data_ff/id_prop.json` for a concrete example. + +A Colab notebook for compiling VASP `vasprun.xml` files into `id_prop.json` is available +[here](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb). + +## Train/val/test split + +By default the dataset is split **80/10/10** (set by `train_ratio`, `val_ratio`, +`test_ratio` in the config). + +To control the split manually, set `n_train`, `n_val`, `n_test` and +`keep_data_order: true` in the config — this disables the random shuffle so you can train +on one set and validate/test on another. + +## Configuration file + +All hyperparameters live in a JSON config. Start from the examples: + +- [`config_example.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/config_example.json) — property prediction +- [`config_example_atomwise.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data_ff/config_example_atomwise.json) — force-field + +!!! tip + The example configs ship with `batch_size: 2` so the tests run fast. Bump to **32** + or **64** for real trainings, otherwise training will be slow and under-performing. diff --git a/docs/training/force-field.md b/docs/training/force-field.md new file mode 100644 index 0000000..cb36f6f --- /dev/null +++ b/docs/training/force-field.md @@ -0,0 +1,56 @@ +# Force-Field / MLIP Training + +ALIGNN-FF is trained with the same `train_alignn.py` script but using the +`atomwise_alignn` model. Instead of `id_prop.csv`, the script reads `id_prop.json` with +entries for `jid`, `energy`, `forces`, and `stress`. + +!!! warning "Energy is per atom" + Energy values in `id_prop.json` must be stored **per atom**, not per structure. + +## Train from scratch + +```bash +train_alignn.py \ + --root_dir "alignn/examples/sample_data_ff" \ + --config "alignn/examples/sample_data_ff/config_example_atomwise.json" \ + --output_dir="temp" +``` + +## Fine-tune a pretrained ALIGNN-FF + +Pass `--restart_model_path` with a checkpoint that uses the same model configuration: + +```bash +train_alignn.py \ + --root_dir "alignn/examples/sample_data_ff" \ + --restart_model_path "temp/best_model.pt" \ + --config "alignn/examples/sample_data_ff/config_example_atomwise.json" \ + --output_dir="temp1" +``` + +## Building an id_prop.json from VASP runs + +A reference notebook compiles a directory of `vasprun.xml` files into an +`id_prop.json`: + +→ [Colab: `make_id_prop`](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) + +## Joint training with extra targets (v2024.10.30+) + +From **v2024.10.30**, you can jointly train energy + forces + stress along with +atomwise non-gradient properties (charges, magnetic moments) and optional fingerprint +features. See [Multi-output regression](multi-output-regression.md) and the end-to-end configs under +[`alignn/examples/`](https://github.com/atomgptlab/alignn/tree/main/alignn/examples). + +## Recommended Colab walkthrough + +The Silicon example is the best starting point before training a new force-field on +your own data: + +→ [Colab: Train ALIGNN-FF on Mlearn (Silicon)](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) + +## See also + +- [Using a pretrained ALIGNN-FF](../pretrained/alignn-ff.md) +- [ASE calculator](../usage/ase-calculator.md) +- [Multi-GPU training](multi-gpu.md) diff --git a/docs/training/multi-gpu.md b/docs/training/multi-gpu.md new file mode 100644 index 0000000..4fcaa08 --- /dev/null +++ b/docs/training/multi-gpu.md @@ -0,0 +1,51 @@ +# Multi-GPU & High-Throughput Training + +## Multi-GPU (`torchrun`) + +ALIGNN supports PyTorch `DistributedDataParallel` launched via `torchrun`: + +```bash +torchrun --nproc_per_node=4 train_alignn.py \ + --root_dir DataDir \ + --config config.json \ + --output_dir temp +``` + +!!! note "Experimental" + Multi-GPU training is not yet thoroughly tested. Please report issues on GitHub. + +### SLURM example + +``` +#SBATCH -n 4 +#SBATCH -N 1 +#SBATCH --gres=gpu:4 + +torchrun --nproc_per_node=4 train_alignn.py \ + --root_dir DataDir --config config.json --output_dir temp +``` + +Make sure `--nproc_per_node` matches the GPUs requested from the scheduler. + +## High-throughput training + +For running the same training pipeline across many public datasets, see +[`alignn/scripts/train_*.py`](https://github.com/atomgptlab/alignn/tree/main/alignn/scripts). +These scripts: + +- Download datasets via + [jarvis-tools databases](https://jarvis-tools.readthedocs.io/en/master/databases.html) + (JARVIS-DFT, Materials Project, QM9_JCTC, …) +- Generate the `id_prop.csv` and per-target configs automatically +- Submit one training job per target property + +Adapt the scheduler-specific lines (`sbatch`, `qsub`, …) at the top of each script for +your cluster. + +## When to use which + +| Situation | Use | +|---|---| +| Single dataset, single target, one node | plain `train_alignn.py` | +| Single dataset, multiple GPUs on one node | `torchrun --nproc_per_node=N` | +| Many datasets / targets, cluster queue | `alignn/scripts/train_*.py` | diff --git a/docs/training/multi-output-regression.md b/docs/training/multi-output-regression.md new file mode 100644 index 0000000..7389bbc --- /dev/null +++ b/docs/training/multi-output-regression.md @@ -0,0 +1,45 @@ +# Multi-Output Regression + +Train a single model that jointly predicts multiple scalar or vector targets — e.g. +formation energy + bandgap + total energy, or full electron / phonon DOS. + +## Command + +```bash +train_alignn.py \ + --root_dir "alignn/examples/sample_data_multi_prop" \ + --config "alignn/examples/sample_data/config_example.json" \ + --output_dir=temp +``` + +## Dataset + +Extra target columns are appended to each `id_prop.csv` row (order must match the config). +A generator script lives at +[`alignn/examples/sample_data_multi_prop/scripts/`](https://github.com/atomgptlab/alignn/tree/main/alignn/examples/sample_data_multi_prop). + +For vector targets like DOS, each target value is a list of floats; see the DOS example +folder under `alignn/examples/`. + +## Why multi-output? + +- Shared representations often improve individual target accuracy. +- One forward pass predicts all targets, which is useful for screening workflows. +- Correlated targets (e.g. energy and its components) benefit from joint loss signals. + +## Global atomwise + graph-wise training (v2024.10.30+) + +Starting with **v2024.10.30**, multi-output training can mix output *kinds* in one model: + +- **Graph-wise** scalar/vector outputs (e.g. energy, bandgap) +- **Atomwise gradient** outputs (forces — derivatives of energy w.r.t. positions) +- **Atomwise non-gradient** outputs (charges, magnetic moments) + +Optional per-atom or per-graph fingerprint features can be concatenated into the input +graph. See end-to-end configs in +[`alignn/examples/`](https://github.com/atomgptlab/alignn/tree/main/alignn/examples). + +## See also + +- [Force-field training](force-field.md) — atomwise forces & stress +- [Single-output regression](single-output-regression.md) diff --git a/docs/training/single-output-regression.md b/docs/training/single-output-regression.md new file mode 100644 index 0000000..cf0752d --- /dev/null +++ b/docs/training/single-output-regression.md @@ -0,0 +1,55 @@ +# Single-Output Regression + +Train a model that predicts a single scalar property (formation energy, bandgap, bulk +modulus, …). + +## Command + +```bash +train_alignn.py \ + --root_dir "alignn/examples/sample_data" \ + --config "alignn/examples/sample_data/config_example.json" \ + --output_dir=temp +``` + +## What you need + +- A dataset directory with structures and `id_prop.csv` — + see [Dataset Format](dataset-format.md). +- A config JSON, e.g. + [`config_example.json`](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/config_example.json). + +The bundled `sample_data` trains OPT-bandgap on 50 JARVIS-DFT materials. The data was +generated by +[`generate_sample_data_reg.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/scripts/generate_sample_data_reg.py) — +adapt it to build a larger dataset or use your own structures. + +## Output + +After training finishes, `--output_dir` contains: + +- `best_model.pt` — the model checkpoint with the lowest validation loss +- `history_train.json`, `history_val.json` — per-epoch metrics +- `prediction_results_test_set.csv` — predictions on the held-out test set + +## CLI reference + +```bash +train_alignn.py -h +``` + +## Tips + +- Use `batch_size: 32` or `64` in the config for real trainings. +- If you need to train on one set and evaluate on another, set `n_train`, `n_val`, + `n_test` and `keep_data_order: true`. +- For datasets hosted by `jarvis-tools`, the + [`alignn/scripts/`](https://github.com/atomgptlab/alignn/tree/main/alignn/scripts) + directory has ready-to-run scripts that download and train in one step. + +## See also + +- [Classification](classification.md) — same workflow with a threshold +- [Multi-output regression](multi-output-regression.md) — multiple targets at once +- [Colab: 2D exfoliation energy](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) +- [Colab: superconductor Tc](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) diff --git a/docs/usage/ase-calculator.md b/docs/usage/ase-calculator.md new file mode 100644 index 0000000..5d2ea19 --- /dev/null +++ b/docs/usage/ase-calculator.md @@ -0,0 +1,95 @@ +# ALIGNN-FF ASE Calculator + +The [ASE calculator](https://wiki.fysik.dtu.dk/ase/ase/calculators/calculators.html) +interface lets you drop ALIGNN-FF into any ASE-based workflow — optimization, phonons, +MD, equation-of-state scans, interfaces, etc. + +![ALIGNN-FF animation](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/animation.gif?raw=true) + +## Minimal usage + +```python +from alignn.ff.ff import AlignnAtomwiseCalculator, default_path + +calc = AlignnAtomwiseCalculator(path=default_path()) +ase_atoms.calc = calc +energy = ase_atoms.get_potential_energy() +forces = ase_atoms.get_forces() +``` + +## Full example — relax + EV curve on Silicon + +```python +from alignn.ff.ff import AlignnAtomwiseCalculator, default_path +from jarvis.io.vasp.inputs import Poscar +from jarvis.core.atoms import ase_to_atoms +from ase.constraints import ExpCellFilter +from ase.optimize.fire import FIRE +import numpy as np +import matplotlib.pyplot as plt + +poscar = """Si2 +1.0 +3.3641499856336465 -2.5027128e-09 1.94229273881412 +1.121382991333525 3.1717517190189715 1.9422927388141193 +-2.5909987e-09 -1.8321133e-09 3.884586486670313 +Si +2 +Cartesian +3.92483875 2.77528125 6.7980237500000005 +0.56069125 0.39646875 0.9711462500000001 +""" + +calc = AlignnAtomwiseCalculator(path=default_path()) + +def general_relaxer(atoms, calculator, fmax=0.05, steps=150, relax=True): + ase_atoms = atoms.ase_converter() + ase_atoms.calc = calculator + if not relax: + return ase_atoms.get_potential_energy() + ase_atoms = ExpCellFilter(ase_atoms) + FIRE(ase_atoms).run(fmax=fmax, steps=steps) + return ase_to_atoms(ase_atoms.atoms) + +atoms = Poscar.from_string(poscar).atoms +atoms = general_relaxer(atoms, calc) + +energies, volumes = [], [] +for strain in np.arange(-0.1, 0.1, 0.01): + s = atoms.strain_atoms(strain) + a = s.ase_converter() + a.calc = calc + energies.append(a.get_potential_energy()) + volumes.append(s.volume) + +plt.plot(volumes, energies, "-o") +plt.xlabel(r"Volume ($\AA^3$)") +plt.ylabel("Total energy (eV)") +``` + +## Choosing a model + +`default_path()` returns the currently recommended checkpoint. To pin a specific one, +pass the path directly: + +```python +calc = AlignnAtomwiseCalculator(path="alignn/ff/v12.2.2024_dft_3d_307k") +``` + +See the [pretrained ALIGNN-FF page](../pretrained/alignn-ff.md) for available +checkpoints. + +## What you can do from here + +Anything ASE supports — e.g.: + +- `ase.optimize.fire.FIRE`, `BFGS` for structure relaxation +- `ase.phonons.Phonons` for finite-difference phonons (or use `phonopy`) +- `ase.md.VelocityVerlet`, `Langevin`, `NPT` for molecular dynamics +- `ase.neb.NEB` for minimum energy paths + +## See also + +- [Pretrained ALIGNN-FF checkpoints](../pretrained/alignn-ff.md) +- [CLI `run_alignn_ff.py`](../pretrained/alignn-ff.md#cli) +- [Colab: relaxation + phonons + interfaces](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) diff --git a/docs/usage/webapps.md b/docs/usage/webapps.md new file mode 100644 index 0000000..7252ed8 --- /dev/null +++ b/docs/usage/webapps.md @@ -0,0 +1,23 @@ +# Web Apps + +Two hosted web apps let you try ALIGNN without writing any code. + +## JARVIS-ALIGNN property predictor + +→ [atomgpt.org/alignn](https://atomgpt.org/alignn) + +Paste a POSCAR and get predicted formation energy, total energy per atom, bandgap, and +several other properties trained on JARVIS-DFT. + +![JARVIS-ALIGNN](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/jalignn.PNG?raw=true) + +## ALIGNN-FF structure optimizer + +→ [atomgpt.org/alignn_ff_dynamics](https://atomgpt.org/alignn_ff_dynamics) + +Upload a structure and have ALIGNN-FF relax it in the browser. + +## When to use the apps vs. the library + +- **Apps:** ad-hoc single-structure checks, teaching demos, showing colleagues +- **Library:** batch jobs, fine-tuning, MD, custom properties, anything reproducible diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..765ba99 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,90 @@ +site_name: ALIGNN +site_description: Atomistic Line Graph Neural Network for materials property prediction and force-fields +site_url: https://atomgptlab.github.io/alignn/ +repo_url: https://github.com/atomgptlab/alignn +repo_name: atomgptlab/alignn +edit_uri: edit/develop/docs/ + +theme: + name: material + features: + - navigation.instant + - navigation.tabs + - navigation.sections + - navigation.top + - content.code.copy + - content.action.edit + - search.suggest + - search.highlight + - toc.follow + palette: + - media: "(prefers-color-scheme: light)" + scheme: default + primary: indigo + accent: indigo + toggle: + icon: material/weather-night + name: Switch to dark mode + - media: "(prefers-color-scheme: dark)" + scheme: slate + primary: indigo + accent: indigo + toggle: + icon: material/weather-sunny + name: Switch to light mode + icon: + repo: fontawesome/brands/github + +markdown_extensions: + - admonition + - attr_list + - md_in_html + - tables + - toc: + permalink: true + - pymdownx.details + - pymdownx.superfences + - pymdownx.tabbed: + alternate_style: true + - pymdownx.highlight: + anchor_linenums: true + line_spans: __span + pygments_lang_class: true + - pymdownx.inlinehilite + - pymdownx.snippets + - pymdownx.arithmatex: + generic: true + +plugins: + - search + +extra: + social: + - icon: fontawesome/brands/github + link: https://github.com/atomgptlab/alignn + - icon: fontawesome/brands/python + link: https://pypi.org/project/alignn/ + +nav: + - Home: index.md + - Installation: installation.md + - Training: + - Dataset Format: training/dataset-format.md + - Single-Output Regression: training/single-output-regression.md + - Classification: training/classification.md + - Multi-Output Regression: training/multi-output-regression.md + - Force-Field / MLIP: training/force-field.md + - Multi-GPU & High-Throughput: training/multi-gpu.md + - Pre-trained Models: + - Overview: pretrained/index.md + - Property Predictors: pretrained/property-predictor.md + - ALIGNN-FF: pretrained/alignn-ff.md + - Usage: + - ASE Calculator: usage/ase-calculator.md + - Web Apps: usage/webapps.md + - Reference: + - Package Overview: api.md + - Performance: performance.md + - Notes & FAQ: notes.md + - References: references.md + - Contributing: contributing.md