Develop

.github/workflows/docs.yml

@@ -2,7 +2,7 @@ name: docs
 
 on:
   push:
-    branches: [main, develop]
+    branches: [main]
     paths:
       - docs/**
       - mkdocs.yml

README.md

@@ -12,6 +12,7 @@
 * [Introduction](#intro)
 * [Installation](#install)
 * [Examples](#example)
+* [Colab notebooks](#colab)
 * [Pre-trained models](#pretrained)
 * [JARVIS-ALIGNN webapp](#webapp)
 * [ALIGNN-FF & ASE Calculator](#alignnff)
@@ -48,7 +49,24 @@ See [docs/training/](docs/training/) for dataset format and training examples:
 - [Force-field training](docs/training/force-field.md)
 - [Multi-GPU training](docs/training/multi-gpu.md)
 
-Google Colab notebooks are linked from [docs/index.md](docs/index.md).
+<a name="colab"></a>
+## Colab notebooks
+
+Ready-to-run notebooks covering property prediction, force-field
+training, and pretrained-model usage. Click a badge to open in Colab.
+
+[colab-badge]: https://colab.research.google.com/assets/colab-badge.svg
+
+| Notebook | Open in Colab | Description |
+| --- | --- | --- |
+| Regression task (graph-wise prediction) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Single-output regression for 2D-material exfoliation energies. |
+| ML force-field training from scratch | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Train an ALIGNN-FF force field for Silicon. |
+| ALIGNN-FF: relaxation, EV curve, phonons, interfaces | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Pretrained ALIGNN-FF for relaxation, EV curves, phonons, and interfaces. |
+| Scaling / timing comparison | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | Scaling/timing analysis of universal MLFFs. |
+| Melt-Quench MD | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | Generate amorphous structures via molecular dynamics. |
+| Miscellaneous training tasks | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | Single-output, multi-output (phonon/electron DOS), classification, and pretrained usage. |
+| Superconductor Tc | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | Train a model for superconductor transition temperature. |
+| Build `id_prop.json` from VASP runs | [![Open In Colab][colab-badge]](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | Compile `vasprun.xml` files into `id_prop.json` for ALIGNN-FF training. |
 
 <a name="pretrained"></a>
 ## Using pre-trained models
@@ -66,7 +84,26 @@ See [docs/usage/webapps.md](docs/usage/webapps.md). Direct links: [AtomGPT ALIGN
 <a name="alignnff"></a>
 ## ALIGNN-FF ASE Calculator
 
-See [docs/usage/ase-calculator.md](docs/usage/ase-calculator.md) for example usage.
+```python
+from ase.build import bulk
+from alignn.ff.unified_calculator import (
+    AlignnUnifiedCalculator, AlignnUnifiedConfig)
+
+cfg = AlignnUnifiedConfig(
+    energy=True, forces=True, stress=True,
+    properties=["formation_energy_peratom", "optb88vdw_bandgap"],
+)
+calc = AlignnUnifiedCalculator(cfg)          # models loaded once, reused
+
+si = bulk("Si", "diamond", a=5.43); si.calc = calc
+si.get_potential_energy(); si.get_forces(); si.get_stress()
+print(calc.predictions())                    # extra property predictors
+```
+
+A single pydantic config selects the outputs (force-field
+energy/forces/stress plus any pretrained scalar property predictors).
+See [docs/usage/ase-calculator.md](docs/usage/ase-calculator.md) for more,
+and the ASE docs page *Calculators → ALIGNN*.
 
 <a name="performances"></a>
 ## Performances

alignn/__init__.py

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

alignn/config.py

@@ -7,6 +7,7 @@
 from alignn.utils import BaseSettings
 from alignn.models.alignn import ALIGNNConfig
 from alignn.models.alignn_atomwise import ALIGNNAtomWiseConfig
+from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePureConfig
 from alignn.models.ealignn_atomwise import eALIGNNAtomWiseConfig
 
 # import torch
@@ -155,7 +156,13 @@ class TrainingConfig(BaseSettings):
     target: TARGET_ENUM = "exfoliation_energy"
     atom_features: Literal["basic", "atomic_number", "cfid", "cgcnn"] = "cgcnn"
     neighbor_strategy: Literal[
-        "k-nearest", "voronoi", "radius_graph", "radius_graph_jarvis"
+        "k-nearest",
+        "voronoi",
+        "radius_graph",
+        "radius_graph_jarvis",
+        "fast_graph",
+        "torch_graph",
+        "pure_torch",
     ] = "k-nearest"
     id_tag: Literal["jid", "id", "_oqmd_entry_id"] = "jid"
 
@@ -173,6 +180,14 @@ class TrainingConfig(BaseSettings):
     target_multiplication_factor: Optional[float] = None
     epochs: int = 300
     batch_size: int = 64
+    gpu_memory_fraction: Optional[float] = None
+    use_amp: bool = False  # bf16 mixed precision (A100/H100/RTX 30+)
+    # DDP tuning. find_unused_parameters=True is slow; enable only if
+    # your model has conditional branches whose gradients vary per step.
+    ddp_find_unused_parameters: bool = False
+    # When True, forces cuDNN determinism (slower). Decoupled from seed so
+    # you can seed for reproducibility without paying the speed cost.
+    deterministic: bool = False
     weight_decay: float = 0
     learning_rate: float = 1e-2
     filename: str = "sample"
@@ -191,16 +206,23 @@ class TrainingConfig(BaseSettings):
     use_canonize: bool = True
     compute_line_graph: bool = True
     num_workers: int = 4
-    cutoff: float = 8.0
+    cutoff: float = 5.0
     cutoff_extra: float = 3.0
-    max_neighbors: int = 12
+    # Separate 3-body cutoff used by neighbor_strategy="pure_torch".
+    # When None, defaults to `cutoff`. Must be <= cutoff.
+    three_body_cutoff: Optional[float] = 3.5
+    max_neighbors: Optional[int] = 12
     keep_data_order: bool = True
     normalize_graph_level_loss: bool = False
     distributed: bool = False
     data_parallel: bool = False
     n_early_stopping: Optional[int] = None  # typically 50
     output_dir: str = os.path.abspath(".")
     use_lmdb: bool = True
+    # When True, reuse an existing LMDB cache on disk (fast, but the
+    # cache must match the current model backend + graph config).
+    # Default False: always rebuild from scratch for safety.
+    read_existing: bool = False
     # alignn_layers: int = 4
     # gcn_layers: int =4
     # edge_input_features: int= 80
@@ -212,5 +234,6 @@ class TrainingConfig(BaseSettings):
     model: Union[
         ALIGNNConfig,
         ALIGNNAtomWiseConfig,
+        ALIGNNAtomWisePureConfig,
         eALIGNNAtomWiseConfig,
     ] = ALIGNNAtomWiseConfig(name="alignn_atomwise")

alignn/data.py

@@ -10,7 +10,18 @@
 from tqdm import tqdm
 import math
 from jarvis.db.jsonutils import dumpjson
-from dgl.dataloading import GraphDataLoader
+
+try:
+    from dgl.dataloading import GraphDataLoader
+except Exception:  # pure-torch path; fall back to stdlib DataLoader
+    from torch.utils.data import DataLoader as _TorchDataLoader
+
+    def GraphDataLoader(*args, use_ddp=False, **kwargs):
+        """Fall back to torch DataLoader when DGL is unavailable."""
+        kwargs.pop("use_ddp", None)
+        return _TorchDataLoader(*args, **kwargs)
+
+
 import pickle as pk
 from sklearn.preprocessing import StandardScaler
 
@@ -145,6 +156,7 @@ def get_train_val_loaders(
     cutoff: float = 8.0,
     cutoff_extra: float = 3.0,
     max_neighbors: int = 12,
+    three_body_cutoff: Optional[float] = None,
     classification_threshold: Optional[float] = None,
     target_multiplication_factor: Optional[float] = None,
     standard_scalar_and_pca=False,
@@ -154,10 +166,15 @@ def get_train_val_loaders(
     world_size=0,
     rank=0,
     use_lmdb: bool = True,
+    use_pure_torch: bool = False,
+    read_existing: bool = False,
     dtype="float32",
 ):
     """Help function to set up JARVIS train and val dataloaders."""
-    if use_lmdb:
+    if use_pure_torch:
+        print("Using pure-torch LMDB dataset (no DGL).")
+        from alignn.pure_lmdb_dataset import get_torch_dataset
+    elif use_lmdb:
         print("Using LMDB dataset.")
         from alignn.lmdb_dataset import get_torch_dataset
     else:
@@ -383,6 +400,8 @@ def get_train_val_loaders(
             cutoff=cutoff,
             cutoff_extra=cutoff_extra,
             max_neighbors=max_neighbors,
+            three_body_cutoff=three_body_cutoff,
+            read_existing=read_existing,
             classification=classification_threshold is not None,
             output_dir=output_dir,
             sampler=train_sampler,
@@ -409,6 +428,7 @@ def get_train_val_loaders(
                 cutoff_extra=cutoff_extra,
                 sampler=val_sampler,
                 max_neighbors=max_neighbors,
+                three_body_cutoff=three_body_cutoff,
                 classification=classification_threshold is not None,
                 output_dir=output_dir,
                 tmp_name=tmp_name,
@@ -462,7 +482,20 @@ def get_train_val_loaders(
             num_workers=workers,
             pin_memory=pin_memory,
             use_ddp=use_ddp,
+            persistent_workers=(workers > 0),
+            prefetch_factor=2 if workers > 0 else None,
         )
+        # train_loader = GraphDataLoader(
+        #    # train_loader = DataLoader(
+        #    train_data,
+        #    batch_size=batch_size,
+        #    shuffle=True,
+        #    collate_fn=collate_fn,
+        #    drop_last=True,
+        #    num_workers=workers,
+        #    pin_memory=pin_memory,
+        #    use_ddp=use_ddp,
+        # )
 
         val_loader = GraphDataLoader(
             # val_loader = DataLoader(
@@ -474,7 +507,20 @@ def get_train_val_loaders(
             num_workers=workers,
             pin_memory=pin_memory,
             use_ddp=use_ddp,
+            persistent_workers=(workers > 0),
+            prefetch_factor=2 if workers > 0 else None,
         )
+        # val_loader = GraphDataLoader(
+        #    # val_loader = DataLoader(
+        #    val_data,
+        #    batch_size=batch_size,
+        #    shuffle=False,
+        #    collate_fn=collate_fn,
+        #    drop_last=True,
+        #    num_workers=workers,
+        #    pin_memory=pin_memory,
+        #    use_ddp=use_ddp,
+        # )
 
         test_loader = (
             GraphDataLoader(
@@ -488,10 +534,27 @@ def get_train_val_loaders(
                 pin_memory=pin_memory,
                 use_ddp=use_ddp,
             )
-            if len(dataset_test) > 0
+            if test_data is not None
             else None
         )
 
+        # All loaders constructed; safe to free the in-memory dataset
+        # lists. With 1.5M-entry datasets these can hold tens of GB of
+        # full atomic-structure dicts, and they are no longer needed
+        # once the LMDB caches are built and the loaders reference
+        # them via the Dataset wrappers.
+        del dataset_train
+        del dataset_val
+        del dataset_test
+        del dat
+        try:
+            del d
+        except NameError:
+            pass
+        import gc
+
+        gc.collect()
+
         if save_dataloader:
             torch.save(train_loader, train_sample)
             if val_loader is not None:

alignn/dataset.py

@@ -4,7 +4,10 @@
 from typing import Optional
 import os
 import torch
-import dgl
+try:
+    import dgl
+except ImportError:
+    dgl = None
 import numpy as np
 import pandas as pd
 from jarvis.core.atoms import Atoms
@@ -23,6 +26,7 @@ def load_graphs(
     cutoff: float = 8,
     cutoff_extra: float = 3,
     max_neighbors: int = 12,
+    three_body_cutoff: Optional[float] = None,
     cachedir: Optional[Path] = None,
     use_canonize: bool = False,
     id_tag="jid",
@@ -87,6 +91,7 @@ def atoms_to_graph(atoms):
                 compute_line_graph=False,
                 use_canonize=use_canonize,
                 neighbor_strategy=neighbor_strategy,
+                three_body_cutoff=three_body_cutoff,
                 id=i[id_tag],
                 dtype=dtype,
             )
@@ -125,10 +130,12 @@ def get_torch_dataset(
     cutoff=8.0,
     cutoff_extra=3.0,
     max_neighbors=12,
+    three_body_cutoff=None,
     classification=False,
     output_dir=".",
     tmp_name="dataset",
     sampler=None,
+    read_existing=False,  # accepted for API parity; no-op (no cache here)
     dtype="float32",
 ):
     """Get Torch Dataset."""
@@ -152,6 +159,7 @@ def get_torch_dataset(
         cutoff=cutoff,
         cutoff_extra=cutoff_extra,
         max_neighbors=max_neighbors,
+        three_body_cutoff=three_body_cutoff,
         id_tag=id_tag,
         dtype=dtype,
     )