diff --git a/README.md b/README.md
index 66ffbe1..2d1d0aa 100644
--- a/README.md
+++ b/README.md
@@ -67,6 +67,7 @@ training, and pretrained-model usage. Click a badge to open in Colab.
| 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. |
+| LAMMPS MD with ALIGNN-FF (`pair_alignn`) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_FF_LAMMPS_Colab.ipynb) | Build LAMMPS with the native `pair_alignn` style and run NVE / melt-quench MD with the default ALIGNN-FF `mps` force field. |
## Using pre-trained models
diff --git a/alignn/config.py b/alignn/config.py
index 6904417..ae8f2a0 100644
--- a/alignn/config.py
+++ b/alignn/config.py
@@ -8,6 +8,9 @@
from alignn.models.alignn import ALIGNNConfig
from alignn.models.alignn_atomwise import ALIGNNAtomWiseConfig
from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePureConfig
+from alignn.models.alignn_atomwise_pure_smooth import (
+ ALIGNNAtomWisePureSmoothConfig,
+)
from alignn.models.ealignn_atomwise import eALIGNNAtomWiseConfig
# import torch
@@ -235,5 +238,6 @@ class TrainingConfig(BaseSettings):
ALIGNNConfig,
ALIGNNAtomWiseConfig,
ALIGNNAtomWisePureConfig,
+ ALIGNNAtomWisePureSmoothConfig,
eALIGNNAtomWiseConfig,
] = ALIGNNAtomWiseConfig(name="alignn_atomwise")
diff --git a/alignn/examples/lammps/README.md b/alignn/examples/lammps/README.md
new file mode 100644
index 0000000..49f542a
--- /dev/null
+++ b/alignn/examples/lammps/README.md
@@ -0,0 +1,56 @@
+# LAMMPS + ALIGNN-FF (pure-torch) examples
+
+These examples drive LAMMPS with the **native `pair_alignn`** style backed
+by the pure-PyTorch ALIGNN-FF model (TorchScripted to `alignn_ff.pt`).
+No Python callback is involved — forces/stresses come directly from the
+C++ pair style calling libtorch.
+
+## Prerequisites
+
+1. Build LAMMPS with `pair_alignn` linked against libtorch:
+ ```bash
+ bash ../../scripts/torch/build_lammps_alignn.sh
+ ```
+2. Download the default ALIGNN-FF `mps` model from figshare and export
+ it to TorchScript (`alignn_ff.pt`) into this directory:
+ ```bash
+ python get_model.py
+ ```
+ To use a different pretrained model, edit `default_path()` in
+ `get_model.py` (e.g. `model_name="v12.2.2024_dft_3d_307k"`).
+
+## Examples
+
+| File | Purpose |
+| --- | --- |
+| `nve_si_stability.in` + `build_si.py` | NVE stability test on Si (216 atoms). Run this first to confirm the model conserves energy. |
+| `melt_quench_si.in` + `build_si.py` | Melt-quench Si: 300 K → 3500 K → quench → 300 K. Produces amorphous Si. |
+| `melt_quench_sio2.in` + `build_sio2.py` | Melt-quench α-quartz SiO₂: 300 K → 4000 K → quench → 300 K. Produces amorphous silica. |
+
+## How to run
+
+```bash
+# 1. Build the initial structure
+python build_si.py # writes si.data
+
+# 2. NVE stability check (look for flat etotal in the log)
+lmp -in nve_si_stability.in
+
+# 3. Melt-quench
+lmp -in melt_quench_si.in
+```
+
+For SiO₂:
+```bash
+python build_sio2.py # writes sio2.data
+lmp -in melt_quench_sio2.in
+```
+
+## NVE stability — what to look for
+
+Open `log.lammps` and check `etotal`. Over a 5 ps NVE run the total
+energy should drift by less than a few meV/atom. Large drift (>10 meV/atom)
+indicates an undertrained model, too-aggressive timestep, or a cutoff
+mismatch between the model and the LAMMPS neighbor list.
+
+The script also dumps `nve_si.lammpstrj`; visualize with OVITO.
diff --git a/alignn/examples/lammps/build_si.py b/alignn/examples/lammps/build_si.py
new file mode 100644
index 0000000..518d21e
--- /dev/null
+++ b/alignn/examples/lammps/build_si.py
@@ -0,0 +1,38 @@
+"""Build a Si supercell, relax it with the ALIGNN-FF `mps` model
+(so the LAMMPS run doesn't start at MBar-scale residual stress),
+and write `si.data`.
+"""
+from ase.io import write
+from ase.optimize import FIRE
+from jarvis.db.figshare import get_jid_data
+from jarvis.core.atoms import Atoms
+
+from alignn.ff.ff import default_path
+from alignn.ff.calculators import AlignnAtomwiseCalculator
+
+# ExpCellFilter moved to ase.filters (and is deprecated in favor of
+# FrechetCellFilter) in recent ASE; fall back across versions.
+try:
+ from ase.filters import FrechetCellFilter as CellFilter
+except ImportError:
+ try:
+ from ase.filters import ExpCellFilter as CellFilter
+ except ImportError:
+ from ase.constraints import ExpCellFilter as CellFilter
+
+
+a = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"])
+# JVASP-1002 is the 2-atom primitive FCC cell; a supercell of it is
+# triclinic. Use the conventional (cubic) cell so the box stays orthogonal
+# -> a 2x2x2 supercell is a clean 64-atom cubic system for the MD examples.
+ase_atoms = a.get_conventional_atoms.make_supercell([2, 2, 2]).ase_converter()
+print(f"initial: {len(ase_atoms)} atoms, V = {ase_atoms.get_volume():.2f} ų")
+
+ase_atoms.calc = AlignnAtomwiseCalculator(path=default_path())
+ecf = CellFilter(ase_atoms)
+FIRE(ecf, logfile="-").run(fmax=0.05, steps=200)
+print(f"relaxed: V = {ase_atoms.get_volume():.2f} ų, "
+ f"E = {ase_atoms.get_potential_energy():.3f} eV")
+
+write("si.data", ase_atoms, format="lammps-data", specorder=["Si"])
+print(f"wrote si.data: {len(ase_atoms)} atoms")
diff --git a/alignn/examples/lammps/build_sio2.py b/alignn/examples/lammps/build_sio2.py
new file mode 100644
index 0000000..09108a3
--- /dev/null
+++ b/alignn/examples/lammps/build_sio2.py
@@ -0,0 +1,37 @@
+"""Build an α-quartz SiO2 supercell, relax it with the ALIGNN-FF `mps`
+model, and write `sio2.data`.
+
+JVASP-41 is α-quartz. A 3x3x2 supercell gives ~162 atoms.
+"""
+from ase.io import write
+from ase.optimize import FIRE
+from jarvis.db.figshare import get_jid_data
+from jarvis.core.atoms import Atoms
+
+from alignn.ff.ff import default_path
+from alignn.ff.calculators import AlignnAtomwiseCalculator
+
+# ExpCellFilter moved to ase.filters (and is deprecated in favor of
+# FrechetCellFilter) in recent ASE; fall back across versions.
+try:
+ from ase.filters import FrechetCellFilter as CellFilter
+except ImportError:
+ try:
+ from ase.filters import ExpCellFilter as CellFilter
+ except ImportError:
+ from ase.constraints import ExpCellFilter as CellFilter
+
+
+a = Atoms.from_dict(get_jid_data(jid="JVASP-41", dataset="dft_3d")["atoms"])
+ase_atoms = a.make_supercell([3, 3, 2]).ase_converter()
+print(f"initial: {len(ase_atoms)} atoms, V = {ase_atoms.get_volume():.2f} ų")
+
+ase_atoms.calc = AlignnAtomwiseCalculator(path=default_path())
+ecf = CellFilter(ase_atoms)
+FIRE(ecf, logfile="-").run(fmax=0.1, steps=200)
+print(f"relaxed: V = {ase_atoms.get_volume():.2f} ų, "
+ f"E = {ase_atoms.get_potential_energy():.3f} eV")
+
+# specorder fixes type-id mapping: type 1 = Si, type 2 = O.
+write("sio2.data", ase_atoms, format="lammps-data", specorder=["Si", "O"])
+print(f"wrote sio2.data: {len(ase_atoms)} atoms")
diff --git a/alignn/examples/lammps/export_smooth.py b/alignn/examples/lammps/export_smooth.py
new file mode 100644
index 0000000..90d39f7
--- /dev/null
+++ b/alignn/examples/lammps/export_smooth.py
@@ -0,0 +1,71 @@
+"""Export a trained ALIGNN-atomwise-pure-SMOOTH checkpoint to TorchScript.
+
+The stock `export_torchscript` script hardcodes the non-smooth class;
+this is the smooth-variant analogue.
+
+Usage:
+ python export_smooth.py --model-dir \
+ --out alignn_ff.pt
+"""
+import argparse
+import json
+import torch
+
+from alignn.models.alignn_atomwise_pure_smooth import (
+ ALIGNNAtomWisePureSmooth, ALIGNNAtomWisePureSmoothConfig,
+)
+
+
+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")
+ args = ap.parse_args()
+
+ cfg_raw = json.load(open(f"{args.model_dir}/config.json"))
+ mcfg = dict(cfg_raw["model"])
+ mcfg["name"] = "alignn_atomwise_pure_smooth"
+ mcfg["calculate_gradient"] = True
+ if mcfg.get("stresswise_weight", 0) == 0:
+ mcfg["stresswise_weight"] = 0.01
+ cfg = ALIGNNAtomWisePureSmoothConfig(**mcfg)
+
+ model = ALIGNNAtomWisePureSmooth(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)}")
+ if missing:
+ print(f" first missing: {missing[:3]}")
+ if unexpected:
+ print(f" first unexpected: {unexpected[:3]}")
+
+ model.register_species_table(atom_features=args.atom_features)
+ model = model.to(torch.float32).eval()
+
+ scripted = torch.jit.script(model)
+
+ with torch.no_grad():
+ pos = torch.tensor([[0., 0., 0.], [1.35, 1.35, 1.35]],
+ dtype=torch.float32)
+ pos.requires_grad_(True)
+ lat = torch.eye(3, dtype=torch.float32) * 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=torch.float32)
+ out = scripted.forward_tensors_z(pos, lat, Z, src, dst, shift, True)
+ print("smoke test OK:",
+ {k: (v.shape if hasattr(v, "shape") else v) for k, v in out.items()})
+
+ torch.jit.save(scripted, args.out)
+ n = sum(p.numel() for p in model.parameters())
+ print(f"wrote {args.out} ({n:,} params)")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/examples/lammps/get_model.py b/alignn/examples/lammps/get_model.py
new file mode 100644
index 0000000..68729a0
--- /dev/null
+++ b/alignn/examples/lammps/get_model.py
@@ -0,0 +1,41 @@
+"""Download the default ALIGNN-FF `mps` model from figshare and export
+it to TorchScript (`alignn_ff.pt`) for the native LAMMPS `pair_alignn`.
+
+Run once before the LAMMPS examples:
+ python get_model.py
+"""
+import os
+import subprocess
+import sys
+
+from alignn.ff.ff import get_figshare_model_ff
+
+
+def default_path():
+ """Default ALIGNN-FF model path (downloads on first call)."""
+ return get_figshare_model_ff(model_name="mps")
+
+
+def main():
+ model_dir = default_path()
+ print(f"model dir: {model_dir}")
+ assert os.path.exists(os.path.join(model_dir, "best_model.pt")), (
+ f"best_model.pt not found in {model_dir}"
+ )
+ assert os.path.exists(os.path.join(model_dir, "config.json")), (
+ f"config.json not found in {model_dir}"
+ )
+
+ out = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+ "alignn_ff.pt")
+ cmd = [sys.executable,
+ "-m", "alignn.scripts.torch.export_torchscript",
+ "--model-dir", model_dir,
+ "--out", out]
+ print("running:", " ".join(cmd))
+ subprocess.check_call(cmd)
+ print(f"wrote {out}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/alignn/examples/lammps/handsoff.md b/alignn/examples/lammps/handsoff.md
new file mode 100644
index 0000000..b2cfb4f
--- /dev/null
+++ b/alignn/examples/lammps/handsoff.md
@@ -0,0 +1,107 @@
+# LAMMPS + ALIGNN-FF (pure-torch) — handoff notes
+
+## What's here
+
+| File | Purpose |
+| --- | --- |
+| `get_model.py` | Downloads the figshare `mps` ALIGNN-FF checkpoint and exports it to TorchScript (`alignn_ff.pt`) using `alignn.scripts.torch.export_torchscript`. |
+| `export_smooth.py` | Exports a **smooth-variant** checkpoint (`alignn_atomwise_pure_smooth`) to TorchScript. The stock exporter hardcodes the non-smooth class. |
+| `build_si.py` | Builds a 54-atom Si supercell from JARVIS JVASP-1002, relaxes it with `AlignnAtomwiseCalculator + ExpCellFilter + FIRE`, writes `si.data`. |
+| `build_sio2.py` | Same for α-quartz JVASP-41 (162 atoms), writes `sio2.data`. |
+| `nve_si_stability.in` | NVE energy-conservation test on Si. |
+| `melt_quench_si.in` | 9 ps Si melt-quench (300 K → 3500 K → 300 K). |
+| `melt_quench_sio2.in` | 13.5 ps SiO₂ melt-quench (300 K → 4000 K → 300 K). |
+| `alignn_ff.pt` | Currently the **smooth** TorchScript model (143k params). |
+| `alignn_ff_mps.pt` | Backup of the non-smooth `mps` export, kept for comparison. |
+
+## How it runs
+
+```bash
+# 1. Build LAMMPS once with pair_alignn (see ../../scripts/torch/build_lammps_alignn.sh)
+
+# 2. Get a model:
+# a) non-smooth mps checkpoint from figshare
+python get_model.py
+# b) OR the smooth checkpoint (currently the active alignn_ff.pt)
+python export_smooth.py \
+ --model-dir ~/Software/ollama311/alignn/MELT_smooth_long/OutputDir \
+ --out alignn_ff.pt
+
+# 3. Build relaxed structures, then run
+python build_si.py && lmp -in nve_si_stability.in
+python build_si.py && lmp -in melt_quench_si.in
+python build_sio2.py && lmp -in melt_quench_sio2.in
+```
+
+## Gotchas discovered while bringing this up
+
+1. **`pair_style alignn` needs explicit cutoff and max_neighbors** —
+ syntax is `pair_style alignn `. These **must
+ match the model's training config**: look at `cutoff` and
+ `max_neighbors` at the top level of `config.json` (not under
+ `model:`). For both checkpoints used here that's `5.0 12`. Passing
+ `8.0 25` puts the model far out of distribution — saw 200 neighbors/atom
+ and immediate blow-up.
+
+2. **Relax with the model before LAMMPS.** The unrelaxed JARVIS
+ geometries (especially JVASP-1002 Si) give ~MBar residual stress
+ under both checkpoints — LAMMPS `fix box/relax` can't escape the
+ minimum, so MD blows up within a few hundred steps. `build_si.py` /
+ `build_sio2.py` now run `ExpCellFilter + FIRE(fmax=0.05)` first.
+
+3. **`reset_timestep` must come before any `dump`** (or you `undump`
+ first). Initial NVE file had the dump active across the reset.
+
+4. **`pair_alignn.cpp settings(...)` reads the args** — confirmed it
+ keeps the `max_neighbors_` closest neighbors within `cutoff_`.
+ So the model receives exactly the graph it was trained on; the
+ instability problems were the geometry and the model itself,
+ not the bridge.
+
+## Stability comparison: non-smooth vs smooth
+
+5 ps NVE on relaxed 54-atom Si, 1 fs timestep, both models with
+identical 143k params, both trained with `atom_input_features=1`,
+`cutoff=5.0`, `max_neighbors=12`:
+
+| Metric | non-smooth (`mps`) | smooth (`MELT_smooth_long`) |
+| --- | --- | --- |
+| Total energy drift rate | +12.3 meV/atom/ps | **+0.47 meV/atom/ps** |
+| Total energy range over 5 ps | 4424 meV | **112 meV** |
+| Temperature trend | 113 K → 279 K (heating) | **89 K → 94 K** (flat) |
+| Temperature range | [89, 345] K | **[76, 105] K** |
+
+The smooth variant's polynomial cutoff envelope (`CutoffPolynomial`,
+coeff=5.0) eliminates force discontinuities at the cutoff. Net effect:
+forces are essentially conservative, so NVE behaves like NVE instead
+of slowly heating up.
+
+The non-smooth `mps` checkpoint is bounded (atoms don't fly off) but
+**not energy-conserving** to MLIP standards. It's clearly a smoke-test
+checkpoint — `stresswise_weight=0` in its training config, so stress
+was never fit.
+
+## Patches to the model source (needed to TorchScript-export the smooth variant)
+
+`alignn/models/alignn_atomwise_pure_smooth.py`:
+
+- `np.sqrt(2.0)` and `np.sqrt(np.pi)` inside `Fourier.forward()` →
+ replaced with the float constants. TorchScript can't take weak
+ references to numpy ufuncs.
+- `self.radial_envelope` was only assigned in the gaussian branch.
+ Added an explicit `Optional[CutoffPolynomial]` class-body
+ annotation (and default `= None` before the basis-kind switch) so
+ TorchScript can compile `_embed_radial` regardless of which branch
+ is taken at construction.
+- Imported `Optional` from `typing`.
+
+## Open / next
+
+- The MD runs use CPU torch (≈30 steps/s for 54 atoms, ≈12 steps/s
+ for 162 atoms). For longer runs export with CUDA — `pair_alignn.cpp`
+ already picks CUDA automatically if `torch::cuda::is_available()`.
+- The SiO₂ melt-quench is short (13.5 ps); for production-quality
+ amorphous silica it should be ≥50 ps total with a slower (~0.2 K/fs)
+ quench.
+- The non-smooth `mps` checkpoint is not worth rerunning for science;
+ keep it only as a TorchScript-export smoke test.
diff --git a/alignn/examples/lammps/melt_quench_si.in b/alignn/examples/lammps/melt_quench_si.in
new file mode 100644
index 0000000..4c3222f
--- /dev/null
+++ b/alignn/examples/lammps/melt_quench_si.in
@@ -0,0 +1,49 @@
+# Melt-quench of crystalline Si using ALIGNN-FF (pure-torch) via pair_alignn.
+# Produces amorphous Si in `si_quenched.data`.
+#
+# Run: python build_si.py && lmp -in melt_quench_si.in
+
+units metal
+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 200 si_melt_quench.lammpstrj id type x y z
+
+velocity all create 300.0 12345 mom yes rot yes dist gaussian
+
+# Stage 1: equilibrate at 300 K (0.2 ps)
+fix nvt1 all nvt temp 300.0 300.0 0.1
+run 200
+unfix nvt1
+
+# Stage 2: heat 300 K -> 3500 K (1 ps)
+fix nvt2 all nvt temp 300.0 3500.0 0.1
+run 1000
+unfix nvt2
+
+# Stage 3: hold molten at 3500 K (2 ps)
+fix nvt3 all nvt temp 3500.0 3500.0 0.1
+run 2000
+unfix nvt3
+
+# Stage 4: quench 3500 K -> 300 K (4 ps, ~0.8 K/fs)
+fix nvt4 all nvt temp 3500.0 300.0 0.1
+run 4000
+unfix nvt4
+
+# Stage 5: anneal at 300 K (2 ps)
+fix nvt5 all nvt temp 300.0 300.0 0.1
+run 2000
+
+write_data si_quenched.data
diff --git a/alignn/examples/lammps/melt_quench_sio2.in b/alignn/examples/lammps/melt_quench_sio2.in
new file mode 100644
index 0000000..8f6880f
--- /dev/null
+++ b/alignn/examples/lammps/melt_quench_sio2.in
@@ -0,0 +1,57 @@
+# Melt-quench of α-quartz SiO2 using ALIGNN-FF (pure-torch `mps`) via
+# pair_alignn. Produces amorphous silica in `sio2_quenched.data`.
+# SiO2 melts at ~1973 K; we use 4000 K for fast diffusion.
+#
+# Schedule is short (~7 ps total) for a demo. Use a longer quench
+# (e.g. 50 ps) for production-quality amorphous silica.
+#
+# Run: python build_sio2.py && lmp -in melt_quench_sio2.in
+
+units metal
+atom_style atomic
+boundary p p p
+read_data sio2.data
+
+# Type ids set by build_sio2.py (specorder=["Si","O"])
+mass 1 28.0855 # Si
+mass 2 15.9994 # O
+
+# Cutoff + max_neighbors must match the trained mps model.
+pair_style alignn 5.0 12
+pair_coeff * * alignn_ff.pt Si O
+
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+
+timestep 0.0005 # 0.5 fs — lighter O atom needs smaller dt
+thermo 50
+thermo_style custom step temp pe ke etotal press vol
+dump traj all custom 500 sio2_melt_quench.lammpstrj id type x y z
+
+velocity all create 300.0 54321 mom yes rot yes dist gaussian
+
+# Stage 1: equilibrate at 300 K (0.25 ps)
+fix nvt1 all nvt temp 300.0 300.0 0.05
+run 500
+unfix nvt1
+
+# Stage 2: heat 300 K -> 4000 K (1 ps)
+fix nvt2 all nvt temp 300.0 4000.0 0.05
+run 2000
+unfix nvt2
+
+# Stage 3: hold molten at 4000 K (1.5 ps)
+fix nvt3 all nvt temp 4000.0 4000.0 0.05
+run 3000
+unfix nvt3
+
+# Stage 4: quench 4000 K -> 300 K (3 ps, ~1.2 K/fs — fast for a demo)
+fix nvt4 all nvt temp 4000.0 300.0 0.05
+run 6000
+unfix nvt4
+
+# Stage 5: anneal at 300 K (1 ps)
+fix nvt5 all nvt temp 300.0 300.0 0.05
+run 2000
+
+write_data sio2_quenched.data
diff --git a/alignn/examples/lammps/nve_si_stability.in b/alignn/examples/lammps/nve_si_stability.in
new file mode 100644
index 0000000..8a614e7
--- /dev/null
+++ b/alignn/examples/lammps/nve_si_stability.in
@@ -0,0 +1,45 @@
+# NVE stability test for ALIGNN-FF (pure-torch `mps` model) on Si.
+# A well-trained model + correct cutoff should conserve etotal to
+# within a few meV/atom over several picoseconds.
+#
+# Run: lmp -in nve_si_stability.in
+
+units metal
+atom_style atomic
+boundary p p p
+read_data si.data
+mass 1 28.0855
+
+# Native pair style: forces/stresses come from the TorchScripted ALIGNN-FF.
+# Cutoff + max_neighbors must match what the model was trained with
+# (mps model: cutoff=5.0 Å, max_neighbors=12).
+pair_style alignn 5.0 12
+pair_coeff * * alignn_ff.pt Si
+
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+
+# --- Relax cell + atoms first so we don't start at MBar-scale stress ---
+thermo 10
+thermo_style custom step temp pe ke etotal press vol
+fix relax all box/relax iso 0.0 vmax 0.001
+min_style cg
+minimize 1.0e-6 1.0e-4 200 1000
+unfix relax
+minimize 1.0e-6 1.0e-4 200 1000
+
+# --- Gentle thermalisation at 100 K ---
+timestep 0.0005 # 0.5 fs
+velocity all create 100.0 12345 mom yes rot yes dist gaussian
+fix eq all nvt temp 100.0 100.0 0.05
+thermo 50
+run 1000
+unfix eq
+
+# --- NVE stability run: 2.5 ps. Watch etotal in the log. ---
+reset_timestep 0
+dump traj all custom 200 nve_si.lammpstrj id type x y z
+fix nve all nve
+run 5000
+
+write_data si_nve_final.data
diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index 2312ef0..0140d1a 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -194,9 +194,9 @@ def __init__(
output_dir=None,
batch_stress=True,
force_mult_natoms=False,
- force_mult_batchsize=True,
+ force_mult_batchsize=False,
force_multiplier=1,
- stress_wt=0.05,
+ stress_wt=1.00,
):
"""Initialize class."""
super(AlignnAtomwiseCalculator, self).__init__(
@@ -410,7 +410,7 @@ def __init__(
ff_model_filename="best_model.pt",
ff_config_filename="config.json",
ff_config=None,
- stress_wt=0.05,
+ stress_wt=1.0,
props=[
"cbm",
"vbm",
diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index eac85f7..da2b28f 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -1203,16 +1203,17 @@ def phonons(
phonon = Phonopy(bulk, [[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]])
phonon.generate_displacements(distance=distance)
# print("Len dis", len(phonon.supercells_with_displacements))
- supercells = phonon.get_supercells_with_displacements()
+ # property form; get_* methods were removed in phonopy 2.20+
+ supercells = phonon.supercells_with_displacements
# Force calculations by calculator
set_of_forces = []
disp = 0
for scell in supercells:
ase_atoms = AseAtoms(
- symbols=scell.get_chemical_symbols(),
- scaled_positions=scell.get_scaled_positions(),
- cell=scell.get_cell(),
+ symbols=scell.symbols, # property form (phonopy 2.10/2.48/4.x)
+ scaled_positions=scell.scaled_positions,
+ cell=scell.cell,
pbc=True,
)
ase_atoms.calc = calc
@@ -1227,7 +1228,7 @@ def phonons(
phonon.produce_force_constants(forces=set_of_forces)
if write_fc:
write_FORCE_CONSTANTS(
- phonon.get_force_constants(), filename="FORCE_CONSTANTS"
+ phonon.force_constants, filename="FORCE_CONSTANTS"
)
lbls = kpoints.labels
@@ -1286,7 +1287,11 @@ def phonons(
tdos = phonon._total_dos
# print('tods',tdos._frequencies.shape)
- freqs, ds = tdos.get_dos()
+ # phonopy >=2.20 removed TotalDos.get_dos(); use properties with fallback
+ if hasattr(tdos, "frequency_points"):
+ freqs, ds = tdos.frequency_points, tdos.dos
+ else:
+ freqs, ds = tdos.get_dos()
freqs = np.array(freqs)
freqs = freqs * freq_conversion_factor
min_freq = min_freq_tol * freq_conversion_factor
@@ -1357,9 +1362,9 @@ def phonons3(
for ii, scell in enumerate(supercells):
print("scell=", ii)
ase_atoms = AseAtoms(
- symbols=scell.get_chemical_symbols(),
- scaled_positions=scell.get_scaled_positions(),
- cell=scell.get_cell(),
+ symbols=scell.symbols, # property form (phonopy 2.10/2.48/4.x)
+ scaled_positions=scell.scaled_positions,
+ cell=scell.cell,
pbc=True,
)
ase_atoms.calc = calc
diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index 25a3fc2..ebeb065 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -16,7 +16,7 @@
dgl = None
fn = None
AvgPooling = None
- print("WARNING: No DGL, might fail")
+ # print("WARNING: No DGL, might fail")
import numpy as np
import torch
diff --git a/alignn/models/alignn_atomwise_pure_smooth.py b/alignn/models/alignn_atomwise_pure_smooth.py
new file mode 100644
index 0000000..2d2683a
--- /dev/null
+++ b/alignn/models/alignn_atomwise_pure_smooth.py
@@ -0,0 +1,717 @@
+"""Pure-PyTorch ALIGNN-atomwise with smooth radial/angular bases.
+
+A drop-in cousin of ``alignn_atomwise_pure`` with:
+
+ * **Learnable Bessel radial basis** (replaces fixed Gaussian RBF for
+ bond distances).
+ * **Learnable Fourier angular basis** (replaces cosine RBF for
+ triplet angles).
+ * **Polynomial smooth cutoff envelope** applied to the radial basis,
+ so bond features (and forces) go continuously to zero at the cutoff
+ (default coeff 5).
+ * **AtomRef composition baseline**: a linear per-element offset is
+ subtracted before the GNN and added back at the end, so the network
+ only has to learn the residual chemistry.
+ * **mlp_first / site-energy readout**: project per-atom features to a
+ scalar site energy and sum (instead of mean-pool then project).
+ Energy is extensive by construction and per-atom contributions are
+ available for defect / segregation analysis.
+
+Same forward output keys as ``ALIGNNAtomWisePure``; same TorchGraph
+input contract; same TorchScript entry point.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Literal, Optional
+
+import numpy as np
+import torch
+from torch import nn
+from torch.autograd import grad
+
+from alignn.models.alignn_atomwise_pure import (
+ ALIGNNConvPure,
+ _bond_cosines,
+ _make_bond_conv,
+ scatter_mean,
+ scatter_sum,
+)
+from alignn.models.utils import MLPLayer
+from alignn.torch_graph_builder import TorchGraph, torchgraph_from_dgl
+from alignn.utils import BaseSettings
+
+
+# =====================================================================
+# Smooth bases
+# =====================================================================
+
+
+class CutoffPolynomial(nn.Module):
+ """Polynomial soft-cutoff envelope.
+
+ Decays smoothly from 1 at r=0 to 0 at r=cutoff with continuous
+ first two derivatives. ``coeff=0`` disables it (returns 1).
+ """
+
+ def __init__(self, cutoff: float = 8.0, coeff: float = 5.0) -> None:
+ super().__init__()
+ self.cutoff = float(cutoff)
+ self.p = float(coeff)
+ self.a = -(self.p + 1) * (self.p + 2) / 2
+ self.b = self.p * (self.p + 2)
+ self.c = -self.p * (self.p + 1) / 2
+
+ def forward(self, r: torch.Tensor) -> torch.Tensor:
+ if self.p == 0:
+ return torch.ones_like(r)
+ rs = r / self.cutoff
+ env = (
+ 1.0
+ + self.a * rs**self.p
+ + self.b * rs ** (self.p + 1)
+ + self.c * rs ** (self.p + 2)
+ )
+ return torch.where(rs < 1.0, env, torch.zeros_like(rs))
+
+
+class RadialBessel(nn.Module):
+ """Learnable Bessel radial basis with smooth polynomial cutoff."""
+
+ def __init__(
+ self,
+ num_radial: int = 80,
+ cutoff: float = 8.0,
+ learnable: bool = True,
+ smooth_cutoff_coeff: float = 5.0,
+ ) -> None:
+ super().__init__()
+ self.num_radial = num_radial
+ self.cutoff = float(cutoff)
+ self.inv_cutoff = 1.0 / float(cutoff)
+ self.norm_const = (2.0 * self.inv_cutoff) ** 0.5
+
+ freq = np.pi * torch.arange(1, num_radial + 1, dtype=torch.float)
+ if learnable:
+ self.frequencies = nn.Parameter(freq, requires_grad=True)
+ else:
+ self.register_buffer("frequencies", freq)
+
+ self.smooth = (
+ CutoffPolynomial(cutoff=cutoff, coeff=smooth_cutoff_coeff)
+ if smooth_cutoff_coeff and smooth_cutoff_coeff > 0
+ else None
+ )
+
+ @property
+ def out_features(self) -> int:
+ return self.num_radial
+
+ def forward(self, dist: torch.Tensor) -> torch.Tensor:
+ # dist: (E,)
+ d = dist.unsqueeze(-1).clamp_min(1e-8)
+ out = (
+ self.norm_const
+ * torch.sin(self.frequencies * d * self.inv_cutoff)
+ / d
+ )
+ if self.smooth is not None:
+ out = out * self.smooth(d)
+ return out
+
+
+class FourierAngular(nn.Module):
+ """Learnable Fourier angular basis on theta in [-pi, pi].
+
+ Input is bond-angle cosine in [-1, 1]; converted to theta via acos.
+ Output features = 1 + 2*order.
+ """
+
+ def __init__(self, order: int = 9, learnable: bool = True) -> None:
+ super().__init__()
+ self.order = order
+ freq = torch.arange(1, order + 1, dtype=torch.float)
+ if learnable:
+ self.frequencies = nn.Parameter(freq, requires_grad=True)
+ else:
+ self.register_buffer("frequencies", freq)
+
+ @property
+ def out_features(self) -> int:
+ return 1 + 2 * self.order
+
+ def forward(self, cos_theta: torch.Tensor) -> torch.Tensor:
+ theta = torch.acos(cos_theta.clamp(-1.0 + 1e-7, 1.0 - 1e-7))
+ out = theta.new_zeros(theta.shape[0], 1 + 2 * self.order)
+ out[:, 0] = 0.7071067811865475 # 1/sqrt(2)
+ tmp = torch.outer(theta, self.frequencies)
+ out[:, 1 : self.order + 1] = torch.sin(tmp)
+ out[:, self.order + 1 :] = torch.cos(tmp)
+ return out / 1.7724538509055159 # sqrt(pi)
+
+
+# =====================================================================
+# Config
+# =====================================================================
+
+
+class ALIGNNAtomWisePureSmoothConfig(BaseSettings):
+ """Mirrors ALIGNNAtomWisePureConfig and adds smooth-basis options."""
+
+ name: Literal["alignn_atomwise_pure_smooth"]
+ 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
+ # Site-energy readout (mlp_first) makes energy extensive natively, so
+ # the multiplicative natoms is OFF by default in this variant.
+ energy_mult_natoms: bool = False
+ include_pos_deriv: 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
+ use_penalty: bool = True
+ extra_features: int = 0
+ exponent: int = 5
+ penalty_factor: float = 0.5
+ penalty_threshold: float = 1.0
+ additional_output_features: int = 0
+ additional_output_weight: float = 0.0
+ conv_type: Literal["alignn", "n_alignn", "t_alignn"] = "alignn"
+ num_heads: int = 1
+
+ # --- smooth-basis additions ---
+ # Smooth polynomial cutoff on the radial basis. Set coeff=0 to disable.
+ radial_cutoff: float = 8.0
+ smooth_cutoff_coeff: float = 5.0
+ # "gaussian"/"cosine" fall back to the original alignn_atomwise_pure
+ # behaviour.
+ radial_basis: Literal["bessel", "gaussian"] = "bessel"
+ angular_basis: Literal["fourier", "cosine"] = "fourier"
+ radial_learnable: bool = True
+ angular_learnable: bool = True
+ # Linear per-element composition baseline (atomref). Learned jointly.
+ use_atomref: bool = True
+ # Per-atom site-energy readout (mlp_first=True semantics):
+ # site_e = fc(x); energy = sum_i site_e_i.
+ mlp_first: bool = True
+
+ class Config:
+ env_prefix = "jv_model"
+
+
+# =====================================================================
+# Model
+# =====================================================================
+
+
+def _as_torchgraph(x):
+ if isinstance(x, TorchGraph):
+ return x
+ return torchgraph_from_dgl(x)
+
+
+class ALIGNNAtomWisePureSmooth(nn.Module):
+ """DGL-free ALIGNN-FF with smooth bases & site energies."""
+
+ def __init__(
+ self,
+ config: ALIGNNAtomWisePureSmoothConfig = (
+ ALIGNNAtomWisePureSmoothConfig(name="alignn_atomwise_pure_smooth")
+ ),
+ ):
+ super().__init__()
+ self.config = config
+ self.classification = config.classification
+ if self.config.gradwise_weight == 0:
+ self.config.calculate_gradient = False
+
+ # --- atom embedding ---
+ self.atom_embedding = MLPLayer(
+ config.atom_input_features, config.hidden_features
+ )
+
+ # --- AtomRef (composition baseline) ---
+ # Linear in atom_features, init to zero so it starts neutral and
+ # learns a per-element offset by gradient descent.
+ if config.use_atomref:
+ self.atomref = nn.Linear(
+ config.atom_input_features, 1, bias=False
+ )
+ nn.init.zeros_(self.atomref.weight)
+ else:
+ self.atomref = None
+
+ # --- radial basis ---
+ self.radial_envelope: Optional[CutoffPolynomial] = None
+ if config.radial_basis == "bessel":
+ self.radial = RadialBessel(
+ num_radial=config.edge_input_features,
+ cutoff=config.radial_cutoff,
+ learnable=config.radial_learnable,
+ smooth_cutoff_coeff=config.smooth_cutoff_coeff,
+ )
+ edge_in_dim = self.radial.out_features
+ self._radial_kind = "bessel"
+ else:
+ from alignn.models.utils import RBFExpansion
+
+ self.radial = RBFExpansion(
+ vmin=0,
+ vmax=config.radial_cutoff,
+ bins=config.edge_input_features,
+ )
+ edge_in_dim = config.edge_input_features
+ self._radial_kind = "gaussian"
+ # Optional standalone smooth envelope multiplied onto the
+ # Gaussian expansion (off by default for this branch).
+ if config.smooth_cutoff_coeff and config.smooth_cutoff_coeff > 0:
+ self.radial_envelope = CutoffPolynomial(
+ cutoff=config.radial_cutoff,
+ coeff=config.smooth_cutoff_coeff,
+ )
+ else:
+ self.radial_envelope = None
+
+ self.edge_embedding = nn.Sequential(
+ MLPLayer(edge_in_dim, config.embedding_features),
+ MLPLayer(config.embedding_features, config.hidden_features),
+ )
+
+ # --- angular basis ---
+ if config.angular_basis == "fourier":
+ # Pick order so the basis has ~triplet_input_features dims.
+ order = max(1, (config.triplet_input_features - 1) // 2)
+ self.angular = FourierAngular(
+ order=order, learnable=config.angular_learnable
+ )
+ angle_in_dim = self.angular.out_features
+ self._angular_kind = "fourier"
+ else:
+ from alignn.models.utils import RBFExpansion
+
+ self.angular = RBFExpansion(
+ vmin=-1, vmax=1.0, bins=config.triplet_input_features
+ )
+ angle_in_dim = config.triplet_input_features
+ self._angular_kind = "cosine"
+
+ self.angle_embedding = nn.Sequential(
+ MLPLayer(angle_in_dim, config.embedding_features),
+ MLPLayer(config.embedding_features, config.hidden_features),
+ )
+
+ # --- conv stack ---
+ 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)
+ ]
+ )
+
+ # --- heads ---
+ 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)
+
+ self.link_name: str = config.link
+ if config.link == "log":
+ self.fc.bias.data = torch.tensor(np.log(0.7), dtype=torch.float)
+
+ self.register_buffer(
+ "_species_table",
+ torch.zeros(120, config.atom_input_features, dtype=torch.float32),
+ )
+
+ # Typed mirrors of pydantic config for TorchScript-friendly paths.
+ 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)
+ self.mlp_first: bool = bool(config.mlp_first)
+
+ # ----- species lookup (LAMMPS path) -----
+
+ @torch.jit.ignore
+ def register_species_table(
+ self, atom_features: str = "cgcnn", max_z: int = 119
+ ) -> None:
+ from ase.data import chemical_symbols
+ from jarvis.core.specie import get_node_attributes
+
+ F_dim = int(self.config.atom_input_features)
+ rows = np.zeros((max_z + 1, F_dim), 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_dim:
+ continue
+ rows[z] = feat
+ except Exception:
+ continue
+ self._species_table.resize_(rows.shape).copy_(
+ torch.as_tensor(rows, dtype=torch.float32)
+ )
+
+ # ----- internal embedding helpers -----
+
+ def _embed_radial(self, bondlength: torch.Tensor) -> torch.Tensor:
+ if self._radial_kind == "bessel":
+ return self.edge_embedding(self.radial(bondlength))
+ feats = self.radial(bondlength)
+ if self.radial_envelope is not None:
+ feats = feats * self.radial_envelope(bondlength).unsqueeze(-1)
+ return self.edge_embedding(feats)
+
+ def _embed_angular(self, cos_theta: torch.Tensor) -> torch.Tensor:
+ return self.angle_embedding(self.angular(cos_theta))
+
+ # ----- tensor-only forward (LAMMPS / TorchScript path) -----
+
+ @torch.jit.export
+ def forward_tensors_z(
+ self,
+ positions: torch.Tensor,
+ lattice: torch.Tensor,
+ atomic_numbers: torch.Tensor,
+ src: torch.Tensor,
+ dst: torch.Tensor,
+ shift: torch.Tensor,
+ compute_stress: bool = False,
+ ) -> Dict[str, torch.Tensor]:
+ 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,
+ lattice: torch.Tensor,
+ atom_features: torch.Tensor,
+ src: torch.Tensor,
+ dst: torch.Tensor,
+ shift: torch.Tensor,
+ compute_stress: bool = False,
+ ) -> Dict[str, torch.Tensor]:
+ 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: A=(u,v) -> 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).
+ 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 (Bessel/Fourier with smooth cutoff if configured).
+ x = self.atom_embedding(atom_features)
+ y = self._embed_radial(bondlength)
+ z = self._embed_angular(h_cos)
+
+ 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)
+
+ # Site-energy readout: fc per atom then sum (extensive).
+ if self.mlp_first:
+ site_e = self.fc(x).squeeze(-1) # (N,)
+ en = site_e.sum().unsqueeze(0) # (1,)
+ else:
+ h_graph = x.mean(dim=0, keepdim=True)
+ en = self.fc(h_graph).squeeze(-1)
+ if self.energy_mult_natoms:
+ en = en * float(num_nodes)
+
+ # AtomRef composition baseline added back as a per-element sum.
+ if self.atomref is not None:
+ en = en + self.atomref(atom_features).sum().unsqueeze(0)
+
+ 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()}
+
+ 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
+
+ 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:
+ 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):
+ 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)
+
+ if self.config.lg_on_fly and len(self.alignn_layers) > 0:
+ h = _bond_cosines(r[lg.src], r[lg.dst])
+ z = self._embed_angular(h)
+ elif len(self.alignn_layers) > 0:
+ z = self._embed_angular(lg.edata["h"])
+ else:
+ z = None
+
+ atom_feats = g.ndata["atom_features"]
+ x = self.atom_embedding(atom_feats)
+ y = self._embed_radial(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)
+
+ node_bid = g.node_batch_id
+ B = g.batch_size
+ natoms = (
+ g.batch_num_nodes
+ if g.batch_num_nodes is not None
+ else torch.tensor([g.num_nodes], device=r.device)
+ )
+
+ # --- energy readout ---
+ if self.config.mlp_first:
+ site_e = self.fc(x) # (N, output_features)
+ en_out = scatter_sum(site_e, node_bid, B) # (B, output_features)
+ if self.config.output_features == 1:
+ en_out = en_out.squeeze(-1)
+ out = en_out # for link-fn / classification compatibility below
+ else:
+ h_graph = scatter_mean(x, node_bid, B)
+ out = self.fc(h_graph)
+ if self.config.output_features == 1:
+ out = out.squeeze(-1)
+ en_out = out
+ if self.config.energy_mult_natoms:
+ en_out = out * natoms.to(out.dtype)
+
+ # --- AtomRef composition baseline (per-graph sum) ---
+ if self.atomref is not None:
+ atomref_e = self.atomref(atom_feats).squeeze(-1) # (N,)
+ en_out = en_out + scatter_sum(atomref_e, node_bid, B)
+
+ if self.config.use_penalty:
+ # Penalty is per-edge; aggregate per-graph via edge_batch_id.
+ penalties = torch.where(
+ bondlength < self.config.penalty_threshold,
+ self.config.penalty_factor
+ * (self.config.penalty_threshold - bondlength),
+ torch.zeros_like(bondlength),
+ )
+ if g.edge_batch_id is not None:
+ en_out = en_out + scatter_sum(penalties, g.edge_batch_id, B)
+ else:
+ 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(
+ scatter_mean(x, node_bid, B)
+ )
+
+ 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
+
+ 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:
+ 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)
+
+ # Keep the same return contract as ALIGNNAtomWisePure.
+ return {
+ "out": en_out if self.config.mlp_first else out,
+ "additional": additional_out,
+ "grad": forces,
+ "stresses": stress,
+ "atomwise_pred": atomwise_pred,
+ }
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index c05a9d5..8fd5eb6 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -235,6 +235,15 @@
+ " construction.",
)
+parser.add_argument(
+ "--pure_torch",
+ action="store_true",
+ help="Use the DGL-free pure-torch model path (ALIGNNAtomWisePure). "
+ "Model names come from all_models_alignn_atomwise.json, e.g. 'mps', "
+ "'formation_energy_peratom', 'mbj_bandgap'. Auto-enabled if DGL is "
+ "not installed.",
+)
+
device = "cpu"
if torch.cuda.is_available():
@@ -331,14 +340,155 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
return model
+def _dgl_available():
+ """Return True if DGL can be imported."""
+ try:
+ import dgl # noqa: F401
+
+ return True
+ except Exception:
+ return False
+
+
+def _atom_features_from_config(config):
+ """Best-effort atom_features string for the pure-torch graph builder."""
+ af = config.get("atom_features")
+ if af:
+ return af
+ n = config.get("model", {}).get("atom_input_features", 92)
+ return "atomic_number" if int(n) == 1 else "cgcnn"
+
+
+def load_pure_torch_model(model_dir, device="cpu"):
+ """Load a pure-torch ALIGNN model (+ config) from a directory.
+
+ The directory must contain ``config.json`` and ``best_model.pt``
+ (the layout produced by ``alignn.ff.ff.get_figshare_model_ff``).
+ """
+ config = loadjson(os.path.join(model_dir, "config.json"))
+ mcfg = dict(config["model"])
+ name = mcfg.get("name", "alignn_atomwise_pure")
+ state = torch.load(
+ os.path.join(model_dir, "best_model.pt"),
+ map_location=device,
+ weights_only=False,
+ )
+ if isinstance(state, dict) and "model" in state:
+ state = state["model"]
+ if name == "alignn_atomwise_pure_smooth":
+ from alignn.models.alignn_atomwise_pure_smooth import (
+ ALIGNNAtomWisePureSmooth,
+ ALIGNNAtomWisePureSmoothConfig,
+ )
+
+ model = ALIGNNAtomWisePureSmooth(
+ ALIGNNAtomWisePureSmoothConfig(**mcfg)
+ )
+ else:
+ from alignn.models.alignn_atomwise_pure import (
+ ALIGNNAtomWisePure,
+ ALIGNNAtomWisePureConfig,
+ )
+
+ mcfg["name"] = "alignn_atomwise_pure"
+ model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(**mcfg))
+ missing, unexpected = model.load_state_dict(state, strict=False)
+ if missing or unexpected:
+ print(
+ f"[pure_torch] load_state_dict: missing={len(missing)} "
+ f"unexpected={len(unexpected)}"
+ )
+ model = model.to(device).eval()
+ return model, config
+
+
+def get_figshare_model_pure(model_name="mps", device="cpu"):
+ """Download a pure-torch ALIGNN model from figshare and load it.
+
+ ``model_name`` must be a key of
+ ``alignn/ff/all_models_alignn_atomwise.json`` (e.g. ``mps``,
+ ``formation_energy_peratom``, ``mbj_bandgap``, ``bulk_modulus_kv``).
+ """
+ from alignn.ff.ff import get_figshare_model_ff
+
+ model_dir = get_figshare_model_ff(model_name=model_name)
+ return load_pure_torch_model(model_dir, device=device)
+
+
+def get_prediction_pure(
+ model_name="mps",
+ atoms=None,
+ cutoff=None,
+ max_neighbors=None,
+ device="cpu",
+):
+ """Single-structure prediction using the DGL-free (pure-torch) path."""
+ from alignn.torch_graph_builder import build_pure_torch_graph
+
+ if os.path.isdir(model_name):
+ model, config = load_pure_torch_model(
+ model_dir=model_name, device=device
+ )
+ else:
+ model, config = get_figshare_model_pure(model_name, device=device)
+
+ # Prefer the model's training cutoff / max_neighbors (correct for an
+ # ML potential); fall back to the supplied values.
+ cut = float(config.get("cutoff", cutoff if cutoff is not None else 8.0))
+ mn = int(
+ config.get(
+ "max_neighbors", max_neighbors if max_neighbors is not None else 12
+ )
+ )
+ atom_features = _atom_features_from_config(config)
+ print(
+ f"[pure_torch] cutoff={cut} max_neighbors={mn} "
+ f"atom_features={atom_features}"
+ )
+ g, lg = build_pure_torch_graph(
+ atoms=atoms,
+ two_body_cutoff=cut,
+ max_neighbors=mn,
+ atom_features=atom_features,
+ compute_line_graph=True,
+ device=device,
+ )
+ lat = (
+ torch.tensor(atoms.lattice_mat)
+ .type(torch.get_default_dtype())
+ .to(device)
+ )
+ out = model([g, lg, lat])
+ out_data = out["out"] if isinstance(out, dict) else out
+ out_data = out_data.detach().cpu().numpy().flatten().tolist()
+ return out_data
+
+
def get_prediction(
model_name="jv_formation_energy_peratom_alignn",
atoms=None,
cutoff=8,
max_neighbors=12,
+ pure_torch=False,
):
- """Get model prediction on a single structure."""
+ """Get model prediction on a single structure.
+
+ If ``pure_torch`` is True, or if DGL is not installed, the DGL-free
+ ``ALIGNNAtomWisePure`` path is used (models from
+ ``all_models_alignn_atomwise.json``, e.g. ``mps``,
+ ``formation_energy_peratom``).
+ """
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ if pure_torch or not _dgl_available():
+ if not pure_torch:
+ print("DGL not available; using pure-torch model path.")
+ return get_prediction_pure(
+ model_name=model_name,
+ atoms=atoms,
+ cutoff=cutoff,
+ max_neighbors=max_neighbors,
+ device=device,
+ )
if os.path.isdir(model_name):
# import torch
@@ -510,6 +660,7 @@ def atoms_to_graph(atoms):
cutoff=float(cutoff),
max_neighbors=int(max_neighbors),
atoms=atoms,
+ pure_torch=args.pure_torch,
)
print("Predicted value:", model_name, file_path, out_data)
diff --git a/alignn/scripts/torch/lammps_package/README.md b/alignn/scripts/torch/lammps_package/README.md
new file mode 100644
index 0000000..60f17ae
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/README.md
@@ -0,0 +1,107 @@
+# ML-ALIGNN — LAMMPS package scaffold
+
+This directory is a **drop-in LAMMPS package** for the native `pair_style alignn`
+(ALIGNN-FF via libtorch). The layout mirrors the LAMMPS source tree so the files
+can be copied straight into a LAMMPS checkout, and it is structured to become an
+upstream pull request to [lammps/lammps](https://github.com/lammps/lammps).
+
+It packages the same `pair_alignn.{cpp,h}` that `build_lammps_alignn.sh` drops
+into `src/`, but as a *proper optional package* (`make yes-ml-alignn` /
+`-D PKG_ML-ALIGNN=yes`) instead of an unconditional edit to the main build.
+
+## Layout
+
+```
+lammps_package/
+├── src/ML-ALIGNN/
+│ ├── pair_alignn.cpp # pair style implementation
+│ ├── pair_alignn.h # pair style header
+│ ├── Install.sh # legacy make-build install hook
+│ └── README # package description (LAMMPS convention)
+├── cmake/Modules/Packages/
+│ └── ML-ALIGNN.cmake # finds + links libtorch when pkg enabled
+├── doc/src/
+│ └── pair_alignn.rst # user documentation
+└── examples/PACKAGES/alignn/
+ ├── in.alignn.si # runnable NVE example
+ └── README # how to generate inputs + run
+```
+
+Each path under here maps 1:1 onto the same path in a LAMMPS checkout.
+
+## Installing into a LAMMPS checkout
+
+From this directory, with `$LAMMPS` pointing at your LAMMPS clone:
+
+```bash
+LAMMPS=~/lammps
+
+cp -r src/ML-ALIGNN "$LAMMPS/src/"
+cp cmake/Modules/Packages/ML-ALIGNN.cmake "$LAMMPS/cmake/Modules/Packages/"
+cp doc/src/pair_alignn.rst "$LAMMPS/doc/src/"
+cp -r examples/PACKAGES/alignn "$LAMMPS/examples/PACKAGES/"
+```
+
+Then register the package name in the CMake package list
+(`$LAMMPS/cmake/CMakeLists.txt`) — add `ML-ALIGNN` to the `STANDARD_PACKAGES`
+set so `-D PKG_ML-ALIGNN=yes` is recognized and the `.cmake` module above is
+auto-included.
+
+### Build (CMake, recommended)
+
+```bash
+cd "$LAMMPS"
+mkdir build && cd build
+TORCH_CMAKE=$(python -c "import torch,os;print(os.path.dirname(torch.__file__))")/share/cmake
+cmake ../cmake \
+ -D PKG_ML-ALIGNN=yes \
+ -D CMAKE_PREFIX_PATH="$TORCH_CMAKE" \
+ -D CMAKE_BUILD_TYPE=Release
+cmake --build . -j
+```
+
+If your libtorch uses the old C++ ABI, also export
+`TORCH_CXX11_ABI=$(python -c "import torch;print(int(torch._C._GLIBCXX_USE_CXX11_ABI))")`
+before configuring (the CMake module reads it).
+
+### Build (legacy make)
+
+```bash
+cd "$LAMMPS/src"
+make yes-ml-alignn
+# add libtorch include/link flags to your Makefile., then:
+make
+```
+
+## Note on the fully-automated path
+
+For local development you usually do **not** need this package layout — the
+script `../build_lammps_alignn.sh` clones LAMMPS, drops the pair style into
+`src/`, links libtorch, builds, and installs the Python module in one shot. This
+directory exists specifically to package the same code for an **upstream LAMMPS
+contribution**.
+
+## Checklist before opening an upstream PR
+
+The files here are complete enough to build, but mainline LAMMPS acceptance
+needs more. Open items:
+
+- [ ] **MPI / domain decomposition.** The current pair style is single-rank: it
+ builds the model graph from local atoms only and does not include ghost
+ atoms as graph nodes (the model's energy is a global pool, so ghosts would
+ inflate it). Multi-rank support is required for upstream acceptance.
+- [ ] Register `ML-ALIGNN` in `cmake/CMakeLists.txt` (`STANDARD_PACKAGES`) and
+ add it to `src/Makefile`'s package lists / `src/.gitignore`.
+- [ ] Add the package + pair style to the doc indices: `doc/src/Packages_*.rst`,
+ `doc/src/Commands_pair.rst`, and the `pair_style.rst` overview table.
+- [ ] Add a regression/unit test under `unittest/` or a YAML test for the pair
+ style, and wire the example into the example test harness.
+- [ ] Make the libtorch dependency robust/optional in CMake the way `ML-IAP`'s
+ PyTorch backend does, and document supported libtorch versions/ABI.
+- [ ] Confirm licensing/attribution headers match LAMMPS conventions (GPL-2.0,
+ contributor credit) on `pair_alignn.{cpp,h}`.
+- [ ] Validate `pair_alignn` against the Python reference with
+ `../validate_pair_alignn.py` and include the numbers in the PR.
+
+See the upstream LAMMPS guide for contributing packages:
+https://docs.lammps.org/Modify_contribute.html
diff --git a/alignn/scripts/torch/lammps_package/cmake/Modules/Packages/ML-ALIGNN.cmake b/alignn/scripts/torch/lammps_package/cmake/Modules/Packages/ML-ALIGNN.cmake
new file mode 100644
index 0000000..7fba9b4
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/cmake/Modules/Packages/ML-ALIGNN.cmake
@@ -0,0 +1,23 @@
+# ML-ALIGNN package: native `pair_style alignn` backed by libtorch.
+#
+# This file is auto-included by the LAMMPS CMake build when the package is
+# enabled with -D PKG_ML-ALIGNN=yes. The pair_alignn.{cpp,h} sources in
+# src/ML-ALIGNN/ are picked up by the standard per-package source glob, so all
+# this module has to do is locate libtorch and link it.
+#
+# Point CMake at libtorch via CMAKE_PREFIX_PATH, e.g.
+# -D CMAKE_PREFIX_PATH=$(python -c "import torch,os;print(os.path.dirname(torch.__file__))")/share/cmake
+# or download the standalone libtorch and pass its root.
+
+find_package(Torch REQUIRED)
+
+target_link_libraries(lammps PRIVATE ${TORCH_LIBRARIES})
+target_compile_features(lammps PRIVATE cxx_std_17)
+
+# The libtorch ABI flag must match how libtorch itself was compiled. The conda
+# / pip PyTorch wheels are built with the new C++11 ABI for recent versions but
+# this has varied; expose it as an env var so users can override.
+if(DEFINED ENV{TORCH_CXX11_ABI})
+ target_compile_definitions(lammps PRIVATE
+ _GLIBCXX_USE_CXX11_ABI=$ENV{TORCH_CXX11_ABI})
+endif()
diff --git a/alignn/scripts/torch/lammps_package/doc/src/pair_alignn.rst b/alignn/scripts/torch/lammps_package/doc/src/pair_alignn.rst
new file mode 100644
index 0000000..114b14f
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/doc/src/pair_alignn.rst
@@ -0,0 +1,123 @@
+.. index:: pair_style alignn
+
+pair_style alignn command
+=========================
+
+Syntax
+""""""
+
+.. code-block:: LAMMPS
+
+ pair_style alignn cutoff keyword
+
+* alignn = name of this pair style
+* cutoff = cutoff distance for the model graph (distance units, Angstroms)
+* keyword = optional maximum number of neighbors per atom
+
+Examples
+""""""""
+
+.. code-block:: LAMMPS
+
+ pair_style alignn 5.0 12
+ pair_coeff * * alignn_ff.pt Si
+
+ pair_style alignn 5.0 12
+ pair_coeff * * alignn_ff.pt Si O
+
+Description
+"""""""""""
+
+Style *alignn* computes interatomic interactions using the ALIGNN-FF
+(Atomistic Line Graph Neural Network Force-Field) machine-learning potential
+:ref:`(Choudhary2023) `. The potential is evaluated by calling a
+TorchScript-exported ALIGNN-FF model through the libtorch (PyTorch C++) API. At
+each timestep the pair style builds the atom graph from the LAMMPS neighbor
+list, passes atomic numbers, positions, the cell, and periodic edge shifts to
+the model, and reads back the total energy, per-atom forces, and (optionally)
+the virial stress.
+
+Because evaluation happens entirely in compiled C++/libtorch there is no Python
+interpreter in the MD loop. If a CUDA-capable GPU and a CUDA build of libtorch
+are available the model runs on the GPU automatically; otherwise it runs on the
+CPU.
+
+The mandatory *cutoff* argument and the optional max-neighbors keyword define
+the graph the model sees and **must match the values the model was trained
+with**. These are stored as the top-level ``cutoff`` and ``max_neighbors``
+fields of the model's ``config.json``. For the default Materials-Project
+``mps`` ALIGNN-FF model these are ``5.0`` and ``12``. Using values that differ
+from training puts the model out of distribution and typically destabilizes the
+simulation.
+
+Only a single :doc:`pair_coeff ` command is used with this style,
+with the form:
+
+.. code-block:: LAMMPS
+
+ pair_coeff * * model.pt Sym1 Sym2 ...
+
+where:
+
+* the two asterisks map to all LAMMPS atom types,
+* ``model.pt`` is the path to the TorchScript model file, and
+* ``Sym1 Sym2 ...`` are chemical element symbols, one per LAMMPS atom type, in
+ type order. These map each numeric atom type to the atomic number the model
+ expects.
+
+The TorchScript model file is produced from a trained ALIGNN-FF checkpoint with
+the ``export_torchscript.py`` utility shipped with the ALIGNN package. The
+default ``mps`` model can be downloaded and exported with the ``get_model.py``
+helper in the ALIGNN ``examples/lammps`` directory.
+
+----------
+
+Mixing, shift, table, tail correction, restart, rRESPA info
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+This pair style does not support mixing. It is a manybody potential defined by
+a single ``pair_coeff`` command as described above.
+
+This pair style does not support the :doc:`pair_modify ` shift,
+table, and tail options.
+
+This pair style does not write its information to :doc:`binary restart files
+`, since it is stored in the external model file. Thus, you need to
+re-specify the pair_style and pair_coeff commands in an input script that reads
+a restart file.
+
+This pair style can only be used via the *pair* keyword of the :doc:`run_style
+respa ` command. It does not support the *inner*, *middle*, *outer*
+keywords.
+
+Restrictions
+""""""""""""
+
+This pair style is part of the ML-ALIGNN package. It is only enabled if LAMMPS
+was built with that package. See the :doc:`Build package ` page
+for more info.
+
+This pair style requires libtorch (the PyTorch C++ distribution) to build and
+run. The libtorch version's C++ ABI must be compatible with the compiler used
+to build LAMMPS.
+
+The current implementation runs on a single MPI rank (it builds the model graph
+from local atoms only). It requires ``newton on`` and atom IDs.
+
+Related commands
+""""""""""""""""
+
+:doc:`pair_coeff `
+
+Default
+"""""""
+
+The optional max-neighbors keyword defaults to 12.
+
+----------
+
+.. _Choudhary2023:
+
+**(Choudhary2023)** Choudhary, DeCost, Major, Butler, Thiyagalingam, Tavazza,
+"Unified graph neural network force-field for the periodic table: solid state
+applications", Digital Discovery, 2, 346-355 (2023).
diff --git a/alignn/scripts/torch/lammps_package/examples/PACKAGES/alignn/README b/alignn/scripts/torch/lammps_package/examples/PACKAGES/alignn/README
new file mode 100644
index 0000000..69297eb
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/examples/PACKAGES/alignn/README
@@ -0,0 +1,36 @@
+ML-ALIGNN example: NVE stability of crystalline Si
+==================================================
+
+This example runs an NVE molecular-dynamics stability test on a Si supercell
+using `pair_style alignn` with the default Materials-Project (`mps`) ALIGNN-FF
+force field.
+
+Files
+-----
+
+ in.alignn.si LAMMPS input script
+
+Two input files must be generated before running (they are not checked in
+because the model is downloaded from figshare and the structure is relaxed by
+the model itself):
+
+ alignn_ff.pt TorchScript-exported ALIGNN-FF model
+ si.data relaxed Si supercell in lammps-data format
+
+Both are produced by helper scripts in the ALIGNN repository
+(alignn/examples/lammps/):
+
+ python get_model.py # downloads the mps model, writes alignn_ff.pt
+ python build_si.py # relaxes a 3x3x3 Si supercell, writes si.data
+
+Then run:
+
+ lmp -in in.alignn.si
+
+Open log.lammps and inspect the `etotal` column over the NVE production run; a
+well-behaved model conserves total energy to within a few meV/atom. The run
+also writes a trajectory (si.lammpstrj) viewable in OVITO.
+
+Note: pair_style alignn requires the cutoff and max_neighbors arguments to
+match the model's training configuration (top-level `cutoff` and
+`max_neighbors` in config.json). For the mps model these are 5.0 and 12.
diff --git a/alignn/scripts/torch/lammps_package/examples/PACKAGES/alignn/in.alignn.si b/alignn/scripts/torch/lammps_package/examples/PACKAGES/alignn/in.alignn.si
new file mode 100644
index 0000000..2c97c68
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/examples/PACKAGES/alignn/in.alignn.si
@@ -0,0 +1,53 @@
+# NVE stability test for the ML-ALIGNN pair style on crystalline Si.
+#
+# A well-trained ALIGNN-FF model evaluated with the correct cutoff conserves
+# total energy to within a few meV/atom over several picoseconds.
+#
+# Before running:
+# 1. Build LAMMPS with the ML-ALIGNN package (requires libtorch).
+# 2. Download + TorchScript-export the default mps model to alignn_ff.pt
+# and build the relaxed si.data structure. The helper scripts in the
+# ALIGNN repo do both:
+# python get_model.py # writes alignn_ff.pt
+# python build_si.py # writes si.data
+#
+# Run:
+# lmp -in in.alignn.si
+
+units metal
+atom_style atomic
+boundary p p p
+read_data si.data
+mass 1 28.0855
+
+# cutoff (5.0 Ang) and max_neighbors (12) MUST match the model's config.json.
+pair_style alignn 5.0 12
+pair_coeff * * alignn_ff.pt Si
+
+neighbor 2.0 bin
+neigh_modify every 1 delay 0 check yes
+
+# Relax cell + atoms first so we don't start at large residual stress.
+thermo 10
+thermo_style custom step temp pe ke etotal press vol
+fix relax all box/relax iso 0.0 vmax 0.001
+min_style cg
+minimize 1.0e-6 1.0e-4 200 1000
+unfix relax
+minimize 1.0e-6 1.0e-4 200 1000
+
+# Gentle thermalisation at 100 K.
+timestep 0.0005
+velocity all create 100.0 12345 mom yes rot yes dist gaussian
+fix eq all nvt temp 100.0 100.0 0.05
+thermo 50
+run 1000
+unfix eq
+
+# NVE production run: watch etotal in the log.
+reset_timestep 0
+dump traj all custom 200 si.lammpstrj id type x y z
+fix nve all nve
+run 5000
+
+write_data si_nve_final.data
diff --git a/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/Install.sh b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/Install.sh
new file mode 100644
index 0000000..62aa79d
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/Install.sh
@@ -0,0 +1,43 @@
+# Install/unInstall package files in LAMMPS
+# mode = 0/1/2 for uninstall/install/update
+#
+# Standard LAMMPS package install script (legacy `make` build system). The
+# CMake build does not use this file. It copies pair_alignn.{cpp,h} into the
+# parent src/ directory on `make yes-ml-alignn` and removes them on
+# `make no-ml-alignn`.
+
+mode=$1
+
+# enforce using portable C locale
+LC_ALL=C
+export LC_ALL
+
+# arg1 = file, arg2 = file it depends on
+
+action () {
+ if (test $mode = 0) then
+ rm -f ../$1
+ elif (! cmp -s $1 ../$1) then
+ if (test -z "$2" || test -e ../$2) then
+ cp $1 ..
+ if (test $mode = 2) then
+ echo " updating src/$1"
+ fi
+ fi
+ elif (test -n "$2") then
+ if (test ! -e ../$2) then
+ rm -f ../$1
+ fi
+ fi
+}
+
+# all package files
+
+for file in *.cpp *.h; do
+ action $file
+done
+
+# NOTE: the legacy make build will also need libtorch on the include/link path.
+# Add libtorch include dirs to the LMP_INC line and -ltorch -ltorch_cpu -lc10
+# (plus CUDA torch libs if applicable) to the LIB line of your Makefile, or use
+# the CMake build (recommended) which handles this via ML-ALIGNN.cmake.
diff --git a/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/README b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/README
new file mode 100644
index 0000000..bf7527b
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/README
@@ -0,0 +1,44 @@
+ML-ALIGNN package
+=================
+
+This package provides pair_style alignn, a native LAMMPS pair style that
+evaluates the ALIGNN-FF machine-learning interatomic potential by calling a
+TorchScript-exported model through the libtorch (PyTorch C++) API. Energies,
+forces and the virial stress are computed by the model at every MD step with no
+Python in the loop.
+
+The ALIGNN-FF model is described in:
+
+ K. Choudhary et al., "Unified graph neural network force-field for the
+ periodic table", Digital Discovery 2, 346 (2023).
+ https://doi.org/10.1039/D2DD00096B
+
+and the underlying ALIGNN architecture in:
+
+ K. Choudhary and B. DeCost, "Atomistic Line Graph Neural Network for improved
+ materials property predictions", npj Computational Materials 7, 185 (2021).
+ https://doi.org/10.1038/s41524-021-00650-1
+
+The package was contributed by the AtomGPT/JARVIS team
+(https://github.com/atomgptlab/alignn). Contact: Kamal Choudhary.
+
+Dependency
+----------
+
+This package requires libtorch (the PyTorch C++ distribution). The easiest way
+to satisfy it is to point CMake at the libtorch shipped inside a PyTorch
+install:
+
+ CMAKE_PREFIX_PATH=<...>/site-packages/torch/share/cmake
+
+See the documentation in doc/src/pair_alignn.rst and the build/run examples in
+examples/PACKAGES/alignn/ for details, and the upstream helper script
+alignn/scripts/torch/build_lammps_alignn.sh for a fully automated build.
+
+Files in this directory
+-----------------------
+
+ pair_alignn.cpp pair style implementation
+ pair_alignn.h pair style header
+ Install.sh legacy make-build install script
+ README this file
diff --git a/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/pair_alignn.cpp b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/pair_alignn.cpp
new file mode 100644
index 0000000..8f5a057
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/pair_alignn.cpp
@@ -0,0 +1,289 @@
+// 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];
+ // Integer cell offset S with disp = S0*a1 + S1*a2 + S2*a3, where
+ // the lattice rows passed to the model are
+ // a1=(hx,0,0), a2=(xy,hy,0), a3=(xz,yz,hz).
+ // Solve by back-substitution — correct for triclinic cells and
+ // reduces to diagonal rounding when xy=xz=yz=0. (The previous
+ // diagonal-only form gave wrong shifts for non-orthogonal boxes,
+ // e.g. a supercell of a primitive FCC cell, which blew up MD.)
+ (void)triclinic;
+ shz = std::round(doz / hz);
+ shy = std::round((doy - shz * yz) / hy);
+ shx = std::round((dox - shy * xy - shz * xz) / hx);
+ }
+ 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/lammps_package/src/ML-ALIGNN/pair_alignn.h b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/pair_alignn.h
new file mode 100644
index 0000000..36c2875
--- /dev/null
+++ b/alignn/scripts/torch/lammps_package/src/ML-ALIGNN/pair_alignn.h
@@ -0,0 +1,56 @@
+// ALIGNN-FF native LAMMPS pair style (ML-ALIGNN package).
+// Lives in LAMMPS src/ML-ALIGNN/ and is built when the package is enabled
+// (-D PKG_ML-ALIGNN=yes, or `make yes-ml-alignn`). Requires libtorch.
+//
+// Input script usage:
+// pair_style alignn [max_neighbors]
+// 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/pair_alignn.cpp b/alignn/scripts/torch/pair_alignn/pair_alignn.cpp
index 92ac76c..8f5a057 100644
--- a/alignn/scripts/torch/pair_alignn/pair_alignn.cpp
+++ b/alignn/scripts/torch/pair_alignn/pair_alignn.cpp
@@ -192,10 +192,17 @@ void PairAlignn::compute(int eflag, int vflag) {
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];
+ // Integer cell offset S with disp = S0*a1 + S1*a2 + S2*a3, where
+ // the lattice rows passed to the model are
+ // a1=(hx,0,0), a2=(xy,hy,0), a3=(xz,yz,hz).
+ // Solve by back-substitution — correct for triclinic cells and
+ // reduces to diagonal rounding when xy=xz=yz=0. (The previous
+ // diagonal-only form gave wrong shifts for non-orthogonal boxes,
+ // e.g. a supercell of a primitive FCC cell, which blew up MD.)
(void)triclinic;
- shx = std::round(dox / hx);
- shy = std::round(doy / hy);
shz = std::round(doz / hz);
+ shy = std::round((doy - shz * yz) / hy);
+ shx = std::round((dox - shy * xy - shz * xz) / hx);
}
buf.push_back({r2, dst_local,
static_cast(shx),
diff --git a/alignn/scripts/torch/validate_pair_alignn.py b/alignn/scripts/torch/validate_pair_alignn.py
index 90eb1cd..a730256 100644
--- a/alignn/scripts/torch/validate_pair_alignn.py
+++ b/alignn/scripts/torch/validate_pair_alignn.py
@@ -16,35 +16,68 @@
--jid JVASP-1002 --supercell 2 2 2
"""
from __future__ import annotations
-import argparse, os, tempfile
+import argparse, os, sys, 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):
+def python_reference_ase(atoms_jarvis, model_dir):
+ """ASE-calculator reference (DGL graph). NOTE: the DGL neighbor graph
+ can differ from the pure-torch k-nearest graph that pair_alignn (and
+ the model's training) uses, so this is *not* the right baseline for a
+ pure-torch model — see python_reference_pure."""
+ from alignn.ff.ff import AlignnAtomwiseCalculator
+
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
+ a.calc = AlignnAtomwiseCalculator(
+ path=model_dir, force_mult_batchsize=False
+ )
+ E = float(np.asarray(a.get_potential_energy()).reshape(-1)[0])
+ F = np.asarray(a.get_forces()) # eV/Å
+ S6 = np.asarray(a.get_stress(voigt=True)).reshape(-1) # eV/ų Voigt
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):
+def python_reference_pure(atoms_jarvis, model_dir, cutoff, max_neighbors):
+ """Pure-torch reference: same k-nearest graph + entry point
+ (forward_tensors_z) that pair_alignn calls and that the model was
+ trained with. This is the correct apples-to-apples baseline."""
+ import torch
+ from alignn.pretrained import load_pure_torch_model
+ from alignn.torch_graph_builder import torch_neighbor_list
+
+ model, cfg = load_pure_torch_model(model_dir, device="cpu")
+ af = cfg.get("atom_features", "cgcnn")
+ model.register_species_table(atom_features=af)
+
+ dt = torch.get_default_dtype()
+ pos = torch.tensor(
+ np.asarray(atoms_jarvis.cart_coords), dtype=dt
+ ).requires_grad_(True)
+ L = torch.tensor(np.asarray(atoms_jarvis.lattice_mat), dtype=dt)
+ Z = torch.tensor(np.asarray(atoms_jarvis.atomic_numbers), dtype=torch.long)
+ src, dst, shift, _ = torch_neighbor_list(
+ pos, L, cutoff=float(cutoff), max_neighbors=int(max_neighbors),
+ atoms=atoms_jarvis, use_matscipy_topology=True,
+ )
+ out = model.forward_tensors_z(pos, L, Z, src, dst, shift, True)
+ E = float(out["energy"].detach())
+ F = out["forces"].detach().cpu().numpy()
+ S = (out["stress"].detach().cpu().numpy()
+ if "stress" in out else np.zeros((3, 3)))
+ return E, F, S
+
+
+def lammps_pair_alignn(ase_atoms, ts_model_path, species_symbols, cutoff=5.0,
+ max_neighbors=12):
"""Run LAMMPS with pair_alignn at step 0, return E/F/S."""
from lammps import lammps
@@ -61,7 +94,7 @@ def lammps_pair_alignn(ase_atoms, ts_model_path, species_symbols, cutoff=5.0):
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_style alignn {cutoff} {max_neighbors}
pair_coeff * * {ts_model_path} {' '.join(species_symbols)}
neighbor 2.0 bin
neigh_modify every 1 delay 0 check yes
@@ -115,11 +148,30 @@ def main():
"(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)")
+ ap.add_argument("--max-neighbors", type=int, default=None,
+ help="k-nearest cap (read from model-dir config if unset). "
+ "Must match pair_style alignn .")
+ ap.add_argument("--fail-force", type=float, default=0.05,
+ help="eV/Å, gate: exit non-zero if max|ΔF| exceeds this.")
+ ap.add_argument("--fail-energy", type=float, default=0.05,
+ help="eV/atom, gate: exit non-zero if |ΔE|/atom exceeds "
+ "this.")
+ ap.add_argument("--reference", choices=["pure", "ase"], default="pure",
+ help="Python baseline. 'pure' (default) uses the same "
+ "pure-torch k-nearest graph + forward_tensors_z that "
+ "pair_alignn implements and the model was trained "
+ "with. 'ase' uses AlignnAtomwiseCalculator (DGL "
+ "graph) which can legitimately differ.")
args = ap.parse_args()
- if args.cutoff is None:
+ if args.cutoff is None or args.max_neighbors 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} Å")
+ _cfg = _j.load(open(f"{args.model_dir}/config.json"))
+ if args.cutoff is None:
+ args.cutoff = float(_cfg["cutoff"])
+ print(f"cutoff (from config): {args.cutoff} Å")
+ if args.max_neighbors is None:
+ args.max_neighbors = int(_cfg.get("max_neighbors", 12))
+ print(f"max_neighbors (from config): {args.max_neighbors}")
# Build system
d = get_jid_data(jid=args.jid, dataset="dft_3d")
@@ -141,14 +193,22 @@ def main():
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)
+ if args.reference == "pure":
+ print("\n→ Python reference (pure-torch: build_pure_torch_graph "
+ "+ forward_tensors_z)...")
+ E_py, F_py, S_py = python_reference_pure(
+ atoms_pert, args.model_dir, args.cutoff, args.max_neighbors
+ )
+ else:
+ print("\n→ Python reference (ASE AlignnAtomwiseCalculator, DGL graph)"
+ " — may differ from pair_alignn's pure-torch graph...")
+ E_py, F_py, S_py = python_reference_ase(atoms_pert, args.model_dir)
# --- 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)
+ cutoff=args.cutoff,
+ max_neighbors=args.max_neighbors)
# --- Compare ---
dE = abs(E_lmp - E_py)
@@ -184,6 +244,27 @@ def main():
E_lmp=E_lmp, F_lmp=F_lmp, S_lmp=S_lmp)
print("\nsaved raw arrays -> validate_pair_alignn.npz")
+ # Gate decision: a *gross* mismatch means pair_alignn and the Python
+ # reference disagree, so any MD with this .pt is untrustworthy. We use a
+ # loose force threshold here (not the strict 1e-3 above) so float32
+ # round-off doesn't trip the gate — only real bugs do.
+ n_atoms = max(len(F_py), 1)
+ dE_per_atom = dE / n_atoms
+ gross = (max_abs_dF > args.fail_force) or (dE_per_atom > args.fail_energy)
+ print(
+ f"\nGATE: max|ΔF|={max_abs_dF:.3e} (limit {args.fail_force}) "
+ f"|ΔE|/atom={dE_per_atom:.3e} (limit {args.fail_energy}) "
+ f"-> {'FAIL' if gross else 'OK'}"
+ )
+ if gross:
+ print(
+ "pair_alignn disagrees with the Python reference. Do NOT run MD "
+ "until this is fixed (check the TorchScript export and the C++ "
+ "neighbor/graph construction)."
+ )
+ sys.exit(1)
+ sys.exit(0)
+
if __name__ == "__main__":
main()
diff --git a/alignn/train.py b/alignn/train.py
index 7ea303f..57b50d6 100644
--- a/alignn/train.py
+++ b/alignn/train.py
@@ -14,6 +14,9 @@
from alignn.config import TrainingConfig
from alignn.models.alignn_atomwise import ALIGNNAtomWise
from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePure
+from alignn.models.alignn_atomwise_pure_smooth import (
+ ALIGNNAtomWisePureSmooth,
+)
from alignn.torch_graph_builder import unbatch as _graph_unbatch
from alignn.models.ealignn_atomwise import eALIGNNAtomWise
from alignn.models.alignn import ALIGNN
@@ -162,6 +165,7 @@ def train_dgl(
_model = {
"alignn_atomwise": ALIGNNAtomWise,
"alignn_atomwise_pure": ALIGNNAtomWisePure,
+ "alignn_atomwise_pure_smooth": ALIGNNAtomWisePureSmooth,
"ealignn_atomwise": eALIGNNAtomWise,
"alignn": ALIGNN,
}
diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py
index 5891e8a..5b321b0 100644
--- a/alignn/train_alignn.py
+++ b/alignn/train_alignn.py
@@ -419,7 +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",
+ use_pure_torch=config.model.name
+ in ("alignn_atomwise_pure", "alignn_atomwise_pure_smooth"),
read_existing=config.read_existing,
dtype=config.dtype,
)