diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 0bf9aee..64d8dc7 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -2,7 +2,7 @@ name: docs
on:
push:
- branches: [main, develop]
+ branches: [main]
paths:
- docs/**
- mkdocs.yml
diff --git a/README.md b/README.md
index ce28751..66ffbe1 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,7 @@
* [Introduction](#intro)
* [Installation](#install)
* [Examples](#example)
+* [Colab notebooks](#colab)
* [Pre-trained models](#pretrained)
* [JARVIS-ALIGNN webapp](#webapp)
* [ALIGNN-FF & ASE Calculator](#alignnff)
@@ -48,7 +49,24 @@ See [docs/training/](docs/training/) for dataset format and training examples:
- [Force-field training](docs/training/force-field.md)
- [Multi-GPU training](docs/training/multi-gpu.md)
-Google Colab notebooks are linked from [docs/index.md](docs/index.md).
+
+## Colab notebooks
+
+Ready-to-run notebooks covering property prediction, force-field
+training, and pretrained-model usage. Click a badge to open in Colab.
+
+[colab-badge]: https://colab.research.google.com/assets/colab-badge.svg
+
+| Notebook | Open in Colab | Description |
+| --- | --- | --- |
+| Regression task (graph-wise prediction) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Single-output regression for 2D-material exfoliation energies. |
+| ML force-field training from scratch | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Train an ALIGNN-FF force field for Silicon. |
+| ALIGNN-FF: relaxation, EV curve, phonons, interfaces | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Pretrained ALIGNN-FF for relaxation, EV curves, phonons, and interfaces. |
+| Scaling / timing comparison | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | Scaling/timing analysis of universal MLFFs. |
+| Melt-Quench MD | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | Generate amorphous structures via molecular dynamics. |
+| Miscellaneous training tasks | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | Single-output, multi-output (phonon/electron DOS), classification, and pretrained usage. |
+| Superconductor Tc | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | Train a model for superconductor transition temperature. |
+| Build `id_prop.json` from VASP runs | [![Open In Colab][colab-badge]](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | Compile `vasprun.xml` files into `id_prop.json` for ALIGNN-FF training. |
## Using pre-trained models
@@ -66,7 +84,26 @@ See [docs/usage/webapps.md](docs/usage/webapps.md). Direct links: [AtomGPT ALIGN
## ALIGNN-FF ASE Calculator
-See [docs/usage/ase-calculator.md](docs/usage/ase-calculator.md) for example usage.
+```python
+from ase.build import bulk
+from alignn.ff.unified_calculator import (
+ AlignnUnifiedCalculator, AlignnUnifiedConfig)
+
+cfg = AlignnUnifiedConfig(
+ energy=True, forces=True, stress=True,
+ properties=["formation_energy_peratom", "optb88vdw_bandgap"],
+)
+calc = AlignnUnifiedCalculator(cfg) # models loaded once, reused
+
+si = bulk("Si", "diamond", a=5.43); si.calc = calc
+si.get_potential_energy(); si.get_forces(); si.get_stress()
+print(calc.predictions()) # extra property predictors
+```
+
+A single pydantic config selects the outputs (force-field
+energy/forces/stress plus any pretrained scalar property predictors).
+See [docs/usage/ase-calculator.md](docs/usage/ase-calculator.md) for more,
+and the ASE docs page *Calculators → ALIGNN*.
## Performances
diff --git a/alignn/__init__.py b/alignn/__init__.py
index 0c35900..e117a37 100644
--- a/alignn/__init__.py
+++ b/alignn/__init__.py
@@ -1,3 +1,3 @@
"""Version number."""
-__version__ = "2026.4.1"
+__version__ = "2026.5.20"
diff --git a/alignn/config.py b/alignn/config.py
index 3a5cad7..6904417 100644
--- a/alignn/config.py
+++ b/alignn/config.py
@@ -7,6 +7,7 @@
from alignn.utils import BaseSettings
from alignn.models.alignn import ALIGNNConfig
from alignn.models.alignn_atomwise import ALIGNNAtomWiseConfig
+from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePureConfig
from alignn.models.ealignn_atomwise import eALIGNNAtomWiseConfig
# import torch
@@ -155,7 +156,13 @@ class TrainingConfig(BaseSettings):
target: TARGET_ENUM = "exfoliation_energy"
atom_features: Literal["basic", "atomic_number", "cfid", "cgcnn"] = "cgcnn"
neighbor_strategy: Literal[
- "k-nearest", "voronoi", "radius_graph", "radius_graph_jarvis"
+ "k-nearest",
+ "voronoi",
+ "radius_graph",
+ "radius_graph_jarvis",
+ "fast_graph",
+ "torch_graph",
+ "pure_torch",
] = "k-nearest"
id_tag: Literal["jid", "id", "_oqmd_entry_id"] = "jid"
@@ -173,6 +180,14 @@ class TrainingConfig(BaseSettings):
target_multiplication_factor: Optional[float] = None
epochs: int = 300
batch_size: int = 64
+ gpu_memory_fraction: Optional[float] = None
+ use_amp: bool = False # bf16 mixed precision (A100/H100/RTX 30+)
+ # 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
+ # When True, forces cuDNN determinism (slower). Decoupled from seed so
+ # you can seed for reproducibility without paying the speed cost.
+ deterministic: bool = False
weight_decay: float = 0
learning_rate: float = 1e-2
filename: str = "sample"
@@ -191,9 +206,12 @@ class TrainingConfig(BaseSettings):
use_canonize: bool = True
compute_line_graph: bool = True
num_workers: int = 4
- cutoff: float = 8.0
+ cutoff: float = 5.0
cutoff_extra: float = 3.0
- max_neighbors: int = 12
+ # Separate 3-body cutoff used by neighbor_strategy="pure_torch".
+ # When None, defaults to `cutoff`. Must be <= cutoff.
+ three_body_cutoff: Optional[float] = 3.5
+ max_neighbors: Optional[int] = 12
keep_data_order: bool = True
normalize_graph_level_loss: bool = False
distributed: bool = False
@@ -201,6 +219,10 @@ class TrainingConfig(BaseSettings):
n_early_stopping: Optional[int] = None # typically 50
output_dir: str = os.path.abspath(".")
use_lmdb: bool = True
+ # When True, reuse an existing LMDB cache on disk (fast, but the
+ # cache must match the current model backend + graph config).
+ # Default False: always rebuild from scratch for safety.
+ read_existing: bool = False
# alignn_layers: int = 4
# gcn_layers: int =4
# edge_input_features: int= 80
@@ -212,5 +234,6 @@ class TrainingConfig(BaseSettings):
model: Union[
ALIGNNConfig,
ALIGNNAtomWiseConfig,
+ ALIGNNAtomWisePureConfig,
eALIGNNAtomWiseConfig,
] = ALIGNNAtomWiseConfig(name="alignn_atomwise")
diff --git a/alignn/data.py b/alignn/data.py
index 39f2fd0..97fcbe8 100644
--- a/alignn/data.py
+++ b/alignn/data.py
@@ -10,7 +10,18 @@
from tqdm import tqdm
import math
from jarvis.db.jsonutils import dumpjson
-from dgl.dataloading import GraphDataLoader
+
+try:
+ from dgl.dataloading import GraphDataLoader
+except Exception: # pure-torch path; fall back to stdlib DataLoader
+ from torch.utils.data import DataLoader as _TorchDataLoader
+
+ def GraphDataLoader(*args, use_ddp=False, **kwargs):
+ """Fall back to torch DataLoader when DGL is unavailable."""
+ kwargs.pop("use_ddp", None)
+ return _TorchDataLoader(*args, **kwargs)
+
+
import pickle as pk
from sklearn.preprocessing import StandardScaler
@@ -145,6 +156,7 @@ def get_train_val_loaders(
cutoff: float = 8.0,
cutoff_extra: float = 3.0,
max_neighbors: int = 12,
+ three_body_cutoff: Optional[float] = None,
classification_threshold: Optional[float] = None,
target_multiplication_factor: Optional[float] = None,
standard_scalar_and_pca=False,
@@ -154,10 +166,15 @@ def get_train_val_loaders(
world_size=0,
rank=0,
use_lmdb: bool = True,
+ use_pure_torch: bool = False,
+ read_existing: bool = False,
dtype="float32",
):
"""Help function to set up JARVIS train and val dataloaders."""
- if use_lmdb:
+ if use_pure_torch:
+ print("Using pure-torch LMDB dataset (no DGL).")
+ from alignn.pure_lmdb_dataset import get_torch_dataset
+ elif use_lmdb:
print("Using LMDB dataset.")
from alignn.lmdb_dataset import get_torch_dataset
else:
@@ -383,6 +400,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,
sampler=train_sampler,
@@ -409,6 +428,7 @@ def get_train_val_loaders(
cutoff_extra=cutoff_extra,
sampler=val_sampler,
max_neighbors=max_neighbors,
+ three_body_cutoff=three_body_cutoff,
classification=classification_threshold is not None,
output_dir=output_dir,
tmp_name=tmp_name,
@@ -462,7 +482,20 @@ def get_train_val_loaders(
num_workers=workers,
pin_memory=pin_memory,
use_ddp=use_ddp,
+ persistent_workers=(workers > 0),
+ prefetch_factor=2 if workers > 0 else None,
)
+ # train_loader = GraphDataLoader(
+ # # train_loader = DataLoader(
+ # train_data,
+ # batch_size=batch_size,
+ # shuffle=True,
+ # collate_fn=collate_fn,
+ # drop_last=True,
+ # num_workers=workers,
+ # pin_memory=pin_memory,
+ # use_ddp=use_ddp,
+ # )
val_loader = GraphDataLoader(
# val_loader = DataLoader(
@@ -474,7 +507,20 @@ def get_train_val_loaders(
num_workers=workers,
pin_memory=pin_memory,
use_ddp=use_ddp,
+ persistent_workers=(workers > 0),
+ prefetch_factor=2 if workers > 0 else None,
)
+ # val_loader = GraphDataLoader(
+ # # val_loader = DataLoader(
+ # val_data,
+ # batch_size=batch_size,
+ # shuffle=False,
+ # collate_fn=collate_fn,
+ # drop_last=True,
+ # num_workers=workers,
+ # pin_memory=pin_memory,
+ # use_ddp=use_ddp,
+ # )
test_loader = (
GraphDataLoader(
@@ -488,10 +534,27 @@ def get_train_val_loaders(
pin_memory=pin_memory,
use_ddp=use_ddp,
)
- if len(dataset_test) > 0
+ if test_data is not None
else None
)
+ # All loaders constructed; safe to free the in-memory dataset
+ # lists. With 1.5M-entry datasets these can hold tens of GB of
+ # full atomic-structure dicts, and they are no longer needed
+ # once the LMDB caches are built and the loaders reference
+ # them via the Dataset wrappers.
+ del dataset_train
+ del dataset_val
+ del dataset_test
+ del dat
+ try:
+ del d
+ except NameError:
+ pass
+ import gc
+
+ gc.collect()
+
if save_dataloader:
torch.save(train_loader, train_sample)
if val_loader is not None:
diff --git a/alignn/dataset.py b/alignn/dataset.py
index 0d2e88b..f9bcb35 100644
--- a/alignn/dataset.py
+++ b/alignn/dataset.py
@@ -4,7 +4,10 @@
from typing import Optional
import os
import torch
-import dgl
+try:
+ import dgl
+except ImportError:
+ dgl = None
import numpy as np
import pandas as pd
from jarvis.core.atoms import Atoms
@@ -23,6 +26,7 @@ def load_graphs(
cutoff: float = 8,
cutoff_extra: float = 3,
max_neighbors: int = 12,
+ three_body_cutoff: Optional[float] = None,
cachedir: Optional[Path] = None,
use_canonize: bool = False,
id_tag="jid",
@@ -87,6 +91,7 @@ def atoms_to_graph(atoms):
compute_line_graph=False,
use_canonize=use_canonize,
neighbor_strategy=neighbor_strategy,
+ three_body_cutoff=three_body_cutoff,
id=i[id_tag],
dtype=dtype,
)
@@ -125,10 +130,12 @@ def get_torch_dataset(
cutoff=8.0,
cutoff_extra=3.0,
max_neighbors=12,
+ three_body_cutoff=None,
classification=False,
output_dir=".",
tmp_name="dataset",
sampler=None,
+ read_existing=False, # accepted for API parity; no-op (no cache here)
dtype="float32",
):
"""Get Torch Dataset."""
@@ -152,6 +159,7 @@ def get_torch_dataset(
cutoff=cutoff,
cutoff_extra=cutoff_extra,
max_neighbors=max_neighbors,
+ three_body_cutoff=three_body_cutoff,
id_tag=id_tag,
dtype=dtype,
)
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json
index a1baace..8eb9da7 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json
@@ -1,65 +1,66 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 31,
- "n_test": 31,
- "n_train": 262,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 13,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "alignn_layers": 2,
- "gcn_layers": 4,
- "atom_input_features": 92,
- "edge_input_features": 80,
- "triplet_input_features": 40,
- "embedding_features": 164,
- "hidden_features": 256,
- "output_features": 1,
- "force_mult_natoms": true,
- "grad_multiplier": -1,
- "calculate_gradient": true,
- "atomwise_output_features": 0,
- "graphwise_weight": 1.0,
- "gradwise_weight": 50.00,
- "stresswise_weight": 0.0,
- "atomwise_weight": 0.0,
- "link": "identity",
- "zero_inflated": false,
- "classification": false
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 31,
+ "n_test": 31,
+ "n_train": 262,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise_pure",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "atom_input_features": 92,
+ "edge_input_features": 80,
+ "triplet_input_features": 40,
+ "embedding_features": 164,
+ "hidden_features": 256,
+ "output_features": 1,
+ "force_mult_natoms": true,
+ "grad_multiplier": -1,
+ "calculate_gradient": true,
+ "atomwise_output_features": 0,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 50.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json
new file mode 100644
index 0000000..6acbfba
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json
@@ -0,0 +1,66 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 31,
+ "n_test": 31,
+ "n_train": 262,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "atom_input_features": 92,
+ "edge_input_features": 80,
+ "triplet_input_features": 40,
+ "embedding_features": 164,
+ "hidden_features": 256,
+ "output_features": 1,
+ "force_mult_natoms": true,
+ "grad_multiplier": -1,
+ "calculate_gradient": true,
+ "atomwise_output_features": 0,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 50.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json
index 303f1fc..dbb5b1d 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json
@@ -1,65 +1,66 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 25,
- "n_test": 25,
- "n_train": 228,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 12,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "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": 10.0,
- "stresswise_weight": 0.0,
- "atomwise_weight": 0.0,
- "link": "identity",
- "force_mult_natoms": true,
- "zero_inflated": false,
- "classification": false
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 25,
+ "n_test": 25,
+ "n_train": 228,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise_pure",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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": 10.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "force_mult_natoms": true,
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json
new file mode 100644
index 0000000..ce13c38
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json
@@ -0,0 +1,66 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 25,
+ "n_test": 25,
+ "n_train": 228,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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": 10.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "force_mult_natoms": true,
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json
index 5efce30..0019f0c 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json
@@ -1,65 +1,66 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 29,
- "n_test": 29,
- "n_train": 241,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 12,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "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": 10.0,
- "stresswise_weight": 0.0,
- "force_mult_natoms": true,
- "atomwise_weight": 0.0,
- "link": "identity",
- "zero_inflated": false,
- "classification": false
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 29,
+ "n_test": 29,
+ "n_train": 241,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise_pure",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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": 10.0,
+ "stresswise_weight": 0.0,
+ "force_mult_natoms": true,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json
new file mode 100644
index 0000000..19ebaa1
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json
@@ -0,0 +1,66 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 29,
+ "n_test": 29,
+ "n_train": 241,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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": 10.0,
+ "stresswise_weight": 0.0,
+ "force_mult_natoms": true,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json
index 639acdb..4742928 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json
@@ -1,65 +1,66 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 23,
- "n_test": 23,
- "n_train": 194,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 12,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "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": 10.0,
- "stresswise_weight": 0.0,
- "atomwise_weight": 0.0,
- "link": "identity",
- "zero_inflated": false,
- "force_mult_natoms": true,
- "classification": false
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 23,
+ "n_test": 23,
+ "n_train": 194,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise_pure",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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": 10.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "force_mult_natoms": true,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json
new file mode 100644
index 0000000..eb26eb1
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json
@@ -0,0 +1,66 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 23,
+ "n_test": 23,
+ "n_train": 194,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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": 10.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "force_mult_natoms": true,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json
index 953555a..c592951 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json
@@ -1,65 +1,66 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 31,
- "n_test": 31,
- "n_train": 263,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 12,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "alignn_layers": 2,
- "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,
- "force_mult_natoms": true,
- "calculate_gradient": true,
- "atomwise_output_features": 0,
- "graphwise_weight": 1.0,
- "gradwise_weight": 1.0,
- "stresswise_weight": 0.0,
- "atomwise_weight": 0.0,
- "link": "identity",
- "zero_inflated": false,
- "classification": false
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 31,
+ "n_test": 31,
+ "n_train": 263,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise_pure",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "atom_input_features": 92,
+ "edge_input_features": 80,
+ "triplet_input_features": 40,
+ "embedding_features": 64,
+ "hidden_features": 256,
+ "output_features": 1,
+ "grad_multiplier": -1,
+ "force_mult_natoms": true,
+ "calculate_gradient": true,
+ "atomwise_output_features": 0,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 1.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json
new file mode 100644
index 0000000..6841663
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json
@@ -0,0 +1,66 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 31,
+ "n_test": 31,
+ "n_train": 263,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "atom_input_features": 92,
+ "edge_input_features": 80,
+ "triplet_input_features": 40,
+ "embedding_features": 64,
+ "hidden_features": 256,
+ "output_features": 1,
+ "grad_multiplier": -1,
+ "force_mult_natoms": true,
+ "calculate_gradient": true,
+ "atomwise_output_features": 0,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 1.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json
deleted file mode 100644
index 24bf1da..0000000
--- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json
+++ /dev/null
@@ -1,68 +0,0 @@
-{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "radius_graph",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 31,
- "n_test": 31,
- "n_train": 263,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 4.0,
- "max_neighbors": 12,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "alignn_layers": 2,
- "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,
- "force_mult_natoms": true,
- "calculate_gradient": true,
- "atomwise_output_features": 0,
- "graphwise_weight": 1.0,
- "gradwise_weight": 10.0,
- "stresswise_weight": 0.0,
- "atomwise_weight": 0.0,
- "link": "identity",
- "use_cutoff_function": true,
- "zero_inflated": false,
- "classification": false,
- "inner_cutoff": 3.0,
- "lg_on_fly": false
- }
-}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json
index 4e9d7f5..b91de03 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json
@@ -1,66 +1,67 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 25,
- "n_test": 25,
- "n_train": 214,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 12,
- "keep_data_order": true,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "./",
- "model": {
- "name": "alignn_atomwise",
- "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.0,
- "atomwise_weight": 0.0,
- "force_mult_natoms": true,
- "link": "identity",
- "zero_inflated": false,
- "classification": false,
- "use_cutoff_function":true
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 25,
+ "n_test": 25,
+ "n_train": 214,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "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.0,
+ "atomwise_weight": 0.0,
+ "force_mult_natoms": true,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false,
+ "use_cutoff_function": true
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json
new file mode 100644
index 0000000..de2122c
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json
@@ -0,0 +1,67 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 25,
+ "n_test": 25,
+ "n_train": 214,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": true,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "./",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "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.0,
+ "atomwise_weight": 0.0,
+ "force_mult_natoms": true,
+ "link": "identity",
+ "zero_inflated": false,
+ "classification": false,
+ "use_cutoff_function": true
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json
index 1529ec8..d927e8f 100644
--- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json
+++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json
@@ -1,68 +1,69 @@
{
- "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
- "dataset": "user_data",
- "target": "target",
- "atom_features": "cgcnn",
- "neighbor_strategy": "k-nearest",
- "id_tag": "jid",
- "random_seed": 123,
- "classification_threshold": null,
- "n_val": 164,
- "n_test": 164,
- "n_train":1402,
- "train_ratio": 0.9,
- "val_ratio": 0.05,
- "test_ratio": 0.05,
- "target_multiplication_factor": null,
- "epochs": 50,
- "batch_size": 5,
- "weight_decay": 1e-05,
- "learning_rate": 0.001,
- "filename": "sample",
- "warmup_steps": 2000,
- "criterion": "l1",
- "optimizer": "adamw",
- "scheduler": "onecycle",
- "pin_memory": false,
- "save_dataloader": false,
- "write_checkpoint": true,
- "write_predictions": true,
- "store_outputs": false,
- "progress": true,
- "log_tensorboard": false,
- "standard_scalar_and_pca": false,
- "use_canonize": false,
- "num_workers": 0,
- "cutoff": 8.0,
- "max_neighbors": 12,
- "keep_data_order": false,
- "normalize_graph_level_loss": false,
- "distributed": false,
- "n_early_stopping": null,
- "output_dir": "out_continue",
- "model": {
- "name": "alignn_atomwise",
- "alignn_layers": 2,
- "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,
- "force_mult_natoms": true,
- "calculate_gradient": true,
- "atomwise_output_features": 0,
- "graphwise_weight": 1.0,
- "gradwise_weight": 50.0,
- "stresswise_weight": 0.0,
- "atomwise_weight": 0.0,
- "link": "identity",
- "zero_inflated": false,
- "use_cutoff_function": true,
- "energy_mult_natoms": false,
- "classification": false,
- "stress_multiplier":1
- }
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "pure_torch",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 164,
+ "n_test": 164,
+ "n_train": 1402,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": false,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "out_continue",
+ "model": {
+ "name": "alignn_atomwise_pure",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "atom_input_features": 92,
+ "edge_input_features": 80,
+ "triplet_input_features": 40,
+ "embedding_features": 64,
+ "hidden_features": 256,
+ "output_features": 1,
+ "grad_multiplier": -1,
+ "force_mult_natoms": true,
+ "calculate_gradient": true,
+ "atomwise_output_features": 0,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 50.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "use_cutoff_function": true,
+ "energy_mult_natoms": false,
+ "classification": false,
+ "stress_multiplier": 1
+ },
+ "three_body_cutoff": 3.5
}
diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json
new file mode 100644
index 0000000..165639e
--- /dev/null
+++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json
@@ -0,0 +1,69 @@
+{
+ "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "radius_graph",
+ "id_tag": "jid",
+ "random_seed": 123,
+ "classification_threshold": null,
+ "n_val": 164,
+ "n_test": 164,
+ "n_train": 1402,
+ "train_ratio": 0.9,
+ "val_ratio": 0.05,
+ "test_ratio": 0.05,
+ "target_multiplication_factor": null,
+ "epochs": 50,
+ "batch_size": 2,
+ "weight_decay": 1e-05,
+ "learning_rate": 0.001,
+ "filename": "sample",
+ "warmup_steps": 2000,
+ "criterion": "l1",
+ "optimizer": "adamw",
+ "scheduler": "onecycle",
+ "pin_memory": false,
+ "save_dataloader": false,
+ "write_checkpoint": true,
+ "write_predictions": true,
+ "store_outputs": false,
+ "progress": true,
+ "log_tensorboard": false,
+ "standard_scalar_and_pca": false,
+ "use_canonize": false,
+ "num_workers": 0,
+ "cutoff": 5.0,
+ "max_neighbors": null,
+ "keep_data_order": false,
+ "normalize_graph_level_loss": false,
+ "distributed": false,
+ "n_early_stopping": null,
+ "output_dir": "out_continue",
+ "model": {
+ "name": "alignn_atomwise",
+ "alignn_layers": 2,
+ "gcn_layers": 2,
+ "atom_input_features": 92,
+ "edge_input_features": 80,
+ "triplet_input_features": 40,
+ "embedding_features": 64,
+ "hidden_features": 256,
+ "output_features": 1,
+ "grad_multiplier": -1,
+ "force_mult_natoms": true,
+ "calculate_gradient": true,
+ "atomwise_output_features": 0,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 50.0,
+ "stresswise_weight": 0.0,
+ "atomwise_weight": 0.0,
+ "link": "identity",
+ "zero_inflated": false,
+ "use_cutoff_function": true,
+ "energy_mult_natoms": false,
+ "classification": false,
+ "stress_multiplier": 1
+ },
+ "three_body_cutoff": 3.5
+}
diff --git a/alignn/ff/all_models_alignn_atomwise.json b/alignn/ff/all_models_alignn_atomwise.json
index 54a0165..f0c1465 100644
--- a/alignn/ff/all_models_alignn_atomwise.json
+++ b/alignn/ff/all_models_alignn_atomwise.json
@@ -1,4 +1,33 @@
{
+ "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",
diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index ebc1e6c..2312ef0 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -14,6 +14,10 @@
eALIGNNAtomWiseConfig,
)
from alignn.models.alignn import ALIGNN, ALIGNNConfig
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure,
+ ALIGNNAtomWisePureConfig,
+)
import matplotlib.pyplot as plt # noqa
import zipfile
import numpy as np
@@ -248,17 +252,24 @@ def __init__(
)
if self.model is None:
- if self.config["model"]["name"] == "alignn_atomwise":
+ model = None
+ mname = self.config["model"]["name"]
+ if mname == "alignn_atomwise":
model = ALIGNNAtomWise(
ALIGNNAtomWiseConfig(**self.config["model"])
)
- elif self.config["model"]["name"] == "alignn":
+ elif mname == "alignn_atomwise_pure":
+ model = ALIGNNAtomWisePure(
+ ALIGNNAtomWisePureConfig(**self.config["model"])
+ )
+ elif mname == "alignn":
model = ALIGNN(ALIGNNConfig(**self.config["model"]))
- elif self.config["model"]["name"] == "ealignn_atomwise":
+ elif mname == "ealignn_atomwise":
model = eALIGNNAtomWise(
eALIGNNAtomWiseConfig(**self.config["model"])
)
- model.state_dict()
+ if model is None:
+ raise ValueError(f"Unsupported model name '{mname}' in config")
if "atomwise" in self.config["model"]["name"]:
model.load_state_dict(
torch.load(
@@ -362,7 +373,8 @@ def calculate(self, atoms, properties=None, system_changes=None):
# / 160.21766208
# )
if "atomwise" in self.config["model"]["name"]:
- energy = result["out"].detach().cpu().numpy()
+ # energy = result["out"].squeeze().detach().cpu().numpy()
+ energy = result["out"].detach().cpu().item()
else:
energy = result.detach().cpu().numpy()
if self.intensive:
@@ -378,6 +390,7 @@ def calculate(self, atoms, properties=None, system_changes=None):
"forces": forces,
"stress": stress,
}
+ # print("self.results",self.results)
class iAlignnAtomwiseCalculator(ase.calculators.calculator.Calculator):
diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index b065e4f..eac85f7 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -48,6 +48,7 @@
from sklearn.metrics import mean_absolute_error
from tqdm import tqdm
from jarvis.core.utils import get_cache_dir
+
# import torch
from alignn.ff.calculators import (
AlignnAtomwiseCalculator,
@@ -207,7 +208,8 @@ def get_figshare_model_prop(
def default_path():
"""Get default model path."""
- dpath = get_figshare_model_ff(model_name="v12.2.2024_dft_3d_307k")
+ dpath = get_figshare_model_ff(model_name="mps")
+ # 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")
# dpath = get_figshare_model_ff(model_name="alignnff_wt10")
diff --git a/alignn/ff/lammps_bridge.py b/alignn/ff/lammps_bridge.py
index cb8cf05..3bebad4 100644
--- a/alignn/ff/lammps_bridge.py
+++ b/alignn/ff/lammps_bridge.py
@@ -133,7 +133,7 @@ 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}
+ fix tfix all langevin {args.temp} {args.temp} 0.1 {args.seed} # noqa: E501
thermo 10
thermo_style custom step temp pe ke etotal press
"""
diff --git a/alignn/ff/unified_calculator.py b/alignn/ff/unified_calculator.py
new file mode 100644
index 0000000..d144354
--- /dev/null
+++ b/alignn/ff/unified_calculator.py
@@ -0,0 +1,247 @@
+"""Unified ALIGNN calculator driven by a pydantic config.
+
+One calculator, one config object. The config selects which outputs are
+produced:
+
+ * force-field properties : energy / forces / stress (one ALIGNN-FF
+ model)
+ * scalar property predictors : formation energy, band gap, bulk
+ modulus, ... (one pretrained ALIGNN regression model per property)
+
+All models are loaded once and reused for every structure / every call
+(no per-call reload).
+
+Example
+-------
+ from alignn.ff.unified_calculator import (
+ AlignnUnifiedCalculator, AlignnUnifiedConfig)
+ from ase.build import bulk
+
+ cfg = AlignnUnifiedConfig(
+ energy=True, forces=True, stress=True,
+ properties=["formation_energy_peratom",
+ "optb88vdw_bandgap", "bulk_modulus_kv"],
+ )
+ calc = AlignnUnifiedCalculator(cfg)
+
+ si = bulk("Si", "diamond", a=5.43); si.calc = calc
+ si.get_potential_energy(); si.get_forces(); si.get_stress()
+ print(calc.results["formation_energy_peratom"],
+ calc.results["optb88vdw_bandgap"],
+ calc.results["bulk_modulus_kv"])
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+import torch
+from ase.calculators.calculator import Calculator, all_changes
+from pydantic import BaseModel, Field, field_validator
+
+from alignn.ff.calculators import (
+ AlignnAtomwiseCalculator,
+ 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)."
+ )
+
+
+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"
+ energy: bool = True
+ forces: bool = True
+ stress: bool = True
+ charges: bool = False # only if the FF model emits them
+
+ # scalar property predictors to also evaluate (friendly names)
+ properties: List[str] = Field(default_factory=list)
+
+ # shared knobs
+ device: Optional[str] = None
+ prop_cutoff: float = 8.0
+ prop_max_neighbors: int = 12
+
+ model_config = {"extra": "forbid"}
+
+ @field_validator("properties")
+ @classmethod
+ def _check_props(cls, v: List[str]) -> List[str]:
+ for p in v:
+ _resolve_prop(p) # raises on unknown
+ return v
+
+ @classmethod
+ def from_file(cls, path: str) -> "AlignnUnifiedConfig":
+ """Load from a JSON file."""
+ import json
+
+ with open(path) as fh:
+ return cls(**json.load(fh))
+
+
+class AlignnUnifiedCalculator(Calculator):
+ """ASE calculator: ALIGNN-FF + selected ALIGNN property predictors.
+
+ Args:
+ config: an ``AlignnUnifiedConfig``, a dict, a path to a JSON
+ config, or None (defaults: energy/forces/stress only).
+ ff_path: optional local path/dir for the FF model (overrides
+ ``config.ff_model``).
+ """
+
+ def __init__(self, config=None, ff_path=None, **kw):
+ Calculator.__init__(self, **kw)
+
+ if config is None:
+ config = AlignnUnifiedConfig()
+ elif isinstance(config, str):
+ config = AlignnUnifiedConfig.from_file(config)
+ elif isinstance(config, dict):
+ config = AlignnUnifiedConfig(**config)
+ self.cfg: AlignnUnifiedConfig = config
+
+ self.device = config.device or (
+ "cuda" if torch.cuda.is_available() else "cpu"
+ )
+
+ props = []
+ if config.energy:
+ props.append("energy")
+ if config.forces:
+ props.append("forces")
+ if config.stress:
+ props.append("stress")
+ self.implemented_properties = props
+
+ # --- 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
+ )
+ 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] = {}
+ for friendly in config.properties:
+ self._prop_models[friendly] = get_figshare_model(
+ _resolve_prop(friendly)
+ )
+
+ # -- scalar property forward (reuses cached model, no reload) -------
+ def _predict_scalar(self, model, j_atoms) -> float:
+ g, lg = Graph.atom_dgl_multigraph(
+ j_atoms,
+ cutoff=float(self.cfg.prop_cutoff),
+ max_neighbors=self.cfg.prop_max_neighbors,
+ )
+ lat = torch.tensor(j_atoms.lattice_mat)
+ with torch.no_grad():
+ out = model(
+ [
+ g.to(self.device),
+ lg.to(self.device),
+ lat.to(self.device),
+ ]
+ )
+ if isinstance(out, dict):
+ out = out["out"]
+ arr = out.detach().cpu().numpy().flatten()
+ return float(arr[0]) if arr.size == 1 else arr.tolist()
+
+ def calculate(
+ self, atoms=None, properties=("energy",), system_changes=all_changes
+ ):
+ Calculator.calculate(self, atoms, properties, system_changes)
+ self.results = {}
+
+ # force-field block
+ if self._ff is not None:
+ self._ff.calculate(
+ self.atoms, properties=properties,
+ system_changes=system_changes,
+ )
+ if self.cfg.energy:
+ self.results["energy"] = self._ff.results["energy"]
+ self.results["free_energy"] = self._ff.results["energy"]
+ if self.cfg.forces:
+ self.results["forces"] = self._ff.results["forces"]
+ if self.cfg.stress:
+ self.results["stress"] = self._ff.results["stress"]
+ if self.cfg.charges:
+ ch = self._ff.results.get("charges")
+ if ch is not None:
+ self.results["charges"] = ch
+ else:
+ self.results.setdefault("warnings", []).append(
+ "charges requested but the FF model does not "
+ "emit them"
+ )
+
+ # 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
+ )
+
+ # convenience
+ def predictions(self) -> Dict[str, object]:
+ """All scalar property predictions from the last calculate()."""
+ return {
+ k: self.results[k]
+ for k in self.cfg.properties
+ if k in self.results
+ }
+
+
+def _ff_model_path(name_or_path: str) -> str:
+ """Accept a local dir or a known FF figshare model name."""
+ import os
+
+ if os.path.isdir(name_or_path):
+ return name_or_path
+ from alignn.ff.calculators import get_figshare_model_ff
+
+ return get_figshare_model_ff(model_name=name_or_path)
diff --git a/alignn/graphs.py b/alignn/graphs.py
index a017aa2..0146d2e 100644
--- a/alignn/graphs.py
+++ b/alignn/graphs.py
@@ -1,5 +1,7 @@
"""Module to generate networkx graphs."""
+from __future__ import annotations
+
from jarvis.core.atoms import get_supercell_dims
from jarvis.core.specie import Specie
from jarvis.core.utils import random_colors
@@ -11,9 +13,22 @@
import math
from collections import defaultdict
from typing import List, Tuple, Sequence, Optional
-from dgl.data import DGLDataset
import torch
-import dgl
+
+try:
+ from dgl.data import DGLDataset
+ import dgl
+except Exception: # DGL optional; DGL-dependent helpers fail lazily
+ dgl = None
+
+ class DGLDataset: # noqa: D101 - minimal stub for subclasses
+ """Minimal stub so subclasses can be defined without DGL."""
+
+ def __init__(self, *args, **kwargs):
+ """Accept and ignore any arguments."""
+ pass
+
+
from tqdm import tqdm
from jarvis.core.atoms import Atoms
@@ -264,6 +279,152 @@ def build_undirected_edgedata(
return u, v, r, all_images
+def fast_graph(
+ atoms=None,
+ cutoff=5.0,
+ max_neighbors=None,
+ id=None,
+):
+ """Fast radius neighbor list using matscipy (C) on CPU.
+
+ Returns (u, v, r, images) in the same format as radius_graph so it
+ can be consumed by Graph.atom_dgl_multigraph downstream.
+ """
+ from matscipy.neighbours import neighbour_list as _matscipy_neighbours
+
+ ase_atoms = atoms.ase_converter()
+ i_np, j_np, D_np, S_np = _matscipy_neighbours(
+ "ijDS", ase_atoms, float(cutoff)
+ )
+
+ if max_neighbors is not None and max_neighbors > 0 and i_np.size > 0:
+ d_np = np.linalg.norm(D_np, axis=1)
+ # rank edges by distance globally, then keep first K per source
+ order = np.lexsort((d_np, i_np))
+ i_np = i_np[order]
+ j_np = j_np[order]
+ D_np = D_np[order]
+ S_np = S_np[order]
+ N = int(ase_atoms.positions.shape[0])
+ src_start = np.searchsorted(i_np, np.arange(N))
+ within = np.arange(i_np.shape[0]) - src_start[i_np]
+ sel = within < int(max_neighbors)
+ i_np = i_np[sel]
+ j_np = j_np[sel]
+ D_np = D_np[sel]
+ S_np = S_np[sel]
+
+ dtype = torch.get_default_dtype()
+ u = torch.from_numpy(np.ascontiguousarray(i_np)).long()
+ v = torch.from_numpy(np.ascontiguousarray(j_np)).long()
+ r = torch.from_numpy(np.ascontiguousarray(D_np)).type(dtype)
+ images = torch.from_numpy(np.ascontiguousarray(S_np)).type(dtype)
+ return u, v, r, images
+
+
+# =====================================================================
+# Fully-torch autograd-compatible neighbor list & graph builder
+#
+# Rationale:
+# Existing radius_graph / fast_graph return `r` as a fresh tensor
+# rebuilt from numpy, severing its link to atom positions and the
+# lattice. For autograd-based forces (dE/dx) and stress (dE/dL), `r`
+# must be a torch function of both. Here topology is computed once
+# (discrete, no grad), and displacement vectors are then recomputed
+# as `r = pos[dst] - pos[src] + shift @ lattice`, which carries
+# gradients. The line graph uses `shared=True` so bond-angle cosines
+# inherit the autograd graph automatically.
+# =====================================================================
+
+
+# Neighbor-list primitives now live in alignn.torch_graph_builder so
+# the pure-torch path doesn't transitively import DGL. Re-export them
+# here for backward compatibility.
+from alignn.torch_graph_builder import ( # noqa: E402,F401
+ _torch_periodic_shifts,
+ _topk_per_source,
+ torch_neighbor_list,
+)
+
+
+def torch_graph(
+ atoms=None,
+ cutoff=5.0,
+ max_neighbors=None,
+ atom_features="cgcnn",
+ compute_line_graph=True,
+ use_lattice_prop=False,
+ positions: Optional[torch.Tensor] = None,
+ lattice: Optional[torch.Tensor] = None,
+ device=None,
+ use_matscipy_topology: bool = True,
+ id=None,
+):
+ """End-to-end autograd-capable DGL (graph, line_graph) builder.
+
+ Returns ``(g, lg)`` if ``compute_line_graph`` else ``g``. When
+ ``positions`` / ``lattice`` are provided as torch leaves with
+ ``requires_grad=True``, ``g.edata['r']`` and ``lg.edata['h']``
+ carry gradients back to them.
+ """
+ dtype = torch.get_default_dtype()
+ if positions is None:
+ positions = torch.as_tensor(np.asarray(atoms.cart_coords), dtype=dtype)
+ if lattice is None:
+ lattice = torch.as_tensor(np.asarray(atoms.lattice_mat), dtype=dtype)
+ if device is not None:
+ positions = positions.to(device)
+ lattice = lattice.to(device)
+ device = positions.device
+ n_atoms = len(atoms.elements)
+
+ src, dst, shift, r = torch_neighbor_list(
+ positions=positions,
+ lattice=lattice,
+ cutoff=float(cutoff),
+ max_neighbors=max_neighbors,
+ atoms=atoms,
+ use_matscipy_topology=use_matscipy_topology,
+ )
+
+ sps_features = np.array(
+ [
+ list(get_node_attributes(s, atom_features=atom_features))
+ for s in atoms.elements
+ ]
+ )
+ node_features = torch.as_tensor(sps_features, dtype=dtype, device=device)
+
+ g = dgl.graph((src, dst), num_nodes=n_atoms)
+ if g.device != device:
+ g = g.to(device)
+ g.ndata["atom_features"] = node_features
+ g.edata["r"] = r # differentiable w.r.t. positions & lattice
+ g.edata["images"] = shift # integer cell offsets
+ vol = torch.abs(torch.det(lattice))
+ g.ndata["V"] = vol.expand(n_atoms)
+ frac = torch.as_tensor(
+ np.asarray(atoms.frac_coords), dtype=dtype, device=device
+ )
+ g.ndata["frac_coords"] = frac
+ if use_lattice_prop:
+ lp = np.array(
+ [atoms.lattice.lat_lengths(), atoms.lattice.lat_angles()]
+ ).flatten()
+ g.ndata["extra_features"] = (
+ torch.as_tensor(lp, dtype=dtype, device=device)
+ .unsqueeze(0)
+ .expand(n_atoms, -1)
+ .contiguous()
+ )
+
+ if compute_line_graph:
+ lg = g.line_graph(shared=True)
+ lg.apply_edges(compute_bond_cosines)
+ return g, lg
+ return g
+
+
def radius_graph(
atoms=None,
cutoff=5,
@@ -472,7 +633,7 @@ def __init__(
def atom_dgl_multigraph(
atoms=None,
neighbor_strategy="k-nearest",
- cutoff=8.0,
+ cutoff=5.0,
max_neighbors=12,
atom_features="cgcnn",
max_attempts=3,
@@ -483,6 +644,7 @@ def atom_dgl_multigraph(
use_lattice_prop: bool = False,
cutoff_extra=3.5,
dtype="float32",
+ three_body_cutoff: Optional[float] = 3.5,
):
"""Obtain a DGLGraph for Atoms object."""
# print('id',id)
@@ -503,6 +665,56 @@ def atom_dgl_multigraph(
u, v, r, images = radius_graph(
atoms, cutoff=cutoff, cutoff_extra=cutoff_extra
)
+ elif neighbor_strategy == "fast_graph":
+ u, v, r, images = fast_graph(
+ atoms,
+ cutoff=cutoff,
+ max_neighbors=max_neighbors,
+ id=id,
+ )
+ elif neighbor_strategy == "torch_graph":
+ # Returns a fully-built (g, lg) with autograd-capable r.
+ # Short-circuit to avoid the generic post-processing below
+ # which would detach r via torch.tensor(np.array(r)).
+ return torch_graph(
+ atoms=atoms,
+ cutoff=cutoff,
+ max_neighbors=max_neighbors,
+ atom_features=atom_features,
+ compute_line_graph=compute_line_graph,
+ use_lattice_prop=use_lattice_prop,
+ id=id,
+ )
+ elif neighbor_strategy == "pure_torch":
+ # Build with the pure-torch (non-DGL) builder, then
+ # materialize DGL graphs at the boundary. Downstream
+ # consumers (dataset/model) see DGL graphs as usual, and
+ # `edata["r"]` stays autograd-connected to positions and
+ # lattice because TorchGraph.to_dgl() assigns by reference.
+ from alignn.torch_graph_builder import (
+ build_pure_torch_graph as _build_pure_torch_graph,
+ )
+
+ _r3 = (
+ three_body_cutoff if three_body_cutoff is not None else cutoff
+ )
+ _out = _build_pure_torch_graph(
+ atoms=atoms,
+ two_body_cutoff=cutoff,
+ three_body_cutoff=_r3,
+ max_neighbors=max_neighbors,
+ atom_features=atom_features,
+ use_lattice_prop=use_lattice_prop,
+ compute_line_graph=compute_line_graph,
+ )
+ # Return TorchGraph directly when DGL is unavailable; downstream
+ # pure-torch datasets / models accept TorchGraph natively.
+ if dgl is None:
+ return _out
+ if compute_line_graph:
+ _tg, _tlg = _out
+ return _tg.to_dgl(), _tlg.to_dgl()
+ return _out.to_dgl()
elif neighbor_strategy == "radius_graph_jarvis":
g, lg = radius_graph_jarvis(
atoms,
@@ -844,6 +1056,18 @@ def prepare_line_graph_batch(
# return tuple(x.to(device) for x in batch)
+# Re-export the pure-torch (non-DGL) builder for convenience.
+# Defined in alignn.torch_graph_builder to keep that path importable
+# without pulling the full DGL-oriented graphs module.
+def build_pure_torch_graph(*args, **kwargs):
+ """See alignn.torch_graph_builder.build_pure_torch_graph."""
+ from alignn.torch_graph_builder import (
+ build_pure_torch_graph as _impl,
+ )
+
+ return _impl(*args, **kwargs)
+
+
def compute_bond_cosines(edges):
"""Compute bond angle cosines from bond displacement vectors."""
# line graph edge: (a, b), (b, c)
@@ -1062,7 +1286,7 @@ def collate(samples: List[Tuple[dgl.DGLGraph, torch.Tensor]]):
@staticmethod
def collate_line_graph(
- samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]]
+ samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]],
):
"""Dataloader helper to batch graphs cross `samples`."""
graphs, line_graphs, lattices, labels = map(list, zip(*samples))
diff --git a/alignn/lmdb_dataset.py b/alignn/lmdb_dataset.py
index 08e507c..c4f86f8 100644
--- a/alignn/lmdb_dataset.py
+++ b/alignn/lmdb_dataset.py
@@ -1,5 +1,7 @@
"""Module to prepare LMDB ALIGNN dataset."""
+from __future__ import annotations
+
import os
import numpy as np
import lmdb
@@ -11,7 +13,17 @@
import torch
from tqdm import tqdm
from typing import List, Tuple
-import dgl
+
+try:
+ import dgl
+except ImportError:
+ dgl = None
+
+
+# How often to commit the LMDB write transaction during a fresh build.
+# Holding a single transaction across all entries keeps every dirty page
+# resident in RAM until commit; flushing periodically caps peak memory.
+LMDB_COMMIT_EVERY = 5000
def prepare_line_graph_batch(
@@ -86,7 +98,7 @@ def collate(samples: List[Tuple[dgl.DGLGraph, torch.Tensor]]):
@staticmethod
def collate_line_graph(
- samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]]
+ samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]],
):
"""Dataloader helper to batch graphs cross `samples`."""
graphs, line_graphs, lattices, labels = map(list, zip(*samples))
@@ -108,6 +120,16 @@ def collate_line_graph(
)
+def _mem_gb():
+ """Resident memory of the current process in GB (for diagnostics)."""
+ try:
+ import psutil
+
+ return psutil.Process(os.getpid()).memory_info().rss / 1e9
+ except Exception:
+ return float("nan")
+
+
def get_torch_dataset(
dataset=[],
id_tag="jid",
@@ -124,12 +146,13 @@ def get_torch_dataset(
cutoff=8.0,
cutoff_extra=3.0,
max_neighbors=12,
+ three_body_cutoff=None,
classification=False,
sampler=None,
output_dir=".",
tmp_name="dataset",
map_size=1e12,
- read_existing=True,
+ read_existing=False,
dtype="float32",
):
"""Get Torch Dataset with LMDB."""
@@ -142,21 +165,59 @@ def get_torch_dataset(
line = "Min=" + str(np.min(vals)) + "\n"
f.write(line)
f.close()
+
+ print(f"[MEM] {tmp_name}: entry: {_mem_gb():.2f} GB")
+
ids = []
if os.path.exists(tmp_name) and read_existing:
+ # Validate the cache backend matches (DGL). A previous run with
+ # model.name="alignn_atomwise_pure" would have left TorchGraph
+ # pickles here, which dgl.batch can't consume.
+ _env = lmdb.open(tmp_name, readonly=True, lock=False)
+ with _env.begin() as _txn:
+ _probe = _txn.get(b"0")
+ _env.close()
+ if _probe is not None:
+ _sample = pk.loads(_probe)
+ _graph = _sample[0]
+ if dgl is None or not isinstance(_graph, dgl.DGLGraph):
+ raise RuntimeError(
+ f"LMDB cache at '{tmp_name}' contains "
+ f"{type(_graph).__name__} records, not DGLGraph. "
+ "Delete the stale cache (e.g. `rm -rf "
+ f"{tmp_name}`) and rerun — it was built for a "
+ "different model backend."
+ )
for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)):
ids.append(d[id_tag])
dat = TorchLMDBDataset(
lmdb_path=tmp_name, line_graph=line_graph, ids=ids
)
print("Reading dataset", tmp_name)
+ print(
+ f"[MEM] {tmp_name}: after TorchLMDBDataset (cached): "
+ f"{_mem_gb():.2f} GB"
+ )
return dat
+
ids = []
+ # Fresh build: wipe any pre-existing cache dir to avoid mixing old
+ # records (possibly from a different model backend) with new writes.
+ if os.path.exists(tmp_name):
+ import shutil
+
+ shutil.rmtree(tmp_name)
env = lmdb.open(tmp_name, map_size=int(map_size))
- with env.begin(write=True) as txn:
+
+ print(f"[MEM] {tmp_name}: before write loop: {_mem_gb():.2f} GB")
+
+ # Use periodic commits instead of one giant transaction. A single
+ # write txn around 1.5M put() calls keeps every dirty page in RAM
+ # until commit and was the cause of OOM on 62 GB hosts.
+ txn = env.begin(write=True)
+ try:
for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)):
ids.append(d[id_tag])
- # g, lg = Graph.atom_dgl_multigraph(
atoms = Atoms.from_dict(d["atoms"])
g = Graph.atom_dgl_multigraph(
atoms,
@@ -167,6 +228,7 @@ def get_torch_dataset(
use_canonize=use_canonize,
cutoff_extra=cutoff_extra,
neighbor_strategy=neighbor_strategy,
+ three_body_cutoff=three_body_cutoff,
dtype=dtype,
)
if line_graph:
@@ -176,10 +238,8 @@ def get_torch_dataset(
)
label = torch.tensor(d[target]).type(torch.get_default_dtype())
natoms = len(d["atoms"]["elements"])
- # print('label',label,label.view(-1).long())
if classification:
label = label.long()
- # label = label.view(-1).long()
if "extra_features" in d:
g.ndata["extra_features"] = torch.tensor(
[d["extra_features"] for n in range(natoms)]
@@ -189,8 +249,6 @@ def get_torch_dataset(
np.array(d[target_atomwise])
).type(torch.get_default_dtype())
if target_grad is not None and target_grad != "":
- # print('grad', np.array(d[target_grad]))
- # print('grad shape',np.array(d[target_grad]).shape)
arr = np.array(d[target_grad])
try:
g.ndata[target_grad] = torch.tensor(arr).type(
@@ -201,7 +259,6 @@ def get_torch_dataset(
g.ndata[target_grad] = torch.tensor(arr).type(
torch.get_default_dtype()
)
- # print('arr',arr.shape)
if target_stress is not None and target_stress != "":
stress = np.array(d[target_stress])
g.ndata[target_stress] = torch.tensor(
@@ -216,17 +273,38 @@ def get_torch_dataset(
([additional_output for ii in range(natoms)])
).type(torch.get_default_dtype())
- # labels.append(label)
if line_graph:
serialized_data = pk.dumps((g, lg, lattice, label))
else:
serialized_data = pk.dumps((g, lattice, label))
txn.put(f"{idx}".encode(), serialized_data)
+ # Flush dirty pages to disk every LMDB_COMMIT_EVERY entries.
+ if (idx + 1) % LMDB_COMMIT_EVERY == 0:
+ txn.commit()
+ txn = env.begin(write=True)
+ if (idx + 1) % (LMDB_COMMIT_EVERY * 20) == 0:
+ print(
+ f"[MEM] {tmp_name}: after {idx + 1} entries: "
+ f"{_mem_gb():.2f} GB"
+ )
+ finally:
+ # Commit whatever is left in the final partial batch.
+ txn.commit()
+
env.close()
+
+ print(f"[MEM] {tmp_name}: after write loop: {_mem_gb():.2f} GB")
+
lmdb_dataset = TorchLMDBDataset(
lmdb_path=tmp_name, line_graph=line_graph, ids=ids
)
+
+ print(
+ f"[MEM] {tmp_name}: after TorchLMDBDataset wrap: "
+ f"{_mem_gb():.2f} GB"
+ )
+
return lmdb_dataset
diff --git a/alignn/md/__init__.py b/alignn/md/__init__.py
new file mode 100644
index 0000000..dab1aeb
--- /dev/null
+++ b/alignn/md/__init__.py
@@ -0,0 +1,13 @@
+"""Pure-PyTorch MD integrators for ALIGNN-FF."""
+from alignn.md.integrators import (
+ VelocityVerlet, Langevin, NVTBerendsen, NVTBussi, NVTNoseHooverChain,
+ run, maxwell_boltzmann,
+)
+from alignn.md.forces import AlignnForces
+from alignn.md.relax import FIRE
+__all__ = [
+ "VelocityVerlet", "Langevin",
+ "NVTBerendsen", "NVTBussi", "NVTNoseHooverChain",
+ "FIRE",
+ "run", "maxwell_boltzmann", "AlignnForces",
+]
diff --git a/alignn/md/forces.py b/alignn/md/forces.py
new file mode 100644
index 0000000..bd8ac60
--- /dev/null
+++ b/alignn/md/forces.py
@@ -0,0 +1,79 @@
+"""Wrap ALIGNN-FF (pure PyTorch variant) as an on-device forces function.
+
+A `forces_fn` has signature:
+ energy, forces = forces_fn(positions) # positions: (N, 3) on GPU
+ # energy: scalar tensor
+ # forces: (N, 3) tensor
+
+Current implementation rebuilds the graph every step via the existing
+jarvis/DGL pipeline. That's CPU-heavy and will be the bottleneck for large
+systems — replace with an on-device neighbor list when you scale up.
+"""
+from __future__ import annotations
+import numpy as np
+import torch
+from ase import Atoms as AseAtoms
+from jarvis.core.atoms import ase_to_atoms as ase_to_jarvis
+
+from alignn.graphs import Graph
+from alignn.torch_graph_builder import torchgraph_from_dgl
+
+
+class AlignnForces:
+ def __init__(
+ self,
+ model,
+ atomic_numbers: np.ndarray,
+ cell: np.ndarray, # (3, 3), constant for NVT/NVE
+ cutoff: float = 8.0,
+ max_neighbors: int = 12,
+ device: str | torch.device = "cuda",
+ dtype: torch.dtype = torch.float32,
+ ):
+ self.model = model.eval().to(device).to(dtype)
+ self.Z = np.asarray(atomic_numbers)
+ self.cell_np = np.asarray(cell, dtype=float)
+ self.cell = torch.tensor(self.cell_np, dtype=dtype, device=device)
+ self.cutoff = cutoff
+ self.max_neighbors = max_neighbors
+ self.device = torch.device(device)
+ self.dtype = dtype
+
+ def _build_graph(self, positions_np):
+ ase_atoms = AseAtoms(
+ numbers=self.Z, positions=positions_np, cell=self.cell_np, pbc=True
+ )
+ j = ase_to_jarvis(ase_atoms)
+ g, lg = Graph.atom_dgl_multigraph(
+ j, neighbor_strategy="k-nearest",
+ cutoff=self.cutoff, max_neighbors=self.max_neighbors,
+ atom_features="cgcnn", use_canonize=True,
+ )
+ g = g.to(self.device)
+ lg = lg.to(self.device)
+ tg = torchgraph_from_dgl(g)
+ tlg = torchgraph_from_dgl(lg)
+ for d in (tg.ndata, tg.edata, tlg.ndata, tlg.edata):
+ for k, v in list(d.items()):
+ if v.is_floating_point():
+ d[k] = v.to(self.dtype)
+ return tg, tlg
+
+ @torch.enable_grad()
+ def __call__(self, positions: torch.Tensor):
+ positions = positions.detach()
+ tg, tlg = self._build_graph(positions.cpu().numpy())
+ # model sets r.requires_grad_(True) internally; energy grad wrt r
+ # gives per-edge forces, aggregated to per-atom in the model's head.
+ out = self.model((tg, tlg, self.cell))
+ energy = (
+ out["out"].sum()
+ if "out" in out
+ else next(iter(out.values())).sum()
+ )
+ if "grad" in out:
+ forces = out["grad"].detach()
+ else:
+ # fallback: autograd wrt edge vectors is handled inside the model
+ forces = torch.zeros_like(positions)
+ return energy.detach(), forces.to(positions.dtype)
diff --git a/alignn/md/integrators.py b/alignn/md/integrators.py
new file mode 100644
index 0000000..2611363
--- /dev/null
+++ b/alignn/md/integrators.py
@@ -0,0 +1,382 @@
+"""Pure-PyTorch MD integrators: Velocity-Verlet and Langevin (BAOAB).
+
+State kept entirely on-device. Unit system follows ASE:
+ energy eV, length Å, mass amu, time fs, temperature K.
+With these units, the acceleration factor `F/m` needs the conversion
+ 1 eV/(Å·amu) = 9.6485e-3 Å/fs² (i.e., ASE's `ase.units.fs` scaling)
+which we bake in as `EV_AMU_A_PER_FS2`.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Callable, Optional
+import math
+import time
+import torch
+
+KB_EV = 8.617333262e-5 # Boltzmann, eV/K
+EV_AMU_A_PER_FS2 = 9.6485332e-3 # eV/(Å·amu) -> Å/fs²
+
+ForcesFn = Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]]
+
+
+def wrap_pbc(positions: torch.Tensor, cell: torch.Tensor) -> torch.Tensor:
+ """Wrap positions into the primary unit cell (3x3 cell, row vectors)."""
+ inv = torch.linalg.inv(cell)
+ frac = positions @ inv
+ frac = frac - torch.floor(frac)
+ return frac @ cell
+
+
+def kinetic_energy(
+ masses: torch.Tensor, velocities: torch.Tensor
+) -> torch.Tensor:
+ """Return kinetic energy in eV (m in amu, v in Å/fs)."""
+ # 0.5 * m * v², converted from amu·Å²/fs² to eV
+ return (
+ 0.5 * (masses * (velocities ** 2).sum(dim=-1)).sum()
+ / EV_AMU_A_PER_FS2
+ )
+
+
+def instantaneous_temperature(masses, velocities):
+ ke = kinetic_energy(masses, velocities)
+ n_dof = 3 * masses.numel() - 3 # remove COM
+ return 2.0 * ke / (n_dof * KB_EV)
+
+
+def maxwell_boltzmann(
+ masses: torch.Tensor, T: float, generator=None
+) -> torch.Tensor:
+ # sigma_v = sqrt(kT/m), in Å/fs after unit conversion
+ sigma = torch.sqrt(KB_EV * T * EV_AMU_A_PER_FS2 / masses).unsqueeze(-1)
+ v = torch.randn((masses.numel(), 3), device=masses.device,
+ dtype=masses.dtype, generator=generator) * sigma
+ # remove COM drift
+ v -= (masses.unsqueeze(-1) * v).sum(dim=0, keepdim=True) / masses.sum()
+ # Rescale to exact target T (removes finite-sample noise)
+ T_now = instantaneous_temperature(masses, v)
+ v = v * torch.sqrt(torch.tensor(T, device=v.device, dtype=v.dtype) / T_now)
+ return v
+
+
+# ---------------------------------------------------------------------------
+# NVE: Velocity-Verlet
+# ---------------------------------------------------------------------------
+@dataclass
+class VelocityVerlet:
+ forces_fn: ForcesFn
+ masses: torch.Tensor # (N,) amu
+ dt: float = 1.0 # fs
+
+ def __post_init__(self):
+ self._forces = None
+
+ def step(self, positions, velocities):
+ if self._forces is None:
+ _, self._forces = self.forces_fn(positions)
+ m = self.masses.unsqueeze(-1)
+ a = EV_AMU_A_PER_FS2 * self._forces / m
+ # Half-kick
+ velocities = velocities + 0.5 * self.dt * a
+ # Drift
+ positions = positions + self.dt * velocities
+ # New forces
+ _, new_forces = self.forces_fn(positions)
+ a_new = EV_AMU_A_PER_FS2 * new_forces / m
+ # Half-kick
+ velocities = velocities + 0.5 * self.dt * a_new
+ self._forces = new_forces
+ return positions, velocities
+
+
+# ---------------------------------------------------------------------------
+# NVT: Langevin via BAOAB splitting (Leimkuhler & Matthews 2013)
+# ---------------------------------------------------------------------------
+@dataclass
+class Langevin:
+ forces_fn: ForcesFn
+ masses: torch.Tensor
+ dt: float = 1.0 # fs
+ T: float = 300.0 # K
+ friction: float = 0.01 # 1/fs (typical: 0.01 = 100 fs damping)
+ generator: Optional[torch.Generator] = None
+
+ def __post_init__(self):
+ self._forces = None
+ self.c1 = math.exp(-self.friction * self.dt)
+ self.c3 = math.sqrt(1.0 - self.c1 ** 2)
+
+ def step(self, positions, velocities):
+ if self._forces is None:
+ _, self._forces = self.forces_fn(positions)
+ m = self.masses.unsqueeze(-1)
+ # B: half-kick
+ a = EV_AMU_A_PER_FS2 * self._forces / m
+ velocities = velocities + 0.5 * self.dt * a
+ # A: half-drift
+ positions = positions + 0.5 * self.dt * velocities
+ # O: Ornstein-Uhlenbeck in velocity
+ sigma = torch.sqrt(
+ KB_EV * self.T * EV_AMU_A_PER_FS2 / self.masses
+ ).unsqueeze(-1)
+ if self.generator is None:
+ noise = torch.randn_like(velocities)
+ else:
+ noise = torch.randn(
+ velocities.shape,
+ device=velocities.device,
+ dtype=velocities.dtype,
+ generator=self.generator,
+ )
+ velocities = self.c1 * velocities + self.c3 * sigma * noise
+ # A: half-drift
+ positions = positions + 0.5 * self.dt * velocities
+ # B: half-kick with new forces
+ _, new_forces = self.forces_fn(positions)
+ a_new = EV_AMU_A_PER_FS2 * new_forces / m
+ velocities = velocities + 0.5 * self.dt * a_new
+ self._forces = new_forces
+ return positions, velocities
+
+
+# ---------------------------------------------------------------------------
+# NVT: Berendsen weak-coupling thermostat
+# (Not a true ensemble — fine for equilibration / melt-quench, not for
+# equilibrium sampling or fluctuation-based observables.)
+# ---------------------------------------------------------------------------
+@dataclass
+class NVTBerendsen:
+ forces_fn: ForcesFn
+ masses: torch.Tensor
+ dt: float = 1.0 # fs
+ T: float = 300.0 # K (target; mutable via set_temperature)
+ taut: float = 20.0 # fs (relaxation time)
+
+ def __post_init__(self):
+ self._forces = None
+
+ def set_temperature(self, T: float):
+ self.T = float(T)
+
+ def step(self, positions, velocities):
+ if self._forces is None:
+ _, self._forces = self.forces_fn(positions)
+ m = self.masses.unsqueeze(-1)
+ # VV half-kick
+ a = EV_AMU_A_PER_FS2 * self._forces / m
+ velocities = velocities + 0.5 * self.dt * a
+ # drift
+ positions = positions + self.dt * velocities
+ # new forces
+ _, new_forces = self.forces_fn(positions)
+ a_new = EV_AMU_A_PER_FS2 * new_forces / m
+ velocities = velocities + 0.5 * self.dt * a_new
+ # Berendsen velocity rescale
+ T_now = instantaneous_temperature(self.masses, velocities)
+ # Guard against T_now == 0 early on
+ scale = torch.sqrt(
+ 1.0
+ + (self.dt / self.taut)
+ * (self.T / T_now.clamp_min(1e-6) - 1.0)
+ )
+ velocities = velocities * scale
+ self._forces = new_forces
+ return positions, velocities
+
+
+# ---------------------------------------------------------------------------
+# NVT: Bussi-Donadio-Parrinello stochastic velocity rescaling (CSVR)
+# Bussi, Donadio & Parrinello, J. Chem. Phys. 126, 014101 (2007).
+# Samples the true canonical distribution with Berendsen-like relaxation.
+# ---------------------------------------------------------------------------
+@dataclass
+class NVTBussi:
+ forces_fn: ForcesFn
+ masses: torch.Tensor
+ dt: float = 1.0 # fs
+ T: float = 300.0 # K
+ taut: float = 20.0 # fs (same meaning as Berendsen τ_T)
+ remove_com: bool = True # subtract 3 DOF if true
+
+ def __post_init__(self):
+ self._forces = None
+
+ def set_temperature(self, T: float):
+ self.T = float(T)
+
+ def _rescale_factor(self, K_now: torch.Tensor) -> torch.Tensor:
+ """Bussi 2007 Eq. A7 — returns α s.t. v ← α v samples canonical KE."""
+ Nf = 3 * self.masses.numel() - (3 if self.remove_com else 0)
+ K_bar = 0.5 * Nf * KB_EV * self.T # target KE (eV)
+ # convert K_now from "integrator-internal" kinetic units:
+ # K_now already in eV since our kinetic_energy() returns eV.
+ c = math.exp(-self.dt / self.taut)
+ device, dtype = K_now.device, K_now.dtype
+ # one N(0,1) and one chi²(Nf-1) draw
+ R1 = torch.randn((), device=device, dtype=dtype)
+ # chi²(k) sampled as sum of k squared normals — cheap for small/med N;
+ # for huge N switch to torch.distributions.Gamma.
+ k = max(Nf - 1, 1)
+ S = (torch.randn((k,), device=device, dtype=dtype) ** 2).sum()
+ factor = K_bar / (Nf * K_now.clamp_min(1e-12))
+ alpha_sq = (
+ c
+ + (1.0 - c) * (S + R1 * R1) * factor
+ + 2.0 * R1 * torch.sqrt(c * (1.0 - c) * factor)
+ )
+ # sign: α has the sign of (R1 + sqrt(c*Nf*K/(factor*(1-c)))) per BDP;
+ # in practice α² > 0 and we take positive root.
+ return torch.sqrt(alpha_sq.clamp_min(0.0))
+
+ def step(self, positions, velocities):
+ if self._forces is None:
+ _, self._forces = self.forces_fn(positions)
+ m = self.masses.unsqueeze(-1)
+ a = EV_AMU_A_PER_FS2 * self._forces / m
+ velocities = velocities + 0.5 * self.dt * a
+ positions = positions + self.dt * velocities
+ _, new_forces = self.forces_fn(positions)
+ a_new = EV_AMU_A_PER_FS2 * new_forces / m
+ velocities = velocities + 0.5 * self.dt * a_new
+ # Stochastic canonical rescale
+ K_now = kinetic_energy(self.masses, velocities)
+ alpha = self._rescale_factor(K_now)
+ velocities = velocities * alpha
+ self._forces = new_forces
+ return positions, velocities
+
+
+# ---------------------------------------------------------------------------
+# NVT: Nosé-Hoover chain (Martyna-Tobias-Klein integrator)
+# Chains of length M ≥ 3 fix the non-ergodicity of a single Nosé-Hoover.
+# Reference: Martyna, Tobias & Klein, J. Chem. Phys. 101, 4177 (1994);
+# Tuckerman, "Statistical Mechanics: Theory and Molecular
+# Simulation", Ch. 4.
+# ---------------------------------------------------------------------------
+@dataclass
+class NVTNoseHooverChain:
+ forces_fn: ForcesFn
+ masses: torch.Tensor
+ dt: float = 1.0 # fs
+ T: float = 300.0 # K
+ taut: float = 20.0 # fs — chain relaxation time; Q1 = Nf·kT·τ²
+ chain_length: int = 3 # M; use ≥3 for ergodicity
+ remove_com: bool = True
+
+ def __post_init__(self):
+ device, dtype = self.masses.device, self.masses.dtype
+ M = self.chain_length
+ self.v_xi = torch.zeros(M, device=device, dtype=dtype)
+ self._forces = None
+ self._set_masses()
+
+ def _set_masses(self):
+ device, dtype = self.masses.device, self.masses.dtype
+ Nf = 3 * self.masses.numel() - (3 if self.remove_com else 0)
+ kT = KB_EV * self.T
+ self._Nf = Nf
+ self._kT = kT
+ tau2 = self.taut ** 2
+ Q = torch.full(
+ (self.chain_length,), kT * tau2, device=device, dtype=dtype
+ )
+ Q[0] = Nf * kT * tau2
+ self._Q = Q
+
+ def set_temperature(self, T: float):
+ self.T = float(T)
+ self._set_masses() # rescale Q so relaxation stays ~τ
+
+ def _apply_chain_half(self, velocities):
+ """Apply half a chain step (dt_half = dt/2 of MD timestep)."""
+ dt_half = self.dt / 2
+ dthalf = dt_half / 2
+ dtqtr = dt_half / 4
+ Q, v_xi = self._Q, self.v_xi
+ M = self.chain_length
+ K = kinetic_energy(self.masses, velocities) # eV
+ # --- reverse pass (top of chain down) ---
+ if M >= 2:
+ G_M = (Q[M - 2] * v_xi[M - 2] ** 2 - self._kT) / Q[M - 1]
+ else:
+ G_M = (2 * K - self._Nf * self._kT) / Q[0]
+ v_xi[M - 1] = v_xi[M - 1] + G_M * dthalf
+ for i in range(M - 2, -1, -1):
+ factor = torch.exp(-dtqtr * v_xi[i + 1])
+ v_xi[i] = v_xi[i] * factor
+ if i == 0:
+ G_i = (2 * K - self._Nf * self._kT) / Q[0]
+ else:
+ G_i = (Q[i - 1] * v_xi[i - 1] ** 2 - self._kT) / Q[i]
+ v_xi[i] = v_xi[i] + G_i * dthalf
+ v_xi[i] = v_xi[i] * factor
+ # --- rescale particle velocities ---
+ scale = torch.exp(-dt_half * v_xi[0])
+ velocities = velocities * scale
+ K = K * scale * scale
+ # --- forward pass (bottom up) ---
+ for i in range(0, M - 1):
+ factor = torch.exp(-dtqtr * v_xi[i + 1])
+ v_xi[i] = v_xi[i] * factor
+ if i == 0:
+ G_i = (2 * K - self._Nf * self._kT) / Q[0]
+ else:
+ G_i = (Q[i - 1] * v_xi[i - 1] ** 2 - self._kT) / Q[i]
+ v_xi[i] = v_xi[i] + G_i * dthalf
+ v_xi[i] = v_xi[i] * factor
+ if M >= 2:
+ G_M = (Q[M - 2] * v_xi[M - 2] ** 2 - self._kT) / Q[M - 1]
+ v_xi[M - 1] = v_xi[M - 1] + G_M * dthalf
+ return velocities
+
+ def step(self, positions, velocities):
+ if self._forces is None:
+ _, self._forces = self.forces_fn(positions)
+ m = self.masses.unsqueeze(-1)
+ # Chain half-step
+ velocities = self._apply_chain_half(velocities)
+ # VV: kick-drift-kick
+ a = EV_AMU_A_PER_FS2 * self._forces / m
+ velocities = velocities + 0.5 * self.dt * a
+ positions = positions + self.dt * velocities
+ _, new_forces = self.forces_fn(positions)
+ a_new = EV_AMU_A_PER_FS2 * new_forces / m
+ velocities = velocities + 0.5 * self.dt * a_new
+ # Chain half-step
+ velocities = self._apply_chain_half(velocities)
+ self._forces = new_forces
+ return positions, velocities
+
+
+# ---------------------------------------------------------------------------
+# Driver loop
+# ---------------------------------------------------------------------------
+def run(
+ integrator,
+ positions: torch.Tensor,
+ velocities: torch.Tensor,
+ nsteps: int,
+ log_every: int = 100,
+ callback: Optional[Callable[[int, dict], None]] = None,
+):
+ positions = positions.clone()
+ velocities = velocities.clone()
+ hist = []
+ t0 = time.time()
+ for i in range(nsteps):
+ positions, velocities = integrator.step(positions, velocities)
+ if (i % log_every) == 0 or i == nsteps - 1:
+ ke = kinetic_energy(integrator.masses, velocities).item()
+ T = instantaneous_temperature(integrator.masses, velocities).item()
+ row = {"step": i, "time_fs": i * integrator.dt, "T_K": T,
+ "KE_eV": ke, "wall_s": time.time() - t0}
+ hist.append(row)
+ if callback is not None:
+ callback(i, row)
+ else:
+ print(
+ f"step {i:6d} t={row['time_fs']:8.1f} fs "
+ f"T={T:7.1f} K KE={ke:10.4f} eV "
+ f"wall={row['wall_s']:6.1f}s"
+ )
+ return positions, velocities, hist
diff --git a/alignn/md/relax.py b/alignn/md/relax.py
new file mode 100644
index 0000000..bc3db8d
--- /dev/null
+++ b/alignn/md/relax.py
@@ -0,0 +1,104 @@
+"""FIRE structural relaxation (Bitzek et al., PRL 97, 170201, 2006).
+
+ASE-compatible defaults. Drop-in for `ase.optimize.FIRE` when you want
+to stay on-device and avoid the ASE roundtrip every step.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Callable, Optional
+import time
+import torch
+
+from alignn.md.integrators import ForcesFn, EV_AMU_A_PER_FS2
+
+
+@dataclass
+class FIRE:
+ forces_fn: ForcesFn
+ masses: torch.Tensor # (N,) amu
+ dt: float = 0.1 # fs (initial timestep)
+ dt_max: float = 1.0 # fs
+ max_step: float = 0.2 # Å (per-atom cap on Δx)
+ N_min: int = 5
+ f_inc: float = 1.1
+ f_dec: float = 0.5
+ alpha_start: float = 0.1
+ f_alpha: float = 0.99
+
+ def run(
+ self,
+ positions: torch.Tensor,
+ fmax: float = 0.05, # eV/Å convergence criterion
+ max_steps: int = 500,
+ log_every: int = 10,
+ callback: Optional[Callable[[int, dict], None]] = None,
+ ):
+ positions = positions.clone()
+ velocities = torch.zeros_like(positions)
+ m = self.masses.unsqueeze(-1)
+ dt = self.dt
+ alpha = self.alpha_start
+ n_pos = 0
+ history = []
+ t0 = time.time()
+
+ # initial force
+ energy, forces = self.forces_fn(positions)
+ for step in range(max_steps):
+ fnorm_max = forces.norm(dim=-1).max().item()
+ row = {
+ "step": step, "E_eV": float(energy),
+ "fmax_eV_A": fnorm_max, "dt_fs": dt, "alpha": alpha,
+ "wall_s": time.time() - t0,
+ }
+ if step % log_every == 0 or fnorm_max < fmax:
+ history.append(row)
+ if callback is not None:
+ callback(step, row)
+ else:
+ print(f"FIRE step={step:4d} E={row['E_eV']:12.4f} eV "
+ f"|F|max={fnorm_max:9.4f} eV/Å dt={dt:5.3f} fs "
+ f"α={alpha:5.3f}")
+ if fnorm_max < fmax:
+ print(
+ f"converged at step {step}: "
+ f"|F|max = {fnorm_max:.5f} < {fmax}"
+ )
+ break
+
+ # FIRE velocity mix: v ← (1-α) v + α |v| F̂
+ f_hat = forces / (forces.norm() + 1e-30)
+ v_norm = velocities.norm()
+ velocities = (1.0 - alpha) * velocities + alpha * v_norm * f_hat
+
+ # Power P = F·v (sum over all atoms)
+ P = (forces * velocities).sum()
+ if P.item() > 0.0:
+ n_pos += 1
+ if n_pos > self.N_min:
+ dt = min(dt * self.f_inc, self.dt_max)
+ alpha = alpha * self.f_alpha
+ else:
+ n_pos = 0
+ dt = dt * self.f_dec
+ alpha = self.alpha_start
+ velocities = torch.zeros_like(velocities)
+
+ # Velocity Verlet step (with max_step clipping)
+ a = EV_AMU_A_PER_FS2 * forces / m
+ velocities = velocities + 0.5 * dt * a
+ dx = dt * velocities
+ # Per-atom displacement cap
+ dx_norm = dx.norm(dim=-1, keepdim=True)
+ scale = torch.where(
+ dx_norm > self.max_step,
+ self.max_step / dx_norm.clamp_min(1e-30),
+ torch.ones_like(dx_norm),
+ )
+ dx = dx * scale
+ positions = positions + dx
+ energy, new_forces = self.forces_fn(positions)
+ a_new = EV_AMU_A_PER_FS2 * new_forces / m
+ velocities = velocities + 0.5 * dt * a_new
+ forces = new_forces
+ return positions, history
diff --git a/alignn/models/alignn.py b/alignn/models/alignn.py
index 69c14ad..6d44954 100644
--- a/alignn/models/alignn.py
+++ b/alignn/models/alignn.py
@@ -3,12 +3,20 @@
A prototype crystal line graph network dgl implementation.
"""
+from __future__ import annotations
+
from typing import Tuple, Union
-import dgl
-import dgl.function as fn
+
+try:
+ import dgl
+ import dgl.function as fn
+ from dgl.nn import AvgPooling
+except Exception:
+ dgl = None
+ fn = None
+ AvgPooling = None
import numpy as np
import torch
-from dgl.nn import AvgPooling
from typing import Literal
from torch import nn
from torch.nn import functional as F
diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index aff0243..25a3fc2 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -3,12 +3,21 @@
A prototype crystal line graph network dgl implementation.
"""
+from __future__ import annotations
+
from typing import Tuple, Union
from torch.autograd import grad
-import dgl
-import dgl.function as fn
+
+try:
+ import dgl
+ import dgl.function as fn
+ from dgl.nn import AvgPooling
+except Exception: # DGL optional; module only usable if DGL installed
+ dgl = None
+ fn = None
+ AvgPooling = None
+ print("WARNING: No DGL, might fail")
import numpy as np
-from dgl.nn import AvgPooling
import torch
# from dgl.nn.functional import edge_softmax
diff --git a/alignn/models/alignn_atomwise_pure.py b/alignn/models/alignn_atomwise_pure.py
new file mode 100644
index 0000000..006cdd2
--- /dev/null
+++ b/alignn/models/alignn_atomwise_pure.py
@@ -0,0 +1,909 @@
+"""Pure-PyTorch ALIGNN-atomwise (no DGL).
+
+This is a line-for-line reimplementation of ``ALIGNNAtomWise`` with all
+DGL ops (``dgl.function``, ``dgl.nn``, ``dgl.batch``, ``dgl.reverse``)
+replaced by index/scatter primitives on ``TorchGraph`` tensors. Same
+config fields, same forward output keys — drop-in for training and
+inference without depending on DGL at model-evaluation time.
+
+Accepts either:
+ * a ``(TorchGraph, TorchGraph, lattice)`` triple (preferred), or
+ * a ``(DGLGraph, DGLGraph, lattice)`` triple — converted at the
+ boundary via ``torchgraph_from_dgl`` so you can flip models without
+ touching the dataloader.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Literal, Tuple
+
+import numpy as np
+import torch
+from torch import nn
+from torch.autograd import grad
+from torch.nn import functional as F
+
+from alignn.models.utils import MLPLayer, RBFExpansion
+from alignn.torch_graph_builder import TorchGraph, torchgraph_from_dgl
+from alignn.utils import BaseSettings
+
+# =====================================================================
+# Primitives
+# =====================================================================
+
+
+def scatter_sum(
+ src: torch.Tensor, index: torch.Tensor, dim_size: int
+) -> torch.Tensor:
+ """Sum ``src`` rows into ``dim_size`` buckets indexed by ``index``.
+
+ Script-safe: builds output shape as an explicit ``List[int]`` and
+ broadcasts ``index`` via repeated ``unsqueeze`` (avoids unpacking).
+ """
+ out_shape: List[int] = [dim_size]
+ for i in range(1, src.dim()):
+ out_shape.append(src.size(i))
+ out = torch.zeros(out_shape, dtype=src.dtype, device=src.device)
+ idx = index
+ while idx.dim() < src.dim():
+ idx = idx.unsqueeze(-1)
+ idx = idx.expand_as(src)
+ out.scatter_add_(0, idx, src)
+ return out
+
+
+def scatter_mean(
+ src: torch.Tensor, index: torch.Tensor, dim_size: int
+) -> torch.Tensor:
+ total = scatter_sum(src, index, dim_size)
+ ones = torch.ones(src.shape[0], device=src.device, dtype=src.dtype)
+ count = torch.zeros(
+ dim_size, device=src.device, dtype=src.dtype
+ ).scatter_add_(0, index, ones)
+ denom = count.clamp_min(1.0)
+ while denom.dim() < total.dim():
+ denom = denom.unsqueeze(-1)
+ return total / denom
+
+
+# =====================================================================
+# Config
+# =====================================================================
+
+
+class ALIGNNAtomWisePureConfig(BaseSettings):
+ """Hyperparameter schema — mirrors ALIGNNAtomWiseConfig."""
+
+ name: Literal["alignn_atomwise_pure"]
+ alignn_layers: int = 2
+ gcn_layers: int = 2
+ atom_input_features: int = 1
+ edge_input_features: int = 80
+ triplet_input_features: int = 40
+ embedding_features: int = 64
+ hidden_features: int = 64
+ output_features: int = 1
+ grad_multiplier: int = -1
+ calculate_gradient: bool = True
+ atomwise_output_features: int = 0
+ graphwise_weight: float = 1.0
+ gradwise_weight: float = 1.0
+ stresswise_weight: float = 0.0
+ atomwise_weight: float = 0.0
+ link: Literal["identity", "log", "logit"] = "identity"
+ zero_inflated: bool = False
+ classification: bool = False
+ force_mult_natoms: bool = False
+ energy_mult_natoms: bool = True
+ # Accepted for config-schema compatibility with ALIGNNAtomWiseConfig;
+ # the pure-torch forward doesn't implement this branch.
+ include_pos_deriv: bool = False
+ use_cutoff_function: bool = False
+ inner_cutoff: float = 3.0
+ stress_multiplier: float = 1.0
+ add_reverse_forces: bool = True
+ lg_on_fly: bool = True
+ batch_stress: bool = True
+ multiply_cutoff: bool = False
+ use_penalty: bool = True
+ extra_features: int = 0
+ exponent: int = 5
+ penalty_factor: float = 0.5
+ penalty_threshold: float = 1.0
+ # penalty_factor: float = 0.1
+ # penalty_threshold: float = 1.0
+ additional_output_features: int = 0
+ additional_output_weight: float = 0.0
+ # Attention variants (Shao et al., Adv. Theory Simul. 2026):
+ # "alignn" - original edge-gated conv
+ # "n_alignn" - Node-Attention Layer (NAL), per-node learnable weights
+ # "t_alignn" - Self-Attention Layer (SAL), Transformer-style Q/K/V
+ conv_type: Literal["alignn", "n_alignn", "t_alignn"] = "alignn"
+ num_heads: int = 1
+
+ class Config:
+ env_prefix = "jv_model"
+
+
+# =====================================================================
+# Conv layers
+# =====================================================================
+
+
+class EdgeGatedGraphConvPure(nn.Module):
+ """Edge-gated graph convolution, pure torch / scatter-based."""
+
+ def __init__(
+ self, input_features: int, output_features: int, residual: bool = True
+ ):
+ super().__init__()
+ self.residual = residual
+ self.src_gate = nn.Linear(input_features, output_features)
+ self.dst_gate = nn.Linear(input_features, output_features)
+ self.edge_gate = nn.Linear(input_features, output_features)
+ self.bn_edges = nn.LayerNorm(output_features)
+ self.src_update = nn.Linear(input_features, output_features)
+ self.dst_update = nn.Linear(input_features, output_features)
+ self.bn_nodes = nn.LayerNorm(output_features)
+
+ def forward_tensors(
+ self,
+ src: torch.Tensor,
+ dst: torch.Tensor,
+ num_nodes: int,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Script-friendly core: takes plain tensors instead of a graph."""
+ # Edge update: m = e_src[src] + e_dst[dst] + edge_gate(y)
+ e_src = self.src_gate(x)
+ e_dst = self.dst_gate(x)
+ m = e_src[src] + e_dst[dst] + self.edge_gate(y)
+ sigma = torch.sigmoid(m)
+
+ # Node update aggregated at dst:
+ # h = sum_{(i,j)} sigma * Bh[src] / (sum sigma + eps).
+ Bh = self.dst_update(x)
+ msg_h = Bh[src] * sigma
+ sum_sigma_h = scatter_sum(msg_h, dst, num_nodes)
+ sum_sigma = scatter_sum(sigma, dst, num_nodes)
+ h = sum_sigma_h / (sum_sigma + 1e-6)
+ x_new = self.src_update(x) + h
+
+ x_new = F.silu(self.bn_nodes(x_new))
+ y_new = F.silu(self.bn_edges(m))
+
+ if self.residual:
+ x_new = x + x_new
+ y_new = y + y_new
+ return x_new, y_new
+
+ @torch.jit.ignore
+ def forward(
+ self, g: TorchGraph, x: torch.Tensor, y: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ return self.forward_tensors(g.src, g.dst, g.num_nodes, x, y)
+
+
+def _scatter_softmax(
+ logits: torch.Tensor, dst: torch.Tensor, num_nodes: int
+) -> torch.Tensor:
+ """Per-destination softmax over edges (pure torch, scatter-based)."""
+ max_per_dst = torch.full(
+ (num_nodes,) + logits.shape[1:],
+ float("-inf"),
+ dtype=logits.dtype,
+ device=logits.device,
+ )
+ idx = dst
+ while idx.dim() < logits.dim():
+ idx = idx.unsqueeze(-1)
+ idx_e = idx.expand_as(logits)
+ max_per_dst.scatter_reduce_(
+ 0, idx_e, logits, reduce="amax", include_self=True
+ )
+ shifted = logits - max_per_dst[dst]
+ exp_l = torch.exp(shifted)
+ sum_exp = scatter_sum(exp_l, dst, num_nodes)
+ return exp_l / (sum_exp[dst] + 1e-8)
+
+
+class NodeAttentionGraphConvPure(nn.Module):
+ """Node-Attention Layer (NAL) — pure torch.
+
+ Paper Eq. (8):
+ m_ij = A_src_i * L_src h_i + sum_j (A_dst_j * L_dst h_j + L_e e_ij)
+ with A_src, A_dst as per-node learned sigmoids over node features.
+ """
+
+ def __init__(
+ self,
+ input_features: int,
+ output_features: int,
+ residual: bool = True,
+ num_heads: int = 1,
+ ):
+ super().__init__()
+ self.residual = residual
+ self.num_heads = num_heads
+ self.output_features = output_features
+ assert output_features % num_heads == 0
+ self.head_dim = output_features // num_heads
+
+ self.src_gate = nn.Linear(input_features, output_features)
+ self.dst_gate = nn.Linear(input_features, output_features)
+ self.edge_gate = nn.Linear(input_features, output_features)
+ self.bn_edges = nn.LayerNorm(output_features)
+ self.src_update = nn.Linear(input_features, output_features)
+ self.dst_update = nn.Linear(input_features, output_features)
+ self.bn_nodes = nn.LayerNorm(output_features)
+
+ # Per-node attention heads: sigmoid(fc(h)) -> [N, num_heads]
+ self.attn_src = nn.Linear(input_features, num_heads)
+ self.attn_dst = nn.Linear(input_features, num_heads)
+
+ def _apply_attn(self, proj: torch.Tensor, a: torch.Tensor) -> torch.Tensor:
+ # proj: [N, F]; a: [N, heads]
+ N = proj.shape[0]
+ ph = proj.view(N, self.num_heads, self.head_dim)
+ return (a.unsqueeze(-1) * ph).view(N, self.output_features)
+
+ @torch.jit.ignore
+ def forward(
+ self, g: TorchGraph, x: torch.Tensor, y: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ src, dst, N = g.src, g.dst, g.num_nodes
+
+ a_src = torch.sigmoid(self.attn_src(x))
+ a_dst = torch.sigmoid(self.attn_dst(x))
+ e_src = self._apply_attn(self.src_gate(x), a_src)
+ e_dst = self._apply_attn(self.dst_gate(x), a_dst)
+
+ m = e_src[src] + e_dst[dst] + self.edge_gate(y)
+ sigma = torch.sigmoid(m)
+
+ Bh = self.dst_update(x)
+ sum_sigma_h = scatter_sum(Bh[src] * sigma, dst, N)
+ sum_sigma = scatter_sum(sigma, dst, N)
+ h = sum_sigma_h / (sum_sigma + 1e-6)
+ x_new = self.src_update(x) + h
+
+ x_new = F.silu(self.bn_nodes(x_new))
+ y_new = F.silu(self.bn_edges(m))
+ if self.residual:
+ x_new = x + x_new
+ y_new = y + y_new
+ return x_new, y_new
+
+
+class SelfAttentionGraphConvPure(nn.Module):
+ """Self-Attention Layer (SAL) — pure torch.
+
+ Paper Eqs. (9)-(11):
+ score = softmax((q_i . k_j . k_e) / sqrt(d))
+ h_i' = fc(h_i) + SiLU(Norm(sum_j score . v_j))
+ e_ij' = fc(e_ij) + SiLU(Norm(score))
+ """
+
+ def __init__(
+ self,
+ input_features: int,
+ output_features: int,
+ residual: bool = True,
+ num_heads: int = 1,
+ ):
+ super().__init__()
+ self.residual = residual
+ self.num_heads = num_heads
+ self.output_features = output_features
+ assert output_features % num_heads == 0
+ self.head_dim = output_features // num_heads
+ self.scale = self.head_dim**0.5
+
+ self.W_q = nn.Linear(input_features, output_features)
+ self.W_k = nn.Linear(input_features, output_features)
+ self.W_v = nn.Linear(input_features, output_features)
+ self.W_ke = nn.Linear(input_features, output_features)
+
+ self.fc_node = nn.Linear(input_features, output_features)
+ self.fc_edge = nn.Linear(input_features, output_features)
+ self.bn_nodes = nn.LayerNorm(output_features)
+ self.bn_edges = nn.LayerNorm(output_features)
+
+ @torch.jit.ignore
+ def forward(
+ self, g: TorchGraph, x: torch.Tensor, y: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ src, dst, N = g.src, g.dst, g.num_nodes
+ E = y.shape[0]
+ H, D = self.num_heads, self.head_dim
+
+ q = self.W_q(x).view(N, H, D)
+ k = self.W_k(x).view(N, H, D)
+ v = self.W_v(x).view(N, H, D)
+ k_e = self.W_ke(y).view(E, H, D)
+
+ # Attention logits per edge, per head: [E, H, 1]
+ attn_logits = (q[dst] * k[src] * k_e).sum(dim=-1, keepdim=True) / (
+ self.scale
+ )
+ # Softmax over incoming edges per destination, per head.
+ attn = _scatter_softmax(attn_logits.squeeze(-1), dst, N).unsqueeze(
+ -1
+ ) # [E, H, 1]
+
+ # Weighted value aggregation per destination node.
+ weighted = (attn * v[src]).view(E, self.output_features)
+ h_agg = scatter_sum(weighted, dst, N)
+ x_new = self.fc_node(x) + F.silu(self.bn_nodes(h_agg))
+
+ # Edge update uses the attention scores themselves (paper Eq. 11).
+ score_e = attn.expand(E, H, D).reshape(E, self.output_features)
+ y_new = self.fc_edge(y) + F.silu(self.bn_edges(score_e))
+ return x_new, y_new
+
+
+def _make_bond_conv(
+ in_features: int,
+ out_features: int,
+ conv_type: str,
+ num_heads: int,
+) -> nn.Module:
+ if conv_type == "alignn":
+ return EdgeGatedGraphConvPure(in_features, out_features)
+ if conv_type == "n_alignn":
+ return NodeAttentionGraphConvPure(
+ in_features, out_features, num_heads=num_heads
+ )
+ if conv_type == "t_alignn":
+ return SelfAttentionGraphConvPure(
+ in_features, out_features, num_heads=num_heads
+ )
+ raise ValueError(f"Unknown conv_type: {conv_type!r}")
+
+
+class ALIGNNConvPure(nn.Module):
+ """Line-graph-aware ALIGNN update."""
+
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ conv_type: str = "alignn",
+ num_heads: int = 1,
+ ):
+ super().__init__()
+ self.node_update = _make_bond_conv(
+ in_features, out_features, conv_type, num_heads
+ )
+ self.edge_update = _make_bond_conv(
+ out_features, out_features, conv_type, num_heads
+ )
+
+ def forward_tensors(
+ self,
+ g_src: torch.Tensor,
+ g_dst: torch.Tensor,
+ g_num_nodes: int,
+ lg_src: torch.Tensor,
+ lg_dst: torch.Tensor,
+ lg_num_nodes: int,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ z: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ x, m = self.node_update.forward_tensors(
+ g_src, g_dst, g_num_nodes, x, y
+ )
+ y, z = self.edge_update.forward_tensors(
+ lg_src, lg_dst, lg_num_nodes, m, z
+ )
+ return x, y, z
+
+ @torch.jit.ignore
+ def forward(
+ self,
+ g: TorchGraph,
+ lg: TorchGraph,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ z: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ # Attention convs (NAL/SAL) don't have forward_tensors — dispatch
+ # through the graph-object forward, which handles all conv types.
+ if not isinstance(self.node_update, EdgeGatedGraphConvPure):
+ x, m = self.node_update(g, x, y)
+ y, z = self.edge_update(lg, m, z)
+ return x, y, z
+ return self.forward_tensors(
+ g.src,
+ g.dst,
+ g.num_nodes,
+ lg.src,
+ lg.dst,
+ lg.num_nodes,
+ x,
+ y,
+ z,
+ )
+
+
+# =====================================================================
+# Helpers
+# =====================================================================
+
+
+def cutoff_function_based_edges(
+ r: torch.Tensor, inner_cutoff: float = 4.0, exponent: int = 3
+) -> torch.Tensor:
+ ratio = r / inner_cutoff
+ c1 = -(exponent + 1) * (exponent + 2) / 2
+ c2 = exponent * (exponent + 2)
+ c3 = -exponent * (exponent + 1) / 2
+ envelope = (
+ 1.0
+ + c1 * ratio**exponent
+ + c2 * ratio ** (exponent + 1)
+ + c3 * ratio ** (exponent + 2)
+ )
+ return torch.where(r <= inner_cutoff, envelope, torch.zeros_like(r))
+
+
+def _bond_cosines(r_ij: torch.Tensor, r_jk: torch.Tensor) -> torch.Tensor:
+ num = -(r_ij * r_jk).sum(dim=-1)
+ denom = r_ij.norm(dim=-1) * r_jk.norm(dim=-1)
+ return (num / denom.clamp_min(1e-12)).clamp(-1.0, 1.0)
+
+
+def _as_torchgraph(x):
+ """Accept TorchGraph or DGLGraph transparently."""
+ if isinstance(x, TorchGraph):
+ return x
+ return torchgraph_from_dgl(x)
+
+
+# =====================================================================
+# Model
+# =====================================================================
+
+
+class ALIGNNAtomWisePure(nn.Module):
+ """DGL-free ALIGNN atomwise."""
+
+ def __init__(
+ self,
+ config: ALIGNNAtomWisePureConfig = ALIGNNAtomWisePureConfig(
+ name="alignn_atomwise_pure"
+ ),
+ ):
+ super().__init__()
+ self.config = config
+ self.classification = config.classification
+ if self.config.gradwise_weight == 0:
+ self.config.calculate_gradient = False
+
+ self.atom_embedding = MLPLayer(
+ config.atom_input_features, config.hidden_features
+ )
+ self.edge_embedding = nn.Sequential(
+ RBFExpansion(vmin=0, vmax=8.0, bins=config.edge_input_features),
+ MLPLayer(config.edge_input_features, config.embedding_features),
+ MLPLayer(config.embedding_features, config.hidden_features),
+ )
+ self.angle_embedding = nn.Sequential(
+ RBFExpansion(
+ vmin=-1, vmax=1.0, bins=config.triplet_input_features
+ ),
+ MLPLayer(config.triplet_input_features, config.embedding_features),
+ MLPLayer(config.embedding_features, config.hidden_features),
+ )
+ self.alignn_layers = nn.ModuleList(
+ [
+ ALIGNNConvPure(
+ config.hidden_features,
+ config.hidden_features,
+ conv_type=config.conv_type,
+ num_heads=config.num_heads,
+ )
+ for _ in range(config.alignn_layers)
+ ]
+ )
+ self.gcn_layers = nn.ModuleList(
+ [
+ _make_bond_conv(
+ config.hidden_features,
+ config.hidden_features,
+ config.conv_type,
+ config.num_heads,
+ )
+ for _ in range(config.gcn_layers)
+ ]
+ )
+ if config.atomwise_output_features > 0:
+ self.fc_atomwise = nn.Linear(
+ config.hidden_features, config.atomwise_output_features
+ )
+ if config.additional_output_features:
+ self.fc_additional_output = nn.Linear(
+ config.hidden_features, config.additional_output_features
+ )
+ if self.classification:
+ self.fc = nn.Linear(config.hidden_features, 2)
+ self.softmax = nn.LogSoftmax(dim=1)
+ else:
+ self.fc = nn.Linear(config.hidden_features, config.output_features)
+ # Link is applied inline in the forward paths (kept as a string so
+ # `torch.jit.script` doesn't choke on a Python lambda attribute).
+ self.link_name: str = config.link
+ if config.link == "log":
+ self.fc.bias.data = torch.tensor(np.log(0.7), dtype=torch.float)
+
+ # Species feature lookup for LAMMPS / TorchScript entry point.
+ # Registered empty here; ``register_species_table`` populates it
+ # before export. Stays unused on the training path.
+ self.register_buffer(
+ "_species_table",
+ torch.zeros(120, config.atom_input_features, dtype=torch.float32),
+ )
+
+ # Expose scalar config fields that `forward_tensors` reads as
+ # typed Python attributes — the pydantic ``self.config`` object
+ # is not TorchScript-compatible.
+ self.energy_mult_natoms: bool = bool(config.energy_mult_natoms)
+ self.use_penalty: bool = bool(config.use_penalty)
+ self.penalty_threshold: float = float(config.penalty_threshold)
+ self.penalty_factor: float = float(config.penalty_factor)
+ self.grad_multiplier: int = int(config.grad_multiplier)
+ self.add_reverse_forces: bool = bool(config.add_reverse_forces)
+ self.stress_multiplier: float = float(config.stress_multiplier)
+
+ # ----- species lookup (LAMMPS hands us atomic numbers, not features) -----
+
+ @torch.jit.ignore
+ def register_species_table(
+ self, atom_features: str = "cgcnn", max_z: int = 119
+ ) -> None:
+ """Build a (max_z+1, F) feature table indexed by atomic number.
+
+ Once registered, ``forward_tensors_z`` can be used from TorchScript
+ so callers (e.g. LAMMPS) don't need to ship the feature table
+ alongside the model. Invalid / unknown atomic numbers map to
+ zero vectors.
+ """
+ from ase.data import chemical_symbols
+ from jarvis.core.specie import get_node_attributes
+
+ F = int(self.config.atom_input_features)
+ rows = np.zeros((max_z + 1, F), dtype=np.float32)
+ for z in range(1, max_z + 1):
+ if z >= len(chemical_symbols):
+ continue
+ symbol = chemical_symbols[z]
+ try:
+ feat = np.asarray(
+ get_node_attributes(symbol, atom_features=atom_features),
+ dtype=np.float32,
+ )
+ if feat.shape[0] != F:
+ continue
+ rows[z] = feat
+ except Exception:
+ continue
+ # Buffer was allocated in __init__; overwrite the tensor contents
+ # (register_buffer would complain about duplicate registration).
+ self._species_table.resize_(rows.shape).copy_(
+ torch.as_tensor(rows, dtype=torch.float32)
+ )
+
+ # ----- tensor-only forward (LAMMPS / TorchScript path) -----
+
+ @torch.jit.export
+ def forward_tensors_z(
+ self,
+ positions: torch.Tensor, # (N, 3), requires_grad=True for forces
+ lattice: torch.Tensor, # (3, 3)
+ atomic_numbers: torch.Tensor, # (N,) long
+ src: torch.Tensor, # (E,) long
+ dst: torch.Tensor, # (E,) long
+ shift: torch.Tensor, # (E, 3)
+ compute_stress: bool = False,
+ ) -> Dict[str, torch.Tensor]:
+ """Forward with internal atomic-number → feature lookup.
+
+ Same as ``forward_tensors`` but accepts atomic numbers directly.
+ Requires ``register_species_table`` to have been called
+ beforehand. LAMMPS-style callers only ship atomic numbers.
+ """
+ atom_features = self._species_table.index_select(0, atomic_numbers)
+ return self.forward_tensors(
+ positions, lattice, atom_features, src, dst, shift, compute_stress
+ )
+
+ @torch.jit.export
+ def forward_tensors(
+ self,
+ positions: torch.Tensor, # (N, 3), requires_grad=True for forces
+ lattice: torch.Tensor, # (3, 3), requires_grad=True for stress
+ atom_features: torch.Tensor, # (N, atom_input_features)
+ src: torch.Tensor, # (E,) long
+ dst: torch.Tensor, # (E,) long
+ shift: torch.Tensor, # (E, 3) integer cell offsets (float dtype)
+ compute_stress: bool = False,
+ ) -> Dict[str, torch.Tensor]:
+ """Single-system forward driven by plain tensors.
+
+ Intended for LAMMPS-style MD where each step supplies an
+ atom list, lattice, and neighbor list. No batching, no class
+ pooling machinery. Returns ``{"energy", "forces", ["stress"]}``.
+ """
+ num_nodes = positions.shape[0]
+
+ # Differentiable edge vectors.
+ r = (
+ positions.index_select(0, dst)
+ - positions.index_select(0, src)
+ + torch.matmul(shift, lattice)
+ )
+ bondlength = torch.linalg.vector_norm(r, dim=1)
+
+ # Inline line-graph: connect parent edge A=(u,v) to B=(v,w).
+ E = src.shape[0]
+ order = torch.argsort(src, stable=True)
+ sorted_src = src.index_select(0, order)
+ node_range = torch.arange(num_nodes, device=src.device)
+ bucket_start = torch.searchsorted(sorted_src, node_range)
+ bucket_end = torch.searchsorted(sorted_src, node_range, right=True)
+
+ A_v = dst
+ starts = bucket_start.index_select(0, A_v)
+ ends = bucket_end.index_select(0, A_v)
+ counts = ends - starts
+
+ total = int(counts.sum().item())
+ A_ids = torch.arange(E, device=src.device)
+ lg_src = torch.repeat_interleave(A_ids, counts)
+ cum = torch.cumsum(counts, dim=0)
+ row_start = cum - counts
+ offsets = torch.arange(
+ total, device=src.device
+ ) - torch.repeat_interleave(row_start, counts)
+ pos_idx = torch.repeat_interleave(starts, counts) + offsets
+ lg_dst = order.index_select(0, pos_idx)
+ lg_num_nodes = E
+
+ # Angle cosines (differentiable through r -> positions, lattice).
+ r_ij = r.index_select(0, lg_src)
+ r_jk = r.index_select(0, lg_dst)
+ num = -(r_ij * r_jk).sum(dim=-1)
+ denom = torch.linalg.vector_norm(r_ij, dim=-1).clamp_min(
+ 1e-12
+ ) * torch.linalg.vector_norm(r_jk, dim=-1).clamp_min(1e-12)
+ h_cos = (num / denom).clamp(-1.0, 1.0)
+
+ # Embeddings.
+ x = self.atom_embedding(atom_features)
+ y = self.edge_embedding(bondlength)
+ z = self.angle_embedding(h_cos)
+
+ # ALIGNN and GCN layers via tensor-only paths.
+ for layer in self.alignn_layers:
+ x, y, z = layer.forward_tensors(
+ src, dst, num_nodes, lg_src, lg_dst, lg_num_nodes, x, y, z
+ )
+ for layer in self.gcn_layers:
+ x, y = layer.forward_tensors(src, dst, num_nodes, x, y)
+
+ # Mean readout (single graph) and final projection.
+ h_graph = x.mean(dim=0, keepdim=True) # (1, hidden)
+ out = self.fc(h_graph).squeeze(-1) # (1,) for output_features=1
+
+ en = out
+ if self.energy_mult_natoms:
+ en = out * float(num_nodes)
+
+ if self.use_penalty:
+ pen = torch.where(
+ bondlength < self.penalty_threshold,
+ self.penalty_factor * (self.penalty_threshold - bondlength),
+ torch.zeros_like(bondlength),
+ )
+ en = en + pen.sum()
+
+ result: Dict[str, torch.Tensor] = {"energy": en.squeeze()}
+
+ # Autograd forces through r -> positions.
+ grad_outs = torch.autograd.grad(
+ outputs=[en.sum()],
+ inputs=[r],
+ create_graph=self.training,
+ retain_graph=True,
+ )
+ pair_forces = grad_outs[0]
+ assert pair_forces is not None
+ pair_forces = self.grad_multiplier * pair_forces
+
+ # Atom forces = -dE/dx = sum of dE/dr over incoming edges
+ # minus sum over outgoing (add_reverse_forces convention).
+ forces_ji = scatter_sum(pair_forces, dst, num_nodes)
+ if self.add_reverse_forces:
+ forces_ij = scatter_sum(pair_forces, src, num_nodes)
+ forces = forces_ji - forces_ij
+ else:
+ forces = forces_ji
+ result["forces"] = forces
+
+ if compute_stress:
+ # Virial stress: -160.217 * r.T @ pair_forces / V.
+ V = torch.abs(torch.det(lattice))
+ stress = -160.21766208 * (torch.matmul(r.T, pair_forces) / V)
+ result["stress"] = self.stress_multiplier * stress
+
+ return result
+
+ # ----- dataclass-based forward (training pipeline) -----
+
+ @torch.jit.ignore
+ def forward(self, g):
+ # Unpack and normalize input.
+ gg, lg, lat = g
+ g = _as_torchgraph(gg)
+ lg = _as_torchgraph(lg)
+
+ r = g.edata["r"]
+ if self.config.calculate_gradient:
+ r.requires_grad_(True)
+
+ bondlength = torch.norm(r, dim=1)
+
+ # Recompute angle cosines inside the autograd graph when lg_on_fly.
+ if self.config.lg_on_fly and len(self.alignn_layers) > 0:
+ h = _bond_cosines(r[lg.src], r[lg.dst])
+ z = self.angle_embedding(h)
+ elif len(self.alignn_layers) > 0:
+ z = self.angle_embedding(lg.edata["h"])
+ else:
+ z = None # unused
+
+ x = self.atom_embedding(g.ndata["atom_features"])
+
+ if self.config.use_cutoff_function:
+ if self.config.multiply_cutoff:
+ c_off = cutoff_function_based_edges(
+ bondlength,
+ inner_cutoff=self.config.inner_cutoff,
+ exponent=self.config.exponent,
+ ).unsqueeze(-1)
+ y = self.edge_embedding(bondlength) * c_off
+ else:
+ bondlength_eff = cutoff_function_based_edges(
+ bondlength,
+ inner_cutoff=self.config.inner_cutoff,
+ exponent=self.config.exponent,
+ )
+ y = self.edge_embedding(bondlength_eff)
+ else:
+ y = self.edge_embedding(bondlength)
+
+ for layer in self.alignn_layers:
+ x, y, z = layer(g, lg, x, y, z)
+ for layer in self.gcn_layers:
+ x, y = layer(g, x, y)
+
+ # Per-graph readout (mean over nodes).
+ node_bid = g.node_batch_id
+ h_graph = scatter_mean(x, node_bid, g.batch_size)
+ out = self.fc(h_graph)
+ if self.config.output_features == 1:
+ out = out.squeeze(-1)
+
+ natoms = (
+ g.batch_num_nodes
+ if g.batch_num_nodes is not None
+ else torch.tensor([g.num_nodes], device=r.device)
+ )
+ en_out = out
+ if self.config.energy_mult_natoms:
+ en_out = out * natoms.to(out.dtype)
+
+ if self.config.use_penalty:
+ penalties = torch.where(
+ bondlength < self.config.penalty_threshold,
+ self.config.penalty_factor
+ * (self.config.penalty_threshold - bondlength),
+ torch.zeros_like(bondlength),
+ )
+ en_out = en_out + penalties.sum()
+
+ atomwise_pred = torch.empty(1, device=r.device)
+ if (
+ self.config.atomwise_output_features > 0
+ and self.config.atomwise_weight != 0
+ ):
+ atomwise_pred = self.fc_atomwise(x)
+
+ additional_out = torch.empty(1, device=r.device)
+ if self.config.additional_output_features > 0:
+ additional_out = self.fc_additional_output(h_graph)
+
+ forces = torch.empty(1, device=r.device)
+ stress = torch.empty(1, device=r.device)
+ if self.config.calculate_gradient:
+ pair_forces = (
+ self.config.grad_multiplier
+ * grad(
+ en_out.sum(),
+ r,
+ grad_outputs=torch.ones_like(en_out.sum()),
+ create_graph=True,
+ retain_graph=True,
+ )[0]
+ )
+ if self.config.force_mult_natoms:
+ pair_forces = pair_forces * g.num_nodes
+
+ # force_i from incoming edges (dst == i): sum pair_forces at dst.
+ forces_ji = scatter_sum(pair_forces, g.dst, g.num_nodes)
+ if self.config.add_reverse_forces:
+ forces_ij = scatter_sum(pair_forces, g.src, g.num_nodes)
+ forces = (forces_ji - forces_ij).squeeze()
+ else:
+ forces = forces_ji.squeeze()
+
+ if self.config.stresswise_weight != 0:
+ # Per-graph stress = -160.21766208 * (r.T @ pair_forces) / V.
+ # V is stored per-node; take V at the first node of each graph.
+ B = g.batch_size
+ V_per_node = g.ndata["V"]
+ if g.batch_num_nodes is not None:
+ node_offsets = torch.zeros(
+ B, dtype=torch.long, device=r.device
+ )
+ node_offsets[1:] = torch.cumsum(
+ g.batch_num_nodes[:-1], dim=0
+ )
+ edge_offsets = torch.zeros(
+ B, dtype=torch.long, device=r.device
+ )
+ edge_offsets[1:] = torch.cumsum(
+ g.batch_num_edges[:-1], dim=0
+ )
+ else:
+ node_offsets = torch.zeros(
+ 1, dtype=torch.long, device=r.device
+ )
+ edge_offsets = torch.zeros(
+ 1, dtype=torch.long, device=r.device
+ )
+
+ stresses = []
+ for b in range(B):
+ e0 = int(edge_offsets[b].item())
+ e1 = e0 + (
+ int(g.batch_num_edges[b].item())
+ if g.batch_num_edges is not None
+ else g.num_edges
+ )
+ n0 = int(node_offsets[b].item())
+ V_b = V_per_node[n0]
+ st = -160.21766208 * (
+ torch.matmul(r[e0:e1].T, pair_forces[e0:e1]) / V_b
+ )
+ stresses.append(st)
+ stress = self.config.stress_multiplier * torch.stack(stresses)
+
+ if self.link_name == "log":
+ out = torch.exp(out)
+ elif self.link_name == "logit":
+ out = torch.sigmoid(out)
+ if self.classification:
+ if out.dim() == 1:
+ out = out.unsqueeze(0)
+ out = self.softmax(out)
+
+ return {
+ "out": out,
+ "additional": additional_out,
+ "grad": forces,
+ "stresses": stress,
+ "atomwise_pred": atomwise_pred,
+ }
diff --git a/alignn/models/ealignn_atomwise.py b/alignn/models/ealignn_atomwise.py
index d501090..a7d8088 100644
--- a/alignn/models/ealignn_atomwise.py
+++ b/alignn/models/ealignn_atomwise.py
@@ -3,13 +3,21 @@
A prototype crystal line graph network dgl implementation.
"""
+from __future__ import annotations
+
from typing import Tuple, Union
from torch.autograd import grad
-import dgl
-import dgl.function as fn
+
+try:
+ import dgl
+ import dgl.function as fn
+ from dgl.nn import AvgPooling
+except Exception:
+ dgl = None
+ fn = None
+ AvgPooling = None
# import numpy as np
-from dgl.nn import AvgPooling
import torch
# from dgl.nn.functional import edge_softmax
diff --git a/alignn/models/model_zoo.py b/alignn/models/model_zoo.py
new file mode 100644
index 0000000..64b0be8
--- /dev/null
+++ b/alignn/models/model_zoo.py
@@ -0,0 +1,246 @@
+"""Bundle multiple pretrained models behind one nn.Module.
+
+Each entry is a (model, output_spec) pair. ``output_spec`` says how to map
+the underlying model's raw output dict to a single tensor for that
+property — e.g. EFS returns three tensors, formation energy returns one.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+import torch
+from torch import nn
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure,
+ ALIGNNAtomWisePureConfig,
+)
+
+
+@dataclass
+class PropertySpec:
+ """How to invoke one member model and what to extract from it."""
+
+ name: str
+ model: nn.Module
+ # Map raw model output (whatever ``model(graph)`` returns) → dict of
+ # named tensors for this property. Default pulls "out" from an
+ # ALIGNN-Pure dict.
+ extract: Callable[[object], Dict[str, torch.Tensor]] = field(
+ default=lambda d: {"value": d["out"]}
+ )
+
+
+def _efs_extract(d):
+ return {"energy": d["out"], "forces": d["grad"], "stress": d["stresses"]}
+
+
+def _scalar_extract(d):
+ return {"value": d["out"]}
+
+
+class ModelZoo(nn.Module):
+ """Registry of pretrained models exposed as one nn.Module.
+
+ All members live in a single ``nn.ModuleDict`` so ``state_dict``,
+ ``.to(device)``, ``torch.save`` etc. work on the whole collection.
+ """
+
+ def __init__(self, specs: Optional[List[PropertySpec]] = None):
+ super().__init__()
+ self._models = nn.ModuleDict()
+ self._extractors: Dict[str, Callable] = {}
+ for spec in specs or []:
+ self.register(spec)
+
+ def register(self, spec: PropertySpec) -> None:
+ if spec.name in self._models:
+ raise KeyError(f"property {spec.name!r} already registered")
+ self._models[spec.name] = spec.model
+ self._extractors[spec.name] = spec.extract
+
+ def properties(self) -> List[str]:
+ return list(self._models.keys())
+
+ def predict(self, graph, prop: str) -> Dict[str, torch.Tensor]:
+ out = self._models[prop](graph)
+ return self._extractors[prop](out)
+
+ def predict_all(self, graph) -> Dict[str, Dict[str, torch.Tensor]]:
+ return {p: self.predict(graph, p) for p in self._models}
+
+ # default forward = predict_all so torch.compile sees one entry point
+ def forward(self, graph) -> Dict[str, Dict[str, torch.Tensor]]:
+ return self.predict_all(graph)
+
+ # ----- convenience constructors -----
+
+ # ----- manifest-based lazy loading -----
+
+ @classmethod
+ def from_manifest(
+ cls,
+ manifest_path: str,
+ map_location: str = "cpu",
+ eager: bool = False,
+ ) -> "LazyModelZoo":
+ """Build a zoo whose members live on disk and load on first use.
+
+ Manifest schema (JSON):
+ {
+ "base_dir": "optional; paths are resolved relative to this",
+ "properties": {
+ "": {
+ "checkpoint": "",
+ "config": { ... ALIGNNAtomWisePureConfig kwargs ... },
+ "is_efs": false
+ },
+ ...
+ }
+ }
+ """
+ with open(manifest_path) as f:
+ spec = json.load(f)
+ base = Path(spec.get("base_dir") or Path(manifest_path).parent)
+ zoo = LazyModelZoo(map_location=map_location)
+ for name, entry in spec["properties"].items():
+ ckpt = base / entry["checkpoint"]
+ cfg_kwargs = dict(entry.get("config") or {})
+ cfg = ALIGNNAtomWisePureConfig(**cfg_kwargs)
+ zoo.add_lazy(
+ name=name,
+ checkpoint=str(ckpt),
+ config=cfg,
+ is_efs=bool(entry.get("is_efs", False)),
+ )
+ if eager:
+ for name in zoo.properties():
+ zoo._materialize(name)
+ return zoo
+
+ @classmethod
+ def from_alignn_checkpoints(
+ cls,
+ ckpts: Dict[str, str],
+ configs: Optional[Dict[str, ALIGNNAtomWisePureConfig]] = None,
+ efs_keys: Optional[List[str]] = None,
+ map_location: str = "cpu",
+ ) -> "ModelZoo":
+ """Build a zoo from {property_name: checkpoint_path}.
+
+ ``efs_keys`` lists names whose model returns energy/forces/stress
+ (everything else is treated as a single scalar head).
+ ``configs`` lets you pass a per-model config; missing entries fall
+ back to the default ``ALIGNNAtomWisePureConfig()``.
+ """
+ configs = configs or {}
+ efs = set(efs_keys or [])
+ specs: List[PropertySpec] = []
+ for name, path in ckpts.items():
+ cfg = configs.get(name, ALIGNNAtomWisePureConfig(name=name))
+ model = ALIGNNAtomWisePure(cfg)
+ state = torch.load(path, map_location=map_location)
+ if isinstance(state, dict) and "model" in state:
+ state = state["model"]
+ model.load_state_dict(state)
+ model.eval()
+ specs.append(
+ PropertySpec(
+ name=name,
+ model=model,
+ extract=_efs_extract if name in efs else _scalar_extract,
+ )
+ )
+ return cls(specs)
+
+
+class LazyModelZoo(ModelZoo):
+ """Zoo whose members are loaded from disk on first use.
+
+ Only one ALIGNN trunk's weights occupy memory until you actually
+ request a property. Materialized members are cached, so repeated
+ calls hit RAM, not disk.
+ """
+
+ def __init__(self, map_location: str = "cpu"):
+ super().__init__()
+ self._map_location = map_location
+ # name -> {checkpoint, config, is_efs}
+ self._lazy: Dict[str, dict] = {}
+
+ def add_lazy(
+ self,
+ name: str,
+ checkpoint: str,
+ config: ALIGNNAtomWisePureConfig,
+ is_efs: bool = False,
+ ) -> None:
+ if name in self._lazy or name in self._models:
+ raise KeyError(f"property {name!r} already registered")
+ self._lazy[name] = {
+ "checkpoint": checkpoint,
+ "config": config,
+ "is_efs": is_efs,
+ }
+ self._extractors[name] = _efs_extract if is_efs else _scalar_extract
+
+ def properties(self) -> List[str]:
+ return list(self._lazy.keys()) + [
+ n for n in self._models.keys() if n not in self._lazy
+ ]
+
+ def _materialize(self, name: str) -> nn.Module:
+ if name in self._models:
+ return self._models[name]
+ if name not in self._lazy:
+ raise KeyError(name)
+ entry = self._lazy[name]
+ model = ALIGNNAtomWisePure(entry["config"])
+ state = torch.load(
+ entry["checkpoint"], map_location=self._map_location
+ )
+ if isinstance(state, dict) and "model" in state:
+ state = state["model"]
+ model.load_state_dict(state)
+ model.eval()
+ self._models[name] = model
+ return model
+
+ def predict(self, graph, prop: str) -> Dict[str, torch.Tensor]:
+ self._materialize(prop)
+ return super().predict(graph, prop)
+
+ def evict(self, name: str) -> None:
+ """Drop a materialized member back to disk-only state."""
+ if name in self._models and name in self._lazy:
+ del self._models[name]
+
+
+# ----- example: registering a non-ALIGNN model (e.g. SlakoNet) ---------------
+#
+# from slakonet import SlakoNetModel
+#
+# zoo = ModelZoo.from_alignn_checkpoints(
+# {
+# "efs": "ff/alignnff.pt",
+# "formation_energy": "jv_formation_energy_peratom_alignn.pt",
+# "mbj_bandgap": "jv_mbj_bandgap_alignn.pt",
+# "bulk_modulus": "jv_bulk_modulus_kv_alignn.pt",
+# },
+# efs_keys=["efs"],
+# )
+#
+# slakonet = SlakoNetModel.load("slakonet.pt").eval()
+# zoo.register(PropertySpec(
+# name="bandstructure",
+# model=slakonet,
+# extract=lambda d: {"bandgap": d["bandgap"], "bands": d["bands"]},
+# ))
+#
+# torch.save(zoo.state_dict(), "alignn_zoo.pt")
+# preds = zoo.predict_all(graph) # {prop: {field: tensor}}
+# fe = zoo.predict(graph, "formation_energy")["value"]
diff --git a/alignn/models/utils.py b/alignn/models/utils.py
index cc02dc0..b68c53a 100644
--- a/alignn/models/utils.py
+++ b/alignn/models/utils.py
@@ -1,12 +1,17 @@
"""Shared model-building components."""
-from typing import Optional, Callable
+from typing import Optional, Callable, TYPE_CHECKING
import numpy as np
import torch
import torch.nn as nn
-import dgl
from typing import Tuple
+# Lazy-import dgl: only DGL-graph-specific helpers below need it, and
+# we want importing this module to stay DGL-free so the pure-torch
+# model can be loaded in environments without DGL.
+if TYPE_CHECKING: # pragma: no cover
+ import dgl
+
class RBFExpansion(nn.Module):
"""Expand interatomic distances with radial basis functions."""
@@ -28,14 +33,14 @@ def __init__(
)
if lengthscale is None:
- # SchNet-style
- # set lengthscales relative to granularity of RBF expansion
- self.lengthscale = np.diff(self.centers).mean()
- self.gamma = 1 / self.lengthscale
-
+ # SchNet-style: set lengthscales relative to RBF granularity.
+ # Cast to Python float so TorchScript can type-infer the
+ # attribute (numpy scalars aren't recognized).
+ self.lengthscale = float(np.diff(self.centers).mean())
+ self.gamma = 1.0 / self.lengthscale
else:
- self.lengthscale = lengthscale
- self.gamma = 1 / (lengthscale**2)
+ self.lengthscale = float(lengthscale)
+ self.gamma = 1.0 / (self.lengthscale**2)
def forward(self, distance: torch.Tensor) -> torch.Tensor:
"""Apply RBF expansion to interatomic distance tensor."""
@@ -44,7 +49,7 @@ def forward(self, distance: torch.Tensor) -> torch.Tensor:
)
-def compute_pair_vector_and_distance(g: dgl.DGLGraph):
+def compute_pair_vector_and_distance(g: "dgl.DGLGraph"):
"""Calculate bond vectors and distances using dgl graphs."""
# print('g.edges()',g.ndata["cart_coords"][g.edges()[1]].shape,g.edata["pbc_offshift"].shape)
dst_pos = g.ndata["cart_coords"][g.edges()[1]] + g.edata["images"]
@@ -127,10 +132,10 @@ def compute_cartesian_coordinates(g, lattice, dtype=torch.float32):
def lightweight_line_graph(
- input_graph: dgl.DGLGraph,
+ input_graph: "dgl.DGLGraph",
feature_name: str,
filter_condition: Callable[[torch.Tensor], torch.Tensor],
-) -> dgl.DGLGraph:
+) -> "dgl.DGLGraph":
"""Make the line graphs lightweight with preserved node ordering.
Handles both batched and unbatched graphs.
@@ -143,6 +148,8 @@ def lightweight_line_graph(
New DGL graph with filtered edges preserving original node ordering
"""
# Check if graph is batched
+ import dgl
+
is_batched = (
hasattr(input_graph, "batch_size") and input_graph.batch_size > 1
)
@@ -223,10 +230,10 @@ def lightweight_line_graph(
def lightweight_line_graph1(
- input_graph: dgl.DGLGraph,
+ input_graph: "dgl.DGLGraph",
feature_name: str,
filter_condition: Callable[[torch.Tensor], torch.Tensor],
-) -> dgl.DGLGraph:
+) -> "dgl.DGLGraph":
"""Make the line graphs lightweight with preserved node ordering.
Args:
@@ -237,6 +244,8 @@ def lightweight_line_graph1(
Returns:
Filtered edges while preserving original node ordering
"""
+ import dgl
+
# Get active edges based on filter condition
active_edges = torch.logical_not(
filter_condition(input_graph.edata[feature_name])
@@ -317,7 +326,7 @@ def compute_net_torque(
def remove_net_torque(
- g: dgl.DGLGraph,
+ g: "dgl.DGLGraph",
forces: torch.Tensor,
n_nodes: torch.Tensor,
) -> torch.Tensor:
diff --git a/alignn/pure_lmdb_dataset.py b/alignn/pure_lmdb_dataset.py
new file mode 100644
index 0000000..35e2159
--- /dev/null
+++ b/alignn/pure_lmdb_dataset.py
@@ -0,0 +1,345 @@
+"""LMDB-backed dataset yielding TorchGraph objects (no DGL).
+
+Mirrors the API of ``alignn.lmdb_dataset`` but serializes
+``TorchGraph`` dataclasses via pickle, and batches them with
+``batch_torch_graph_pairs`` at the DataLoader boundary.
+"""
+
+from __future__ import annotations
+
+import os
+import pickle as pk
+from typing import List, Tuple
+
+import lmdb
+import numpy as np
+import torch
+from jarvis.core.atoms import Atoms
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+from alignn.graphs import Graph
+from alignn.torch_graph_builder import (
+ TorchGraph,
+ batch_torch_graph_pairs,
+ batch_torch_graphs,
+ torchgraph_from_dgl,
+)
+
+# How often to commit the LMDB write transaction during a fresh build.
+# Holding a single transaction across all entries keeps every dirty page
+# resident in RAM until commit; flushing periodically caps peak memory.
+LMDB_COMMIT_EVERY = 100000
+
+
+def _mem_gb():
+ """Resident memory of the current process in GB (for diagnostics)."""
+ try:
+ import psutil
+
+ return psutil.Process(os.getpid()).memory_info().rss / 1e9
+ except Exception:
+ return float("nan")
+
+
+def prepare_pure_batch(batch, device=None, non_blocking=False):
+ """Move a pure-torch batch to device; mirrors prepare_line_graph_batch."""
+ g, lg, t, id = batch
+ return (
+ (g.to(device), lg.to(device)),
+ t.to(device, non_blocking=non_blocking),
+ )
+
+
+def _graph_to_torchgraph(g, lg=None):
+ """Accept DGL or TorchGraph; always return TorchGraph (pair)."""
+ if isinstance(g, TorchGraph):
+ return g, (lg if lg is not None else None)
+ tg = torchgraph_from_dgl(g)
+ tlg = torchgraph_from_dgl(lg) if lg is not None else None
+ return tg, tlg
+
+
+class PureTorchLMDBDataset(Dataset):
+ """Dataset of crystal TorchGraphs using LMDB."""
+
+ def __init__(self, lmdb_path: str = "", line_graph: bool = True, ids=None):
+ """Open the LMDB at ``lmdb_path`` read-only."""
+ super().__init__()
+ 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:
+ self.length = txn.stat()["entries"]
+ self.prepare_batch = prepare_pure_batch
+
+ 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:
+ serialized = txn.get(f"{idx}".encode())
+ if self.line_graph:
+ g, lg, lattice, label = pk.loads(serialized)
+ return g, lg, lattice, label
+ g, lattice, label = pk.loads(serialized)
+ return g, lattice, label
+
+ def close(self):
+ """Close the LMDB environment (idempotent)."""
+ try:
+ self.env.close()
+ except Exception:
+ pass
+
+ def __del__(self):
+ """Ensure the LMDB handle is released on GC."""
+ self.close()
+
+ @staticmethod
+ def collate(samples: List[Tuple[TorchGraph, torch.Tensor, torch.Tensor]]):
+ """Collate ``(g, lattice, label)`` samples into a batched triple."""
+ graphs, lattices, labels = map(list, zip(*samples))
+ batched = batch_torch_graphs(graphs)
+ if labels[0].dim() > 0:
+ return batched, torch.stack(lattices), torch.stack(labels)
+ return batched, torch.stack(lattices), torch.tensor(labels)
+
+ @staticmethod
+ def collate_line_graph(
+ samples: List[
+ Tuple[TorchGraph, TorchGraph, torch.Tensor, torch.Tensor]
+ ],
+ ):
+ """Collate ``(g, lg, lattice, label)`` samples into a batched tuple."""
+ graphs, line_graphs, lattices, labels = map(list, zip(*samples))
+ g_b, lg_b = batch_torch_graph_pairs(list(zip(graphs, line_graphs)))
+ if labels[0].dim() > 0:
+ return g_b, lg_b, torch.stack(lattices), torch.stack(labels)
+ return g_b, lg_b, torch.stack(lattices), torch.tensor(labels)
+
+
+def _attach_node_payload(
+ g: TorchGraph, key: str, value: np.ndarray, natoms: int
+):
+ """Tile a per-structure tensor across nodes, like the DGL loader."""
+ 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)
+
+
+def _lmdb_is_usable(tmp_name: str) -> bool:
+ """Return True if the LMDB at ``tmp_name`` contains at least one entry.
+
+ Guards against the empty-stub left behind when a previous build was
+ interrupted (OOM, Ctrl-C, etc.). Such stubs have an 8KB data.mdb but
+ zero records, and the read_existing fast path would silently return
+ a zero-length dataset if not caught.
+ """
+ if not os.path.exists(tmp_name):
+ return False
+ try:
+ env = lmdb.open(tmp_name, readonly=True, lock=False)
+ with env.begin() as txn:
+ n_entries = txn.stat()["entries"]
+ env.close()
+ return n_entries > 0
+ except Exception:
+ return False
+
+
+def get_torch_dataset(
+ dataset=None,
+ id_tag="jid",
+ target="",
+ target_atomwise="",
+ target_grad="",
+ target_stress="",
+ target_additional_output="",
+ neighbor_strategy="pure_torch",
+ atom_features="cgcnn",
+ use_canonize="",
+ name="",
+ line_graph=True,
+ cutoff=8.0,
+ cutoff_extra=3.0,
+ max_neighbors=12,
+ three_body_cutoff=None,
+ classification=False,
+ sampler=None,
+ output_dir=".",
+ tmp_name="dataset",
+ map_size=1e12,
+ read_existing=False,
+ dtype="float32",
+):
+ """Build or load a PureTorchLMDBDataset from a list of records."""
+ dataset = dataset or []
+ vals = np.array([ii[target] for ii in dataset])
+ print("data range", np.max(vals), np.min(vals))
+ print("line_graph", line_graph)
+ os.makedirs(output_dir, exist_ok=True)
+ with open(os.path.join(output_dir, tmp_name + "_data_range"), "w") as f:
+ f.write(f"Max={np.max(vals)}\nMin={np.min(vals)}\n")
+
+ print(f"[MEM] {tmp_name}: entry: {_mem_gb():.2f} GB")
+
+ # Fast path: reuse a previously built cache, but only if it actually
+ # contains records. An empty stub (8KB data.mdb) from an interrupted
+ # build would otherwise be returned as a zero-length dataset.
+ if read_existing and _lmdb_is_usable(tmp_name):
+ # Validate the cache was built for the pure-torch backend.
+ _env = lmdb.open(tmp_name, readonly=True, lock=False)
+ with _env.begin() as _txn:
+ _probe = _txn.get(b"0")
+ _env.close()
+ if _probe is not None:
+ _sample = pk.loads(_probe)
+ if not isinstance(_sample[0], TorchGraph):
+ raise RuntimeError(
+ f"LMDB cache at '{tmp_name}' contains "
+ f"{type(_sample[0]).__name__} records, not "
+ "TorchGraph. Delete the stale cache (e.g. "
+ f"`rm -rf {tmp_name}`) and rerun — it was built "
+ "for a different model backend."
+ )
+ ids = [d[id_tag] for d in dataset]
+ print("Reading dataset", tmp_name)
+ ds = PureTorchLMDBDataset(
+ lmdb_path=tmp_name, line_graph=line_graph, ids=ids
+ )
+ print(
+ f"[MEM] {tmp_name}: after PureTorchLMDBDataset (cached): "
+ f"{_mem_gb():.2f} GB"
+ )
+ return ds
+
+ # If read_existing was requested but the cache is unusable (empty
+ # stub or missing), warn so the user knows we're rebuilding.
+ if read_existing and os.path.exists(tmp_name):
+ print(
+ f"[WARN] {tmp_name}: read_existing=True but cache is "
+ "empty or unreadable; rebuilding from scratch."
+ )
+
+ ids = []
+ # Fresh build: wipe any pre-existing cache dir to avoid mixing old
+ # records (possibly from a different model backend) with new writes.
+ if os.path.exists(tmp_name):
+ import shutil
+
+ shutil.rmtree(tmp_name)
+ env = lmdb.open(tmp_name, map_size=int(map_size))
+
+ print(f"[MEM] {tmp_name}: before write loop: {_mem_gb():.2f} GB")
+
+ # Use periodic commits instead of a single giant write transaction.
+ # A single txn around millions of put() calls keeps every dirty page
+ # in RAM until commit and was the cause of OOM on memory-tight hosts.
+ txn = env.begin(write=True)
+ try:
+ for idx, d in tqdm(enumerate(dataset), total=len(dataset)):
+ ids.append(d[id_tag])
+ atoms = Atoms.from_dict(d["atoms"])
+ raw = Graph.atom_dgl_multigraph(
+ atoms,
+ cutoff=float(cutoff),
+ max_neighbors=max_neighbors,
+ atom_features=atom_features,
+ compute_line_graph=line_graph,
+ use_canonize=use_canonize,
+ cutoff_extra=cutoff_extra,
+ neighbor_strategy=neighbor_strategy,
+ three_body_cutoff=three_body_cutoff,
+ dtype=dtype,
+ )
+ if line_graph:
+ g, lg = raw
+ g, lg = _graph_to_torchgraph(g, lg)
+ else:
+ g = raw
+ g, _ = _graph_to_torchgraph(g)
+
+ lattice = torch.as_tensor(
+ atoms.lattice_mat, dtype=torch.get_default_dtype()
+ )
+ label = torch.as_tensor(d[target], dtype=torch.get_default_dtype())
+ natoms = len(d["atoms"]["elements"])
+ if classification:
+ label = label.long()
+ if "extra_features" in d:
+ _attach_node_payload(
+ g,
+ "extra_features",
+ np.asarray(d["extra_features"]),
+ natoms,
+ )
+ if target_atomwise:
+ g.ndata[target_atomwise] = torch.as_tensor(
+ np.asarray(d[target_atomwise]),
+ dtype=torch.get_default_dtype(),
+ )
+ if target_grad:
+ arr = np.asarray(d[target_grad])
+ g.ndata[target_grad] = torch.as_tensor(
+ arr, dtype=torch.get_default_dtype()
+ )
+ if target_stress:
+ stress = np.asarray(d[target_stress])
+ _attach_node_payload(g, target_stress, stress, natoms)
+ if target_additional_output:
+ _attach_node_payload(
+ g,
+ target_additional_output,
+ np.asarray(d[target_additional_output]),
+ natoms,
+ )
+
+ if line_graph:
+ txn.put(f"{idx}".encode(), pk.dumps((g, lg, lattice, label)))
+ else:
+ txn.put(f"{idx}".encode(), pk.dumps((g, lattice, label)))
+
+ # Drop this entry's full atomic-structure dict now that it's
+ # serialized into LMDB. The caller's list survives but its
+ # contents are released, freeing ~10-30 KB per entry. For
+ # 1.5M-entry datasets this is the difference between sitting
+ # at 24+ GB during the loop and staying at 3-4 GB.
+ dataset[idx] = None
+
+ # Flush dirty pages to disk every LMDB_COMMIT_EVERY entries.
+ if (idx + 1) % LMDB_COMMIT_EVERY == 0:
+ txn.commit()
+ txn = env.begin(write=True)
+ if (idx + 1) % (LMDB_COMMIT_EVERY * 5) == 0:
+ print(
+ f"[MEM] {tmp_name}: after {idx + 1} entries: "
+ f"{_mem_gb():.2f} GB"
+ )
+ finally:
+ # Commit whatever is left in the final partial batch.
+ txn.commit()
+
+ env.close()
+
+ print(f"[MEM] {tmp_name}: after write loop: {_mem_gb():.2f} GB")
+
+ ds = PureTorchLMDBDataset(
+ lmdb_path=tmp_name, line_graph=line_graph, ids=ids
+ )
+
+ print(
+ f"[MEM] {tmp_name}: after PureTorchLMDBDataset wrap: "
+ f"{_mem_gb():.2f} GB"
+ )
+
+ return ds
diff --git a/alignn/scripts/bench_dgl_vs_pure.py b/alignn/scripts/bench_dgl_vs_pure.py
new file mode 100644
index 0000000..900e607
--- /dev/null
+++ b/alignn/scripts/bench_dgl_vs_pure.py
@@ -0,0 +1,345 @@
+"""Scaling benchmark: DGL ALIGNN vs pure-torch ALIGNN.
+
+Builds Cu-FCC supercells of increasing size, runs one forward + backward
+(energy + forces + stress) through each backend, and reports time and
+peak memory.
+
+Both backends receive the *same* graph topology: we build a single DGL
+graph via matscipy (``neighbor_strategy="fast_graph"``) and feed it to
+``ALIGNNAtomWise`` directly; the pure-torch model gets the same graph
+converted to ``TorchGraph`` at the model boundary. So the numbers
+below isolate **model-level** cost (message passing, autograd,
+readout) — not the neighbor-list step.
+
+Usage
+-----
+ python alignn/scripts/bench_dgl_vs_pure.py
+ python alignn/scripts/bench_dgl_vs_pure.py \\
+ --sizes 1,2,3,4,5 --cutoff 4.0 --hidden 64 \\
+ --alignn-layers 1 --gcn-layers 1 --repeat 3
+
+Writes a printed table; with ``--output foo`` also saves ``foo.npz``
+and, if matplotlib is available, ``foo.png``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import time
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+import torch
+from ase.build.supercells import make_supercell
+from jarvis.io.vasp.inputs import Poscar
+
+from alignn.graphs import Graph
+from alignn.models.alignn_atomwise import (
+ ALIGNNAtomWise,
+ ALIGNNAtomWiseConfig,
+)
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure,
+ ALIGNNAtomWisePureConfig,
+)
+from alignn.torch_graph_builder import torchgraph_from_dgl
+
+
+CU_POSCAR = """Cu
+1.0
+3.6 0.0 0.0
+0.0 3.6 0.0
+0.0 0.0 3.6
+Cu
+4
+direct
+0.0 0.0 0.0
+0.0 0.5 0.5
+0.5 0.0 0.5
+0.5 0.5 0.0
+"""
+
+
+def _build_models(args, device) -> "tuple[torch.nn.Module, torch.nn.Module]":
+ """Construct both backends from the same hyperparameters and weights."""
+ common = dict(
+ alignn_layers=args.alignn_layers,
+ gcn_layers=args.gcn_layers,
+ atom_input_features=92,
+ hidden_features=args.hidden,
+ output_features=1,
+ calculate_gradient=True,
+ graphwise_weight=0.85,
+ gradwise_weight=0.05,
+ stresswise_weight=0.05,
+ atomwise_weight=0.0,
+ add_reverse_forces=True,
+ lg_on_fly=True,
+ use_penalty=True,
+ )
+ torch.manual_seed(0)
+ m_dgl = ALIGNNAtomWise(
+ ALIGNNAtomWiseConfig(name="alignn_atomwise", **common)
+ ).to(device)
+ torch.manual_seed(0)
+ m_pure = ALIGNNAtomWisePure(
+ ALIGNNAtomWisePureConfig(name="alignn_atomwise_pure", **common)
+ ).to(device)
+ # Copy DGL weights into pure so any difference is purely mechanical.
+ sd_dgl = m_dgl.state_dict()
+ sd_pure = m_pure.state_dict()
+ for k in sd_pure:
+ if k in sd_dgl and sd_dgl[k].shape == sd_pure[k].shape:
+ sd_pure[k] = sd_dgl[k].clone()
+ m_pure.load_state_dict(sd_pure, strict=False)
+ m_dgl.train()
+ m_pure.train()
+ return m_dgl, m_pure
+
+
+def _build_graph(atoms, cutoff: float, max_neighbors: int):
+ g, lg = Graph.atom_dgl_multigraph(
+ atoms,
+ neighbor_strategy="fast_graph",
+ cutoff=cutoff,
+ max_neighbors=max_neighbors,
+ atom_features="cgcnn",
+ )
+ return g, lg
+
+
+def _run_once(
+ model,
+ graphs,
+ lat,
+ device: torch.device,
+ is_pure: bool,
+):
+ """One forward + backward. Returns elapsed seconds and peak memory (bytes)."""
+ g, lg = graphs
+ if is_pure:
+ g_in, lg_in = torchgraph_from_dgl(g), torchgraph_from_dgl(lg)
+ else:
+ g_in, lg_in = g, lg
+ lat_in = lat
+
+ # Sync + reset peak memory.
+ if device.type == "cuda":
+ torch.cuda.synchronize(device)
+ torch.cuda.reset_peak_memory_stats(device)
+
+ t0 = time.perf_counter()
+ out = model((g_in, lg_in, lat_in))
+ loss = (
+ out["out"].sum()
+ + out["grad"].pow(2).sum()
+ + out["stresses"].pow(2).sum()
+ )
+ loss.backward()
+ if device.type == "cuda":
+ torch.cuda.synchronize(device)
+ dt = time.perf_counter() - t0
+
+ peak = (
+ torch.cuda.max_memory_allocated(device)
+ if device.type == "cuda"
+ else 0
+ )
+ # Free autograd graph / grads before next iter.
+ model.zero_grad(set_to_none=True)
+ del out, loss
+ gc.collect()
+ if device.type == "cuda":
+ torch.cuda.empty_cache()
+ return dt, int(peak)
+
+
+def _bench_one_size(
+ model,
+ graphs,
+ lat,
+ device,
+ is_pure: bool,
+ warmup: int,
+ repeat: int,
+):
+ """Warm up then time ``repeat`` iterations; return median time + peak mem."""
+ for _ in range(warmup):
+ _run_once(model, graphs, lat, device, is_pure)
+ times, peaks = [], []
+ for _ in range(repeat):
+ t, p = _run_once(model, graphs, lat, device, is_pure)
+ times.append(t)
+ peaks.append(p)
+ return float(np.median(times)), max(peaks)
+
+
+def run(args) -> Dict[str, np.ndarray]:
+ device = torch.device(args.device)
+ print(f"device: {device}")
+ if device.type == "cuda":
+ print(f" {torch.cuda.get_device_name(device)}")
+
+ m_dgl, m_pure = _build_models(args, device)
+ n_params = sum(p.numel() for p in m_dgl.parameters())
+ print(f"model parameters: {n_params}")
+
+ base = Poscar.from_string(CU_POSCAR).atoms.ase_converter()
+ sizes = [int(s) for s in args.sizes.split(",")]
+
+ rows: List[dict] = []
+ print()
+ hdr = (
+ f"{'N':>3} {'atoms':>6} {'edges':>8} "
+ f"{'DGL ms':>10} {'pure ms':>10} {'speedup':>8} "
+ f"{'DGL MB':>10} {'pure MB':>10} {'mem ratio':>10}"
+ )
+ print(hdr)
+ print("-" * len(hdr))
+
+ for n in sizes:
+ sc = make_supercell(base, [[n, 0, 0], [0, n, 0], [0, 0, n]])
+ # Convert to jarvis Atoms for the neighbor-list pipeline.
+ from jarvis.core.atoms import Atoms as JAtoms
+
+ jatoms = JAtoms(
+ lattice_mat=np.asarray(sc.cell),
+ coords=np.asarray(sc.positions),
+ elements=list(sc.get_chemical_symbols()),
+ cartesian=True,
+ )
+ try:
+ g, lg = _build_graph(jatoms, args.cutoff, args.max_neighbors)
+ except Exception as exc:
+ print(f" graph build failed at N={n}: {exc}")
+ break
+ lat = torch.as_tensor(
+ np.asarray(jatoms.lattice_mat), dtype=torch.float32, device=device
+ )
+ g = g.to(device)
+ lg = lg.to(device)
+
+ try:
+ t_dgl, m_dgl_peak = _bench_one_size(
+ m_dgl, (g, lg), lat, device, False, args.warmup, args.repeat
+ )
+ except (RuntimeError, torch.cuda.OutOfMemoryError) as exc:
+ print(f" DGL OOM/error at N={n}: {type(exc).__name__}")
+ break
+ try:
+ t_pure, m_pure_peak = _bench_one_size(
+ m_pure, (g, lg), lat, device, True, args.warmup, args.repeat
+ )
+ except (RuntimeError, torch.cuda.OutOfMemoryError) as exc:
+ print(f" pure OOM/error at N={n}: {type(exc).__name__}")
+ break
+
+ n_atoms = int(sc.positions.shape[0])
+ n_edges = int(g.num_edges())
+ speedup = t_dgl / t_pure if t_pure > 0 else float("nan")
+ mem_ratio = (
+ m_pure_peak / m_dgl_peak if m_dgl_peak > 0 else float("nan")
+ )
+ print(
+ f"{n:>3} {n_atoms:>6} {n_edges:>8} "
+ f"{t_dgl * 1000:>10.2f} {t_pure * 1000:>10.2f} "
+ f"{speedup:>8.2f}x "
+ f"{m_dgl_peak / 1024**2:>10.1f} "
+ f"{m_pure_peak / 1024**2:>10.1f} "
+ f"{mem_ratio:>10.2f}x"
+ )
+ rows.append(
+ dict(
+ n=n,
+ atoms=n_atoms,
+ edges=n_edges,
+ dgl_time_s=t_dgl,
+ pure_time_s=t_pure,
+ dgl_peak_bytes=m_dgl_peak,
+ pure_peak_bytes=m_pure_peak,
+ )
+ )
+
+ result = {k: np.array([r[k] for r in rows]) for k in rows[0]} if rows else {}
+
+ if args.output and rows:
+ out = Path(args.output)
+ out.parent.mkdir(parents=True, exist_ok=True)
+ np.savez(str(out.with_suffix(".npz")), **result)
+ print(f"\nSaved data to {out.with_suffix('.npz')}")
+ try:
+ import matplotlib
+
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+
+ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.5))
+ n_arr = result["atoms"]
+ ax1.plot(
+ n_arr, result["dgl_time_s"] * 1000, "-o", label="DGL"
+ )
+ ax1.plot(
+ n_arr, result["pure_time_s"] * 1000, "-s", label="pure-torch"
+ )
+ ax1.set(
+ xlabel="# atoms", ylabel="time / iter (ms)",
+ title="Forward + backward (E+F+σ)",
+ )
+ ax1.set_xscale("log")
+ ax1.set_yscale("log")
+ ax1.grid(True, which="both", alpha=0.3)
+ ax1.legend()
+
+ ax2.plot(n_arr, result["dgl_peak_bytes"] / 1024**2, "-o", label="DGL")
+ ax2.plot(
+ n_arr, result["pure_peak_bytes"] / 1024**2, "-s",
+ label="pure-torch",
+ )
+ ax2.set(
+ xlabel="# atoms", ylabel="peak memory (MB)",
+ title="Peak GPU memory",
+ )
+ ax2.set_xscale("log")
+ ax2.set_yscale("log")
+ ax2.grid(True, which="both", alpha=0.3)
+ ax2.legend()
+
+ plt.tight_layout()
+ png = out.with_suffix(".png")
+ plt.savefig(str(png), dpi=160)
+ print(f"Saved plot to {png}")
+ except ImportError:
+ pass
+ return result
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser()
+ ap.add_argument(
+ "--sizes", default="1,2,3,4,5",
+ help="Comma-separated supercell factors (N -> 4*N^3 atoms).",
+ )
+ ap.add_argument("--cutoff", type=float, default=4.0)
+ ap.add_argument("--max-neighbors", type=int, default=12)
+ ap.add_argument("--hidden", type=int, default=64)
+ ap.add_argument("--alignn-layers", type=int, default=1)
+ ap.add_argument("--gcn-layers", type=int, default=1)
+ ap.add_argument("--warmup", type=int, default=2)
+ ap.add_argument("--repeat", type=int, default=3)
+ ap.add_argument(
+ "--device",
+ default="cuda" if torch.cuda.is_available() else "cpu",
+ )
+ ap.add_argument(
+ "--output",
+ default=None,
+ help="Path stem; writes .npz and (if matplotlib) .png.",
+ )
+ args = ap.parse_args()
+ run(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/export_torchscript.py b/alignn/scripts/export_torchscript.py
new file mode 100644
index 0000000..d01f4a3
--- /dev/null
+++ b/alignn/scripts/export_torchscript.py
@@ -0,0 +1,194 @@
+"""Export a trained pure-torch ALIGNN model to TorchScript.
+
+The saved ``.pt`` file is self-contained: it holds the model weights,
+architecture code (scripted), and an entry-point ``forward(...)``
+that takes plain tensors (positions, lattice, atom_features, src, dst,
+shift) and returns ``{"energy", "forces", ["stress"]}``.
+
+This is what a LAMMPS ``pair_alignn`` plugin (or any libtorch C++ host)
+would call each MD step — the only runtime dependency is libtorch.
+
+Usage
+-----
+ 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
+
+Design notes
+------------
+Only the ``forward_tensors`` path is scripted (a thin Wrapper module
+exposes it as ``forward``). The dataclass-based training forward is
+left alone. Classification and extra-features branches are not
+exported — they're training-only features.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Dict
+
+import numpy as np
+import torch
+from torch import nn
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure,
+ ALIGNNAtomWisePureConfig,
+)
+
+
+class _ScriptWrapper(nn.Module):
+ """LAMMPS-style entry point for libtorch.
+
+ Takes atomic numbers (not pre-computed features) — the feature
+ table is baked into the inner model via ``register_species_table``
+ so the scripted ``.pt`` is fully self-contained.
+ """
+
+ def __init__(self, inner: ALIGNNAtomWisePure):
+ super().__init__()
+ self.inner = inner
+
+ def forward(
+ self,
+ positions: torch.Tensor, # (N, 3), requires_grad=True for forces
+ lattice: torch.Tensor, # (3, 3)
+ atomic_numbers: torch.Tensor, # (N,) long
+ src: torch.Tensor, # (E,) long
+ dst: torch.Tensor, # (E,) long
+ shift: torch.Tensor, # (E, 3) integer cell offsets
+ compute_stress: bool = False,
+ ) -> Dict[str, torch.Tensor]:
+ return self.inner.forward_tensors_z(
+ positions, lattice, atomic_numbers, src, dst, shift, compute_stress
+ )
+
+
+def build_model(
+ checkpoint_path: Path,
+ config_path: Path,
+ atom_features: str = "cgcnn",
+) -> ALIGNNAtomWisePure:
+ config = json.load(open(config_path))
+ mcfg = config["model"] if "model" in config else config
+ if mcfg.get("name") != "alignn_atomwise_pure":
+ # Accept a DGL-trained checkpoint's config by overriding the name.
+ # Architecture is identical between alignn_atomwise and
+ # alignn_atomwise_pure, so weights are state-dict compatible.
+ mcfg = dict(mcfg, name="alignn_atomwise_pure")
+ model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(**mcfg))
+ state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+ if isinstance(state, dict) and "state_dict" in state:
+ state = state["state_dict"]
+ # Drop any previously-saved species table — it'll be rebuilt fresh.
+ state = {k: v for k, v in state.items() if k != "_species_table"}
+ model.load_state_dict(state, strict=False)
+ model.register_species_table(atom_features=atom_features)
+ model.eval()
+ return model
+
+
+def _smoke_test(scripted: torch.jit.ScriptModule) -> None:
+ """Run a tiny Cu-FCC box through the scripted model (atomic-number input)."""
+ from jarvis.io.vasp.inputs import Poscar
+ from alignn.graphs import Graph
+ from alignn.torch_graph_builder import torchgraph_from_dgl
+
+ atoms = Poscar.from_string(
+ """Cu
+1.0
+3.6 0.0 0.0
+0.0 3.6 0.0
+0.0 0.0 3.6
+Cu
+4
+direct
+0.0 0.0 0.0
+0.0 0.5 0.5
+0.5 0.0 0.5
+0.5 0.5 0.0
+"""
+ ).atoms
+ g, _ = Graph.atom_dgl_multigraph(
+ atoms,
+ neighbor_strategy="fast_graph",
+ cutoff=4.0,
+ max_neighbors=12,
+ atom_features="cgcnn",
+ )
+ tg = torchgraph_from_dgl(g)
+
+ pos = torch.as_tensor(np.asarray(atoms.cart_coords), dtype=torch.float32)
+ pos.requires_grad_(True)
+ lat = torch.as_tensor(np.asarray(atoms.lattice_mat), dtype=torch.float32)
+ # LAMMPS-style input: atomic numbers, not feature vectors.
+ z_tensor = torch.as_tensor(atoms.atomic_numbers, dtype=torch.long)
+ shift = tg.edata["images"].float()
+
+ out = scripted(pos, lat, z_tensor, tg.src, tg.dst, shift, True)
+ print("Scripted smoke test (atomic-number input):")
+ print(f" energy: {out['energy'].item():.6f}")
+ print(
+ f" forces: shape={tuple(out['forces'].shape)} "
+ f"max|F|={out['forces'].abs().max().item():.4e} "
+ f"sum(F)={out['forces'].detach().sum(0).tolist()}"
+ )
+ print(
+ f" stress: shape={tuple(out['stress'].shape)} "
+ f"max|σ|={out['stress'].abs().max().item():.4e}"
+ )
+
+
+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("--output", required=True, type=Path)
+ ap.add_argument(
+ "--atom-features",
+ default=None,
+ help="Feature scheme for the atomic-number lookup table "
+ "(cgcnn/atomic_number/basic/cfid). Defaults to the value "
+ "found in the training config, or 'cgcnn'.",
+ )
+ ap.add_argument(
+ "--test",
+ action="store_true",
+ help="Run the scripted model on a tiny Cu-FCC box after saving.",
+ )
+ args = ap.parse_args()
+
+ # Resolve atom_features: CLI > training config > 'cgcnn'.
+ cfg = json.load(open(args.config))
+ atom_features = args.atom_features or cfg.get("atom_features", "cgcnn")
+ print(f"Loading model from {args.checkpoint} + {args.config}")
+ print(f" atom_features (species lookup) = {atom_features!r}")
+ model = build_model(
+ args.checkpoint, args.config, atom_features=atom_features
+ )
+
+ print("Wrapping + scripting...")
+ wrapper = _ScriptWrapper(model).eval()
+ scripted = torch.jit.script(wrapper)
+
+ scripted.save(str(args.output))
+ print(f"Saved scripted model -> {args.output}")
+ print(
+ " Input signature: (positions[N,3], lattice[3,3], "
+ "atomic_numbers[N], src[E], dst[E], shift[E,3], compute_stress: bool)"
+ )
+ print(' Output: {"energy": scalar, "forces": [N,3], ["stress": [3,3]]}')
+
+ if args.test:
+ scripted = torch.jit.load(str(args.output), map_location="cpu")
+ _smoke_test(scripted)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/parity_dgl_vs_pure.py b/alignn/scripts/parity_dgl_vs_pure.py
new file mode 100644
index 0000000..23da677
--- /dev/null
+++ b/alignn/scripts/parity_dgl_vs_pure.py
@@ -0,0 +1,229 @@
+"""Parity sweep: DGL ALIGNN vs pure-torch ALIGNN across the doc examples.
+
+Runs each scenario from ``docs/training/*.md`` twice -- once with the
+DGL-backed ``alignn_atomwise`` model and once with the DGL-free
+``alignn_atomwise_pure`` model -- and compares final train/val/test
+losses side-by-side. Graph construction and LMDB storage flip with the
+model: the DGL run uses ``alignn/lmdb_dataset.py``, the pure run uses
+``alignn/pure_lmdb_dataset.py`` (both selected automatically by
+``train_alignn.py`` based on ``config.model.name``).
+
+Scenarios covered (match the docs):
+ 1. Single-output regression (alignn/examples/sample_data)
+ 2. Classification (sample_data + classification_threshold)
+ 3. Multi-output regression (alignn/examples/sample_data_multi_prop)
+ 4. Force field (alignn/examples/sample_data_ff)
+
+Usage
+-----
+ python alignn/scripts/parity_dgl_vs_pure.py # all scenarios
+ python alignn/scripts/parity_dgl_vs_pure.py --only ff # one scenario
+ python alignn/scripts/parity_dgl_vs_pure.py --epochs 1 # quick smoke
+
+Exits 0 on success; prints a table of metrics and |Δ| between backends.
+
+Notes
+-----
+* Graph caches ({train,val,test}_data) are deleted before each
+ run so the LMDB store always matches the current model backend.
+* Residual drift of ~1e-2 on train/val is expected from scatter_add
+ nondeterminism on CUDA (accumulation order differs from DGL's kernels
+ and run-to-run). Set ``torch.use_deterministic_algorithms(True)`` to
+ eliminate at the cost of speed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+import subprocess
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+# Repo root = two parents up from this file (alignn/scripts/ -> repo root).
+REPO = Path(__file__).resolve().parents[2]
+FLOAT = r"[-+]?\d+\.\d+(?:[eE][-+]?\d+)?"
+
+
+SCENARIOS: List[Dict] = [
+ {
+ "key": "single",
+ "name": "single-output regression (sample_data)",
+ "base_config": "alignn/examples/sample_data/config_example.json",
+ "root_dir": "alignn/examples/sample_data",
+ "extra": None,
+ },
+ {
+ "key": "classification",
+ "name": "classification (sample_data, threshold 0.01)",
+ "base_config": "alignn/examples/sample_data/config_example.json",
+ "root_dir": "alignn/examples/sample_data",
+ "extra": {"classification_threshold": 0.01},
+ },
+ {
+ "key": "multi",
+ "name": "multi-output regression (sample_data_multi_prop)",
+ "base_config": "alignn/examples/sample_data/config_example.json",
+ "root_dir": "alignn/examples/sample_data_multi_prop",
+ "extra": None,
+ },
+ {
+ "key": "ff",
+ "name": "force field (sample_data_ff)",
+ "base_config": "alignn/examples/sample_data_ff/config_example_atomwise.json",
+ "root_dir": "alignn/examples/sample_data_ff",
+ "extra": None,
+ },
+]
+
+
+def _patch_cfg(
+ base_path: Path,
+ out_path: Path,
+ *,
+ pure: bool,
+ extra: Optional[dict] = None,
+ epochs: Optional[int] = None,
+) -> None:
+ cfg = json.load(open(base_path))
+ cfg["use_lmdb"] = True
+ cfg["read_existing"] = False
+ cfg["num_workers"] = 0
+ if pure:
+ cfg["model"]["name"] = "alignn_atomwise_pure"
+ if extra:
+ cfg.update(extra)
+ if epochs is not None:
+ cfg["epochs"] = int(epochs)
+ json.dump(cfg, open(out_path, "w"))
+
+
+def _run(cmd: List[str], log_path: Path) -> "tuple[int, float]":
+ t0 = time.time()
+ with open(log_path, "w") as f:
+ p = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT)
+ return p.returncode, time.time() - t0
+
+
+def _extract(log_path: Path) -> Dict:
+ txt = log_path.read_text()
+ train = re.findall(rf"Train Loss:.*\n.*?(\d+)\s+({FLOAT})", txt)
+ val = re.findall(rf"Val Loss:.*\n.*?(\d+)\s+({FLOAT})", txt)
+ test = re.search(rf"TestLoss\s+\d+\s+({FLOAT})", txt)
+ return {
+ "train_last": float(train[-1][1]) if train else None,
+ "val_last": float(val[-1][1]) if val else None,
+ "test": float(test.group(1)) if test else None,
+ "final_epoch": int(train[-1][0]) if train else None,
+ }
+
+
+def _wipe_graph_caches(filename_stem: str) -> None:
+ """Remove any stale LMDB cache dirs tied to the cfg's ``filename`` stem."""
+ for suffix in ("train_data", "val_data", "test_data"):
+ p = REPO / f"{filename_stem}{suffix}"
+ if p.exists():
+ shutil.rmtree(p)
+
+
+def run_scenario(s: Dict, epochs: Optional[int] = None) -> Dict[str, Dict]:
+ base_cfg = REPO / s["base_config"]
+ stem = json.load(open(base_cfg)).get("filename", "")
+ results: Dict[str, Dict] = {}
+ for mode in ("dgl", "pure"):
+ _wipe_graph_caches(stem)
+ patched_cfg = Path(f"/tmp/parity_{s['key']}_{mode}.json")
+ _patch_cfg(
+ base_cfg,
+ patched_cfg,
+ pure=(mode == "pure"),
+ extra=s["extra"],
+ epochs=epochs,
+ )
+ out_dir = Path(f"/tmp/parity_out_{s['key']}_{mode}")
+ shutil.rmtree(out_dir, ignore_errors=True)
+ log = Path(f"/tmp/parity_{s['key']}_{mode}.log")
+ rc, dt = _run(
+ [
+ "python",
+ "alignn/train_alignn.py",
+ "--root_dir",
+ s["root_dir"],
+ "--config",
+ str(patched_cfg),
+ "--output_dir",
+ str(out_dir),
+ ],
+ log,
+ )
+ res = _extract(log) if rc == 0 else {
+ "train_last": None,
+ "val_last": None,
+ "test": None,
+ "final_epoch": None,
+ }
+ res.update(rc=rc, time_s=dt, log=str(log))
+ results[mode] = res
+ status = "OK" if rc == 0 else f"FAIL rc={rc}"
+ print(
+ f" [{mode:4s}] {status:10s} "
+ f"train={res['train_last']} val={res['val_last']} "
+ f"test={res['test']} ({dt:.1f}s)"
+ )
+ return results
+
+
+def _fmt(x: Optional[float], width: int = 14) -> str:
+ if x is None:
+ return f"{'n/a':>{width}s}"
+ return f"{x:>{width}.4f}"
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser()
+ ap.add_argument(
+ "--only",
+ choices=[s["key"] for s in SCENARIOS],
+ help="Run just one scenario.",
+ )
+ ap.add_argument(
+ "--epochs",
+ type=int,
+ default=None,
+ help="Override config.epochs (default: use JSON value).",
+ )
+ args = ap.parse_args()
+
+ os.chdir(REPO)
+ scenarios = (
+ [s for s in SCENARIOS if s["key"] == args.only]
+ if args.only
+ else SCENARIOS
+ )
+
+ all_results: Dict[str, Dict[str, Dict]] = {}
+ for s in scenarios:
+ print(f"\n=== {s['name']} ===")
+ all_results[s["name"]] = run_scenario(s, epochs=args.epochs)
+
+ print("\n=== Parity summary ===")
+ hdr = f"{'scenario':52s} {'metric':8s} {'DGL':>14s} {'pure':>14s} {'|Δ|':>12s}"
+ print(hdr)
+ print("-" * len(hdr))
+ for name, d in all_results.items():
+ for metric in ("train_last", "val_last", "test"):
+ a = d.get("dgl", {}).get(metric)
+ b = d.get("pure", {}).get(metric)
+ delta = abs(a - b) if (a is not None and b is not None) else None
+ print(
+ f"{name:52s} {metric:8s} "
+ f"{_fmt(a)} {_fmt(b)} {_fmt(delta, width=12)}"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/build_lammps_alignn.sh b/alignn/scripts/torch/build_lammps_alignn.sh
new file mode 100755
index 0000000..f0f4753
--- /dev/null
+++ b/alignn/scripts/torch/build_lammps_alignn.sh
@@ -0,0 +1,267 @@
+#!/bin/bash
+# ---------------------------------------------------------------------------
+# build_lammps_alignn.sh
+#
+# One-shot build of LAMMPS stable_29Aug2024_update1 with the native
+# pair_alignn style linked against the conda-env libtorch.
+#
+# Does:
+# 1. Discovers PyTorch/CUDA/ABI in the active conda env
+# 2. (Optional) installs matching cuda-toolkit + MKL if missing
+# 3. Clones LAMMPS to $LAMMPS_DIR (default ~/lammps-alignn)
+# 4. Copies pair_alignn.{cpp,h} from this repo into LAMMPS src/
+# 5. Appends find_package(Torch) + target_link_libraries to main CMakeLists.txt
+# 6. Configures + builds with Ninja
+# 7. Installs the lammps Python module into the active env
+# 8. Overwrites any pre-existing liblammps.so.0 in the Python prefix with the
+# pair_alignn build (fixes the standalone `lmp` binary's RPATH picking up
+# a stale .so).
+#
+# Works in a conda env (uses mamba for CUDA/MKL) OR in system Python like
+# Colab (falls back to `pip install nvidia-cuda-*-cu12 mkl mkl-devel mkl-include`).
+# 9. Verifies "alignn" appears in both Python module and the standalone binary
+#
+# Env vars:
+# LAMMPS_DIR (default: $HOME/lammps-alignn)
+# LAMMPS_TAG (default: stable_29Aug2024_update1)
+# SKIP_CUDA_MKL set to 1 to skip the mamba install step
+# JOBS (default: $(nproc))
+# ---------------------------------------------------------------------------
+set -euo pipefail
+
+LAMMPS_TAG="${LAMMPS_TAG:-stable_29Aug2024_update1}"
+LAMMPS_DIR="${LAMMPS_DIR:-$HOME/lammps-alignn}"
+JOBS="${JOBS:-$(nproc)}"
+
+# Locate the alignn repo root from *this* script's path (so it's portable).
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PAIR_SRC="$SCRIPT_DIR/pair_alignn"
+ALIGNN_REPO="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+
+echo "── paths ────────────────────────────────"
+echo "script dir : $SCRIPT_DIR"
+echo "alignn repo : $ALIGNN_REPO"
+echo "pair sources : $PAIR_SRC"
+echo "LAMMPS target: $LAMMPS_DIR (tag $LAMMPS_TAG)"
+echo
+
+# ── detect Python install prefix (conda env OR system python, e.g. Colab) ──
+PY_PREFIX=$(python -c "import sys; print(sys.prefix)")
+PREFIX="${CONDA_PREFIX:-$PY_PREFIX}"
+HAS_MAMBA=0
+command -v mamba >/dev/null 2>&1 && HAS_MAMBA=1
+command -v cmake >/dev/null || {
+ echo "ERROR: cmake not found. Install it (apt / conda / pip) and rerun." >&2
+ exit 1
+}
+command -v ninja >/dev/null || pip install -q ninja
+
+# ── discover torch + CUDA ABI ────────────────────────────────────────────
+TORCH_DIR=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))")
+TORCH_CMAKE="$TORCH_DIR/share/cmake/Torch"
+TORCH_CXX11_ABI=$(python -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))")
+TORCH_CUDA=$(python -c "import torch; print(torch.version.cuda or 'cpu')")
+
+echo "── environment ──────────────────────────"
+echo "python = $(which python)"
+echo "torch = $(python -c 'import torch; print(torch.__version__)')"
+echo "torch CUDA = $TORCH_CUDA"
+echo "torch CXX11 ABI = $TORCH_CXX11_ABI"
+echo "cmake = $(cmake --version | head -1)"
+echo "prefix = $PREFIX"
+echo "has mamba = $HAS_MAMBA"
+echo
+
+# ── ensure CUDA toolkit + MKL match libtorch ─────────────────────────────
+if [ "${SKIP_CUDA_MKL:-0}" != "1" ] && [ "$TORCH_CUDA" != "cpu" ]; then
+ # CUDA toolkit check: nvcc matches torch CUDA?
+ NVCC_VER="none"
+ if command -v nvcc >/dev/null; then
+ NVCC_VER=$(nvcc --version | awk '/release/{gsub(",","",$5); print $5}')
+ fi
+ NEED_CUDA=0
+ [ "$NVCC_VER" != "$TORCH_CUDA" ] && NEED_CUDA=1
+
+ # MKL headers check — look in common places
+ MKL_H=""
+ for d in "$PREFIX/include" "$PY_PREFIX/include" /usr/include /usr/local/include; do
+ [ -f "$d/mkl.h" ] && MKL_H="$d/mkl.h" && break
+ done
+ NEED_MKL=0
+ [ -z "$MKL_H" ] && NEED_MKL=1
+
+ if [ $HAS_MAMBA -eq 1 ]; then
+ if [ $NEED_CUDA -eq 1 ] || [ $NEED_MKL -eq 1 ]; then
+ echo "── installing cuda-toolkit=$TORCH_CUDA + mkl-devel via mamba ──"
+ mamba install -y -c nvidia -c conda-forge \
+ "cuda-toolkit=$TORCH_CUDA" mkl-devel mkl-include
+ fi
+ else
+ # Pip fallback (Colab / system python). Grab CUDA + MKL from pip wheels.
+ PIP_ARGS=()
+ if [ $NEED_CUDA -eq 1 ]; then
+ cu_short=$(echo "$TORCH_CUDA" | tr -d '.')
+ echo "── installing CUDA ${TORCH_CUDA} dev headers via pip (cu${cu_short}) ──"
+ PIP_ARGS+=(\
+ "nvidia-cuda-nvcc-cu${cu_short%??}==${TORCH_CUDA}.*" \
+ "nvidia-cuda-runtime-cu${cu_short%??}==${TORCH_CUDA}.*" \
+ "nvidia-cuda-cupti-cu${cu_short%??}==${TORCH_CUDA}.*" \
+ ) || true
+ fi
+ if [ $NEED_MKL -eq 1 ]; then
+ echo "── installing mkl via pip ──"
+ PIP_ARGS+=(mkl mkl-devel mkl-include)
+ fi
+ [ ${#PIP_ARGS[@]} -gt 0 ] && pip install -q "${PIP_ARGS[@]}" || true
+
+ # Re-locate MKL headers after install
+ for d in "$PREFIX/include" "$PY_PREFIX/include" /usr/include /usr/local/include; do
+ [ -f "$d/mkl.h" ] && MKL_H="$d/mkl.h" && break
+ done
+
+ # Colab ships nvcc in /usr/local/cuda; export for CMake's detection
+ if [ -x /usr/local/cuda/bin/nvcc ] && ! command -v nvcc >/dev/null; then
+ export PATH="/usr/local/cuda/bin:$PATH"
+ export CUDA_HOME=/usr/local/cuda
+ fi
+ fi
+fi
+
+# ── clone LAMMPS ─────────────────────────────────────────────────────────
+if [ ! -d "$LAMMPS_DIR" ]; then
+ echo "[1/5] cloning LAMMPS $LAMMPS_TAG..."
+ git clone --depth=1 --branch "$LAMMPS_TAG" \
+ https://github.com/lammps/lammps.git "$LAMMPS_DIR"
+else
+ echo "[1/5] LAMMPS dir exists at $LAMMPS_DIR, reusing"
+fi
+
+# ── drop pair_alignn into src/ (no USER- package machinery) ──────────────
+echo "[2/5] installing pair_alignn.{cpp,h} into LAMMPS src/..."
+cp "$PAIR_SRC/pair_alignn.cpp" "$LAMMPS_DIR/src/pair_alignn.cpp"
+cp "$PAIR_SRC/pair_alignn.h" "$LAMMPS_DIR/src/pair_alignn.h"
+
+# ── append unconditional libtorch link to the main CMakeLists.txt ────────
+MAIN_CMAKE="$LAMMPS_DIR/cmake/CMakeLists.txt"
+HOOK_MARKER="# ALIGNN-FF libtorch hook"
+if ! grep -q "$HOOK_MARKER" "$MAIN_CMAKE"; then
+ cat >> "$MAIN_CMAKE" </dev/null || echo "") \
+ 2>&1 | tail -8
+
+echo
+echo "[4/5] compiling (this takes 5–15 min)..."
+cmake --build . -j"$JOBS" 2>&1 | tail -8
+
+# ── install the lammps Python module ─────────────────────────────────────
+echo "[5/5] installing lammps Python module into $PREFIX..."
+install_ok=1
+if cmake --build . --target install-python 2>&1 | tail -5; then
+ # `install-python` may report success but actually fail silently on Colab
+ # when its internal venv bootstrap breaks. Verify by importing.
+ python -c "import lammps" 2>/dev/null || install_ok=0
+else
+ install_ok=0
+fi
+
+if [ $install_ok -eq 0 ]; then
+ echo
+ echo " install-python target failed (common on Colab / system Python)."
+ echo " Falling back to manual copy of python/lammps → site-packages..."
+ SITE_PKG=$(python -c "import site; print(site.getsitepackages()[0])")
+ # Uninstall any stale lammps wheel first
+ pip uninstall -y lammps 2>/dev/null || true
+ rm -rf "$SITE_PKG/lammps"
+ cp -r "$LAMMPS_DIR/python/lammps" "$SITE_PKG/lammps"
+ # The Python module locates liblammps.so via rpath / LD_LIBRARY_PATH.
+ # Put the .so next to the package so ctypes.CDLL finds it.
+ cp "$BUILD_DIR/liblammps.so.0" "$SITE_PKG/lammps/liblammps.so"
+ # Also place a copy in PY_PREFIX/lib so the standalone lmp binary's RPATH works
+ mkdir -p "$PY_PREFIX/lib"
+ cp "$BUILD_DIR/liblammps.so.0" "$PY_PREFIX/lib/liblammps.so.0"
+ echo " installed to $SITE_PKG/lammps"
+fi
+
+# ── fix standalone binary's RPATH by overwriting any pre-existing liblammps ──
+for lib_dir in "$PREFIX/lib" "$PY_PREFIX/lib" /usr/lib /usr/local/lib; do
+ if [ -f "$lib_dir/liblammps.so.0" ]; then
+ echo " syncing liblammps.so.0 → $lib_dir/ (for standalone lmp binary)"
+ cp "$BUILD_DIR/liblammps.so.0" "$lib_dir/liblammps.so.0"
+ fi
+done
+
+# ── verify ───────────────────────────────────────────────────────────────
+echo
+echo "── verification ─────────────────────────"
+
+PY_OK=$(python - <<'PYEOF'
+from lammps import lammps
+l = lammps()
+ok = "alignn" in l.available_styles("pair")
+print("YES" if ok else "NO")
+l.close()
+PYEOF
+)
+if [ "$PY_OK" = "YES" ]; then
+ echo "✓ Python module has pair_alignn"
+else
+ echo "✗ Python module does NOT have pair_alignn"
+fi
+
+if "$BUILD_DIR/lmp" -help 2>&1 | grep -q alignn; then
+ echo "✓ Standalone binary has pair_alignn: $BUILD_DIR/lmp"
+else
+ echo "✗ Standalone binary does NOT have pair_alignn"
+fi
+
+echo
+echo "──────────────────────────────────────────"
+echo "Done. Try:"
+echo " export_torchscript.py --model-dir --out alignn_ff.pt"
+echo " $BUILD_DIR/lmp -in "
+echo "──────────────────────────────────────────"
diff --git a/alignn/scripts/torch/export_torchscript.py b/alignn/scripts/torch/export_torchscript.py
new file mode 100644
index 0000000..7c4f008
--- /dev/null
+++ b/alignn/scripts/torch/export_torchscript.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+"""Export a trained ALIGNN-FF model as a TorchScript `.pt` file for the
+native LAMMPS pair style (pair_alignn).
+
+Installed as the `export_torchscript.py` console script after
+`pip install -e .` (see setup.py).
+
+Usage:
+ export_torchscript.py --model-dir OutputDir --out alignn_ff.pt
+
+Or during development:
+ python alignn/scripts/torch/export_torchscript.py \\
+ --model-dir OutputDir --out alignn_ff.pt
+
+The saved module exposes:
+ forward_tensors_z(positions, lattice, atomic_numbers,
+ src, dst, shift, compute_stress) -> {"energy","forces","stress"}
+The C++ pair style calls this entry point each MD step.
+"""
+from __future__ import annotations
+import argparse, json
+import torch
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+
+
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--model-dir", required=True)
+ ap.add_argument("--out", default="alignn_ff.pt")
+ ap.add_argument("--atom-features", default="atomic_number",
+ help="must match what was used at training time")
+ ap.add_argument("--dtype", default="float32",
+ choices=["float32", "float64", "float16", "bfloat16"],
+ help="float32 is the safe default; bfloat16 halves memory "
+ "with minimal accuracy loss; float16 is risky for "
+ "MD forces (may NaN on small gradient magnitudes).")
+ args = ap.parse_args()
+
+ dtype_map = {"float32": torch.float32, "float64": torch.float64,
+ "float16": torch.float16, "bfloat16": torch.bfloat16}
+ dtype = dtype_map[args.dtype]
+ if args.dtype == "float16":
+ print("[warn] fp16 forces can NaN during long MD runs. "
+ "Consider --dtype bfloat16 instead.")
+
+ cfg_raw = json.load(open(f"{args.model_dir}/config.json"))
+ mcfg = dict(cfg_raw["model"])
+ mcfg["name"] = "alignn_atomwise_pure"
+ mcfg["calculate_gradient"] = True
+ if mcfg.get("stresswise_weight", 0) == 0:
+ mcfg["stresswise_weight"] = 0.01
+ cfg = ALIGNNAtomWisePureConfig(**mcfg)
+
+ model = ALIGNNAtomWisePure(cfg)
+ sd = torch.load(f"{args.model_dir}/best_model.pt", map_location="cpu",
+ weights_only=False)
+ if isinstance(sd, dict) and "model" in sd:
+ sd = sd["model"]
+ missing, unexpected = model.load_state_dict(sd, strict=False)
+ if missing or unexpected:
+ print(f"[warn] missing={len(missing)} unexpected={len(unexpected)}")
+
+ # Baked-in CGCNN species table so C++ only ships atomic numbers (Z).
+ model.register_species_table(atom_features=args.atom_features)
+ model = model.to(dtype).eval()
+
+ # TorchScript compile — this is the critical step. If it fails here,
+ # you'll need to fix the offending annotation in the model source.
+ scripted = torch.jit.script(model)
+
+ # Smoke test the scripted model on a minimal 2-atom example.
+ with torch.no_grad():
+ pos = torch.tensor([[0., 0., 0.], [1.35, 1.35, 1.35]], dtype=dtype)
+ pos.requires_grad_(True)
+ lat = torch.eye(3, dtype=dtype) * 5.43
+ Z = torch.tensor([14, 14], dtype=torch.long)
+ src = torch.tensor([0, 1], dtype=torch.long)
+ dst = torch.tensor([1, 0], dtype=torch.long)
+ shift = torch.zeros(2, 3, dtype=dtype)
+ # torch.det() isn't implemented for bf16/fp16 on CPU, so skip stress
+ # in the smoke test for reduced-precision exports. It works fine on CUDA.
+ needs_fp32_det = dtype in (torch.bfloat16, torch.float16)
+ out = scripted.forward_tensors_z(pos, lat, Z, src, dst, shift,
+ not needs_fp32_det)
+ print("smoke test OK:",
+ {k: (v.shape if hasattr(v, 'shape') else v) for k, v in out.items()})
+ if needs_fp32_det:
+ print(" (stress path skipped on CPU — bf16/fp16 det only works on CUDA)")
+
+ torch.jit.save(scripted, args.out)
+ print(f"wrote {args.out} ({sum(p.numel() for p in model.parameters()):,} params)")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/kappa_plotly.py b/alignn/scripts/torch/kappa_plotly.py
new file mode 100644
index 0000000..d97906b
--- /dev/null
+++ b/alignn/scripts/torch/kappa_plotly.py
@@ -0,0 +1,206 @@
+"""Lattice thermal conductivity κ(T) and 3-phonon properties from ALIGNN-FF.
+
+Uses phono3py for anharmonic (3rd-order) force constants, then solves the
+linearized Boltzmann transport equation under the relaxation-time
+approximation (RTA).
+
+WARNING: this is expensive. For Si primitive + 2x2x2 supercell (16 atoms),
+phono3py generates ~111 displacements for fc3 + ~1 for fc2 ~~ hundreds
+of force evaluations. Allow 10-30 min on a single GPU.
+
+Usage:
+ python kappa_plotly.py --model-dir OutputDir --jid JVASP-1002 \\
+ --dim 2 2 2 --mesh 11 11 11
+"""
+from __future__ import annotations
+import argparse
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+from ase import Atoms as AseAtoms
+from ase.filters import ExpCellFilter
+from ase.optimize import FIRE
+from phono3py import Phono3py
+from phonopy.structure.atoms import PhonopyAtoms
+
+from jarvis.core.atoms import Atoms, ase_to_atoms
+from jarvis.db.figshare import get_jid_data
+
+from alignn.ff.ff import AlignnAtomwiseCalculator
+
+
+def relax(atoms, calc, fmax=0.05, steps=100):
+ a = atoms.ase_converter()
+ a.calc = calc
+ FIRE(ExpCellFilter(a), logfile=None).run(fmax=fmax, steps=steps)
+ return ase_to_atoms(a)
+
+
+def forces_on_supercell(sc, calc):
+ a = AseAtoms(
+ symbols=list(sc.symbols),
+ scaled_positions=np.array(sc.scaled_positions),
+ cell=np.array(sc.cell),
+ pbc=True,
+ )
+ a.calc = calc
+ F = np.array(a.get_forces())
+ F -= F.mean(axis=0)
+ return F
+
+
+def run_phono3py(atoms, calc, dim=(2, 2, 2), dim_fc2=None, distance=0.03):
+ """Build Phono3py, compute fc3 (and fc2) via ALIGNN-FF."""
+ bulk = PhonopyAtoms(
+ symbols=atoms.elements,
+ scaled_positions=atoms.frac_coords,
+ cell=atoms.lattice_mat,
+ )
+ # fc2 supercell can be larger than fc3 supercell for a better harmonic part
+ if dim_fc2 is None:
+ dim_fc2 = dim
+ ph3 = Phono3py(
+ bulk,
+ supercell_matrix=[[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]],
+ phonon_supercell_matrix=[[dim_fc2[0], 0, 0], [0, dim_fc2[1], 0], [0, 0, dim_fc2[2]]],
+ )
+ ph3.generate_displacements(distance=distance)
+
+ # fc3 supercells
+ sc3 = ph3.supercells_with_displacements
+ n3 = len(sc3)
+ print(f"fc3 displacements: {n3}")
+ fset3 = []
+ for i, sc in enumerate(sc3):
+ if sc is None: # symmetry-reduced skip
+ fset3.append(None); continue
+ F = forces_on_supercell(sc, calc)
+ fset3.append(F)
+ if (i + 1) % 10 == 0 or i == n3 - 1:
+ print(f" fc3 {i+1}/{n3}")
+ ph3.forces = fset3
+
+ # fc2 supercells (phonon supercells)
+ sc2 = ph3.phonon_supercells_with_displacements
+ n2 = len(sc2)
+ print(f"fc2 displacements: {n2}")
+ fset2 = []
+ for i, sc in enumerate(sc2):
+ if sc is None:
+ fset2.append(None); continue
+ F = forces_on_supercell(sc, calc)
+ fset2.append(F)
+ ph3.phonon_forces = fset2
+
+ ph3.produce_fc2(symmetrize_fc2=True)
+ ph3.produce_fc3(symmetrize_fc3r=True)
+ return ph3
+
+
+def compute_kappa(ph3, mesh=(11, 11, 11), t_min=100, t_max=1000, t_step=50,
+ boundary_mfp=None):
+ """Run RTA thermal conductivity. Returns (T, kappa_xx_yy_zz_xy_xz_yz)."""
+ ph3.mesh_numbers = list(mesh)
+ ph3.init_phph_interaction()
+ temperatures = np.arange(t_min, t_max + 1, t_step)
+ ph3.run_thermal_conductivity(
+ temperatures=temperatures,
+ is_isotope=False,
+ boundary_mfp=boundary_mfp, # µm; None = infinite sample
+ write_kappa=False,
+ )
+ tc = ph3.thermal_conductivity
+ T = np.asarray(tc.temperatures)
+ kappa = np.asarray(tc.kappa) # shape (1, nT, 6): xx yy zz yz xz xy
+ return T, kappa[0]
+
+
+def plot_kappa(T, kappa, thermo=None, title=""):
+ # trace of kappa = average of diagonal components
+ kappa_avg = kappa[:, :3].mean(axis=1)
+
+ fig = make_subplots(
+ rows=1, cols=2,
+ subplot_titles=("Thermal conductivity κ(T)",
+ "Log-log view"),
+ horizontal_spacing=0.12,
+ )
+ for i, lbl in enumerate(["κ_xx", "κ_yy", "κ_zz"]):
+ fig.add_trace(go.Scatter(x=T, y=kappa[:, i], mode="lines+markers",
+ name=lbl, line=dict(width=1.5)),
+ row=1, col=1)
+ fig.add_trace(go.Scatter(x=T, y=kappa_avg, mode="lines",
+ name="κ_avg", line=dict(color="black", dash="dash")),
+ row=1, col=1)
+
+ fig.add_trace(go.Scatter(x=T, y=kappa_avg, mode="lines+markers",
+ name="κ_avg (log-log)",
+ line=dict(color="steelblue", width=1.5),
+ showlegend=False),
+ row=1, col=2)
+ fig.update_xaxes(type="log", title_text="T (K)", row=1, col=2)
+ fig.update_yaxes(type="log", title_text="κ (W/m·K)", row=1, col=2)
+
+ fig.update_xaxes(title_text="T (K)", row=1, col=1)
+ fig.update_yaxes(title_text="κ (W/m·K)", row=1, col=1)
+ fig.update_layout(template="plotly_white", title=title,
+ width=1100, height=450)
+ return fig
+
+
+def print_summary(T, kappa, name=""):
+ kappa_avg = kappa[:, :3].mean(axis=1)
+ i300 = int(np.argmin(abs(T - 300)))
+ print(f"\n── κ summary {name} ──")
+ print(f"κ_xx, κ_yy, κ_zz (300 K) = {kappa[i300,0]:.2f}, "
+ f"{kappa[i300,1]:.2f}, {kappa[i300,2]:.2f} W/m·K")
+ print(f"κ_avg(300 K) = {kappa_avg[i300]:.2f} W/m·K")
+ # T-scaling: fit log(κ) = a - p*log(T) in high-T regime
+ hi = T > 300
+ if hi.sum() > 2:
+ slope = np.polyfit(np.log(T[hi]), np.log(kappa_avg[hi]), 1)[0]
+ print(f"high-T exponent p (κ ~ T^-p): p ≈ {-slope:.2f} "
+ f"(3-phonon expects ~1)")
+
+
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--model-dir", default="OutputDir")
+ ap.add_argument("--jid", default="JVASP-1002")
+ ap.add_argument("--dim", type=int, nargs=3, default=[2, 2, 2],
+ help="fc3 supercell")
+ ap.add_argument("--dim-fc2", type=int, nargs=3, default=None,
+ help="fc2 supercell (default = fc3 supercell)")
+ ap.add_argument("--mesh", type=int, nargs=3, default=[11, 11, 11])
+ ap.add_argument("--t-max", type=float, default=1000.0)
+ ap.add_argument("--out", default="kappa.html")
+ args = ap.parse_args()
+
+ calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False)
+ d = get_jid_data(jid=args.jid, dataset="dft_3d")
+ atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms
+ formula = atoms.composition.reduced_formula
+ print(f"loaded {args.jid} {formula} N={atoms.num_atoms}")
+
+ relaxed = relax(atoms, calc)
+ prim = relaxed.get_primitive_atoms
+ print(f"relaxed → primitive N={prim.num_atoms}")
+
+ ph3 = run_phono3py(
+ prim, calc,
+ dim=tuple(args.dim),
+ dim_fc2=tuple(args.dim_fc2) if args.dim_fc2 else None,
+ )
+ T, kappa = compute_kappa(ph3, mesh=tuple(args.mesh), t_max=args.t_max)
+ print_summary(T, kappa, formula)
+
+ fig = plot_kappa(T, kappa,
+ title=f"{formula} ({args.jid}) — κ(T) via phono3py+ALIGNN-FF "
+ f"[fc3 {args.dim}, mesh {args.mesh}]")
+ fig.write_html(args.out)
+ print(f"\nwrote {args.out}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/lammps_bridge_torch.py b/alignn/scripts/torch/lammps_bridge_torch.py
new file mode 100644
index 0000000..d12ddc8
--- /dev/null
+++ b/alignn/scripts/torch/lammps_bridge_torch.py
@@ -0,0 +1,197 @@
+"""Drive LAMMPS with ALIGNN-FF via direct torch calls (no ASE wrapper).
+
+Loads `ALIGNNAtomWisePure` from OutputDir/ and registers a LAMMPS
+fix-external callback that:
+ 1. Takes positions/cell from LAMMPS each step
+ 2. Builds the atom + line graph via jarvis (CPU, one-shot per step)
+ 3. Runs the pure-PyTorch model on GPU
+ 4. Returns forces (+ virial) to LAMMPS
+
+This bypasses AlignnAtomwiseCalculator + ASE entirely. Next stop for
+speed: torch.jit.script the model (see `jit_compile()` below) or write
+a C++ pair style.
+
+Usage:
+ python scripts/torch/lammps_bridge_torch.py \\
+ --model-dir OutputDir --input melt_quench.in --types Si
+"""
+from __future__ import annotations
+import argparse
+import json
+import numpy as np
+import torch
+
+from ase import Atoms as AseAtoms
+from jarvis.core.atoms import ase_to_atoms as ase_to_jarvis
+
+from alignn.graphs import Graph
+from alignn.torch_graph_builder import torchgraph_from_dgl
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+
+# eV/ų -> bar (LAMMPS 'metal' units)
+EV_PER_A3_TO_BAR = 1.602176634e6
+
+
+# ---------------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------------
+def load_pure_model(model_dir: str, device, dtype=torch.float32):
+ cfg_raw = json.load(open(f"{model_dir}/config.json"))
+ mcfg = dict(cfg_raw["model"])
+ mcfg["name"] = "alignn_atomwise_pure"
+ # make sure forces + stress autograd paths are on
+ mcfg["calculate_gradient"] = True
+ if "stresswise_weight" in mcfg and mcfg["stresswise_weight"] == 0:
+ mcfg["stresswise_weight"] = 0.01
+ cfg = ALIGNNAtomWisePureConfig(**mcfg)
+ model = ALIGNNAtomWisePure(cfg)
+ sd = torch.load(f"{model_dir}/best_model.pt", map_location=device,
+ weights_only=False)
+ if isinstance(sd, dict) and "model" in sd:
+ sd = sd["model"]
+ missing, unexpected = model.load_state_dict(sd, strict=False)
+ if missing or unexpected:
+ print(f"[warn] state_dict missing={len(missing)} unexpected={len(unexpected)}")
+ model = model.to(device).to(dtype).eval()
+ return model, cfg_raw
+
+
+def jit_compile(model):
+ """Optional: torch.jit.script the model for Path B."""
+ try:
+ return torch.jit.script(model)
+ except Exception as e:
+ print(f"[warn] jit.script failed: {e}\n falling back to eager")
+ return model
+
+
+# ---------------------------------------------------------------------------
+# Graph build — currently jarvis/DGL, the CPU bottleneck.
+# For large systems, replace with an on-device neighbor list.
+# ---------------------------------------------------------------------------
+def build_graph_on_device(symbols, positions, cell, cutoff, max_neighbors,
+ device, dtype):
+ ase_atoms = AseAtoms(symbols=symbols, positions=positions,
+ cell=cell, pbc=True)
+ j = ase_to_jarvis(ase_atoms)
+ g, lg = Graph.atom_dgl_multigraph(
+ j, neighbor_strategy="k-nearest",
+ cutoff=cutoff, max_neighbors=max_neighbors,
+ atom_features="cgcnn", use_canonize=True,
+ )
+ g = g.to(device); lg = lg.to(device)
+ tg = torchgraph_from_dgl(g); tlg = torchgraph_from_dgl(lg)
+ for d in (tg.ndata, tg.edata, tlg.ndata, tlg.edata):
+ for k, v in list(d.items()):
+ if v.is_floating_point():
+ d[k] = v.to(dtype)
+ return tg, tlg
+
+
+# ---------------------------------------------------------------------------
+# Callback factory
+# ---------------------------------------------------------------------------
+def make_callback(model, cfg, symbols_by_type, device, dtype, compute_virial=True):
+ cutoff = float(cfg.get("cutoff", 8.0))
+ max_neighbors = int(cfg.get("max_neighbors", 12))
+ step_counter = [0]
+
+ def callback(lmp, ntimestep, nlocal, tag, x, f):
+ # --- unpack box from LAMMPS ---
+ boxlo, boxhi, xy, yz, xz, *_ = lmp.extract_box()
+ cell = np.array([
+ [boxhi[0] - boxlo[0], 0.0, 0.0],
+ [xy, boxhi[1] - boxlo[1], 0.0],
+ [xz, yz, boxhi[2] - boxlo[2]],
+ ])
+ # --- unpack atom types + positions ---
+ types = np.ctypeslib.as_array(lmp.extract_atom("type"), shape=(nlocal,))
+ symbols = [symbols_by_type[t] for t in types]
+ positions = np.array(x, dtype=np.float64)
+
+ # --- build graph ---
+ tg, tlg = build_graph_on_device(symbols, positions, cell,
+ cutoff, max_neighbors, device, dtype)
+ cell_t = torch.tensor(cell, dtype=dtype, device=device)
+
+ # --- forward ---
+ with torch.enable_grad():
+ out = model((tg, tlg, cell_t))
+
+ # energy + forces
+ energy = out["out"].sum().item()
+ forces = out["grad"].detach().cpu().numpy()
+ f[:] = forces
+
+ # virial (optional)
+ if compute_virial and "stresses" in out:
+ stress_voigt = out["stresses"].detach().cpu().numpy().reshape(-1)[:6]
+ # stresses returned in eV/ų, Voigt order xx yy zz yz xz xy
+ volume = float(np.linalg.det(cell))
+ virial_bar = -stress_voigt * volume * EV_PER_A3_TO_BAR
+ v6 = [
+ float(virial_bar[0]), float(virial_bar[1]), float(virial_bar[2]),
+ float(virial_bar[5]), float(virial_bar[4]), float(virial_bar[3]),
+ ]
+ try:
+ lmp.fix_external_set_virial_global("alignn", v6)
+ except Exception:
+ pass
+
+ step_counter[0] += 1
+ if step_counter[0] % 50 == 0:
+ fmax = float(np.max(np.abs(forces)))
+ print(f"[torch-bridge] step={step_counter[0]:6d} "
+ f"E={energy: .4f} eV |F|max={fmax:.3f} eV/Å")
+
+ return callback
+
+
+# ---------------------------------------------------------------------------
+# Runner
+# ---------------------------------------------------------------------------
+def run(args):
+ from lammps import lammps
+
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ dtype = torch.float32
+ model, cfg_raw = load_pure_model(args.model_dir, device, dtype)
+ if args.jit:
+ model = jit_compile(model)
+
+ symbols_by_type = {i + 1: s for i, s in enumerate(args.types.split(","))}
+
+ lmp = lammps()
+ if args.input:
+ with open(args.input) as fh:
+ lines = fh.readlines()
+ cb = make_callback(model, cfg_raw, symbols_by_type, device, dtype,
+ compute_virial=not args.no_virial)
+ registered = False
+ for line in lines:
+ stripped = line.strip()
+ if (not registered and stripped.startswith("run ")
+ and "alignn" in "".join(lines[: lines.index(line)])):
+ lmp.set_fix_external_callback("alignn", cb, lmp)
+ registered = True
+ lmp.command(line.rstrip("\n"))
+ else:
+ raise SystemExit("Provide --input with a LAMMPS script")
+ lmp.close()
+
+
+def main():
+ p = argparse.ArgumentParser()
+ p.add_argument("--model-dir", required=True, help="dir with best_model.pt + config.json")
+ p.add_argument("--input", required=True, help="LAMMPS input script")
+ p.add_argument("--types", required=True, help="comma-separated element symbols by LAMMPS type")
+ p.add_argument("--jit", action="store_true", help="try torch.jit.script the model")
+ p.add_argument("--no-virial", action="store_true")
+ args = p.parse_args()
+ run(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/md_demo_cu.py b/alignn/scripts/torch/md_demo_cu.py
new file mode 100644
index 0000000..c0fd07d
--- /dev/null
+++ b/alignn/scripts/torch/md_demo_cu.py
@@ -0,0 +1,68 @@
+"""Demo: NVE (Velocity-Verlet) and NVT (Langevin) MD on a Cu supercell
+using pure-PyTorch ALIGNN-FF + pure-PyTorch integrators.
+
+This uses a random-init model for the integrator correctness check
+(energy conservation under NVE, temperature stabilization under NVT).
+Swap in trained weights for production.
+"""
+import numpy as np, torch
+from ase.build import bulk
+from ase.data import atomic_masses
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+from alignn.md import AlignnForces, VelocityVerlet, Langevin, run
+from alignn.md.integrators import maxwell_boltzmann
+
+
+def make_model():
+ cfg = ALIGNNAtomWisePureConfig(
+ name="alignn_atomwise_pure",
+ calculate_gradient=True,
+ atomwise_output_features=0,
+ atom_input_features=92,
+ )
+ return ALIGNNAtomWisePure(cfg)
+
+
+def setup(n=4, T0=300.0, device="cuda", dtype=torch.float32):
+ ase_atoms = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((n, n, n))
+ N = len(ase_atoms)
+ Z = ase_atoms.get_atomic_numbers()
+ cell = np.array(ase_atoms.cell)
+ pos = torch.tensor(ase_atoms.get_positions(), dtype=dtype, device=device)
+ masses = torch.tensor(
+ [atomic_masses[z] for z in Z], dtype=dtype, device=device
+ )
+ model = make_model()
+ forces_fn = AlignnForces(model, Z, cell, device=device, dtype=dtype)
+ g = torch.Generator(device=device).manual_seed(0)
+ vel = maxwell_boltzmann(masses, T=T0, generator=g)
+ print(f"N atoms = {N}")
+ return pos, vel, masses, forces_fn
+
+
+def demo_nve(nsteps=200, n=3):
+ pos, vel, masses, fn = setup(n=n, T0=300.0)
+ integ = VelocityVerlet(forces_fn=fn, masses=masses, dt=1.0)
+ print("\n--- NVE (Velocity-Verlet) ---")
+ _, _, hist = run(integ, pos, vel, nsteps=nsteps, log_every=20)
+ # energy conservation metric
+ T = [r["T_K"] for r in hist]
+ print(f"\nT drift: min={min(T):.1f} max={max(T):.1f} range={max(T)-min(T):.1f} K")
+
+
+def demo_nvt(nsteps=400, n=3, T=300.0):
+ pos, vel, masses, fn = setup(n=n, T0=T)
+ integ = Langevin(forces_fn=fn, masses=masses, dt=1.0, T=T, friction=0.01)
+ print(f"\n--- NVT Langevin (target {T} K) ---")
+ _, _, hist = run(integ, pos, vel, nsteps=nsteps, log_every=40)
+ T_tail = np.mean([r["T_K"] for r in hist[-5:]])
+ print(f"\nmean T over last 5 log points: {T_tail:.1f} K (target {T} K)")
+
+
+if __name__ == "__main__":
+ torch.manual_seed(0)
+ demo_nve(nsteps=100, n=3)
+ demo_nvt(nsteps=200, n=3, T=300.0)
diff --git a/alignn/scripts/torch/md_melt_quench.py b/alignn/scripts/torch/md_melt_quench.py
new file mode 100644
index 0000000..a01cf03
--- /dev/null
+++ b/alignn/scripts/torch/md_melt_quench.py
@@ -0,0 +1,153 @@
+"""Melt-quench MD using pure-PyTorch ALIGNN-FF + in-house NVTBerendsen.
+
+Drop-in replacement for the ASE NVTBerendsen workflow. Same two-stage
+schedule: heat at T0 (melt), cool to T1 (quench), write final POSCAR.
+
+Usage (defaults to a small Si supercell so it runs anywhere):
+ python md_melt_quench.py
+
+For a trained model, pass --model-dir pointing to a directory with
+`best_model.pt` + `config.json` that was trained with the
+`alignn_atomwise` model. The state_dict is loaded into the pure-PyTorch
+variant (layer names match by construction).
+"""
+from __future__ import annotations
+import argparse, json, os, time
+import numpy as np
+import torch
+from ase import Atoms as AseAtoms
+from ase.build import bulk
+from ase.data import atomic_masses
+from ase.io import write as ase_write
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+from alignn.md import AlignnForces, NVTBerendsen, run, maxwell_boltzmann
+from alignn.md.integrators import wrap_pbc, instantaneous_temperature
+
+
+def ensure_cell_size(ase_atoms: AseAtoms, min_size: float):
+ L = ase_atoms.get_cell().lengths()
+ return [max(1, int(np.ceil(min_size / Li))) for Li in L]
+
+
+def build_model(model_dir: str | None, device, dtype):
+ if model_dir and os.path.exists(os.path.join(model_dir, "config.json")):
+ cfg_raw = json.load(open(os.path.join(model_dir, "config.json")))
+ mcfg = cfg_raw["model"]
+ # force pure variant name so the config validates
+ mcfg = {**mcfg, "name": "alignn_atomwise_pure"}
+ model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(**mcfg))
+ sd_path = os.path.join(model_dir, "best_model.pt")
+ sd = torch.load(sd_path, map_location=device, weights_only=False)
+ if isinstance(sd, dict) and "model" in sd:
+ sd = sd["model"]
+ missing, unexpected = model.load_state_dict(sd, strict=False)
+ if missing or unexpected:
+ print(f"[warn] missing={len(missing)} unexpected={len(unexpected)}")
+ print(f"loaded trained weights from {sd_path}")
+ else:
+ print("[note] no model_dir given — using random-init model")
+ model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(
+ name="alignn_atomwise_pure",
+ calculate_gradient=True,
+ atomwise_output_features=0,
+ atom_input_features=92,
+ ))
+ return model.to(device).to(dtype)
+
+
+def make_starting_atoms(args) -> AseAtoms:
+ # Try JARVIS figshare if jid given; otherwise a built-in Si cubic cell.
+ if args.jid:
+ from jarvis.db.figshare import get_jid_data
+ from jarvis.core.atoms import Atoms as JAtoms
+ d = get_jid_data(jid=args.jid, dataset="dft_3d")
+ ja = JAtoms.from_dict(d["atoms"]).get_conventional_atoms
+ ase_atoms = ja.ase_converter()
+ else:
+ ase_atoms = bulk("Si", "diamond", a=5.43, cubic=True)
+ dims = ensure_cell_size(ase_atoms, args.min_size)
+ ase_atoms = ase_atoms.repeat(dims)
+ print(f"starting cell: {len(ase_atoms)} atoms, dims={dims}")
+ return ase_atoms
+
+
+def write_poscar(positions_t, Z, cell_np, path):
+ ase_atoms = AseAtoms(numbers=Z, positions=positions_t.detach().cpu().numpy(),
+ cell=cell_np, pbc=True)
+ ase_write(path, ase_atoms, format="vasp")
+ print(f"wrote {path}")
+
+
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--jid", default=None, help="JARVIS id, e.g. JVASP-1002")
+ ap.add_argument("--model-dir", default=None)
+ ap.add_argument("--outdir", default="output_melt")
+ ap.add_argument("--min-size", type=float, default=10.0)
+ ap.add_argument("--dt", type=float, default=1.0, help="fs")
+ ap.add_argument("--temp0", type=float, default=3500.0)
+ ap.add_argument("--nsteps0", type=int, default=200)
+ ap.add_argument("--temp1", type=float, default=300.0)
+ ap.add_argument("--nsteps1", type=int, default=400)
+ ap.add_argument("--taut", type=float, default=20.0, help="fs")
+ ap.add_argument("--log-every", type=int, default=20)
+ args = ap.parse_args()
+
+ os.makedirs(args.outdir, exist_ok=True)
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ dtype = torch.float32
+
+ ase_atoms = make_starting_atoms(args)
+ Z = ase_atoms.get_atomic_numbers()
+ cell_np = np.array(ase_atoms.cell)
+ cell = torch.tensor(cell_np, dtype=dtype, device=device)
+ positions = torch.tensor(ase_atoms.get_positions(), dtype=dtype, device=device)
+ masses = torch.tensor([atomic_masses[z] for z in Z], dtype=dtype, device=device)
+
+ model = build_model(args.model_dir, device, dtype)
+ forces_fn = AlignnForces(model, Z, cell_np, device=device, dtype=dtype)
+
+ # Initialize velocities at T0
+ gen = torch.Generator(device=device).manual_seed(0)
+ velocities = maxwell_boltzmann(masses, T=args.temp0, generator=gen)
+
+ # Wrap integrator step to also PBC-wrap positions (the model's graph
+ # builder uses the ASE atoms object, which handles PBC internally; we
+ # wrap here mostly to keep output POSCARs clean).
+ integ = NVTBerendsen(forces_fn=forces_fn, masses=masses, dt=args.dt,
+ T=args.temp0, taut=args.taut)
+
+ t_start = time.time()
+ jid_tag = args.jid or "Si"
+
+ # --- Stage 1: melt at T0 ---
+ print(f"\n--- MELT @ {args.temp0} K for {args.nsteps0} steps ---")
+ def cb(i, row): print(f"[melt] step={row['step']:5d} t={row['time_fs']:7.1f} fs "
+ f"T={row['T_K']:7.1f} K wall={row['wall_s']:6.1f} s")
+ positions, velocities, _ = run(integ, positions, velocities,
+ nsteps=args.nsteps0,
+ log_every=args.log_every, callback=cb)
+ positions = wrap_pbc(positions, cell)
+
+ # --- Stage 2: quench to T1 ---
+ integ.set_temperature(args.temp1)
+ print(f"\n--- QUENCH to {args.temp1} K for {args.nsteps1} steps ---")
+ def cb2(i, row): print(f"[quench] step={row['step']:5d} t={row['time_fs']:7.1f} fs "
+ f"T={row['T_K']:7.1f} K wall={row['wall_s']:6.1f} s")
+ positions, velocities, _ = run(integ, positions, velocities,
+ nsteps=args.nsteps1,
+ log_every=args.log_every, callback=cb2)
+ positions = wrap_pbc(positions, cell)
+
+ T_final = instantaneous_temperature(masses, velocities).item()
+ print(f"\nDone. final T = {T_final:.1f} K total wall = {time.time()-t_start:.1f} s")
+
+ poscar = os.path.join(args.outdir, f"POSCAR_{jid_tag}_quenched_alignn.vasp")
+ write_poscar(positions, Z, cell_np, poscar)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/md_test_bussi.py b/alignn/scripts/torch/md_test_bussi.py
new file mode 100644
index 0000000..528c9cf
--- /dev/null
+++ b/alignn/scripts/torch/md_test_bussi.py
@@ -0,0 +1,66 @@
+"""Compare NVTBerendsen vs NVTBussi: relaxation + KE fluctuations.
+
+Canonical KE variance: <ΔKE²> = (3N/2) (kT)² => σ(KE)/⟨KE⟩ = sqrt(2/3N).
+Berendsen should under-fluctuate; Bussi should match.
+"""
+import numpy as np, torch
+from ase.build import bulk
+from ase.data import atomic_masses
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+from alignn.md import (
+ AlignnForces, NVTBerendsen, NVTBussi, NVTNoseHooverChain,
+ run, maxwell_boltzmann,
+)
+from alignn.md.integrators import kinetic_energy
+
+
+def setup(n=3, T=300.0, device="cuda", dtype=torch.float32):
+ a = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((n, n, n))
+ Z = a.get_atomic_numbers()
+ cell = np.array(a.cell)
+ pos = torch.tensor(a.get_positions(), dtype=dtype, device=device)
+ masses = torch.tensor([atomic_masses[z] for z in Z], dtype=dtype, device=device)
+ model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(
+ name="alignn_atomwise_pure", calculate_gradient=True,
+ atomwise_output_features=0, atom_input_features=92))
+ fn = AlignnForces(model, Z, cell, device=device, dtype=dtype)
+ g = torch.Generator(device=device).manual_seed(0)
+ vel = maxwell_boltzmann(masses, T=T, generator=g)
+ return pos, vel, masses, fn
+
+
+def sample_KE(integ, pos, vel, nsteps=400, skip=100):
+ """Run, collect KE samples after equilibration."""
+ Ks = []
+ for i in range(nsteps):
+ pos, vel = integ.step(pos, vel)
+ if i >= skip:
+ Ks.append(kinetic_energy(integ.masses, vel).item())
+ return np.array(Ks)
+
+
+def main():
+ torch.manual_seed(1)
+ T = 300.0
+ pos, vel, masses, fn = setup(n=3, T=T) # 108 atoms
+ N = masses.numel()
+ KE_mean_expected = 0.5 * (3*N - 3) * 8.617e-5 * T
+ sigma_expected = KE_mean_expected * np.sqrt(2.0 / (3*N - 3))
+ print(f"N={N} expected ⟨KE⟩ = {KE_mean_expected:.3f} eV "
+ f"expected σ(KE) = {sigma_expected:.4f} eV")
+
+ for name, integ in [
+ ("Berendsen", NVTBerendsen(forces_fn=fn, masses=masses, dt=1.0, T=T, taut=20.0)),
+ ("Bussi", NVTBussi (forces_fn=fn, masses=masses, dt=1.0, T=T, taut=20.0)),
+ ("NH-chain3", NVTNoseHooverChain(forces_fn=fn, masses=masses, dt=1.0, T=T, taut=20.0, chain_length=3)),
+ ]:
+ Ks = sample_KE(integ, pos.clone(), vel.clone(), nsteps=400, skip=100)
+ print(f"{name:10s} ⟨KE⟩={Ks.mean():.3f} eV σ={Ks.std():.4f} eV "
+ f"σ/σ_expected={Ks.std()/sigma_expected:.2f}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/md_test_fire.py b/alignn/scripts/torch/md_test_fire.py
new file mode 100644
index 0000000..805fcf5
--- /dev/null
+++ b/alignn/scripts/torch/md_test_fire.py
@@ -0,0 +1,46 @@
+"""Perturb a Cu FCC cell and relax with FIRE."""
+import numpy as np, torch
+from ase.build import bulk
+from ase.data import atomic_masses
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+from alignn.md import AlignnForces, FIRE
+
+
+def main():
+ torch.manual_seed(0)
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ dtype = torch.float32
+
+ a = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((2, 2, 2)) # 32 atoms
+ Z = a.get_atomic_numbers()
+ cell = np.array(a.cell)
+ pos = torch.tensor(a.get_positions(), dtype=dtype, device=device)
+ # random perturbation ±0.15 Å
+ pos = pos + 0.15 * torch.randn_like(pos)
+ masses = torch.tensor([atomic_masses[z] for z in Z], dtype=dtype, device=device)
+
+ # Toy harmonic forces_fn: V = 0.5 k (x - x0)², F = -k(x - x0).
+ # This validates FIRE independently of the (untrained) ALIGNN-FF model.
+ x0 = pos.clone().detach()
+ k = 2.0 # eV/Ų
+ def fn(positions):
+ dx = positions - x0
+ E = 0.5 * k * (dx * dx).sum()
+ F = -k * dx
+ return E, F
+ # perturb AFTER defining x0 so forces are nonzero
+ pos = pos + 0.30 * torch.randn_like(pos)
+
+ relax = FIRE(forces_fn=fn, masses=masses, dt=0.1, dt_max=1.0, max_step=0.2)
+ _, hist = relax.run(pos, fmax=0.05, max_steps=500, log_every=20)
+
+ print(f"\ntotal logged points: {len(hist)}")
+ print(f"initial |F|max = {hist[0]['fmax_eV_A']:.4f} eV/Å")
+ print(f"final |F|max = {hist[-1]['fmax_eV_A']:.4f} eV/Å")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt b/alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt
new file mode 100644
index 0000000..75f3348
--- /dev/null
+++ b/alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt
@@ -0,0 +1,27 @@
+# Add to LAMMPS cmake/CMakeLists.txt to enable pair_alignn.
+# Or keep as a standalone external package.
+
+option(PKG_USER-ALIGNN "Build USER-ALIGNN with libtorch" OFF)
+
+if(PKG_USER-ALIGNN)
+ find_package(Torch REQUIRED)
+
+ file(GLOB ALIGNN_SOURCES
+ ${LAMMPS_SOURCE_DIR}/USER-ALIGNN/pair_alignn.cpp)
+ target_sources(lammps PRIVATE ${ALIGNN_SOURCES})
+ target_include_directories(lammps PRIVATE ${LAMMPS_SOURCE_DIR}/USER-ALIGNN)
+ target_link_libraries(lammps PRIVATE ${TORCH_LIBRARIES})
+ target_compile_features(lammps PRIVATE cxx_std_17)
+
+ # libtorch ABI flag — must match how libtorch was built.
+ if(DEFINED ENV{TORCH_CXX11_ABI})
+ target_compile_definitions(lammps PRIVATE _GLIBCXX_USE_CXX11_ABI=$ENV{TORCH_CXX11_ABI})
+ endif()
+endif()
+
+# Build:
+# cd build
+# cmake ../cmake -DPKG_USER-ALIGNN=ON \
+# -DCMAKE_PREFIX_PATH=/path/to/libtorch \
+# -DCMAKE_BUILD_TYPE=Release
+# make -j
diff --git a/alignn/scripts/torch/pair_alignn/pair_alignn.cpp b/alignn/scripts/torch/pair_alignn/pair_alignn.cpp
new file mode 100644
index 0000000..92ac76c
--- /dev/null
+++ b/alignn/scripts/torch/pair_alignn/pair_alignn.cpp
@@ -0,0 +1,282 @@
+// Native LAMMPS pair style for ALIGNN-FF (via libtorch).
+// Single-rank MVP: uses only LOCAL atoms + shift-based edges. The model's
+// energy readout is a global pool over all graph nodes, so ghost atoms
+// must NOT be passed as graph nodes — they'd inflate the sum.
+
+#include "pair_alignn.h"
+
+#include "atom.h"
+#include "comm.h"
+#include "domain.h"
+#include "error.h"
+#include "force.h"
+#include "memory.h"
+#include "neigh_list.h"
+#include "neigh_request.h"
+#include "neighbor.h"
+
+#include
+#include
+#include
+#include
+
+using namespace LAMMPS_NS;
+
+static int element_to_Z(const std::string &sym) {
+ static const std::vector periodic = {
+ "H","He","Li","Be","B","C","N","O","F","Ne","Na","Mg","Al","Si","P",
+ "S","Cl","Ar","K","Ca","Sc","Ti","V","Cr","Mn","Fe","Co","Ni","Cu",
+ "Zn","Ga","Ge","As","Se","Br","Kr","Rb","Sr","Y","Zr","Nb","Mo","Tc",
+ "Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I","Xe","Cs","Ba",
+ };
+ for (size_t i = 0; i < periodic.size(); ++i)
+ if (periodic[i] == sym) return static_cast(i + 1);
+ return -1;
+}
+
+PairAlignn::PairAlignn(LAMMPS *lmp) : Pair(lmp) {
+ single_enable = 0;
+ restartinfo = 0;
+ one_coeff = 1;
+ manybody_flag = 1;
+ no_virial_fdotr_compute = 1;
+}
+
+PairAlignn::~PairAlignn() {
+ if (allocated) {
+ memory->destroy(setflag);
+ memory->destroy(cutsq);
+ }
+}
+
+void PairAlignn::allocate() {
+ allocated = 1;
+ int n = atom->ntypes;
+ memory->create(setflag, n + 1, n + 1, "pair:setflag");
+ memory->create(cutsq, n + 1, n + 1, "pair:cutsq");
+ for (int i = 1; i <= n; i++)
+ for (int j = i; j <= n; j++) setflag[i][j] = 0;
+}
+
+void PairAlignn::settings(int narg, char **arg) {
+ if (narg < 1 || narg > 2)
+ error->all(FLERR, "pair_style alignn: cutoff (Å) [max_neighbors]");
+ cutoff_ = std::stod(arg[0]);
+ if (cutoff_ <= 0.0)
+ error->all(FLERR, "pair_style alignn: cutoff must be positive");
+ if (narg == 2) max_neighbors_ = std::stoi(arg[1]);
+}
+
+void PairAlignn::coeff(int narg, char **arg) {
+ if (!allocated) allocate();
+ if (narg != 3 + atom->ntypes)
+ error->all(FLERR,
+ "pair_coeff * * model.pt [ ...]");
+ if (strcmp(arg[0], "*") != 0 || strcmp(arg[1], "*") != 0)
+ error->all(FLERR, "pair_coeff alignn requires * *");
+
+ const std::string model_path = arg[2];
+ type_to_Z_.assign(atom->ntypes + 1, 0);
+ for (int t = 1; t <= atom->ntypes; ++t) {
+ std::string sym = arg[2 + t];
+ int Z = element_to_Z(sym);
+ if (Z < 0) error->all(FLERR,
+ (std::string("Unknown element: ") + sym).c_str());
+ type_to_Z_[t] = Z;
+ }
+
+ device_ = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU;
+ try {
+ model_ = torch::jit::load(model_path, device_);
+ model_.eval();
+ } catch (const c10::Error &e) {
+ error->all(FLERR,
+ (std::string("Failed to load TorchScript model: ") + e.what()).c_str());
+ }
+ model_loaded_ = true;
+ for (int i = 1; i <= atom->ntypes; i++)
+ for (int j = i; j <= atom->ntypes; j++) setflag[i][j] = 1;
+}
+
+void PairAlignn::init_style() {
+ if (!model_loaded_)
+ error->all(FLERR, "pair_style alignn requires pair_coeff * * ...");
+ if (atom->tag_enable == 0)
+ error->all(FLERR, "pair_style alignn requires atom IDs");
+ if (force->newton_pair == 0)
+ error->all(FLERR, "pair_style alignn requires newton on");
+ if (atom->map_style == Atom::MAP_NONE) atom->map_init();
+ neighbor->add_request(this, NeighConst::REQ_FULL | NeighConst::REQ_GHOST);
+}
+
+double PairAlignn::init_one(int /*i*/, int /*j*/) { return cutoff_; }
+
+// ---------------------------------------------------------------------------
+void PairAlignn::compute(int eflag, int vflag) {
+ ev_init(eflag, vflag);
+
+ const int nlocal = atom->nlocal;
+ double **x = atom->x;
+ double **f = atom->f;
+ int *type = atom->type;
+ tagint *tag = atom->tag;
+
+ int *ilist = list->ilist;
+ int *numneigh = list->numneigh;
+ int **firstneigh = list->firstneigh;
+ const int inum = list->inum;
+
+ const double hx = domain->xprd;
+ const double hy = domain->yprd;
+ const double hz = domain->zprd;
+ const double xy = domain->xy;
+ const double xz = domain->xz;
+ const double yz = domain->yz;
+ const bool triclinic = (xy != 0.0) || (xz != 0.0) || (yz != 0.0);
+
+ // --- positions / atomic numbers: LOCAL atoms only ---
+ auto t_opts = torch::TensorOptions().dtype(dtype_).device(torch::kCPU);
+ auto l_opts = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
+ torch::Tensor pos = torch::empty({nlocal, 3}, t_opts);
+ torch::Tensor Z = torch::empty({nlocal}, l_opts);
+ {
+ auto pa = pos.accessor();
+ auto Za = Z.accessor();
+ for (int i = 0; i < nlocal; ++i) {
+ pa[i][0] = static_cast(x[i][0]);
+ pa[i][1] = static_cast(x[i][1]);
+ pa[i][2] = static_cast(x[i][2]);
+ Za[i] = static_cast(type_to_Z_[type[i]]);
+ }
+ }
+
+ // --- edges: src/dst both LOCAL; shift carries PBC image offset ---
+ std::vector src_vec, dst_vec;
+ std::vector shift_flat;
+ src_vec.reserve(nlocal * 16);
+ dst_vec.reserve(nlocal * 16);
+ shift_flat.reserve(nlocal * 16 * 3);
+ const double cut2 = cutoff_ * cutoff_;
+
+ // k-nearest (same convention as jarvis `neighbor_strategy="k-nearest"`):
+ // for each local i, keep up to max_neighbors_ closest j's within cutoff.
+ struct Cand { double r2; int dst_local; float shx, shy, shz; };
+ std::vector buf;
+ buf.reserve(256);
+
+ for (int ii = 0; ii < inum; ++ii) {
+ const int i = ilist[ii];
+ if (i >= nlocal) continue;
+ const double xi = x[i][0], yi = x[i][1], zi = x[i][2];
+ const int jnum = numneigh[i];
+ const int *jlist = firstneigh[i];
+ buf.clear();
+ for (int jj = 0; jj < jnum; ++jj) {
+ int j = jlist[jj];
+ j &= NEIGHMASK;
+ const double dxr = x[j][0] - xi;
+ const double dyr = x[j][1] - yi;
+ const double dzr = x[j][2] - zi;
+ const double r2 = dxr*dxr + dyr*dyr + dzr*dzr;
+ if (r2 >= cut2) continue;
+
+ int dst_local;
+ double shx = 0.0, shy = 0.0, shz = 0.0;
+ if (j < nlocal) {
+ dst_local = j;
+ } else {
+ tagint jtag = tag[j];
+ int owner = atom->map(jtag);
+ if (owner < 0 || owner >= nlocal) continue;
+ dst_local = owner;
+ const double dox = x[j][0] - x[owner][0];
+ const double doy = x[j][1] - x[owner][1];
+ const double doz = x[j][2] - x[owner][2];
+ (void)triclinic;
+ shx = std::round(dox / hx);
+ shy = std::round(doy / hy);
+ shz = std::round(doz / hz);
+ }
+ buf.push_back({r2, dst_local,
+ static_cast(shx),
+ static_cast(shy),
+ static_cast(shz)});
+ }
+ // Sort by r2 ascending, keep first max_neighbors_.
+ if (static_cast(buf.size()) > max_neighbors_) {
+ std::partial_sort(buf.begin(),
+ buf.begin() + max_neighbors_,
+ buf.end(),
+ [](const Cand &a, const Cand &b) { return a.r2 < b.r2; });
+ buf.resize(max_neighbors_);
+ }
+ for (const auto &c : buf) {
+ src_vec.push_back(i);
+ dst_vec.push_back(c.dst_local);
+ shift_flat.push_back(c.shx);
+ shift_flat.push_back(c.shy);
+ shift_flat.push_back(c.shz);
+ }
+ }
+ const long nedges = static_cast(src_vec.size());
+ torch::Tensor src = torch::from_blob(src_vec.data(), {nedges}, torch::kLong).clone();
+ torch::Tensor dst = torch::from_blob(dst_vec.data(), {nedges}, torch::kLong).clone();
+ torch::Tensor shift = torch::from_blob(shift_flat.data(),{nedges, 3}, torch::kFloat32)
+ .to(dtype_).clone();
+
+ torch::Tensor lat = torch::zeros({3, 3}, t_opts);
+ {
+ auto la = lat.accessor();
+ la[0][0] = static_cast(hx);
+ la[1][1] = static_cast(hy);
+ la[2][2] = static_cast(hz);
+ la[1][0] = static_cast(xy);
+ la[2][0] = static_cast(xz);
+ la[2][1] = static_cast(yz);
+ }
+
+ pos = pos.to(device_).set_requires_grad(true);
+ lat = lat.to(device_);
+ Z = Z.to(device_);
+ src = src.to(device_);
+ dst = dst.to(device_);
+ shift = shift.to(device_);
+
+ std::vector inputs;
+ inputs.emplace_back(pos);
+ inputs.emplace_back(lat);
+ inputs.emplace_back(Z);
+ inputs.emplace_back(src);
+ inputs.emplace_back(dst);
+ inputs.emplace_back(shift);
+ inputs.emplace_back(static_cast(vflag_global));
+
+ c10::IValue result;
+ {
+ torch::AutoGradMode grad(true);
+ result = model_.get_method("forward_tensors_z")(inputs);
+ }
+ auto dict = result.toGenericDict();
+ torch::Tensor energy_t = dict.at("energy").toTensor().to(torch::kCPU).to(torch::kDouble);
+ torch::Tensor forces_t = dict.at("forces").toTensor().to(torch::kCPU).to(torch::kDouble);
+
+ auto fa = forces_t.accessor();
+ for (int i = 0; i < nlocal; ++i) {
+ f[i][0] += fa[i][0];
+ f[i][1] += fa[i][1];
+ f[i][2] += fa[i][2];
+ }
+ if (eflag_global) eng_vdwl += energy_t.item();
+
+ if (vflag_global && dict.contains("stress")) {
+ torch::Tensor s = dict.at("stress").toTensor().to(torch::kCPU).to(torch::kDouble);
+ auto sa = s.accessor();
+ const double V = hx * hy * hz;
+ virial[0] += -sa[0][0] * V;
+ virial[1] += -sa[1][1] * V;
+ virial[2] += -sa[2][2] * V;
+ virial[3] += -sa[0][1] * V;
+ virial[4] += -sa[0][2] * V;
+ virial[5] += -sa[1][2] * V;
+ }
+}
diff --git a/alignn/scripts/torch/pair_alignn/pair_alignn.h b/alignn/scripts/torch/pair_alignn/pair_alignn.h
new file mode 100644
index 0000000..ec45e5f
--- /dev/null
+++ b/alignn/scripts/torch/pair_alignn/pair_alignn.h
@@ -0,0 +1,55 @@
+// ALIGNN-FF native LAMMPS pair style.
+// Drop this header + pair_alignn.cpp into LAMMPS src/USER-ALIGNN/ and rebuild.
+//
+// Input script usage:
+// pair_style alignn
+// pair_coeff * * alignn_ff.pt Si # TorchScript model path, then element
+// # symbols ordered by LAMMPS atom type.
+
+#ifdef PAIR_CLASS
+// clang-format off
+PairStyle(alignn,PairAlignn)
+// clang-format on
+#else
+
+#ifndef LMP_PAIR_ALIGNN_H
+#define LMP_PAIR_ALIGNN_H
+
+#include "pair.h"
+
+#include
+#include
+#include
+#include
+
+namespace LAMMPS_NS {
+
+class PairAlignn : public Pair {
+ public:
+ PairAlignn(class LAMMPS *);
+ ~PairAlignn() override;
+
+ void compute(int, int) override;
+ void settings(int, char **) override;
+ void coeff(int, char **) override;
+ void init_style() override;
+ double init_one(int, int) override;
+ void allocate();
+
+ protected:
+ torch::jit::script::Module model_;
+ torch::Device device_ = torch::kCPU;
+ torch::ScalarType dtype_ = torch::kFloat32;
+
+ double cutoff_ = 8.0; // must match trained cutoff
+ int max_neighbors_ = 12; // k-nearest cap per atom
+ std::vector type_to_Z_; // LAMMPS type (1-indexed) -> atomic number
+
+ bool model_loaded_ = false;
+ bool debug_ = false;
+};
+
+} // namespace LAMMPS_NS
+
+#endif
+#endif
diff --git a/alignn/scripts/torch/pair_alignn/si_native.in b/alignn/scripts/torch/pair_alignn/si_native.in
new file mode 100644
index 0000000..de048bb
--- /dev/null
+++ b/alignn/scripts/torch/pair_alignn/si_native.in
@@ -0,0 +1,40 @@
+# Native-pair-style melt-quench for Si using pair_alignn.
+
+units metal
+atom_style atomic
+boundary p p p
+read_data si.data
+mass 1 28.0855
+
+pair_style alignn
+pair_coeff * * alignn_ff.pt Si
+
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+
+timestep 0.001
+thermo 50
+thermo_style custom step temp pe ke etotal press vol
+dump traj all custom 100 si.lammpstrj id type x y z
+
+velocity all create 300.0 12345 mom yes rot yes
+fix nvt1 all nvt temp 300.0 300.0 0.1
+run 200
+unfix nvt1
+
+fix nvt2 all nvt temp 300.0 3500.0 0.1
+run 1000
+unfix nvt2
+
+fix nvt3 all nvt temp 3500.0 3500.0 0.1
+run 2000
+unfix nvt3
+
+fix nvt4 all nvt temp 3500.0 300.0 0.1
+run 4000
+unfix nvt4
+
+fix nvt5 all nvt temp 300.0 300.0 0.1
+run 2000
+
+write_data si_quenched.data
diff --git a/alignn/scripts/torch/phonons_plotly.py b/alignn/scripts/torch/phonons_plotly.py
new file mode 100644
index 0000000..9bd7f85
--- /dev/null
+++ b/alignn/scripts/torch/phonons_plotly.py
@@ -0,0 +1,131 @@
+"""Phonon bandstructure from a trained ALIGNN-FF model, plotted with Plotly.
+
+Usage:
+ python phonons_plotly.py --model-dir OutputDir --jid JVASP-1002
+"""
+from __future__ import annotations
+import argparse
+import numpy as np
+import plotly.graph_objects as go
+from ase import Atoms as AseAtoms
+from ase.filters import ExpCellFilter
+from ase.optimize import FIRE
+from phonopy import Phonopy
+from phonopy.structure.atoms import PhonopyAtoms
+
+from jarvis.core.atoms import Atoms, ase_to_atoms
+from jarvis.core.kpoints import Kpoints3D as Kpoints
+from jarvis.db.figshare import get_jid_data
+
+from alignn.ff.ff import AlignnAtomwiseCalculator
+
+
+def relax(atoms: Atoms, calc, fmax=0.05, steps=100):
+ a = atoms.ase_converter()
+ a.calc = calc
+ FIRE(ExpCellFilter(a), logfile=None).run(fmax=fmax, steps=steps)
+ return ase_to_atoms(a)
+
+
+def compute_phonons(atoms: Atoms, calc, dim=(2, 2, 2), distance=0.03, line_density=20):
+ bulk = PhonopyAtoms(
+ symbols=atoms.elements,
+ scaled_positions=atoms.frac_coords,
+ cell=atoms.lattice_mat,
+ )
+ ph = Phonopy(bulk, [[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]])
+ ph.generate_displacements(distance=distance)
+ # evaluate forces on each displaced supercell
+ forces_set = []
+ for sc in ph.get_supercells_with_displacements():
+ a = AseAtoms(symbols=sc.get_chemical_symbols(),
+ scaled_positions=sc.get_scaled_positions(),
+ cell=sc.get_cell(), pbc=True)
+ a.calc = calc
+ F = np.array(a.get_forces())
+ F -= F.mean(axis=0) # remove drift
+ forces_set.append(F)
+ ph.produce_force_constants(forces=forces_set)
+
+ # walk the standard k-path (jarvis builds it from the lattice)
+ kp = Kpoints().kpath(atoms, line_density=line_density)
+ q_points = np.array(kp.kpts)
+ labels_raw = kp.labels
+
+ freqs, distances, tick_pos, tick_labels = [], [], [], []
+ last_kstr = None
+ d = 0.0
+ for i, (q, lbl) in enumerate(zip(q_points, labels_raw)):
+ if i > 0:
+ d += np.linalg.norm(q_points[i] - q_points[i - 1])
+ distances.append(d)
+ freqs.append(ph.get_frequencies(q)) # THz
+ kstr = ",".join(f"{x:.6f}" for x in q)
+ if lbl and kstr != last_kstr:
+ tick_pos.append(d)
+ tick_labels.append(lbl)
+ last_kstr = kstr
+ return np.array(distances), np.array(freqs), tick_pos, tick_labels
+
+
+def plot(distances, freqs, tick_pos, tick_labels, title=""):
+ fig = go.Figure()
+ nbands = freqs.shape[1]
+ for b in range(nbands):
+ fig.add_trace(go.Scatter(
+ x=distances, y=freqs[:, b], mode="lines",
+ line=dict(color="steelblue", width=1.5),
+ showlegend=False, hovertemplate="d=%{x:.3f}
ν=%{y:.3f} THz",
+ ))
+ fig.add_hline(y=0, line=dict(color="black", dash="dot", width=1))
+ for xp in tick_pos[1:-1]:
+ fig.add_vline(x=xp, line=dict(color="lightgray", width=1))
+ fig.update_layout(
+ title=title, width=800, height=500,
+ xaxis=dict(
+ title="Wave vector",
+ tickmode="array",
+ tickvals=tick_pos,
+ ticktext=[l.replace("\\Gamma", "Γ") for l in tick_labels],
+ showgrid=False, zeroline=False,
+ ),
+ yaxis=dict(title="Frequency (THz)", gridcolor="lightgray"),
+ template="plotly_white",
+ )
+ return fig
+
+
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--model-dir", default="OutputDir")
+ ap.add_argument("--jid", default="JVASP-1002")
+ ap.add_argument("--dim", type=int, nargs=3, default=[2, 2, 2])
+ ap.add_argument("--out", default="phonons.html")
+ args = ap.parse_args()
+
+ calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False)
+
+ d = get_jid_data(jid=args.jid, dataset="dft_3d")
+ atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms
+ print(f"loaded {args.jid} {atoms.composition.reduced_formula} N={atoms.num_atoms}")
+
+ relaxed = relax(atoms, calc)
+ prim = relaxed.get_primitive_atoms
+ print(f"relaxed → primitive N={prim.num_atoms}")
+
+ distances, freqs, tick_pos, tick_labels = compute_phonons(
+ prim, calc, dim=tuple(args.dim)
+ )
+ n_imag = int((freqs < -0.05).sum())
+ print(f"bands computed: {freqs.shape[1]} "
+ f"imag modes (<-0.05 THz): {n_imag}")
+
+ fig = plot(distances, freqs, tick_pos, tick_labels,
+ title=f"Phonon bands — {atoms.composition.reduced_formula} "
+ f"({args.jid}) supercell {args.dim}")
+ fig.write_html(args.out)
+ print(f"wrote {args.out}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/scaling_test_cu.py b/alignn/scripts/torch/scaling_test_cu.py
new file mode 100644
index 0000000..61aa6ba
--- /dev/null
+++ b/alignn/scripts/torch/scaling_test_cu.py
@@ -0,0 +1,167 @@
+"""ALIGNN-FF (pure PyTorch) scaling test on a Cu FCC supercell.
+
+Sweeps supercell size 1x1x1 .. NxNxN across fp32 / bf16-autocast / fp16-autocast.
+Records peak GPU memory and wall time for a single forward+backward (E, F).
+
+Usage:
+ python scaling_test_cu.py --max 10 --model-dir alignn/ff/alignnff_wt01
+"""
+
+from __future__ import annotations
+import argparse, gc, json, os, time, traceback
+import numpy as np
+import torch
+from ase.build import bulk
+from jarvis.core.atoms import ase_to_atoms as ase_to_jarvis
+
+from alignn.graphs import Graph
+from alignn.torch_graph_builder import torchgraph_from_dgl
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+
+
+def build_cu_supercell(n: int):
+ a = bulk("Cu", "fcc", a=3.615, cubic=True) # 4 atoms / cubic cell
+ return a.repeat((n, n, n))
+
+
+def build_graphs(ase_atoms, cutoff=8.0, max_neighbors=12, device="cuda"):
+ j = ase_to_jarvis(ase_atoms)
+ g, lg = Graph.atom_dgl_multigraph(
+ j, neighbor_strategy="k-nearest",
+ cutoff=cutoff, max_neighbors=max_neighbors,
+ atom_features="cgcnn", use_canonize=True,
+ )
+ g = g.to(device); lg = lg.to(device)
+ tg = torchgraph_from_dgl(g)
+ tlg = torchgraph_from_dgl(lg)
+ return tg, tlg
+
+
+def make_model(device):
+ cfg = ALIGNNAtomWisePureConfig(
+ name="alignn_atomwise_pure",
+ calculate_gradient=True,
+ atomwise_output_features=0,
+ atom_input_features=92,
+ )
+ m = ALIGNNAtomWisePure(cfg).to(device)
+ m.eval()
+ return m
+
+
+def _cast_graph(tg, dtype):
+ for k, v in list(tg.ndata.items()):
+ if v.is_floating_point():
+ tg.ndata[k] = v.to(dtype)
+ for k, v in list(tg.edata.items()):
+ if v.is_floating_point():
+ tg.edata[k] = v.to(dtype)
+
+
+def run_one(model, tg, tlg, cell, device, precision, model_fp32_state, fp32_snapshot):
+ torch.cuda.empty_cache(); gc.collect()
+ torch.cuda.reset_peak_memory_stats(device)
+ model.to(torch.float32)
+ model.load_state_dict(model_fp32_state)
+ # restore graph tensors from fp32 snapshot
+ ndata0, edata0, lndata0, ledata0 = fp32_snapshot
+ tg.ndata = {k: v.detach().clone() for k, v in ndata0.items()}
+ tg.edata = {k: v.detach().clone() for k, v in edata0.items()}
+ tlg.ndata = {k: v.detach().clone() for k, v in lndata0.items()}
+ tlg.edata = {k: v.detach().clone() for k, v in ledata0.items()}
+ cell = cell.to(torch.float32)
+ torch.cuda.synchronize()
+ t0 = time.time()
+ try:
+ if precision == "fp32":
+ out = model((tg, tlg, cell))
+ elif precision == "bf16":
+ with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+ out = model((tg, tlg, cell))
+ elif precision == "fp16":
+ with torch.autocast(device_type="cuda", dtype=torch.float16):
+ out = model((tg, tlg, cell))
+ elif precision == "bf16_pure":
+ model.to(torch.bfloat16)
+ _cast_graph(tg, torch.bfloat16)
+ _cast_graph(tlg, torch.bfloat16)
+ out = model((tg, tlg, cell.to(torch.bfloat16)))
+ elif precision == "fp16_pure":
+ model.to(torch.float16)
+ _cast_graph(tg, torch.float16)
+ _cast_graph(tlg, torch.float16)
+ out = model((tg, tlg, cell.to(torch.float16)))
+ else:
+ raise ValueError(precision)
+ # Force backward to capture gradient memory
+ e = out["out"].sum() if "out" in out else out.get("total_energy", None)
+ if e is None: # fallback
+ e = next(iter(out.values())).sum()
+ e.backward()
+ torch.cuda.synchronize()
+ dt = time.time() - t0
+ peak = torch.cuda.max_memory_allocated(device) / 1e9
+ return {"ok": True, "time_s": dt, "peak_gb": peak}
+ except torch.cuda.OutOfMemoryError as ex:
+ return {"ok": False, "err": "OOM"}
+ except Exception as ex:
+ return {"ok": False, "err": f"{type(ex).__name__}: {ex}"}
+
+
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--max", type=int, default=10)
+ ap.add_argument("--min", type=int, default=1)
+ ap.add_argument("--out", default="scaling_cu_results.json")
+ args = ap.parse_args()
+
+ device = torch.device("cuda")
+ gpu = torch.cuda.get_device_name(0)
+ total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+ print(f"GPU: {gpu} total={total_gb:.1f} GB")
+
+ results = {"gpu": gpu, "total_gb": total_gb, "runs": []}
+ model = make_model(device)
+ n_params = sum(p.numel() for p in model.parameters())
+ print(f"model params: {n_params/1e6:.2f} M")
+ import copy
+ model_fp32_state = copy.deepcopy(model.state_dict())
+
+ for n in range(args.min, args.max + 1):
+ atoms = build_cu_supercell(n)
+ N = len(atoms)
+ try:
+ tg, tlg = build_graphs(atoms, device=device)
+ except torch.cuda.OutOfMemoryError:
+ print(f"[{n}^3 N={N}] graph build OOM"); break
+ except Exception as ex:
+ print(f"[{n}^3 N={N}] graph build failed: {ex}"); break
+ n_edges = tg.src.numel(); n_triplets = tlg.src.numel()
+ cell = torch.tensor(np.array(atoms.cell), dtype=torch.float32, device=device)
+ print(f"\n=== {n}x{n}x{n} N={N} edges={n_edges} triplets={n_triplets} ===")
+ row = {"n": n, "N": N, "edges": n_edges, "triplets": n_triplets}
+ snap = (
+ {k: v.detach().clone() for k, v in tg.ndata.items()},
+ {k: v.detach().clone() for k, v in tg.edata.items()},
+ {k: v.detach().clone() for k, v in tlg.ndata.items()},
+ {k: v.detach().clone() for k, v in tlg.edata.items()},
+ )
+ for prec in ["fp32", "bf16", "fp16", "bf16_pure", "fp16_pure"]:
+ r = run_one(model, tg, tlg, cell, device, prec, model_fp32_state, snap)
+ row[prec] = r
+ tag = f"{r['peak_gb']:.2f} GB {r['time_s']*1000:.0f} ms" if r["ok"] else r.get("err","?")
+ print(f" {prec:5s}: {tag}")
+ results["runs"].append(row)
+ with open(args.out, "w") as f:
+ json.dump(results, f, indent=2)
+ # stop if all precisions OOM
+ if not any(row[p]["ok"] for p in ["fp32","bf16","fp16","bf16_pure","fp16_pure"]):
+ print("all precisions OOM — stopping"); break
+ del tg, tlg, cell; torch.cuda.empty_cache(); gc.collect()
+
+ print(f"\nSaved -> {args.out}")
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/si_melt_quench.in b/alignn/scripts/torch/si_melt_quench.in
new file mode 100644
index 0000000..cad37ba
--- /dev/null
+++ b/alignn/scripts/torch/si_melt_quench.in
@@ -0,0 +1,49 @@
+# LAMMPS melt-quench for Si with ALIGNN-FF via torch bridge.
+# Invoke: python scripts/torch/lammps_bridge_torch.py \
+# --model-dir OutputDir --input scripts/torch/si_melt_quench.in --types Si
+
+units metal
+atom_style atomic
+boundary p p p
+
+read_data si.data
+mass 1 28.0855
+
+# ALIGNN-FF hook — callback provides forces + virial every step
+pair_style zero 8.0
+pair_coeff * *
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+fix alignn all external pf/callback 1 1
+
+timestep 0.001 # 1 fs (metal units: ps)
+thermo 50
+thermo_style custom step temp pe ke etotal press vol
+dump traj all custom 100 si_melt_quench.lammpstrj id type x y z
+
+# ── Stage 1 — equilibrate 300 K ──
+velocity all create 300.0 12345 mom yes rot yes
+fix nvt1 all nvt temp 300.0 300.0 0.1
+run 200
+unfix nvt1
+
+# ── Stage 2 — heat to 3500 K over 1 ps ──
+fix nvt2 all nvt temp 300.0 3500.0 0.1
+run 1000
+unfix nvt2
+
+# ── Stage 3 — hold molten ──
+fix nvt3 all nvt temp 3500.0 3500.0 0.1
+run 2000
+unfix nvt3
+
+# ── Stage 4 — quench to 300 K over 4 ps ──
+fix nvt4 all nvt temp 3500.0 300.0 0.1
+run 4000
+unfix nvt4
+
+# ── Stage 5 — anneal 300 K ──
+fix nvt5 all nvt temp 300.0 300.0 0.1
+run 2000
+
+write_data si_quenched.data
diff --git a/alignn/scripts/torch/thermal_plotly.py b/alignn/scripts/torch/thermal_plotly.py
new file mode 100644
index 0000000..c1729e2
--- /dev/null
+++ b/alignn/scripts/torch/thermal_plotly.py
@@ -0,0 +1,240 @@
+"""Phonon + thermal properties from trained ALIGNN-FF, plotted with Plotly.
+
+Computes:
+ * phonon bandstructure
+ * phonon DOS
+ * harmonic thermodynamics: Cv(T), S(T), F(T), ZPE
+ * Dulong-Petit comparison
+
+Outputs a single HTML with 4 panels.
+
+Usage:
+ python thermal_plotly.py --model-dir OutputDir --jid JVASP-1002
+"""
+from __future__ import annotations
+import argparse
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+from ase import Atoms as AseAtoms
+from ase.filters import ExpCellFilter
+from ase.optimize import FIRE
+from phonopy import Phonopy
+from phonopy.structure.atoms import PhonopyAtoms
+
+from jarvis.core.atoms import Atoms, ase_to_atoms
+from jarvis.core.kpoints import Kpoints3D as Kpoints
+from jarvis.db.figshare import get_jid_data
+
+from alignn.ff.ff import AlignnAtomwiseCalculator
+
+
+# ---------------------------------------------------------------------------
+# Core calculations
+# ---------------------------------------------------------------------------
+def relax(atoms, calc, fmax=0.05, steps=100):
+ a = atoms.ase_converter()
+ a.calc = calc
+ FIRE(ExpCellFilter(a), logfile=None).run(fmax=fmax, steps=steps)
+ return ase_to_atoms(a)
+
+
+def run_phonopy(atoms, calc, dim=(2, 2, 2), distance=0.03):
+ bulk = PhonopyAtoms(
+ symbols=atoms.elements,
+ scaled_positions=atoms.frac_coords,
+ cell=atoms.lattice_mat,
+ )
+ ph = Phonopy(bulk, [[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]])
+ ph.generate_displacements(distance=distance)
+
+ forces_set = []
+ for sc in ph.supercells_with_displacements:
+ a = AseAtoms(
+ symbols=list(sc.symbols),
+ scaled_positions=np.array(sc.scaled_positions),
+ cell=np.array(sc.cell),
+ pbc=True,
+ )
+ a.calc = calc
+ F = np.array(a.get_forces())
+ F -= F.mean(axis=0)
+ forces_set.append(F)
+ ph.produce_force_constants(forces=forces_set)
+ return ph
+
+
+def bandstructure(ph, atoms, line_density=20):
+ kp = Kpoints().kpath(atoms, line_density=line_density)
+ q_points = np.array(kp.kpts)
+ labels_raw = kp.labels
+
+ freqs, distances, tick_pos, tick_labels = [], [], [], []
+ last_kstr = None
+ d = 0.0
+ for i, (q, lbl) in enumerate(zip(q_points, labels_raw)):
+ if i > 0:
+ d += np.linalg.norm(q_points[i] - q_points[i - 1])
+ distances.append(d)
+ freqs.append(ph.get_frequencies(q))
+ kstr = ",".join(f"{x:.6f}" for x in q)
+ if lbl and kstr != last_kstr:
+ tick_pos.append(d)
+ tick_labels.append(lbl)
+ last_kstr = kstr
+ return np.array(distances), np.array(freqs), tick_pos, tick_labels
+
+
+def thermal_properties(ph, mesh=(20, 20, 20), t_min=0, t_max=1000, t_step=10):
+ ph.run_mesh(list(mesh))
+ ph.run_thermal_properties(t_min=t_min, t_max=t_max, t_step=t_step)
+ td = ph.get_thermal_properties_dict()
+ N = len(ph.primitive)
+ return {
+ "T": td["temperatures"],
+ "Cv": td["heat_capacity"] / N, # J/K/mol-atom
+ "S": td["entropy"] / N, # J/K/mol-atom
+ "F": td["free_energy"] / N, # kJ/mol-atom
+ "ZPE_per_cell_kJmol": float(td["free_energy"][0]),
+ "ZPE_per_atom_eV": float(td["free_energy"][0]) * 1000.0 / 96.485 / N,
+ }
+
+
+def dos(ph):
+ ph.run_total_dos()
+ d = ph.get_total_dos_dict()
+ return np.array(d["frequency_points"]), np.array(d["total_dos"])
+
+
+# ---------------------------------------------------------------------------
+# Plotting
+# ---------------------------------------------------------------------------
+def plot_all(bands, dos_data, thermo, title=""):
+ distances, freqs, tick_pos, tick_labels = bands
+ dos_freq, dos_val = dos_data
+
+ fig = make_subplots(
+ rows=2, cols=3,
+ specs=[
+ [{"colspan": 2}, None, {}],
+ [{}, {}, {}],
+ ],
+ subplot_titles=(
+ "Phonon bandstructure", "Phonon DOS",
+ "Heat capacity Cv", "Entropy S", "Free energy F",
+ ),
+ horizontal_spacing=0.08, vertical_spacing=0.15,
+ )
+
+ # Bandstructure (top-left, 2 cols wide)
+ for b in range(freqs.shape[1]):
+ fig.add_trace(
+ go.Scatter(x=distances, y=freqs[:, b], mode="lines",
+ line=dict(color="steelblue", width=1.2),
+ showlegend=False,
+ hovertemplate="d=%{x:.2f}
ν=%{y:.2f} THz"),
+ row=1, col=1,
+ )
+ for xp in tick_pos[1:-1]:
+ fig.add_vline(x=xp, line=dict(color="lightgray", width=1), row=1, col=1)
+ fig.add_hline(y=0, line=dict(color="black", dash="dot", width=1), row=1, col=1)
+
+ # DOS (top-right)
+ fig.add_trace(
+ go.Scatter(x=dos_val, y=dos_freq, mode="lines",
+ fill="tozerox", line=dict(color="firebrick"),
+ showlegend=False,
+ hovertemplate="DOS=%{x:.2f}
ν=%{y:.2f} THz"),
+ row=1, col=3,
+ )
+
+ # Cv, S, F (bottom row)
+ T = thermo["T"]
+ fig.add_trace(go.Scatter(x=T, y=thermo["Cv"], mode="lines",
+ line=dict(color="seagreen"), showlegend=False),
+ row=2, col=1)
+ fig.add_hline(y=24.94, line=dict(dash="dot", color="gray"),
+ annotation_text="3R (Dulong-Petit)",
+ annotation_position="top right", row=2, col=1)
+ fig.add_trace(go.Scatter(x=T, y=thermo["S"], mode="lines",
+ line=dict(color="darkorange"), showlegend=False),
+ row=2, col=2)
+ fig.add_trace(go.Scatter(x=T, y=thermo["F"], mode="lines",
+ line=dict(color="mediumpurple"), showlegend=False),
+ row=2, col=3)
+
+ # Axes labels
+ ticktext = [l.replace("\\Gamma", "Γ") for l in tick_labels]
+ fig.update_xaxes(title_text="Wave vector",
+ tickmode="array", tickvals=tick_pos, ticktext=ticktext,
+ showgrid=False, row=1, col=1)
+ fig.update_yaxes(title_text="Frequency (THz)", row=1, col=1)
+ fig.update_xaxes(title_text="DOS (states/THz/cell)", row=1, col=3)
+ fig.update_yaxes(title_text="Frequency (THz)", row=1, col=3)
+
+ for col, ylab in [(1, "Cv (J/K/mol-atom)"),
+ (2, "S (J/K/mol-atom)"),
+ (3, "F (kJ/mol-atom)")]:
+ fig.update_xaxes(title_text="T (K)", row=2, col=col)
+ fig.update_yaxes(title_text=ylab, row=2, col=col)
+
+ fig.update_layout(template="plotly_white", title=title,
+ width=1250, height=800)
+ return fig
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--model-dir", default="OutputDir")
+ ap.add_argument("--jid", default="JVASP-1002")
+ ap.add_argument("--dim", type=int, nargs=3, default=[2, 2, 2])
+ ap.add_argument("--mesh", type=int, nargs=3, default=[20, 20, 20])
+ ap.add_argument("--t-max", type=float, default=1000.0)
+ ap.add_argument("--out", default="thermal.html")
+ args = ap.parse_args()
+
+ calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False)
+
+ d = get_jid_data(jid=args.jid, dataset="dft_3d")
+ atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms
+ formula = atoms.composition.reduced_formula
+ print(f"loaded {args.jid} {formula} N={atoms.num_atoms}")
+
+ relaxed = relax(atoms, calc)
+ prim = relaxed.get_primitive_atoms
+ print(f"relaxed → primitive N={prim.num_atoms}")
+
+ ph = run_phonopy(prim, calc, dim=tuple(args.dim))
+ bands = bandstructure(ph, prim)
+ dos_data = dos(ph)
+ thermo = thermal_properties(ph, mesh=tuple(args.mesh), t_max=args.t_max)
+
+ # Summary
+ n_imag = int((bands[1] < -0.05).sum())
+ T = thermo["T"]
+ i300 = int(np.argmin(abs(T - 300)))
+ print(f"\n── Results ────────────────────────────────────────────")
+ print(f"bands : {bands[1].shape[1]} imag modes: {n_imag}")
+ print(f"ZPE : {thermo['ZPE_per_cell_kJmol']:.3f} kJ/mol-cell"
+ f" ({thermo['ZPE_per_atom_eV']:.4f} eV/atom)")
+ print(f"Cv(300 K) : {thermo['Cv'][i300]:.3f} J/K/mol-atom "
+ f"(Dulong-Petit: 24.94)")
+ print(f"S(300 K) : {thermo['S'][i300]:.3f} J/K/mol-atom")
+ print(f"F(300 K) - F(0) : {thermo['F'][i300] - thermo['F'][0]:.3f} kJ/mol-atom")
+
+ fig = plot_all(bands, dos_data, thermo,
+ title=f"{formula} ({args.jid}) — ALIGNN-FF phonons "
+ f"[supercell {args.dim}, mesh {args.mesh}]")
+ fig.write_html(args.out)
+ print(f"\nwrote {args.out}")
+
+ # In a Jupyter cell, uncomment:
+ # fig.show()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/scripts/torch/validate_pair_alignn.py b/alignn/scripts/torch/validate_pair_alignn.py
new file mode 100644
index 0000000..90eb1cd
--- /dev/null
+++ b/alignn/scripts/torch/validate_pair_alignn.py
@@ -0,0 +1,189 @@
+"""Validate pair_alignn against the Python reference (AlignnAtomwiseCalculator).
+
+Does a single-point `run 0` in LAMMPS with pair_alignn, then computes the
+same quantities via the Python calculator, and reports component-wise
+differences. Run this immediately after you get pair_alignn compiling and
+loading the .pt file — it's the gate between "compiles" and "correct."
+
+Pass criteria (typical):
+ |ΔE| < 1e-4 eV
+ max|ΔF| < 1e-3 eV/Å
+ max|Δσ| < 1e-4 eV/ų
+
+Usage:
+ python scripts/torch/validate_pair_alignn.py \\
+ --model-dir OutputDir --ts-model alignn_ff.pt \\
+ --jid JVASP-1002 --supercell 2 2 2
+"""
+from __future__ import annotations
+import argparse, os, tempfile
+import numpy as np
+from ase.io import write as ase_write
+from jarvis.core.atoms import Atoms
+from jarvis.db.figshare import get_jid_data
+
+# Reference: Python ASE calculator path
+from alignn.ff.ff import AlignnAtomwiseCalculator
+
+
+# eV/ų ↔ bar conversion (LAMMPS `metal` units pressure)
+EV_PER_A3_TO_BAR = 1.602176634e6
+
+
+def python_reference(atoms_jarvis, calc):
+ a = atoms_jarvis.ase_converter()
+ a.calc = calc
+ E_raw = a.get_potential_energy()
+ E = float(np.asarray(E_raw).reshape(-1)[0]) # calc may return shape-(1,)
+ F = np.asarray(a.get_forces()) # eV/Å
+ S6 = np.asarray(a.get_stress(voigt=True)).reshape(-1) # eV/ų (Voigt)
+ # Reorder ASE Voigt (xx,yy,zz,yz,xz,xy) -> plain 3x3
+ S = np.array([[S6[0], S6[5], S6[4]],
+ [S6[5], S6[1], S6[3]],
+ [S6[4], S6[3], S6[2]]])
+ return E, F, S
+
+
+def lammps_pair_alignn(ase_atoms, ts_model_path, species_symbols, cutoff=5.0):
+ """Run LAMMPS with pair_alignn at step 0, return E/F/S."""
+ from lammps import lammps
+
+ tmpdir = tempfile.mkdtemp(prefix="pair_alignn_val_")
+ data_path = os.path.join(tmpdir, "sys.data")
+ ase_write(data_path, ase_atoms, format="lammps-data",
+ specorder=list(species_symbols))
+
+ lmp = lammps()
+ lmp.commands_string(f"""
+ units metal
+ atom_style atomic
+ boundary p p p
+ read_data {data_path}
+ {'mass ' + ' '.join(str(i+1)+' '+str(_atomic_mass(sym))
+ for i,sym in enumerate(species_symbols))}
+ pair_style alignn {cutoff}
+ pair_coeff * * {ts_model_path} {' '.join(species_symbols)}
+ neighbor 2.0 bin
+ neigh_modify every 1 delay 0 check yes
+ compute PE all pe
+ compute VIR all pressure NULL virial
+ thermo 1
+ thermo_style custom step pe c_VIR[1] c_VIR[2] c_VIR[3] c_VIR[4] c_VIR[5] c_VIR[6]
+ run 0
+ """)
+
+ nlocal = lmp.get_natoms()
+ # Gather forces by atom ID so ordering matches the ASE/jarvis object
+ F_flat = lmp.gather_atoms("f", 1, 3)
+ F = np.array(F_flat[:]).reshape(nlocal, 3)
+
+ # Energy (metal units → eV already)
+ E = float(lmp.get_thermo("pe"))
+
+ # Pressure tensor in LAMMPS = sum of kinetic + virial. At step 0 with
+ # zero velocities, kinetic term is 0 so p = virial/V. Flip sign to
+ # match ASE stress convention (stress = -1/V * dE/dε).
+ V = lmp.get_thermo("vol") # ų
+ p_bar = np.array([lmp.get_thermo(name) for name in
+ ("c_VIR[1]","c_VIR[2]","c_VIR[3]",
+ "c_VIR[4]","c_VIR[5]","c_VIR[6]")])
+ # LAMMPS virial order: xx yy zz xy xz yz; pressure = virial/V (bar)
+ p_ev_A3 = p_bar / EV_PER_A3_TO_BAR # bar -> eV/ų
+ S = np.array([[ p_ev_A3[0], p_ev_A3[3], p_ev_A3[4]],
+ [ p_ev_A3[3], p_ev_A3[1], p_ev_A3[5]],
+ [ p_ev_A3[4], p_ev_A3[5], p_ev_A3[2]]])
+ # LAMMPS pressure sign convention is opposite of stress
+ S = -S
+ lmp.close()
+ return E, F, S
+
+
+def _atomic_mass(sym):
+ from ase.data import atomic_masses, atomic_numbers
+ return atomic_masses[atomic_numbers[sym]]
+
+
+# ---------------------------------------------------------------------------
+def main():
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--model-dir", required=True, help="dir with best_model.pt + config.json for the Python reference")
+ ap.add_argument("--ts-model", required=True, help="TorchScript .pt exported via export_torchscript.py")
+ ap.add_argument("--jid", default="JVASP-1002")
+ ap.add_argument("--supercell", type=int, nargs=3, default=[2,2,2])
+ ap.add_argument("--perturb", type=float, default=0.03,
+ help="Å, random atomic displacement to break symmetry "
+ "(so forces are non-zero and meaningful)")
+ ap.add_argument("--cutoff", type=float, default=None,
+ help="Å, model cutoff (read from model-dir config if unset)")
+ args = ap.parse_args()
+ if args.cutoff is None:
+ import json as _j
+ args.cutoff = float(_j.load(open(f"{args.model_dir}/config.json"))["cutoff"])
+ print(f"cutoff (from config): {args.cutoff} Å")
+
+ # Build system
+ d = get_jid_data(jid=args.jid, dataset="dft_3d")
+ atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms
+ atoms = atoms.make_supercell(args.supercell)
+ ase_atoms = atoms.ase_converter()
+
+ # Small perturbation — exact equilibrium would give all-zero forces,
+ # which is not a useful validation.
+ rng = np.random.default_rng(0)
+ ase_atoms.positions += args.perturb * rng.standard_normal(ase_atoms.positions.shape)
+
+ species = sorted(set(ase_atoms.get_chemical_symbols()))
+ print(f"System: {atoms.composition.reduced_formula} x{args.supercell} "
+ f"N={len(ase_atoms)} species={species}")
+
+ # Re-wrap into jarvis for the Python reference
+ from jarvis.core.atoms import ase_to_atoms as ase2j
+ atoms_pert = ase2j(ase_atoms)
+
+ # --- Python reference ---
+ print("\n→ Python reference (AlignnAtomwiseCalculator)...")
+ calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False)
+ E_py, F_py, S_py = python_reference(atoms_pert, calc)
+
+ # --- LAMMPS pair_alignn ---
+ print("→ LAMMPS pair_alignn (run 0)...")
+ E_lmp, F_lmp, S_lmp = lammps_pair_alignn(ase_atoms, args.ts_model, species,
+ cutoff=args.cutoff)
+
+ # --- Compare ---
+ dE = abs(E_lmp - E_py)
+ dF = F_lmp - F_py
+ dS = S_lmp - S_py
+ print("\n── Comparison ──────────────────────────────────────────")
+ print(f"E (Python): {E_py: .6f} eV")
+ print(f"E (LAMMPS): {E_lmp: .6f} eV")
+ print(f"|ΔE|: {dE: .3e} eV "
+ f"[{'PASS' if dE < 1e-4 else 'FAIL'}] (threshold 1e-4)")
+
+ max_abs_dF = float(np.max(np.abs(dF)))
+ rms_dF = float(np.sqrt(np.mean(dF**2)))
+ max_abs_F = float(np.max(np.abs(F_py)))
+ print(f"max|ΔF|: {max_abs_dF: .3e} eV/Å "
+ f"[{'PASS' if max_abs_dF < 1e-3 else 'FAIL'}] (threshold 1e-3)")
+ print(f"rms ΔF: {rms_dF: .3e} eV/Å")
+ print(f"max|F_ref|: {max_abs_F: .3e} eV/Å (scale for context)")
+
+ max_abs_dS = float(np.max(np.abs(dS)))
+ print(f"max|Δσ|: {max_abs_dS: .3e} eV/ų "
+ f"[{'PASS' if max_abs_dS < 1e-4 else 'FAIL'}] (threshold 1e-4)")
+
+ print("\n── Per-atom ΔF (eV/Å) — first 5 atoms ──")
+ print("idx ΔF_x ΔF_y ΔF_z |ΔF|")
+ for i in range(min(5, len(F_py))):
+ d = dF[i]
+ print(f"{i:3d} {d[0]:+10.3e} {d[1]:+10.3e} {d[2]:+10.3e} {np.linalg.norm(d):10.3e}")
+
+ # Write out both force arrays for post-hoc inspection
+ np.savez("validate_pair_alignn.npz",
+ E_py=E_py, F_py=F_py, S_py=S_py,
+ E_lmp=E_lmp, F_lmp=F_lmp, S_lmp=S_lmp)
+ print("\nsaved raw arrays -> validate_pair_alignn.npz")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/torch_graph_builder.py b/alignn/torch_graph_builder.py
new file mode 100644
index 0000000..3d16f7a
--- /dev/null
+++ b/alignn/torch_graph_builder.py
@@ -0,0 +1,689 @@
+"""Pure-torch graph and line-graph builder (no DGL).
+
+Produces a crystal neighbor graph together with its atomistic line
+graph as plain torch data, with separate two-body and three-body
+cutoffs. Edge displacement vectors and bond angle cosines are torch
+functions of atomic positions and the lattice, so autograd flows back
+to both (forces via -dE/dx, stress via dE/dL).
+
+Typical use
+-----------
+
+ g, lg = build_pure_torch_graph(
+ atoms=jarvis_atoms,
+ two_body_cutoff=5.0,
+ three_body_cutoff=4.0, # defaults to two_body_cutoff if None
+ max_neighbors=12,
+ )
+ # g.edata["r"] : (E, 3) differentiable displacements
+ # g.edata["images"]: (E, 3) integer cell offsets
+ # g.ndata["atom_features"], g.ndata["frac_coords"], g.ndata["V"]
+ # lg.src / lg.dst : (T,) indices into g's edges (i.e. line-graph
+ # nodes *are* parent edges)
+ # lg.edata["h"] : (T,) bond-angle cosines at the shared atom
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+# ---------------------------------------------------------------------
+# Neighbor-list primitives (moved here from alignn.graphs so importing
+# this module stays DGL-free).
+# ---------------------------------------------------------------------
+
+
+def _torch_periodic_shifts(
+ lattice: torch.Tensor, cutoff: float
+) -> torch.Tensor:
+ """Integer shift vectors (K, 3) covering a ``cutoff`` sphere."""
+ with torch.no_grad():
+ recip = 2 * math.pi * torch.linalg.inv(lattice).T
+ recip_len = torch.linalg.norm(recip, dim=1)
+ n_max = torch.ceil(cutoff * recip_len / (2 * math.pi)).to(torch.long)
+ ranges = [
+ torch.arange(-int(n.item()), int(n.item()) + 1, device=lattice.device)
+ for n in n_max
+ ]
+ return torch.cartesian_prod(*ranges).to(lattice.dtype)
+
+
+def _topk_per_source(
+ src: torch.Tensor,
+ keys: torch.Tensor,
+ max_neighbors: int,
+ num_nodes: int,
+) -> torch.Tensor:
+ """Return indices keeping the K smallest ``keys`` per src node."""
+ device = src.device
+ order_key = torch.argsort(keys)
+ src_k = src[order_key]
+ order_src = torch.argsort(src_k, stable=True)
+ perm = order_key[order_src]
+ src_sorted = src[perm]
+ E = perm.numel()
+ starts = torch.searchsorted(
+ src_sorted, torch.arange(num_nodes, device=device)
+ )
+ within = torch.arange(E, device=device) - starts[src_sorted]
+ return perm[within < max_neighbors]
+
+
+def torch_neighbor_list(
+ positions: torch.Tensor,
+ lattice: torch.Tensor,
+ cutoff: float,
+ max_neighbors: Optional[int] = None,
+ atoms=None,
+ use_matscipy_topology: bool = False,
+ self_tol: float = 1e-8,
+ chunk_size: int = 512,
+):
+ """Torch-native periodic neighbor list with differentiable edges.
+
+ Memory-chunked over source atoms: peak memory is O(K * chunk * N)
+ instead of O(K * N^2), which also sidesteps torch's INT_MAX limit
+ on torch.where for very large boolean tensors.
+ """
+ dtype = positions.dtype
+ device = positions.device
+ num_nodes = int(positions.shape[0])
+
+ used_matscipy = False
+ if use_matscipy_topology and atoms is not None:
+ try:
+ from matscipy.neighbours import neighbour_list as _mnl
+
+ i_np, j_np, S_np = _mnl(
+ "ijS", atoms.ase_converter(), float(cutoff)
+ )
+ src = torch.from_numpy(np.ascontiguousarray(i_np)).to(
+ device=device, dtype=torch.long
+ )
+ dst = torch.from_numpy(np.ascontiguousarray(j_np)).to(
+ device=device, dtype=torch.long
+ )
+ shift = torch.from_numpy(np.ascontiguousarray(S_np)).to(
+ device=device, dtype=dtype
+ )
+ used_matscipy = True
+ except ImportError:
+ pass
+
+ if not used_matscipy:
+ shifts = _torch_periodic_shifts(lattice, cutoff) # (K, 3)
+ with torch.no_grad():
+ offs = shifts @ lattice # (K, 3) cartesian
+ c2 = float(cutoff) * float(cutoff)
+
+ # Dynamically shrink chunk for very large systems so that
+ # (K * chunk * N) bool tensor stays well under INT_MAX.
+ K = int(shifts.shape[0])
+ max_elems = 2**30 # ~1.07e9, safe
+ max_chunk_by_int = max(1, max_elems // max(K * num_nodes, 1))
+ eff_chunk = max(1, min(chunk_size, max_chunk_by_int))
+
+ src_chunks, dst_chunks, shift_chunks = [], [], []
+ for i0 in range(0, num_nodes, eff_chunk):
+ i1 = min(i0 + eff_chunk, num_nodes)
+ # (K, chunk, N, 3)
+ rvec = (
+ positions[None, None, :, :]
+ + offs[:, None, None, :]
+ - positions[None, i0:i1, None, :]
+ )
+ dist2 = rvec.pow(2).sum(-1) # (K, chunk, N)
+ mask = (dist2 <= c2) & (dist2 > self_tol)
+ del rvec, dist2
+ k_idx, i_local, j_idx = torch.where(mask)
+ del mask
+ src_chunks.append((i_local + i0).to(torch.long))
+ dst_chunks.append(j_idx.to(torch.long))
+ shift_chunks.append(shifts[k_idx])
+ del k_idx, i_local, j_idx
+
+ src = (
+ torch.cat(src_chunks)
+ if src_chunks
+ else torch.empty(0, dtype=torch.long, device=device)
+ )
+ dst = (
+ torch.cat(dst_chunks)
+ if dst_chunks
+ else torch.empty(0, dtype=torch.long, device=device)
+ )
+ shift = (
+ torch.cat(shift_chunks)
+ if shift_chunks
+ else torch.empty((0, 3), dtype=dtype, device=device)
+ )
+
+ # Differentiable displacement vectors — this is the autograd bridge
+ r = positions[dst] - positions[src] + shift @ lattice
+
+ if max_neighbors is not None and max_neighbors > 0 and src.numel() > 0:
+ with torch.no_grad():
+ dist = r.norm(dim=1)
+ keep = _topk_per_source(src, dist, int(max_neighbors), num_nodes)
+ src, dst, shift, r = src[keep], dst[keep], shift[keep], r[keep]
+
+ return src, dst, shift, r
+
+
+def torch_neighbor_list_old(
+ positions: torch.Tensor,
+ lattice: torch.Tensor,
+ cutoff: float,
+ max_neighbors: Optional[int] = None,
+ atoms=None,
+ use_matscipy_topology: bool = True,
+ self_tol: float = 1e-8,
+):
+ """Torch-native periodic neighbor list with differentiable edges.
+
+ Topology (src, dst, integer shift) is computed without gradient
+ tracking (matscipy when available, pure-torch fallback otherwise).
+ Displacement vectors are then a torch function of positions and
+ lattice: ``r = positions[dst] - positions[src] + shift @ lattice``.
+ """
+ dtype = positions.dtype
+ device = positions.device
+ num_nodes = positions.shape[0]
+
+ used_matscipy = False
+ if use_matscipy_topology and atoms is not None:
+ try:
+ from matscipy.neighbours import neighbour_list as _mnl
+
+ i_np, j_np, S_np = _mnl(
+ "ijS", atoms.ase_converter(), float(cutoff)
+ )
+ src = torch.from_numpy(np.ascontiguousarray(i_np)).to(
+ device=device, dtype=torch.long
+ )
+ dst = torch.from_numpy(np.ascontiguousarray(j_np)).to(
+ device=device, dtype=torch.long
+ )
+ shift = torch.from_numpy(np.ascontiguousarray(S_np)).to(
+ device=device, dtype=dtype
+ )
+ used_matscipy = True
+ except ImportError:
+ pass
+
+ if not used_matscipy:
+ shifts = _torch_periodic_shifts(lattice, cutoff)
+ with torch.no_grad():
+ offs = shifts @ lattice
+ rvec_full = (
+ positions[None, None, :, :]
+ + offs[:, None, None, :]
+ - positions[None, :, None, :]
+ )
+ dist2 = rvec_full.pow(2).sum(-1)
+ mask = (dist2 <= cutoff * cutoff) & (dist2 > self_tol)
+ k_idx, i_idx, j_idx = torch.where(mask)
+ src = i_idx.to(torch.long)
+ dst = j_idx.to(torch.long)
+ shift = shifts[k_idx]
+
+ r = positions[dst] - positions[src] + shift @ lattice
+
+ if max_neighbors is not None and max_neighbors > 0 and src.numel() > 0:
+ with torch.no_grad():
+ dist = r.norm(dim=1)
+ keep = _topk_per_source(src, dist, int(max_neighbors), num_nodes)
+ src, dst, shift, r = src[keep], dst[keep], shift[keep], r[keep]
+
+ return src, dst, shift, r
+
+
+@dataclass
+class TorchGraph:
+ """Dict-based graph container. Edge list is directed.
+
+ When a graph is the result of batching, ``batch_num_nodes`` /
+ ``batch_num_edges`` carry the per-subgraph counts so downstream
+ code (pooling, per-graph stress, etc.) can segment correctly.
+ """
+
+ num_nodes: int
+ src: torch.Tensor
+ dst: torch.Tensor
+ ndata: Dict[str, torch.Tensor] = field(default_factory=dict)
+ edata: Dict[str, torch.Tensor] = field(default_factory=dict)
+ batch_num_nodes: Optional[torch.Tensor] = None
+ batch_num_edges: Optional[torch.Tensor] = None
+
+ @property
+ def num_edges(self) -> int:
+ """Return the number of directed edges."""
+ return int(self.src.shape[0])
+
+ @property
+ def device(self) -> torch.device:
+ """Return the device of the underlying tensors."""
+ return self.src.device
+
+ @property
+ def batch_size(self) -> int:
+ """Return the number of subgraphs in this (possibly batched) graph."""
+ if self.batch_num_nodes is None:
+ return 1
+ return int(self.batch_num_nodes.shape[0])
+
+ @property
+ def node_batch_id(self) -> torch.Tensor:
+ """(N,) long tensor mapping each node to its subgraph index."""
+ dev = self.src.device
+ if self.batch_num_nodes is None:
+ return torch.zeros(self.num_nodes, dtype=torch.long, device=dev)
+ return torch.repeat_interleave(
+ torch.arange(self.batch_size, device=dev),
+ self.batch_num_nodes,
+ )
+
+ @property
+ def edge_batch_id(self) -> torch.Tensor:
+ """(E,) long tensor mapping each edge to its subgraph index."""
+ dev = self.src.device
+ if self.batch_num_edges is None:
+ return torch.zeros(self.num_edges, dtype=torch.long, device=dev)
+ return torch.repeat_interleave(
+ torch.arange(self.batch_size, device=dev),
+ self.batch_num_edges,
+ )
+
+ def to(self, device) -> "TorchGraph":
+ """Return a copy of this graph moved to ``device``."""
+ return TorchGraph(
+ num_nodes=self.num_nodes,
+ src=self.src.to(device),
+ dst=self.dst.to(device),
+ ndata={k: v.to(device) for k, v in self.ndata.items()},
+ edata={k: v.to(device) for k, v in self.edata.items()},
+ batch_num_nodes=(
+ self.batch_num_nodes.to(device)
+ if self.batch_num_nodes is not None
+ else None
+ ),
+ batch_num_edges=(
+ self.batch_num_edges.to(device)
+ if self.batch_num_edges is not None
+ else None
+ ),
+ )
+
+ def to_dgl(self):
+ """Materialize as a ``dgl.DGLGraph`` preserving ndata / edata.
+
+ Autograd-safe: tensor values are assigned by reference, so
+ gradients on ``edata['r']`` continue to flow back to any leaves
+ it was derived from (e.g. positions / lattice).
+ """
+ import dgl as _dgl
+
+ g = _dgl.graph((self.src, self.dst), num_nodes=self.num_nodes)
+ for k, v in self.ndata.items():
+ g.ndata[k] = v
+ for k, v in self.edata.items():
+ g.edata[k] = v
+ return g
+
+
+def torch_bond_cosines(
+ r_ij: torch.Tensor, r_jk: torch.Tensor, eps: float = 1e-12
+) -> torch.Tensor:
+ """Cosine of the bond angle at atom j for triplets i -> j -> k.
+
+ Matches ALIGNN's convention: negate the first bond so vectors
+ point away from j, i.e. cos = (-r_ij) . r_jk / (|r_ij| |r_jk|).
+ """
+ num = -(r_ij * r_jk).sum(dim=-1)
+ denom = r_ij.norm(dim=-1).clamp_min(eps) * r_jk.norm(dim=-1).clamp_min(eps)
+ return (num / denom).clamp(-1.0, 1.0)
+
+
+def _line_graph_edges(
+ src: torch.Tensor,
+ dst: torch.Tensor,
+ num_nodes: int,
+ allowed: Optional[torch.Tensor] = None,
+):
+ """Build line-graph edges from directed parent edges.
+
+ For each parent edge A = (u, v), emit a line-graph edge (A, B) for
+ every parent edge B = (v, w). When ``allowed`` is given (bool mask
+ of length E), only edges with allowed=True may act as A *or* B.
+
+ Returns
+ -------
+ lg_src, lg_dst : long tensors, indices into the parent edge list.
+ """
+ E = int(src.shape[0])
+ device = src.device
+
+ if allowed is None:
+ allowed_ids = torch.arange(E, device=device)
+ else:
+ allowed_ids = torch.nonzero(allowed, as_tuple=False).squeeze(-1)
+
+ # Sort allowed parent edges by their src node so we can find all
+ # outgoing edges from a node v as a contiguous slice.
+ sub_src = src[allowed_ids]
+ order = torch.argsort(sub_src, stable=True)
+ sorted_edge_ids = allowed_ids[order]
+ sorted_src = sub_src[order]
+
+ node_range = torch.arange(num_nodes, device=device)
+ bucket_start = torch.searchsorted(sorted_src, node_range)
+ bucket_end = torch.searchsorted(sorted_src, node_range, right=True)
+
+ # For each allowed A, count successors = (# allowed edges out of A.dst).
+ A_ids = allowed_ids
+ A_v = dst[A_ids]
+ starts = bucket_start[A_v]
+ ends = bucket_end[A_v]
+ counts = ends - starts
+
+ total = int(counts.sum().item())
+ if total == 0:
+ empty = torch.empty(0, dtype=torch.long, device=device)
+ return empty, empty
+
+ lg_src = torch.repeat_interleave(A_ids, counts)
+
+ # Expand per-A ranges [start_A, start_A + count_A) into flat indices.
+ cum = torch.cumsum(counts, dim=0)
+ row_start = cum - counts
+ offsets = torch.arange(total, device=device) - torch.repeat_interleave(
+ row_start, counts
+ )
+ pos = torch.repeat_interleave(starts, counts) + offsets
+ lg_dst = sorted_edge_ids[pos]
+ return lg_src, lg_dst
+
+
+def build_pure_torch_graph(
+ atoms=None,
+ two_body_cutoff: float = 5.0,
+ three_body_cutoff: Optional[float] = None,
+ max_neighbors: Optional[int] = None,
+ atom_features: str = "cgcnn",
+ use_lattice_prop: bool = False,
+ positions: Optional[torch.Tensor] = None,
+ lattice: Optional[torch.Tensor] = None,
+ device=None,
+ use_matscipy_topology: bool = True,
+ compute_line_graph: bool = True,
+):
+ """Build (g, lg) as pure-torch TorchGraph objects.
+
+ Parameters
+ ----------
+ two_body_cutoff : float
+ Cutoff for the main neighbor graph (edges / pair distances).
+ three_body_cutoff : float or None
+ Cutoff for triplets in the line graph. Defaults to
+ ``two_body_cutoff`` when None. If smaller, only edges with
+ ``|r| <= three_body_cutoff`` are allowed to take part in
+ triplets; the main graph still carries all edges up to
+ ``two_body_cutoff``.
+ max_neighbors : int or None
+ Per-source cap on the pair graph (keeps K closest).
+ positions, lattice : optional torch leaves; when ``requires_grad``
+ is set, gradients flow through ``g.edata['r']`` and
+ ``lg.edata['h']`` back to these.
+
+ Returns
+ -------
+ (g, lg) if compute_line_graph else g.
+ """
+ from jarvis.core.specie import get_node_attributes
+
+ if three_body_cutoff is None:
+ three_body_cutoff = two_body_cutoff
+ if three_body_cutoff > two_body_cutoff:
+ raise ValueError(
+ f"three_body_cutoff ({three_body_cutoff}) must be <= "
+ f"two_body_cutoff ({two_body_cutoff})."
+ )
+
+ dtype = torch.get_default_dtype()
+ if positions is None:
+ positions = torch.as_tensor(np.asarray(atoms.cart_coords), dtype=dtype)
+ if lattice is None:
+ lattice = torch.as_tensor(np.asarray(atoms.lattice_mat), dtype=dtype)
+ if device is not None:
+ positions = positions.to(device)
+ lattice = lattice.to(device)
+ device = positions.device
+ n_atoms = int(positions.shape[0])
+
+ src, dst, shift, r = torch_neighbor_list(
+ positions=positions,
+ lattice=lattice,
+ cutoff=float(two_body_cutoff),
+ max_neighbors=max_neighbors,
+ atoms=atoms,
+ use_matscipy_topology=use_matscipy_topology,
+ )
+
+ sps = np.array(
+ [
+ list(get_node_attributes(s, atom_features=atom_features))
+ for s in atoms.elements
+ ]
+ )
+ node_features = torch.as_tensor(sps, dtype=dtype, device=device)
+ frac = torch.as_tensor(
+ np.asarray(atoms.frac_coords), dtype=dtype, device=device
+ )
+ vol = torch.abs(torch.det(lattice))
+
+ g = TorchGraph(
+ num_nodes=n_atoms,
+ src=src,
+ dst=dst,
+ ndata={
+ "atom_features": node_features,
+ "frac_coords": frac,
+ "V": vol.expand(n_atoms),
+ },
+ edata={"r": r, "images": shift},
+ )
+ if use_lattice_prop:
+ lp = np.array(
+ [atoms.lattice.lat_lengths(), atoms.lattice.lat_angles()]
+ ).flatten()
+ g.ndata["extra_features"] = (
+ torch.as_tensor(lp, dtype=dtype, device=device)
+ .unsqueeze(0)
+ .expand(n_atoms, -1)
+ .contiguous()
+ )
+
+ if not compute_line_graph:
+ return g
+
+ # Three-body cutoff filter: only edges short enough join triplets.
+ if three_body_cutoff < two_body_cutoff:
+ with torch.no_grad():
+ allowed = r.norm(dim=1) <= three_body_cutoff
+ else:
+ allowed = None
+
+ lg_src, lg_dst = _line_graph_edges(src, dst, n_atoms, allowed=allowed)
+ # Angle cosines at the shared atom, differentiable through r.
+ h = torch_bond_cosines(r[lg_src], r[lg_dst])
+
+ lg = TorchGraph(
+ num_nodes=g.num_edges, # line-graph nodes == parent edges
+ src=lg_src,
+ dst=lg_dst,
+ # Share parent edata as ndata (like DGL's shared line graph).
+ ndata={"r": r, "images": shift},
+ edata={"h": h},
+ )
+ return g, lg
+
+
+# =====================================================================
+# Batching and DGL interop
+# =====================================================================
+
+
+def _batch_concat(
+ graphs: List[TorchGraph],
+ node_offset_source: Optional[List[int]] = None,
+) -> TorchGraph:
+ """Concatenate graphs into one, offsetting edge indices.
+
+ ``node_offset_source`` is a list whose cumulative sum gives the
+ offset applied to ``src`` / ``dst`` for each subgraph. When None,
+ uses each subgraph's ``num_nodes`` (standard pair-graph batching).
+ Set it to parent-graph edge counts when batching *line graphs*.
+ """
+ if len(graphs) == 0:
+ raise ValueError("batch_torch_graphs: empty graph list.")
+ device = graphs[0].src.device
+ if node_offset_source is None:
+ node_offset_source = [g.num_nodes for g in graphs]
+
+ offsets = torch.zeros(len(graphs), dtype=torch.long, device=device)
+ if len(graphs) > 1:
+ offsets[1:] = torch.cumsum(
+ torch.tensor(
+ node_offset_source[:-1], dtype=torch.long, device=device
+ ),
+ dim=0,
+ )
+
+ srcs, dsts = [], []
+ for g, off in zip(graphs, offsets):
+ srcs.append(g.src + off)
+ dsts.append(g.dst + off)
+
+ # Only keep keys present in *all* subgraphs (avoid partial batches).
+ ndata_keys = set(graphs[0].ndata)
+ edata_keys = set(graphs[0].edata)
+ for g in graphs[1:]:
+ ndata_keys &= set(g.ndata)
+ edata_keys &= set(g.edata)
+ ndata = {
+ k: torch.cat([g.ndata[k] for g in graphs], dim=0) for k in ndata_keys
+ }
+ edata = {
+ k: torch.cat([g.edata[k] for g in graphs], dim=0) for k in edata_keys
+ }
+
+ return TorchGraph(
+ num_nodes=sum(g.num_nodes for g in graphs),
+ src=torch.cat(srcs),
+ dst=torch.cat(dsts),
+ ndata=ndata,
+ edata=edata,
+ batch_num_nodes=torch.tensor(
+ [g.num_nodes for g in graphs], dtype=torch.long, device=device
+ ),
+ batch_num_edges=torch.tensor(
+ [g.num_edges for g in graphs], dtype=torch.long, device=device
+ ),
+ )
+
+
+def batch_torch_graphs(graphs: List[TorchGraph]) -> TorchGraph:
+ """Batch pair graphs with offsets = cumulative ``num_nodes``."""
+ return _batch_concat(graphs, node_offset_source=None)
+
+
+def batch_torch_graph_pairs(
+ pairs: List,
+) -> "tuple[TorchGraph, TorchGraph]":
+ """Batch a list of ``(g, lg)`` pairs with consistent offsetting.
+
+ The line graph's src/dst index into the *parent* edge list, so its
+ node-offset source is the list of parent num_edges, not num_nodes.
+ """
+ parents = [p[0] for p in pairs]
+ lgs = [p[1] for p in pairs]
+ g_batched = batch_torch_graphs(parents)
+ lg_batched = _batch_concat(
+ lgs, node_offset_source=[p.num_edges for p in parents]
+ )
+ return g_batched, lg_batched
+
+
+# ---- DGL adapter (so the pure-torch model can consume DGL input) ----
+
+
+def unbatch(g):
+ """Polymorphic unbatch — works for TorchGraph *and* DGL graph.
+
+ For a batched ``TorchGraph``, returns a list of single-graph
+ ``TorchGraph`` views with ``src`` / ``dst`` reindexed to the local
+ node range and ndata / edata sliced accordingly. For a DGL graph,
+ delegates to ``dgl.unbatch``.
+ """
+ if isinstance(g, TorchGraph):
+ if g.batch_num_nodes is None:
+ return [g]
+ dev = g.src.device
+ B = g.batch_size
+ node_off = torch.zeros(B + 1, dtype=torch.long, device=dev)
+ node_off[1:] = torch.cumsum(g.batch_num_nodes, dim=0)
+ edge_off = torch.zeros(B + 1, dtype=torch.long, device=dev)
+ edge_off[1:] = torch.cumsum(g.batch_num_edges, dim=0)
+ out = []
+ for b in range(B):
+ n0 = int(node_off[b].item())
+ n1 = int(node_off[b + 1].item())
+ e0 = int(edge_off[b].item())
+ e1 = int(edge_off[b + 1].item())
+ out.append(
+ TorchGraph(
+ num_nodes=n1 - n0,
+ src=g.src[e0:e1] - n0,
+ dst=g.dst[e0:e1] - n0,
+ ndata={k: v[n0:n1] for k, v in g.ndata.items()},
+ edata={k: v[e0:e1] for k, v in g.edata.items()},
+ )
+ )
+ return out
+ import dgl as _dgl
+
+ return _dgl.unbatch(g)
+
+
+def torchgraph_from_dgl(g) -> TorchGraph:
+ """Wrap a ``dgl.DGLGraph`` as a ``TorchGraph`` (zero-copy where possible).
+
+ Batching metadata is preserved when the DGL graph is itself a batch.
+ """
+ src, dst = g.edges()
+ ndata = {k: v for k, v in g.ndata.items()}
+ edata = {k: v for k, v in g.edata.items()}
+ bnn = None
+ bne = None
+ try:
+ bnn_t = g.batch_num_nodes()
+ bne_t = g.batch_num_edges()
+ if bnn_t.numel() > 1:
+ bnn = bnn_t.to(torch.long)
+ bne = bne_t.to(torch.long)
+ except Exception:
+ pass
+ return TorchGraph(
+ num_nodes=int(g.num_nodes()),
+ src=src.to(torch.long),
+ dst=dst.to(torch.long),
+ ndata=ndata,
+ edata=edata,
+ batch_num_nodes=bnn,
+ batch_num_edges=bne,
+ )
diff --git a/alignn/train.py b/alignn/train.py
index 6af1951..7ea303f 100644
--- a/alignn/train.py
+++ b/alignn/train.py
@@ -1,6 +1,7 @@
"""Module for training script."""
from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
from functools import partial
from typing import Any, Dict, Union
import torch
@@ -12,6 +13,8 @@
from alignn.data import get_train_val_loaders
from alignn.config import TrainingConfig
from alignn.models.alignn_atomwise import ALIGNNAtomWise
+from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePure
+from alignn.torch_graph_builder import unbatch as _graph_unbatch
from alignn.models.ealignn_atomwise import eALIGNNAtomWise
from alignn.models.alignn import ALIGNN
from jarvis.db.jsonutils import dumpjson
@@ -22,19 +25,28 @@
import time
from sklearn.metrics import roc_auc_score
from alignn.utils import (
- # activated_output_transform,
- # make_standard_scalar_and_pca,
- # thresholded_output_transform,
group_decay,
setup_optimizer,
print_train_val_loss,
)
-import dgl
-
-# from sklearn.metrics import log_loss
warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+def _ddp_mean(value: float, use_ddp: bool) -> float:
+ """All-reduce a Python scalar across ranks and return the mean."""
+ if not use_ddp or not dist.is_available() or not dist.is_initialized():
+ return float(value)
+ t = torch.tensor(float(value), device=torch.cuda.current_device())
+ dist.all_reduce(t, op=dist.ReduceOp.SUM)
+ return (t / dist.get_world_size()).item()
+
+
+def _unwrap(net):
+ """Return underlying module from a DDP-wrapped model, else net."""
+ return net.module if isinstance(net, DDP) else net
+
+
# torch.autograd.detect_anomaly()
figlet_alignn = """
@@ -49,21 +61,16 @@
def train_dgl(
config: Union[TrainingConfig, Dict[str, Any]],
model: nn.Module = None,
- # checkpoint_dir: Path = Path("./"),
train_val_test_loaders=[],
rank=0,
world_size=0,
- # log_tensorboard: bool = False,
):
"""Training entry point for DGL networks.
`config` should conform to alignn.conf.TrainingConfig, and
if passed as a dict with matching keys, pydantic validation is used
"""
- # print("rank", rank)
- # setup(rank, world_size)
if rank == 0:
- # print(config)
if type(config) is dict:
try:
print("Trying to convert dictionary.")
@@ -74,16 +81,16 @@ def train_dgl(
if not os.path.exists(config.output_dir):
os.makedirs(config.output_dir)
- # checkpoint_dir = os.path.join(config.output_dir)
- # deterministic = False
classification = False
+ is_main = rank == 0
tmp = config.dict()
- f = open(os.path.join(config.output_dir, "config.json"), "w")
- f.write(json.dumps(tmp, indent=4))
- f.close()
+ if is_main:
+ f = open(os.path.join(config.output_dir, "config.json"), "w")
+ f.write(json.dumps(tmp, indent=4))
+ f.close()
global tmp_output_dir
tmp_output_dir = config.output_dir
- pprint.pprint(tmp) # , sort_dicts=False)
+ pprint.pprint(tmp)
if config.classification_threshold is not None:
classification = True
TORCH_DTYPES = {
@@ -95,7 +102,6 @@ def train_dgl(
torch.set_default_dtype(TORCH_DTYPES[config.dtype])
line_graph = False
if config.compute_line_graph > 0:
- # if config.model.alignn_layers > 0:
line_graph = True
if world_size > 1:
use_ddp = True
@@ -105,10 +111,6 @@ def train_dgl(
if torch.cuda.is_available():
device = torch.device("cuda")
if not train_val_test_loaders:
- # use input standardization for all real-valued feature sets
- # print("config.neighbor_strategy",config.neighbor_strategy)
- # import sys
- # sys.exit()
(
train_loader,
val_loader,
@@ -136,6 +138,8 @@ def train_dgl(
filename=config.filename,
cutoff=config.cutoff,
max_neighbors=config.max_neighbors,
+ three_body_cutoff=config.three_body_cutoff,
+ read_existing=config.read_existing,
output_features=config.model.output_features,
classification_threshold=config.classification_threshold,
target_multiplication_factor=config.target_multiplication_factor,
@@ -150,7 +154,6 @@ def train_dgl(
val_loader = train_val_test_loaders[1]
test_loader = train_val_test_loaders[2]
prepare_batch = train_val_test_loaders[3]
- # rank=0
if use_ddp:
device = torch.device(f"cuda:{rank}")
prepare_batch = partial(prepare_batch, device=device)
@@ -158,6 +161,7 @@ def train_dgl(
config.model.classification = True
_model = {
"alignn_atomwise": ALIGNNAtomWise,
+ "alignn_atomwise_pure": ALIGNNAtomWisePure,
"ealignn_atomwise": eALIGNNAtomWise,
"alignn": ALIGNN,
}
@@ -172,11 +176,14 @@ def train_dgl(
xm.set_rng_state(config.random_seed)
except ImportError:
pass
- torch.backends.cudnn.deterministic = True
- torch.backends.cudnn.benchmark = False
os.environ["PYTHONHASHSEED"] = str(config.random_seed)
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = str(":4096:8")
- torch.use_deterministic_algorithms(True)
+ if getattr(config, "deterministic", False):
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = str(":4096:8")
+ torch.use_deterministic_algorithms(True)
+ else:
+ torch.backends.cudnn.benchmark = True
if model is None:
net = _model.get(config.model.name)(config.model)
else:
@@ -201,52 +208,59 @@ def train_dgl(
print(statistics)
except Exception:
pass
- # print("device", device)
net.to(device)
if use_ddp:
- net = DDP(net, device_ids=[rank], find_unused_parameters=True)
- # group parameters to skip weight decay for bias and batchnorm
+ net = DDP(
+ net,
+ device_ids=[rank],
+ find_unused_parameters=bool(
+ getattr(config, "ddp_find_unused_parameters", False)
+ ),
+ )
+ # ------------------------------------------------------------------
+ # Build optimizer ONCE and bind the scheduler to it.
+ # (Previously this block created the optimizer twice — once here and
+ # once again inside the `if "alignn_" in config.model.name:` block —
+ # which left the scheduler bound to a stale optimizer that never
+ # stepped, triggering the
+ # "lr_scheduler.step() before optimizer.step()" warning and producing
+ # a broken LR schedule.)
+ # ------------------------------------------------------------------
params = group_decay(net)
optimizer = setup_optimizer(params, config)
+
if config.scheduler == "none":
# always return multiplier of 1 (i.e. do nothing)
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer, lambda epoch: 1.0
)
-
elif config.scheduler == "onecycle":
steps_per_epoch = len(train_loader)
- # pct_start = config.warmup_steps / (config.epochs * steps_per_epoch)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=config.learning_rate,
epochs=config.epochs,
steps_per_epoch=steps_per_epoch,
- # pct_start=pct_start,
pct_start=0.3,
)
elif config.scheduler == "step":
- # pct_start = config.warmup_steps / (config.epochs * steps_per_epoch)
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
)
- # if (
- # config.model.name == "alignn_atomwise"
- # or config.model.name == "alignn_ff2"
- # ):
+ # OneCycleLR expects scheduler.step() per batch; other schedulers we
+ # step per epoch.
+ _step_scheduler_per_batch = config.scheduler == "onecycle"
+
if "alignn_" in config.model.name:
best_loss = np.inf
criterion = nn.L1Loss()
if classification:
criterion = nn.NLLLoss()
- params = group_decay(net)
- optimizer = setup_optimizer(params, config)
- # optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
+ # NOTE: optimizer / scheduler intentionally NOT recreated here.
history_train = []
history_val = []
for e in range(config.epochs):
- # optimizer.zero_grad()
train_init_time = time.time()
running_loss = 0
running_loss1 = 0
@@ -257,21 +271,24 @@ def train_dgl(
train_result = []
for dats, jid in zip(train_loader, train_loader.dataset.ids):
info = {}
- # info["id"] = jid
optimizer.zero_grad()
- if (config.compute_line_graph) > 0:
- # if (config.model.alignn_layers) > 0:
- result = net(
- [
- dats[0].to(device),
- dats[1].to(device),
- dats[2].to(device),
- ]
- )
-
- else:
- result = net([dats[0].to(device), dats[1].to(device)])
- # info = {}
+ _amp_ctx = torch.autocast(
+ device_type="cuda",
+ dtype=torch.bfloat16,
+ enabled=bool(getattr(config, "use_amp", False))
+ and torch.cuda.is_available(),
+ )
+ with _amp_ctx:
+ if (config.compute_line_graph) > 0:
+ result = net(
+ [
+ dats[0].to(device),
+ dats[1].to(device),
+ dats[2].to(device),
+ ]
+ )
+ else:
+ result = net([dats[0].to(device), dats[1].to(device)])
info["target_out"] = []
info["pred_out"] = []
info["target_atomwise_pred"] = []
@@ -289,23 +306,17 @@ def train_dgl(
loss4 = 0 # Such as stresses
loss5 = 0 # Such as dos
if config.model.output_features is not None:
- # print('criterion',criterion)
- # print('result["out"]',result["out"])
- # print('dats[-1]',dats[-1])
loss1 = config.model.graphwise_weight * criterion(
result["out"],
dats[-1].to(device),
- # result["out"], dats[2].to(device)
)
info["target_out"] = dats[-1].cpu().numpy().tolist()
- # info["target_out"] = dats[2].cpu().numpy().tolist()
info["pred_out"] = (
result["out"].cpu().detach().numpy().tolist()
)
running_loss1 += loss1.item()
if (
config.model.atomwise_output_features > 0
- # config.model.atomwise_output_features is not None
and config.model.atomwise_weight != 0
):
loss2 = config.model.atomwise_weight * criterion(
@@ -333,69 +344,65 @@ def train_dgl(
)
running_loss3 += loss3.item()
if config.model.stresswise_weight != 0:
- # print('unbatch',dgl.unbatch(dats[0]))
-
targ_stress = torch.stack(
[
gg.ndata["stresses"][0]
- for gg in dgl.unbatch(dats[0])
+ for gg in _graph_unbatch(dats[0])
]
).to(device)
pred_stress = result["stresses"]
- # print('targ_stress',targ_stress,targ_stress.shape)
- # print('pred_stress',pred_stress,pred_stress.shape)
loss4 = config.model.stresswise_weight * criterion(
pred_stress.to(device),
targ_stress.to(device),
)
- info["target_stress"] = (
- targ_stress.cpu()
- .numpy()
- .tolist()
- # dats[0].ndata["stresses"][0].cpu().numpy().tolist()
- )
+ info["target_stress"] = targ_stress.cpu().numpy().tolist()
info["pred_stress"] = (
result["stresses"].cpu().detach().numpy().tolist()
)
running_loss4 += loss4.item()
if config.model.additional_output_weight != 0:
- # print('unbatch',dgl.unbatch(dats[0]))
additional_dat = [
gg.ndata["additional"][0]
- for gg in dgl.unbatch(dats[0])
+ for gg in _graph_unbatch(dats[0])
]
- # print('additional_dat',additional_dat,len(additional_dat))
targ = torch.stack(additional_dat).to(device)
- # targ=torch.tensor(additional_dat).to( dats[0].device)
- # print('result["additional"]',result["additional"],result["additional"].shape)
- # print('targ',targ,targ.shape)
- # print('targ device',targ.device)
loss5 = config.model.additional_output_weight * criterion(
(result["additional"]).to(device),
targ,
- # (dats[0].ndata["additional"]).to(device),
)
info["target_additional"] = targ.cpu().numpy().tolist()
info["pred_additional"] = (
result["additional"].cpu().detach().numpy().tolist()
)
running_loss5 += loss5.item()
- # print("target_stress", info["target_stress"][0])
- # print("pred_stress", info["pred_stress"][0])
train_result.append(info)
loss = loss1 + loss2 + loss3 + loss4 + loss5
loss.backward()
optimizer.step()
- # optimizer.zero_grad() #never
+ # Step OneCycleLR per batch (its design assumption); other
+ # schedulers are stepped once per epoch below.
+ if _step_scheduler_per_batch:
+ scheduler.step()
running_loss += loss.item()
- # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors(
- # train_result
- # )
- # dumpjson(filename="Train_results.json", data=train_result)
- scheduler.step()
+ # Normalize running losses by number of batches
+ _n_tr = max(1, len(train_loader))
+ running_loss /= _n_tr
+ running_loss1 /= _n_tr
+ running_loss2 /= _n_tr
+ running_loss3 /= _n_tr
+ running_loss4 /= _n_tr
+ running_loss5 /= _n_tr
+ running_loss = _ddp_mean(running_loss, use_ddp)
+ running_loss1 = _ddp_mean(running_loss1, use_ddp)
+ running_loss2 = _ddp_mean(running_loss2, use_ddp)
+ running_loss3 = _ddp_mean(running_loss3, use_ddp)
+ running_loss4 = _ddp_mean(running_loss4, use_ddp)
+ running_loss5 = _ddp_mean(running_loss5, use_ddp)
+ # Epoch-level scheduler step for non-OneCycle schedulers.
+ if not _step_scheduler_per_batch:
+ scheduler.step()
train_final_time = time.time()
train_ep_time = train_final_time - train_init_time
- # if rank == 0: # or world_size == 1:
history_train.append(
[
running_loss,
@@ -406,10 +413,13 @@ def train_dgl(
running_loss5,
]
)
- dumpjson(
- filename=os.path.join(config.output_dir, "history_train.json"),
- data=history_train,
- )
+ if is_main:
+ dumpjson(
+ filename=os.path.join(
+ config.output_dir, "history_train.json"
+ ),
+ data=history_train,
+ )
val_loss = 0
val_loss1 = 0
val_loss2 = 0
@@ -417,15 +427,11 @@ def train_dgl(
val_loss4 = 0
val_loss5 = 0
val_result = []
- # for dats in val_loader:
val_init_time = time.time()
for dats, jid in zip(val_loader, val_loader.dataset.ids):
info = {}
info["id"] = jid
optimizer.zero_grad()
- # result = net([dats[0].to(device), dats[1].to(device)])
- # if (config.model.alignn_layers) > 0:
- # if (config.create_line_graph) > 0:
if (config.compute_line_graph) > 0:
result = net(
[
@@ -436,8 +442,6 @@ def train_dgl(
)
else:
result = net([dats[0].to(device), dats[1].to(device)])
- # result = net(dats[0].to(device))
- # info = {}
info["target_out"] = []
info["pred_out"] = []
info["target_atomwise_pred"] = []
@@ -446,11 +450,11 @@ def train_dgl(
info["pred_grad"] = []
info["target_stress"] = []
info["pred_stress"] = []
- loss1 = 0 # Such as energy
- loss2 = 0 # Such as bader charges
- loss3 = 0 # Such as forces
- loss4 = 0 # Such as stresses
- loss5 = 0 # Such as stresses
+ loss1 = 0
+ loss2 = 0
+ loss3 = 0
+ loss4 = 0
+ loss5 = 0
if config.model.output_features is not None:
loss1 = config.model.graphwise_weight * criterion(
result["out"], dats[-1].to(device)
@@ -489,30 +493,18 @@ def train_dgl(
)
val_loss3 += loss3.item()
if config.model.stresswise_weight != 0:
- # loss4 = config.model.stresswise_weight * criterion(
- # result["stress"].to(device),
- # dats[0].ndata["stresses"][0].to(device),
- # )
-
targ_stress = torch.stack(
[
gg.ndata["stresses"][0]
- for gg in dgl.unbatch(dats[0])
+ for gg in _graph_unbatch(dats[0])
]
).to(device)
pred_stress = result["stresses"]
- # print('targ_stress',targ_stress,targ_stress.shape)
- # print('pred_stress',pred_stress,pred_stress.shape)
loss4 = config.model.stresswise_weight * criterion(
pred_stress.to(device),
targ_stress.to(device),
)
- info["target_stress"] = (
- targ_stress.cpu()
- .numpy()
- .tolist()
- # dats[0].ndata["stresses"][0].cpu().numpy().tolist()
- )
+ info["target_stress"] = targ_stress.cpu().numpy().tolist()
info["pred_stress"] = (
result["stresses"].cpu().detach().numpy().tolist()
)
@@ -521,18 +513,12 @@ def train_dgl(
if config.model.additional_output_weight != 0:
additional_dat = [
gg.ndata["additional"][0]
- for gg in dgl.unbatch(dats[0])
+ for gg in _graph_unbatch(dats[0])
]
- # print('additional_dat',additional_dat,len(additional_dat))
targ = torch.stack(additional_dat).to(device)
- # targ=torch.tensor(additional_dat).to( dats[0].device)
- # print('result["additional"]',result["additional"],result["additional"].shape)
- # print('targ',targ,targ.shape)
- # print('targ device',targ.device)
loss5 = config.model.additional_output_weight * criterion(
(result["additional"]).to(device),
targ,
- # (dats[0].ndata["additional"]).to(device),
)
info["target_additional"] = targ.cpu().numpy().tolist()
info["pred_additional"] = (
@@ -543,38 +529,49 @@ def train_dgl(
loss = loss1 + loss2 + loss3 + loss4 + loss5
val_result.append(info)
val_loss += loss.item()
- # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors(
- # val_result
- # )
+ _n_vl = max(1, len(val_loader))
+ val_loss /= _n_vl
+ val_loss1 /= _n_vl
+ val_loss2 /= _n_vl
+ val_loss3 /= _n_vl
+ val_loss4 /= _n_vl
+ val_loss5 /= _n_vl
+ val_loss = _ddp_mean(val_loss, use_ddp)
+ val_loss1 = _ddp_mean(val_loss1, use_ddp)
+ val_loss2 = _ddp_mean(val_loss2, use_ddp)
+ val_loss3 = _ddp_mean(val_loss3, use_ddp)
+ val_loss4 = _ddp_mean(val_loss4, use_ddp)
+ val_loss5 = _ddp_mean(val_loss5, use_ddp)
val_fin_time = time.time()
val_ep_time = val_fin_time - val_init_time
current_model_name = "current_model.pt"
- torch.save(
- net.state_dict(),
- os.path.join(config.output_dir, current_model_name),
- )
+ if is_main:
+ torch.save(
+ _unwrap(net).state_dict(),
+ os.path.join(config.output_dir, current_model_name),
+ )
saving_msg = ""
if val_loss < best_loss:
best_loss = val_loss
best_model_name = "best_model.pt"
- torch.save(
- net.state_dict(),
- os.path.join(config.output_dir, best_model_name),
- )
- # print("Saving data for epoch:", e)
- saving_msg = "Saving model"
- dumpjson(
- filename=os.path.join(
- config.output_dir, "Train_results.json"
- ),
- data=train_result,
- )
- dumpjson(
- filename=os.path.join(
- config.output_dir, "Val_results.json"
- ),
- data=val_result,
- )
+ if is_main:
+ torch.save(
+ _unwrap(net).state_dict(),
+ os.path.join(config.output_dir, best_model_name),
+ )
+ saving_msg = "Saving model"
+ dumpjson(
+ filename=os.path.join(
+ config.output_dir, "Train_results.json"
+ ),
+ data=train_result,
+ )
+ dumpjson(
+ filename=os.path.join(
+ config.output_dir, "Val_results.json"
+ ),
+ data=val_result,
+ )
best_model = net
history_val.append(
[
@@ -586,11 +583,13 @@ def train_dgl(
val_loss5,
]
)
- # history_val.append([mean_out, mean_atom, mean_grad, mean_stress])
- dumpjson(
- filename=os.path.join(config.output_dir, "history_val.json"),
- data=history_val,
- )
+ if is_main:
+ dumpjson(
+ filename=os.path.join(
+ config.output_dir, "history_val.json"
+ ),
+ data=history_val,
+ )
if rank == 0:
print_train_val_loss(
e,
@@ -615,14 +614,10 @@ def train_dgl(
test_loss = 0
test_result = []
for dats, jid in zip(test_loader, test_loader.dataset.ids):
- # for dats in test_loader:
info = {}
info["id"] = jid
optimizer.zero_grad()
- # if (config.model.alignn_layers) > 0:
- # if (config.create_line_graph) > 0:
if (config.compute_line_graph) > 0:
- # result = net([dats[0].to(device), dats[1].to(device)])
result = net(
[
dats[0].to(device),
@@ -632,17 +627,14 @@ def train_dgl(
)
else:
result = net([dats[0].to(device), dats[1].to(device)])
- # result = net(dats[0].to(device))
- loss1 = 0 # Such as energy
- loss2 = 0 # Such as bader charges
- loss3 = 0 # Such as forces
- loss4 = 0 # Such as stresses
+ loss1 = 0
+ loss2 = 0
+ loss3 = 0
+ loss4 = 0
if (
config.model.output_features is not None
and not classification
):
- # print('result["out"]',result["out"])
- # print('dats[2]',dats[2])
loss1 = config.model.graphwise_weight * criterion(
result["out"], dats[-1].to(device)
)
@@ -679,22 +671,15 @@ def train_dgl(
targ_stress = torch.stack(
[
gg.ndata["stresses"][0]
- for gg in dgl.unbatch(dats[0])
+ for gg in _graph_unbatch(dats[0])
]
).to(device)
pred_stress = result["stresses"]
- # print('targ_stress',targ_stress,targ_stress.shape)
- # print('pred_stress',pred_stress,pred_stress.shape)
loss4 = config.model.stresswise_weight * criterion(
pred_stress.to(device),
targ_stress.to(device),
)
- info["target_stress"] = (
- targ_stress.cpu()
- .numpy()
- .tolist()
- # dats[0].ndata["stresses"][0].cpu().numpy().tolist()
- )
+ info["target_stress"] = targ_stress.cpu().numpy().tolist()
info["pred_stress"] = (
result["stresses"].cpu().detach().numpy().tolist()
)
@@ -703,21 +688,22 @@ def train_dgl(
loss = loss1 + loss2 + loss3 + loss4
if not classification:
test_loss += loss.item()
- print("TestLoss", e, test_loss)
- dumpjson(
- filename=os.path.join(config.output_dir, "Test_results.json"),
- data=test_result,
- )
- last_model_name = "last_model.pt"
- torch.save(
- net.state_dict(),
- os.path.join(config.output_dir, last_model_name),
- )
- # return test_result
+ if is_main:
+ print("TestLoss", e, test_loss)
+ dumpjson(
+ filename=os.path.join(
+ config.output_dir, "Test_results.json"
+ ),
+ data=test_result,
+ )
+ last_model_name = "last_model.pt"
+ torch.save(
+ _unwrap(net).state_dict(),
+ os.path.join(config.output_dir, last_model_name),
+ )
if rank == 0 or world_size == 1:
if config.write_predictions and classification:
best_model.eval()
- # net.eval()
f = open(
os.path.join(
config.output_dir, "prediction_results_test_set.csv"
@@ -728,16 +714,12 @@ def train_dgl(
targets = []
predictions = []
with torch.no_grad():
- ids = test_loader.dataset.ids # [test_loader.dataset.indices]
+ ids = test_loader.dataset.ids
for dat, id in zip(test_loader, ids):
g, lg, lat, target = dat
out_data = best_model(
[g.to(device), lg.to(device), lat.to(device)]
)["out"]
- # out_data = net([g.to(device), lg.to(device)])["out"]
- # out_data = torch.exp(out_data.cpu())
- # print('target',target)
- # print('out_data',out_data)
top_p, top_class = torch.topk(torch.exp(out_data), k=1)
target = int(target.cpu().numpy().flatten().tolist()[0])
@@ -761,22 +743,20 @@ def train_dgl(
and config.model.output_features > 1
):
best_model.eval()
- # net.eval()
mem = []
with torch.no_grad():
- ids = test_loader.dataset.ids # [test_loader.dataset.indices]
+ ids = test_loader.dataset.ids
for dat, id in zip(test_loader, ids):
g, lg, lat, target = dat
out_data = best_model(
[g.to(device), lg.to(device), lat.to(device)]
)["out"]
- # out_data = net([g.to(device), lg.to(device)])["out"]
out_data = out_data.detach().cpu().numpy().tolist()
if config.standard_scalar_and_pca:
sc = pk.load(open("sc.pkl", "rb"))
out_data = list(
sc.transform(np.array(out_data).reshape(1, -1))[0]
- ) # [0][0]
+ )
target = target.cpu().numpy().flatten().tolist()
info = {}
info["id"] = id
@@ -796,7 +776,6 @@ def train_dgl(
and config.model.gradwise_weight == 0
):
best_model.eval()
- # net.eval()
f = open(
os.path.join(
config.output_dir, "prediction_results_test_set.csv"
@@ -807,13 +786,12 @@ def train_dgl(
targets = []
predictions = []
with torch.no_grad():
- ids = test_loader.dataset.ids # [test_loader.dataset.indices]
+ ids = test_loader.dataset.ids
for dat, id in zip(test_loader, ids):
g, lg, lat, target = dat
out_data = best_model(
[g.to(device), lg.to(device), lat.to(device)]
)["out"]
- # out_data = net([g.to(device), lg.to(device)])["out"]
out_data = out_data.cpu().numpy().tolist()
if config.standard_scalar_and_pca:
sc = pk.load(
@@ -825,7 +803,18 @@ def train_dgl(
target = target.cpu().numpy().flatten().tolist()
if len(target) == 1:
target = target[0]
- f.write("%s, %6f, %6f\n" % (id, target, out_data))
+ # out_data is a (possibly nested) list after .tolist();
+ # unwrap single-output regression to a scalar, mirroring
+ # the target handling above. Multi-output falls back to
+ # a generic format so the write never raises.
+ if isinstance(out_data, list):
+ out_data = np.array(out_data).flatten().tolist()
+ if len(out_data) == 1:
+ out_data = out_data[0]
+ if isinstance(target, list) or isinstance(out_data, list):
+ f.write("%s, %s, %s\n" % (id, target, out_data))
+ else:
+ f.write("%s, %6f, %6f\n" % (id, target, out_data))
targets.append(target)
predictions.append(out_data)
f.close()
@@ -835,7 +824,6 @@ def train_dgl(
mean_absolute_error(np.array(targets), np.array(predictions)),
)
best_model.eval()
- # net.eval()
f = open(
os.path.join(
config.output_dir, "prediction_results_train_set.csv"
@@ -846,13 +834,12 @@ def train_dgl(
targets = []
predictions = []
with torch.no_grad():
- ids = train_loader.dataset.ids # [test_loader.dataset.indices]
+ ids = train_loader.dataset.ids
for dat, id in zip(train_loader, ids):
g, lg, lat, target = dat
out_data = best_model(
[g.to(device), lg.to(device), lat.to(device)]
)["out"]
- # out_data = net([g.to(device), lg.to(device)])["out"]
out_data = out_data.cpu().numpy().tolist()
if config.standard_scalar_and_pca:
sc = pk.load(
@@ -862,10 +849,6 @@ def train_dgl(
np.array(out_data).reshape(-1, 1)
)[0][0]
target = target.cpu().numpy().flatten().tolist()
- # if len(target) == 1:
- # target = target[0]
- # if len(out_data) == 1:
- # out_data = out_data[0]
for ii, jj in zip(target, out_data):
f.write("%6f, %6f\n" % (ii, jj))
targets.append(ii)
diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py
index 95060bd..5891e8a 100644
--- a/alignn/train_alignn.py
+++ b/alignn/train_alignn.py
@@ -1,6 +1,7 @@
#!/usr/bin/env python
"""Module to train for a folder with formatted dataset."""
+
import os
import torch.distributed as dist
import csv
@@ -13,6 +14,10 @@
from jarvis.db.jsonutils import loadjson
import argparse
from alignn.models.alignn_atomwise import ALIGNNAtomWise, ALIGNNAtomWiseConfig
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure,
+ ALIGNNAtomWisePureConfig,
+)
import torch
import time
from jarvis.core.atoms import Atoms
@@ -192,6 +197,12 @@ def train_for_folder(
except Exception as exp:
print("Check", exp)
+ if config.gpu_memory_fraction is not None and torch.cuda.is_available():
+ torch.cuda.set_per_process_memory_fraction(
+ float(config.gpu_memory_fraction),
+ rank if world_size > 1 else 0,
+ )
+
# config.keep_data_order = keep_data_order
if classification_threshold is not None:
config.classification_threshold = float(classification_threshold)
@@ -336,15 +347,18 @@ def train_for_folder(
if restart_model_path is not None:
# Should be best_model.pt file
print("Restarting the model training:", restart_model_path)
- if config.model.name == "alignn_atomwise":
+ if "alignn_" in config.model.name: # == "alignn_atomwise":
rest_config = loadjson(
restart_model_path.replace("current_model.pt", "config.json")
# restart_model_path.replace("best_model.pt", "config.json")
)
-
- tmp = ALIGNNAtomWiseConfig(**rest_config["model"])
+ try:
+ tmp = ALIGNNAtomWiseConfig(**rest_config["model"])
+ model = ALIGNNAtomWise(tmp) # config.model)
+ except Exception:
+ tmp = ALIGNNAtomWisePureConfig(**rest_config["model"])
+ model = ALIGNNAtomWisePure(tmp) # config.model)
print("Rest config", tmp)
- model = ALIGNNAtomWise(tmp) # config.model)
print("model", model)
model.load_state_dict(
torch.load(restart_model_path, map_location=device)
@@ -397,6 +411,7 @@ def train_for_folder(
cutoff=config.cutoff,
cutoff_extra=config.cutoff_extra,
max_neighbors=config.max_neighbors,
+ three_body_cutoff=config.three_body_cutoff,
output_features=config.model.output_features,
classification_threshold=config.classification_threshold,
target_multiplication_factor=config.target_multiplication_factor,
@@ -404,6 +419,8 @@ def train_for_folder(
keep_data_order=config.keep_data_order,
output_dir=config.output_dir,
use_lmdb=config.use_lmdb,
+ use_pure_torch=config.model.name == "alignn_atomwise_pure",
+ read_existing=config.read_existing,
dtype=config.dtype,
)
# print("dataset", dataset[0])
diff --git a/docs/examples-colab.md b/docs/examples-colab.md
new file mode 100644
index 0000000..185b246
--- /dev/null
+++ b/docs/examples-colab.md
@@ -0,0 +1,18 @@
+# Google Colab Notebooks
+
+Ready-to-run Colab notebooks covering installation, property prediction, force-field
+training, and pretrained-model usage. Click the badge in any row to open the notebook in
+Colab.
+
+[colab-badge]: https://colab.research.google.com/assets/colab-badge.svg
+
+| Notebook | Open in Colab | Description |
+| --- | --- | --- |
+| [Regression task (graph-wise prediction)](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Single-output regression model for exfoliation energies of 2D materials. |
+| [ML force-field training from scratch](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Train an ALIGNN-FF machine-learning force-field for Silicon. |
+| [ALIGNN-FF: relaxation, EV curve, phonons, interfaces](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Using the pretrained ALIGNN-FF for structure relaxation, EV curves, phonons, interface gamma surfaces, and interface separation. |
+| [Scaling / timing comparison](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | Analyze scaling/timing of universal MLFFs. |
+| [Melt-Quench MD](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | Generate amorphous structures via molecular dynamics. |
+| [Miscellaneous training tasks](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | Single-output (formation energy, bandgap), multi-output (phonon/electron DOS), classification (metal vs non-metal), and pretrained-model usage. |
+| [Superconductor Tc](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | Train a model for superconductor transition temperature. |
+| [Build `id_prop.json` from VASP runs](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | Compile `vasprun.xml` files into an `id_prop.json` for ALIGNN-FF training. |
diff --git a/docs/index.md b/docs/index.md
index 933776b..c1bd20c 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -28,6 +28,7 @@ dataset (~75,000 materials and 4M+ energy/force entries) and supports any combin
## Quick links
- [Install ALIGNN](installation.md)
+- [Colab notebooks](examples-colab.md)
- [Train your first model](training/single-output-regression.md)
- [Use a pretrained model](pretrained/index.md)
- [Run the ALIGNN-FF ASE calculator](usage/ase-calculator.md)
diff --git a/docs/reference/config-reference.md b/docs/reference/config-reference.md
new file mode 100644
index 0000000..a41a7c4
--- /dev/null
+++ b/docs/reference/config-reference.md
@@ -0,0 +1,335 @@
+# Config Reference
+
+Every training run in ALIGNN is controlled by a single JSON file. The schema is
+enforced by two pydantic classes: `TrainingConfig` (top-level) in
+[`alignn/config.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/config.py)
+and a per-model config (e.g. `ALIGNNAtomWiseConfig`) in
+[`alignn/models/alignn_atomwise.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/models/alignn_atomwise.py).
+
+At runtime the config is validated and any unknown field is rejected, so this
+file is the authoritative list of what's available.
+
+---
+
+## Top-level (`TrainingConfig`)
+
+### Dataset & target
+
+| Field | Type / default | Meaning |
+|---|---|---|
+| `dataset` | one of `dft_3d`, `megnet`, `qm9`, `user_data`, … (see enum) / `"dft_3d"` | Which dataset. Use `"user_data"` for your own `POSCAR`s. |
+| `target` | string / `"exfoliation_energy"` | Column in `id_prop.csv` to predict. |
+| `id_tag` | `"jid"` \| `"id"` \| `"_oqmd_entry_id"` / `"jid"` | Which column holds the ID. |
+| `classification_threshold` | `float?` / `None` | If set, regression targets above it become class 1; enables classification mode. |
+| `target_multiplication_factor` | `float?` / `None` | Scale all targets by this constant (useful when units don't match loss scale). |
+
+### Atoms → graph
+
+| Field | Default | Meaning |
+|---|---|---|
+| `atom_features` | `"cgcnn"` | Node feature scheme. `cgcnn` = 92-d one-hot-ish table, `atomic_number` = scalar Z, `basic` = 11 hand-picked features, `cfid` = 438 chemical descriptors. Must match `model.atom_input_features`. |
+| `neighbor_strategy` | `"k-nearest"` | How to build the neighbor graph. See [Neighbor strategies](#neighbor-strategies). |
+| `cutoff` | `8.0` | Radius cutoff (Å) for pair edges. |
+| `cutoff_extra` | `3.0` | Fallback slack for `radius_graph` if too few neighbors are found. |
+| `three_body_cutoff` | `None` | Separate cutoff (Å) for line-graph / angle triplets. Only honored by `neighbor_strategy="pure_torch"`. `None` → use `cutoff`. Must be ≤ `cutoff`. |
+| `max_neighbors` | `12` | Per-source cap. |
+| `use_canonize` | `True` | Canonicalize k-nearest edge order (reproducibility). |
+| `compute_line_graph` | `True` | Build the bond-angle line graph. |
+
+#### Neighbor strategies
+
+| Value | Backend | Autograd through `r`? | Notes |
+|---|---|---|---|
+| `k-nearest` (default) | Jarvis | no | Fastest for small property-prediction cells. |
+| `radius_graph` | pure torch | no | Honest cutoff; slower. |
+| `radius_graph_jarvis` | Jarvis | no | Uses Jarvis's `NeighborsAnalysis`. |
+| `fast_graph` | matscipy (C) | no | Same topology as `radius_graph` but ~200× faster on large cells. Requires `pip install matscipy`. |
+| `torch_graph` | matscipy + torch | **yes** | Recomputes `r = positions[dst] − positions[src] + shift @ lattice` in torch → gradients flow to positions and lattice. |
+| `pure_torch` | matscipy + torch | **yes** | Like `torch_graph`, plus respects `three_body_cutoff` for the line graph. |
+
+### Splitting
+
+| Field | Default | Meaning |
+|---|---|---|
+| `train_ratio` / `val_ratio` / `test_ratio` | `0.8 / 0.1 / 0.1` | Fractions (used when counts are not set). |
+| `n_train` / `n_val` / `n_test` | `None` | Absolute counts; override the ratios when set. |
+| `keep_data_order` | `True` | If `True`, no shuffling before the split — makes comparisons across runs reproducible. |
+| `random_seed` | `123` | Seed for Python / numpy / torch / CUDA RNGs. |
+
+### Training loop
+
+| Field | Default | Meaning |
+|---|---|---|
+| `epochs` | `300` | Number of training passes. |
+| `batch_size` | `64` | Per-step batch size. |
+| `learning_rate` | `1e-2` | Base LR. |
+| `weight_decay` | `0` | AdamW / SGD weight decay. |
+| `warmup_steps` | `2000` | OneCycle warmup steps (ignored if `scheduler="none"`). |
+| `criterion` | `"mse"` | Regression loss. `"l1"` is what the force-field examples use. `"zig"` / `"poisson"` are specialized. |
+| `optimizer` | `"adamw"` | `"adamw"` or `"sgd"`. |
+| `scheduler` | `"onecycle"` | `"onecycle"` or `"none"`. |
+| `n_early_stopping` | `None` | Stop if val loss hasn't improved in this many epochs (~50 is common). |
+
+### Graph caching
+
+| Field | Default | Meaning |
+|---|---|---|
+| `use_lmdb` | `True` | Cache built graphs in LMDB (fast reload, memory-bounded). |
+| `read_existing` | `False` | If `True`, reuse an existing LMDB cache on disk. **Must match the current model backend** (DGL vs pure-torch). Default `False` rebuilds from scratch every run to avoid footguns. |
+
+### Runtime & distributed
+
+| Field | Default | Meaning |
+|---|---|---|
+| `dtype` | `"float32"` | `"float32"` or `"float64"`. |
+| `num_workers` | `4` | `DataLoader` workers. `0` for debugging. |
+| `pin_memory` | `False` | Pinned host memory for faster H2D copies on GPU. |
+| `distributed` | `False` | Enable DDP multi-process training (paired with `torchrun`). |
+| `data_parallel` | `False` | Legacy single-process multi-GPU. Prefer `distributed`. |
+
+### Logging & outputs
+
+| Field | Default | Meaning |
+|---|---|---|
+| `output_dir` | `.` | Where checkpoints / logs / predictions are written. |
+| `filename` | `"sample"` | Prefix for cache dirs (`train_data`, etc.). |
+| `write_checkpoint` | `True` | Save `best_model.pt` / `current_model.pt`. |
+| `write_predictions` | `True` | Dump per-sample predictions at the end. |
+| `store_outputs` | `True` | Keep intermediate outputs in the history JSONs. |
+| `progress` | `True` | `tqdm` progress bars. |
+| `log_tensorboard` | `False` | Also write a TensorBoard run. |
+| `save_dataloader` | `False` | Pickle the built DataLoader objects (large; rarely useful). |
+| `standard_scalar_and_pca` | `False` | Z-score + PCA the targets (legacy). |
+| `normalize_graph_level_loss` | `False` | Divide the graph-level loss by number of atoms (enable for extensive quantities mixed with intensive ones). |
+
+### `model`
+
+A nested object whose schema is chosen by `model.name`. Currently available:
+
+| `name` | Class | Purpose |
+|---|---|---|
+| `alignn` | `ALIGNNConfig` | Original ALIGNN (scalar regression only). |
+| `alignn_atomwise` | `ALIGNNAtomWiseConfig` | Adds atomwise / gradwise / stresswise outputs — this is what `force-field.md` uses. |
+| `alignn_atomwise_pure` | `ALIGNNAtomWisePureConfig` | **DGL-free**, scatter-based reimplementation of `alignn_atomwise`. Accepts the same config fields; TorchScript-exportable for LAMMPS. See [Pure-torch notes](#pure-torch-notes). |
+| `ealignn_atomwise` | `eALIGNNAtomWiseConfig` | Equivariant variant. |
+
+See [Model config (`alignn_atomwise`)](#model-config-alignn_atomwise) below.
+
+---
+
+## Model config (`alignn_atomwise`)
+
+All fields below apply to both `alignn_atomwise` and `alignn_atomwise_pure`
+unless noted.
+
+### Architecture
+
+| Field | Default | Meaning |
+|---|---|---|
+| `name` | required | `"alignn_atomwise"` or `"alignn_atomwise_pure"`. |
+| `alignn_layers` | `2` | Number of ALIGNN (node + edge + triplet) conv blocks. |
+| `gcn_layers` | `2` | Number of subsequent node + edge conv blocks (no triplets). |
+| `atom_input_features` | `1` | Width of incoming atom features. **Must match** `atom_features` (92 for `cgcnn`, 1 for `atomic_number`, 11 for `basic`, 438 for `cfid`). |
+| `edge_input_features` | `80` | Number of RBF bins for bond lengths. |
+| `triplet_input_features` | `40` | Number of RBF bins for bond-angle cosines. |
+| `embedding_features` | `64` | Width of the two-layer MLP inside the edge/angle embeddings. |
+| `hidden_features` | `64` | Core channel width running through all conv blocks. |
+| `output_features` | `1` | Graph-level output dimension (>1 for multi-output regression). |
+
+### Output heads and loss weights
+
+| Field | Default | Meaning |
+|---|---|---|
+| `calculate_gradient` | `True` | Compute per-atom forces via `∂E/∂r`. |
+| `atomwise_output_features` | `0` | Per-atom prediction dim (e.g. charges, magmoms). `0` disables. |
+| `additional_output_features` | `0` | Extra per-graph head dim. `0` disables. |
+| `graphwise_weight` | `1.0` | Loss weight on the graph-level target. |
+| `gradwise_weight` | `1.0` | Loss weight on forces. Setting to `0` automatically disables `calculate_gradient`. |
+| `stresswise_weight` | `0.0` | Loss weight on the 3×3 stress tensor. |
+| `atomwise_weight` | `0.0` | Loss weight on per-atom outputs (requires `atomwise_output_features > 0`). |
+| `additional_output_weight` | `0.0` | Loss weight on the additional head. |
+
+### Force-field specifics
+
+| Field | Default | Meaning |
+|---|---|---|
+| `grad_multiplier` | `-1` | Multiplier applied to `∂E/∂r` when converting to forces. Physics convention `F = −∂E/∂x` ⇒ `-1`. |
+| `add_reverse_forces` | `True` | Sum `∂E/∂r` over both `(i→j)` **and** `(j→i)` edges (proper Newton's third law on directed graphs). |
+| `lg_on_fly` | `True` | Recompute bond-angle cosines inside the autograd graph every forward (needed for exact force/stress derivatives). Leave `True` unless profiling. |
+| `batch_stress` | `True` | Compute stress per graph in the batch (vs. one stress per whole batch). |
+| `force_mult_natoms` | `False` | Scale forces by number of atoms (rare; use for unit conventions). |
+| `energy_mult_natoms` | `True` | Treat the graph-head output as energy *per atom* and multiply by `natoms` before taking the gradient. Standard for ALIGNN-FF. |
+| `stress_multiplier` | `1.0` | Final scalar on the stress output (for unit fixes). |
+| `include_pos_deriv` | `False` | Alternative force path: differentiate through Cartesian positions directly (rather than through edge displacements). Currently experimental; the `_pure` variant accepts the field but ignores it. |
+
+### Cutoff shaping
+
+| Field | Default | Meaning |
+|---|---|---|
+| `use_cutoff_function` | `False` | Apply a smooth polynomial envelope to distances before RBF. Useful when `cutoff` is tight. |
+| `multiply_cutoff` | `False` | If `use_cutoff_function=True`, multiply the edge embedding by the envelope (instead of the bondlength). |
+| `inner_cutoff` | `3.0` | Envelope radius (Å); envelope = 1 below this, smoothly → 0 up to `cutoff`. |
+| `exponent` | `5` | Exponent of the smooth-cutoff polynomial. |
+
+### Short-distance penalty
+
+| Field | Default | Meaning |
+|---|---|---|
+| `use_penalty` | `True` | Add a repulsive penalty when any bondlength falls below `penalty_threshold`. Prevents the model from hallucinating attractive wells at unphysical distances — important for MD stability. |
+| `penalty_factor` | `0.1` | Strength. |
+| `penalty_threshold` | `1.0` | Distance (Å) below which the penalty kicks in. |
+
+### Miscellaneous
+
+| Field | Default | Meaning |
+|---|---|---|
+| `link` | `"identity"` | Output activation. `"log"` applies `exp()` (positive outputs), `"logit"` applies `sigmoid()`. |
+| `classification` | `False` | Two-class NLL head (set automatically when `classification_threshold` is non-null at the training level). |
+| `zero_inflated` | `False` | Zero-inflated Gaussian output (specialized). |
+| `extra_features` | `0` | If > 0, concatenate this many pre-computed fingerprint features per atom (e.g. Gong et al.). |
+
+---
+
+## Pure-torch notes
+
+`model.name = "alignn_atomwise_pure"` swaps the DGL message-passing core for
+`scatter_add` / `scatter_mean` primitives. The **forward outputs are numerically
+equivalent** to `alignn_atomwise` — parity sweep across all doc examples shows
+`|Δ| ≤ 0.04` on all losses after 3 epochs (force-field scenario is bit-identical,
+see `alignn/scripts/parity_dgl_vs_pure.py`). Residual drift is `scatter_add`
+nondeterminism on CUDA, not a modelling difference.
+
+Not carried over from `alignn_atomwise` (config fields accepted but ignored):
+
+- `include_pos_deriv`
+- Some `extra_features` code paths — experimental in the DGL model too.
+
+When the pure model is selected, the training pipeline automatically:
+
+- Uses `alignn/pure_lmdb_dataset.py` (pickles `TorchGraph` objects, no DGL).
+- Uses `PureTorchLMDBDataset.collate_line_graph` for batching.
+- Leaves dataloader / loss / metrics / checkpoints untouched.
+
+For LAMMPS integration, call
+[`alignn/scripts/torch/export_torchscript.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/scripts/torch/export_torchscript.py) (installed as the `export_torchscript.py` console script) — it scripts the
+`forward_tensors_z(positions, lattice, atomic_numbers, src, dst, shift,
+compute_stress)` entry point and bakes the atomic-number → feature lookup into
+the `.pt` so the C++ host only needs atomic numbers.
+
+---
+
+## Minimal examples
+
+### Single-output regression
+
+```json
+{
+ "dataset": "user_data",
+ "target": "formation_energy_peratom",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "k-nearest",
+ "cutoff": 8.0,
+ "max_neighbors": 12,
+ "epochs": 100,
+ "batch_size": 32,
+ "learning_rate": 1e-3,
+ "criterion": "mse",
+ "model": {
+ "name": "alignn_atomwise",
+ "atom_input_features": 92,
+ "alignn_layers": 4,
+ "gcn_layers": 4,
+ "hidden_features": 256,
+ "output_features": 1,
+ "calculate_gradient": false,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 0.0,
+ "stresswise_weight": 0.0,
+ "use_penalty": false
+ }
+}
+```
+
+### ALIGNN-FF (energy + forces + stress)
+
+```json
+{
+ "dataset": "user_data",
+ "target": "target",
+ "atom_features": "cgcnn",
+ "neighbor_strategy": "fast_graph",
+ "cutoff": 6.0,
+ "max_neighbors": 12,
+ "epochs": 200,
+ "batch_size": 2,
+ "learning_rate": 1e-3,
+ "criterion": "l1",
+ "model": {
+ "name": "alignn_atomwise",
+ "atom_input_features": 92,
+ "alignn_layers": 3,
+ "gcn_layers": 3,
+ "hidden_features": 128,
+ "output_features": 1,
+ "calculate_gradient": true,
+ "graphwise_weight": 0.85,
+ "gradwise_weight": 0.05,
+ "stresswise_weight": 0.05,
+ "add_reverse_forces": true,
+ "lg_on_fly": true,
+ "use_penalty": true
+ }
+}
+```
+
+### Pure-torch FF with separate 3-body cutoff
+
+```json
+{
+ "neighbor_strategy": "pure_torch",
+ "cutoff": 6.0,
+ "three_body_cutoff": 4.0,
+ "max_neighbors": 12,
+ "model": { "name": "alignn_atomwise_pure", "…": "…" }
+}
+```
+
+### Classification
+
+```json
+{
+ "classification_threshold": 0.01,
+ "target": "is_metal",
+ "criterion": "mse",
+ "model": {
+ "name": "alignn_atomwise",
+ "classification": true,
+ "graphwise_weight": 1.0,
+ "gradwise_weight": 0.0,
+ "calculate_gradient": false
+ }
+}
+```
+
+---
+
+## Common gotchas
+
+- **`atom_input_features` vs `atom_features`.** These must match (cgcnn → 92,
+ atomic_number → 1, basic → 11, cfid → 438). A mismatch raises a shape error
+ at the first forward pass, not at config parse.
+- **`read_existing`.** Keep `false` unless you're sure the cache matches the
+ current model backend + graph settings. Backend-mismatch (DGL vs pure) is
+ caught and raises a clear error; parameter drift (different cutoff etc.) is
+ not caught.
+- **`gradwise_weight=0`** silently disables `calculate_gradient` in the model
+ code. If you think forces should be training but see `Grad=0.0000`, check
+ both flags.
+- **`use_penalty=true` with a loose cutoff** can inflate training loss when
+ your data legitimately has short bonds (e.g. H-containing systems). Lower
+ `penalty_threshold` or set `use_penalty=false` in that case.
+- **`batch_stress=false` + large batch** gives one stress for the whole
+ batch, not per-graph — almost never what you want.
+- **`lg_on_fly=false`** caches triplet angle cosines; saves time but forces
+ computed by autograd will be missing the angle-contribution unless you
+ recompute them. Leave `true` for force-field training.
diff --git a/docs/usage/lammps-pair-style.md b/docs/usage/lammps-pair-style.md
new file mode 100644
index 0000000..127a111
--- /dev/null
+++ b/docs/usage/lammps-pair-style.md
@@ -0,0 +1,391 @@
+# LAMMPS native pair style (`pair_alignn`)
+
+`pair_alignn` is a C++ LAMMPS pair style that runs a TorchScript-exported
+ALIGNN-FF model as the force engine. It replaces the Python+ASE bridge
+(`alignn.ff.lammps_bridge`) for production runs:
+
+| Path | Overhead per step | When to use |
+|---|---|---|
+| ASE calculator (Python) | high (~40 ms for 200 atoms) | one-off scripts, debugging |
+| `lammps_bridge_torch.py` | medium (~15 ms) | quick Python-hosted MD |
+| **`pair_alignn` (this page)** | **low (~3–10 ms)** | production MD, large systems |
+
+Everything stays on the GPU; Python is never invoked inside the MD step
+loop.
+
+!!! note "Status"
+ MVP single-rank implementation. Matches the TorchScript Pure model
+ forward to ~10⁻⁶ eV on Si. Multi-rank / domain decomposition is not
+ yet implemented — cross-rank ghost contributions are skipped.
+
+## What you need
+
+1. A trained ALIGNN-FF model directory containing `best_model.pt` + `config.json`
+2. LAMMPS built against **libtorch with the same CUDA/ABI flags** as your
+ Python PyTorch install
+3. `pair_alignn.cpp` / `pair_alignn.h` compiled into that LAMMPS binary
+4. A TorchScript `.pt` of your model exported via `export_torchscript.py`
+
+## Building LAMMPS with `pair_alignn` — from a fresh clone
+
+The repo ships `alignn/scripts/torch/build_lammps_alignn.sh`, a
+self-contained script that handles the full build, including prerequisite
+installs.
+
+### Step 1 — clone and install alignn
+
+```bash
+git clone https://github.com/atomgptlab/alignn.git
+cd alignn
+pip install -e . # installs train_alignn.py, export_torchscript.py, ...
+```
+
+Works in both environments:
+- **Conda env**: the build uses `mamba` to install CUDA + MKL when missing.
+- **System Python (e.g. Google Colab)**: the build falls back to `pip install` for
+ `nvidia-cuda-*-cu12 mkl mkl-devel mkl-include`, and detects CUDA at
+ `/usr/local/cuda` automatically.
+
+### Step 2 — run the build script
+
+```bash
+bash alignn/scripts/torch/build_lammps_alignn.sh
+```
+
+That's it. The script:
+
+1. Discovers the active env's PyTorch + CUDA + libtorch ABI flag
+2. **Auto-installs matching `cuda-toolkit` and `mkl-devel`** via
+ `mamba install -c nvidia -c conda-forge` if they're missing or the
+ version doesn't match libtorch (set `SKIP_CUDA_MKL=1` to disable)
+3. Clones LAMMPS `stable_29Aug2024_update1` to `$LAMMPS_DIR`
+ (default `~/lammps-alignn`)
+4. Drops `pair_alignn.cpp/.h` into LAMMPS's `src/` (picked up by the
+ default source glob — no package-system gymnastics)
+5. Appends a `find_package(Torch)` + `target_link_libraries(... TORCH_LIBRARIES)`
+ block to `cmake/CMakeLists.txt` (idempotent via marker comment)
+6. Configures with Ninja, compiles with all available cores
+7. Runs `make install-python` to wire the libtorch-linked build into your
+ conda env's `lammps` Python module
+8. **Overwrites `$CONDA_PREFIX/lib/liblammps.so.0`** with the new build
+ (the standalone `lmp` binary's RPATH picks the conda lib first, so
+ without this step the binary keeps using a stale .so)
+9. Prints ✓/✗ verification for both the Python module and the standalone
+ binary
+
+Takes ~5–15 minutes depending on machine. Successful output ends with:
+```
+✓ Python module has pair_alignn
+✓ Standalone binary has pair_alignn: ~/lammps-alignn/build/lmp
+```
+
+**Env-var knobs:**
+
+| Variable | Default | Purpose |
+|---|---|---|
+| `LAMMPS_DIR` | `$HOME/lammps-alignn` | where to clone + build LAMMPS |
+| `LAMMPS_TAG` | `stable_29Aug2024_update1` | LAMMPS release to clone |
+| `SKIP_CUDA_MKL` | unset | set to `1` to skip auto-install of CUDA/MKL |
+| `JOBS` | `$(nproc)` | parallel compile jobs |
+
+Example overrides:
+```bash
+LAMMPS_DIR=/opt/lammps-alignn JOBS=16 \
+ bash alignn/scripts/torch/build_lammps_alignn.sh
+```
+
+### Step 3 — verify (script does this automatically; manual fallback)
+
+```bash
+# Standalone binary
+$LAMMPS_DIR/build/lmp -help | grep alignn
+# → alignn (as a registered pair style)
+
+# Python module (what most scripts / validators use)
+python -c "from lammps import lammps; l=lammps(); print('alignn' in l.available_styles('pair'))"
+# → True
+```
+
+### Rebuild after editing `pair_alignn.cpp`
+
+If you modify the C++ source in the alignn repo, just re-run the build
+script — the clone step is skipped automatically if `$LAMMPS_DIR` already
+exists, and everything else runs from a clean `build/`:
+
+```bash
+bash alignn/scripts/torch/build_lammps_alignn.sh
+```
+
+Or do it manually, skipping the full reconfigure:
+
+```bash
+cp alignn/scripts/torch/pair_alignn/pair_alignn.{cpp,h} $LAMMPS_DIR/src/
+cd $LAMMPS_DIR/build
+cmake --build . -j$(nproc)
+cmake --build . --target install-python
+cp liblammps.so.0 $CONDA_PREFIX/lib/liblammps.so.0 # keep standalone binary in sync
+```
+
+### Common build issues
+
+These are all handled automatically by the build script, but listed here
+for context if something goes wrong:
+
+| Symptom | Cause | What the script does (and manual fix if needed) |
+|---|---|---|
+| `PyTorch: CUDA cannot be found` | CUDA toolkit version mismatches libtorch | auto-installs matching `cuda-toolkit=` via mamba |
+| `Imported target "torch" includes non-existent path ".../MKL_INCLUDE_DIR-NOTFOUND"` | MKL not installed | auto-installs `mkl-devel mkl-include` |
+| `fatal error: cuda.h: No such file or directory` | CUDA headers not on include path | adds `-I$CONDA_PREFIX/targets/x86_64-linux/include` to CXX flags |
+| Python module works but `lmp` binary fails with "Unrecognized pair style" | RPATH resolves conda's stale `liblammps.so.0` first | `cp liblammps.so.0 $CONDA_PREFIX/lib/` — done in step 8 |
+| `Could NOT find CUDA` on a CPU-only machine | libtorch was built with CUDA | install a CPU-only libtorch separately, or use a matching CPU-only PyTorch in the env |
+
+## Exporting the TorchScript model
+
+After installing the repo with `pip install -e .` (or `pip install alignn`),
+the `export_torchscript.py` console script is on your `PATH`:
+
+```bash
+export_torchscript.py \
+ --model-dir MELT/OutputDir \
+ --out alignn_ff.pt \
+ --atom-features atomic_number \
+ --dtype float32
+```
+
+(Equivalent to `python alignn/scripts/torch/export_torchscript.py ...` if
+you prefer running the source file directly during development.)
+
+Output:
+```
+smoke test OK: {'energy': torch.Size([]), 'forces': torch.Size([2, 3]), 'stress': torch.Size([3, 3])}
+wrote alignn_ff.pt (143,297 params)
+```
+
+**`--atom-features` must match what was used at training time** — if the
+model was trained with CGCNN features (92-dim), pass `--atom-features
+cgcnn`. The bundled species table is baked into the `.pt` so the C++ side
+never needs `jarvis`.
+
+## Input script syntax
+
+```
+pair_style alignn [max_neighbors]
+pair_coeff * * [ ...]
+```
+
+| Argument | Meaning | Example |
+|---|---|---|
+| `` | Neighbor cutoff in Å. **Must match training** | `5.0` |
+| `[max_neighbors]` | k-nearest cap per atom. Default 12 | `12` |
+| `` | Path to the TorchScript file | `alignn_ff.pt` |
+| `` | Element symbol for LAMMPS atom type N | `Si` |
+
+Example for Si trained with cutoff 5.0 Å, 12 neighbors:
+```
+pair_style alignn 5.0 12
+pair_coeff * * alignn_ff.pt Si
+```
+
+Look up the correct values from your trained model:
+```bash
+python -c "
+import json
+c = json.load(open('OutputDir/config.json'))
+print('cutoff:', c['cutoff'], 'Å')
+print('max_neighbors:', c['max_neighbors'])
+print('atom_features:', c.get('atom_features'))
+"
+```
+
+## A complete melt-quench input
+
+```
+units metal # eV, Å, ps, bar, K
+atom_style atomic
+boundary p p p
+read_data si.data
+mass 1 28.0855
+
+pair_style alignn 5.0 12
+pair_coeff * * alignn_ff.pt Si
+
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+
+timestep 0.001 # 1 fs
+thermo 50
+thermo_style custom step temp pe ke etotal press vol
+dump traj all custom 100 si.lammpstrj id type x y z
+
+# ── Stage 1: equilibrate at 300 K ──
+velocity all create 300.0 12345 mom yes rot yes
+fix nvt1 all nvt temp 300.0 300.0 0.1
+run 200
+unfix nvt1
+
+# ── Stage 2: heat to 3500 K over 1 ps ──
+fix nvt2 all nvt temp 300.0 3500.0 0.1
+run 1000
+unfix nvt2
+
+# ── Stage 3: hold molten ──
+fix nvt3 all nvt temp 3500.0 3500.0 0.1
+run 2000
+unfix nvt3
+
+# ── Stage 4: quench to 300 K over 4 ps ──
+fix nvt4 all nvt temp 3500.0 300.0 0.1
+run 4000
+
+write_data si_quenched.data
+```
+
+Making the Si data file:
+```python
+from jarvis.db.figshare import get_jid_data
+from jarvis.core.atoms import Atoms
+from ase.io import write
+a = Atoms.from_dict(get_jid_data(jid='JVASP-1002', dataset='dft_3d')['atoms'])
+ase = a.make_supercell([3, 3, 3]).ase_converter() # 216 atoms
+write("si.data", ase, format="lammps-data", specorder=["Si"])
+```
+
+## Running
+
+Three equivalent ways:
+
+=== "Custom LAMMPS binary"
+
+ ```bash
+ ~/lammps-alignn/build/lmp -in si_native.in
+ ```
+
+=== "Python-driven"
+
+ ```python
+ from lammps import lammps
+ l = lammps()
+ l.file("si_native.in")
+ ```
+
+=== "In a Jupyter cell"
+
+ ```python
+ !~/lammps-alignn/build/lmp -in si_native.in
+ ```
+
+## Validating against the Python calculator
+
+Before trusting a production run, compare LAMMPS forces against the
+Python reference:
+
+```bash
+python alignn/scripts/torch/validate_pair_alignn.py \
+ --model-dir MELT/OutputDir \
+ --ts-model alignn_ff.pt \
+ --jid JVASP-1002 \
+ --supercell 2 2 2
+```
+
+Expected output shape:
+```
+E (Python): -346.003876 eV
+E (LAMMPS): -345.xxxxxx eV
+|ΔE|: < 1e-4 eV [PASS]
+max|ΔF|: < 1e-3 eV/Å [PASS]
+```
+
+If energies agree to ~1e-6 but forces are off, you have a gradient /
+autograd issue. If everything is off by a consistent factor, it's a unit
+or normalization mismatch. If the C++ energy is off by an O(1) factor vs
+a single-precision residual, check cutoff and atom_features.
+
+## Required LAMMPS settings
+
+- `newton_pair on` — LAMMPS default; the pair style requires it for force
+ reduction on ghost atoms
+- atom IDs must be enabled (default)
+- `atom->map()` must be built — pair_alignn auto-initializes it
+
+## What pair_alignn does per step
+
+1. Copy local atom positions + types to a GPU tensor (`(N_local, 3)`)
+2. Walk LAMMPS's full+ghost neighbor list. For each local `i`, collect all
+ neighbors within cutoff; for ghosts, look up the owner via
+ `atom->map(tag)` and compute the PBC shift vector
+3. Keep only the `max_neighbors` closest per atom (matches jarvis's
+ `neighbor_strategy="k-nearest"` / `"pure_torch"` behavior)
+4. Pack `(src, dst, shift)` tensors and call the TorchScript model's
+ `forward_tensors_z(positions, lattice, atomic_numbers, src, dst, shift,
+ compute_stress)` entry point
+5. Read back `energy` (scalar), `forces` (N×3), and `stress` (3×3)
+6. Scatter forces to `f[i]` for local atoms; accumulate energy and virial
+7. LAMMPS's reverse comm handles any ghost-force reductions (no-op on a
+ single rank)
+
+## Model-side requirements
+
+Your trained model must:
+
+1. Use `name: "alignn_atomwise_pure"` in its config (so
+ `forward_tensors_z` exists). If you trained with `alignn_atomwise`
+ (DGL), retrain or convert weights — the two variants are currently
+ NOT numerically equivalent on this trained checkpoint (known issue).
+2. Be TorchScript-scriptable — the pure variant is designed for this
+ (`@torch.jit.ignore` on DGL-specific methods, `@torch.jit.export` on
+ tensor-only forwards)
+3. Have `calculate_gradient: True` — forces come from autograd
+
+## Known limitations
+
+- **Single MPI rank only.** Cross-rank ghost atoms get skipped silently.
+ For multi-rank, the compute() needs a proper ghost-force reverse_comm
+ handshake that mirrors pair_allegro's pattern.
+- **Orthorhombic box assumption for PBC shifts.** Triclinic boxes work
+ approximately but may produce wrong shift vectors near cell corners.
+ TODO: use the full 3×3 inverse lattice for shift reconstruction.
+- **Pure model currently diverges from DGL model** on the same trained
+ weights by ~10% energy. Root cause: `forward_tensors_z`'s inline line
+ graph construction enumerates ~9× more bond-angle triplets than DGL's
+ `g.line_graph()`. To match DGL exactly, either retrain with the pure
+ model directly or patch the line-graph filter.
+- **fp32 only** in the C++ path. Changing `dtype_` in `pair_alignn.h` to
+ `torch::kFloat64` also requires a matching fp64 TorchScript export and
+ changed `accessor` types.
+- **No `eflag_atom` / `vflag_atom`** (per-atom energy/stress) — only
+ global quantities.
+
+## Speed reference
+
+On a single NVIDIA RTX-class GPU, 216-atom Si at cutoff 5.0 Å:
+
+| Path | ms / step |
+|---|---|
+| AlignnAtomwiseCalculator (Python via ASE) | ~40 |
+| lammps_bridge_torch (Python callback) | ~15 |
+| **pair_alignn (native C++)** | **~3–5** |
+
+Bottleneck in the C++ path is the forward pass itself, not the neighbor
+list building. Memory stays resident on the GPU between steps.
+
+## Files in the repo
+
+```
+alignn/scripts/torch/
+├── export_torchscript.py # Export .pt TorchScript model
+├── build_lammps_alignn.sh # Full LAMMPS+libtorch build
+├── validate_pair_alignn.py # Python↔LAMMPS force/energy check
+├── lammps_bridge_torch.py # Alternative Python bridge (no C++)
+└── pair_alignn/
+ ├── pair_alignn.cpp # Pair style source
+ ├── pair_alignn.h # Pair style header
+ ├── si_native.in # Example LAMMPS melt-quench
+ └── CMakeLists-snippet.txt # Standalone build fragment
+```
+
+## Related
+
+- [MD & Relaxation (pure PyTorch)](md-integrators.md) — on-device
+ integrators that skip LAMMPS entirely
+- [ASE Calculator](ase-calculator.md) — the slower but more flexible
+ Python path
diff --git a/docs/usage/md-integrators.md b/docs/usage/md-integrators.md
new file mode 100644
index 0000000..a88a50b
--- /dev/null
+++ b/docs/usage/md-integrators.md
@@ -0,0 +1,267 @@
+# Pure-PyTorch MD & Relaxation
+
+`alignn.md` provides on-device molecular dynamics integrators and a FIRE
+relaxer that pair with the pure-PyTorch ALIGNN-FF model
+(`ALIGNNAtomWisePure`). Everything runs on GPU without bouncing through
+ASE/NumPy each step, which is the prerequisite for scaling to large
+systems and (eventually) multi-GPU.
+
+## When to use this vs. the ASE calculator
+
+| Use | Pick |
+|---|---|
+| Small systems, quick scripts, trajectory files, ensembles you trust (NPT, surfaces, constraints, free-energy tools) | ASE calculator (`AlignnAtomwiseCalculator`) |
+| Large systems where per-step CPU→GPU transfer dominates | `alignn.md` integrators |
+| Low-precision (bf16/fp16) or multi-GPU MD | `alignn.md` integrators |
+| Differentiable MD / gradients through trajectories | `alignn.md` integrators |
+
+The two paths share the same model code; you can mix them (e.g.,
+equilibrate with ASE, production with `alignn.md`).
+
+## Components
+
+```python
+from alignn.md import (
+ AlignnForces, # wraps a pure-PyTorch model as forces_fn
+ VelocityVerlet, # NVE
+ Langevin, # NVT — stochastic, BAOAB splitting
+ NVTBerendsen, # weak-coupling (equilibration only)
+ NVTBussi, # canonical NVT — stochastic rescale
+ NVTNoseHooverChain, # canonical NVT — deterministic chain
+ FIRE, # structural relaxation
+ run, maxwell_boltzmann,
+)
+```
+
+All integrators share a `step(positions, velocities) -> (positions, velocities)`
+API. Positions are `(N, 3)` Å, velocities `(N, 3)` Å/fs, masses `(N,)` amu,
+time `fs`, energy `eV`, temperature `K` — ASE's unit convention.
+
+## Forces function
+
+Every integrator takes a `forces_fn(positions) -> (energy, forces)`.
+`AlignnForces` wraps the pure-PyTorch model:
+
+```python
+import numpy as np, torch
+from ase.build import bulk
+from ase.data import atomic_masses
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig,
+)
+from alignn.md import AlignnForces
+
+a = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((3, 3, 3))
+Z = a.get_atomic_numbers()
+cell = np.array(a.cell)
+
+model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(
+ name="alignn_atomwise_pure",
+ calculate_gradient=True,
+ atomwise_output_features=0,
+ atom_input_features=92,
+))
+# load trained weights if you have them:
+# model.load_state_dict(torch.load("OutputDir/best_model.pt"), strict=False)
+
+forces_fn = AlignnForces(model, Z, cell, device="cuda", dtype=torch.float32)
+```
+
+!!! warning "Neighbor-list bottleneck"
+ The current `AlignnForces` rebuilds the neighbor/line graph on CPU
+ every step. For large systems this dominates wall time. An on-device
+ neighbor list is planned.
+
+## NVE — `VelocityVerlet`
+
+Symplectic, time-reversible, second-order. Energy conservation is the
+standard sanity check.
+
+```python
+from alignn.md import VelocityVerlet, maxwell_boltzmann, run
+
+masses = torch.tensor([atomic_masses[z] for z in Z],
+ dtype=torch.float32, device="cuda")
+pos = torch.tensor(a.get_positions(), dtype=torch.float32, device="cuda")
+vel = maxwell_boltzmann(masses, T=300.0)
+
+integ = VelocityVerlet(forces_fn=forces_fn, masses=masses, dt=1.0)
+pos, vel, hist = run(integ, pos, vel, nsteps=1000, log_every=100)
+```
+
+## NVT — four options
+
+| Thermostat | Ensemble | Dynamics preserved? | Best for |
+|---|---|---|---|
+| `NVTBerendsen` | **not canonical** (crushes KE fluctuations) | Yes (deterministic) | Equilibration, melt-quench |
+| `Langevin` | Canonical | No (adds friction + noise) | General NVT |
+| `NVTBussi` | Canonical | Nearly — velocity-direction preserved | General NVT; fast relaxation |
+| `NVTNoseHooverChain` | Canonical | Yes (time-reversible) | Transport coefficients, spectra |
+
+### Berendsen
+
+Fast, stable, but **samples a distribution with wrong variance** — safe
+only for equilibration. Exposes `.set_temperature(T)` for ramping.
+
+```python
+from alignn.md import NVTBerendsen
+integ = NVTBerendsen(forces_fn=forces_fn, masses=masses,
+ dt=1.0, T=300.0, taut=20.0)
+```
+
+### Langevin (BAOAB splitting)
+
+Solves $\ddot{x} = F/m - \gamma \dot{x} + \sqrt{2\gamma k_B T / m}\,\eta$
+with the Leimkuhler-Matthews BAOAB scheme.
+
+```python
+from alignn.md import Langevin
+integ = Langevin(forces_fn=forces_fn, masses=masses,
+ dt=1.0, T=300.0, friction=0.01) # 1/fs
+```
+
+Friction 0.01 fs⁻¹ = 100 fs damping time. Use larger γ for faster
+equilibration, smaller γ to perturb dynamics less.
+
+### Bussi-Donadio-Parrinello
+
+Stochastic rescale (Bussi et al., JCP 126, 014101, 2007). Same `taut`
+interface as Berendsen but samples the true canonical distribution. **Use
+this as your default NVT** unless you specifically need deterministic
+dynamics.
+
+```python
+from alignn.md import NVTBussi
+integ = NVTBussi(forces_fn=forces_fn, masses=masses,
+ dt=1.0, T=300.0, taut=20.0)
+```
+
+### Nosé-Hoover chain
+
+Martyna-Tobias-Klein integrator. Deterministic and time-reversible —
+preferred when measuring transport coefficients (diffusion, viscosity,
+VACF) where stochastic thermostats perturb the relevant dynamics.
+Default chain length 3 (never use `chain_length=1` — it's
+non-ergodic).
+
+```python
+from alignn.md import NVTNoseHooverChain
+integ = NVTNoseHooverChain(forces_fn=forces_fn, masses=masses,
+ dt=1.0, T=300.0, taut=20.0,
+ chain_length=3)
+```
+
+### Verifying "true NVT"
+
+Canonical KE should fluctuate with
+$\sigma(KE)/\langle KE \rangle = \sqrt{2/(3N - 3)}$. Measured on a 108-atom
+Cu system at 300 K (random-init model, isolates thermostat behavior):
+
+| Thermostat | ⟨KE⟩ eV | σ(KE) / σ_canonical |
+|---|---|---|
+| Berendsen | 4.143 ✓ | 0.04 (crushed) |
+| Bussi | 4.173 ✓ | 1.13 ✓ |
+| NH-chain | 4.171 ✓ | 0.52 (correlated, converges with longer runs) |
+
+See `md_test_bussi.py` in the repo root.
+
+## FIRE relaxation
+
+Fast Inertial Relaxation Engine (Bitzek et al., PRL 97, 170201, 2006).
+ASE-compatible defaults.
+
+```python
+from alignn.md import FIRE
+
+relax = FIRE(forces_fn=forces_fn, masses=masses,
+ dt=0.1, dt_max=1.0, max_step=0.2)
+relaxed_pos, history = relax.run(
+ positions, fmax=0.05, max_steps=500, log_every=10,
+)
+```
+
+**Defaults match ASE:** `dt=0.1`, `dt_max=1.0`, `max_step=0.2`,
+`N_min=5`, `f_inc=1.1`, `f_dec=0.5`, `α_start=0.1`, `f_α=0.99`.
+Convergence: `max(|F_i|) < fmax` in eV/Å.
+
+**Not yet supported:** variable-cell relaxation (no `ExpCellFilter`
+equivalent), fixed-atom constraints. Use the ASE calculator path for
+those.
+
+## Velocity initialization
+
+```python
+from alignn.md import maxwell_boltzmann
+
+vel = maxwell_boltzmann(masses, T=300.0, generator=torch.Generator(device="cuda").manual_seed(0))
+# COM drift removed and velocities rescaled to hit T exactly.
+```
+
+## The `run` driver
+
+```python
+def cb(step, row):
+ print(f"t={row['time_fs']:.1f} fs T={row['T_K']:.1f} K")
+
+pos, vel, history = run(integ, pos, vel,
+ nsteps=10000, log_every=100, callback=cb)
+```
+
+Each log row contains `step`, `time_fs`, `T_K`, `KE_eV`, `wall_s`.
+
+## Melt-quench workflow (full example)
+
+Drop-in replacement for the ASE `NVTBerendsen` melt-quench pattern:
+
+```python
+from alignn.md import AlignnForces, NVTBerendsen, run, maxwell_boltzmann
+from alignn.md.integrators import wrap_pbc
+
+forces_fn = AlignnForces(model, Z, cell_np, device="cuda")
+vel = maxwell_boltzmann(masses, T=3500.0)
+
+integ = NVTBerendsen(forces_fn=forces_fn, masses=masses,
+ dt=1.0, T=3500.0, taut=20.0)
+
+# Stage 1 — melt
+pos, vel, _ = run(integ, pos, vel, nsteps=1000, log_every=20)
+
+# Stage 2 — quench
+integ.set_temperature(300.0)
+pos, vel, _ = run(integ, pos, vel, nsteps=2000, log_every=20)
+
+pos = wrap_pbc(pos, cell) # tidy coordinates for POSCAR output
+```
+
+See `md_melt_quench.py` in the repo root for a complete script with
+JARVIS loading and POSCAR output.
+
+## Choosing a thermostat — quick guide
+
+- **Equilibrating a new structure:** Berendsen or Bussi with `taut = 20-100 fs`
+- **Production sampling at fixed T:** Bussi (fast, simple) or NHC (if
+ you need dynamics)
+- **Transport coefficients (D, η, κ):** NHC (`chain_length=3-5`) with
+ weak coupling (`taut ≈ 100-500 fs`)
+- **Melt-quench:** Berendsen for the ramp, Bussi for any equilibration at
+ the target T
+- **Free energies, heat capacity, structural fluctuations:** Bussi or NHC
+ — **never** Berendsen
+
+## Known limitations
+
+1. **CPU neighbor list** rebuilt every step (see `AlignnForces._build_graph`).
+ Unblocks once replaced with an on-device version.
+2. **No trajectory file output** — plug an `ase.io.Trajectory` writer
+ into the `callback` parameter of `run()` if needed.
+3. **No NPT** — planned (Parrinello-Rahman / MTK barostat).
+4. **No constraints** (fixed atoms, SHAKE/RATTLE, symmetry).
+5. **Single GPU** — multi-GPU domain decomposition is a separate roadmap
+ item.
+
+## Related
+
+- [ASE Calculator](ase-calculator.md) — the canonical, full-featured
+ path via ASE.
+- `ALIGNNAtomWisePure` in `alignn.models.alignn_atomwise_pure` — the
+ DGL-free model powering these integrators.
diff --git a/docs/usage/melt-quench-tutorial.md b/docs/usage/melt-quench-tutorial.md
new file mode 100644
index 0000000..e7fad4b
--- /dev/null
+++ b/docs/usage/melt-quench-tutorial.md
@@ -0,0 +1,382 @@
+# Tutorial — Train → Export → LAMMPS Melt-Quench → RDF
+
+End-to-end walkthrough: train an ALIGNN-FF model on a Si dataset, export
+it to TorchScript, run a melt-quench simulation in LAMMPS with the native
+`pair_alignn` style, and plot the radial distribution function (RDF) of
+the quenched structure.
+
+Prereqs:
+
+- Clone + install alignn:
+ ```bash
+ git clone https://github.com/atomgptlab/alignn.git
+ cd alignn && pip install -e .
+ ```
+- LAMMPS built with `pair_alignn` compiled in. It's a single command:
+ ```bash
+ bash alignn/scripts/torch/build_lammps_alignn.sh
+ ```
+ The script auto-installs CUDA + MKL into your conda env if they're
+ missing. Full reference:
+ [Building LAMMPS with `pair_alignn`](lammps-pair-style.md#building-lammps-with-pair_alignn-from-a-fresh-clone).
+- Plotly for the RDF plot: `pip install plotly`.
+
+Working directory for this tutorial: `MELT/` (create anywhere).
+
+---
+
+## 1. Train the model
+
+Standard training command — nothing new here:
+
+```bash
+train_alignn.py \
+ --root_dir DataDir/ \
+ --config config.json \
+ --output_dir OutputDir
+```
+
+When training finishes, `OutputDir/` contains:
+
+```
+OutputDir/
+├── best_model.pt # weights
+├── config.json # hyperparameters + cutoff + neighbor_strategy
+├── history_*.json
+└── ...
+```
+
+**Important — requirements for the LAMMPS path:**
+
+1. In `config.json`, set `model.name: "alignn_atomwise_pure"` so
+ `forward_tensors_z` (the TorchScript entry point) is available. If
+ you trained with `alignn_atomwise` (DGL-backed), see the note at the
+ end of this page.
+2. Set `calculate_gradient: true` for forces to be computed.
+3. Note `cutoff` and `max_neighbors` from `config.json` — you'll pass
+ them to `pair_style alignn`.
+
+Quick look-up:
+```bash
+python -c "
+import json
+c = json.load(open('OutputDir/config.json'))
+print('cutoff :', c['cutoff'], 'Å')
+print('max_neighbors :', c['max_neighbors'])
+print('atom_features :', c.get('atom_features'))
+print('model.name :', c['model']['name'])
+"
+```
+
+---
+
+## 2. Export to TorchScript
+
+After `pip install -e .`, the export script is on your PATH:
+
+```bash
+export_torchscript.py \
+ --model-dir OutputDir \
+ --out alignn_ff.pt \
+ --atom-features atomic_number \
+ --dtype float32
+```
+
+Expected output:
+```
+smoke test OK: {'energy': torch.Size([]), 'forces': torch.Size([2, 3]), 'stress': torch.Size([3, 3])}
+wrote alignn_ff.pt (143,297 params)
+```
+
+**`--atom-features` must match training.** For most ALIGNN-FF trainings
+this is `atomic_number` (single scalar per atom) or `cgcnn` (92-dim).
+Check your config's top-level `atom_features` field.
+
+---
+
+## 3. Build a Si structure
+
+LAMMPS needs a data file. Build a 216-atom cubic Si supercell:
+
+```python
+# make_si_data.py
+from ase.build import bulk
+from ase.io import write
+
+# Cubic (orthorhombic) conventional cell, repeated 3×3×3 → 216 atoms.
+# Avoid triclinic primitive cells when possible — pair_alignn's shift
+# computation is most accurate for orthorhombic boxes.
+atoms = bulk("Si", "diamond", a=5.43, cubic=True).repeat((3, 3, 3))
+write("si.data", atoms, format="lammps-data", specorder=["Si"])
+print(f"wrote si.data N={len(atoms)} cell={atoms.cell.lengths()}")
+```
+
+```bash
+python make_si_data.py
+# wrote si.data N=216 cell=[16.29 16.29 16.29]
+```
+
+---
+
+## 4. Write the LAMMPS input script
+
+Five-stage melt-quench: equilibrate → heat → hold molten → quench →
+anneal. Save as `si_melt_quench.in`:
+
+```
+# Si melt-quench via pair_alignn.
+units metal # eV, Å, ps, bar, K
+atom_style atomic
+boundary p p p
+read_data si.data
+mass 1 28.0855
+
+# ── ALIGNN-FF native pair style ──
+pair_style alignn 5.0 12 # cutoff (Å), max_neighbors
+pair_coeff * * alignn_ff.pt Si
+
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+
+timestep 0.001 # 1 fs
+thermo 100
+thermo_style custom step temp pe ke etotal press vol
+dump traj all custom 200 si.lammpstrj id type x y z
+
+# ── Stage 1 — equilibrate 300 K (0.2 ps) ──
+velocity all create 300.0 12345 mom yes rot yes
+fix nvt1 all nvt temp 300.0 300.0 0.1
+run 200
+unfix nvt1
+
+# ── Stage 2 — heat to 3500 K (1 ps) ──
+fix nvt2 all nvt temp 300.0 3500.0 0.1
+run 1000
+unfix nvt2
+
+# ── Stage 3 — hold molten (2 ps) ──
+fix nvt3 all nvt temp 3500.0 3500.0 0.1
+run 2000
+unfix nvt3
+
+# ── Stage 4 — quench to 300 K (4 ps) ──
+fix nvt4 all nvt temp 3500.0 300.0 0.1
+run 4000
+unfix nvt4
+
+# ── Stage 5 — anneal 300 K (2 ps) — RDF measured here ──
+compute rdf all rdf 150 1 1 cutoff 6.0
+fix rdf_out all ave/time 100 10 1000 c_rdf[*] file si_rdf.dat mode vector
+fix nvt5 all nvt temp 300.0 300.0 0.1
+run 2000
+unfix nvt5
+unfix rdf_out
+
+write_data si_quenched.data
+```
+
+The `compute rdf` + `fix ave/time` block samples the radial distribution
+function every 100 steps and writes the averaged result to `si_rdf.dat`.
+
+---
+
+## 5. Run it
+
+```bash
+~/lammps-alignn/build/lmp -in si_melt_quench.in
+```
+
+Or from Python:
+```python
+from lammps import lammps
+lmp = lammps()
+lmp.file("si_melt_quench.in")
+```
+
+Runtime on a single modest GPU: ~10–30 min for the full 9.2 ps. Tail of
+the thermo output will show temperature ramping up and coming back down:
+
+```
+Step Temp PotEng KinEng TotEng Press Volume
+0 300 -987.23 5.58 -981.65 -5.2e4 4318.9
+200 842 -984.11 15.67 -968.44 -2.3e6 4318.9
+...
+3200 3501 -961.22 65.12 -896.10 +3.1e5 4318.9
+...
+9200 301 -985.40 5.60 -979.80 -1.4e4 4318.9
+```
+
+**Known quirk**: your pressure numbers will look extreme if your model
+was trained with `stresswise_weight: 0.0` — stresses aren't learned,
+only energies and forces. That's fine for NVT; don't use NPT with such a
+model.
+
+---
+
+## 6. Plot the RDF
+
+The `si_rdf.dat` file LAMMPS wrote has a specific multi-block format.
+This script parses it and plots with Plotly:
+
+```python
+# plot_rdf.py
+import numpy as np
+import plotly.graph_objects as go
+from pathlib import Path
+
+def parse_lammps_rdf(path: str | Path):
+ """Parse LAMMPS `fix ave/time ... mode vector` RDF output.
+
+ Each averaged block begins with a header line ` `
+ followed by `nbins` lines of `index r g(r) coord(r)`.
+ Returns the LAST block (steady-state RDF).
+ """
+ with open(path) as f:
+ lines = [L.rstrip() for L in f if L.strip() and not L.startswith("#")]
+ # Find block starts (lines with exactly 2 tokens = timestep, nbins)
+ starts = [i for i, L in enumerate(lines)
+ if len(L.split()) == 2 and not L.startswith(" ")]
+ last = starts[-1]
+ nbins = int(lines[last].split()[1])
+ data = np.array([list(map(float, L.split())) for L in lines[last+1:last+1+nbins]])
+ r = data[:, 1]
+ g_r = data[:, 2]
+ coord = data[:, 3]
+ return r, g_r, coord
+
+def plot_rdf(r, g_r, coord, title="Si RDF after melt-quench"):
+ fig = go.Figure()
+ fig.add_trace(go.Scatter(x=r, y=g_r, name="g(r)",
+ line=dict(color="steelblue", width=2)))
+ fig.add_trace(go.Scatter(x=r, y=coord, name="coordination",
+ line=dict(color="firebrick", width=1.5, dash="dash"),
+ yaxis="y2"))
+ fig.update_layout(
+ title=title, template="plotly_white",
+ width=800, height=500,
+ xaxis=dict(title="r (Å)"),
+ yaxis =dict(title="g(r)", titlefont=dict(color="steelblue")),
+ yaxis2=dict(title="coordination", overlaying="y", side="right",
+ titlefont=dict(color="firebrick"), showgrid=False),
+ )
+ return fig
+
+if __name__ == "__main__":
+ r, g_r, coord = parse_lammps_rdf("si_rdf.dat")
+ # Report key features
+ i_peak = np.argmax(g_r)
+ print(f"first-neighbor peak: r = {r[i_peak]:.3f} Å, g(r) = {g_r[i_peak]:.2f}")
+ print(f"coordination at 3.0 Å: {np.interp(3.0, r, coord):.2f} "
+ f"(expected ≈ 4 for crystalline Si, lower for amorphous)")
+ fig = plot_rdf(r, g_r, coord)
+ fig.write_html("si_rdf.html")
+ print("wrote si_rdf.html")
+```
+
+Run it:
+```bash
+python plot_rdf.py
+```
+
+**Expected Si RDF:**
+
+| Feature | Crystalline Si | Amorphous Si (well-relaxed) | Amorphous Si (poorly quenched) |
+|---|---|---|---|
+| First peak position | 2.35 Å | 2.35–2.40 Å | similar |
+| First peak height | sharp, g>10 | broad, g≈3–5 | low, g≈2 |
+| 2nd-neighbor splitting at ~3.8 Å | distinct peak | shoulder or gone | absent |
+| Coordination at 3.0 Å | 4.0 (exact) | ~3.8–4.2 | variable |
+
+A **4 ps quench from 3500 K is very fast** — don't expect a perfectly
+amorphous or perfectly crystalline result. For publication-quality
+glasses, quench over 100+ ps.
+
+---
+
+## 7. Verifying the run was physical
+
+Sanity checks before trusting the results:
+
+1. **Energy conservation during NVE windows.** Add a short NVE block and
+ watch `etotal` — drift should be < 1e-4 eV/atom/ps.
+2. **No model blow-ups.** If temperature suddenly spikes or PE diverges,
+ the force call probably hit an unphysical configuration outside the
+ model's training domain. Use a shorter timestep (0.0005 ps) or a
+ softer heating ramp.
+3. **Compare to the Python calculator.** Run one single-point with
+ `AlignnAtomwiseCalculator` on the starting structure; LAMMPS should
+ match to ~1e-4 eV for energy. If it doesn't, see [Model-side
+ caveats](#model-side-caveats) below.
+4. **Check the trajectory.** Open `si.lammpstrj` in OVITO; atoms
+ shouldn't fly out of the box.
+
+---
+
+## Model-side caveats
+
+!!! warning "`alignn_atomwise` vs `alignn_atomwise_pure`"
+
+ The TorchScript entry point `forward_tensors_z` only exists on
+ `ALIGNNAtomWisePure`. If your training config used `name:
+ alignn_atomwise` (the DGL-backed variant), `export_torchscript.py`
+ will still run by loading the DGL weights into Pure — but the two
+ models are not numerically identical on the same weights (currently
+ a ~10% energy difference due to line-graph construction
+ differences).
+
+ **Clean fix**: retrain with `name: alignn_atomwise_pure` in your
+ training config. All other hyperparameters can stay the same.
+
+!!! warning "Stress not learned"
+
+ If your training data didn't include stresses (or
+ `stresswise_weight: 0.0`), `pair_alignn` still returns virial
+ numbers, but they're noise. Don't use NPT / `fix npt`.
+
+!!! warning "Single MPI rank"
+
+ Current `pair_alignn` silently drops cross-rank ghost contributions.
+ Run with `mpirun -n 1` only. For larger systems, increase
+ parallelism via threads (`OMP_NUM_THREADS`) until multi-rank ghost
+ handling lands.
+
+---
+
+## What to try next
+
+- **Transport**: replace NVT with NVE in stage 5, compute diffusion
+ coefficient from the mean-squared-displacement (LAMMPS
+ `compute msd`).
+- **Cool more slowly**: extend stage 4 to 40 ps and compare the RDF —
+ you should see sharper peaks and a clearer second-neighbor split.
+- **Longer correlation**: rerun with `mode vector 10` (average over 10
+ snapshots per write) for smoother RDFs.
+- **Different material**: retrain on another element / alloy, re-export,
+ swap `Si` in `pair_coeff` for the target species.
+
+---
+
+## Files produced by this tutorial
+
+```
+MELT/
+├── alignn_ff.pt # exported TorchScript model
+├── si.data # LAMMPS input structure
+├── si_melt_quench.in # LAMMPS input script
+├── si.lammpstrj # trajectory (custom dump)
+├── si_rdf.dat # RDF samples (fix ave/time)
+├── si_quenched.data # final structure
+├── plot_rdf.py # RDF plotter
+└── si_rdf.html # interactive RDF plot
+```
+
+---
+
+## See also
+
+- [LAMMPS Pair Style (native C++)](lammps-pair-style.md) — full
+ reference for `pair_alignn`
+- [MD & Relaxation (pure PyTorch)](md-integrators.md) — on-device
+ integrators as an alternative to LAMMPS
+- [ASE Calculator](ase-calculator.md) — the Python-side calculator
+ used for the reference comparison in step 7
diff --git a/mkdocs.yml b/mkdocs.yml
index 765ba99..dd09e74 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -68,6 +68,7 @@ extra:
nav:
- Home: index.md
- Installation: installation.md
+ - Colab Notebooks: examples-colab.md
- Training:
- Dataset Format: training/dataset-format.md
- Single-Output Regression: training/single-output-regression.md
@@ -81,9 +82,13 @@ nav:
- ALIGNN-FF: pretrained/alignn-ff.md
- Usage:
- ASE Calculator: usage/ase-calculator.md
+ - MD & Relaxation (pure PyTorch): usage/md-integrators.md
+ - LAMMPS Pair Style (native C++): usage/lammps-pair-style.md
+ - Tutorial — Melt-Quench → RDF: usage/melt-quench-tutorial.md
- Web Apps: usage/webapps.md
- Reference:
- Package Overview: api.md
+ - Config Reference: reference/config-reference.md
- Performance: performance.md
- Notes & FAQ: notes.md
- References: references.md
diff --git a/setup.py b/setup.py
index ab9bb21..0fcbe58 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@
setuptools.setup(
name="alignn",
- version="2026.4.1",
+ version="2026.5.20",
author="Kamal Choudhary, Brian DeCost",
author_email="kamal.choudhary@nist.gov",
description="alignn",
@@ -22,7 +22,7 @@
"jarvis-tools>=2021.07.19",
"torch>=2.2.1",
# "torch<=2.2.1",
- "mpmath<=1.3.0",
+ "mpmath>=1.3.0",
# "dgl>=0.6.0",
"spglib>=2.0.2",
"scikit-learn>=0.22.2",
@@ -43,6 +43,7 @@
"alignn/pretrained.py",
"alignn/train_alignn.py",
"alignn/run_alignn_ff.py",
+ "alignn/scripts/torch/export_torchscript.py",
],
long_description=long_description,
long_description_content_type="text/markdown",