Develop

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. |
 
 <a name="pretrained"></a>
 ## Using pre-trained models

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")

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.

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")

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")

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 <dir-with-best_model.pt+config.json> \
+                            --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()

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()

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 <cutoff> <max_neighbors>`. 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.