diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index c12b9f3..6b9d9dc 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -24,8 +24,8 @@ jobs: shell: bash -l {0} run: | conda install flake8 pycodestyle pydocstyle - flake8 --ignore E203,W503,W605 --exclude=examples,tests,scripts --statistics --count --exit-zero alignn - pycodestyle --ignore E203,W503,W605 --exclude=examples,tests,scripts alignn + flake8 --ignore E203,W503,W605,E501 --exclude=examples,tests,scripts --statistics --count --exit-zero alignn + pycodestyle --ignore E203,W503,W605,E501 --exclude=examples,tests,scripts alignn pydocstyle --match-dir=core --match-dir=io --match-dir=io --match-dir=ai --match-dir=analysis --match-dir=db --match-dir=tasks --count alignn - name: Run pytest diff --git a/README.md b/README.md index 2d1d0aa..eb884cc 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,8 @@ # Table of Contents * [Introduction](#intro) * [Installation](#install) -* [Examples](#example) +* [Examples — train every model type](#example) +* [Reproducing a JARVIS-Leaderboard contribution](#reproduce) * [Colab notebooks](#colab) * [Pre-trained models](#pretrained) * [JARVIS-ALIGNN webapp](#webapp) @@ -28,26 +29,182 @@ The Atomistic Line Graph Neural Network ([paper](https://www.nature.com/articles/s41524-021-00650-1)) introduces a graph convolution layer that explicitly models both two- and three-body interactions in atomistic systems. The ALIGNN-FF variant ([paper](https://pubs.rsc.org/en/content/articlehtml/2023/dd/d2dd00096b)) extends this to a force-field for structurally and chemically diverse systems across 89 elements. -See [docs/index.md](docs/index.md) for the full introduction. - ![ALIGNN layer schematic](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/schematic_lg.jpg) +> ⚡ **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`](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](docs/installation.md) for conda, GitHub, and pip installation methods. -## Examples +## Examples — train every model type + +All training recipes live on this page. Each one ships a **self-contained, +runnable example** under [`alignn/examples/recipes/`](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. + +| Recipe | Task | Graph | Example dir | +| --- | --- | --- | --- | +| kNN | scalar property | kNN (cutoff 8) | [`recipes/knn`](alignn/examples/recipes/knn) | +| Radius | scalar property (MD-compatible) | radius (cutoff 5) | [`recipes/radius`](alignn/examples/recipes/radius) | +| Tensor | D-dim response tensor | kNN | [`recipes/tensor`](alignn/examples/recipes/tensor) | +| Spectra | DOS / Raman curve | kNN | [`recipes/spectra`](alignn/examples/recipes/spectra) | +| Force field | energy + forces + stress | radius | [`recipes/forcefield`](alignn/examples/recipes/forcefield) | +| Atomwise | per-atom charge / moment | kNN | [`recipes/atomwise`](alignn/examples/recipes/atomwise) | + +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](docs/training/dataset-format.md) for the full spec. + +
+1. kNN graph — scalar property (formation energy, band gap, Tc, …) + +Wider k-nearest-neighbour graph (`cutoff: 8.0`, `max_neighbors: 12`) — the more +accurate choice for property prediction. + +```bash +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 +``` +Key knobs: `cutoff: 8.0`, `model.output_features: 1`, `graphwise_weight: 1.0`, +`calculate_gradient: false`. More: [recipes/knn/README.md](alignn/examples/recipes/knn/README.md). +
+ +
+2. Radius graph — scalar property (MD-compatible neighbour list) + +Same scalar task, but the fixed-radius graph (`cutoff: 5.0`) that is continuous +under displacement — use it when you need MD-consistency. + +```bash +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 +``` +Key knobs: `cutoff: 5.0` (vs 8.0 for kNN). More: [recipes/radius/README.md](alignn/examples/recipes/radius/README.md). +
+ +
+3. Tensor property (dielectric D=9, piezo D=18, elastic D=36) + +Predict a fixed-length response tensor per structure. Target is a length-`D` list. + +```bash +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 +``` +Key knobs: set `model.output_features` to your tensor dimension (9/18/36) and +match `D` in `make_toy_dataset.py`. More: [recipes/tensor/README.md](alignn/examples/recipes/tensor/README.md). +
+ +
+4. Spectra / multi-output curve (eDOS 300, pDOS 200, Raman 200) + +Predict a full curve on a fixed grid. Target is a length-`D` list (one per bin). + +```bash +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 +``` +Key knobs: `model.output_features` = number of bins (200/300); match `D` in the +toy script. More: [recipes/spectra/README.md](alignn/examples/recipes/spectra/README.md). +
-See [docs/training/](docs/training/) for dataset format and training examples: +
+5. Force field (energy + forces + stress, ALIGNN-FF) + +Train an interatomic potential with energy-conserving (gradient) forces and +stress — usable for relaxation, MD, and LAMMPS (`pair_alignn`). + +```bash +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 +``` +Key knobs: `model.calculate_gradient: true`, and the loss mixture +`graphwise_weight` (energy) / `gradwise_weight` (forces) / `stresswise_weight` +(stress). **Energy must be per atom.** More: +[recipes/forcefield/README.md](alignn/examples/recipes/forcefield/README.md). +
+ +
+6. Atomwise property (per-atom charges, magnetic moments) + +Predict one value per atom. Target is a length-`Natoms` list under a per-atom key. + +```bash +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 +``` +Key knobs: `model.atomwise_output_features: 1`, `atomwise_weight: 1.0`, +`graphwise_weight: 0.0`; pass `--atomwise_key charges`. More: +[recipes/atomwise/README.md](alignn/examples/recipes/atomwise/README.md). +
+ +For the historical per-topic docs see also [docs/training/](docs/training/) +(dataset format, classification, multi-GPU). + + +## Reproducing a JARVIS-Leaderboard contribution + +Every ALIGNN entry on the [JARVIS-Leaderboard](https://atomgptlab.github.io/jarvis_leaderboard/) +ships the exact config, split, and `run.sh` used to produce it, so any result can +be reproduced end to end: + +```bash +# 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// +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`). -- [Dataset format](docs/training/dataset-format.md) -- [Single-output regression](docs/training/single-output-regression.md) -- [Classification](docs/training/classification.md) -- [Multi-output regression](docs/training/multi-output-regression.md) -- [Force-field training](docs/training/force-field.md) -- [Multi-GPU training](docs/training/multi-gpu.md) ## Colab notebooks @@ -102,24 +259,165 @@ print(calc.predictions()) # extra property predictors ``` A single pydantic config selects the outputs (force-field -energy/forces/stress plus any pretrained scalar property predictors). +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](docs/usage/ase-calculator.md) for more, and the ASE docs page *Calculators → ALIGNN*. ## Performances -See [docs/performance.md](docs/performance.md) for benchmark tables on JARVIS-DFT, Materials Project, QM9, hMOF, qMOF, OpenCatalyst, and other datasets. Also see [JARVIS-Leaderboard](https://pages.nist.gov/jarvis_leaderboard/). +ALIGNN 2.0 benchmarked across single-property, multi-property (spectra / per-atom / +tensor), and interatomic-force-field tasks. 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](https://atomgptlab.github.io/jarvis_leaderboard/). + +
+Full benchmark table (54 tasks) + +**(a) Single-property prediction — test MAE** + +| # | Task (unit) | N tr/val/te | ALIGNN 2.0 (radius) | ALIGNN 2.0 (kNN) | orig. ALIGNN | CGCNN | Baseline (MAD) | Skill % | +|---|---|---|---|---|---|---|---|---| +| 1 | formation_energy (eV/atom) | 44569/5572/5572 | 0.0316 | **0.0307** | 0.0331 | 0.0551 | 0.876 | 96.5 | +| 2 | optb88vdw_total_energy (eV/atom) | 44569/5572/5572 | 0.0321 | **0.0314** | 0.0367 | 0.0584 | 1.786 | 98.2 | +| 3 | optb88vdw_bandgap (eV) | 44569/5572/5572 | 0.1314 | **0.1306** | 0.1423 | 0.1857 | 0.999 | 86.9 | +| 4 | mbj_bandgap (eV) | 14535/1817/1815 | **0.2721** | 0.2730 | 0.3104 | 0.3261 | 1.765 | 84.6 | +| 5 | QM9 HOMO–LUMO gap (eV) | 110,000/10,000/10,829 | **0.031** | | 0.0345 | | 0.834 | 96.3 | +| 6 | QMOF bandgap (eV) | 16,340/2042/2042 | 0.208 | | **0.202** | | 0.946 | 78.7 | +| 7 | ehull (eV/atom) | 44290/5537/5537 | **0.0576** | 0.0590 | 0.0763 | 0.0590 | 1.148 | 95.0 | +| 8 | bulk_modulus_kv (GPa) | 15744/1968/1968 | 9.885 | **9.302** | 10.399 | 11.015 | 53.76 | 82.7 | +| 9 | shear_modulus_gv (GPa) | 15744/1968/1968 | 9.063 | **8.825** | 9.476 | 10.079 | 27.06 | 67.4 | +| 10 | magmom_oszicar (μ_B) | 41766/5222/5222 | 0.2608 | **0.2567** | 0.2574 | 0.3065 | 1.254 | 79.5 | +| 11 | slme (%) | 7250/906/906 | 4.493 | **4.447** | 4.521 | 5.014 | 11.21 | 60.3 | +| 12 | spillage | 9101/1137/1137 | 0.3527 | **0.3456** | 0.3510 | 0.3844 | 0.518 | 33.3 | +| 13 | kpoint_length_unit (Å) | 44313/5540/5539 | 9.699 | **9.342** | 9.515 | 9.875 | 17.94 | 47.9 | +| 14 | encut (eV) | 44308/5539/5539 | 131.81 | **128.08** | 133.80 | 134.83 | 262.6 | 51.2 | +| 15 | epsx | 35592/4449/4449 | 20.705 | **20.139** | 20.394 | 22.199 | 57.45 | 64.9 | +| 16 | epsy | 35592/4449/4449 | 20.088 | **19.829** | 19.999 | 21.787 | 57.32 | 65.4 | +| 17 | epsz | 35592/4449/4449 | 19.633 | **19.453** | 19.568 | 21.121 | 55.79 | 65.1 | +| 18 | mepsx | 13447/1681/1681 | 24.646 | **23.847** | 24.046 | 26.929 | 63.39 | 62.4 | +| 19 | mepsy | 13447/1681/1681 | 23.823 | 24.044 | **23.648** | 26.556 | 63.68 | 62.6 | +| 20 | mepsz | 13447/1681/1681 | **23.247** | 23.531 | 23.731 | 26.629 | 60.71 | 61.7 | +| 21 | dfpt_piezo_max_dij (pC/N) | 2677/334/334 | 12.603 | **12.498** | 20.570 | 18.392 | 22.69 | 44.9 | +| 22 | dfpt_piezo_max_dielectric | 3764/470/470 | 26.823 | **24.305** | 28.151 | 30.961 | 43.91 | 44.7 | +| 23 | exfoliation_energy (meV/atom) | 650/81/81 | 40.272 | **37.628** | 52.703 | 45.762 | 61.03 | 38.3 | +| 24 | max_efg (10²¹V/m^2) | 9493/1186/1186 | 19.802 | 19.248 | **19.121** | 22.957 | 44.46 | 56.7 | +| 25 | avg_elec_mass (m_e) | 14114/1764/1764 | 0.0837 | **0.0810** | 0.0853 | 0.0921 | 0.225 | 64.1 | +| 26 | avg_hole_mass (m_e) | 14114/1764/1764 | 0.1299 | 0.1240 | **0.1239** | 0.1406 | 0.399 | 68.9 | +| 27 | n_Seebeck (\muV/K) | 18568/2321/2321 | 41.524 | **40.346** | 40.921 | 45.660 | 111.5 | 63.8 | +| 28 | n_powerfact (\muW/mK^2) | 18568/2321/2321 | 469.07 | 451.90 | **442.30** | 485.59 | 709.2 | 36.3 | +| 29 | ph_heat_capacity (J/mol/K) | 9644/1205/1205 | **9.577** | – | 9.606 | 12.936 | 40.16 | 76.2 | +| 30 | Thermal Cond. (log₁₀κ_L) | 3227/–/404 | 0.376 | **0.362** | – | – | 0.597 | 39.4 | +| 31 | Tc_supercon (K) | 556/30/30 | 1.637 | **1.490** | 2.032 | – | 2.723 | 45.3 | +| 32 | Tc_supercon_hydride (K) | 763/95/95 | 9.937 | **9.425** | – | – | 33.56 | 71.9 | +| 33 | Tc_supercon_ hydride_plus_bulk (K) | 1595/199/199 | 8.670 | **8.407** | – | – | 22.33 | 62.3 | +| 34 | alex_supercon Tc (K) | 6592/824/825 | **0.883** | | | | 2.818 | 68.7 | +| 35 | alex_supercon N(E_F) (states/eV) | 6592/824/825 | **0.821** | | | | 1.559 | 47.3 | +| 36 | alex_supercon θ_D (K) | 6592/824/825 | **11.33** | | | | 80.30 | 85.9 | +| 37 | alex_supercon λ | 6592/824/825 | **0.0707** | | | | 0.194 | 63.6 | +| 38 | alex_supercon ω_log (K) | 6592/824/825 | **20.31** | | | | 55.37 | 63.3 | + +**(b) Multi-property — spectra / per-atom / tensor; held-out MAE (col. "radius")** + +| # | Task (unit) | N tr/val/te | ALIGNN 2.0 (radius) | ALIGNN 2.0 (kNN) | orig. ALIGNN | CGCNN | Baseline (MAD) | Skill % | +|---|---|---|---|---|---|---|---|---| +| 39 | eDOS, electronic DOS (D=300) | 4103/227/229 | **0.0138** | | | | 0.0213 | 35.2 | +| 40 | pDOS, phonon DOS (D=200) | 4103/227/229 | 0.0819 | | | | 0.117 | 29.8 | +| 41 | Raman spectrum (D=200) | 4059/507/508 | 0.0378 | **0.0326** | | | 0.0497 | 34.4 | +| 42 | Bader charge, per atom (e) | 75,028/3000/3000 | 0.0192 | | | | 2.124 | 99.1 | +| 43 | Net charge, per atom (e) | 75,033/3000/3000 | 0.0167 | | | | – | – | +| 44 | Magnetic moment, per atom (μ_B) | 89,231/3000/3000 | 0.0256 | | | | 2.063 | 98.8 | +| 45 | Dielectric tensor (D=9) | 4103/227/229 | 1.690 | | | | 3.401 | 50.3 | +| 46 | Born effective charge (e) | 4472/248/249 | 0.234 | | | | – | – | +| 47 | Piezoelectric tensor, C/m^2 (D=18) | 4513/250/252 | 0.077 | | | | 0.089 | 13.9 | +| 48 | Elastic C_{ij} tensor, GPa (D=36) | 15,936/885/886 | 5.593 | | | | 18.73 | 70.1 | + +**(c) Interatomic force fields — `mlearn` per-element energy/force; large sets energy / force** + +| # | Task (unit) | N tr/val/te | ALIGNN 2.0 (radius) | ALIGNN 2.0 (kNN) | orig. ALIGNN | CGCNN | Baseline (MAD) | Skill % | +|---|---|---|---|---|---|---|---|---| +| 49 | `mlearn`-Si, energy (meV/atom) | 214/–/25 | 13.88‡ | | | | – | | +| 50 | `mlearn`-Si, force (eV/Å) | 214/–/25 | 0.0872‡ | | | | – | | +| 51 | ALIGNN-FF-DB (E/F) | 276,401/–/15,355 | 32.4† / 0.0564† | | | | – | | +| 52 | MATPES-PBE (E/F) | 391,241/21,736/– | 40.4 / 0.1475 | | | | – | | +| 53 | FD-FF, 1.1 M (E/F) | 1,097,227/60,957/60,958 | 28.9† / 0.0445† | | | | – | | +| 54 | MPtrj (E/F) | ~1.5 M | 56.7† / 0.0707† | | | | – | | +*Blank cells: not run for that graph/model. `–`: baseline unavailable or ill-defined. +† still training. ‡ `mlearn` MAE pending re-verification against a consistent per-atom +energy convention.* +
## Useful notes -See [docs/notes.md](docs/notes.md) for common pitfalls and FAQs. +
+Tips & FAQ + +**Pure-PyTorch path (no DGL)** + +- ALIGNN 2.0 runs fully in native PyTorch — set the model name to a `*_pure` variant + (e.g. `alignn_atomwise_pure`) and `neighbor_strategy: "pure_torch"`. DGL is optional. +- If you *do* use the legacy DGL path, install a DGL build matching your CUDA runtime; + mismatched builds are the most common install failure. + +**Structure file parsing** + +- Simple `.cif`/`.pdb` are handled by `jarvis-tools` directly. +- For complex CIFs: `pip install cif2cell==2.0.0a3`. For complex PDBs: + `conda install -c ambermd pytraj`. + +**Training hyperparameters** + +- Example configs ship with a small `batch_size`/`epochs` so tests run fast. **Use + `batch_size: 32`–`64` and `epochs: 100`–`300` for real trainings** — otherwise + training is slow and under-performing. +- `pandas >= 1.2.3` required. Since March 2024, `pytorch-ignite` is no longer a dependency. + +**CLIs are importable scripts** + +- `train_alignn.py`, `pretrained.py`, and `run_alignn_ff.py` install as executables in + your environment's `bin/` — just run them by name, no absolute path needed. + +**Known dataset issues** + +- **QM9**: see [issue #54](https://github.com/atomgptlab/alignn/issues/54) for a + data-split discrepancy affecting reproducibility. + +**Getting help** + +- GitHub issues: · Email: `drkamal@jhu.edu` +
## References -See [docs/references.md](docs/references.md) for the publication list. +If ALIGNN or ALIGNN-FF contributed to your work, please cite the relevant papers. + +
+Publication list + +**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) + +A complete list is maintained at [jarvis-tools publications](https://jarvis-tools.readthedocs.io/en/master/publications.html). +
## How to contribute diff --git a/alignn/__init__.py b/alignn/__init__.py index e117a37..2e78fff 100644 --- a/alignn/__init__.py +++ b/alignn/__init__.py @@ -1,3 +1,3 @@ """Version number.""" -__version__ = "2026.5.20" +__version__ = "2026.8.6" diff --git a/alignn/config.py b/alignn/config.py index ae8f2a0..d90094d 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -188,6 +188,16 @@ class TrainingConfig(BaseSettings): # DDP tuning. find_unused_parameters=True is slow; enable only if # your model has conditional branches whose gradients vary per step. ddp_find_unused_parameters: bool = False + # Resume optimizer/scheduler/epoch from /current_state.pt if + # present. Lets a run that hit the walltime continue with a single + # continuous LR schedule instead of restarting it (weights alone are + # restored via --restart_model_path, which does not restore these). + resume_checkpoint: bool = False + # LR-schedule horizon in epochs, independent of how many epochs this job + # runs (`epochs`). Set it to the FINAL target when a run is split into + # resumed segments so OneCycle spans the whole run and stays continuous + # across restarts. None => use `epochs` (the single-job default). + lr_total_epochs: Optional[int] = None # When True, forces cuDNN determinism (slower). Decoupled from seed so # you can seed for reproducibility without paying the speed cost. deterministic: bool = False diff --git a/alignn/data.py b/alignn/data.py index 97fcbe8..26ac604 100644 --- a/alignn/data.py +++ b/alignn/data.py @@ -429,6 +429,7 @@ def get_train_val_loaders( sampler=val_sampler, max_neighbors=max_neighbors, three_body_cutoff=three_body_cutoff, + read_existing=read_existing, classification=classification_threshold is not None, output_dir=output_dir, tmp_name=tmp_name, @@ -456,6 +457,8 @@ def get_train_val_loaders( cutoff=cutoff, cutoff_extra=cutoff_extra, max_neighbors=max_neighbors, + three_body_cutoff=three_body_cutoff, + read_existing=read_existing, classification=classification_threshold is not None, output_dir=output_dir, tmp_name=tmp_name, diff --git a/alignn/examples/recipes/atomwise/README.md b/alignn/examples/recipes/atomwise/README.md new file mode 100644 index 0000000..a94e7e2 --- /dev/null +++ b/alignn/examples/recipes/atomwise/README.md @@ -0,0 +1,42 @@ +# ALIGNN example: Atomwise (per-atom) property + +Predict a **per-atom** scalar — atomic charges (Bader), site magnetic moments, etc. Each atom gets its own prediction, so the target in `id_prop.json` is a list of length `Natoms` under the `charges` key. + +## Run it (CPU, ~1-2 min) + +```bash +# 1) generate a tiny synthetic dataset -> id_prop.json +python make_toy_dataset.py + +# 2) train (10 epochs on the toy data) +train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key target --id_key jid --atomwise_key charges +``` + +You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`, +with results written to `toy_out/` (`history_val.json`, predictions, checkpoint). + +## Key config knobs (`config_example.json`) + +- `model.atomwise_output_features: 1`, `atomwise_weight: 1.0` — turn on the per-atom head. +- `model.graphwise_weight: 0.0` — the graph-level target is unused. +- Pass `--atomwise_key charges` (rename to your per-atom key). + +> `target` (graph-level) is kept as a dummy 0.0 because `graphwise_weight` is 0. + +## ⚠️ This is a smoke test, not a real model + +The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove +the pipeline runs. For a usable model: + +- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data** + (thousands to millions of entries). +- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`. +- Expect training to take much longer and to need a GPU for large datasets. + +## Dataset format (`id_prop.json`) + +A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s): + +```json +[{"jid": "toy-0", "atoms": {...}, "target": ...}] +``` diff --git a/alignn/examples/recipes/atomwise/config_example.json b/alignn/examples/recipes/atomwise/config_example.json new file mode 100644 index 0000000..884aa86 --- /dev/null +++ b/alignn/examples/recipes/atomwise/config_example.json @@ -0,0 +1,83 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "pure_torch", + "id_tag": "jid", + "dtype": "float32", + "random_seed": 123, + "classification_threshold": null, + "n_val": null, + "n_test": null, + "n_train": null, + "train_ratio": 0.8, + "val_ratio": 0.1, + "test_ratio": 0.1, + "target_multiplication_factor": null, + "epochs": 10, + "batch_size": 4, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "A", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": true, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": true, + "num_workers": 0, + "cutoff": 8.0, + "cutoff_extra": 3.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "data_parallel": false, + "n_early_stopping": null, + "output_dir": "temp", + "use_lmdb": true, + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": false, + "atomwise_output_features": 1, + "graphwise_weight": 0.0, + "gradwise_weight": 0.0, + "stresswise_weight": 0.0, + "atomwise_weight": 1.0, + "link": "identity", + "zero_inflated": false, + "classification": false, + "force_mult_natoms": false, + "energy_mult_natoms": false, + "include_pos_deriv": false, + "use_cutoff_function": false, + "inner_cutoff": 3.0, + "stress_multiplier": 1.0, + "add_reverse_forces": true, + "lg_on_fly": true, + "batch_stress": true, + "multiply_cutoff": false, + "use_penalty": true, + "extra_features": 0, + "exponent": 5, + "penalty_factor": 0.1, + "penalty_threshold": 1.0 + } +} \ No newline at end of file diff --git a/alignn/examples/recipes/atomwise/make_toy_dataset.py b/alignn/examples/recipes/atomwise/make_toy_dataset.py new file mode 100644 index 0000000..b13ecbf --- /dev/null +++ b/alignn/examples/recipes/atomwise/make_toy_dataset.py @@ -0,0 +1,33 @@ +"""Generate a TOY atomwise per-atom property (e.g. charges) dataset -> id_prop.json (inline jarvis Atoms dicts). + +CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training +needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below, +supply real targets, and raise `epochs`/`batch_size` in config_example.json. +""" +import json, random +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +random.seed(0) +N = 40 # toy size -- increase to thousands for a real run + +def rattle(a, amp=0.05): + d = a.to_dict() + d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]] + return Atoms.from_dict(d) + +# base crystal (Si); swap for your own structures +base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"]) + +data = [] +for i in range(N): + a = rattle(base) + nat = a.num_atoms + # TOY per-atom target (one value per atom): replace with real charges/magmoms. + charges = [round(random.uniform(-1, 1), 3) for _ in range(nat)] + data.append({"jid": f"toy-{i}", "atoms": a.to_dict(), + "target": 0.0, # graph-level target unused (graphwise_weight=0) + "charges": charges}) # per-atom target, pass via --atomwise_key charges + +json.dump(data, open("id_prop.json", "w")) +print(f"wrote id_prop.json with {len(data)} entries (per-atom charges)") diff --git a/alignn/examples/recipes/forcefield/README.md b/alignn/examples/recipes/forcefield/README.md new file mode 100644 index 0000000..8c4b632 --- /dev/null +++ b/alignn/examples/recipes/forcefield/README.md @@ -0,0 +1,42 @@ +# ALIGNN example: Force field (energy + forces + stress) + +Train an **ALIGNN-FF** interatomic potential that outputs energy, analytic forces (gradients of the energy), and stress — usable for relaxation and molecular dynamics (including LAMMPS via `pair_alignn`). Uses the radius graph (`cutoff: 5.0`). + +## Run it (CPU, ~1-2 min) + +```bash +# 1) generate a tiny synthetic dataset -> id_prop.json +python make_toy_dataset.py + +# 2) train (10 epochs on the toy data) +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 +``` + +You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`, +with results written to `toy_out/` (`history_val.json`, predictions, checkpoint). + +## Key config knobs (`config_example.json`) + +- `model.calculate_gradient: true` — forces are the true gradient of the energy (energy-conserving). +- `graphwise_weight` (energy), `gradwise_weight` (forces), `stresswise_weight` (stress) — the loss mixture. +- Pass `--force_key forces`; stresses are read automatically if present. + +> **Energy must be per atom** (`energy_per_atom`) in `id_prop.json`, not per structure. Forces are `Natoms x 3`, stresses Voigt-6. + +## ⚠️ This is a smoke test, not a real model + +The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove +the pipeline runs. For a usable model: + +- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data** + (thousands to millions of entries). +- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`. +- Expect training to take much longer and to need a GPU for large datasets. + +## Dataset format (`id_prop.json`) + +A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s): + +```json +[{"jid": "toy-0", "atoms": {...}, "energy_per_atom": ...}] +``` diff --git a/alignn/examples/recipes/forcefield/config_example.json b/alignn/examples/recipes/forcefield/config_example.json new file mode 100644 index 0000000..ce65a93 --- /dev/null +++ b/alignn/examples/recipes/forcefield/config_example.json @@ -0,0 +1,83 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "energy_per_atom", + "atom_features": "cgcnn", + "neighbor_strategy": "pure_torch", + "id_tag": "jid", + "dtype": "float32", + "random_seed": 123, + "classification_threshold": null, + "n_val": null, + "n_test": null, + "n_train": null, + "train_ratio": 0.8, + "val_ratio": 0.1, + "test_ratio": 0.1, + "target_multiplication_factor": null, + "epochs": 10, + "batch_size": 4, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "A", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": true, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": true, + "num_workers": 0, + "cutoff": 5.0, + "cutoff_extra": 3.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "data_parallel": false, + "n_early_stopping": null, + "output_dir": "temp", + "use_lmdb": true, + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 1.0, + "stresswise_weight": 0.1, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false, + "force_mult_natoms": false, + "energy_mult_natoms": false, + "include_pos_deriv": false, + "use_cutoff_function": false, + "inner_cutoff": 3.5, + "stress_multiplier": 1.0, + "add_reverse_forces": true, + "lg_on_fly": true, + "batch_stress": true, + "multiply_cutoff": false, + "use_penalty": true, + "extra_features": 0, + "exponent": 5, + "penalty_factor": 0.1, + "penalty_threshold": 1.0 + } +} \ No newline at end of file diff --git a/alignn/examples/recipes/forcefield/make_toy_dataset.py b/alignn/examples/recipes/forcefield/make_toy_dataset.py new file mode 100644 index 0000000..c45f72d --- /dev/null +++ b/alignn/examples/recipes/forcefield/make_toy_dataset.py @@ -0,0 +1,35 @@ +"""Generate a TOY force-field (energy/forces/stress) dataset -> id_prop.json (inline jarvis Atoms dicts). + +CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training +needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below, +supply real targets, and raise `epochs`/`batch_size` in config_example.json. +""" +import json, random +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +random.seed(0) +N = 40 # toy size -- increase to thousands for a real run + +def rattle(a, amp=0.05): + d = a.to_dict() + d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]] + return Atoms.from_dict(d) + +# base crystal (Si); swap for your own structures +base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"]) + +data = [] +for i in range(N): + a = rattle(base) + nat = a.num_atoms + # TOY labels: replace with DFT energy_per_atom (eV/atom), forces (Nx3 eV/A), + # stresses (Voigt-6). Energy MUST be per atom. + energy_per_atom = -5.0 + 0.01 * i + forces = [[random.uniform(-0.1, 0.1) for _ in range(3)] for _ in range(nat)] + stresses = [random.uniform(-0.5, 0.5) for _ in range(6)] + data.append({"jid": f"toy-{i}", "atoms": a.to_dict(), + "energy_per_atom": energy_per_atom, "forces": forces, "stresses": stresses}) + +json.dump(data, open("id_prop.json", "w")) +print(f"wrote id_prop.json with {len(data)} entries (energy_per_atom/forces/stresses)") diff --git a/alignn/examples/recipes/knn/README.md b/alignn/examples/recipes/knn/README.md new file mode 100644 index 0000000..4650f65 --- /dev/null +++ b/alignn/examples/recipes/knn/README.md @@ -0,0 +1,42 @@ +# ALIGNN example: kNN graph — scalar property + +Predict a single graph-level scalar (formation energy, band gap, Tc, ...) using the **wider k-nearest-neighbour graph** (`cutoff: 8.0`, `max_neighbors: 12`). kNN is the more accurate choice for property prediction. + +## Run it (CPU, ~1-2 min) + +```bash +# 1) generate a tiny synthetic dataset -> id_prop.json +python make_toy_dataset.py + +# 2) train (10 epochs on the toy data) +train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key target --id_key jid +``` + +You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`, +with results written to `toy_out/` (`history_val.json`, predictions, checkpoint). + +## Key config knobs (`config_example.json`) + +- `cutoff: 8.0`, `max_neighbors: 12` — the kNN neighbourhood. +- `model.output_features: 1` — single scalar. +- `model.graphwise_weight: 1.0`, `calculate_gradient: false` — pure graph-level regression. + + + +## ⚠️ This is a smoke test, not a real model + +The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove +the pipeline runs. For a usable model: + +- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data** + (thousands to millions of entries). +- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`. +- Expect training to take much longer and to need a GPU for large datasets. + +## Dataset format (`id_prop.json`) + +A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s): + +```json +[{"jid": "toy-0", "atoms": {...}, "target": ...}] +``` diff --git a/alignn/examples/recipes/knn/config_example.json b/alignn/examples/recipes/knn/config_example.json new file mode 100644 index 0000000..6abc1a9 --- /dev/null +++ b/alignn/examples/recipes/knn/config_example.json @@ -0,0 +1,83 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "pure_torch", + "id_tag": "jid", + "dtype": "float32", + "random_seed": 123, + "classification_threshold": null, + "n_val": null, + "n_test": null, + "n_train": null, + "train_ratio": 0.8, + "val_ratio": 0.1, + "test_ratio": 0.1, + "target_multiplication_factor": null, + "epochs": 10, + "batch_size": 4, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "A", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": true, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": true, + "num_workers": 0, + "cutoff": 8.0, + "cutoff_extra": 3.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "data_parallel": false, + "n_early_stopping": null, + "output_dir": "temp", + "use_lmdb": true, + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": false, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 0.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false, + "force_mult_natoms": false, + "energy_mult_natoms": false, + "include_pos_deriv": false, + "use_cutoff_function": false, + "inner_cutoff": 3.0, + "stress_multiplier": 1.0, + "add_reverse_forces": true, + "lg_on_fly": true, + "batch_stress": true, + "multiply_cutoff": false, + "use_penalty": true, + "extra_features": 0, + "exponent": 5, + "penalty_factor": 0.1, + "penalty_threshold": 1.0 + } +} \ No newline at end of file diff --git a/alignn/examples/recipes/knn/make_toy_dataset.py b/alignn/examples/recipes/knn/make_toy_dataset.py new file mode 100644 index 0000000..749823a --- /dev/null +++ b/alignn/examples/recipes/knn/make_toy_dataset.py @@ -0,0 +1,30 @@ +"""Generate a TOY scalar-property (kNN) dataset -> id_prop.json (inline jarvis Atoms dicts). + +CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training +needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below, +supply real targets, and raise `epochs`/`batch_size` in config_example.json. +""" +import json, random +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +random.seed(0) +N = 40 # toy size -- increase to thousands for a real run + +def rattle(a, amp=0.05): + d = a.to_dict() + d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]] + return Atoms.from_dict(d) + +# base crystal (Si); swap for your own structures +base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"]) + +data = [] +for i in range(N): + a = rattle(base) + # TOY target: replace with a real DFT property (e.g. formation energy) + target = -5.0 + 0.1 * i + data.append({"jid": f"toy-{i}", "atoms": a.to_dict(), "target": target}) + +json.dump(data, open("id_prop.json", "w")) +print(f"wrote id_prop.json with {len(data)} entries") diff --git a/alignn/examples/recipes/radius/README.md b/alignn/examples/recipes/radius/README.md new file mode 100644 index 0000000..0d726d4 --- /dev/null +++ b/alignn/examples/recipes/radius/README.md @@ -0,0 +1,41 @@ +# ALIGNN example: Radius graph — scalar property + +Same scalar-property task as the kNN recipe, but with the **radius graph** (`cutoff: 5.0`). The radius neighbour list is continuous under atomic displacement, so it is the construction used for molecular dynamics / force fields. For static property prediction kNN is usually a little more accurate; use radius when you need MD-consistency. + +## Run it (CPU, ~1-2 min) + +```bash +# 1) generate a tiny synthetic dataset -> id_prop.json +python make_toy_dataset.py + +# 2) train (10 epochs on the toy data) +train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key target --id_key jid +``` + +You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`, +with results written to `toy_out/` (`history_val.json`, predictions, checkpoint). + +## Key config knobs (`config_example.json`) + +- `cutoff: 5.0` — fixed radius cutoff (vs 8.0 for kNN). +- `model.output_features: 1`, `graphwise_weight: 1.0`. + + + +## ⚠️ This is a smoke test, not a real model + +The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove +the pipeline runs. For a usable model: + +- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data** + (thousands to millions of entries). +- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`. +- Expect training to take much longer and to need a GPU for large datasets. + +## Dataset format (`id_prop.json`) + +A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s): + +```json +[{"jid": "toy-0", "atoms": {...}, "target": ...}] +``` diff --git a/alignn/examples/recipes/radius/config_example.json b/alignn/examples/recipes/radius/config_example.json new file mode 100644 index 0000000..f214152 --- /dev/null +++ b/alignn/examples/recipes/radius/config_example.json @@ -0,0 +1,83 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "pure_torch", + "id_tag": "jid", + "dtype": "float32", + "random_seed": 123, + "classification_threshold": null, + "n_val": null, + "n_test": null, + "n_train": null, + "train_ratio": 0.8, + "val_ratio": 0.1, + "test_ratio": 0.1, + "target_multiplication_factor": null, + "epochs": 10, + "batch_size": 4, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "A", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": true, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": true, + "num_workers": 0, + "cutoff": 5.0, + "cutoff_extra": 3.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "data_parallel": false, + "n_early_stopping": null, + "output_dir": "temp", + "use_lmdb": true, + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": false, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 0.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false, + "force_mult_natoms": false, + "energy_mult_natoms": false, + "include_pos_deriv": false, + "use_cutoff_function": false, + "inner_cutoff": 3.5, + "stress_multiplier": 1.0, + "add_reverse_forces": true, + "lg_on_fly": true, + "batch_stress": true, + "multiply_cutoff": false, + "use_penalty": true, + "extra_features": 0, + "exponent": 5, + "penalty_factor": 0.1, + "penalty_threshold": 1.0 + } +} \ No newline at end of file diff --git a/alignn/examples/recipes/radius/make_toy_dataset.py b/alignn/examples/recipes/radius/make_toy_dataset.py new file mode 100644 index 0000000..6b62e07 --- /dev/null +++ b/alignn/examples/recipes/radius/make_toy_dataset.py @@ -0,0 +1,29 @@ +"""Generate a TOY scalar-property (radius graph) dataset -> id_prop.json (inline jarvis Atoms dicts). + +CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training +needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below, +supply real targets, and raise `epochs`/`batch_size` in config_example.json. +""" +import json, random +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +random.seed(0) +N = 40 # toy size -- increase to thousands for a real run + +def rattle(a, amp=0.05): + d = a.to_dict() + d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]] + return Atoms.from_dict(d) + +# base crystal (Si); swap for your own structures +base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"]) + +data = [] +for i in range(N): + a = rattle(base) + target = -5.0 + 0.1 * i # TOY scalar target + data.append({"jid": f"toy-{i}", "atoms": a.to_dict(), "target": target}) + +json.dump(data, open("id_prop.json", "w")) +print(f"wrote id_prop.json with {len(data)} entries") diff --git a/alignn/examples/recipes/spectra/README.md b/alignn/examples/recipes/spectra/README.md new file mode 100644 index 0000000..fb79ca4 --- /dev/null +++ b/alignn/examples/recipes/spectra/README.md @@ -0,0 +1,41 @@ +# ALIGNN example: Spectra / multi-output curve + +Predict a full **spectral curve** on a fixed grid — electronic DOS (300 bins), phonon DOS (200 bins), or Raman spectrum (200 bins). The target is a length-`D` list (one value per bin). + +## Run it (CPU, ~1-2 min) + +```bash +# 1) generate a tiny synthetic dataset -> id_prop.json +python make_toy_dataset.py + +# 2) train (10 epochs on the toy data) +train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key target --id_key jid +``` + +You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`, +with results written to `toy_out/` (`history_val.json`, predictions, checkpoint). + +## Key config knobs (`config_example.json`) + +- `model.output_features: 200` — **number of bins** (200 Raman/pDOS, 300 eDOS); match `D` in `make_toy_dataset.py`. +- `criterion: l1` — averaged over bins. + +> Real DOS/Raman curves are smooth; the toy script uses a Gaussian bump as a stand-in. + +## ⚠️ This is a smoke test, not a real model + +The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove +the pipeline runs. For a usable model: + +- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data** + (thousands to millions of entries). +- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`. +- Expect training to take much longer and to need a GPU for large datasets. + +## Dataset format (`id_prop.json`) + +A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s): + +```json +[{"jid": "toy-0", "atoms": {...}, "target": ...}] +``` diff --git a/alignn/examples/recipes/spectra/config_example.json b/alignn/examples/recipes/spectra/config_example.json new file mode 100644 index 0000000..e2236b2 --- /dev/null +++ b/alignn/examples/recipes/spectra/config_example.json @@ -0,0 +1,83 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "pure_torch", + "id_tag": "jid", + "dtype": "float32", + "random_seed": 123, + "classification_threshold": null, + "n_val": null, + "n_test": null, + "n_train": null, + "train_ratio": 0.8, + "val_ratio": 0.1, + "test_ratio": 0.1, + "target_multiplication_factor": null, + "epochs": 10, + "batch_size": 4, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "A", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": true, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": true, + "num_workers": 0, + "cutoff": 8.0, + "cutoff_extra": 3.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "data_parallel": false, + "n_early_stopping": null, + "output_dir": "temp", + "use_lmdb": true, + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 200, + "grad_multiplier": -1, + "calculate_gradient": false, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 0.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false, + "force_mult_natoms": false, + "energy_mult_natoms": false, + "include_pos_deriv": false, + "use_cutoff_function": false, + "inner_cutoff": 3.0, + "stress_multiplier": 1.0, + "add_reverse_forces": true, + "lg_on_fly": true, + "batch_stress": true, + "multiply_cutoff": false, + "use_penalty": true, + "extra_features": 0, + "exponent": 5, + "penalty_factor": 0.1, + "penalty_threshold": 1.0 + } +} \ No newline at end of file diff --git a/alignn/examples/recipes/spectra/make_toy_dataset.py b/alignn/examples/recipes/spectra/make_toy_dataset.py new file mode 100644 index 0000000..6f80879 --- /dev/null +++ b/alignn/examples/recipes/spectra/make_toy_dataset.py @@ -0,0 +1,33 @@ +"""Generate a TOY spectra (D=200 bins) dataset -> id_prop.json (inline jarvis Atoms dicts). + +CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training +needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below, +supply real targets, and raise `epochs`/`batch_size` in config_example.json. +""" +import json, random +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +random.seed(0) +N = 40 # toy size -- increase to thousands for a real run + +def rattle(a, amp=0.05): + d = a.to_dict() + d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]] + return Atoms.from_dict(d) + +# base crystal (Si); swap for your own structures +base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"]) + +import math +D = 200 # number of spectral bins -- match config output_features (200 Raman/pDOS, 300 eDOS) +data = [] +for i in range(N): + a = rattle(base) + # TOY spectrum: a shifted Gaussian bump; replace with a real DOS/Raman curve + c = 40 + i % 120 + target = [math.exp(-((k - c) ** 2) / (2 * 15 ** 2)) for k in range(D)] + data.append({"jid": f"toy-{i}", "atoms": a.to_dict(), "target": target}) + +json.dump(data, open("id_prop.json", "w")) +print(f"wrote id_prop.json with {len(data)} entries, {D} bins") diff --git a/alignn/examples/recipes/tensor/README.md b/alignn/examples/recipes/tensor/README.md new file mode 100644 index 0000000..e479a12 --- /dev/null +++ b/alignn/examples/recipes/tensor/README.md @@ -0,0 +1,41 @@ +# ALIGNN example: Tensor property (D-dimensional) + +Predict a fixed-length response **tensor** per structure — e.g. the 9-component dielectric tensor, 18-component piezoelectric tensor, or 36-component elastic $C_{ij}$. The target in `id_prop.json` is a length-`D` list. + +## Run it (CPU, ~1-2 min) + +```bash +# 1) generate a tiny synthetic dataset -> id_prop.json +python make_toy_dataset.py + +# 2) train (10 epochs on the toy data) +train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key target --id_key jid +``` + +You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`, +with results written to `toy_out/` (`history_val.json`, predictions, checkpoint). + +## Key config knobs (`config_example.json`) + +- `model.output_features: 9` — **set this to your tensor dimension** (9 dielectric, 18 piezo, 36 elastic) and match `D` in `make_toy_dataset.py`. +- `criterion: l1` — mean-absolute-error over components. + +> The toy script emits a length-9 vector; change `D` in both `make_toy_dataset.py` and `output_features` together. + +## ⚠️ This is a smoke test, not a real model + +The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove +the pipeline runs. For a usable model: + +- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data** + (thousands to millions of entries). +- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`. +- Expect training to take much longer and to need a GPU for large datasets. + +## Dataset format (`id_prop.json`) + +A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s): + +```json +[{"jid": "toy-0", "atoms": {...}, "target": ...}] +``` diff --git a/alignn/examples/recipes/tensor/config_example.json b/alignn/examples/recipes/tensor/config_example.json new file mode 100644 index 0000000..fc2e6c6 --- /dev/null +++ b/alignn/examples/recipes/tensor/config_example.json @@ -0,0 +1,83 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "pure_torch", + "id_tag": "jid", + "dtype": "float32", + "random_seed": 123, + "classification_threshold": null, + "n_val": null, + "n_test": null, + "n_train": null, + "train_ratio": 0.8, + "val_ratio": 0.1, + "test_ratio": 0.1, + "target_multiplication_factor": null, + "epochs": 10, + "batch_size": 4, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "A", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": true, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": true, + "num_workers": 0, + "cutoff": 8.0, + "cutoff_extra": 3.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "data_parallel": false, + "n_early_stopping": null, + "output_dir": "temp", + "use_lmdb": true, + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 9, + "grad_multiplier": -1, + "calculate_gradient": false, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 0.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false, + "force_mult_natoms": false, + "energy_mult_natoms": false, + "include_pos_deriv": false, + "use_cutoff_function": false, + "inner_cutoff": 3.0, + "stress_multiplier": 1.0, + "add_reverse_forces": true, + "lg_on_fly": true, + "batch_stress": true, + "multiply_cutoff": false, + "use_penalty": true, + "extra_features": 0, + "exponent": 5, + "penalty_factor": 0.1, + "penalty_threshold": 1.0 + } +} \ No newline at end of file diff --git a/alignn/examples/recipes/tensor/make_toy_dataset.py b/alignn/examples/recipes/tensor/make_toy_dataset.py new file mode 100644 index 0000000..0a3db42 --- /dev/null +++ b/alignn/examples/recipes/tensor/make_toy_dataset.py @@ -0,0 +1,31 @@ +"""Generate a TOY tensor-property (D=9) dataset -> id_prop.json (inline jarvis Atoms dicts). + +CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training +needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below, +supply real targets, and raise `epochs`/`batch_size` in config_example.json. +""" +import json, random +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +random.seed(0) +N = 40 # toy size -- increase to thousands for a real run + +def rattle(a, amp=0.05): + d = a.to_dict() + d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]] + return Atoms.from_dict(d) + +# base crystal (Si); swap for your own structures +base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"]) + +D = 9 # 9 dielectric, 18 piezoelectric, 36 elastic Cij -- match config output_features +data = [] +for i in range(N): + a = rattle(base) + # TOY length-D vector target: replace with a real response tensor (flattened) + target = [float(i % 7) + 0.01 * k for k in range(D)] + data.append({"jid": f"toy-{i}", "atoms": a.to_dict(), "target": target}) + +json.dump(data, open("id_prop.json", "w")) +print(f"wrote id_prop.json with {len(data)} entries, target dim {D}") diff --git a/alignn/examples/sample_data/config_example.json b/alignn/examples/sample_data/config_example.json index 073819a..406f24e 100644 --- a/alignn/examples/sample_data/config_example.json +++ b/alignn/examples/sample_data/config_example.json @@ -2,8 +2,8 @@ "version": "112bbedebdaecf59fb18e11c929080fb2f358246", "dataset": "user_data", "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "atom_features": "atomic_number", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "dtype": "float32", "random_seed": 123, @@ -45,10 +45,10 @@ "output_dir": "temp", "use_lmdb": true, "model": { - "name": "alignn_atomwise", + "name": "alignn_atomwise_pure", "alignn_layers": 4, "gcn_layers": 4, - "atom_input_features": 92, + "atom_input_features": 1, "edge_input_features": 80, "triplet_input_features": 40, "embedding_features": 64, @@ -80,4 +80,4 @@ "penalty_factor": 0.1, "penalty_threshold": 1.0 } -} \ No newline at end of file +} diff --git a/alignn/examples/sample_data_ff/config_example_atomwise.json b/alignn/examples/sample_data_ff/config_example_atomwise.json index fe0caec..e566980 100644 --- a/alignn/examples/sample_data_ff/config_example_atomwise.json +++ b/alignn/examples/sample_data_ff/config_example_atomwise.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "radius_graph", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "dtype": "float32", "random_seed": 123, @@ -40,7 +40,7 @@ "distributed": false, "use_lmdb": true, "model": { - "name": "alignn_atomwise", + "name": "alignn_atomwise_pure", "atom_input_features": 92, "calculate_gradient": true, "atomwise_output_features": 0, diff --git a/alignn/examples/sample_data_ff/econfig_example_atomwise.json b/alignn/examples/sample_data_ff/econfig_example_atomwise.json index 522cfd1..2c8b2e0 100644 --- a/alignn/examples/sample_data_ff/econfig_example_atomwise.json +++ b/alignn/examples/sample_data_ff/econfig_example_atomwise.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "radius_graph", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "dtype": "float32", "random_seed": 123, diff --git a/alignn/examples/sample_data_ff_additional/config.json b/alignn/examples/sample_data_ff_additional/config.json index 69944ed..7b624c5 100644 --- a/alignn/examples/sample_data_ff_additional/config.json +++ b/alignn/examples/sample_data_ff_additional/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "atomic_number", - "neighbor_strategy": "radius_graph", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "dtype": "float32", "random_seed": 123, @@ -45,7 +45,7 @@ "output_dir": "temp", "use_lmdb": true, "model": { - "name": "alignn_atomwise", + "name": "alignn_atomwise_pure", "alignn_layers": 2, "gcn_layers": 2, "atom_input_features": 1, diff --git a/alignn/examples/sample_data_ff_feats/config_example_atomwise.json b/alignn/examples/sample_data_ff_feats/config_example_atomwise.json index 68892d7..60f3236 100644 --- a/alignn/examples/sample_data_ff_feats/config_example_atomwise.json +++ b/alignn/examples/sample_data_ff_feats/config_example_atomwise.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, @@ -38,7 +38,7 @@ "keep_data_order": true, "distributed": false, "model": { - "name": "alignn_atomwise", + "name": "alignn_atomwise_pure", "atom_input_features": 92, "calculate_gradient": true, "atomwise_output_features": 0, diff --git a/alignn/ff/all_models_alignn_atomwise.json b/alignn/ff/all_models_alignn_atomwise.json index f0c1465..3f439f1 100644 --- a/alignn/ff/all_models_alignn_atomwise.json +++ b/alignn/ff/all_models_alignn_atomwise.json @@ -1,57 +1,58 @@ { - "dft_3d_307k": "https://ndownloader.figshare.com/files/64744104", - "mps": "https://ndownloader.figshare.com/files/64744107", - "avg_elec_mass": "https://ndownloader.figshare.com/files/64744110", - "bulk_modulus_kv": "https://ndownloader.figshare.com/files/64744113", - "avg_hole_mass": "https://ndownloader.figshare.com/files/64744116", - "dfpt_piezo_max_dielectric": "https://ndownloader.figshare.com/files/64744119", - "density": "https://ndownloader.figshare.com/files/64744122", - "dfpt_piezo_max_dij": "https://ndownloader.figshare.com/files/64744125", - "dfpt_piezo_max_eij": "https://ndownloader.figshare.com/files/64744128", - "ehull": "https://ndownloader.figshare.com/files/64744131", - "epsx": "https://ndownloader.figshare.com/files/64744134", - "exfoliation_energy": "https://ndownloader.figshare.com/files/64744137", - "formation_energy_peratom": "https://ndownloader.figshare.com/files/64744140", - "magmom_oszicar": "https://ndownloader.figshare.com/files/64744143", - "max_efg": "https://ndownloader.figshare.com/files/64744146", - "max_ir_mode": "https://ndownloader.figshare.com/files/64744149", - "mbj_bandgap": "https://ndownloader.figshare.com/files/64744152", - "mepsx": "https://ndownloader.figshare.com/files/64744155", - "min_ir_mode": "https://ndownloader.figshare.com/files/64744158", - "n-powerfact": "https://ndownloader.figshare.com/files/64744161", - "poisson": "https://ndownloader.figshare.com/files/64744164", - "optb88vdw_bandgap": "https://ndownloader.figshare.com/files/64744167", - "n-Seebeck": "https://ndownloader.figshare.com/files/64744170", - "p-powerfact": "https://ndownloader.figshare.com/files/64744173", - "p-Seebeck": "https://ndownloader.figshare.com/files/64744176", - "shear_modulus_gv": "https://ndownloader.figshare.com/files/64744179", - "slme": "https://ndownloader.figshare.com/files/64744182", - "spillage": "https://ndownloader.figshare.com/files/64744185", - "Tc_supercon": "https://ndownloader.figshare.com/files/64744188", - "v12.2.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50904240", - "v12.2.2024_mp_1.5mill": "https://ndownloader.figshare.com/files/50904783", - "v12.2.2024_mp_187k": "https://ndownloader.figshare.com/files/50904801", - "v2024.12.12_dft_3d_multi_prop": "https://ndownloader.figshare.com/files/52025186", - "alex_band_gap": "https://ndownloader.figshare.com/files/51993641", - "alex_dos_pa": "https://ndownloader.figshare.com/files/51993653", - "alex_e_form": "https://ndownloader.figshare.com/files/51993647", - "alex_e_hull": "https://ndownloader.figshare.com/files/51993644", - "alex_e_total": "https://ndownloader.figshare.com/files/51993650", - "alex_mag_per_vol": "https://ndownloader.figshare.com/files/51993659", - "alex_vol_pa": "https://ndownloader.figshare.com/files/51993656", - "v10.30.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50634327", - "v10.30.2024_mp_168k": "https://ndownloader.figshare.com/files/50634318", - "v8.29.2024_dft_3d": "https://ndownloader.figshare.com/files/48889834", - "v8.29.2024_mpf": "https://ndownloader.figshare.com/files/48889837", - "v5.27.2024": "https://ndownloader.figshare.com/files/47286127", - "alignnff_fmult": "https://ndownloader.figshare.com/files/41583585", - "alignnff_wt10": "https://ndownloader.figshare.com/files/41583594", - "alignnff_fd": "https://ndownloader.figshare.com/files/41583582", - "alignnff_wt01": "https://ndownloader.figshare.com/files/41583588", - "alignnff_wt1": "https://ndownloader.figshare.com/files/41583591", - "fmult_mlearn_only": "https://ndownloader.figshare.com/files/41583597", - "aff_Oct23": "https://ndownloader.figshare.com/files/42880573", - "revised": "https://ndownloader.figshare.com/files/41583600", - "scf_fd_top_10_en_42_fmax_600_wt01": "https://ndownloader.figshare.com/files/41967375", - "scf_fd_top_10_en_42_fmax_600_wt10": "https://ndownloader.figshare.com/files/41967372" -} + "dft_3d_307k": "https://ndownloader.figshare.com/files/64744104", + "mps": "https://ndownloader.figshare.com/files/64744107", + "avg_elec_mass": "https://ndownloader.figshare.com/files/64744110", + "bulk_modulus_kv": "https://ndownloader.figshare.com/files/64744113", + "avg_hole_mass": "https://ndownloader.figshare.com/files/64744116", + "dfpt_piezo_max_dielectric": "https://ndownloader.figshare.com/files/64744119", + "density": "https://ndownloader.figshare.com/files/64744122", + "dfpt_piezo_max_dij": "https://ndownloader.figshare.com/files/64744125", + "dfpt_piezo_max_eij": "https://ndownloader.figshare.com/files/64744128", + "ehull": "https://ndownloader.figshare.com/files/64744131", + "epsx": "https://ndownloader.figshare.com/files/64744134", + "exfoliation_energy": "https://ndownloader.figshare.com/files/64744137", + "formation_energy_peratom": "https://ndownloader.figshare.com/files/64744140", + "magmom_oszicar": "https://ndownloader.figshare.com/files/64744143", + "max_efg": "https://ndownloader.figshare.com/files/64744146", + "max_ir_mode": "https://ndownloader.figshare.com/files/64744149", + "mbj_bandgap": "https://ndownloader.figshare.com/files/64744152", + "mepsx": "https://ndownloader.figshare.com/files/64744155", + "min_ir_mode": "https://ndownloader.figshare.com/files/64744158", + "n-powerfact": "https://ndownloader.figshare.com/files/64744161", + "poisson": "https://ndownloader.figshare.com/files/64744164", + "optb88vdw_bandgap": "https://ndownloader.figshare.com/files/64744167", + "n-Seebeck": "https://ndownloader.figshare.com/files/64744170", + "p-powerfact": "https://ndownloader.figshare.com/files/64744173", + "p-Seebeck": "https://ndownloader.figshare.com/files/64744176", + "shear_modulus_gv": "https://ndownloader.figshare.com/files/64744179", + "slme": "https://ndownloader.figshare.com/files/64744182", + "spillage": "https://ndownloader.figshare.com/files/64744185", + "Tc_supercon": "https://ndownloader.figshare.com/files/64744188", + "v12.2.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50904240", + "v12.2.2024_mp_1.5mill": "https://ndownloader.figshare.com/files/50904783", + "v12.2.2024_mp_187k": "https://ndownloader.figshare.com/files/50904801", + "v2024.12.12_dft_3d_multi_prop": "https://ndownloader.figshare.com/files/52025186", + "alex_band_gap": "https://ndownloader.figshare.com/files/51993641", + "alex_dos_pa": "https://ndownloader.figshare.com/files/51993653", + "alex_e_form": "https://ndownloader.figshare.com/files/51993647", + "alex_e_hull": "https://ndownloader.figshare.com/files/51993644", + "alex_e_total": "https://ndownloader.figshare.com/files/51993650", + "alex_mag_per_vol": "https://ndownloader.figshare.com/files/51993659", + "alex_vol_pa": "https://ndownloader.figshare.com/files/51993656", + "v10.30.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50634327", + "v10.30.2024_mp_168k": "https://ndownloader.figshare.com/files/50634318", + "v8.29.2024_dft_3d": "https://ndownloader.figshare.com/files/48889834", + "v8.29.2024_mpf": "https://ndownloader.figshare.com/files/48889837", + "v5.27.2024": "https://ndownloader.figshare.com/files/47286127", + "alignnff_fmult": "https://ndownloader.figshare.com/files/41583585", + "alignnff_wt10": "https://ndownloader.figshare.com/files/41583594", + "alignnff_fd": "https://ndownloader.figshare.com/files/41583582", + "alignnff_wt01": "https://ndownloader.figshare.com/files/41583588", + "alignnff_wt1": "https://ndownloader.figshare.com/files/41583591", + "fmult_mlearn_only": "https://ndownloader.figshare.com/files/41583597", + "aff_Oct23": "https://ndownloader.figshare.com/files/42880573", + "revised": "https://ndownloader.figshare.com/files/41583600", + "scf_fd_top_10_en_42_fmax_600_wt01": "https://ndownloader.figshare.com/files/41967375", + "scf_fd_top_10_en_42_fmax_600_wt10": "https://ndownloader.figshare.com/files/41967372", + "matpes_smooth": "https://ndownloader.figshare.com/files/67217507" +} \ No newline at end of file diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py index 0140d1a..bb3e07e 100644 --- a/alignn/ff/calculators.py +++ b/alignn/ff/calculators.py @@ -104,8 +104,9 @@ def get_figshare_model_ff( def default_path(): - """Get default model path.""" - dpath = get_figshare_model_ff(model_name="v12.2.2024_dft_3d_307k") + """Get default model path (ALIGNN 2.0 matpes_smooth, 2/2/128 smooth cutoff).""" + dpath = get_figshare_model_ff(model_name="matpes_smooth") + # dpath = get_figshare_model_ff(model_name="v12.2.2024_dft_3d_307k") # previous # dpath = get_figshare_model_ff(model_name="v5.27.2024") # dpath = get_figshare_model_ff(model_name="v8.29.2024_dft_3d") # dpath = get_figshare_model_ff(model_name="alignnff_wt10") @@ -303,6 +304,7 @@ def calculate(self, atoms, properties=None, system_changes=None): max_neighbors=self.config["max_neighbors"], atom_features=self.config["atom_features"], use_canonize=self.config["use_canonize"], + three_body_cutoff=self.config.get("three_body_cutoff", None), ) # print("self.devicee", self.device) # print("g", g.device) @@ -314,14 +316,17 @@ def calculate(self, atoms, properties=None, system_changes=None): ( g.to(self.device), lg.to(self.device), - torch.tensor(atoms.cell) + torch.tensor(np.array(atoms.cell)) .type(torch.get_default_dtype()) .to(self.device), ) ) else: result = self.model( - (g.to(self.device), torch.tensor(atoms.cell).to(self.device)) + ( + g.to(self.device), + torch.tensor(np.array(atoms.cell)).to(self.device), + ) ) # print("result",result) if "atomwise" in self.config["model"]["name"]: diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py index da2b28f..3a30895 100644 --- a/alignn/ff/ff.py +++ b/alignn/ff/ff.py @@ -207,8 +207,9 @@ def get_figshare_model_prop( def default_path(): - """Get default model path.""" - dpath = get_figshare_model_ff(model_name="mps") + """Get default model path (ALIGNN 2.0 matpes_smooth, 2/2/128 smooth cutoff).""" + dpath = get_figshare_model_ff(model_name="matpes_smooth") + # dpath = get_figshare_model_ff(model_name="mps") # previous default # dpath = get_figshare_model_ff(model_name="v12.2.2024_dft_3d_307k") # dpath = get_figshare_model_ff(model_name="v5.27.2024") # dpath = get_figshare_model_ff(model_name="v8.29.2024_dft_3d") @@ -351,9 +352,9 @@ def example_print(self): except Exception: pass line += ( - f"a={self.atoms.get_cell()[0,0]: 3.3f} Ang " - + f"b={self.atoms.get_cell()[1,1]: 3.3f} Ang " - + f"c={self.atoms.get_cell()[2,2]: 3.3f} Ang " + f"a={self.atoms.get_cell()[0, 0]: 3.3f} Ang " + + f"b={self.atoms.get_cell()[1, 1]: 3.3f} Ang " + + f"c={self.atoms.get_cell()[2, 2]: 3.3f} Ang " + f"Volume={self.atoms.get_volume(): 3.3f} amu/a3 " + f"PE={self.atoms.get_potential_energy(): 5.5f} eV " + f"KE={self.atoms.get_kinetic_energy(): 5.5f} eV " diff --git a/alignn/ff/lammps_bridge.py b/alignn/ff/lammps_bridge.py index 3bebad4..f73a6ab 100644 --- a/alignn/ff/lammps_bridge.py +++ b/alignn/ff/lammps_bridge.py @@ -7,9 +7,10 @@ Usage: python -m alignn.ff.lammps_bridge \\ --data system.data \\ - --model-path alignn/ff/v12.2.2024_dft_3d_307k \\ --types Si,O \\ --steps 1000 --timestep 0.001 --temp 300 + (uses the default ALIGNN-FF model, default_path(); pass --model-path to + override with your own model directory.) Notes: - Requires the LAMMPS Python module (`pip install lammps` or build with @@ -25,8 +26,7 @@ from ase import Atoms from ase.stress import voigt_6_to_full_3x3_stress -from alignn.ff.calculators import AlignnAtomwiseCalculator - +from alignn.ff.calculators import AlignnAtomwiseCalculator, default_path # eV/A^3 -> bar (LAMMPS `metal` pressure unit) EV_PER_A3_TO_BAR = 1.602176634e6 @@ -89,7 +89,7 @@ def callback(lmp, ntimestep, nlocal, tag, x, f): def run(args): from lammps import lammps - calc = AlignnAtomwiseCalculator(path=args.model_path) + calc = AlignnAtomwiseCalculator(path=args.model_path or default_path()) symbols_by_type = { i + 1: sym for i, sym in enumerate(args.types.split(",")) } @@ -119,8 +119,7 @@ def run(args): registered = True lmp.command(line.rstrip("\n")) else: - lmp.commands_string( - f""" + lmp.commands_string(f""" units metal atom_style atomic boundary p p p @@ -133,11 +132,10 @@ def run(args): timestep {args.timestep} velocity all create {args.temp} {args.seed} mom yes rot yes fix nve all nve - fix tfix all langevin {args.temp} {args.temp} 0.1 {args.seed} # noqa: E501 + fix tfix all langevin {args.temp} {args.temp} 0.1 {args.seed} thermo 10 thermo_style custom step temp pe ke etotal press - """ - ) + """) cb = make_callback(calc, symbols_by_type) lmp.set_fix_external_callback("alignn", cb, lmp) lmp.command(f"run {args.steps}") @@ -153,7 +151,12 @@ def main(): help="Path to a LAMMPS input script. Must define `fix alignn ... " "external pf/callback` before any `run` command. Overrides --data.", ) - p.add_argument("--model-path", required=True, help="ALIGNN-FF model dir") + p.add_argument( + "--model-path", + default=None, + help="ALIGNN-FF model dir (default: the bundled default_path() model, " + "i.e. matpes_smooth)", + ) p.add_argument( "--types", required=True, diff --git a/alignn/ff/unified_calculator.py b/alignn/ff/unified_calculator.py index d144354..b2481cb 100644 --- a/alignn/ff/unified_calculator.py +++ b/alignn/ff/unified_calculator.py @@ -44,46 +44,67 @@ ase_to_atoms, ) from alignn.graphs import Graph -from alignn.pretrained import get_figshare_model - -# friendly name -> figshare model name (extend freely; raw figshare -# names ending in "_alignn" are also accepted as-is) -PROP_ALIASES: Dict[str, str] = { - "formation_energy_peratom": "jv_formation_energy_peratom_alignn", - "total_energy": "jv_optb88vdw_total_energy_alignn", - "optb88vdw_bandgap": "jv_optb88vdw_bandgap_alignn", - "mbj_bandgap": "jv_mbj_bandgap_alignn", - "bulk_modulus_kv": "jv_bulk_modulus_kv_alignn", - "shear_modulus_gv": "jv_shear_modulus_gv_alignn", - "ehull": "jv_ehull_alignn", - "spillage": "jv_spillage_alignn", - "slme": "jv_slme_alignn", - "magmom_oszicar": "jv_magmom_oszicar_alignn", - "exfoliation_energy": "jv_exfoliation_energy_alignn", - "supercon_tc": "jv_supercon_tc_alignn", - "epsx": "jv_epsx_alignn", - "n_seebeck": "jv_n-Seebeck_alignn", - "n_powerfact": "jv_n-powerfact_alignn", -} - - -def _resolve_prop(name: str) -> str: - """friendly or raw -> figshare model name.""" - if name in PROP_ALIASES: - return PROP_ALIASES[name] - if name.endswith("_alignn"): # raw figshare name passed through - return name - raise ValueError( - f"Unknown property '{name}'. Known: " - f"{sorted(PROP_ALIASES)} (or a raw *_alignn figshare name)." +from alignn.pretrained2 import ( + get_alignn2_model, + resolve_by_target, + ALIGNN2_MODELS, +) +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, + ALIGNNAtomWisePureConfig, +) +import json as _json + + +def _prop2_name(friendly, graph): + """Resolve a friendly property name to a pretrained2 model key, preferring the + requested graph. Handles all three registry conventions: a direct model name + (``elastic_tensor``), ``{name}_{graph}`` (``ir_radius``), and target lookup + (``formation_energy_peratom`` -> ``..._radius``). Falls back to the other graph + when only one variant exists (e.g. ``raman`` -> ``raman_knn``).""" + if friendly in ALIGNN2_MODELS: # direct model name + return friendly + other = "knn" if graph == "radius" else "radius" + for g in (graph, other): # {name}_{graph} + if "{}_{}".format(friendly, g) in ALIGNN2_MODELS: + return "{}_{}".format(friendly, g) + cands = resolve_by_target(friendly) # by training target + if cands: + pref = [m for m in cands if m.endswith("_" + graph)] + return (pref or cands)[0] + raise KeyError("No pretrained2 property model for '{}'".format(friendly)) + + +def _load_prop2_model(friendly, graph, device): + """Load a pure-PyTorch ALIGNN 2.0 property model (pretrained2) for `friendly` + on the requested `graph` ("radius"/"knn"). Scalar, spectra (D>1) and tensor + outputs are all supported. Returns a dict with the model and its own + graph-construction settings (cutoff/max_neighbors/atom_features).""" + name = _prop2_name(friendly, graph) + paths = get_alignn2_model(name) + cfg = _json.load(open(paths["config.json"])) + model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(**cfg["model"])) + model.load_state_dict( + torch.load( + paths["best_model.pt"], map_location=device, weights_only=False + ) ) + model.to(device).eval() + return { + "model": model, + "name": name, + "cutoff": float(cfg.get("cutoff", 5.0)), + "max_neighbors": int(cfg.get("max_neighbors", 12)), + "atom_features": cfg.get("atom_features", "cgcnn"), + "use_canonize": bool(cfg.get("use_canonize", False)), + } class AlignnUnifiedConfig(BaseModel): """Declarative spec of what the calculator should output.""" # force-field model (energy/forces/stress source) - ff_model: str = "v12.2.2024_dft_3d_307k" + ff_model: str = "matpes_smooth" energy: bool = True forces: bool = True stress: bool = True @@ -94,8 +115,10 @@ class AlignnUnifiedConfig(BaseModel): # shared knobs device: Optional[str] = None - prop_cutoff: float = 8.0 - prop_max_neighbors: int = 12 + # property-predictor graph: "radius" (default, FF-compatible) or "knn". + # Uses the pure-PyTorch ALIGNN 2.0 models from pretrained2; each carries its + # own cutoff/max_neighbors from its training config. + prop_graph: str = "radius" model_config = {"extra": "forbid"} @@ -103,7 +126,10 @@ class AlignnUnifiedConfig(BaseModel): @classmethod def _check_props(cls, v: List[str]) -> List[str]: for p in v: - _resolve_prop(p) # raises on unknown + try: # scalar, spectra or tensor ALIGNN 2.0 model must exist + _prop2_name(p, "radius") + except KeyError as exc: + raise ValueError(str(exc)) return v @classmethod @@ -152,40 +178,42 @@ def __init__(self, config=None, ff_path=None, **kw): # --- force-field calculator (loaded ONCE) --- self._ff = None if config.energy or config.forces or config.stress: - ff_kwargs = dict( - include_stress=config.stress, device=self.device - ) + ff_kwargs = dict(include_stress=config.stress, device=self.device) if ff_path is not None: ff_kwargs["path"] = ff_path elif config.ff_model: ff_kwargs["path"] = _ff_model_path(config.ff_model) self._ff = AlignnAtomwiseCalculator(**ff_kwargs) - # --- property predictor models (each loaded ONCE) --- - self._prop_models: Dict[str, object] = {} + # --- property predictor models (pure-torch ALIGNN 2.0, loaded ONCE) --- + self._prop_models: Dict[str, dict] = {} for friendly in config.properties: - self._prop_models[friendly] = get_figshare_model( - _resolve_prop(friendly) + self._prop_models[friendly] = _load_prop2_model( + friendly, config.prop_graph, self.device ) # -- scalar property forward (reuses cached model, no reload) ------- - def _predict_scalar(self, model, j_atoms) -> float: + def _predict_scalar(self, info, j_atoms) -> float: + # pure-torch graph at the property model's own cutoff/neighbors g, lg = Graph.atom_dgl_multigraph( j_atoms, - cutoff=float(self.cfg.prop_cutoff), - max_neighbors=self.cfg.prop_max_neighbors, + neighbor_strategy="pure_torch", + cutoff=info["cutoff"], + max_neighbors=info["max_neighbors"], + atom_features=info["atom_features"], + use_canonize=info["use_canonize"], ) - lat = torch.tensor(j_atoms.lattice_mat) + lat = torch.tensor(j_atoms.lattice_mat).type(torch.get_default_dtype()) with torch.no_grad(): - out = model( - [ + out = info["model"]( + ( g.to(self.device), lg.to(self.device), lat.to(self.device), - ] + ) ) if isinstance(out, dict): - out = out["out"] + out = out.get("out", out.get("energy", next(iter(out.values())))) arr = out.detach().cpu().numpy().flatten() return float(arr[0]) if arr.size == 1 else arr.tolist() @@ -198,7 +226,8 @@ def calculate( # force-field block if self._ff is not None: self._ff.calculate( - self.atoms, properties=properties, + self.atoms, + properties=properties, system_changes=system_changes, ) if self.cfg.energy: @@ -221,10 +250,8 @@ def calculate( # scalar property predictors if self._prop_models: j_atoms = ase_to_atoms(self.atoms) - for friendly, model in self._prop_models.items(): - self.results[friendly] = self._predict_scalar( - model, j_atoms - ) + for friendly, info in self._prop_models.items(): + self.results[friendly] = self._predict_scalar(info, j_atoms) # convenience def predictions(self) -> Dict[str, object]: diff --git a/alignn/pretrained2.py b/alignn/pretrained2.py new file mode 100644 index 0000000..78d34ff --- /dev/null +++ b/alignn/pretrained2.py @@ -0,0 +1,1464 @@ +"""ALIGNN 2.0 pretrained-model registry (glossary). + +Analogous to :mod:`jarvis.db.figshare`: a single dict maps a model name to its +Figshare artifacts and metadata, and :func:`get_alignn2_model` downloads +(and caches) the ``config.json``, ``best_model.pt``, and +``ids_train_val_test.json`` for that model. + +All models live in the Figshare **ALIGNN2** project +(https://figshare.com/projects/ALIGNN2, id ``279395``). Each variant (e.g. +``formation_energy_peratom_radius`` and ``..._knn``) is its OWN article holding +ONE flat zip (best_model.pt + config.json + ids_train_val_test.json) with its +own ``url`` (``https://ndownloader.figshare.com/files/``) -- mirroring the +per-model download URLs in the original ``alignn/pretrained.py``. + +Note: Figshare *draft* files are not publicly downloadable; the ``url`` loader +activates once the article is published (token needed meanwhile). +""" + +import os +import zipfile +import requests + +FIGSHARE_PROJECT_ID = 279395 +FIGSHARE_PROJECT_URL = "https://figshare.com/projects/ALIGNN2/{}".format( + FIGSHARE_PROJECT_ID +) + +# --- registry ----------------------------------------------------------------- +# name -> metadata. Each entry carries its own `url` (one flat zip per variant). +ALIGNN2_MODELS = { + "alex_supercon_Tc": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "alex_Tc", + "unit": "", + "test_mae": 0.883, + "figshare_article_id": 33135179, + "url": "https://ndownloader.figshare.com/files/67163582", + "description": "Alexandria superconductor Tc (radius).", + }, + "alex_supercon_debye": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "alex_debye", + "unit": "", + "test_mae": 11.33, + "figshare_article_id": 33135185, + "url": "https://ndownloader.figshare.com/files/67163588", + "description": "Alexandria superconductor debye (radius).", + }, + "alex_supercon_dosef": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "alex_dosef", + "unit": "", + "test_mae": 0.821, + "figshare_article_id": 33135182, + "url": "https://ndownloader.figshare.com/files/67163585", + "description": "Alexandria superconductor dosef (radius).", + }, + "alex_supercon_la": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "alex_la", + "unit": "", + "test_mae": 0.0707, + "figshare_article_id": 33135188, + "url": "https://ndownloader.figshare.com/files/67163591", + "description": "Alexandria superconductor la (radius).", + }, + "alex_supercon_wlog": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "alex_wlog", + "unit": "", + "test_mae": 20.31, + "figshare_article_id": 33135191, + "url": "https://ndownloader.figshare.com/files/67163597", + "description": "Alexandria superconductor wlog (radius).", + }, + "avg_elec_mass_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "avg_elec_mass", + "unit": "", + "test_mae": None, + "figshare_article_id": 33134975, + "url": "https://ndownloader.figshare.com/files/67163339", + "description": "JARVIS-DFT avg_elec_mass (radius graph).", + }, + "avg_hole_mass_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "avg_hole_mass", + "unit": "", + "test_mae": None, + "figshare_article_id": 33134981, + "url": "https://ndownloader.figshare.com/files/67163345", + "description": "JARVIS-DFT avg_hole_mass (radius graph).", + }, + "bulk_modulus_kv_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "bulk_modulus_kv", + "unit": "", + "test_mae": 9.8854, + "figshare_article_id": 33134987, + "url": "https://ndownloader.figshare.com/files/67163351", + "description": "JARVIS-DFT bulk_modulus_kv (radius graph).", + }, + "c2db_gap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "gap", + "unit": "", + "test_mae": 0.0971, + "figshare_article_id": 33135236, + "url": "https://ndownloader.figshare.com/files/67163642", + "description": "c2db gap (radius).", + }, + "dfpt_piezo_max_dielectric_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "dfpt_piezo_max_dielectric", + "unit": "", + "test_mae": None, + "figshare_article_id": 33134993, + "url": "https://ndownloader.figshare.com/files/67163357", + "description": "JARVIS-DFT dfpt_piezo_max_dielectric (radius graph).", + }, + "dfpt_piezo_max_dij_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "dfpt_piezo_max_dij", + "unit": "", + "test_mae": None, + "figshare_article_id": 33134999, + "url": "https://ndownloader.figshare.com/files/67163363", + "description": "JARVIS-DFT dfpt_piezo_max_dij (radius graph).", + }, + "ehull_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "ehull", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135005, + "url": "https://ndownloader.figshare.com/files/67163369", + "description": "JARVIS-DFT ehull (radius graph).", + }, + "encut_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "encut", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135011, + "url": "https://ndownloader.figshare.com/files/67163375", + "description": "JARVIS-DFT encut (radius graph).", + }, + "epsx_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "epsx", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135017, + "url": "https://ndownloader.figshare.com/files/67163381", + "description": "JARVIS-DFT epsx (radius graph).", + }, + "epsy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "epsy", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135023, + "url": "https://ndownloader.figshare.com/files/67163387", + "description": "JARVIS-DFT epsy (radius graph).", + }, + "epsz_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "epsz", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135029, + "url": "https://ndownloader.figshare.com/files/67163393", + "description": "JARVIS-DFT epsz (radius graph).", + }, + "exfoliation_energy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "exfoliation_energy", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135041, + "url": "https://ndownloader.figshare.com/files/67163405", + "description": "JARVIS-DFT exfoliation_energy (radius graph).", + }, + "formation_energy_peratom_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "formation_energy_peratom", + "unit": "", + "test_mae": 0.0316, + "figshare_article_id": 33135047, + "url": "https://ndownloader.figshare.com/files/67163411", + "description": "JARVIS-DFT formation_energy_peratom (radius graph).", + }, + "kpoint_length_unit_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "kpoint_length_unit", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135053, + "url": "https://ndownloader.figshare.com/files/67163420", + "description": "JARVIS-DFT kpoint_length_unit (radius graph).", + }, + "magmom_oszicar_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "magmom_oszicar", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135059, + "url": "https://ndownloader.figshare.com/files/67163426", + "description": "JARVIS-DFT magmom_oszicar (radius graph).", + }, + "max_efg_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "max_efg", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135065, + "url": "https://ndownloader.figshare.com/files/67163450", + "description": "JARVIS-DFT max_efg (radius graph).", + }, + "mbj_bandgap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "mbj_bandgap", + "unit": "", + "test_mae": 0.2721, + "figshare_article_id": 33135071, + "url": "https://ndownloader.figshare.com/files/67163456", + "description": "JARVIS-DFT mbj_bandgap (radius graph).", + }, + "mepsx_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "mepsx", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135077, + "url": "https://ndownloader.figshare.com/files/67163462", + "description": "JARVIS-DFT mepsx (radius graph).", + }, + "mepsy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "mepsy", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135083, + "url": "https://ndownloader.figshare.com/files/67163468", + "description": "JARVIS-DFT mepsy (radius graph).", + }, + "mepsz_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "mepsz", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135089, + "url": "https://ndownloader.figshare.com/files/67163474", + "description": "JARVIS-DFT mepsz (radius graph).", + }, + "mxene275_formation_energy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "formation_energy", + "unit": "eV/atom", + "test_mae": 0.0348, + "figshare_article_id": 33135221, + "url": "https://ndownloader.figshare.com/files/67163627", + "description": "mxene275 formation energy (radius).", + }, + "n_Seebeck_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "n_Seebeck", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135095, + "url": "https://ndownloader.figshare.com/files/67163480", + "description": "JARVIS-DFT n_Seebeck (radius graph).", + }, + "n_powerfact_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "n_powerfact", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135101, + "url": "https://ndownloader.figshare.com/files/67163486", + "description": "JARVIS-DFT n_powerfact (radius graph).", + }, + "optb88vdw_bandgap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "optb88vdw_bandgap", + "unit": "", + "test_mae": 0.1314, + "figshare_article_id": 33135107, + "url": "https://ndownloader.figshare.com/files/67163492", + "description": "JARVIS-DFT optb88vdw_bandgap (radius graph).", + }, + "optb88vdw_total_energy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "optb88vdw_total_energy", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135113, + "url": "https://ndownloader.figshare.com/files/67163501", + "description": "JARVIS-DFT optb88vdw_total_energy (radius graph).", + }, + "polymer_genome_gga_gap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "gga_gap", + "unit": "", + "test_mae": 0.2274, + "figshare_article_id": 33135230, + "url": "https://ndownloader.figshare.com/files/67163636", + "description": "polymer_genome gga_gap (radius).", + }, + "shear_modulus_gv_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "shear_modulus_gv", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135119, + "url": "https://ndownloader.figshare.com/files/67163507", + "description": "JARVIS-DFT shear_modulus_gv (radius graph).", + }, + "slme_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "slme", + "unit": "", + "test_mae": 4.4929, + "figshare_article_id": 33135125, + "url": "https://ndownloader.figshare.com/files/67163516", + "description": "JARVIS-DFT slme (radius graph).", + }, + "spillage_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "spillage", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135131, + "url": "https://ndownloader.figshare.com/files/67163522", + "description": "JARVIS-DFT spillage (radius graph).", + }, + "tc_supercon_hydride_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "Tc_supercon_hydride", + "unit": "K", + "test_mae": 10.06, + "figshare_article_id": 33135215, + "url": "https://ndownloader.figshare.com/files/67163621", + "description": "Tc_supercon_hydride (radius, retrained pure-torch).", + }, + "tc_supercon_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "Tc_supercon", + "unit": "", + "test_mae": 1.637, + "figshare_article_id": 33135209, + "url": "https://ndownloader.figshare.com/files/67163615", + "description": "Tc_supercon (radius).", + }, + "thermal_cond_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "ltc", + "unit": "", + "test_mae": 0.386, + "figshare_article_id": 33135194, + "url": "https://ndownloader.figshare.com/files/67163600", + "description": "Lattice thermal cond. log10(kL) (radius).", + }, + "twod_matpd_bandgap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "bandgap", + "unit": "", + "test_mae": 0.3802, + "figshare_article_id": 33135242, + "url": "https://ndownloader.figshare.com/files/67163648", + "description": "twod_matpd bandgap (radius).", + }, + "avg_elec_mass_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "avg_elec_mass", + "unit": "", + "test_mae": 0.081, + "figshare_article_id": 33134978, + "url": "https://ndownloader.figshare.com/files/67163342", + "description": "JARVIS-DFT avg_elec_mass (knn graph).", + }, + "avg_hole_mass_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "avg_hole_mass", + "unit": "", + "test_mae": 0.124, + "figshare_article_id": 33134984, + "url": "https://ndownloader.figshare.com/files/67163348", + "description": "JARVIS-DFT avg_hole_mass (knn graph).", + }, + "bulk_modulus_kv_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "bulk_modulus_kv", + "unit": "", + "test_mae": 9.302, + "figshare_article_id": 33134990, + "url": "https://ndownloader.figshare.com/files/67163354", + "description": "JARVIS-DFT bulk_modulus_kv (knn graph).", + }, + "c2db_gap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "gap", + "unit": "", + "test_mae": 0.0802, + "figshare_article_id": 33135239, + "url": "https://ndownloader.figshare.com/files/67163645", + "description": "c2db gap (knn).", + }, + "dfpt_piezo_max_dielectric_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "dfpt_piezo_max_dielectric", + "unit": "", + "test_mae": 24.3052, + "figshare_article_id": 33134996, + "url": "https://ndownloader.figshare.com/files/67163360", + "description": "JARVIS-DFT dfpt_piezo_max_dielectric (knn graph).", + }, + "dfpt_piezo_max_dij_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "dfpt_piezo_max_dij", + "unit": "", + "test_mae": 12.4983, + "figshare_article_id": 33135002, + "url": "https://ndownloader.figshare.com/files/67163366", + "description": "JARVIS-DFT dfpt_piezo_max_dij (knn graph).", + }, + "ehull_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "ehull", + "unit": "", + "test_mae": 0.059, + "figshare_article_id": 33135008, + "url": "https://ndownloader.figshare.com/files/67163372", + "description": "JARVIS-DFT ehull (knn graph).", + }, + "encut_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "encut", + "unit": "", + "test_mae": 128.0826, + "figshare_article_id": 33135014, + "url": "https://ndownloader.figshare.com/files/67163378", + "description": "JARVIS-DFT encut (knn graph).", + }, + "epsx_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "epsx", + "unit": "", + "test_mae": 20.1393, + "figshare_article_id": 33135020, + "url": "https://ndownloader.figshare.com/files/67163384", + "description": "JARVIS-DFT epsx (knn graph).", + }, + "epsy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "epsy", + "unit": "", + "test_mae": 19.8292, + "figshare_article_id": 33135026, + "url": "https://ndownloader.figshare.com/files/67163390", + "description": "JARVIS-DFT epsy (knn graph).", + }, + "epsz_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "epsz", + "unit": "", + "test_mae": 19.4531, + "figshare_article_id": 33135032, + "url": "https://ndownloader.figshare.com/files/67163396", + "description": "JARVIS-DFT epsz (knn graph).", + }, + "exfoliation_energy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "exfoliation_energy", + "unit": "", + "test_mae": 37.6285, + "figshare_article_id": 33135044, + "url": "https://ndownloader.figshare.com/files/67163408", + "description": "JARVIS-DFT exfoliation_energy (knn graph).", + }, + "formation_energy_peratom_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "formation_energy_peratom", + "unit": "", + "test_mae": 0.0307, + "figshare_article_id": 33135050, + "url": "https://ndownloader.figshare.com/files/67163414", + "description": "JARVIS-DFT formation_energy_peratom (knn graph).", + }, + "hmof_co2": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "hmof_co2", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135173, + "url": "https://ndownloader.figshare.com/files/67163576", + "description": "hMOF CO2 uptake (kNN).", + }, + "kpoint_length_unit_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "kpoint_length_unit", + "unit": "", + "test_mae": 9.342, + "figshare_article_id": 33135056, + "url": "https://ndownloader.figshare.com/files/67163423", + "description": "JARVIS-DFT kpoint_length_unit (knn graph).", + }, + "magmom_oszicar_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "magmom_oszicar", + "unit": "", + "test_mae": 0.2567, + "figshare_article_id": 33135062, + "url": "https://ndownloader.figshare.com/files/67163429", + "description": "JARVIS-DFT magmom_oszicar (knn graph).", + }, + "max_efg_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "max_efg", + "unit": "", + "test_mae": 19.2483, + "figshare_article_id": 33135068, + "url": "https://ndownloader.figshare.com/files/67163453", + "description": "JARVIS-DFT max_efg (knn graph).", + }, + "mbj_bandgap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "mbj_bandgap", + "unit": "", + "test_mae": 0.273, + "figshare_article_id": 33135074, + "url": "https://ndownloader.figshare.com/files/67163459", + "description": "JARVIS-DFT mbj_bandgap (knn graph).", + }, + "mepsx_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "mepsx", + "unit": "", + "test_mae": 23.8474, + "figshare_article_id": 33135080, + "url": "https://ndownloader.figshare.com/files/67163465", + "description": "JARVIS-DFT mepsx (knn graph).", + }, + "mepsy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "mepsy", + "unit": "", + "test_mae": 24.0439, + "figshare_article_id": 33135086, + "url": "https://ndownloader.figshare.com/files/67163471", + "description": "JARVIS-DFT mepsy (knn graph).", + }, + "mepsz_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "mepsz", + "unit": "", + "test_mae": 23.5306, + "figshare_article_id": 33135092, + "url": "https://ndownloader.figshare.com/files/67163477", + "description": "JARVIS-DFT mepsz (knn graph).", + }, + "mxene275_formation_energy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "formation_energy", + "unit": "eV/atom", + "test_mae": 0.0343, + "figshare_article_id": 33135224, + "url": "https://ndownloader.figshare.com/files/67163630", + "description": "mxene275 formation energy (kNN).", + }, + "n_Seebeck_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "n_Seebeck", + "unit": "", + "test_mae": 40.3457, + "figshare_article_id": 33135098, + "url": "https://ndownloader.figshare.com/files/67163483", + "description": "JARVIS-DFT n_Seebeck (knn graph).", + }, + "n_powerfact_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "n_powerfact", + "unit": "", + "test_mae": 451.904, + "figshare_article_id": 33135104, + "url": "https://ndownloader.figshare.com/files/67163489", + "description": "JARVIS-DFT n_powerfact (knn graph).", + }, + "optb88vdw_bandgap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "optb88vdw_bandgap", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135110, + "url": "https://ndownloader.figshare.com/files/67163495", + "description": "JARVIS-DFT optb88vdw_bandgap (knn graph).", + }, + "optb88vdw_total_energy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "optb88vdw_total_energy", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135116, + "url": "https://ndownloader.figshare.com/files/67163504", + "description": "JARVIS-DFT optb88vdw_total_energy (knn graph).", + }, + "polymer_genome_gga_gap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "gga_gap", + "unit": "", + "test_mae": 0.2273, + "figshare_article_id": 33135233, + "url": "https://ndownloader.figshare.com/files/67163639", + "description": "polymer_genome gga_gap (knn).", + }, + "qm9_gap": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "qm9_gap", + "unit": "eV", + "test_mae": 0.031, + "figshare_article_id": 33135167, + "url": "https://ndownloader.figshare.com/files/67163570", + "description": "QM9 HOMO-LUMO gap (kNN).", + }, + "qmof_bandgap": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "qmof_bandgap", + "unit": "eV", + "test_mae": 0.208, + "figshare_article_id": 33135170, + "url": "https://ndownloader.figshare.com/files/67163573", + "description": "QMOF bandgap (kNN).", + }, + "shear_modulus_gv_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "shear_modulus_gv", + "unit": "", + "test_mae": 8.8254, + "figshare_article_id": 33135122, + "url": "https://ndownloader.figshare.com/files/67163510", + "description": "JARVIS-DFT shear_modulus_gv (knn graph).", + }, + "slme_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "slme", + "unit": "", + "test_mae": 4.4475, + "figshare_article_id": 33135128, + "url": "https://ndownloader.figshare.com/files/67163519", + "description": "JARVIS-DFT slme (knn graph).", + }, + "spillage_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "spillage", + "unit": "", + "test_mae": 0.3456, + "figshare_article_id": 33135134, + "url": "https://ndownloader.figshare.com/files/67163525", + "description": "JARVIS-DFT spillage (knn graph).", + }, + "tc_supercon_hydride_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "Tc_supercon_hydride", + "unit": "K", + "test_mae": 10.07, + "figshare_article_id": 33135218, + "url": "https://ndownloader.figshare.com/files/67163624", + "description": "Tc_supercon_hydride (kNN, retrained pure-torch).", + }, + "tc_supercon_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "Tc_supercon", + "unit": "", + "test_mae": 1.49, + "figshare_article_id": 33135212, + "url": "https://ndownloader.figshare.com/files/67163618", + "description": "Tc_supercon (kNN).", + }, + "thermal_cond_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "ltc", + "unit": "", + "test_mae": 0.375, + "figshare_article_id": 33135197, + "url": "https://ndownloader.figshare.com/files/67163603", + "description": "Lattice thermal cond. log10(kL) (kNN).", + }, + "twod_matpd_bandgap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "bandgap", + "unit": "", + "test_mae": 0.366, + "figshare_article_id": 33135245, + "url": "https://ndownloader.figshare.com/files/67163651", + "description": "twod_matpd bandgap (knn).", + }, + "dielectric_knn": { + "category": "spectra", + "graph": "knn", + "cutoff": 8.0, + "output_features": 300, + "target": "dielectric", + "unit": "arb", + "test_mae": 0.0005, + "figshare_article_id": 33135266, + "url": "https://ndownloader.figshare.com/files/67163663", + "description": "TBmBJ dielectric function imag_xx 300-bin 0-15 eV (knn graph).", + }, + "dielectric_radius": { + "category": "spectra", + "graph": "radius", + "cutoff": 5.0, + "output_features": 300, + "target": "dielectric", + "unit": "arb", + "test_mae": 0.0005, + "figshare_article_id": 33135254, + "url": "https://ndownloader.figshare.com/files/67163660", + "description": "TBmBJ dielectric function imag_xx 300-bin 0-15 eV (radius graph).", + }, + "edos": { + "category": "spectra", + "graph": "radius", + "cutoff": 5.0, + "output_features": 300, + "target": "edos", + "unit": "states/eV", + "test_mae": 0.0138, + "figshare_article_id": 33135158, + "url": "https://ndownloader.figshare.com/files/67163561", + "description": "300-bin electronic DOS (radius graph).", + }, + "ir_knn": { + "category": "spectra", + "graph": "knn", + "cutoff": 8.0, + "output_features": 200, + "target": "ir", + "unit": "arb", + "test_mae": 0.023, + "figshare_article_id": 33135251, + "url": "https://ndownloader.figshare.com/files/67163657", + "description": "IR spectrum 200-bin 0-2000 cm^-1 (knn graph).", + }, + "ir_radius": { + "category": "spectra", + "graph": "radius", + "cutoff": 5.0, + "output_features": 200, + "target": "ir", + "unit": "arb", + "test_mae": 0.0232, + "figshare_article_id": 33135248, + "url": "https://ndownloader.figshare.com/files/67163654", + "description": "IR spectrum 200-bin 0-2000 cm^-1 (radius graph).", + }, + "pdos": { + "category": "spectra", + "graph": "radius", + "cutoff": 5.0, + "output_features": 200, + "target": "pdos", + "unit": "states/THz", + "test_mae": 0.0819, + "figshare_article_id": 33135164, + "url": "https://ndownloader.figshare.com/files/67163567", + "description": "200-bin phonon DOS (radius graph).", + }, + "raman_knn": { + "category": "spectra", + "graph": "knn", + "cutoff": 8.0, + "output_features": 200, + "target": "raman_intensity", + "unit": "arb", + "test_mae": 0.0326, + "figshare_article_id": 33134963, + "url": "https://ndownloader.figshare.com/files/67163327", + "description": "200-bin Raman spectrum (kNN graph).", + }, + "born_tensor": { + "category": "tensor", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "born", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135155, + "url": "https://ndownloader.figshare.com/files/67163558", + "description": "born response tensor (D=1).", + }, + "dielectric_tensor": { + "category": "tensor", + "graph": "radius", + "cutoff": 5.0, + "output_features": 9, + "target": "dielectric", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135143, + "url": "https://ndownloader.figshare.com/files/67163546", + "description": "dielectric response tensor (D=9).", + }, + "elastic_tensor": { + "category": "tensor", + "graph": "radius", + "cutoff": 5.0, + "output_features": 36, + "target": "elastic", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135152, + "url": "https://ndownloader.figshare.com/files/67163555", + "description": "elastic response tensor (D=36).", + }, + "piezo_tensor": { + "category": "tensor", + "graph": "radius", + "cutoff": 5.0, + "output_features": 18, + "target": "piezo", + "unit": "", + "test_mae": None, + "figshare_article_id": 33135149, + "url": "https://ndownloader.figshare.com/files/67163552", + "description": "piezo response tensor (D=18).", + }, + "alignn_ff_db": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "", + "test_mae": 0.0564, + "figshare_article_id": 33135200, + "url": "https://ndownloader.figshare.com/files/67163606", + "description": "ALIGNN-FF-DB force field (radius).", + }, + "fd_ff_base": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "", + "test_mae": 0.0445, + "figshare_article_id": 33135203, + "url": "https://ndownloader.figshare.com/files/67163609", + "description": "FD-FF 1.1M base force field (radius).", + }, + "fd_ff_ev": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "eV/atom", + "test_mae": None, + "figshare_article_id": 33134966, + "url": "https://ndownloader.figshare.com/files/67163330", + "description": "FD-FF + EV/vacancy/surface/interface augmented force field (energy/forces/stress, radius graph).", + }, + "matpes_smooth_ff": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "eV/atom", + "test_mae": 0.218, + "figshare_article_id": 33148208, + "url": "https://ndownloader.figshare.com/files/67217507", + "description": "DEFAULT ALIGNN 2.0 force field: 2/2/128, smooth cutoff (multiply_cutoff, inner_cutoff 4.0), nbr52, MATPES-PBE ep100; NVE-stable (Si/MgO/Cu ~CHGNet).", + }, + "matpes_ff": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "eV/atom", + "test_mae": 0.113, + "figshare_article_id": 33134972, + "url": "https://ndownloader.figshare.com/files/67163336", + "description": "Force field trained on MATPES-PBE (keep_data_order, ep100).", + }, + "mlearn_si": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "", + "test_mae": 0.0872, + "figshare_article_id": 33135206, + "url": "https://ndownloader.figshare.com/files/67163612", + "description": "mlearn Si force field.", + }, + "mptrj_ff": { + "category": "forcefield", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "energy_per_atom", + "unit": "eV/atom", + "test_mae": 0.0707, + "figshare_article_id": 33134969, + "url": "https://ndownloader.figshare.com/files/67163333", + "description": "Universal force field trained on MPtrj (~1.5M configs, ep46).", + }, + "charge_atomwise": { + "category": "atomwise", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "charge", + "unit": "e", + "test_mae": None, + "figshare_article_id": 33135137, + "url": "https://ndownloader.figshare.com/files/67163528", + "description": "Per-atom charge (atomwise head).", + }, + "magmom_atomwise": { + "category": "atomwise", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "magmom", + "unit": "muB", + "test_mae": None, + "figshare_article_id": 33135140, + "url": "https://ndownloader.figshare.com/files/67163543", + "description": "Per-atom magmom (atomwise head).", + }, + "net_charge_atomwise": { + "category": "atomwise", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "net_charge", + "unit": "e", + "test_mae": 0.0167, + "figshare_article_id": 33135227, + "url": "https://ndownloader.figshare.com/files/67163633", + "description": "Per-atom net charge (atomwise head).", + }, + "omdb_bandgap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "bandgap", + "unit": "", + "test_mae": 0.2428, + "figshare_article_id": 33135329, + "url": "https://ndownloader.figshare.com/files/67163978", + "description": "omdb bandgap (radius).", + }, + "omdb_bandgap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "bandgap", + "unit": "", + "test_mae": 0.2411, + "figshare_article_id": 33135332, + "url": "https://ndownloader.figshare.com/files/67163981", + "description": "omdb bandgap (knn).", + }, + "pdbbind_binding_affinity_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "binding_affinity", + "unit": "", + "test_mae": 1.5833, + "figshare_article_id": 33135335, + "url": "https://ndownloader.figshare.com/files/67163984", + "description": "pdbbind binding_affinity (radius).", + }, + "pdbbind_binding_affinity_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "binding_affinity", + "unit": "", + "test_mae": 2.7265, + "figshare_article_id": 33135338, + "url": "https://ndownloader.figshare.com/files/67163987", + "description": "pdbbind binding_affinity (knn).", + }, + "halide_peroskites_PBE_gap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "PBE_gap", + "unit": "", + "test_mae": 0.1010, + "figshare_article_id": 33135395, + "url": "https://ndownloader.figshare.com/files/67164095", + "description": "halide_peroskites PBE_gap (radius).", + }, + "halide_peroskites_PBE_gap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "PBE_gap", + "unit": "", + "test_mae": 0.0867, + "figshare_article_id": 33135398, + "url": "https://ndownloader.figshare.com/files/67164098", + "description": "halide_peroskites PBE_gap (knn).", + }, + "halide_peroskites_HSE_gap_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "HSE_gap", + "unit": "", + "test_mae": 0.1397, + "figshare_article_id": 33135401, + "url": "https://ndownloader.figshare.com/files/67164101", + "description": "halide_peroskites HSE_gap (radius).", + }, + "halide_peroskites_HSE_gap_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "HSE_gap", + "unit": "", + "test_mae": 0.1383, + "figshare_article_id": 33135404, + "url": "https://ndownloader.figshare.com/files/67164104", + "description": "halide_peroskites HSE_gap (knn).", + }, + "halide_peroskites_Ref_ind_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "Ref_ind", + "unit": "", + "test_mae": 0.0127, + "figshare_article_id": 33135416, + "url": "https://ndownloader.figshare.com/files/67164140", + "description": "halide_peroskites Ref_ind (radius).", + }, + "halide_peroskites_Ref_ind_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "Ref_ind", + "unit": "", + "test_mae": 0.0127, + "figshare_article_id": 33135419, + "url": "https://ndownloader.figshare.com/files/67164143", + "description": "halide_peroskites Ref_ind (knn).", + }, + "halide_peroskites_PBE_decomp_energy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "PBE_decomp_energy", + "unit": "", + "test_mae": 0.0265, + "figshare_article_id": 33135425, + "url": "https://ndownloader.figshare.com/files/67164152", + "description": "halide_peroskites PBE_decomp_energy (radius).", + }, + "halide_peroskites_PBE_decomp_energy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "PBE_decomp_energy", + "unit": "", + "test_mae": 0.0273, + "figshare_article_id": 33135428, + "url": "https://ndownloader.figshare.com/files/67164155", + "description": "halide_peroskites PBE_decomp_energy (knn).", + }, + "halide_peroskites_HSE_decomp_energy_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "HSE_decomp_energy", + "unit": "", + "test_mae": 0.0233, + "figshare_article_id": 33135434, + "url": "https://ndownloader.figshare.com/files/67164215", + "description": "halide_peroskites HSE_decomp_energy (radius).", + }, + "halide_peroskites_HSE_decomp_energy_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "HSE_decomp_energy", + "unit": "", + "test_mae": 0.0233, + "figshare_article_id": 33135437, + "url": "https://ndownloader.figshare.com/files/67164218", + "description": "halide_peroskites HSE_decomp_energy (knn).", + }, + "snumat_Band_gap_HSE_radius": { + "category": "radius", + "graph": "radius", + "cutoff": 5.0, + "output_features": 1, + "target": "Band_gap_HSE", + "unit": "", + "test_mae": 0.3680, + "figshare_article_id": 33135974, + "url": "https://ndownloader.figshare.com/files/67166069", + "description": "snumat Band_gap_HSE (radius).", + }, + "snumat_Band_gap_HSE_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "Band_gap_HSE", + "unit": "", + "test_mae": 0.3483, + "figshare_article_id": 33135977, + "url": "https://ndownloader.figshare.com/files/67166072", + "description": "snumat Band_gap_HSE (knn).", + }, + "hmof_co2_knn": { + "category": "knn", + "graph": "knn", + "cutoff": 8.0, + "output_features": 1, + "target": "co2_uptake", + "unit": "mol/kg", + "test_mae": 0.4687, + "figshare_article_id": 33137036, + "url": "https://ndownloader.figshare.com/files/67174418", + "description": "hMOF CO2 uptake (knn graph, epoch-99).", + }, + # --- add new models here (one entry each) --- +} + +_API = "https://api.figshare.com/v2" + + +def list_alignn2_models(category=None): + """Return the registry, optionally filtered by category.""" + return { + k: v + for k, v in ALIGNN2_MODELS.items() + if category is None or v["category"] == category + } + + +def resolve_by_target(target, category=None): + """Return registry keys whose ``target`` matches (e.g. a property name).""" + return [ + k + for k, v in ALIGNN2_MODELS.items() + if v.get("target") == target + and (category is None or v["category"] == category) + ] + + +def _cache_dir(): + d = os.path.join(os.path.expanduser("~"), ".alignn2_models") + os.makedirs(d, exist_ok=True) + return d + + +def get_alignn2_model(name, download=True, cache_dir=None): + """Return local paths to a model's artifacts, downloading+caching on first use. + + Downloads the variant's single flat Figshare zip (``config.json``, + ``best_model.pt``, ``ids_train_val_test.json``) via its ``url`` and extracts + it to ``~/.alignn2_models//``. + """ + if name not in ALIGNN2_MODELS: + raise KeyError( + "Unknown model '{}'. Available: {}".format( + name, sorted(ALIGNN2_MODELS) + ) + ) + meta = ALIGNN2_MODELS[name] + dest = os.path.join(cache_dir or _cache_dir(), name) + triplet = ["config.json", "best_model.pt", "ids_train_val_test.json"] + have = os.path.exists( + os.path.join(dest, "best_model.pt") + ) and os.path.exists(os.path.join(dest, "config.json")) + if download and not have: + os.makedirs(dest, exist_ok=True) + url = meta.get("url") + if not url: # unmigrated fallback: fetch the article's zip file url + r = requests.get( + "{}/articles/{}".format(_API, meta["figshare_article_id"]) + ) + r.raise_for_status() + files = { + f["name"]: f["download_url"] for f in r.json().get("files", []) + } + url = files.get(meta.get("zip")) or next(iter(files.values())) + r = requests.get(url, stream=True) + r.raise_for_status() + zp = os.path.join(dest, "_model.zip") + with open(zp, "wb") as o: + for c in r.iter_content(1 << 20): + o.write(c) + with zipfile.ZipFile(zp) as z: + members = z.namelist() + sub = meta.get("subdir", "") + for base in triplet + ["multi_out_predictions.json"]: + hit = [ + n for n in members if n == base or n.endswith("/" + base) + ] + if sub: + hit = [ + n for n in hit if ("%s/%s" % (sub, base)) in n + ] or hit + if hit: + with z.open(hit[0]) as fsrc, open( + os.path.join(dest, base), "wb" + ) as fdst: + fdst.write(fsrc.read()) + os.remove(zp) + return { + f: os.path.join(dest, f) + for f in triplet + if os.path.exists(os.path.join(dest, f)) + } diff --git a/alignn/pure_lmdb_dataset.py b/alignn/pure_lmdb_dataset.py index 35e2159..28b471c 100644 --- a/alignn/pure_lmdb_dataset.py +++ b/alignn/pure_lmdb_dataset.py @@ -69,18 +69,35 @@ def __init__(self, lmdb_path: str = "", line_graph: bool = True, ids=None): self.lmdb_path = lmdb_path self.ids = ids or [] self.line_graph = line_graph - self.env = lmdb.open(self.lmdb_path, readonly=True, lock=False) - with self.env.begin() as txn: + # Open lazily: an lmdb.Environment cannot be pickled, so holding one + # here breaks DataLoader(num_workers>0), which pickles the dataset to + # each worker. Each process opens its own handle on first access. + self.env = None + _env = lmdb.open(self.lmdb_path, readonly=True, lock=False) + with _env.begin() as txn: self.length = txn.stat()["entries"] + _env.close() self.prepare_batch = prepare_pure_batch + def _get_env(self): + """Return this process's LMDB handle, opening it on first use.""" + if self.env is None: + self.env = lmdb.open(self.lmdb_path, readonly=True, lock=False) + return self.env + + def __getstate__(self): + """Drop the unpicklable LMDB handle when sent to a worker.""" + state = self.__dict__.copy() + state["env"] = None + return state + def __len__(self): """Return the number of records in the LMDB.""" return self.length def __getitem__(self, idx): """Load and unpickle the ``idx``-th record.""" - with self.env.begin() as txn: + with self._get_env().begin() as txn: serialized = txn.get(f"{idx}".encode()) if self.line_graph: g, lg, lattice, label = pk.loads(serialized) @@ -125,14 +142,18 @@ def collate_line_graph( def _attach_node_payload( g: TorchGraph, key: str, value: np.ndarray, natoms: int ): - """Tile a per-structure tensor across nodes, like the DGL loader.""" + """Tile a per-structure (global) tensor across nodes, like the DGL loader. + + All callers pass a per-structure quantity (stress 3x3, extra_features, + additional_output) that must be broadcast to every node; genuine per-node + arrays (forces, atomwise targets) are assigned to ``g.ndata`` directly. + We therefore always tile. A previous ``shape[0] == natoms`` early-return + mis-fired for a 3x3 stress on a 3-atom cell (3 == 3), storing it as a + 2-D per-node array and crashing batch collation (``got 2 and 3``) once + the batch mixed 3-atom and non-3-atom structures. + """ dtype = torch.get_default_dtype() arr = np.asarray(value) - # Per-node array already: first dim matches num nodes. - if arr.ndim >= 1 and arr.shape[0] == natoms: - g.ndata[key] = torch.as_tensor(arr, dtype=dtype) - return - # Otherwise, broadcast / tile across nodes. tiled = np.broadcast_to(arr, (natoms,) + arr.shape).copy() g.ndata[key] = torch.as_tensor(tiled, dtype=dtype) diff --git a/alignn/scripts/export_torchscript.py b/alignn/scripts/export_torchscript.py index d01f4a3..7a4e184 100644 --- a/alignn/scripts/export_torchscript.py +++ b/alignn/scripts/export_torchscript.py @@ -10,13 +10,13 @@ Usage ----- + # default: exports the bundled default_path() model (matpes_smooth) + python alignn/scripts/export_torchscript.py --output alignn_scripted.pt + # or point at your own checkpoint/config python alignn/scripts/export_torchscript.py \ --checkpoint /path/to/best_model.pt \ --config /path/to/config.json \ - --output alignn_scripted.pt - # test the scripted file on a random structure - python alignn/scripts/export_torchscript.py \ - --checkpoint ... --config ... --output out.pt --test + --output alignn_scripted.pt --test Design notes ------------ @@ -100,8 +100,7 @@ def _smoke_test(scripted: torch.jit.ScriptModule) -> None: from alignn.graphs import Graph from alignn.torch_graph_builder import torchgraph_from_dgl - atoms = Poscar.from_string( - """Cu + atoms = Poscar.from_string("""Cu 1.0 3.6 0.0 0.0 0.0 3.6 0.0 @@ -113,8 +112,7 @@ def _smoke_test(scripted: torch.jit.ScriptModule) -> None: 0.0 0.5 0.5 0.5 0.0 0.5 0.5 0.5 0.0 -""" - ).atoms +""").atoms g, _ = Graph.atom_dgl_multigraph( atoms, neighbor_strategy="fast_graph", @@ -147,8 +145,20 @@ def _smoke_test(scripted: torch.jit.ScriptModule) -> None: def main() -> None: ap = argparse.ArgumentParser() - ap.add_argument("--checkpoint", required=True, type=Path) - ap.add_argument("--config", required=True, type=Path) + ap.add_argument( + "--checkpoint", + default=None, + type=Path, + help="Model checkpoint (.pt). Defaults to the bundled default_path() " + "model (best_model.pt).", + ) + ap.add_argument( + "--config", + default=None, + type=Path, + help="Training config.json. Defaults to the bundled default_path() " + "model (config.json).", + ) ap.add_argument("--output", required=True, type=Path) ap.add_argument( "--atom-features", @@ -164,6 +174,16 @@ def main() -> None: ) args = ap.parse_args() + # Default to the bundled default_path() model when not given. + if args.checkpoint is None or args.config is None: + from alignn.ff.calculators import default_path + + d = Path(default_path()) + if args.checkpoint is None: + args.checkpoint = d / "best_model.pt" + if args.config is None: + args.config = d / "config.json" + # Resolve atom_features: CLI > training config > 'cgcnn'. cfg = json.load(open(args.config)) atom_features = args.atom_features or cfg.get("atom_features", "cgcnn") diff --git a/alignn/tests/test_alignn_ff.py b/alignn/tests/test_alignn_ff.py index 5855484..d8e13f4 100644 --- a/alignn/tests/test_alignn_ff.py +++ b/alignn/tests/test_alignn_ff.py @@ -66,10 +66,10 @@ def test_graph_builder(): atoms = Poscar.from_string(pos).atoms old_g = Graph.from_atoms(atoms=atoms) - g, lg = Graph.atom_dgl_multigraph(atoms) - g, lg = Graph.atom_dgl_multigraph(atoms, neighbor_strategy="radius_graph") + g, lg = Graph.atom_dgl_multigraph(atoms, neighbor_strategy="pure_torch") + g, lg = Graph.atom_dgl_multigraph(atoms, neighbor_strategy="pure_torch") g, lg = Graph.atom_dgl_multigraph( - atoms, neighbor_strategy="radius_graph_jarvis" + atoms, neighbor_strategy="pure_torch" ) g = radius_graph_old(atoms) diff --git a/alignn/tests/test_eprop.py b/alignn/tests/test_eprop.py index 0be1af8..640a3b5 100644 --- a/alignn/tests/test_eprop.py +++ b/alignn/tests/test_eprop.py @@ -28,7 +28,7 @@ "n_test": 4, "n_val": 4, "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "epochs": 2, "batch_size": 2, "model": { diff --git a/alignn/tests/test_force_reduction.py b/alignn/tests/test_force_reduction.py index f15935e..0215f2e 100644 --- a/alignn/tests/test_force_reduction.py +++ b/alignn/tests/test_force_reduction.py @@ -4,15 +4,31 @@ from torch.nn import functional as F from jarvis.core.atoms import Atoms -import dgl -import dgl.function as fn -from dgl.nn import SumPooling - -from alignn.models.alignn import EdgeGatedGraphConv +# Pure-PyTorch equivalents (no DGL required). ``EdgeGatedGraphConvPure`` is +# the scatter-based edge-gated conv, and ``scatter_sum`` provides the +# segment-sum used for message reduction (replacing dgl.update_all / +# dgl.reverse). The radius graph is built natively in torch below. +from alignn.models.alignn_atomwise_pure import ( + EdgeGatedGraphConvPure, + scatter_sum, +) # double precision for gradient checking torch.set_default_dtype(torch.float64) + +def torch_radius_graph(positions, cutoff): + """Non-periodic radius graph as (src, dst) edge tensors (no self-loops). + + Pure-torch replacement for ``dgl.radius_graph(positions, cutoff)``: + an edge (i, j) exists when 0 < ||r_i - r_j|| <= cutoff. + """ + diff = positions.unsqueeze(1) - positions.unsqueeze(0) # [N, N, 3] + dist = torch.norm(diff, dim=-1) # [N, N] + mask = (dist <= cutoff) & (dist > 0) + src, dst = torch.where(mask) + return src, dst + jvasp_98225_data = { "lattice_mat": [ [7.2963518353359165, 0.0, 0.0], @@ -137,43 +153,39 @@ def __init__(self, cutoff=8, width=16): self.width = width self.edge_embedding = nn.Linear(1, width) - self.hidden1 = EdgeGatedGraphConv(width, width) - self.hidden2 = EdgeGatedGraphConv(width, width) + self.hidden1 = EdgeGatedGraphConvPure(width, width) + self.hidden2 = EdgeGatedGraphConvPure(width, width) self.fc = nn.Linear(width, 1) - self.readout = SumPooling() - def forward(self, positions, autograd_forces=False): # make sure positions are included in the autograd graph if autograd_forces: positions.requires_grad_(True) - # non-periodic radius graph construction - g = dgl.radius_graph(positions, self.cutoff) - g.ndata["r"] = positions + # non-periodic radius graph construction (pure torch) + src, dst = torch_radius_graph(positions, self.cutoff) + num_nodes = positions.shape[0] - # compute bond displacement vectors - g.apply_edges(fn.v_sub_u("r", "r", "bondvec")) - bondvec = g.edata.pop("bondvec") + # compute bond displacement vectors: r_dst - r_src + bondvec = positions[dst] - positions[src] bondlength = torch.norm(bondvec, dim=1).squeeze() # expand bond length basis functions y = self.edge_embedding(bondlength.unsqueeze(-1)) - g.edata["y"] = y # constant node features - x = torch.ones(g.num_nodes(), self.width) + x = torch.ones(num_nodes, self.width) - # graph convolution layers - x, y = self.hidden1(g, x, y) - x, y = self.hidden2(g, x, y) + # graph convolution layers (scatter-based, edge index tensors) + x, y = self.hidden1.forward_tensors(src, dst, num_nodes, x, y) + x, y = self.hidden2.forward_tensors(src, dst, num_nodes, x, y) # node-wise prediction energy = self.fc(x) - # reduction - total_energy = torch.squeeze(self.readout(g, energy)) + # reduction (sum pooling over the single graph) + total_energy = torch.squeeze(energy.sum(dim=0)) if not autograd_forces: return total_energy @@ -190,21 +202,15 @@ def forward(self, positions, autograd_forces=False): # combine r_{ji} and r_{ij} pairwise_forces = -torch.autograd.grad(total_energy, bondvec)[0] - # reduce over bonds to get forces on each atom - g.edata["pairwise_forces"] = pairwise_forces - g.update_all( - fn.copy_e("pairwise_forces", "m"), fn.sum("m", "forces_ji") - ) - - # reduce over reverse edges too! - rg = dgl.reverse(g, copy_edata=True) - rg.update_all( - fn.copy_e("pairwise_forces", "m"), fn.sum("m", "forces_ij") - ) - - forces_vec = torch.squeeze( - g.ndata["forces_ji"] - rg.ndata["forces_ij"] - ) + # reduce over bonds to get forces on each atom. + # forces_ji: sum of pairwise forces over edges arriving at each node + # (dst) -- replaces g.update_all(copy_e, sum). + forces_ji = scatter_sum(pairwise_forces, dst, num_nodes) + # forces_ij: sum over reverse edges (src) -- replaces the + # dgl.reverse(g) + update_all path. + forces_ij = scatter_sum(pairwise_forces, src, num_nodes) + + forces_vec = torch.squeeze(forces_ji - forces_ij) return total_energy, forces_x, forces_vec diff --git a/alignn/tests/test_prop.py b/alignn/tests/test_prop.py index bd68672..c0600a8 100644 --- a/alignn/tests/test_prop.py +++ b/alignn/tests/test_prop.py @@ -28,11 +28,11 @@ "n_test": 4, "n_val": 4, "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "epochs": 2, "batch_size": 2, "model": { - "name": "alignn_atomwise", + "name": "alignn_atomwise_pure", "calculate_gradient": False, "energy_mult_natoms": False, "atom_input_features": 92, diff --git a/alignn/tests/test_unified_calculator.py b/alignn/tests/test_unified_calculator.py new file mode 100644 index 0000000..823d790 --- /dev/null +++ b/alignn/tests/test_unified_calculator.py @@ -0,0 +1,116 @@ +"""Tests for the unified ALIGNN calculator (FF + pure-torch property predictors). + +Covers the force field (energy/forces/stress via matpes_smooth default), a scalar +property (formation_energy_peratom), a spectrum (edos, D=300), a tensor +(elastic_tensor, D=36), the radius/knn graph switch, and config validation. +Models are downloaded+cached on first use (like the other pretrained tests). +""" +import math + +import numpy as np +from ase.build import bulk + +from alignn.ff.unified_calculator import ( + AlignnUnifiedCalculator, + AlignnUnifiedConfig, + _prop2_name, +) + + +def _finite(x): + return all(math.isfinite(v) for v in np.asarray(x).ravel().tolist()) + + +def test_unified_ff_scalar_spectrum_tensor(): + """One calculator returning FF + scalar + spectrum + tensor for Si.""" + cfg = AlignnUnifiedConfig( + energy=True, + forces=True, + stress=True, + properties=[ + "formation_energy_peratom", # scalar + "edos", # spectrum, D=300 + "elastic_tensor", # tensor, D=36 + ], + ) + calc = AlignnUnifiedCalculator(cfg) + + si = bulk("Si", "diamond", a=5.43) + si.calc = calc + + energy = si.get_potential_energy() + forces = si.get_forces() + stress = si.get_stress() + assert math.isfinite(energy) + assert forces.shape == (len(si), 3) and _finite(forces) + assert stress.shape == (6,) and _finite(stress) + # relaxed diamond Si: forces are ~zero by symmetry + assert np.abs(forces).max() < 1e-3 + + preds = calc.predictions() + # scalar: elemental Si formation energy is ~0 + fe = preds["formation_energy_peratom"] + assert isinstance(fe, float) and abs(fe) < 0.5 + # spectrum + assert isinstance(preds["edos"], list) and len(preds["edos"]) == 300 + assert _finite(preds["edos"]) + # tensor (6x6 elastic flattened); C11 in a physical range for Si + et = preds["elastic_tensor"] + assert isinstance(et, list) and len(et) == 36 and _finite(et) + assert 80.0 < et[0] < 260.0 # C11 ~ 160 GPa + + +def test_unified_ff_only_default_model(): + """FF-only path uses the pure-torch matpes_smooth default (no dgl).""" + cfg = AlignnUnifiedConfig() # defaults: energy/forces/stress, no properties + assert cfg.ff_model == "matpes_smooth" + calc = AlignnUnifiedCalculator(cfg) + si = bulk("Si", "diamond", a=5.43) + si.calc = calc + assert math.isfinite(si.get_potential_energy()) + assert calc.predictions() == {} + + +def test_prop_name_resolution_radius_and_knn(): + """Name resolver handles direct names, {name}_{graph}, and graph fallback.""" + assert _prop2_name("formation_energy_peratom", "radius") == \ + "formation_energy_peratom_radius" + assert _prop2_name("formation_energy_peratom", "knn") == \ + "formation_energy_peratom_knn" + assert _prop2_name("elastic_tensor", "radius") == "elastic_tensor" # direct + assert _prop2_name("ir", "knn") == "ir_knn" # {name}_{graph} + assert _prop2_name("raman", "radius") == "raman_knn" # knn-only fallback + + +def test_unified_knn_switch(): + """prop_graph='knn' loads the knn property variant.""" + cfg = AlignnUnifiedConfig( + energy=True, forces=True, stress=True, + prop_graph="knn", properties=["formation_energy_peratom"], + ) + calc = AlignnUnifiedCalculator(cfg) + assert calc._prop_models["formation_energy_peratom"]["name"] == \ + "formation_energy_peratom_knn" + si = bulk("Si", "diamond", a=5.43) + si.calc = calc + si.get_potential_energy() + assert math.isfinite(calc.predictions()["formation_energy_peratom"]) + + +def test_unified_unknown_property_raises(): + """An unknown property is rejected at config time.""" + try: + AlignnUnifiedConfig(properties=["not_a_real_property"]) + except Exception as exc: # pydantic ValidationError wraps the ValueError + assert "not_a_real_property" in str(exc) + else: + raise AssertionError("expected validation error for unknown property") + + +if __name__ == "__main__": + test_prop_name_resolution_radius_and_knn() + test_unified_unknown_property_raises() + test_unified_ff_only_default_model() + test_unified_ff_scalar_spectrum_tensor() + test_unified_knn_switch() + print("all unified calculator tests passed") diff --git a/alignn/train.py b/alignn/train.py index 57b50d6..dfcc0bd 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -158,7 +158,11 @@ def train_dgl( test_loader = train_val_test_loaders[2] prepare_batch = train_val_test_loaders[3] if use_ddp: - device = torch.device(f"cuda:{rank}") + # `rank` is GLOBAL: on a multi-node launch it exceeds the per-node + # device count (and with --gpu-bind=closest each rank sees a single + # GPU), so cuda:{rank} raises "invalid device ordinal". setup() has + # already selected this process's device -- read it back. + device = torch.device(f"cuda:{torch.cuda.current_device()}") prepare_batch = partial(prepare_batch, device=device) if classification: config.model.classification = True @@ -216,7 +220,8 @@ def train_dgl( if use_ddp: net = DDP( net, - device_ids=[rank], + # local device index, not the global rank (see `device` above) + device_ids=[torch.cuda.current_device()], find_unused_parameters=bool( getattr(config, "ddp_find_unused_parameters", False) ), @@ -243,7 +248,9 @@ def train_dgl( scheduler = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=config.learning_rate, - epochs=config.epochs, + # Schedule spans the whole run; lr_total_epochs lets a resumed + # chain keep one continuous cycle (falls back to epochs). + epochs=(config.lr_total_epochs or config.epochs), steps_per_epoch=steps_per_epoch, pct_start=0.3, ) @@ -264,7 +271,26 @@ def train_dgl( # NOTE: optimizer / scheduler intentionally NOT recreated here. history_train = [] history_val = [] - for e in range(config.epochs): + # Resume optimizer/scheduler/epoch (weights are restored separately + # via --restart_model_path). current_state.pt is written each epoch + # below, alongside the pure-weights current_model.pt. + start_epoch = 0 + if config.resume_checkpoint: + state_path = os.path.join(config.output_dir, "current_state.pt") + if os.path.exists(state_path): + ckpt = torch.load(state_path, map_location=device) + optimizer.load_state_dict(ckpt["optimizer"]) + scheduler.load_state_dict(ckpt["scheduler"]) + best_loss = ckpt.get("best_loss", best_loss) + start_epoch = ckpt["epoch"] + if rank == 0: + print( + "Resuming from epoch", + start_epoch, + "best_loss", + best_loss, + ) + for e in range(start_epoch, config.epochs): train_init_time = time.time() running_loss = 0 running_loss1 = 0 @@ -310,7 +336,9 @@ def train_dgl( loss4 = 0 # Such as stresses loss5 = 0 # Such as dos if config.model.output_features is not None: - loss1 = config.model.graphwise_weight * criterion( + loss1 = getattr( + config.model, "graphwise_weight", 1.0 + ) * criterion( result["out"], dats[-1].to(device), ) @@ -321,9 +349,11 @@ def train_dgl( running_loss1 += loss1.item() if ( config.model.atomwise_output_features > 0 - and config.model.atomwise_weight != 0 + and getattr(config.model, "atomwise_weight", 0.0) != 0 ): - loss2 = config.model.atomwise_weight * criterion( + loss2 = getattr( + config.model, "atomwise_weight", 0.0 + ) * criterion( result["atomwise_pred"].to(device), dats[0].ndata["atomwise_target"].to(device), ) @@ -335,8 +365,10 @@ def train_dgl( ) running_loss2 += loss2.item() - if config.model.calculate_gradient: - loss3 = config.model.gradwise_weight * criterion( + if getattr(config.model, "calculate_gradient", False): + loss3 = getattr( + config.model, "gradwise_weight", 0.0 + ) * criterion( result["grad"].to(device), dats[0].ndata["atomwise_grad"].to(device), ) @@ -347,7 +379,7 @@ def train_dgl( result["grad"].cpu().detach().numpy().tolist() ) running_loss3 += loss3.item() - if config.model.stresswise_weight != 0: + if getattr(config.model, "stresswise_weight", 0.0) != 0: targ_stress = torch.stack( [ gg.ndata["stresses"][0] @@ -355,7 +387,9 @@ def train_dgl( ] ).to(device) pred_stress = result["stresses"] - loss4 = config.model.stresswise_weight * criterion( + loss4 = getattr( + config.model, "stresswise_weight", 0.0 + ) * criterion( pred_stress.to(device), targ_stress.to(device), ) @@ -460,9 +494,9 @@ def train_dgl( loss4 = 0 loss5 = 0 if config.model.output_features is not None: - loss1 = config.model.graphwise_weight * criterion( - result["out"], dats[-1].to(device) - ) + loss1 = getattr( + config.model, "graphwise_weight", 1.0 + ) * criterion(result["out"], dats[-1].to(device)) info["target_out"] = dats[-1].cpu().numpy().tolist() info["pred_out"] = ( result["out"].cpu().detach().numpy().tolist() @@ -471,9 +505,11 @@ def train_dgl( if ( config.model.atomwise_output_features > 0 - and config.model.atomwise_weight != 0 + and getattr(config.model, "atomwise_weight", 0.0) != 0 ): - loss2 = config.model.atomwise_weight * criterion( + loss2 = getattr( + config.model, "atomwise_weight", 0.0 + ) * criterion( result["atomwise_pred"].to(device), dats[0].ndata["atomwise_target"].to(device), ) @@ -484,8 +520,10 @@ def train_dgl( result["atomwise_pred"].cpu().detach().numpy().tolist() ) val_loss2 += loss2.item() - if config.model.calculate_gradient: - loss3 = config.model.gradwise_weight * criterion( + if getattr(config.model, "calculate_gradient", False): + loss3 = getattr( + config.model, "gradwise_weight", 0.0 + ) * criterion( result["grad"].to(device), dats[0].ndata["atomwise_grad"].to(device), ) @@ -496,7 +534,7 @@ def train_dgl( result["grad"].cpu().detach().numpy().tolist() ) val_loss3 += loss3.item() - if config.model.stresswise_weight != 0: + if getattr(config.model, "stresswise_weight", 0.0) != 0: targ_stress = torch.stack( [ gg.ndata["stresses"][0] @@ -504,7 +542,9 @@ def train_dgl( ] ).to(device) pred_stress = result["stresses"] - loss4 = config.model.stresswise_weight * criterion( + loss4 = getattr( + config.model, "stresswise_weight", 0.0 + ) * criterion( pred_stress.to(device), targ_stress.to(device), ) @@ -554,6 +594,17 @@ def train_dgl( _unwrap(net).state_dict(), os.path.join(config.output_dir, current_model_name), ) + # Resume state (optimizer/scheduler/next-epoch/best_loss), + # kept separate so current_model.pt stays a pure state_dict. + torch.save( + { + "epoch": e + 1, + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "best_loss": best_loss, + }, + os.path.join(config.output_dir, "current_state.pt"), + ) saving_msg = "" if val_loss < best_loss: best_loss = val_loss @@ -615,6 +666,11 @@ def train_dgl( ) if rank == 0 or world_size == 1: + # This block runs on rank 0 only. Forward through the underlying + # module, not the DDP wrapper: a DDP forward performs collectives + # (buffer broadcast) that the other ranks never reach here, which + # deadlocks rank 0 until the NCCL watchdog times out. + eval_net = _unwrap(net) test_loss = 0 test_result = [] for dats, jid in zip(test_loader, test_loader.dataset.ids): @@ -622,7 +678,7 @@ def train_dgl( info["id"] = jid optimizer.zero_grad() if (config.compute_line_graph) > 0: - result = net( + result = eval_net( [ dats[0].to(device), dats[1].to(device), @@ -630,7 +686,7 @@ def train_dgl( ] ) else: - result = net([dats[0].to(device), dats[1].to(device)]) + result = eval_net([dats[0].to(device), dats[1].to(device)]) loss1 = 0 loss2 = 0 loss3 = 0 @@ -639,16 +695,18 @@ def train_dgl( config.model.output_features is not None and not classification ): - loss1 = config.model.graphwise_weight * criterion( - result["out"], dats[-1].to(device) - ) + loss1 = getattr( + config.model, "graphwise_weight", 1.0 + ) * criterion(result["out"], dats[-1].to(device)) info["target_out"] = dats[-1].cpu().numpy().tolist() info["pred_out"] = ( result["out"].cpu().detach().numpy().tolist() ) if config.model.atomwise_output_features > 0: - loss2 = config.model.atomwise_weight * criterion( + loss2 = getattr( + config.model, "atomwise_weight", 0.0 + ) * criterion( result["atomwise_pred"].to(device), dats[0].ndata["atomwise_target"].to(device), ) @@ -659,8 +717,10 @@ def train_dgl( result["atomwise_pred"].cpu().detach().numpy().tolist() ) - if config.model.calculate_gradient: - loss3 = config.model.gradwise_weight * criterion( + if getattr(config.model, "calculate_gradient", False): + loss3 = getattr( + config.model, "gradwise_weight", 0.0 + ) * criterion( result["grad"].to(device), dats[0].ndata["atomwise_grad"].to(device), ) @@ -670,7 +730,7 @@ def train_dgl( info["pred_grad"] = ( result["grad"].cpu().detach().numpy().tolist() ) - if config.model.stresswise_weight != 0: + if getattr(config.model, "stresswise_weight", 0.0) != 0: targ_stress = torch.stack( [ @@ -679,7 +739,9 @@ def train_dgl( ] ).to(device) pred_stress = result["stresses"] - loss4 = config.model.stresswise_weight * criterion( + loss4 = getattr( + config.model, "stresswise_weight", 0.0 + ) * criterion( pred_stress.to(device), targ_stress.to(device), ) @@ -777,7 +839,7 @@ def train_dgl( config.write_predictions and not classification and config.model.output_features == 1 - and config.model.gradwise_weight == 0 + and getattr(config.model, "gradwise_weight", 0.0) == 0 ): best_model.eval() f = open( diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py index 5b321b0..91765c7 100644 --- a/alignn/train_alignn.py +++ b/alignn/train_alignn.py @@ -3,6 +3,7 @@ """Module to train for a folder with formatted dataset.""" import os +import subprocess import torch.distributed as dist import csv import sys @@ -29,18 +30,55 @@ device = torch.device("cuda") +def _local_rank(rank): + """GPU index on THIS node. + + Under a multi-node launch the global rank exceeds the per-node device + count, so it cannot be used to select a device. Slurm exports + SLURM_LOCALID; fall back to the single-node case where global == local. + """ + n = torch.cuda.device_count() or 1 + # SLURM_LOCALID is only meaningful when there is one process per task + # (the multi-node launch). Under mp.spawn there is a single srun task + # with LOCALID=0 but 8 processes, so `rank` is the local index and + # trusting LOCALID would put every rank on GPU 0. + if int(os.environ.get("SLURM_NTASKS", "1")) > 1: + return int(os.environ.get("SLURM_LOCALID", 0)) % n + return rank % n + + +def _master_addr(): + """Hostname of the coordinating rank. + + Single node -> localhost (unchanged behaviour). Multi-node -> first host + of the allocation, which every rank resolves identically. + """ + nodelist = os.environ.get("SLURM_NODELIST", "") + if int(os.environ.get("SLURM_NNODES", "1")) <= 1 or not nodelist: + return "localhost" + try: # expand e.g. "frontier[01-04]" -> first hostname + out = subprocess.run( + ["scontrol", "show", "hostnames", nodelist], + capture_output=True, text=True, check=True, + ) + return out.stdout.split()[0] + except Exception: + return nodelist.split(",")[0] + + def setup(rank=0, world_size=0, port="12356"): """Set up multi GPU rank.""" # "12356" if port == "": port = str(random.randint(10000, 99999)) if world_size > 1: - os.environ["MASTER_ADDR"] = "localhost" - os.environ["MASTER_PORT"] = port - # os.environ["MASTER_PORT"] = "12355" + os.environ.setdefault("MASTER_ADDR", _master_addr()) + os.environ.setdefault("MASTER_PORT", port) # Initialize the distributed environment. dist.init_process_group("nccl", rank=rank, world_size=world_size) - torch.cuda.set_device(rank) + # Device is selected by LOCAL rank: global rank >= 8 would index a + # GPU that does not exist on this node. + torch.cuda.set_device(_local_rank(rank)) def cleanup(world_size): @@ -449,6 +487,43 @@ def train_for_folder( if __name__ == "__main__": args = parser.parse_args(sys.argv[1:]) + # Multi-node: one process per GCD, launched by srun/torchrun, so ranks + # come from the environment instead of mp.spawn (which cannot cross + # nodes). Requires SLURM_NTASKS > 1, i.e. `srun -n `; + # a single-task launch falls through to the unchanged path below. + _ntasks = int(os.environ.get("SLURM_NTASKS", "1")) + if _ntasks > 1 and "SLURM_PROCID" in os.environ: + rank = int(os.environ["SLURM_PROCID"]) + world_size = _ntasks + print( + f"multi-node launch: global rank {rank}/{world_size} " + f"local {os.environ.get('SLURM_LOCALID', '?')} " + f"on {os.environ.get('SLURMD_NODENAME', '?')}", + flush=True, + ) + train_for_folder( + rank, + world_size, + args.root_dir, + args.config_name, + args.classification_threshold, + args.batch_size, + args.epochs, + args.id_key, + args.target_key, + args.atomwise_key, + args.force_key, + args.stresswise_key, + args.additional_output_key, + args.file_format, + args.restart_model_path, + args.output_dir, + ) + try: + cleanup(world_size) + except Exception: + pass + sys.exit(0) world_size = int(torch.cuda.device_count()) print("world_size", world_size) if world_size > 1: diff --git a/docs/performance.md b/docs/performance.md index 1db02be..10bfaee 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -1,7 +1,7 @@ # 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/). +up-to-date numbers see [JARVIS-Leaderboard](https://atomgptlab.github.io/jarvis_leaderboard/). ## JARVIS-DFT 2021 — classification @@ -87,6 +87,6 @@ OMDB, HOPV, QETB. --- Claims of *best* performance should be verified on the latest -[JARVIS-Leaderboard](https://pages.nist.gov/jarvis_leaderboard/). Numbers from models +[JARVIS-Leaderboard](https://atomgptlab.github.io/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/property-predictor.md b/docs/pretrained/property-predictor.md index 84905ea..f0aa940 100644 --- a/docs/pretrained/property-predictor.md +++ b/docs/pretrained/property-predictor.md @@ -57,15 +57,31 @@ Pass `--file_format` matching your structure file: ## 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 +from ase.build import bulk +from alignn.ff.unified_calculator import ( + AlignnUnifiedCalculator, AlignnUnifiedConfig) + +# Any pretrained ALIGNN 2.0 predictor(s): scalar (formation_energy_peratom, +# optb88vdw_bandgap, ...), spectra (edos, pdos, ir, raman) or tensor +# (dielectric_tensor, elastic_tensor, piezo_tensor). prop_graph is "radius" +# (default, force-field-compatible) or "knn". +cfg = AlignnUnifiedConfig( + energy=False, forces=False, stress=False, # property-only (no force field) + properties=["formation_energy_peratom", "optb88vdw_bandgap"], ) -print(prediction) +calc = AlignnUnifiedCalculator(cfg) # models downloaded+cached, loaded once + +atoms = bulk("Si", "diamond", a=5.43) +calc.calculate(atoms) # single forward pass +print(calc.predictions()) +# {'formation_energy_peratom': 0.0026, 'optb88vdw_bandgap': 0.7725} ``` +The pure-PyTorch ALIGNN 2.0 models are pulled from +[`pretrained2.py`](https://github.com/atomgptlab/alignn/blob/develop/alignn/pretrained2.py) +(`list_alignn2_models()` shows the full registry). Set `energy=True` to also get +force-field energy/forces/stress from the same calculator. + ## See also - [ALIGNN-FF pretrained models](alignn-ff.md) diff --git a/setup.py b/setup.py index 0fcbe58..e5c3771 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setuptools.setup( name="alignn", - version="2026.5.20", + version="2026.8.6", author="Kamal Choudhary, Brian DeCost", author_email="kamal.choudhary@nist.gov", description="alignn",