Develop

.github/workflows/main.yml

@@ -24,8 +24,8 @@ jobs:
         shell: bash -l {0}
         run: |
             conda install flake8 pycodestyle pydocstyle
-            flake8 --ignore E203,W503,W605 --exclude=examples,tests,scripts --statistics --count --exit-zero alignn
-            pycodestyle --ignore E203,W503,W605 --exclude=examples,tests,scripts alignn
+            flake8 --ignore E203,W503,W605,E501 --exclude=examples,tests,scripts --statistics --count --exit-zero alignn
+            pycodestyle --ignore E203,W503,W605,E501 --exclude=examples,tests,scripts alignn
             pydocstyle --match-dir=core --match-dir=io --match-dir=io --match-dir=ai --match-dir=analysis --match-dir=db --match-dir=tasks --count alignn
 
       - name: Run pytest

alignn/__init__.py

@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2026.5.20"
+__version__ = "2026.8.6"

alignn/config.py

@@ -188,6 +188,16 @@ class TrainingConfig(BaseSettings):
     # DDP tuning. find_unused_parameters=True is slow; enable only if
     # your model has conditional branches whose gradients vary per step.
     ddp_find_unused_parameters: bool = False
+    # Resume optimizer/scheduler/epoch from <output_dir>/current_state.pt if
+    # present. Lets a run that hit the walltime continue with a single
+    # continuous LR schedule instead of restarting it (weights alone are
+    # restored via --restart_model_path, which does not restore these).
+    resume_checkpoint: bool = False
+    # LR-schedule horizon in epochs, independent of how many epochs this job
+    # runs (`epochs`). Set it to the FINAL target when a run is split into
+    # resumed segments so OneCycle spans the whole run and stays continuous
+    # across restarts. None => use `epochs` (the single-job default).
+    lr_total_epochs: Optional[int] = None
     # When True, forces cuDNN determinism (slower). Decoupled from seed so
     # you can seed for reproducibility without paying the speed cost.
     deterministic: bool = False

alignn/data.py

@@ -429,6 +429,7 @@ def get_train_val_loaders(
                 sampler=val_sampler,
                 max_neighbors=max_neighbors,
                 three_body_cutoff=three_body_cutoff,
+                read_existing=read_existing,
                 classification=classification_threshold is not None,
                 output_dir=output_dir,
                 tmp_name=tmp_name,
@@ -456,6 +457,8 @@ def get_train_val_loaders(
                 cutoff=cutoff,
                 cutoff_extra=cutoff_extra,
                 max_neighbors=max_neighbors,
+                three_body_cutoff=three_body_cutoff,
+                read_existing=read_existing,
                 classification=classification_threshold is not None,
                 output_dir=output_dir,
                 tmp_name=tmp_name,

alignn/examples/recipes/atomwise/README.md

@@ -0,0 +1,42 @@
+# ALIGNN example: Atomwise (per-atom) property
+
+Predict a **per-atom** scalar — atomic charges (Bader), site magnetic moments, etc. Each atom gets its own prediction, so the target in `id_prop.json` is a list of length `Natoms` under the `charges` key.
+
+## Run it (CPU, ~1-2 min)
+
+```bash
+# 1) generate a tiny synthetic dataset -> id_prop.json
+python make_toy_dataset.py
+
+# 2) train (10 epochs on the toy data)
+train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key target --id_key jid --atomwise_key charges
+```
+
+You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`,
+with results written to `toy_out/` (`history_val.json`, predictions, checkpoint).
+
+## Key config knobs (`config_example.json`)
+
+- `model.atomwise_output_features: 1`, `atomwise_weight: 1.0` — turn on the per-atom head.
+- `model.graphwise_weight: 0.0` — the graph-level target is unused.
+- Pass `--atomwise_key charges` (rename to your per-atom key).
+
+> `target` (graph-level) is kept as a dummy 0.0 because `graphwise_weight` is 0.
+
+## ⚠️ This is a smoke test, not a real model
+
+The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove
+the pipeline runs. For a usable model:
+
+- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data**
+  (thousands to millions of entries).
+- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`.
+- Expect training to take much longer and to need a GPU for large datasets.
+
+## Dataset format (`id_prop.json`)
+
+A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s):
+
+```json
+[{"jid": "toy-0", "atoms": {...}, "target": ...}]
+```

alignn/examples/recipes/atomwise/config_example.json

@@ -0,0 +1,83 @@
+{
+  "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+  "dataset": "user_data",
+  "target": "target",
+  "atom_features": "cgcnn",
+  "neighbor_strategy": "pure_torch",
+  "id_tag": "jid",
+  "dtype": "float32",
+  "random_seed": 123,
+  "classification_threshold": null,
+  "n_val": null,
+  "n_test": null,
+  "n_train": null,
+  "train_ratio": 0.8,
+  "val_ratio": 0.1,
+  "test_ratio": 0.1,
+  "target_multiplication_factor": null,
+  "epochs": 10,
+  "batch_size": 4,
+  "weight_decay": 1e-05,
+  "learning_rate": 0.001,
+  "filename": "A",
+  "warmup_steps": 2000,
+  "criterion": "l1",
+  "optimizer": "adamw",
+  "scheduler": "onecycle",
+  "pin_memory": false,
+  "save_dataloader": false,
+  "write_checkpoint": true,
+  "write_predictions": true,
+  "store_outputs": true,
+  "progress": true,
+  "log_tensorboard": false,
+  "standard_scalar_and_pca": false,
+  "use_canonize": true,
+  "num_workers": 0,
+  "cutoff": 8.0,
+  "cutoff_extra": 3.0,
+  "max_neighbors": 12,
+  "keep_data_order": true,
+  "normalize_graph_level_loss": false,
+  "distributed": false,
+  "data_parallel": false,
+  "n_early_stopping": null,
+  "output_dir": "temp",
+  "use_lmdb": true,
+  "model": {
+    "name": "alignn_atomwise_pure",
+    "alignn_layers": 4,
+    "gcn_layers": 4,
+    "atom_input_features": 92,
+    "edge_input_features": 80,
+    "triplet_input_features": 40,
+    "embedding_features": 64,
+    "hidden_features": 256,
+    "output_features": 1,
+    "grad_multiplier": -1,
+    "calculate_gradient": false,
+    "atomwise_output_features": 1,
+    "graphwise_weight": 0.0,
+    "gradwise_weight": 0.0,
+    "stresswise_weight": 0.0,
+    "atomwise_weight": 1.0,
+    "link": "identity",
+    "zero_inflated": false,
+    "classification": false,
+    "force_mult_natoms": false,
+    "energy_mult_natoms": false,
+    "include_pos_deriv": false,
+    "use_cutoff_function": false,
+    "inner_cutoff": 3.0,
+    "stress_multiplier": 1.0,
+    "add_reverse_forces": true,
+    "lg_on_fly": true,
+    "batch_stress": true,
+    "multiply_cutoff": false,
+    "use_penalty": true,
+    "extra_features": 0,
+    "exponent": 5,
+    "penalty_factor": 0.1,
+    "penalty_threshold": 1.0
+  }
+}

alignn/examples/recipes/atomwise/make_toy_dataset.py

@@ -0,0 +1,33 @@
+"""Generate a TOY atomwise per-atom property (e.g. charges) dataset -> id_prop.json (inline jarvis Atoms dicts).
+
+CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training
+needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below,
+supply real targets, and raise `epochs`/`batch_size` in config_example.json.
+"""
+import json, random
+from jarvis.core.atoms import Atoms
+from jarvis.db.figshare import get_jid_data
+
+random.seed(0)
+N = 40  # toy size -- increase to thousands for a real run
+
+def rattle(a, amp=0.05):
+    d = a.to_dict()
+    d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]]
+    return Atoms.from_dict(d)
+
+# base crystal (Si); swap for your own structures
+base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"])
+
+data = []
+for i in range(N):
+    a = rattle(base)
+    nat = a.num_atoms
+    # TOY per-atom target (one value per atom): replace with real charges/magmoms.
+    charges = [round(random.uniform(-1, 1), 3) for _ in range(nat)]
+    data.append({"jid": f"toy-{i}", "atoms": a.to_dict(),
+                 "target": 0.0,          # graph-level target unused (graphwise_weight=0)
+                 "charges": charges})    # per-atom target, pass via --atomwise_key charges
+
+json.dump(data, open("id_prop.json", "w"))
+print(f"wrote id_prop.json with {len(data)} entries (per-atom charges)")

alignn/examples/recipes/forcefield/README.md

@@ -0,0 +1,42 @@
+# ALIGNN example: Force field (energy + forces + stress)
+
+Train an **ALIGNN-FF** interatomic potential that outputs energy, analytic forces (gradients of the energy), and stress — usable for relaxation and molecular dynamics (including LAMMPS via `pair_alignn`). Uses the radius graph (`cutoff: 5.0`).
+
+## Run it (CPU, ~1-2 min)
+
+```bash
+# 1) generate a tiny synthetic dataset -> id_prop.json
+python make_toy_dataset.py
+
+# 2) train (10 epochs on the toy data)
+train_alignn.py --root_dir . --config_name config_example.json --output_dir toy_out --target_key energy_per_atom --force_key forces --id_key jid
+```
+
+You should see per-epoch `Train Loss` / `Val Loss` lines and a final `Test MAE`,
+with results written to `toy_out/` (`history_val.json`, predictions, checkpoint).
+
+## Key config knobs (`config_example.json`)
+
+- `model.calculate_gradient: true` — forces are the true gradient of the energy (energy-conserving).
+- `graphwise_weight` (energy), `gradwise_weight` (forces), `stresswise_weight` (stress) — the loss mixture.
+- Pass `--force_key forces`; stresses are read automatically if present.
+
+> **Energy must be per atom** (`energy_per_atom`) in `id_prop.json`, not per structure. Forces are `Natoms x 3`, stresses Voigt-6.
+
+## ⚠️ This is a smoke test, not a real model
+
+The synthetic dataset (40 rattled Si cells with fake targets) exists only to prove
+the pipeline runs. For a usable model:
+
+- Replace the toy structures/labels in `make_toy_dataset.py` with **real DFT data**
+  (thousands to millions of entries).
+- Raise `epochs` to **100-300** and `batch_size` to **32-64** in `config_example.json`.
+- Expect training to take much longer and to need a GPU for large datasets.
+
+## Dataset format (`id_prop.json`)
+
+A JSON list; each entry has a `jid`, an inline jarvis `Atoms` dict, and the target(s):
+
+```json
+[{"jid": "toy-0", "atoms": {...}, "energy_per_atom": ...}]
+```

alignn/examples/recipes/forcefield/config_example.json

@@ -0,0 +1,83 @@
+{
+  "version": "112bbedebdaecf59fb18e11c929080fb2f358246",
+  "dataset": "user_data",
+  "target": "energy_per_atom",
+  "atom_features": "cgcnn",
+  "neighbor_strategy": "pure_torch",
+  "id_tag": "jid",
+  "dtype": "float32",
+  "random_seed": 123,
+  "classification_threshold": null,
+  "n_val": null,
+  "n_test": null,
+  "n_train": null,
+  "train_ratio": 0.8,
+  "val_ratio": 0.1,
+  "test_ratio": 0.1,
+  "target_multiplication_factor": null,
+  "epochs": 10,
+  "batch_size": 4,
+  "weight_decay": 1e-05,
+  "learning_rate": 0.001,
+  "filename": "A",
+  "warmup_steps": 2000,
+  "criterion": "l1",
+  "optimizer": "adamw",
+  "scheduler": "onecycle",
+  "pin_memory": false,
+  "save_dataloader": false,
+  "write_checkpoint": true,
+  "write_predictions": true,
+  "store_outputs": true,
+  "progress": true,
+  "log_tensorboard": false,
+  "standard_scalar_and_pca": false,
+  "use_canonize": true,
+  "num_workers": 0,
+  "cutoff": 5.0,
+  "cutoff_extra": 3.0,
+  "max_neighbors": 12,
+  "keep_data_order": true,
+  "normalize_graph_level_loss": false,
+  "distributed": false,
+  "data_parallel": false,
+  "n_early_stopping": null,
+  "output_dir": "temp",
+  "use_lmdb": true,
+  "model": {
+    "name": "alignn_atomwise_pure",
+    "alignn_layers": 4,
+    "gcn_layers": 4,
+    "atom_input_features": 92,
+    "edge_input_features": 80,
+    "triplet_input_features": 40,
+    "embedding_features": 64,
+    "hidden_features": 256,
+    "output_features": 1,
+    "grad_multiplier": -1,
+    "calculate_gradient": true,
+    "atomwise_output_features": 0,
+    "graphwise_weight": 1.0,
+    "gradwise_weight": 1.0,
+    "stresswise_weight": 0.1,
+    "atomwise_weight": 0.0,
+    "link": "identity",
+    "zero_inflated": false,
+    "classification": false,
+    "force_mult_natoms": false,
+    "energy_mult_natoms": false,
+    "include_pos_deriv": false,
+    "use_cutoff_function": false,
+    "inner_cutoff": 3.5,
+    "stress_multiplier": 1.0,
+    "add_reverse_forces": true,
+    "lg_on_fly": true,
+    "batch_stress": true,
+    "multiply_cutoff": false,
+    "use_penalty": true,
+    "extra_features": 0,
+    "exponent": 5,
+    "penalty_factor": 0.1,
+    "penalty_threshold": 1.0
+  }
+}

alignn/examples/recipes/forcefield/make_toy_dataset.py

@@ -0,0 +1,35 @@
+"""Generate a TOY force-field (energy/forces/stress) dataset -> id_prop.json (inline jarvis Atoms dicts).
+
+CAUTION: this is a tiny synthetic dataset for a smoke-test only. Real training
+needs thousands+ of DFT-labeled structures and 100-300 epochs. Bump N below,
+supply real targets, and raise `epochs`/`batch_size` in config_example.json.
+"""
+import json, random
+from jarvis.core.atoms import Atoms
+from jarvis.db.figshare import get_jid_data
+
+random.seed(0)
+N = 40  # toy size -- increase to thousands for a real run
+
+def rattle(a, amp=0.05):
+    d = a.to_dict()
+    d["coords"] = [[c + random.uniform(-amp, amp) for c in xyz] for xyz in d["coords"]]
+    return Atoms.from_dict(d)
+
+# base crystal (Si); swap for your own structures
+base = Atoms.from_dict(get_jid_data(jid="JVASP-1002", dataset="dft_3d")["atoms"])
+
+data = []
+for i in range(N):
+    a = rattle(base)
+    nat = a.num_atoms
+    # TOY labels: replace with DFT energy_per_atom (eV/atom), forces (Nx3 eV/A),
+    # stresses (Voigt-6). Energy MUST be per atom.
+    energy_per_atom = -5.0 + 0.01 * i
+    forces = [[random.uniform(-0.1, 0.1) for _ in range(3)] for _ in range(nat)]
+    stresses = [random.uniform(-0.5, 0.5) for _ in range(6)]
+    data.append({"jid": f"toy-{i}", "atoms": a.to_dict(),
+                 "energy_per_atom": energy_per_atom, "forces": forces, "stresses": stresses})
+
+json.dump(data, open("id_prop.json", "w"))
+print(f"wrote id_prop.json with {len(data)} entries (energy_per_atom/forces/stresses)")