From fdf64a3546c6fca79aa428a69253cb2210649983 Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 21:09:14 -0400 Subject: [PATCH 01/34] Pure torch --- alignn/config.py | 17 +- alignn/data.py | 11 +- alignn/dataset.py | 5 + alignn/ff/lammps_bridge.py | 2 +- alignn/graphs.py | 207 +++++++- alignn/lmdb_dataset.py | 30 +- alignn/models/alignn_atomwise_pure.py | 701 ++++++++++++++++++++++++++ alignn/models/utils.py | 39 +- alignn/scripts/export_torchscript.py | 194 +++++++ alignn/scripts/parity_dgl_vs_pure.py | 229 +++++++++ alignn/torch_graph_builder.py | 589 ++++++++++++++++++++++ alignn/train.py | 16 +- alignn/train_alignn.py | 3 + docs/reference/config-reference.md | 335 ++++++++++++ mkdocs.yml | 1 + 15 files changed, 2352 insertions(+), 27 deletions(-) create mode 100644 alignn/models/alignn_atomwise_pure.py create mode 100644 alignn/scripts/export_torchscript.py create mode 100644 alignn/scripts/parity_dgl_vs_pure.py create mode 100644 alignn/torch_graph_builder.py create mode 100644 docs/reference/config-reference.md diff --git a/alignn/config.py b/alignn/config.py index 3a5cad7..2ea317b 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -7,6 +7,7 @@ from alignn.utils import BaseSettings from alignn.models.alignn import ALIGNNConfig from alignn.models.alignn_atomwise import ALIGNNAtomWiseConfig +from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePureConfig from alignn.models.ealignn_atomwise import eALIGNNAtomWiseConfig # import torch @@ -155,7 +156,13 @@ class TrainingConfig(BaseSettings): target: TARGET_ENUM = "exfoliation_energy" atom_features: Literal["basic", "atomic_number", "cfid", "cgcnn"] = "cgcnn" neighbor_strategy: Literal[ - "k-nearest", "voronoi", "radius_graph", "radius_graph_jarvis" + "k-nearest", + "voronoi", + "radius_graph", + "radius_graph_jarvis", + "fast_graph", + "torch_graph", + "pure_torch", ] = "k-nearest" id_tag: Literal["jid", "id", "_oqmd_entry_id"] = "jid" @@ -193,6 +200,9 @@ class TrainingConfig(BaseSettings): num_workers: int = 4 cutoff: float = 8.0 cutoff_extra: float = 3.0 + # Separate 3-body cutoff used by neighbor_strategy="pure_torch". + # When None, defaults to `cutoff`. Must be <= cutoff. + three_body_cutoff: Optional[float] = None max_neighbors: int = 12 keep_data_order: bool = True normalize_graph_level_loss: bool = False @@ -201,6 +211,10 @@ class TrainingConfig(BaseSettings): n_early_stopping: Optional[int] = None # typically 50 output_dir: str = os.path.abspath(".") use_lmdb: bool = True + # When True, reuse an existing LMDB cache on disk (fast, but the + # cache must match the current model backend + graph config). + # Default False: always rebuild from scratch for safety. + read_existing: bool = False # alignn_layers: int = 4 # gcn_layers: int =4 # edge_input_features: int= 80 @@ -212,5 +226,6 @@ class TrainingConfig(BaseSettings): model: Union[ ALIGNNConfig, ALIGNNAtomWiseConfig, + ALIGNNAtomWisePureConfig, eALIGNNAtomWiseConfig, ] = ALIGNNAtomWiseConfig(name="alignn_atomwise") diff --git a/alignn/data.py b/alignn/data.py index 39f2fd0..24f9ef3 100644 --- a/alignn/data.py +++ b/alignn/data.py @@ -145,6 +145,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 +155,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 +389,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 +417,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, diff --git a/alignn/dataset.py b/alignn/dataset.py index 0d2e88b..0f425b6 100644 --- a/alignn/dataset.py +++ b/alignn/dataset.py @@ -23,6 +23,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 +88,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 +127,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 +156,7 @@ def get_torch_dataset( cutoff=cutoff, cutoff_extra=cutoff_extra, max_neighbors=max_neighbors, + three_body_cutoff=three_body_cutoff, id_tag=id_tag, dtype=dtype, ) diff --git a/alignn/ff/lammps_bridge.py b/alignn/ff/lammps_bridge.py index cb8cf05..3bebad4 100644 --- a/alignn/ff/lammps_bridge.py +++ b/alignn/ff/lammps_bridge.py @@ -133,7 +133,7 @@ def run(args): timestep {args.timestep} velocity all create {args.temp} {args.seed} mom yes rot yes fix nve all nve - fix tfix all langevin {args.temp} {args.temp} 0.1 {args.seed} + fix tfix all langevin {args.temp} {args.temp} 0.1 {args.seed} # noqa: E501 thermo 10 thermo_style custom step temp pe ke etotal press """ diff --git a/alignn/graphs.py b/alignn/graphs.py index a017aa2..250a16e 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -264,6 +264,152 @@ def build_undirected_edgedata( return u, v, r, all_images +def fast_graph( + atoms=None, + cutoff=5.0, + max_neighbors=None, + id=None, +): + """Fast radius neighbor list using matscipy (C) on CPU. + + Returns (u, v, r, images) in the same format as radius_graph so it + can be consumed by Graph.atom_dgl_multigraph downstream. + """ + from matscipy.neighbours import neighbour_list as _matscipy_neighbours + + ase_atoms = atoms.ase_converter() + i_np, j_np, D_np, S_np = _matscipy_neighbours( + "ijDS", ase_atoms, float(cutoff) + ) + + if max_neighbors is not None and max_neighbors > 0 and i_np.size > 0: + d_np = np.linalg.norm(D_np, axis=1) + # rank edges by distance globally, then keep first K per source + order = np.lexsort((d_np, i_np)) + i_np = i_np[order] + j_np = j_np[order] + D_np = D_np[order] + S_np = S_np[order] + N = int(ase_atoms.positions.shape[0]) + src_start = np.searchsorted(i_np, np.arange(N)) + within = np.arange(i_np.shape[0]) - src_start[i_np] + sel = within < int(max_neighbors) + i_np = i_np[sel] + j_np = j_np[sel] + D_np = D_np[sel] + S_np = S_np[sel] + + dtype = torch.get_default_dtype() + u = torch.from_numpy(np.ascontiguousarray(i_np)).long() + v = torch.from_numpy(np.ascontiguousarray(j_np)).long() + r = torch.from_numpy(np.ascontiguousarray(D_np)).type(dtype) + images = torch.from_numpy(np.ascontiguousarray(S_np)).type(dtype) + return u, v, r, images + + +# ===================================================================== +# Fully-torch autograd-compatible neighbor list & graph builder +# +# Rationale: +# Existing radius_graph / fast_graph return `r` as a fresh tensor +# rebuilt from numpy, severing its link to atom positions and the +# lattice. For autograd-based forces (dE/dx) and stress (dE/dL), `r` +# must be a torch function of both. Here topology is computed once +# (discrete, no grad), and displacement vectors are then recomputed +# as `r = pos[dst] - pos[src] + shift @ lattice`, which carries +# gradients. The line graph uses `shared=True` so bond-angle cosines +# inherit the autograd graph automatically. +# ===================================================================== + + +# Neighbor-list primitives now live in alignn.torch_graph_builder so +# the pure-torch path doesn't transitively import DGL. Re-export them +# here for backward compatibility. +from alignn.torch_graph_builder import ( # noqa: E402,F401 + _torch_periodic_shifts, + _topk_per_source, + torch_neighbor_list, +) + + +def torch_graph( + atoms=None, + cutoff=5.0, + max_neighbors=None, + atom_features="cgcnn", + compute_line_graph=True, + use_lattice_prop=False, + positions: Optional[torch.Tensor] = None, + lattice: Optional[torch.Tensor] = None, + device=None, + use_matscipy_topology: bool = True, + id=None, +): + """End-to-end autograd-capable DGL (graph, line_graph) builder. + + Returns ``(g, lg)`` if ``compute_line_graph`` else ``g``. When + ``positions`` / ``lattice`` are provided as torch leaves with + ``requires_grad=True``, ``g.edata['r']`` and ``lg.edata['h']`` + carry gradients back to them. + """ + dtype = torch.get_default_dtype() + if positions is None: + positions = torch.as_tensor(np.asarray(atoms.cart_coords), dtype=dtype) + if lattice is None: + lattice = torch.as_tensor(np.asarray(atoms.lattice_mat), dtype=dtype) + if device is not None: + positions = positions.to(device) + lattice = lattice.to(device) + device = positions.device + n_atoms = len(atoms.elements) + + src, dst, shift, r = torch_neighbor_list( + positions=positions, + lattice=lattice, + cutoff=float(cutoff), + max_neighbors=max_neighbors, + atoms=atoms, + use_matscipy_topology=use_matscipy_topology, + ) + + sps_features = np.array( + [ + list(get_node_attributes(s, atom_features=atom_features)) + for s in atoms.elements + ] + ) + node_features = torch.as_tensor(sps_features, dtype=dtype, device=device) + + g = dgl.graph((src, dst), num_nodes=n_atoms) + if g.device != device: + g = g.to(device) + g.ndata["atom_features"] = node_features + g.edata["r"] = r # differentiable w.r.t. positions & lattice + g.edata["images"] = shift # integer cell offsets + vol = torch.abs(torch.det(lattice)) + g.ndata["V"] = vol.expand(n_atoms) + frac = torch.as_tensor( + np.asarray(atoms.frac_coords), dtype=dtype, device=device + ) + g.ndata["frac_coords"] = frac + if use_lattice_prop: + lp = np.array( + [atoms.lattice.lat_lengths(), atoms.lattice.lat_angles()] + ).flatten() + g.ndata["extra_features"] = ( + torch.as_tensor(lp, dtype=dtype, device=device) + .unsqueeze(0) + .expand(n_atoms, -1) + .contiguous() + ) + + if compute_line_graph: + lg = g.line_graph(shared=True) + lg.apply_edges(compute_bond_cosines) + return g, lg + return g + + def radius_graph( atoms=None, cutoff=5, @@ -483,6 +629,7 @@ def atom_dgl_multigraph( use_lattice_prop: bool = False, cutoff_extra=3.5, dtype="float32", + three_body_cutoff: Optional[float] = None, ): """Obtain a DGLGraph for Atoms object.""" # print('id',id) @@ -503,6 +650,52 @@ def atom_dgl_multigraph( u, v, r, images = radius_graph( atoms, cutoff=cutoff, cutoff_extra=cutoff_extra ) + elif neighbor_strategy == "fast_graph": + u, v, r, images = fast_graph( + atoms, + cutoff=cutoff, + max_neighbors=max_neighbors, + id=id, + ) + elif neighbor_strategy == "torch_graph": + # Returns a fully-built (g, lg) with autograd-capable r. + # Short-circuit to avoid the generic post-processing below + # which would detach r via torch.tensor(np.array(r)). + return torch_graph( + atoms=atoms, + cutoff=cutoff, + max_neighbors=max_neighbors, + atom_features=atom_features, + compute_line_graph=compute_line_graph, + use_lattice_prop=use_lattice_prop, + id=id, + ) + elif neighbor_strategy == "pure_torch": + # Build with the pure-torch (non-DGL) builder, then + # materialize DGL graphs at the boundary. Downstream + # consumers (dataset/model) see DGL graphs as usual, and + # `edata["r"]` stays autograd-connected to positions and + # lattice because TorchGraph.to_dgl() assigns by reference. + from alignn.torch_graph_builder import ( + build_pure_torch_graph as _build_pure_torch_graph, + ) + + _r3 = ( + three_body_cutoff if three_body_cutoff is not None else cutoff + ) + _out = _build_pure_torch_graph( + atoms=atoms, + two_body_cutoff=cutoff, + three_body_cutoff=_r3, + max_neighbors=max_neighbors, + atom_features=atom_features, + use_lattice_prop=use_lattice_prop, + compute_line_graph=compute_line_graph, + ) + if compute_line_graph: + _tg, _tlg = _out + return _tg.to_dgl(), _tlg.to_dgl() + return _out.to_dgl() elif neighbor_strategy == "radius_graph_jarvis": g, lg = radius_graph_jarvis( atoms, @@ -844,6 +1037,18 @@ def prepare_line_graph_batch( # return tuple(x.to(device) for x in batch) +# Re-export the pure-torch (non-DGL) builder for convenience. +# Defined in alignn.torch_graph_builder to keep that path importable +# without pulling the full DGL-oriented graphs module. +def build_pure_torch_graph(*args, **kwargs): + """See alignn.torch_graph_builder.build_pure_torch_graph.""" + from alignn.torch_graph_builder import ( + build_pure_torch_graph as _impl, + ) + + return _impl(*args, **kwargs) + + def compute_bond_cosines(edges): """Compute bond angle cosines from bond displacement vectors.""" # line graph edge: (a, b), (b, c) @@ -1062,7 +1267,7 @@ def collate(samples: List[Tuple[dgl.DGLGraph, torch.Tensor]]): @staticmethod def collate_line_graph( - samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]] + samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]], ): """Dataloader helper to batch graphs cross `samples`.""" graphs, line_graphs, lattices, labels = map(list, zip(*samples)) diff --git a/alignn/lmdb_dataset.py b/alignn/lmdb_dataset.py index 08e507c..3b540a2 100644 --- a/alignn/lmdb_dataset.py +++ b/alignn/lmdb_dataset.py @@ -86,7 +86,7 @@ def collate(samples: List[Tuple[dgl.DGLGraph, torch.Tensor]]): @staticmethod def collate_line_graph( - samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]] + samples: List[Tuple[dgl.DGLGraph, dgl.DGLGraph, torch.Tensor]], ): """Dataloader helper to batch graphs cross `samples`.""" graphs, line_graphs, lattices, labels = map(list, zip(*samples)) @@ -124,12 +124,13 @@ def get_torch_dataset( cutoff=8.0, cutoff_extra=3.0, max_neighbors=12, + three_body_cutoff=None, classification=False, sampler=None, output_dir=".", tmp_name="dataset", map_size=1e12, - read_existing=True, + read_existing=False, dtype="float32", ): """Get Torch Dataset with LMDB.""" @@ -144,6 +145,24 @@ def get_torch_dataset( f.close() ids = [] if os.path.exists(tmp_name) and read_existing: + # Validate the cache backend matches (DGL). A previous run with + # model.name="alignn_atomwise_pure" would have left TorchGraph + # pickles here, which dgl.batch can't consume. + _env = lmdb.open(tmp_name, readonly=True, lock=False) + with _env.begin() as _txn: + _probe = _txn.get(b"0") + _env.close() + if _probe is not None: + _sample = pk.loads(_probe) + _graph = _sample[0] + if not isinstance(_graph, dgl.DGLGraph): + raise RuntimeError( + f"LMDB cache at '{tmp_name}' contains " + f"{type(_graph).__name__} records, not DGLGraph. " + "Delete the stale cache (e.g. `rm -rf " + f"{tmp_name}`) and rerun — it was built for a " + "different model backend." + ) for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)): ids.append(d[id_tag]) dat = TorchLMDBDataset( @@ -152,6 +171,12 @@ def get_torch_dataset( print("Reading dataset", tmp_name) return dat ids = [] + # Fresh build: wipe any pre-existing cache dir to avoid mixing old + # records (possibly from a different model backend) with new writes. + if os.path.exists(tmp_name): + import shutil + + shutil.rmtree(tmp_name) env = lmdb.open(tmp_name, map_size=int(map_size)) with env.begin(write=True) as txn: for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)): @@ -167,6 +192,7 @@ def get_torch_dataset( use_canonize=use_canonize, cutoff_extra=cutoff_extra, neighbor_strategy=neighbor_strategy, + three_body_cutoff=three_body_cutoff, dtype=dtype, ) if line_graph: diff --git a/alignn/models/alignn_atomwise_pure.py b/alignn/models/alignn_atomwise_pure.py new file mode 100644 index 0000000..9c1a8d9 --- /dev/null +++ b/alignn/models/alignn_atomwise_pure.py @@ -0,0 +1,701 @@ +"""Pure-PyTorch ALIGNN-atomwise (no DGL). + +This is a line-for-line reimplementation of ``ALIGNNAtomWise`` with all +DGL ops (``dgl.function``, ``dgl.nn``, ``dgl.batch``, ``dgl.reverse``) +replaced by index/scatter primitives on ``TorchGraph`` tensors. Same +config fields, same forward output keys — drop-in for training and +inference without depending on DGL at model-evaluation time. + +Accepts either: + * a ``(TorchGraph, TorchGraph, lattice)`` triple (preferred), or + * a ``(DGLGraph, DGLGraph, lattice)`` triple — converted at the + boundary via ``torchgraph_from_dgl`` so you can flip models without + touching the dataloader. +""" + +from __future__ import annotations + +from typing import Dict, List, Literal, Tuple + +import numpy as np +import torch +from torch import nn +from torch.autograd import grad +from torch.nn import functional as F + +from alignn.models.utils import MLPLayer, RBFExpansion +from alignn.torch_graph_builder import TorchGraph, torchgraph_from_dgl +from alignn.utils import BaseSettings + + +# ===================================================================== +# Primitives +# ===================================================================== + + +def scatter_sum( + src: torch.Tensor, index: torch.Tensor, dim_size: int +) -> torch.Tensor: + """Sum ``src`` rows into ``dim_size`` buckets indexed by ``index``. + + Script-safe: builds output shape as an explicit ``List[int]`` and + broadcasts ``index`` via repeated ``unsqueeze`` (avoids unpacking). + """ + out_shape: List[int] = [dim_size] + for i in range(1, src.dim()): + out_shape.append(src.size(i)) + out = torch.zeros(out_shape, dtype=src.dtype, device=src.device) + idx = index + while idx.dim() < src.dim(): + idx = idx.unsqueeze(-1) + idx = idx.expand_as(src) + out.scatter_add_(0, idx, src) + return out + + +def scatter_mean( + src: torch.Tensor, index: torch.Tensor, dim_size: int +) -> torch.Tensor: + total = scatter_sum(src, index, dim_size) + ones = torch.ones(src.shape[0], device=src.device, dtype=src.dtype) + count = torch.zeros( + dim_size, device=src.device, dtype=src.dtype + ).scatter_add_(0, index, ones) + denom = count.clamp_min(1.0) + while denom.dim() < total.dim(): + denom = denom.unsqueeze(-1) + return total / denom + + +# ===================================================================== +# Config +# ===================================================================== + + +class ALIGNNAtomWisePureConfig(BaseSettings): + """Hyperparameter schema — mirrors ALIGNNAtomWiseConfig.""" + + name: Literal["alignn_atomwise_pure"] + alignn_layers: int = 2 + gcn_layers: int = 2 + atom_input_features: int = 1 + edge_input_features: int = 80 + triplet_input_features: int = 40 + embedding_features: int = 64 + hidden_features: int = 64 + output_features: int = 1 + grad_multiplier: int = -1 + calculate_gradient: bool = True + atomwise_output_features: int = 0 + graphwise_weight: float = 1.0 + gradwise_weight: float = 1.0 + stresswise_weight: float = 0.0 + atomwise_weight: float = 0.0 + link: Literal["identity", "log", "logit"] = "identity" + zero_inflated: bool = False + classification: bool = False + force_mult_natoms: bool = False + energy_mult_natoms: bool = True + # Accepted for config-schema compatibility with ALIGNNAtomWiseConfig; + # the pure-torch forward doesn't implement this branch. + include_pos_deriv: bool = False + use_cutoff_function: bool = False + inner_cutoff: float = 3.0 + stress_multiplier: float = 1.0 + add_reverse_forces: bool = True + lg_on_fly: bool = True + batch_stress: bool = True + multiply_cutoff: bool = False + use_penalty: bool = True + extra_features: int = 0 + exponent: int = 5 + penalty_factor: float = 0.1 + penalty_threshold: float = 1.0 + additional_output_features: int = 0 + additional_output_weight: float = 0.0 + + class Config: + env_prefix = "jv_model" + + +# ===================================================================== +# Conv layers +# ===================================================================== + + +class EdgeGatedGraphConvPure(nn.Module): + """Edge-gated graph convolution, pure torch / scatter-based.""" + + def __init__( + self, input_features: int, output_features: int, residual: bool = True + ): + super().__init__() + self.residual = residual + self.src_gate = nn.Linear(input_features, output_features) + self.dst_gate = nn.Linear(input_features, output_features) + self.edge_gate = nn.Linear(input_features, output_features) + self.bn_edges = nn.LayerNorm(output_features) + self.src_update = nn.Linear(input_features, output_features) + self.dst_update = nn.Linear(input_features, output_features) + self.bn_nodes = nn.LayerNorm(output_features) + + def forward_tensors( + self, + src: torch.Tensor, + dst: torch.Tensor, + num_nodes: int, + x: torch.Tensor, + y: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Script-friendly core: takes plain tensors instead of a graph.""" + # Edge update: m = e_src[src] + e_dst[dst] + edge_gate(y) + e_src = self.src_gate(x) + e_dst = self.dst_gate(x) + m = e_src[src] + e_dst[dst] + self.edge_gate(y) + sigma = torch.sigmoid(m) + + # Node update aggregated at dst: + # h = sum_{(i,j)} sigma * Bh[src] / (sum sigma + eps). + Bh = self.dst_update(x) + msg_h = Bh[src] * sigma + sum_sigma_h = scatter_sum(msg_h, dst, num_nodes) + sum_sigma = scatter_sum(sigma, dst, num_nodes) + h = sum_sigma_h / (sum_sigma + 1e-6) + x_new = self.src_update(x) + h + + x_new = F.silu(self.bn_nodes(x_new)) + y_new = F.silu(self.bn_edges(m)) + + if self.residual: + x_new = x + x_new + y_new = y + y_new + return x_new, y_new + + @torch.jit.ignore + def forward( + self, g: TorchGraph, x: torch.Tensor, y: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + return self.forward_tensors(g.src, g.dst, g.num_nodes, x, y) + + +class ALIGNNConvPure(nn.Module): + """Line-graph-aware ALIGNN update.""" + + def __init__(self, in_features: int, out_features: int): + super().__init__() + self.node_update = EdgeGatedGraphConvPure(in_features, out_features) + self.edge_update = EdgeGatedGraphConvPure(out_features, out_features) + + def forward_tensors( + self, + g_src: torch.Tensor, + g_dst: torch.Tensor, + g_num_nodes: int, + lg_src: torch.Tensor, + lg_dst: torch.Tensor, + lg_num_nodes: int, + x: torch.Tensor, + y: torch.Tensor, + z: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x, m = self.node_update.forward_tensors( + g_src, g_dst, g_num_nodes, x, y + ) + y, z = self.edge_update.forward_tensors( + lg_src, lg_dst, lg_num_nodes, m, z + ) + return x, y, z + + @torch.jit.ignore + def forward( + self, + g: TorchGraph, + lg: TorchGraph, + x: torch.Tensor, + y: torch.Tensor, + z: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.forward_tensors( + g.src, + g.dst, + g.num_nodes, + lg.src, + lg.dst, + lg.num_nodes, + x, + y, + z, + ) + + +# ===================================================================== +# Helpers +# ===================================================================== + + +def cutoff_function_based_edges( + r: torch.Tensor, inner_cutoff: float = 4.0, exponent: int = 3 +) -> torch.Tensor: + ratio = r / inner_cutoff + c1 = -(exponent + 1) * (exponent + 2) / 2 + c2 = exponent * (exponent + 2) + c3 = -exponent * (exponent + 1) / 2 + envelope = ( + 1.0 + + c1 * ratio**exponent + + c2 * ratio ** (exponent + 1) + + c3 * ratio ** (exponent + 2) + ) + return torch.where(r <= inner_cutoff, envelope, torch.zeros_like(r)) + + +def _bond_cosines(r_ij: torch.Tensor, r_jk: torch.Tensor) -> torch.Tensor: + num = -(r_ij * r_jk).sum(dim=-1) + denom = r_ij.norm(dim=-1) * r_jk.norm(dim=-1) + return (num / denom.clamp_min(1e-12)).clamp(-1.0, 1.0) + + +def _as_torchgraph(x): + """Accept TorchGraph or DGLGraph transparently.""" + if isinstance(x, TorchGraph): + return x + return torchgraph_from_dgl(x) + + +# ===================================================================== +# Model +# ===================================================================== + + +class ALIGNNAtomWisePure(nn.Module): + """DGL-free ALIGNN atomwise.""" + + def __init__( + self, + config: ALIGNNAtomWisePureConfig = ALIGNNAtomWisePureConfig( + name="alignn_atomwise_pure" + ), + ): + super().__init__() + self.config = config + self.classification = config.classification + if self.config.gradwise_weight == 0: + self.config.calculate_gradient = False + + self.atom_embedding = MLPLayer( + config.atom_input_features, config.hidden_features + ) + self.edge_embedding = nn.Sequential( + RBFExpansion(vmin=0, vmax=8.0, bins=config.edge_input_features), + MLPLayer(config.edge_input_features, config.embedding_features), + MLPLayer(config.embedding_features, config.hidden_features), + ) + self.angle_embedding = nn.Sequential( + RBFExpansion( + vmin=-1, vmax=1.0, bins=config.triplet_input_features + ), + MLPLayer(config.triplet_input_features, config.embedding_features), + MLPLayer(config.embedding_features, config.hidden_features), + ) + self.alignn_layers = nn.ModuleList( + [ + ALIGNNConvPure(config.hidden_features, config.hidden_features) + for _ in range(config.alignn_layers) + ] + ) + self.gcn_layers = nn.ModuleList( + [ + EdgeGatedGraphConvPure( + config.hidden_features, config.hidden_features + ) + for _ in range(config.gcn_layers) + ] + ) + if config.atomwise_output_features > 0: + self.fc_atomwise = nn.Linear( + config.hidden_features, config.atomwise_output_features + ) + if config.additional_output_features: + self.fc_additional_output = nn.Linear( + config.hidden_features, config.additional_output_features + ) + if self.classification: + self.fc = nn.Linear(config.hidden_features, 2) + self.softmax = nn.LogSoftmax(dim=1) + else: + self.fc = nn.Linear(config.hidden_features, config.output_features) + # Link is applied inline in the forward paths (kept as a string so + # `torch.jit.script` doesn't choke on a Python lambda attribute). + self.link_name: str = config.link + if config.link == "log": + self.fc.bias.data = torch.tensor(np.log(0.7), dtype=torch.float) + + # Species feature lookup for LAMMPS / TorchScript entry point. + # Registered empty here; ``register_species_table`` populates it + # before export. Stays unused on the training path. + self.register_buffer( + "_species_table", + torch.zeros(120, config.atom_input_features, dtype=torch.float32), + ) + + # Expose scalar config fields that `forward_tensors` reads as + # typed Python attributes — the pydantic ``self.config`` object + # is not TorchScript-compatible. + self.energy_mult_natoms: bool = bool(config.energy_mult_natoms) + self.use_penalty: bool = bool(config.use_penalty) + self.penalty_threshold: float = float(config.penalty_threshold) + self.penalty_factor: float = float(config.penalty_factor) + self.grad_multiplier: int = int(config.grad_multiplier) + self.add_reverse_forces: bool = bool(config.add_reverse_forces) + self.stress_multiplier: float = float(config.stress_multiplier) + + # ----- species lookup (LAMMPS hands us atomic numbers, not features) ----- + + @torch.jit.ignore + def register_species_table( + self, atom_features: str = "cgcnn", max_z: int = 119 + ) -> None: + """Build a (max_z+1, F) feature table indexed by atomic number. + + Once registered, ``forward_tensors_z`` can be used from TorchScript + so callers (e.g. LAMMPS) don't need to ship the feature table + alongside the model. Invalid / unknown atomic numbers map to + zero vectors. + """ + from ase.data import chemical_symbols + from jarvis.core.specie import get_node_attributes + + F = int(self.config.atom_input_features) + rows = np.zeros((max_z + 1, F), dtype=np.float32) + for z in range(1, max_z + 1): + if z >= len(chemical_symbols): + continue + symbol = chemical_symbols[z] + try: + feat = np.asarray( + get_node_attributes(symbol, atom_features=atom_features), + dtype=np.float32, + ) + if feat.shape[0] != F: + continue + rows[z] = feat + except Exception: + continue + # Buffer was allocated in __init__; overwrite the tensor contents + # (register_buffer would complain about duplicate registration). + self._species_table.resize_(rows.shape).copy_( + torch.as_tensor(rows, dtype=torch.float32) + ) + + # ----- tensor-only forward (LAMMPS / TorchScript path) ----- + + @torch.jit.export + def forward_tensors_z( + self, + positions: torch.Tensor, # (N, 3), requires_grad=True for forces + lattice: torch.Tensor, # (3, 3) + atomic_numbers: torch.Tensor, # (N,) long + src: torch.Tensor, # (E,) long + dst: torch.Tensor, # (E,) long + shift: torch.Tensor, # (E, 3) + compute_stress: bool = False, + ) -> Dict[str, torch.Tensor]: + """Forward with internal atomic-number → feature lookup. + + Same as ``forward_tensors`` but accepts atomic numbers directly. + Requires ``register_species_table`` to have been called + beforehand. LAMMPS-style callers only ship atomic numbers. + """ + atom_features = self._species_table.index_select(0, atomic_numbers) + return self.forward_tensors( + positions, lattice, atom_features, src, dst, shift, compute_stress + ) + + @torch.jit.export + def forward_tensors( + self, + positions: torch.Tensor, # (N, 3), requires_grad=True for forces + lattice: torch.Tensor, # (3, 3), requires_grad=True for stress + atom_features: torch.Tensor, # (N, atom_input_features) + src: torch.Tensor, # (E,) long + dst: torch.Tensor, # (E,) long + shift: torch.Tensor, # (E, 3) integer cell offsets (float dtype) + compute_stress: bool = False, + ) -> Dict[str, torch.Tensor]: + """Single-system forward driven by plain tensors. + + Intended for LAMMPS-style MD where each step supplies an + atom list, lattice, and neighbor list. No batching, no class + pooling machinery. Returns ``{"energy", "forces", ["stress"]}``. + """ + num_nodes = positions.shape[0] + + # Differentiable edge vectors. + r = ( + positions.index_select(0, dst) + - positions.index_select(0, src) + + torch.matmul(shift, lattice) + ) + bondlength = torch.linalg.vector_norm(r, dim=1) + + # Inline line-graph: connect parent edge A=(u,v) to B=(v,w). + E = src.shape[0] + order = torch.argsort(src, stable=True) + sorted_src = src.index_select(0, order) + node_range = torch.arange(num_nodes, device=src.device) + bucket_start = torch.searchsorted(sorted_src, node_range) + bucket_end = torch.searchsorted(sorted_src, node_range, right=True) + + A_v = dst + starts = bucket_start.index_select(0, A_v) + ends = bucket_end.index_select(0, A_v) + counts = ends - starts + + total = int(counts.sum().item()) + A_ids = torch.arange(E, device=src.device) + lg_src = torch.repeat_interleave(A_ids, counts) + cum = torch.cumsum(counts, dim=0) + row_start = cum - counts + offsets = torch.arange( + total, device=src.device + ) - torch.repeat_interleave(row_start, counts) + pos_idx = torch.repeat_interleave(starts, counts) + offsets + lg_dst = order.index_select(0, pos_idx) + lg_num_nodes = E + + # Angle cosines (differentiable through r -> positions, lattice). + r_ij = r.index_select(0, lg_src) + r_jk = r.index_select(0, lg_dst) + num = -(r_ij * r_jk).sum(dim=-1) + denom = torch.linalg.vector_norm(r_ij, dim=-1).clamp_min( + 1e-12 + ) * torch.linalg.vector_norm(r_jk, dim=-1).clamp_min(1e-12) + h_cos = (num / denom).clamp(-1.0, 1.0) + + # Embeddings. + x = self.atom_embedding(atom_features) + y = self.edge_embedding(bondlength) + z = self.angle_embedding(h_cos) + + # ALIGNN and GCN layers via tensor-only paths. + for layer in self.alignn_layers: + x, y, z = layer.forward_tensors( + src, dst, num_nodes, lg_src, lg_dst, lg_num_nodes, x, y, z + ) + for layer in self.gcn_layers: + x, y = layer.forward_tensors(src, dst, num_nodes, x, y) + + # Mean readout (single graph) and final projection. + h_graph = x.mean(dim=0, keepdim=True) # (1, hidden) + out = self.fc(h_graph).squeeze(-1) # (1,) for output_features=1 + + en = out + if self.energy_mult_natoms: + en = out * float(num_nodes) + + if self.use_penalty: + pen = torch.where( + bondlength < self.penalty_threshold, + self.penalty_factor * (self.penalty_threshold - bondlength), + torch.zeros_like(bondlength), + ) + en = en + pen.sum() + + result: Dict[str, torch.Tensor] = {"energy": en.squeeze()} + + # Autograd forces through r -> positions. + grad_outs = torch.autograd.grad( + outputs=[en.sum()], + inputs=[r], + create_graph=self.training, + retain_graph=True, + ) + pair_forces = grad_outs[0] + assert pair_forces is not None + pair_forces = self.grad_multiplier * pair_forces + + # Atom forces = -dE/dx = sum of dE/dr over incoming edges + # minus sum over outgoing (add_reverse_forces convention). + forces_ji = scatter_sum(pair_forces, dst, num_nodes) + if self.add_reverse_forces: + forces_ij = scatter_sum(pair_forces, src, num_nodes) + forces = forces_ji - forces_ij + else: + forces = forces_ji + result["forces"] = forces + + if compute_stress: + # Virial stress: -160.217 * r.T @ pair_forces / V. + V = torch.abs(torch.det(lattice)) + stress = -160.21766208 * (torch.matmul(r.T, pair_forces) / V) + result["stress"] = self.stress_multiplier * stress + + return result + + # ----- dataclass-based forward (training pipeline) ----- + + @torch.jit.ignore + def forward(self, g): + # Unpack and normalize input. + gg, lg, lat = g + g = _as_torchgraph(gg) + lg = _as_torchgraph(lg) + + r = g.edata["r"] + if self.config.calculate_gradient: + r.requires_grad_(True) + + bondlength = torch.norm(r, dim=1) + + # Recompute angle cosines inside the autograd graph when lg_on_fly. + if self.config.lg_on_fly and len(self.alignn_layers) > 0: + h = _bond_cosines(r[lg.src], r[lg.dst]) + z = self.angle_embedding(h) + elif len(self.alignn_layers) > 0: + z = self.angle_embedding(lg.edata["h"]) + else: + z = None # unused + + x = self.atom_embedding(g.ndata["atom_features"]) + + if self.config.use_cutoff_function: + if self.config.multiply_cutoff: + c_off = cutoff_function_based_edges( + bondlength, + inner_cutoff=self.config.inner_cutoff, + exponent=self.config.exponent, + ).unsqueeze(-1) + y = self.edge_embedding(bondlength) * c_off + else: + bondlength_eff = cutoff_function_based_edges( + bondlength, + inner_cutoff=self.config.inner_cutoff, + exponent=self.config.exponent, + ) + y = self.edge_embedding(bondlength_eff) + else: + y = self.edge_embedding(bondlength) + + for layer in self.alignn_layers: + x, y, z = layer(g, lg, x, y, z) + for layer in self.gcn_layers: + x, y = layer(g, x, y) + + # Per-graph readout (mean over nodes). + node_bid = g.node_batch_id + h_graph = scatter_mean(x, node_bid, g.batch_size) + out = self.fc(h_graph) + if self.config.output_features == 1: + out = out.squeeze(-1) + + natoms = ( + g.batch_num_nodes + if g.batch_num_nodes is not None + else torch.tensor([g.num_nodes], device=r.device) + ) + en_out = out + if self.config.energy_mult_natoms: + en_out = out * natoms.to(out.dtype) + + if self.config.use_penalty: + penalties = torch.where( + bondlength < self.config.penalty_threshold, + self.config.penalty_factor + * (self.config.penalty_threshold - bondlength), + torch.zeros_like(bondlength), + ) + en_out = en_out + penalties.sum() + + atomwise_pred = torch.empty(1, device=r.device) + if ( + self.config.atomwise_output_features > 0 + and self.config.atomwise_weight != 0 + ): + atomwise_pred = self.fc_atomwise(x) + + additional_out = torch.empty(1, device=r.device) + if self.config.additional_output_features > 0: + additional_out = self.fc_additional_output(h_graph) + + forces = torch.empty(1, device=r.device) + stress = torch.empty(1, device=r.device) + if self.config.calculate_gradient: + pair_forces = ( + self.config.grad_multiplier + * grad( + en_out.sum(), + r, + grad_outputs=torch.ones_like(en_out.sum()), + create_graph=True, + retain_graph=True, + )[0] + ) + if self.config.force_mult_natoms: + pair_forces = pair_forces * g.num_nodes + + # force_i from incoming edges (dst == i): sum pair_forces at dst. + forces_ji = scatter_sum(pair_forces, g.dst, g.num_nodes) + if self.config.add_reverse_forces: + forces_ij = scatter_sum(pair_forces, g.src, g.num_nodes) + forces = (forces_ji - forces_ij).squeeze() + else: + forces = forces_ji.squeeze() + + if self.config.stresswise_weight != 0: + # Per-graph stress = -160.21766208 * (r.T @ pair_forces) / V. + # V is stored per-node; take V at the first node of each graph. + B = g.batch_size + V_per_node = g.ndata["V"] + if g.batch_num_nodes is not None: + node_offsets = torch.zeros( + B, dtype=torch.long, device=r.device + ) + node_offsets[1:] = torch.cumsum( + g.batch_num_nodes[:-1], dim=0 + ) + edge_offsets = torch.zeros( + B, dtype=torch.long, device=r.device + ) + edge_offsets[1:] = torch.cumsum( + g.batch_num_edges[:-1], dim=0 + ) + else: + node_offsets = torch.zeros( + 1, dtype=torch.long, device=r.device + ) + edge_offsets = torch.zeros( + 1, dtype=torch.long, device=r.device + ) + + stresses = [] + for b in range(B): + e0 = int(edge_offsets[b].item()) + e1 = e0 + ( + int(g.batch_num_edges[b].item()) + if g.batch_num_edges is not None + else g.num_edges + ) + n0 = int(node_offsets[b].item()) + V_b = V_per_node[n0] + st = -160.21766208 * ( + torch.matmul(r[e0:e1].T, pair_forces[e0:e1]) / V_b + ) + stresses.append(st) + stress = self.config.stress_multiplier * torch.stack(stresses) + + if self.link_name == "log": + out = torch.exp(out) + elif self.link_name == "logit": + out = torch.sigmoid(out) + if self.classification: + if out.dim() == 1: + out = out.unsqueeze(0) + out = self.softmax(out) + + return { + "out": out, + "additional": additional_out, + "grad": forces, + "stresses": stress, + "atomwise_pred": atomwise_pred, + } diff --git a/alignn/models/utils.py b/alignn/models/utils.py index cc02dc0..b68c53a 100644 --- a/alignn/models/utils.py +++ b/alignn/models/utils.py @@ -1,12 +1,17 @@ """Shared model-building components.""" -from typing import Optional, Callable +from typing import Optional, Callable, TYPE_CHECKING import numpy as np import torch import torch.nn as nn -import dgl from typing import Tuple +# Lazy-import dgl: only DGL-graph-specific helpers below need it, and +# we want importing this module to stay DGL-free so the pure-torch +# model can be loaded in environments without DGL. +if TYPE_CHECKING: # pragma: no cover + import dgl + class RBFExpansion(nn.Module): """Expand interatomic distances with radial basis functions.""" @@ -28,14 +33,14 @@ def __init__( ) if lengthscale is None: - # SchNet-style - # set lengthscales relative to granularity of RBF expansion - self.lengthscale = np.diff(self.centers).mean() - self.gamma = 1 / self.lengthscale - + # SchNet-style: set lengthscales relative to RBF granularity. + # Cast to Python float so TorchScript can type-infer the + # attribute (numpy scalars aren't recognized). + self.lengthscale = float(np.diff(self.centers).mean()) + self.gamma = 1.0 / self.lengthscale else: - self.lengthscale = lengthscale - self.gamma = 1 / (lengthscale**2) + self.lengthscale = float(lengthscale) + self.gamma = 1.0 / (self.lengthscale**2) def forward(self, distance: torch.Tensor) -> torch.Tensor: """Apply RBF expansion to interatomic distance tensor.""" @@ -44,7 +49,7 @@ def forward(self, distance: torch.Tensor) -> torch.Tensor: ) -def compute_pair_vector_and_distance(g: dgl.DGLGraph): +def compute_pair_vector_and_distance(g: "dgl.DGLGraph"): """Calculate bond vectors and distances using dgl graphs.""" # print('g.edges()',g.ndata["cart_coords"][g.edges()[1]].shape,g.edata["pbc_offshift"].shape) dst_pos = g.ndata["cart_coords"][g.edges()[1]] + g.edata["images"] @@ -127,10 +132,10 @@ def compute_cartesian_coordinates(g, lattice, dtype=torch.float32): def lightweight_line_graph( - input_graph: dgl.DGLGraph, + input_graph: "dgl.DGLGraph", feature_name: str, filter_condition: Callable[[torch.Tensor], torch.Tensor], -) -> dgl.DGLGraph: +) -> "dgl.DGLGraph": """Make the line graphs lightweight with preserved node ordering. Handles both batched and unbatched graphs. @@ -143,6 +148,8 @@ def lightweight_line_graph( New DGL graph with filtered edges preserving original node ordering """ # Check if graph is batched + import dgl + is_batched = ( hasattr(input_graph, "batch_size") and input_graph.batch_size > 1 ) @@ -223,10 +230,10 @@ def lightweight_line_graph( def lightweight_line_graph1( - input_graph: dgl.DGLGraph, + input_graph: "dgl.DGLGraph", feature_name: str, filter_condition: Callable[[torch.Tensor], torch.Tensor], -) -> dgl.DGLGraph: +) -> "dgl.DGLGraph": """Make the line graphs lightweight with preserved node ordering. Args: @@ -237,6 +244,8 @@ def lightweight_line_graph1( Returns: Filtered edges while preserving original node ordering """ + import dgl + # Get active edges based on filter condition active_edges = torch.logical_not( filter_condition(input_graph.edata[feature_name]) @@ -317,7 +326,7 @@ def compute_net_torque( def remove_net_torque( - g: dgl.DGLGraph, + g: "dgl.DGLGraph", forces: torch.Tensor, n_nodes: torch.Tensor, ) -> torch.Tensor: diff --git a/alignn/scripts/export_torchscript.py b/alignn/scripts/export_torchscript.py new file mode 100644 index 0000000..d01f4a3 --- /dev/null +++ b/alignn/scripts/export_torchscript.py @@ -0,0 +1,194 @@ +"""Export a trained pure-torch ALIGNN model to TorchScript. + +The saved ``.pt`` file is self-contained: it holds the model weights, +architecture code (scripted), and an entry-point ``forward(...)`` +that takes plain tensors (positions, lattice, atom_features, src, dst, +shift) and returns ``{"energy", "forces", ["stress"]}``. + +This is what a LAMMPS ``pair_alignn`` plugin (or any libtorch C++ host) +would call each MD step — the only runtime dependency is libtorch. + +Usage +----- + python alignn/scripts/export_torchscript.py \ + --checkpoint /path/to/best_model.pt \ + --config /path/to/config.json \ + --output alignn_scripted.pt + # test the scripted file on a random structure + python alignn/scripts/export_torchscript.py \ + --checkpoint ... --config ... --output out.pt --test + +Design notes +------------ +Only the ``forward_tensors`` path is scripted (a thin Wrapper module +exposes it as ``forward``). The dataclass-based training forward is +left alone. Classification and extra-features branches are not +exported — they're training-only features. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Dict + +import numpy as np +import torch +from torch import nn + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, + ALIGNNAtomWisePureConfig, +) + + +class _ScriptWrapper(nn.Module): + """LAMMPS-style entry point for libtorch. + + Takes atomic numbers (not pre-computed features) — the feature + table is baked into the inner model via ``register_species_table`` + so the scripted ``.pt`` is fully self-contained. + """ + + def __init__(self, inner: ALIGNNAtomWisePure): + super().__init__() + self.inner = inner + + def forward( + self, + positions: torch.Tensor, # (N, 3), requires_grad=True for forces + lattice: torch.Tensor, # (3, 3) + atomic_numbers: torch.Tensor, # (N,) long + src: torch.Tensor, # (E,) long + dst: torch.Tensor, # (E,) long + shift: torch.Tensor, # (E, 3) integer cell offsets + compute_stress: bool = False, + ) -> Dict[str, torch.Tensor]: + return self.inner.forward_tensors_z( + positions, lattice, atomic_numbers, src, dst, shift, compute_stress + ) + + +def build_model( + checkpoint_path: Path, + config_path: Path, + atom_features: str = "cgcnn", +) -> ALIGNNAtomWisePure: + config = json.load(open(config_path)) + mcfg = config["model"] if "model" in config else config + if mcfg.get("name") != "alignn_atomwise_pure": + # Accept a DGL-trained checkpoint's config by overriding the name. + # Architecture is identical between alignn_atomwise and + # alignn_atomwise_pure, so weights are state-dict compatible. + mcfg = dict(mcfg, name="alignn_atomwise_pure") + model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(**mcfg)) + state = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + if isinstance(state, dict) and "state_dict" in state: + state = state["state_dict"] + # Drop any previously-saved species table — it'll be rebuilt fresh. + state = {k: v for k, v in state.items() if k != "_species_table"} + model.load_state_dict(state, strict=False) + model.register_species_table(atom_features=atom_features) + model.eval() + return model + + +def _smoke_test(scripted: torch.jit.ScriptModule) -> None: + """Run a tiny Cu-FCC box through the scripted model (atomic-number input).""" + from jarvis.io.vasp.inputs import Poscar + from alignn.graphs import Graph + from alignn.torch_graph_builder import torchgraph_from_dgl + + atoms = Poscar.from_string( + """Cu +1.0 +3.6 0.0 0.0 +0.0 3.6 0.0 +0.0 0.0 3.6 +Cu +4 +direct +0.0 0.0 0.0 +0.0 0.5 0.5 +0.5 0.0 0.5 +0.5 0.5 0.0 +""" + ).atoms + g, _ = Graph.atom_dgl_multigraph( + atoms, + neighbor_strategy="fast_graph", + cutoff=4.0, + max_neighbors=12, + atom_features="cgcnn", + ) + tg = torchgraph_from_dgl(g) + + pos = torch.as_tensor(np.asarray(atoms.cart_coords), dtype=torch.float32) + pos.requires_grad_(True) + lat = torch.as_tensor(np.asarray(atoms.lattice_mat), dtype=torch.float32) + # LAMMPS-style input: atomic numbers, not feature vectors. + z_tensor = torch.as_tensor(atoms.atomic_numbers, dtype=torch.long) + shift = tg.edata["images"].float() + + out = scripted(pos, lat, z_tensor, tg.src, tg.dst, shift, True) + print("Scripted smoke test (atomic-number input):") + print(f" energy: {out['energy'].item():.6f}") + print( + f" forces: shape={tuple(out['forces'].shape)} " + f"max|F|={out['forces'].abs().max().item():.4e} " + f"sum(F)={out['forces'].detach().sum(0).tolist()}" + ) + print( + f" stress: shape={tuple(out['stress'].shape)} " + f"max|σ|={out['stress'].abs().max().item():.4e}" + ) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--checkpoint", required=True, type=Path) + ap.add_argument("--config", required=True, type=Path) + ap.add_argument("--output", required=True, type=Path) + ap.add_argument( + "--atom-features", + default=None, + help="Feature scheme for the atomic-number lookup table " + "(cgcnn/atomic_number/basic/cfid). Defaults to the value " + "found in the training config, or 'cgcnn'.", + ) + ap.add_argument( + "--test", + action="store_true", + help="Run the scripted model on a tiny Cu-FCC box after saving.", + ) + args = ap.parse_args() + + # Resolve atom_features: CLI > training config > 'cgcnn'. + cfg = json.load(open(args.config)) + atom_features = args.atom_features or cfg.get("atom_features", "cgcnn") + print(f"Loading model from {args.checkpoint} + {args.config}") + print(f" atom_features (species lookup) = {atom_features!r}") + model = build_model( + args.checkpoint, args.config, atom_features=atom_features + ) + + print("Wrapping + scripting...") + wrapper = _ScriptWrapper(model).eval() + scripted = torch.jit.script(wrapper) + + scripted.save(str(args.output)) + print(f"Saved scripted model -> {args.output}") + print( + " Input signature: (positions[N,3], lattice[3,3], " + "atomic_numbers[N], src[E], dst[E], shift[E,3], compute_stress: bool)" + ) + print(' Output: {"energy": scalar, "forces": [N,3], ["stress": [3,3]]}') + + if args.test: + scripted = torch.jit.load(str(args.output), map_location="cpu") + _smoke_test(scripted) + + +if __name__ == "__main__": + main() diff --git a/alignn/scripts/parity_dgl_vs_pure.py b/alignn/scripts/parity_dgl_vs_pure.py new file mode 100644 index 0000000..23da677 --- /dev/null +++ b/alignn/scripts/parity_dgl_vs_pure.py @@ -0,0 +1,229 @@ +"""Parity sweep: DGL ALIGNN vs pure-torch ALIGNN across the doc examples. + +Runs each scenario from ``docs/training/*.md`` twice -- once with the +DGL-backed ``alignn_atomwise`` model and once with the DGL-free +``alignn_atomwise_pure`` model -- and compares final train/val/test +losses side-by-side. Graph construction and LMDB storage flip with the +model: the DGL run uses ``alignn/lmdb_dataset.py``, the pure run uses +``alignn/pure_lmdb_dataset.py`` (both selected automatically by +``train_alignn.py`` based on ``config.model.name``). + +Scenarios covered (match the docs): + 1. Single-output regression (alignn/examples/sample_data) + 2. Classification (sample_data + classification_threshold) + 3. Multi-output regression (alignn/examples/sample_data_multi_prop) + 4. Force field (alignn/examples/sample_data_ff) + +Usage +----- + python alignn/scripts/parity_dgl_vs_pure.py # all scenarios + python alignn/scripts/parity_dgl_vs_pure.py --only ff # one scenario + python alignn/scripts/parity_dgl_vs_pure.py --epochs 1 # quick smoke + +Exits 0 on success; prints a table of metrics and |Δ| between backends. + +Notes +----- +* Graph caches ({train,val,test}_data) are deleted before each + run so the LMDB store always matches the current model backend. +* Residual drift of ~1e-2 on train/val is expected from scatter_add + nondeterminism on CUDA (accumulation order differs from DGL's kernels + and run-to-run). Set ``torch.use_deterministic_algorithms(True)`` to + eliminate at the cost of speed. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import subprocess +import time +from pathlib import Path +from typing import Dict, List, Optional + +# Repo root = two parents up from this file (alignn/scripts/ -> repo root). +REPO = Path(__file__).resolve().parents[2] +FLOAT = r"[-+]?\d+\.\d+(?:[eE][-+]?\d+)?" + + +SCENARIOS: List[Dict] = [ + { + "key": "single", + "name": "single-output regression (sample_data)", + "base_config": "alignn/examples/sample_data/config_example.json", + "root_dir": "alignn/examples/sample_data", + "extra": None, + }, + { + "key": "classification", + "name": "classification (sample_data, threshold 0.01)", + "base_config": "alignn/examples/sample_data/config_example.json", + "root_dir": "alignn/examples/sample_data", + "extra": {"classification_threshold": 0.01}, + }, + { + "key": "multi", + "name": "multi-output regression (sample_data_multi_prop)", + "base_config": "alignn/examples/sample_data/config_example.json", + "root_dir": "alignn/examples/sample_data_multi_prop", + "extra": None, + }, + { + "key": "ff", + "name": "force field (sample_data_ff)", + "base_config": "alignn/examples/sample_data_ff/config_example_atomwise.json", + "root_dir": "alignn/examples/sample_data_ff", + "extra": None, + }, +] + + +def _patch_cfg( + base_path: Path, + out_path: Path, + *, + pure: bool, + extra: Optional[dict] = None, + epochs: Optional[int] = None, +) -> None: + cfg = json.load(open(base_path)) + cfg["use_lmdb"] = True + cfg["read_existing"] = False + cfg["num_workers"] = 0 + if pure: + cfg["model"]["name"] = "alignn_atomwise_pure" + if extra: + cfg.update(extra) + if epochs is not None: + cfg["epochs"] = int(epochs) + json.dump(cfg, open(out_path, "w")) + + +def _run(cmd: List[str], log_path: Path) -> "tuple[int, float]": + t0 = time.time() + with open(log_path, "w") as f: + p = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT) + return p.returncode, time.time() - t0 + + +def _extract(log_path: Path) -> Dict: + txt = log_path.read_text() + train = re.findall(rf"Train Loss:.*\n.*?(\d+)\s+({FLOAT})", txt) + val = re.findall(rf"Val Loss:.*\n.*?(\d+)\s+({FLOAT})", txt) + test = re.search(rf"TestLoss\s+\d+\s+({FLOAT})", txt) + return { + "train_last": float(train[-1][1]) if train else None, + "val_last": float(val[-1][1]) if val else None, + "test": float(test.group(1)) if test else None, + "final_epoch": int(train[-1][0]) if train else None, + } + + +def _wipe_graph_caches(filename_stem: str) -> None: + """Remove any stale LMDB cache dirs tied to the cfg's ``filename`` stem.""" + for suffix in ("train_data", "val_data", "test_data"): + p = REPO / f"{filename_stem}{suffix}" + if p.exists(): + shutil.rmtree(p) + + +def run_scenario(s: Dict, epochs: Optional[int] = None) -> Dict[str, Dict]: + base_cfg = REPO / s["base_config"] + stem = json.load(open(base_cfg)).get("filename", "") + results: Dict[str, Dict] = {} + for mode in ("dgl", "pure"): + _wipe_graph_caches(stem) + patched_cfg = Path(f"/tmp/parity_{s['key']}_{mode}.json") + _patch_cfg( + base_cfg, + patched_cfg, + pure=(mode == "pure"), + extra=s["extra"], + epochs=epochs, + ) + out_dir = Path(f"/tmp/parity_out_{s['key']}_{mode}") + shutil.rmtree(out_dir, ignore_errors=True) + log = Path(f"/tmp/parity_{s['key']}_{mode}.log") + rc, dt = _run( + [ + "python", + "alignn/train_alignn.py", + "--root_dir", + s["root_dir"], + "--config", + str(patched_cfg), + "--output_dir", + str(out_dir), + ], + log, + ) + res = _extract(log) if rc == 0 else { + "train_last": None, + "val_last": None, + "test": None, + "final_epoch": None, + } + res.update(rc=rc, time_s=dt, log=str(log)) + results[mode] = res + status = "OK" if rc == 0 else f"FAIL rc={rc}" + print( + f" [{mode:4s}] {status:10s} " + f"train={res['train_last']} val={res['val_last']} " + f"test={res['test']} ({dt:.1f}s)" + ) + return results + + +def _fmt(x: Optional[float], width: int = 14) -> str: + if x is None: + return f"{'n/a':>{width}s}" + return f"{x:>{width}.4f}" + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--only", + choices=[s["key"] for s in SCENARIOS], + help="Run just one scenario.", + ) + ap.add_argument( + "--epochs", + type=int, + default=None, + help="Override config.epochs (default: use JSON value).", + ) + args = ap.parse_args() + + os.chdir(REPO) + scenarios = ( + [s for s in SCENARIOS if s["key"] == args.only] + if args.only + else SCENARIOS + ) + + all_results: Dict[str, Dict[str, Dict]] = {} + for s in scenarios: + print(f"\n=== {s['name']} ===") + all_results[s["name"]] = run_scenario(s, epochs=args.epochs) + + print("\n=== Parity summary ===") + hdr = f"{'scenario':52s} {'metric':8s} {'DGL':>14s} {'pure':>14s} {'|Δ|':>12s}" + print(hdr) + print("-" * len(hdr)) + for name, d in all_results.items(): + for metric in ("train_last", "val_last", "test"): + a = d.get("dgl", {}).get(metric) + b = d.get("pure", {}).get(metric) + delta = abs(a - b) if (a is not None and b is not None) else None + print( + f"{name:52s} {metric:8s} " + f"{_fmt(a)} {_fmt(b)} {_fmt(delta, width=12)}" + ) + + +if __name__ == "__main__": + main() diff --git a/alignn/torch_graph_builder.py b/alignn/torch_graph_builder.py new file mode 100644 index 0000000..fbcb258 --- /dev/null +++ b/alignn/torch_graph_builder.py @@ -0,0 +1,589 @@ +"""Pure-torch graph and line-graph builder (no DGL). + +Produces a crystal neighbor graph together with its atomistic line +graph as plain torch data, with separate two-body and three-body +cutoffs. Edge displacement vectors and bond angle cosines are torch +functions of atomic positions and the lattice, so autograd flows back +to both (forces via -dE/dx, stress via dE/dL). + +Typical use +----------- + + g, lg = build_pure_torch_graph( + atoms=jarvis_atoms, + two_body_cutoff=5.0, + three_body_cutoff=4.0, # defaults to two_body_cutoff if None + max_neighbors=12, + ) + # g.edata["r"] : (E, 3) differentiable displacements + # g.edata["images"]: (E, 3) integer cell offsets + # g.ndata["atom_features"], g.ndata["frac_coords"], g.ndata["V"] + # lg.src / lg.dst : (T,) indices into g's edges (i.e. line-graph + # nodes *are* parent edges) + # lg.edata["h"] : (T,) bond-angle cosines at the shared atom +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +import numpy as np +import torch + + +# --------------------------------------------------------------------- +# Neighbor-list primitives (moved here from alignn.graphs so importing +# this module stays DGL-free). +# --------------------------------------------------------------------- + + +def _torch_periodic_shifts( + lattice: torch.Tensor, cutoff: float +) -> torch.Tensor: + """Integer shift vectors (K, 3) covering a ``cutoff`` sphere.""" + with torch.no_grad(): + recip = 2 * math.pi * torch.linalg.inv(lattice).T + recip_len = torch.linalg.norm(recip, dim=1) + n_max = torch.ceil(cutoff * recip_len / (2 * math.pi)).to(torch.long) + ranges = [ + torch.arange(-int(n.item()), int(n.item()) + 1, device=lattice.device) + for n in n_max + ] + return torch.cartesian_prod(*ranges).to(lattice.dtype) + + +def _topk_per_source( + src: torch.Tensor, + keys: torch.Tensor, + max_neighbors: int, + num_nodes: int, +) -> torch.Tensor: + """Return indices keeping the K smallest ``keys`` per src node.""" + device = src.device + order_key = torch.argsort(keys) + src_k = src[order_key] + order_src = torch.argsort(src_k, stable=True) + perm = order_key[order_src] + src_sorted = src[perm] + E = perm.numel() + starts = torch.searchsorted( + src_sorted, torch.arange(num_nodes, device=device) + ) + within = torch.arange(E, device=device) - starts[src_sorted] + return perm[within < max_neighbors] + + +def torch_neighbor_list( + positions: torch.Tensor, + lattice: torch.Tensor, + cutoff: float, + max_neighbors: Optional[int] = None, + atoms=None, + use_matscipy_topology: bool = True, + self_tol: float = 1e-8, +): + """Torch-native periodic neighbor list with differentiable edges. + + Topology (src, dst, integer shift) is computed without gradient + tracking (matscipy when available, pure-torch fallback otherwise). + Displacement vectors are then a torch function of positions and + lattice: ``r = positions[dst] - positions[src] + shift @ lattice``. + """ + dtype = positions.dtype + device = positions.device + num_nodes = positions.shape[0] + + used_matscipy = False + if use_matscipy_topology and atoms is not None: + try: + from matscipy.neighbours import neighbour_list as _mnl + + i_np, j_np, S_np = _mnl( + "ijS", atoms.ase_converter(), float(cutoff) + ) + src = torch.from_numpy(np.ascontiguousarray(i_np)).to( + device=device, dtype=torch.long + ) + dst = torch.from_numpy(np.ascontiguousarray(j_np)).to( + device=device, dtype=torch.long + ) + shift = torch.from_numpy(np.ascontiguousarray(S_np)).to( + device=device, dtype=dtype + ) + used_matscipy = True + except ImportError: + pass + + if not used_matscipy: + shifts = _torch_periodic_shifts(lattice, cutoff) + with torch.no_grad(): + offs = shifts @ lattice + rvec_full = ( + positions[None, None, :, :] + + offs[:, None, None, :] + - positions[None, :, None, :] + ) + dist2 = rvec_full.pow(2).sum(-1) + mask = (dist2 <= cutoff * cutoff) & (dist2 > self_tol) + k_idx, i_idx, j_idx = torch.where(mask) + src = i_idx.to(torch.long) + dst = j_idx.to(torch.long) + shift = shifts[k_idx] + + r = positions[dst] - positions[src] + shift @ lattice + + if max_neighbors is not None and max_neighbors > 0 and src.numel() > 0: + with torch.no_grad(): + dist = r.norm(dim=1) + keep = _topk_per_source(src, dist, int(max_neighbors), num_nodes) + src, dst, shift, r = src[keep], dst[keep], shift[keep], r[keep] + + return src, dst, shift, r + + +@dataclass +class TorchGraph: + """Dict-based graph container. Edge list is directed. + + When a graph is the result of batching, ``batch_num_nodes`` / + ``batch_num_edges`` carry the per-subgraph counts so downstream + code (pooling, per-graph stress, etc.) can segment correctly. + """ + + num_nodes: int + src: torch.Tensor + dst: torch.Tensor + ndata: Dict[str, torch.Tensor] = field(default_factory=dict) + edata: Dict[str, torch.Tensor] = field(default_factory=dict) + batch_num_nodes: Optional[torch.Tensor] = None + batch_num_edges: Optional[torch.Tensor] = None + + @property + def num_edges(self) -> int: + """Return the number of directed edges.""" + return int(self.src.shape[0]) + + @property + def device(self) -> torch.device: + """Return the device of the underlying tensors.""" + return self.src.device + + @property + def batch_size(self) -> int: + """Return the number of subgraphs in this (possibly batched) graph.""" + if self.batch_num_nodes is None: + return 1 + return int(self.batch_num_nodes.shape[0]) + + @property + def node_batch_id(self) -> torch.Tensor: + """(N,) long tensor mapping each node to its subgraph index.""" + dev = self.src.device + if self.batch_num_nodes is None: + return torch.zeros(self.num_nodes, dtype=torch.long, device=dev) + return torch.repeat_interleave( + torch.arange(self.batch_size, device=dev), + self.batch_num_nodes, + ) + + @property + def edge_batch_id(self) -> torch.Tensor: + """(E,) long tensor mapping each edge to its subgraph index.""" + dev = self.src.device + if self.batch_num_edges is None: + return torch.zeros(self.num_edges, dtype=torch.long, device=dev) + return torch.repeat_interleave( + torch.arange(self.batch_size, device=dev), + self.batch_num_edges, + ) + + def to(self, device) -> "TorchGraph": + """Return a copy of this graph moved to ``device``.""" + return TorchGraph( + num_nodes=self.num_nodes, + src=self.src.to(device), + dst=self.dst.to(device), + ndata={k: v.to(device) for k, v in self.ndata.items()}, + edata={k: v.to(device) for k, v in self.edata.items()}, + batch_num_nodes=( + self.batch_num_nodes.to(device) + if self.batch_num_nodes is not None + else None + ), + batch_num_edges=( + self.batch_num_edges.to(device) + if self.batch_num_edges is not None + else None + ), + ) + + def to_dgl(self): + """Materialize as a ``dgl.DGLGraph`` preserving ndata / edata. + + Autograd-safe: tensor values are assigned by reference, so + gradients on ``edata['r']`` continue to flow back to any leaves + it was derived from (e.g. positions / lattice). + """ + import dgl as _dgl + + g = _dgl.graph((self.src, self.dst), num_nodes=self.num_nodes) + for k, v in self.ndata.items(): + g.ndata[k] = v + for k, v in self.edata.items(): + g.edata[k] = v + return g + + +def torch_bond_cosines( + r_ij: torch.Tensor, r_jk: torch.Tensor, eps: float = 1e-12 +) -> torch.Tensor: + """Cosine of the bond angle at atom j for triplets i -> j -> k. + + Matches ALIGNN's convention: negate the first bond so vectors + point away from j, i.e. cos = (-r_ij) . r_jk / (|r_ij| |r_jk|). + """ + num = -(r_ij * r_jk).sum(dim=-1) + denom = r_ij.norm(dim=-1).clamp_min(eps) * r_jk.norm(dim=-1).clamp_min(eps) + return (num / denom).clamp(-1.0, 1.0) + + +def _line_graph_edges( + src: torch.Tensor, + dst: torch.Tensor, + num_nodes: int, + allowed: Optional[torch.Tensor] = None, +): + """Build line-graph edges from directed parent edges. + + For each parent edge A = (u, v), emit a line-graph edge (A, B) for + every parent edge B = (v, w). When ``allowed`` is given (bool mask + of length E), only edges with allowed=True may act as A *or* B. + + Returns + ------- + lg_src, lg_dst : long tensors, indices into the parent edge list. + """ + E = int(src.shape[0]) + device = src.device + + if allowed is None: + allowed_ids = torch.arange(E, device=device) + else: + allowed_ids = torch.nonzero(allowed, as_tuple=False).squeeze(-1) + + # Sort allowed parent edges by their src node so we can find all + # outgoing edges from a node v as a contiguous slice. + sub_src = src[allowed_ids] + order = torch.argsort(sub_src, stable=True) + sorted_edge_ids = allowed_ids[order] + sorted_src = sub_src[order] + + node_range = torch.arange(num_nodes, device=device) + bucket_start = torch.searchsorted(sorted_src, node_range) + bucket_end = torch.searchsorted(sorted_src, node_range, right=True) + + # For each allowed A, count successors = (# allowed edges out of A.dst). + A_ids = allowed_ids + A_v = dst[A_ids] + starts = bucket_start[A_v] + ends = bucket_end[A_v] + counts = ends - starts + + total = int(counts.sum().item()) + if total == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty + + lg_src = torch.repeat_interleave(A_ids, counts) + + # Expand per-A ranges [start_A, start_A + count_A) into flat indices. + cum = torch.cumsum(counts, dim=0) + row_start = cum - counts + offsets = torch.arange(total, device=device) - torch.repeat_interleave( + row_start, counts + ) + pos = torch.repeat_interleave(starts, counts) + offsets + lg_dst = sorted_edge_ids[pos] + return lg_src, lg_dst + + +def build_pure_torch_graph( + atoms=None, + two_body_cutoff: float = 5.0, + three_body_cutoff: Optional[float] = None, + max_neighbors: Optional[int] = None, + atom_features: str = "cgcnn", + use_lattice_prop: bool = False, + positions: Optional[torch.Tensor] = None, + lattice: Optional[torch.Tensor] = None, + device=None, + use_matscipy_topology: bool = True, + compute_line_graph: bool = True, +): + """Build (g, lg) as pure-torch TorchGraph objects. + + Parameters + ---------- + two_body_cutoff : float + Cutoff for the main neighbor graph (edges / pair distances). + three_body_cutoff : float or None + Cutoff for triplets in the line graph. Defaults to + ``two_body_cutoff`` when None. If smaller, only edges with + ``|r| <= three_body_cutoff`` are allowed to take part in + triplets; the main graph still carries all edges up to + ``two_body_cutoff``. + max_neighbors : int or None + Per-source cap on the pair graph (keeps K closest). + positions, lattice : optional torch leaves; when ``requires_grad`` + is set, gradients flow through ``g.edata['r']`` and + ``lg.edata['h']`` back to these. + + Returns + ------- + (g, lg) if compute_line_graph else g. + """ + from jarvis.core.specie import get_node_attributes + + if three_body_cutoff is None: + three_body_cutoff = two_body_cutoff + if three_body_cutoff > two_body_cutoff: + raise ValueError( + f"three_body_cutoff ({three_body_cutoff}) must be <= " + f"two_body_cutoff ({two_body_cutoff})." + ) + + dtype = torch.get_default_dtype() + if positions is None: + positions = torch.as_tensor(np.asarray(atoms.cart_coords), dtype=dtype) + if lattice is None: + lattice = torch.as_tensor(np.asarray(atoms.lattice_mat), dtype=dtype) + if device is not None: + positions = positions.to(device) + lattice = lattice.to(device) + device = positions.device + n_atoms = int(positions.shape[0]) + + src, dst, shift, r = torch_neighbor_list( + positions=positions, + lattice=lattice, + cutoff=float(two_body_cutoff), + max_neighbors=max_neighbors, + atoms=atoms, + use_matscipy_topology=use_matscipy_topology, + ) + + sps = np.array( + [ + list(get_node_attributes(s, atom_features=atom_features)) + for s in atoms.elements + ] + ) + node_features = torch.as_tensor(sps, dtype=dtype, device=device) + frac = torch.as_tensor( + np.asarray(atoms.frac_coords), dtype=dtype, device=device + ) + vol = torch.abs(torch.det(lattice)) + + g = TorchGraph( + num_nodes=n_atoms, + src=src, + dst=dst, + ndata={ + "atom_features": node_features, + "frac_coords": frac, + "V": vol.expand(n_atoms), + }, + edata={"r": r, "images": shift}, + ) + if use_lattice_prop: + lp = np.array( + [atoms.lattice.lat_lengths(), atoms.lattice.lat_angles()] + ).flatten() + g.ndata["extra_features"] = ( + torch.as_tensor(lp, dtype=dtype, device=device) + .unsqueeze(0) + .expand(n_atoms, -1) + .contiguous() + ) + + if not compute_line_graph: + return g + + # Three-body cutoff filter: only edges short enough join triplets. + if three_body_cutoff < two_body_cutoff: + with torch.no_grad(): + allowed = r.norm(dim=1) <= three_body_cutoff + else: + allowed = None + + lg_src, lg_dst = _line_graph_edges(src, dst, n_atoms, allowed=allowed) + # Angle cosines at the shared atom, differentiable through r. + h = torch_bond_cosines(r[lg_src], r[lg_dst]) + + lg = TorchGraph( + num_nodes=g.num_edges, # line-graph nodes == parent edges + src=lg_src, + dst=lg_dst, + # Share parent edata as ndata (like DGL's shared line graph). + ndata={"r": r, "images": shift}, + edata={"h": h}, + ) + return g, lg + + +# ===================================================================== +# Batching and DGL interop +# ===================================================================== + + +def _batch_concat( + graphs: List[TorchGraph], + node_offset_source: Optional[List[int]] = None, +) -> TorchGraph: + """Concatenate graphs into one, offsetting edge indices. + + ``node_offset_source`` is a list whose cumulative sum gives the + offset applied to ``src`` / ``dst`` for each subgraph. When None, + uses each subgraph's ``num_nodes`` (standard pair-graph batching). + Set it to parent-graph edge counts when batching *line graphs*. + """ + if len(graphs) == 0: + raise ValueError("batch_torch_graphs: empty graph list.") + device = graphs[0].src.device + if node_offset_source is None: + node_offset_source = [g.num_nodes for g in graphs] + + offsets = torch.zeros(len(graphs), dtype=torch.long, device=device) + if len(graphs) > 1: + offsets[1:] = torch.cumsum( + torch.tensor( + node_offset_source[:-1], dtype=torch.long, device=device + ), + dim=0, + ) + + srcs, dsts = [], [] + for g, off in zip(graphs, offsets): + srcs.append(g.src + off) + dsts.append(g.dst + off) + + # Only keep keys present in *all* subgraphs (avoid partial batches). + ndata_keys = set(graphs[0].ndata) + edata_keys = set(graphs[0].edata) + for g in graphs[1:]: + ndata_keys &= set(g.ndata) + edata_keys &= set(g.edata) + ndata = { + k: torch.cat([g.ndata[k] for g in graphs], dim=0) for k in ndata_keys + } + edata = { + k: torch.cat([g.edata[k] for g in graphs], dim=0) for k in edata_keys + } + + return TorchGraph( + num_nodes=sum(g.num_nodes for g in graphs), + src=torch.cat(srcs), + dst=torch.cat(dsts), + ndata=ndata, + edata=edata, + batch_num_nodes=torch.tensor( + [g.num_nodes for g in graphs], dtype=torch.long, device=device + ), + batch_num_edges=torch.tensor( + [g.num_edges for g in graphs], dtype=torch.long, device=device + ), + ) + + +def batch_torch_graphs(graphs: List[TorchGraph]) -> TorchGraph: + """Batch pair graphs with offsets = cumulative ``num_nodes``.""" + return _batch_concat(graphs, node_offset_source=None) + + +def batch_torch_graph_pairs( + pairs: List, +) -> "tuple[TorchGraph, TorchGraph]": + """Batch a list of ``(g, lg)`` pairs with consistent offsetting. + + The line graph's src/dst index into the *parent* edge list, so its + node-offset source is the list of parent num_edges, not num_nodes. + """ + parents = [p[0] for p in pairs] + lgs = [p[1] for p in pairs] + g_batched = batch_torch_graphs(parents) + lg_batched = _batch_concat( + lgs, node_offset_source=[p.num_edges for p in parents] + ) + return g_batched, lg_batched + + +# ---- DGL adapter (so the pure-torch model can consume DGL input) ---- + + +def unbatch(g): + """Polymorphic unbatch — works for TorchGraph *and* DGL graph. + + For a batched ``TorchGraph``, returns a list of single-graph + ``TorchGraph`` views with ``src`` / ``dst`` reindexed to the local + node range and ndata / edata sliced accordingly. For a DGL graph, + delegates to ``dgl.unbatch``. + """ + if isinstance(g, TorchGraph): + if g.batch_num_nodes is None: + return [g] + dev = g.src.device + B = g.batch_size + node_off = torch.zeros(B + 1, dtype=torch.long, device=dev) + node_off[1:] = torch.cumsum(g.batch_num_nodes, dim=0) + edge_off = torch.zeros(B + 1, dtype=torch.long, device=dev) + edge_off[1:] = torch.cumsum(g.batch_num_edges, dim=0) + out = [] + for b in range(B): + n0 = int(node_off[b].item()) + n1 = int(node_off[b + 1].item()) + e0 = int(edge_off[b].item()) + e1 = int(edge_off[b + 1].item()) + out.append( + TorchGraph( + num_nodes=n1 - n0, + src=g.src[e0:e1] - n0, + dst=g.dst[e0:e1] - n0, + ndata={k: v[n0:n1] for k, v in g.ndata.items()}, + edata={k: v[e0:e1] for k, v in g.edata.items()}, + ) + ) + return out + import dgl as _dgl + + return _dgl.unbatch(g) + + +def torchgraph_from_dgl(g) -> TorchGraph: + """Wrap a ``dgl.DGLGraph`` as a ``TorchGraph`` (zero-copy where possible). + + Batching metadata is preserved when the DGL graph is itself a batch. + """ + src, dst = g.edges() + ndata = {k: v for k, v in g.ndata.items()} + edata = {k: v for k, v in g.edata.items()} + bnn = None + bne = None + try: + bnn_t = g.batch_num_nodes() + bne_t = g.batch_num_edges() + if bnn_t.numel() > 1: + bnn = bnn_t.to(torch.long) + bne = bne_t.to(torch.long) + except Exception: + pass + return TorchGraph( + num_nodes=int(g.num_nodes()), + src=src.to(torch.long), + dst=dst.to(torch.long), + ndata=ndata, + edata=edata, + batch_num_nodes=bnn, + batch_num_edges=bne, + ) diff --git a/alignn/train.py b/alignn/train.py index 6af1951..7672998 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -12,6 +12,8 @@ from alignn.data import get_train_val_loaders from alignn.config import TrainingConfig from alignn.models.alignn_atomwise import ALIGNNAtomWise +from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePure +from alignn.torch_graph_builder import unbatch as _graph_unbatch from alignn.models.ealignn_atomwise import eALIGNNAtomWise from alignn.models.alignn import ALIGNN from jarvis.db.jsonutils import dumpjson @@ -29,7 +31,6 @@ setup_optimizer, print_train_val_loss, ) -import dgl # from sklearn.metrics import log_loss @@ -136,6 +137,8 @@ def train_dgl( filename=config.filename, cutoff=config.cutoff, max_neighbors=config.max_neighbors, + three_body_cutoff=config.three_body_cutoff, + read_existing=config.read_existing, output_features=config.model.output_features, classification_threshold=config.classification_threshold, target_multiplication_factor=config.target_multiplication_factor, @@ -158,6 +161,7 @@ def train_dgl( config.model.classification = True _model = { "alignn_atomwise": ALIGNNAtomWise, + "alignn_atomwise_pure": ALIGNNAtomWisePure, "ealignn_atomwise": eALIGNNAtomWise, "alignn": ALIGNN, } @@ -338,7 +342,7 @@ def train_dgl( targ_stress = torch.stack( [ gg.ndata["stresses"][0] - for gg in dgl.unbatch(dats[0]) + for gg in _graph_unbatch(dats[0]) ] ).to(device) pred_stress = result["stresses"] @@ -362,7 +366,7 @@ def train_dgl( # print('unbatch',dgl.unbatch(dats[0])) additional_dat = [ gg.ndata["additional"][0] - for gg in dgl.unbatch(dats[0]) + for gg in _graph_unbatch(dats[0]) ] # print('additional_dat',additional_dat,len(additional_dat)) targ = torch.stack(additional_dat).to(device) @@ -497,7 +501,7 @@ def train_dgl( targ_stress = torch.stack( [ gg.ndata["stresses"][0] - for gg in dgl.unbatch(dats[0]) + for gg in _graph_unbatch(dats[0]) ] ).to(device) pred_stress = result["stresses"] @@ -521,7 +525,7 @@ def train_dgl( if config.model.additional_output_weight != 0: additional_dat = [ gg.ndata["additional"][0] - for gg in dgl.unbatch(dats[0]) + for gg in _graph_unbatch(dats[0]) ] # print('additional_dat',additional_dat,len(additional_dat)) targ = torch.stack(additional_dat).to(device) @@ -679,7 +683,7 @@ def train_dgl( targ_stress = torch.stack( [ gg.ndata["stresses"][0] - for gg in dgl.unbatch(dats[0]) + for gg in _graph_unbatch(dats[0]) ] ).to(device) pred_stress = result["stresses"] diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py index 95060bd..54f2721 100644 --- a/alignn/train_alignn.py +++ b/alignn/train_alignn.py @@ -397,6 +397,7 @@ def train_for_folder( cutoff=config.cutoff, cutoff_extra=config.cutoff_extra, max_neighbors=config.max_neighbors, + three_body_cutoff=config.three_body_cutoff, output_features=config.model.output_features, classification_threshold=config.classification_threshold, target_multiplication_factor=config.target_multiplication_factor, @@ -404,6 +405,8 @@ def train_for_folder( keep_data_order=config.keep_data_order, output_dir=config.output_dir, use_lmdb=config.use_lmdb, + use_pure_torch=config.model.name == "alignn_atomwise_pure", + read_existing=config.read_existing, dtype=config.dtype, ) # print("dataset", dataset[0]) diff --git a/docs/reference/config-reference.md b/docs/reference/config-reference.md new file mode 100644 index 0000000..e9245bf --- /dev/null +++ b/docs/reference/config-reference.md @@ -0,0 +1,335 @@ +# Config Reference + +Every training run in ALIGNN is controlled by a single JSON file. The schema is +enforced by two pydantic classes: `TrainingConfig` (top-level) in +[`alignn/config.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/config.py) +and a per-model config (e.g. `ALIGNNAtomWiseConfig`) in +[`alignn/models/alignn_atomwise.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/models/alignn_atomwise.py). + +At runtime the config is validated and any unknown field is rejected, so this +file is the authoritative list of what's available. + +--- + +## Top-level (`TrainingConfig`) + +### Dataset & target + +| Field | Type / default | Meaning | +|---|---|---| +| `dataset` | one of `dft_3d`, `megnet`, `qm9`, `user_data`, … (see enum) / `"dft_3d"` | Which dataset. Use `"user_data"` for your own `POSCAR`s. | +| `target` | string / `"exfoliation_energy"` | Column in `id_prop.csv` to predict. | +| `id_tag` | `"jid"` \| `"id"` \| `"_oqmd_entry_id"` / `"jid"` | Which column holds the ID. | +| `classification_threshold` | `float?` / `None` | If set, regression targets above it become class 1; enables classification mode. | +| `target_multiplication_factor` | `float?` / `None` | Scale all targets by this constant (useful when units don't match loss scale). | + +### Atoms → graph + +| Field | Default | Meaning | +|---|---|---| +| `atom_features` | `"cgcnn"` | Node feature scheme. `cgcnn` = 92-d one-hot-ish table, `atomic_number` = scalar Z, `basic` = 11 hand-picked features, `cfid` = 438 chemical descriptors. Must match `model.atom_input_features`. | +| `neighbor_strategy` | `"k-nearest"` | How to build the neighbor graph. See [Neighbor strategies](#neighbor-strategies). | +| `cutoff` | `8.0` | Radius cutoff (Å) for pair edges. | +| `cutoff_extra` | `3.0` | Fallback slack for `radius_graph` if too few neighbors are found. | +| `three_body_cutoff` | `None` | Separate cutoff (Å) for line-graph / angle triplets. Only honored by `neighbor_strategy="pure_torch"`. `None` → use `cutoff`. Must be ≤ `cutoff`. | +| `max_neighbors` | `12` | Per-source cap. | +| `use_canonize` | `True` | Canonicalize k-nearest edge order (reproducibility). | +| `compute_line_graph` | `True` | Build the bond-angle line graph. | + +#### Neighbor strategies + +| Value | Backend | Autograd through `r`? | Notes | +|---|---|---|---| +| `k-nearest` (default) | Jarvis | no | Fastest for small property-prediction cells. | +| `radius_graph` | pure torch | no | Honest cutoff; slower. | +| `radius_graph_jarvis` | Jarvis | no | Uses Jarvis's `NeighborsAnalysis`. | +| `fast_graph` | matscipy (C) | no | Same topology as `radius_graph` but ~200× faster on large cells. Requires `pip install matscipy`. | +| `torch_graph` | matscipy + torch | **yes** | Recomputes `r = positions[dst] − positions[src] + shift @ lattice` in torch → gradients flow to positions and lattice. | +| `pure_torch` | matscipy + torch | **yes** | Like `torch_graph`, plus respects `three_body_cutoff` for the line graph. | + +### Splitting + +| Field | Default | Meaning | +|---|---|---| +| `train_ratio` / `val_ratio` / `test_ratio` | `0.8 / 0.1 / 0.1` | Fractions (used when counts are not set). | +| `n_train` / `n_val` / `n_test` | `None` | Absolute counts; override the ratios when set. | +| `keep_data_order` | `True` | If `True`, no shuffling before the split — makes comparisons across runs reproducible. | +| `random_seed` | `123` | Seed for Python / numpy / torch / CUDA RNGs. | + +### Training loop + +| Field | Default | Meaning | +|---|---|---| +| `epochs` | `300` | Number of training passes. | +| `batch_size` | `64` | Per-step batch size. | +| `learning_rate` | `1e-2` | Base LR. | +| `weight_decay` | `0` | AdamW / SGD weight decay. | +| `warmup_steps` | `2000` | OneCycle warmup steps (ignored if `scheduler="none"`). | +| `criterion` | `"mse"` | Regression loss. `"l1"` is what the force-field examples use. `"zig"` / `"poisson"` are specialized. | +| `optimizer` | `"adamw"` | `"adamw"` or `"sgd"`. | +| `scheduler` | `"onecycle"` | `"onecycle"` or `"none"`. | +| `n_early_stopping` | `None` | Stop if val loss hasn't improved in this many epochs (~50 is common). | + +### Graph caching + +| Field | Default | Meaning | +|---|---|---| +| `use_lmdb` | `True` | Cache built graphs in LMDB (fast reload, memory-bounded). | +| `read_existing` | `False` | If `True`, reuse an existing LMDB cache on disk. **Must match the current model backend** (DGL vs pure-torch). Default `False` rebuilds from scratch every run to avoid footguns. | + +### Runtime & distributed + +| Field | Default | Meaning | +|---|---|---| +| `dtype` | `"float32"` | `"float32"` or `"float64"`. | +| `num_workers` | `4` | `DataLoader` workers. `0` for debugging. | +| `pin_memory` | `False` | Pinned host memory for faster H2D copies on GPU. | +| `distributed` | `False` | Enable DDP multi-process training (paired with `torchrun`). | +| `data_parallel` | `False` | Legacy single-process multi-GPU. Prefer `distributed`. | + +### Logging & outputs + +| Field | Default | Meaning | +|---|---|---| +| `output_dir` | `.` | Where checkpoints / logs / predictions are written. | +| `filename` | `"sample"` | Prefix for cache dirs (`train_data`, etc.). | +| `write_checkpoint` | `True` | Save `best_model.pt` / `current_model.pt`. | +| `write_predictions` | `True` | Dump per-sample predictions at the end. | +| `store_outputs` | `True` | Keep intermediate outputs in the history JSONs. | +| `progress` | `True` | `tqdm` progress bars. | +| `log_tensorboard` | `False` | Also write a TensorBoard run. | +| `save_dataloader` | `False` | Pickle the built DataLoader objects (large; rarely useful). | +| `standard_scalar_and_pca` | `False` | Z-score + PCA the targets (legacy). | +| `normalize_graph_level_loss` | `False` | Divide the graph-level loss by number of atoms (enable for extensive quantities mixed with intensive ones). | + +### `model` + +A nested object whose schema is chosen by `model.name`. Currently available: + +| `name` | Class | Purpose | +|---|---|---| +| `alignn` | `ALIGNNConfig` | Original ALIGNN (scalar regression only). | +| `alignn_atomwise` | `ALIGNNAtomWiseConfig` | Adds atomwise / gradwise / stresswise outputs — this is what `force-field.md` uses. | +| `alignn_atomwise_pure` | `ALIGNNAtomWisePureConfig` | **DGL-free**, scatter-based reimplementation of `alignn_atomwise`. Accepts the same config fields; TorchScript-exportable for LAMMPS. See [Pure-torch notes](#pure-torch-notes). | +| `ealignn_atomwise` | `eALIGNNAtomWiseConfig` | Equivariant variant. | + +See [Model config (`alignn_atomwise`)](#model-config-alignn_atomwise) below. + +--- + +## Model config (`alignn_atomwise`) + +All fields below apply to both `alignn_atomwise` and `alignn_atomwise_pure` +unless noted. + +### Architecture + +| Field | Default | Meaning | +|---|---|---| +| `name` | required | `"alignn_atomwise"` or `"alignn_atomwise_pure"`. | +| `alignn_layers` | `2` | Number of ALIGNN (node + edge + triplet) conv blocks. | +| `gcn_layers` | `2` | Number of subsequent node + edge conv blocks (no triplets). | +| `atom_input_features` | `1` | Width of incoming atom features. **Must match** `atom_features` (92 for `cgcnn`, 1 for `atomic_number`, 11 for `basic`, 438 for `cfid`). | +| `edge_input_features` | `80` | Number of RBF bins for bond lengths. | +| `triplet_input_features` | `40` | Number of RBF bins for bond-angle cosines. | +| `embedding_features` | `64` | Width of the two-layer MLP inside the edge/angle embeddings. | +| `hidden_features` | `64` | Core channel width running through all conv blocks. | +| `output_features` | `1` | Graph-level output dimension (>1 for multi-output regression). | + +### Output heads and loss weights + +| Field | Default | Meaning | +|---|---|---| +| `calculate_gradient` | `True` | Compute per-atom forces via `∂E/∂r`. | +| `atomwise_output_features` | `0` | Per-atom prediction dim (e.g. charges, magmoms). `0` disables. | +| `additional_output_features` | `0` | Extra per-graph head dim. `0` disables. | +| `graphwise_weight` | `1.0` | Loss weight on the graph-level target. | +| `gradwise_weight` | `1.0` | Loss weight on forces. Setting to `0` automatically disables `calculate_gradient`. | +| `stresswise_weight` | `0.0` | Loss weight on the 3×3 stress tensor. | +| `atomwise_weight` | `0.0` | Loss weight on per-atom outputs (requires `atomwise_output_features > 0`). | +| `additional_output_weight` | `0.0` | Loss weight on the additional head. | + +### Force-field specifics + +| Field | Default | Meaning | +|---|---|---| +| `grad_multiplier` | `-1` | Multiplier applied to `∂E/∂r` when converting to forces. Physics convention `F = −∂E/∂x` ⇒ `-1`. | +| `add_reverse_forces` | `True` | Sum `∂E/∂r` over both `(i→j)` **and** `(j→i)` edges (proper Newton's third law on directed graphs). | +| `lg_on_fly` | `True` | Recompute bond-angle cosines inside the autograd graph every forward (needed for exact force/stress derivatives). Leave `True` unless profiling. | +| `batch_stress` | `True` | Compute stress per graph in the batch (vs. one stress per whole batch). | +| `force_mult_natoms` | `False` | Scale forces by number of atoms (rare; use for unit conventions). | +| `energy_mult_natoms` | `True` | Treat the graph-head output as energy *per atom* and multiply by `natoms` before taking the gradient. Standard for ALIGNN-FF. | +| `stress_multiplier` | `1.0` | Final scalar on the stress output (for unit fixes). | +| `include_pos_deriv` | `False` | Alternative force path: differentiate through Cartesian positions directly (rather than through edge displacements). Currently experimental; the `_pure` variant accepts the field but ignores it. | + +### Cutoff shaping + +| Field | Default | Meaning | +|---|---|---| +| `use_cutoff_function` | `False` | Apply a smooth polynomial envelope to distances before RBF. Useful when `cutoff` is tight. | +| `multiply_cutoff` | `False` | If `use_cutoff_function=True`, multiply the edge embedding by the envelope (instead of the bondlength). | +| `inner_cutoff` | `3.0` | Envelope radius (Å); envelope = 1 below this, smoothly → 0 up to `cutoff`. | +| `exponent` | `5` | Exponent of the smooth-cutoff polynomial. | + +### Short-distance penalty + +| Field | Default | Meaning | +|---|---|---| +| `use_penalty` | `True` | Add a repulsive penalty when any bondlength falls below `penalty_threshold`. Prevents the model from hallucinating attractive wells at unphysical distances — important for MD stability. | +| `penalty_factor` | `0.1` | Strength. | +| `penalty_threshold` | `1.0` | Distance (Å) below which the penalty kicks in. | + +### Miscellaneous + +| Field | Default | Meaning | +|---|---|---| +| `link` | `"identity"` | Output activation. `"log"` applies `exp()` (positive outputs), `"logit"` applies `sigmoid()`. | +| `classification` | `False` | Two-class NLL head (set automatically when `classification_threshold` is non-null at the training level). | +| `zero_inflated` | `False` | Zero-inflated Gaussian output (specialized). | +| `extra_features` | `0` | If > 0, concatenate this many pre-computed fingerprint features per atom (e.g. Gong et al.). | + +--- + +## Pure-torch notes + +`model.name = "alignn_atomwise_pure"` swaps the DGL message-passing core for +`scatter_add` / `scatter_mean` primitives. The **forward outputs are numerically +equivalent** to `alignn_atomwise` — parity sweep across all doc examples shows +`|Δ| ≤ 0.04` on all losses after 3 epochs (force-field scenario is bit-identical, +see `alignn/scripts/parity_dgl_vs_pure.py`). Residual drift is `scatter_add` +nondeterminism on CUDA, not a modelling difference. + +Not carried over from `alignn_atomwise` (config fields accepted but ignored): + +- `include_pos_deriv` +- Some `extra_features` code paths — experimental in the DGL model too. + +When the pure model is selected, the training pipeline automatically: + +- Uses `alignn/pure_lmdb_dataset.py` (pickles `TorchGraph` objects, no DGL). +- Uses `PureTorchLMDBDataset.collate_line_graph` for batching. +- Leaves dataloader / loss / metrics / checkpoints untouched. + +For LAMMPS integration, call +[`alignn/scripts/export_torchscript.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/scripts/export_torchscript.py) — it scripts the +`forward_tensors_z(positions, lattice, atomic_numbers, src, dst, shift, +compute_stress)` entry point and bakes the atomic-number → feature lookup into +the `.pt` so the C++ host only needs atomic numbers. + +--- + +## Minimal examples + +### Single-output regression + +```json +{ + "dataset": "user_data", + "target": "formation_energy_peratom", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "cutoff": 8.0, + "max_neighbors": 12, + "epochs": 100, + "batch_size": 32, + "learning_rate": 1e-3, + "criterion": "mse", + "model": { + "name": "alignn_atomwise", + "atom_input_features": 92, + "alignn_layers": 4, + "gcn_layers": 4, + "hidden_features": 256, + "output_features": 1, + "calculate_gradient": false, + "graphwise_weight": 1.0, + "gradwise_weight": 0.0, + "stresswise_weight": 0.0, + "use_penalty": false + } +} +``` + +### ALIGNN-FF (energy + forces + stress) + +```json +{ + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "fast_graph", + "cutoff": 6.0, + "max_neighbors": 12, + "epochs": 200, + "batch_size": 2, + "learning_rate": 1e-3, + "criterion": "l1", + "model": { + "name": "alignn_atomwise", + "atom_input_features": 92, + "alignn_layers": 3, + "gcn_layers": 3, + "hidden_features": 128, + "output_features": 1, + "calculate_gradient": true, + "graphwise_weight": 0.85, + "gradwise_weight": 0.05, + "stresswise_weight": 0.05, + "add_reverse_forces": true, + "lg_on_fly": true, + "use_penalty": true + } +} +``` + +### Pure-torch FF with separate 3-body cutoff + +```json +{ + "neighbor_strategy": "pure_torch", + "cutoff": 6.0, + "three_body_cutoff": 4.0, + "max_neighbors": 12, + "model": { "name": "alignn_atomwise_pure", "…": "…" } +} +``` + +### Classification + +```json +{ + "classification_threshold": 0.01, + "target": "is_metal", + "criterion": "mse", + "model": { + "name": "alignn_atomwise", + "classification": true, + "graphwise_weight": 1.0, + "gradwise_weight": 0.0, + "calculate_gradient": false + } +} +``` + +--- + +## Common gotchas + +- **`atom_input_features` vs `atom_features`.** These must match (cgcnn → 92, + atomic_number → 1, basic → 11, cfid → 438). A mismatch raises a shape error + at the first forward pass, not at config parse. +- **`read_existing`.** Keep `false` unless you're sure the cache matches the + current model backend + graph settings. Backend-mismatch (DGL vs pure) is + caught and raises a clear error; parameter drift (different cutoff etc.) is + not caught. +- **`gradwise_weight=0`** silently disables `calculate_gradient` in the model + code. If you think forces should be training but see `Grad=0.0000`, check + both flags. +- **`use_penalty=true` with a loose cutoff** can inflate training loss when + your data legitimately has short bonds (e.g. H-containing systems). Lower + `penalty_threshold` or set `use_penalty=false` in that case. +- **`batch_stress=false` + large batch** gives one stress for the whole + batch, not per-graph — almost never what you want. +- **`lg_on_fly=false`** caches triplet angle cosines; saves time but forces + computed by autograd will be missing the angle-contribution unless you + recompute them. Leave `true` for force-field training. diff --git a/mkdocs.yml b/mkdocs.yml index 765ba99..ce97b26 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -84,6 +84,7 @@ nav: - Web Apps: usage/webapps.md - Reference: - Package Overview: api.md + - Config Reference: reference/config-reference.md - Performance: performance.md - Notes & FAQ: notes.md - References: references.md From 8a584d5cee5c84a672a1b6e0def7679afde5a098 Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 21:45:55 -0400 Subject: [PATCH 02/34] Config --- .../sample_data_ff/mlearn_data/Cu/config.json | 127 ++++++++--------- .../sample_data_ff/mlearn_data/Ge/config.json | 127 ++++++++--------- .../sample_data_ff/mlearn_data/Li/config.json | 127 ++++++++--------- .../sample_data_ff/mlearn_data/Mo/config.json | 127 ++++++++--------- .../sample_data_ff/mlearn_data/Ni/config.json | 127 ++++++++--------- .../mlearn_data/Ni/config_rad.json | 68 --------- .../sample_data_ff/mlearn_data/Si/config.json | 129 +++++++++--------- 7 files changed, 385 insertions(+), 447 deletions(-) delete mode 100644 alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json index a1baace..6b0b5a0 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json @@ -1,65 +1,66 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 31, - "n_test": 31, - "n_train": 262, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 13, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 2, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 164, - "hidden_features": 256, - "output_features": 1, - "force_mult_natoms": true, - "grad_multiplier": -1, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 50.00, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "link": "identity", - "zero_inflated": false, - "classification": false - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 31, + "n_test": 31, + "n_train": 262, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 5, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 13, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 2, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 164, + "hidden_features": 256, + "output_features": 1, + "force_mult_natoms": true, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 50.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json index 303f1fc..809133f 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json @@ -1,65 +1,66 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 25, - "n_test": 25, - "n_train": 228, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 12, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 4, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 10.0, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "link": "identity", - "force_mult_natoms": true, - "zero_inflated": false, - "classification": false - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 25, + "n_test": 25, + "n_train": 228, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 5, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise_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": 10.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "force_mult_natoms": true, + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json index 5efce30..dcdafad 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json @@ -1,65 +1,66 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 29, - "n_test": 29, - "n_train": 241, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 12, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 4, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 10.0, - "stresswise_weight": 0.0, - "force_mult_natoms": true, - "atomwise_weight": 0.0, - "link": "identity", - "zero_inflated": false, - "classification": false - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 29, + "n_test": 29, + "n_train": 241, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 5, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise_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": 10.0, + "stresswise_weight": 0.0, + "force_mult_natoms": true, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json index 639acdb..1d661fe 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json @@ -1,65 +1,66 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 23, - "n_test": 23, - "n_train": 194, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 12, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 4, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 10.0, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "link": "identity", - "zero_inflated": false, - "force_mult_natoms": true, - "classification": false - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 23, + "n_test": 23, + "n_train": 194, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 5, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise_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": 10.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "force_mult_natoms": true, + "classification": false + }, + "three_body_cutoff": 4.0 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json index 953555a..c0a1962 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json @@ -1,65 +1,66 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 31, - "n_test": 31, - "n_train": 263, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 12, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 2, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "force_mult_natoms": true, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 1.0, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "link": "identity", - "zero_inflated": false, - "classification": false - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 31, + "n_test": 31, + "n_train": 263, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 5, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 2, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "force_mult_natoms": true, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 1.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json deleted file mode 100644 index 24bf1da..0000000 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_rad.json +++ /dev/null @@ -1,68 +0,0 @@ -{ - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "radius_graph", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 31, - "n_test": 31, - "n_train": 263, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 4.0, - "max_neighbors": 12, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 2, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "force_mult_natoms": true, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 10.0, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "link": "identity", - "use_cutoff_function": true, - "zero_inflated": false, - "classification": false, - "inner_cutoff": 3.0, - "lg_on_fly": false - } -} diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json index 4e9d7f5..804416d 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json @@ -1,66 +1,67 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 25, - "n_test": 25, - "n_train": 214, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 12, - "keep_data_order": true, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "./", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 4, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 1.0, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "force_mult_natoms": true, - "link": "identity", - "zero_inflated": false, - "classification": false, - "use_cutoff_function":true - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 25, + "n_test": 25, + "n_train": 214, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 5, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 4, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 1.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "force_mult_natoms": true, + "link": "identity", + "zero_inflated": false, + "classification": false, + "use_cutoff_function": true + }, + "three_body_cutoff": 4.0 } From 075f4636588169fbfcd0b6025ffad51e06630ac9 Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 21:57:39 -0400 Subject: [PATCH 03/34] Update docs --- alignn/pure_lmdb_dataset.py | 258 +++++++++++++++++++++ alignn/scripts/bench_dgl_vs_pure.py | 345 ++++++++++++++++++++++++++++ docs/examples-colab.md | 18 ++ docs/index.md | 1 + mkdocs.yml | 1 + 5 files changed, 623 insertions(+) create mode 100644 alignn/pure_lmdb_dataset.py create mode 100644 alignn/scripts/bench_dgl_vs_pure.py create mode 100644 docs/examples-colab.md diff --git a/alignn/pure_lmdb_dataset.py b/alignn/pure_lmdb_dataset.py new file mode 100644 index 0000000..b361a3a --- /dev/null +++ b/alignn/pure_lmdb_dataset.py @@ -0,0 +1,258 @@ +"""LMDB-backed dataset yielding TorchGraph objects (no DGL). + +Mirrors the API of ``alignn.lmdb_dataset`` but serializes +``TorchGraph`` dataclasses via pickle, and batches them with +``batch_torch_graph_pairs`` at the DataLoader boundary. +""" + +from __future__ import annotations + +import os +import pickle as pk +from typing import List, Tuple + +import lmdb +import numpy as np +import torch +from jarvis.core.atoms import Atoms +from torch.utils.data import Dataset +from tqdm import tqdm + +from alignn.graphs import Graph +from alignn.torch_graph_builder import ( + TorchGraph, + batch_torch_graph_pairs, + batch_torch_graphs, + torchgraph_from_dgl, +) + + +def prepare_pure_batch(batch, device=None, non_blocking=False): + """Move a pure-torch batch to device; mirrors prepare_line_graph_batch.""" + g, lg, t, id = batch + return ( + (g.to(device), lg.to(device)), + t.to(device, non_blocking=non_blocking), + ) + + +def _graph_to_torchgraph(g, lg=None): + """Accept DGL or TorchGraph; always return TorchGraph (pair).""" + if isinstance(g, TorchGraph): + return g, (lg if lg is not None else None) + tg = torchgraph_from_dgl(g) + tlg = torchgraph_from_dgl(lg) if lg is not None else None + return tg, tlg + + +class PureTorchLMDBDataset(Dataset): + """Dataset of crystal TorchGraphs using LMDB.""" + + def __init__( + self, lmdb_path: str = "", line_graph: bool = True, ids=None + ): + """Open the LMDB at ``lmdb_path`` read-only.""" + super().__init__() + self.lmdb_path = lmdb_path + self.ids = ids or [] + self.line_graph = line_graph + self.env = lmdb.open(self.lmdb_path, readonly=True, lock=False) + with self.env.begin() as txn: + self.length = txn.stat()["entries"] + self.prepare_batch = prepare_pure_batch + + def __len__(self): + """Return the number of records in the LMDB.""" + return self.length + + def __getitem__(self, idx): + """Load and unpickle the ``idx``-th record.""" + with self.env.begin() as txn: + serialized = txn.get(f"{idx}".encode()) + if self.line_graph: + g, lg, lattice, label = pk.loads(serialized) + return g, lg, lattice, label + g, lattice, label = pk.loads(serialized) + return g, lattice, label + + def close(self): + """Close the LMDB environment (idempotent).""" + try: + self.env.close() + except Exception: + pass + + def __del__(self): + """Ensure the LMDB handle is released on GC.""" + self.close() + + @staticmethod + def collate(samples: List[Tuple[TorchGraph, torch.Tensor, torch.Tensor]]): + """Collate ``(g, lattice, label)`` samples into a batched triple.""" + graphs, lattices, labels = map(list, zip(*samples)) + batched = batch_torch_graphs(graphs) + if labels[0].dim() > 0: + return batched, torch.stack(lattices), torch.stack(labels) + return batched, torch.stack(lattices), torch.tensor(labels) + + @staticmethod + def collate_line_graph( + samples: List[ + Tuple[TorchGraph, TorchGraph, torch.Tensor, torch.Tensor] + ], + ): + """Collate ``(g, lg, lattice, label)`` samples into a batched tuple.""" + graphs, line_graphs, lattices, labels = map(list, zip(*samples)) + g_b, lg_b = batch_torch_graph_pairs(list(zip(graphs, line_graphs))) + if labels[0].dim() > 0: + return g_b, lg_b, torch.stack(lattices), torch.stack(labels) + return g_b, lg_b, torch.stack(lattices), torch.tensor(labels) + + +def _attach_node_payload( + g: TorchGraph, key: str, value: np.ndarray, natoms: int +): + """Tile a per-structure tensor across nodes, like the DGL loader.""" + dtype = torch.get_default_dtype() + arr = np.asarray(value) + # Per-node array already: first dim matches num nodes. + if arr.ndim >= 1 and arr.shape[0] == natoms: + g.ndata[key] = torch.as_tensor(arr, dtype=dtype) + return + # Otherwise, broadcast / tile across nodes. + tiled = np.broadcast_to(arr, (natoms,) + arr.shape).copy() + g.ndata[key] = torch.as_tensor(tiled, dtype=dtype) + + +def get_torch_dataset( + dataset=None, + id_tag="jid", + target="", + target_atomwise="", + target_grad="", + target_stress="", + target_additional_output="", + neighbor_strategy="pure_torch", + atom_features="cgcnn", + use_canonize="", + name="", + line_graph=True, + cutoff=8.0, + cutoff_extra=3.0, + max_neighbors=12, + three_body_cutoff=None, + classification=False, + sampler=None, + output_dir=".", + tmp_name="dataset", + map_size=1e12, + read_existing=False, + dtype="float32", +): + """Build or load a PureTorchLMDBDataset from a list of records.""" + dataset = dataset or [] + vals = np.array([ii[target] for ii in dataset]) + print("data range", np.max(vals), np.min(vals)) + print("line_graph", line_graph) + os.makedirs(output_dir, exist_ok=True) + with open(os.path.join(output_dir, tmp_name + "_data_range"), "w") as f: + f.write(f"Max={np.max(vals)}\nMin={np.min(vals)}\n") + + if os.path.exists(tmp_name) and read_existing: + # Validate the cache was built for the pure-torch backend. + _env = lmdb.open(tmp_name, readonly=True, lock=False) + with _env.begin() as _txn: + _probe = _txn.get(b"0") + _env.close() + if _probe is not None: + _sample = pk.loads(_probe) + if not isinstance(_sample[0], TorchGraph): + raise RuntimeError( + f"LMDB cache at '{tmp_name}' contains " + f"{type(_sample[0]).__name__} records, not " + "TorchGraph. Delete the stale cache (e.g. " + f"`rm -rf {tmp_name}`) and rerun — it was built " + "for a different model backend." + ) + ids = [d[id_tag] for d in dataset] + print("Reading dataset", tmp_name) + return PureTorchLMDBDataset( + lmdb_path=tmp_name, line_graph=line_graph, ids=ids + ) + + ids = [] + # Fresh build: wipe any pre-existing cache dir to avoid mixing old + # records (possibly from a different model backend) with new writes. + if os.path.exists(tmp_name): + import shutil + + shutil.rmtree(tmp_name) + env = lmdb.open(tmp_name, map_size=int(map_size)) + with env.begin(write=True) as txn: + for idx, d in tqdm(enumerate(dataset), total=len(dataset)): + ids.append(d[id_tag]) + atoms = Atoms.from_dict(d["atoms"]) + raw = Graph.atom_dgl_multigraph( + atoms, + cutoff=float(cutoff), + max_neighbors=max_neighbors, + atom_features=atom_features, + compute_line_graph=line_graph, + use_canonize=use_canonize, + cutoff_extra=cutoff_extra, + neighbor_strategy=neighbor_strategy, + three_body_cutoff=three_body_cutoff, + dtype=dtype, + ) + if line_graph: + g, lg = raw + g, lg = _graph_to_torchgraph(g, lg) + else: + g = raw + g, _ = _graph_to_torchgraph(g) + + lattice = torch.as_tensor( + atoms.lattice_mat, dtype=torch.get_default_dtype() + ) + label = torch.as_tensor( + d[target], dtype=torch.get_default_dtype() + ) + natoms = len(d["atoms"]["elements"]) + if classification: + label = label.long() + if "extra_features" in d: + _attach_node_payload( + g, + "extra_features", + np.asarray(d["extra_features"]), + natoms, + ) + if target_atomwise: + g.ndata[target_atomwise] = torch.as_tensor( + np.asarray(d[target_atomwise]), + dtype=torch.get_default_dtype(), + ) + if target_grad: + arr = np.asarray(d[target_grad]) + g.ndata[target_grad] = torch.as_tensor( + arr, dtype=torch.get_default_dtype() + ) + if target_stress: + stress = np.asarray(d[target_stress]) + _attach_node_payload(g, target_stress, stress, natoms) + if target_additional_output: + _attach_node_payload( + g, + target_additional_output, + np.asarray(d[target_additional_output]), + natoms, + ) + + if line_graph: + txn.put(f"{idx}".encode(), pk.dumps((g, lg, lattice, label))) + else: + txn.put(f"{idx}".encode(), pk.dumps((g, lattice, label))) + env.close() + return PureTorchLMDBDataset( + lmdb_path=tmp_name, line_graph=line_graph, ids=ids + ) diff --git a/alignn/scripts/bench_dgl_vs_pure.py b/alignn/scripts/bench_dgl_vs_pure.py new file mode 100644 index 0000000..900e607 --- /dev/null +++ b/alignn/scripts/bench_dgl_vs_pure.py @@ -0,0 +1,345 @@ +"""Scaling benchmark: DGL ALIGNN vs pure-torch ALIGNN. + +Builds Cu-FCC supercells of increasing size, runs one forward + backward +(energy + forces + stress) through each backend, and reports time and +peak memory. + +Both backends receive the *same* graph topology: we build a single DGL +graph via matscipy (``neighbor_strategy="fast_graph"``) and feed it to +``ALIGNNAtomWise`` directly; the pure-torch model gets the same graph +converted to ``TorchGraph`` at the model boundary. So the numbers +below isolate **model-level** cost (message passing, autograd, +readout) — not the neighbor-list step. + +Usage +----- + python alignn/scripts/bench_dgl_vs_pure.py + python alignn/scripts/bench_dgl_vs_pure.py \\ + --sizes 1,2,3,4,5 --cutoff 4.0 --hidden 64 \\ + --alignn-layers 1 --gcn-layers 1 --repeat 3 + +Writes a printed table; with ``--output foo`` also saves ``foo.npz`` +and, if matplotlib is available, ``foo.png``. +""" + +from __future__ import annotations + +import argparse +import gc +import time +from pathlib import Path +from typing import Dict, List + +import numpy as np +import torch +from ase.build.supercells import make_supercell +from jarvis.io.vasp.inputs import Poscar + +from alignn.graphs import Graph +from alignn.models.alignn_atomwise import ( + ALIGNNAtomWise, + ALIGNNAtomWiseConfig, +) +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, + ALIGNNAtomWisePureConfig, +) +from alignn.torch_graph_builder import torchgraph_from_dgl + + +CU_POSCAR = """Cu +1.0 +3.6 0.0 0.0 +0.0 3.6 0.0 +0.0 0.0 3.6 +Cu +4 +direct +0.0 0.0 0.0 +0.0 0.5 0.5 +0.5 0.0 0.5 +0.5 0.5 0.0 +""" + + +def _build_models(args, device) -> "tuple[torch.nn.Module, torch.nn.Module]": + """Construct both backends from the same hyperparameters and weights.""" + common = dict( + alignn_layers=args.alignn_layers, + gcn_layers=args.gcn_layers, + atom_input_features=92, + hidden_features=args.hidden, + output_features=1, + calculate_gradient=True, + graphwise_weight=0.85, + gradwise_weight=0.05, + stresswise_weight=0.05, + atomwise_weight=0.0, + add_reverse_forces=True, + lg_on_fly=True, + use_penalty=True, + ) + torch.manual_seed(0) + m_dgl = ALIGNNAtomWise( + ALIGNNAtomWiseConfig(name="alignn_atomwise", **common) + ).to(device) + torch.manual_seed(0) + m_pure = ALIGNNAtomWisePure( + ALIGNNAtomWisePureConfig(name="alignn_atomwise_pure", **common) + ).to(device) + # Copy DGL weights into pure so any difference is purely mechanical. + sd_dgl = m_dgl.state_dict() + sd_pure = m_pure.state_dict() + for k in sd_pure: + if k in sd_dgl and sd_dgl[k].shape == sd_pure[k].shape: + sd_pure[k] = sd_dgl[k].clone() + m_pure.load_state_dict(sd_pure, strict=False) + m_dgl.train() + m_pure.train() + return m_dgl, m_pure + + +def _build_graph(atoms, cutoff: float, max_neighbors: int): + g, lg = Graph.atom_dgl_multigraph( + atoms, + neighbor_strategy="fast_graph", + cutoff=cutoff, + max_neighbors=max_neighbors, + atom_features="cgcnn", + ) + return g, lg + + +def _run_once( + model, + graphs, + lat, + device: torch.device, + is_pure: bool, +): + """One forward + backward. Returns elapsed seconds and peak memory (bytes).""" + g, lg = graphs + if is_pure: + g_in, lg_in = torchgraph_from_dgl(g), torchgraph_from_dgl(lg) + else: + g_in, lg_in = g, lg + lat_in = lat + + # Sync + reset peak memory. + if device.type == "cuda": + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + + t0 = time.perf_counter() + out = model((g_in, lg_in, lat_in)) + loss = ( + out["out"].sum() + + out["grad"].pow(2).sum() + + out["stresses"].pow(2).sum() + ) + loss.backward() + if device.type == "cuda": + torch.cuda.synchronize(device) + dt = time.perf_counter() - t0 + + peak = ( + torch.cuda.max_memory_allocated(device) + if device.type == "cuda" + else 0 + ) + # Free autograd graph / grads before next iter. + model.zero_grad(set_to_none=True) + del out, loss + gc.collect() + if device.type == "cuda": + torch.cuda.empty_cache() + return dt, int(peak) + + +def _bench_one_size( + model, + graphs, + lat, + device, + is_pure: bool, + warmup: int, + repeat: int, +): + """Warm up then time ``repeat`` iterations; return median time + peak mem.""" + for _ in range(warmup): + _run_once(model, graphs, lat, device, is_pure) + times, peaks = [], [] + for _ in range(repeat): + t, p = _run_once(model, graphs, lat, device, is_pure) + times.append(t) + peaks.append(p) + return float(np.median(times)), max(peaks) + + +def run(args) -> Dict[str, np.ndarray]: + device = torch.device(args.device) + print(f"device: {device}") + if device.type == "cuda": + print(f" {torch.cuda.get_device_name(device)}") + + m_dgl, m_pure = _build_models(args, device) + n_params = sum(p.numel() for p in m_dgl.parameters()) + print(f"model parameters: {n_params}") + + base = Poscar.from_string(CU_POSCAR).atoms.ase_converter() + sizes = [int(s) for s in args.sizes.split(",")] + + rows: List[dict] = [] + print() + hdr = ( + f"{'N':>3} {'atoms':>6} {'edges':>8} " + f"{'DGL ms':>10} {'pure ms':>10} {'speedup':>8} " + f"{'DGL MB':>10} {'pure MB':>10} {'mem ratio':>10}" + ) + print(hdr) + print("-" * len(hdr)) + + for n in sizes: + sc = make_supercell(base, [[n, 0, 0], [0, n, 0], [0, 0, n]]) + # Convert to jarvis Atoms for the neighbor-list pipeline. + from jarvis.core.atoms import Atoms as JAtoms + + jatoms = JAtoms( + lattice_mat=np.asarray(sc.cell), + coords=np.asarray(sc.positions), + elements=list(sc.get_chemical_symbols()), + cartesian=True, + ) + try: + g, lg = _build_graph(jatoms, args.cutoff, args.max_neighbors) + except Exception as exc: + print(f" graph build failed at N={n}: {exc}") + break + lat = torch.as_tensor( + np.asarray(jatoms.lattice_mat), dtype=torch.float32, device=device + ) + g = g.to(device) + lg = lg.to(device) + + try: + t_dgl, m_dgl_peak = _bench_one_size( + m_dgl, (g, lg), lat, device, False, args.warmup, args.repeat + ) + except (RuntimeError, torch.cuda.OutOfMemoryError) as exc: + print(f" DGL OOM/error at N={n}: {type(exc).__name__}") + break + try: + t_pure, m_pure_peak = _bench_one_size( + m_pure, (g, lg), lat, device, True, args.warmup, args.repeat + ) + except (RuntimeError, torch.cuda.OutOfMemoryError) as exc: + print(f" pure OOM/error at N={n}: {type(exc).__name__}") + break + + n_atoms = int(sc.positions.shape[0]) + n_edges = int(g.num_edges()) + speedup = t_dgl / t_pure if t_pure > 0 else float("nan") + mem_ratio = ( + m_pure_peak / m_dgl_peak if m_dgl_peak > 0 else float("nan") + ) + print( + f"{n:>3} {n_atoms:>6} {n_edges:>8} " + f"{t_dgl * 1000:>10.2f} {t_pure * 1000:>10.2f} " + f"{speedup:>8.2f}x " + f"{m_dgl_peak / 1024**2:>10.1f} " + f"{m_pure_peak / 1024**2:>10.1f} " + f"{mem_ratio:>10.2f}x" + ) + rows.append( + dict( + n=n, + atoms=n_atoms, + edges=n_edges, + dgl_time_s=t_dgl, + pure_time_s=t_pure, + dgl_peak_bytes=m_dgl_peak, + pure_peak_bytes=m_pure_peak, + ) + ) + + result = {k: np.array([r[k] for r in rows]) for k in rows[0]} if rows else {} + + if args.output and rows: + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + np.savez(str(out.with_suffix(".npz")), **result) + print(f"\nSaved data to {out.with_suffix('.npz')}") + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.5)) + n_arr = result["atoms"] + ax1.plot( + n_arr, result["dgl_time_s"] * 1000, "-o", label="DGL" + ) + ax1.plot( + n_arr, result["pure_time_s"] * 1000, "-s", label="pure-torch" + ) + ax1.set( + xlabel="# atoms", ylabel="time / iter (ms)", + title="Forward + backward (E+F+σ)", + ) + ax1.set_xscale("log") + ax1.set_yscale("log") + ax1.grid(True, which="both", alpha=0.3) + ax1.legend() + + ax2.plot(n_arr, result["dgl_peak_bytes"] / 1024**2, "-o", label="DGL") + ax2.plot( + n_arr, result["pure_peak_bytes"] / 1024**2, "-s", + label="pure-torch", + ) + ax2.set( + xlabel="# atoms", ylabel="peak memory (MB)", + title="Peak GPU memory", + ) + ax2.set_xscale("log") + ax2.set_yscale("log") + ax2.grid(True, which="both", alpha=0.3) + ax2.legend() + + plt.tight_layout() + png = out.with_suffix(".png") + plt.savefig(str(png), dpi=160) + print(f"Saved plot to {png}") + except ImportError: + pass + return result + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--sizes", default="1,2,3,4,5", + help="Comma-separated supercell factors (N -> 4*N^3 atoms).", + ) + ap.add_argument("--cutoff", type=float, default=4.0) + ap.add_argument("--max-neighbors", type=int, default=12) + ap.add_argument("--hidden", type=int, default=64) + ap.add_argument("--alignn-layers", type=int, default=1) + ap.add_argument("--gcn-layers", type=int, default=1) + ap.add_argument("--warmup", type=int, default=2) + ap.add_argument("--repeat", type=int, default=3) + ap.add_argument( + "--device", + default="cuda" if torch.cuda.is_available() else "cpu", + ) + ap.add_argument( + "--output", + default=None, + help="Path stem; writes .npz and (if matplotlib) .png.", + ) + args = ap.parse_args() + run(args) + + +if __name__ == "__main__": + main() diff --git a/docs/examples-colab.md b/docs/examples-colab.md new file mode 100644 index 0000000..185b246 --- /dev/null +++ b/docs/examples-colab.md @@ -0,0 +1,18 @@ +# Google Colab Notebooks + +Ready-to-run Colab notebooks covering installation, property prediction, force-field +training, and pretrained-model usage. Click the badge in any row to open the notebook in +Colab. + +[colab-badge]: https://colab.research.google.com/assets/colab-badge.svg + +| Notebook | Open in Colab | Description | +| --- | --- | --- | +| [Regression task (graph-wise prediction)](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Single-output regression model for exfoliation energies of 2D materials. | +| [ML force-field training from scratch](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Train an ALIGNN-FF machine-learning force-field for Silicon. | +| [ALIGNN-FF: relaxation, EV curve, phonons, interfaces](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Using the pretrained ALIGNN-FF for structure relaxation, EV curves, phonons, interface gamma surfaces, and interface separation. | +| [Scaling / timing comparison](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | Analyze scaling/timing of universal MLFFs. | +| [Melt-Quench MD](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | Generate amorphous structures via molecular dynamics. | +| [Miscellaneous training tasks](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | Single-output (formation energy, bandgap), multi-output (phonon/electron DOS), classification (metal vs non-metal), and pretrained-model usage. | +| [Superconductor Tc](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | Train a model for superconductor transition temperature. | +| [Build `id_prop.json` from VASP runs](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | [![Open In Colab][colab-badge]](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | Compile `vasprun.xml` files into an `id_prop.json` for ALIGNN-FF training. | diff --git a/docs/index.md b/docs/index.md index 933776b..c1bd20c 100644 --- a/docs/index.md +++ b/docs/index.md @@ -28,6 +28,7 @@ dataset (~75,000 materials and 4M+ energy/force entries) and supports any combin ## Quick links - [Install ALIGNN](installation.md) +- [Colab notebooks](examples-colab.md) - [Train your first model](training/single-output-regression.md) - [Use a pretrained model](pretrained/index.md) - [Run the ALIGNN-FF ASE calculator](usage/ase-calculator.md) diff --git a/mkdocs.yml b/mkdocs.yml index ce97b26..f57bf7d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -68,6 +68,7 @@ extra: nav: - Home: index.md - Installation: installation.md + - Colab Notebooks: examples-colab.md - Training: - Dataset Format: training/dataset-format.md - Single-Output Regression: training/single-output-regression.md From 704fc22bf30c60dbf828544f38afa9a53fb814cd Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 22:02:37 -0400 Subject: [PATCH 04/34] Update docs --- .github/workflows/docs.yml | 2 +- alignn/__init__.py | 2 +- .../sample_data_ff/mlearn_data/Cu/config.json | 2 +- .../sample_data_ff/mlearn_data/Ge/config.json | 2 +- .../sample_data_ff/mlearn_data/Li/config.json | 2 +- .../sample_data_ff/mlearn_data/Mo/config.json | 2 +- .../sample_data_ff/mlearn_data/Ni/config.json | 2 +- .../sample_data_ff/mlearn_data/Si/config.json | 2 +- .../mlearn_data/all/config_example.json | 133 +++++++++--------- setup.py | 2 +- 10 files changed, 76 insertions(+), 75 deletions(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 0bf9aee..64d8dc7 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -2,7 +2,7 @@ name: docs on: push: - branches: [main, develop] + branches: [main] paths: - docs/** - mkdocs.yml diff --git a/alignn/__init__.py b/alignn/__init__.py index 0c35900..76db32a 100644 --- a/alignn/__init__.py +++ b/alignn/__init__.py @@ -1,3 +1,3 @@ """Version number.""" -__version__ = "2026.4.1" +__version__ = "2026.4.3" diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json index 6b0b5a0..7183476 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json @@ -15,7 +15,7 @@ "test_ratio": 0.05, "target_multiplication_factor": null, "epochs": 50, - "batch_size": 5, + "batch_size": 2, "weight_decay": 1e-05, "learning_rate": 0.001, "filename": "sample", diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json index 809133f..f5a2308 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json @@ -15,7 +15,7 @@ "test_ratio": 0.05, "target_multiplication_factor": null, "epochs": 50, - "batch_size": 5, + "batch_size": 2, "weight_decay": 1e-05, "learning_rate": 0.001, "filename": "sample", diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json index dcdafad..8330df5 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json @@ -15,7 +15,7 @@ "test_ratio": 0.05, "target_multiplication_factor": null, "epochs": 50, - "batch_size": 5, + "batch_size": 2, "weight_decay": 1e-05, "learning_rate": 0.001, "filename": "sample", diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json index 1d661fe..ab7e6c0 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json @@ -15,7 +15,7 @@ "test_ratio": 0.05, "target_multiplication_factor": null, "epochs": 50, - "batch_size": 5, + "batch_size": 2, "weight_decay": 1e-05, "learning_rate": 0.001, "filename": "sample", diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json index c0a1962..b0136f0 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json @@ -15,7 +15,7 @@ "test_ratio": 0.05, "target_multiplication_factor": null, "epochs": 50, - "batch_size": 5, + "batch_size": 2, "weight_decay": 1e-05, "learning_rate": 0.001, "filename": "sample", diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json index 804416d..0926fbe 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json @@ -15,7 +15,7 @@ "test_ratio": 0.05, "target_multiplication_factor": null, "epochs": 50, - "batch_size": 5, + "batch_size": 2, "weight_decay": 1e-05, "learning_rate": 0.001, "filename": "sample", diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json index 1529ec8..9d117a3 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json @@ -1,68 +1,69 @@ { - "version": "112bbedebdaecf59fb18e11c929080fb2f358246", - "dataset": "user_data", - "target": "target", - "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", - "id_tag": "jid", - "random_seed": 123, - "classification_threshold": null, - "n_val": 164, - "n_test": 164, - "n_train":1402, - "train_ratio": 0.9, - "val_ratio": 0.05, - "test_ratio": 0.05, - "target_multiplication_factor": null, - "epochs": 50, - "batch_size": 5, - "weight_decay": 1e-05, - "learning_rate": 0.001, - "filename": "sample", - "warmup_steps": 2000, - "criterion": "l1", - "optimizer": "adamw", - "scheduler": "onecycle", - "pin_memory": false, - "save_dataloader": false, - "write_checkpoint": true, - "write_predictions": true, - "store_outputs": false, - "progress": true, - "log_tensorboard": false, - "standard_scalar_and_pca": false, - "use_canonize": false, - "num_workers": 0, - "cutoff": 8.0, - "max_neighbors": 12, - "keep_data_order": false, - "normalize_graph_level_loss": false, - "distributed": false, - "n_early_stopping": null, - "output_dir": "out_continue", - "model": { - "name": "alignn_atomwise", - "alignn_layers": 2, - "gcn_layers": 4, - "atom_input_features": 92, - "edge_input_features": 80, - "triplet_input_features": 40, - "embedding_features": 64, - "hidden_features": 256, - "output_features": 1, - "grad_multiplier": -1, - "force_mult_natoms": true, - "calculate_gradient": true, - "atomwise_output_features": 0, - "graphwise_weight": 1.0, - "gradwise_weight": 50.0, - "stresswise_weight": 0.0, - "atomwise_weight": 0.0, - "link": "identity", - "zero_inflated": false, - "use_cutoff_function": true, - "energy_mult_natoms": false, - "classification": false, - "stress_multiplier":1 - } + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "k-nearest", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 164, + "n_test": 164, + "n_train": 1402, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": false, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "out_continue", + "model": { + "name": "alignn_atomwise_pure", + "alignn_layers": 2, + "gcn_layers": 4, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "force_mult_natoms": true, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 50.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "use_cutoff_function": true, + "energy_mult_natoms": false, + "classification": false, + "stress_multiplier": 1 + }, + "three_body_cutoff": 4.0 } diff --git a/setup.py b/setup.py index ab9bb21..e72cc31 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setuptools.setup( name="alignn", - version="2026.4.1", + version="2026.4.3", author="Kamal Choudhary, Brian DeCost", author_email="kamal.choudhary@nist.gov", description="alignn", From 0f6b77e0963151f56ef2c1111b2efc1ddc9c082e Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 22:14:08 -0400 Subject: [PATCH 05/34] Update docs --- alignn/config.py | 1 + .../examples/sample_data_ff/mlearn_data/Cu/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Ge/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Li/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Mo/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Ni/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Si/config.json | 2 +- .../sample_data_ff/mlearn_data/all/config_example.json | 2 +- alignn/train_alignn.py | 9 +++++++++ 9 files changed, 17 insertions(+), 7 deletions(-) diff --git a/alignn/config.py b/alignn/config.py index 2ea317b..3f68d85 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -180,6 +180,7 @@ class TrainingConfig(BaseSettings): target_multiplication_factor: Optional[float] = None epochs: int = 300 batch_size: int = 64 + gpu_memory_fraction: Optional[float] = None weight_decay: float = 0 learning_rate: float = 1e-2 filename: str = "sample" diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json index 7183476..390c260 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json index f5a2308..cce6696 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json index 8330df5..3c4d2f0 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json index ab7e6c0..7e30290 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json index b0136f0..3f5b1dc 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json index 0926fbe..4e32de1 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json index 9d117a3..915d2a4 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json @@ -3,7 +3,7 @@ "dataset": "user_data", "target": "target", "atom_features": "cgcnn", - "neighbor_strategy": "k-nearest", + "neighbor_strategy": "pure_torch", "id_tag": "jid", "random_seed": 123, "classification_threshold": null, diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py index 54f2721..b0d55aa 100644 --- a/alignn/train_alignn.py +++ b/alignn/train_alignn.py @@ -192,6 +192,15 @@ def train_for_folder( except Exception as exp: print("Check", exp) + if ( + config.gpu_memory_fraction is not None + and torch.cuda.is_available() + ): + torch.cuda.set_per_process_memory_fraction( + float(config.gpu_memory_fraction), + rank if world_size > 1 else 0, + ) + # config.keep_data_order = keep_data_order if classification_threshold is not None: config.classification_threshold = float(classification_threshold) From bb04c1c56101f8a38184d0961dfe8d84846fc02d Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 22:50:32 -0400 Subject: [PATCH 06/34] Config --- .../sample_data_ff/mlearn_data/Cu/config.json | 2 +- .../mlearn_data/Cu/config_dgl.json | 66 ++++++++++++++++++ .../sample_data_ff/mlearn_data/Ge/config.json | 4 +- .../mlearn_data/Ge/config_dgl.json | 66 ++++++++++++++++++ .../sample_data_ff/mlearn_data/Li/config.json | 4 +- .../mlearn_data/Li/config_dgl.json | 66 ++++++++++++++++++ .../sample_data_ff/mlearn_data/Mo/config.json | 4 +- .../mlearn_data/Mo/config_dgl.json | 66 ++++++++++++++++++ .../sample_data_ff/mlearn_data/Ni/config.json | 2 +- .../mlearn_data/Ni/config_dgl.json | 66 ++++++++++++++++++ .../sample_data_ff/mlearn_data/Si/config.json | 4 +- .../mlearn_data/Si/config_dgl.json | 67 ++++++++++++++++++ .../mlearn_data/all/config_example.json | 2 +- .../mlearn_data/all/config_example_dgl.json | 69 +++++++++++++++++++ 14 files changed, 477 insertions(+), 11 deletions(-) create mode 100644 alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json create mode 100644 alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json create mode 100644 alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json create mode 100644 alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json create mode 100644 alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json create mode 100644 alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json create mode 100644 alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json index 390c260..46a9b36 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json @@ -43,7 +43,7 @@ "model": { "name": "alignn_atomwise_pure", "alignn_layers": 2, - "gcn_layers": 4, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json new file mode 100644 index 0000000..291050a --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json @@ -0,0 +1,66 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 31, + "n_test": 31, + "n_train": 262, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 13, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 164, + "hidden_features": 256, + "output_features": 1, + "force_mult_natoms": true, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 50.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 +} diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json index cce6696..7611b4a 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json @@ -42,8 +42,8 @@ "output_dir": "./", "model": { "name": "alignn_atomwise_pure", - "alignn_layers": 4, - "gcn_layers": 4, + "alignn_layers": 2, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json new file mode 100644 index 0000000..a9ec063 --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json @@ -0,0 +1,66 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 25, + "n_test": 25, + "n_train": 228, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 10.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "force_mult_natoms": true, + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 +} diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json index 3c4d2f0..4d01755 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json @@ -42,8 +42,8 @@ "output_dir": "./", "model": { "name": "alignn_atomwise_pure", - "alignn_layers": 4, - "gcn_layers": 4, + "alignn_layers": 2, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json new file mode 100644 index 0000000..1a1120b --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json @@ -0,0 +1,66 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 29, + "n_test": 29, + "n_train": 241, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 10.0, + "stresswise_weight": 0.0, + "force_mult_natoms": true, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 +} diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json index 7e30290..fb2f7aa 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json @@ -42,8 +42,8 @@ "output_dir": "./", "model": { "name": "alignn_atomwise_pure", - "alignn_layers": 4, - "gcn_layers": 4, + "alignn_layers": 2, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json new file mode 100644 index 0000000..db5ba5e --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json @@ -0,0 +1,66 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 23, + "n_test": 23, + "n_train": 194, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 10.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "force_mult_natoms": true, + "classification": false + }, + "three_body_cutoff": 4.0 +} diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json index 3f5b1dc..c9e3d11 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json @@ -43,7 +43,7 @@ "model": { "name": "alignn_atomwise_pure", "alignn_layers": 2, - "gcn_layers": 4, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json new file mode 100644 index 0000000..e037778 --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json @@ -0,0 +1,66 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 31, + "n_test": 31, + "n_train": 263, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "force_mult_natoms": true, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 1.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "classification": false + }, + "three_body_cutoff": 4.0 +} diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json index 4e32de1..977504e 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json @@ -42,8 +42,8 @@ "output_dir": "./", "model": { "name": "alignn_atomwise_pure", - "alignn_layers": 4, - "gcn_layers": 4, + "alignn_layers": 2, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json new file mode 100644 index 0000000..257c3a9 --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json @@ -0,0 +1,67 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 25, + "n_test": 25, + "n_train": 214, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": true, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "./", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 1.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "force_mult_natoms": true, + "link": "identity", + "zero_inflated": false, + "classification": false, + "use_cutoff_function": true + }, + "three_body_cutoff": 4.0 +} diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json index 915d2a4..36f581f 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json @@ -43,7 +43,7 @@ "model": { "name": "alignn_atomwise_pure", "alignn_layers": 2, - "gcn_layers": 4, + "gcn_layers": 2, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json new file mode 100644 index 0000000..72c12fe --- /dev/null +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json @@ -0,0 +1,69 @@ +{ + "version": "112bbedebdaecf59fb18e11c929080fb2f358246", + "dataset": "user_data", + "target": "target", + "atom_features": "cgcnn", + "neighbor_strategy": "radius_graph", + "id_tag": "jid", + "random_seed": 123, + "classification_threshold": null, + "n_val": 164, + "n_test": 164, + "n_train": 1402, + "train_ratio": 0.9, + "val_ratio": 0.05, + "test_ratio": 0.05, + "target_multiplication_factor": null, + "epochs": 50, + "batch_size": 2, + "weight_decay": 1e-05, + "learning_rate": 0.001, + "filename": "sample", + "warmup_steps": 2000, + "criterion": "l1", + "optimizer": "adamw", + "scheduler": "onecycle", + "pin_memory": false, + "save_dataloader": false, + "write_checkpoint": true, + "write_predictions": true, + "store_outputs": false, + "progress": true, + "log_tensorboard": false, + "standard_scalar_and_pca": false, + "use_canonize": false, + "num_workers": 0, + "cutoff": 5.0, + "max_neighbors": 12, + "keep_data_order": false, + "normalize_graph_level_loss": false, + "distributed": false, + "n_early_stopping": null, + "output_dir": "out_continue", + "model": { + "name": "alignn_atomwise", + "alignn_layers": 2, + "gcn_layers": 2, + "atom_input_features": 92, + "edge_input_features": 80, + "triplet_input_features": 40, + "embedding_features": 64, + "hidden_features": 256, + "output_features": 1, + "grad_multiplier": -1, + "force_mult_natoms": true, + "calculate_gradient": true, + "atomwise_output_features": 0, + "graphwise_weight": 1.0, + "gradwise_weight": 50.0, + "stresswise_weight": 0.0, + "atomwise_weight": 0.0, + "link": "identity", + "zero_inflated": false, + "use_cutoff_function": true, + "energy_mult_natoms": false, + "classification": false, + "stress_multiplier": 1 + }, + "three_body_cutoff": 4.0 +} From dba428b1fc26748884cef0c12053039fe78715c0 Mon Sep 17 00:00:00 2001 From: user Date: Sat, 18 Apr 2026 23:53:14 -0400 Subject: [PATCH 07/34] Calc --- alignn/data.py | 9 ++++++++- alignn/dataset.py | 5 ++++- alignn/ff/calculators.py | 8 ++++++++ alignn/graphs.py | 11 +++++++++-- alignn/lmdb_dataset.py | 5 ++++- alignn/models/alignn.py | 11 ++++++++--- alignn/models/alignn_atomwise.py | 11 ++++++++--- alignn/models/ealignn_atomwise.py | 11 ++++++++--- 8 files changed, 57 insertions(+), 14 deletions(-) diff --git a/alignn/data.py b/alignn/data.py index 24f9ef3..68b0278 100644 --- a/alignn/data.py +++ b/alignn/data.py @@ -10,7 +10,14 @@ from tqdm import tqdm import math from jarvis.db.jsonutils import dumpjson -from dgl.dataloading import GraphDataLoader +try: + from dgl.dataloading import GraphDataLoader +except ImportError: # pure-torch path; fall back to stdlib DataLoader + from torch.utils.data import DataLoader as _TorchDataLoader + + def GraphDataLoader(*args, use_ddp=False, **kwargs): + kwargs.pop("use_ddp", None) + return _TorchDataLoader(*args, **kwargs) import pickle as pk from sklearn.preprocessing import StandardScaler diff --git a/alignn/dataset.py b/alignn/dataset.py index 0f425b6..f9bcb35 100644 --- a/alignn/dataset.py +++ b/alignn/dataset.py @@ -4,7 +4,10 @@ from typing import Optional import os import torch -import dgl +try: + import dgl +except ImportError: + dgl = None import numpy as np import pandas as pd from jarvis.core.atoms import Atoms diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py index ebc1e6c..2201f01 100644 --- a/alignn/ff/calculators.py +++ b/alignn/ff/calculators.py @@ -14,6 +14,10 @@ eALIGNNAtomWiseConfig, ) from alignn.models.alignn import ALIGNN, ALIGNNConfig +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, + ALIGNNAtomWisePureConfig, +) import matplotlib.pyplot as plt # noqa import zipfile import numpy as np @@ -252,6 +256,10 @@ def __init__( model = ALIGNNAtomWise( ALIGNNAtomWiseConfig(**self.config["model"]) ) + elif self.config["model"]["name"] == "alignn_atomwise_pure": + model = ALIGNNAtomWisePure( + ALIGNNAtomWisePureConfig(**self.config["model"]) + ) elif self.config["model"]["name"] == "alignn": model = ALIGNN(ALIGNNConfig(**self.config["model"])) elif self.config["model"]["name"] == "ealignn_atomwise": diff --git a/alignn/graphs.py b/alignn/graphs.py index 250a16e..7e861fd 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -11,9 +11,16 @@ import math from collections import defaultdict from typing import List, Tuple, Sequence, Optional -from dgl.data import DGLDataset import torch -import dgl +try: + from dgl.data import DGLDataset + import dgl +except ImportError: # DGL optional; DGL-dependent helpers will fail lazily + dgl = None + + class DGLDataset: # minimal stub so subclasses can still be defined + def __init__(self, *args, **kwargs): + pass from tqdm import tqdm from jarvis.core.atoms import Atoms diff --git a/alignn/lmdb_dataset.py b/alignn/lmdb_dataset.py index 3b540a2..3c9fcde 100644 --- a/alignn/lmdb_dataset.py +++ b/alignn/lmdb_dataset.py @@ -11,7 +11,10 @@ import torch from tqdm import tqdm from typing import List, Tuple -import dgl +try: + import dgl +except ImportError: + dgl = None def prepare_line_graph_batch( diff --git a/alignn/models/alignn.py b/alignn/models/alignn.py index 69c14ad..3cd5fb4 100644 --- a/alignn/models/alignn.py +++ b/alignn/models/alignn.py @@ -4,11 +4,16 @@ """ from typing import Tuple, Union -import dgl -import dgl.function as fn +try: + import dgl + import dgl.function as fn + from dgl.nn import AvgPooling +except ImportError: + dgl = None + fn = None + AvgPooling = None import numpy as np import torch -from dgl.nn import AvgPooling from typing import Literal from torch import nn from torch.nn import functional as F diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py index aff0243..c342be9 100644 --- a/alignn/models/alignn_atomwise.py +++ b/alignn/models/alignn_atomwise.py @@ -5,10 +5,15 @@ from typing import Tuple, Union from torch.autograd import grad -import dgl -import dgl.function as fn +try: + import dgl + import dgl.function as fn + from dgl.nn import AvgPooling +except ImportError: # DGL optional; this module only usable if DGL installed + dgl = None + fn = None + AvgPooling = None import numpy as np -from dgl.nn import AvgPooling import torch # from dgl.nn.functional import edge_softmax diff --git a/alignn/models/ealignn_atomwise.py b/alignn/models/ealignn_atomwise.py index d501090..f73fc60 100644 --- a/alignn/models/ealignn_atomwise.py +++ b/alignn/models/ealignn_atomwise.py @@ -5,11 +5,16 @@ from typing import Tuple, Union from torch.autograd import grad -import dgl -import dgl.function as fn +try: + import dgl + import dgl.function as fn + from dgl.nn import AvgPooling +except ImportError: + dgl = None + fn = None + AvgPooling = None # import numpy as np -from dgl.nn import AvgPooling import torch # from dgl.nn.functional import edge_softmax From d1af65645d9a6e00ad8d674a0662ba15271c7a94 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:00:09 -0400 Subject: [PATCH 08/34] Calc --- alignn/ff/calculators.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py index 2201f01..0c3ca1c 100644 --- a/alignn/ff/calculators.py +++ b/alignn/ff/calculators.py @@ -252,21 +252,26 @@ def __init__( ) if self.model is None: - if self.config["model"]["name"] == "alignn_atomwise": + model = None + mname = self.config["model"]["name"] + if mname == "alignn_atomwise": model = ALIGNNAtomWise( ALIGNNAtomWiseConfig(**self.config["model"]) ) - elif self.config["model"]["name"] == "alignn_atomwise_pure": + elif mname == "alignn_atomwise_pure": model = ALIGNNAtomWisePure( ALIGNNAtomWisePureConfig(**self.config["model"]) ) - elif self.config["model"]["name"] == "alignn": + elif mname == "alignn": model = ALIGNN(ALIGNNConfig(**self.config["model"])) - elif self.config["model"]["name"] == "ealignn_atomwise": + elif mname == "ealignn_atomwise": model = eALIGNNAtomWise( eALIGNNAtomWiseConfig(**self.config["model"]) ) - model.state_dict() + if model is None: + raise ValueError( + f"Unsupported model name '{mname}' in config" + ) if "atomwise" in self.config["model"]["name"]: model.load_state_dict( torch.load( From a137dfbae673ce5e3c6d1020790a51b9544898be Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:15:32 -0400 Subject: [PATCH 09/34] Calc --- alignn/models/alignn_atomwise_pure.py | 221 +++++++++++++++++++++++++- 1 file changed, 215 insertions(+), 6 deletions(-) diff --git a/alignn/models/alignn_atomwise_pure.py b/alignn/models/alignn_atomwise_pure.py index 9c1a8d9..5e319df 100644 --- a/alignn/models/alignn_atomwise_pure.py +++ b/alignn/models/alignn_atomwise_pure.py @@ -113,6 +113,12 @@ class ALIGNNAtomWisePureConfig(BaseSettings): penalty_threshold: float = 1.0 additional_output_features: int = 0 additional_output_weight: float = 0.0 + # Attention variants (Shao et al., Adv. Theory Simul. 2026): + # "alignn" - original edge-gated conv + # "n_alignn" - Node-Attention Layer (NAL), per-node learnable weights + # "t_alignn" - Self-Attention Layer (SAL), Transformer-style Q/K/V + conv_type: Literal["alignn", "n_alignn", "t_alignn"] = "alignn" + num_heads: int = 1 class Config: env_prefix = "jv_model" @@ -178,13 +184,202 @@ def forward( return self.forward_tensors(g.src, g.dst, g.num_nodes, x, y) +def _scatter_softmax( + logits: torch.Tensor, dst: torch.Tensor, num_nodes: int +) -> torch.Tensor: + """Per-destination softmax over edges (pure torch, scatter-based).""" + max_per_dst = torch.full( + (num_nodes,) + logits.shape[1:], + float("-inf"), + dtype=logits.dtype, + device=logits.device, + ) + idx = dst + while idx.dim() < logits.dim(): + idx = idx.unsqueeze(-1) + idx_e = idx.expand_as(logits) + max_per_dst.scatter_reduce_( + 0, idx_e, logits, reduce="amax", include_self=True + ) + shifted = logits - max_per_dst[dst] + exp_l = torch.exp(shifted) + sum_exp = scatter_sum(exp_l, dst, num_nodes) + return exp_l / (sum_exp[dst] + 1e-8) + + +class NodeAttentionGraphConvPure(nn.Module): + """Node-Attention Layer (NAL) — pure torch. + + Paper Eq. (8): + m_ij = A_src_i * L_src h_i + sum_j (A_dst_j * L_dst h_j + L_e e_ij) + with A_src, A_dst as per-node learned sigmoids over node features. + """ + + def __init__( + self, + input_features: int, + output_features: int, + residual: bool = True, + num_heads: int = 1, + ): + super().__init__() + self.residual = residual + self.num_heads = num_heads + self.output_features = output_features + assert output_features % num_heads == 0 + self.head_dim = output_features // num_heads + + self.src_gate = nn.Linear(input_features, output_features) + self.dst_gate = nn.Linear(input_features, output_features) + self.edge_gate = nn.Linear(input_features, output_features) + self.bn_edges = nn.LayerNorm(output_features) + self.src_update = nn.Linear(input_features, output_features) + self.dst_update = nn.Linear(input_features, output_features) + self.bn_nodes = nn.LayerNorm(output_features) + + # Per-node attention heads: sigmoid(fc(h)) -> [N, num_heads] + self.attn_src = nn.Linear(input_features, num_heads) + self.attn_dst = nn.Linear(input_features, num_heads) + + def _apply_attn( + self, proj: torch.Tensor, a: torch.Tensor + ) -> torch.Tensor: + # proj: [N, F]; a: [N, heads] + N = proj.shape[0] + ph = proj.view(N, self.num_heads, self.head_dim) + return (a.unsqueeze(-1) * ph).view(N, self.output_features) + + @torch.jit.ignore + def forward( + self, g: TorchGraph, x: torch.Tensor, y: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + src, dst, N = g.src, g.dst, g.num_nodes + + a_src = torch.sigmoid(self.attn_src(x)) + a_dst = torch.sigmoid(self.attn_dst(x)) + e_src = self._apply_attn(self.src_gate(x), a_src) + e_dst = self._apply_attn(self.dst_gate(x), a_dst) + + m = e_src[src] + e_dst[dst] + self.edge_gate(y) + sigma = torch.sigmoid(m) + + Bh = self.dst_update(x) + sum_sigma_h = scatter_sum(Bh[src] * sigma, dst, N) + sum_sigma = scatter_sum(sigma, dst, N) + h = sum_sigma_h / (sum_sigma + 1e-6) + x_new = self.src_update(x) + h + + x_new = F.silu(self.bn_nodes(x_new)) + y_new = F.silu(self.bn_edges(m)) + if self.residual: + x_new = x + x_new + y_new = y + y_new + return x_new, y_new + + +class SelfAttentionGraphConvPure(nn.Module): + """Self-Attention Layer (SAL) — pure torch. + + Paper Eqs. (9)-(11): + score = softmax((q_i . k_j . k_e) / sqrt(d)) + h_i' = fc(h_i) + SiLU(Norm(sum_j score . v_j)) + e_ij' = fc(e_ij) + SiLU(Norm(score)) + """ + + def __init__( + self, + input_features: int, + output_features: int, + residual: bool = True, + num_heads: int = 1, + ): + super().__init__() + self.residual = residual + self.num_heads = num_heads + self.output_features = output_features + assert output_features % num_heads == 0 + self.head_dim = output_features // num_heads + self.scale = self.head_dim**0.5 + + self.W_q = nn.Linear(input_features, output_features) + self.W_k = nn.Linear(input_features, output_features) + self.W_v = nn.Linear(input_features, output_features) + self.W_ke = nn.Linear(input_features, output_features) + + self.fc_node = nn.Linear(input_features, output_features) + self.fc_edge = nn.Linear(input_features, output_features) + self.bn_nodes = nn.LayerNorm(output_features) + self.bn_edges = nn.LayerNorm(output_features) + + @torch.jit.ignore + def forward( + self, g: TorchGraph, x: torch.Tensor, y: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + src, dst, N = g.src, g.dst, g.num_nodes + E = y.shape[0] + H, D = self.num_heads, self.head_dim + + q = self.W_q(x).view(N, H, D) + k = self.W_k(x).view(N, H, D) + v = self.W_v(x).view(N, H, D) + k_e = self.W_ke(y).view(E, H, D) + + # Attention logits per edge, per head: [E, H, 1] + attn_logits = (q[dst] * k[src] * k_e).sum(dim=-1, keepdim=True) / ( + self.scale + ) + # Softmax over incoming edges per destination, per head. + attn = _scatter_softmax( + attn_logits.squeeze(-1), dst, N + ).unsqueeze(-1) # [E, H, 1] + + # Weighted value aggregation per destination node. + weighted = (attn * v[src]).view(E, self.output_features) + h_agg = scatter_sum(weighted, dst, N) + x_new = self.fc_node(x) + F.silu(self.bn_nodes(h_agg)) + + # Edge update uses the attention scores themselves (paper Eq. 11). + score_e = attn.expand(E, H, D).reshape(E, self.output_features) + y_new = self.fc_edge(y) + F.silu(self.bn_edges(score_e)) + return x_new, y_new + + +def _make_bond_conv( + in_features: int, + out_features: int, + conv_type: str, + num_heads: int, +) -> nn.Module: + if conv_type == "alignn": + return EdgeGatedGraphConvPure(in_features, out_features) + if conv_type == "n_alignn": + return NodeAttentionGraphConvPure( + in_features, out_features, num_heads=num_heads + ) + if conv_type == "t_alignn": + return SelfAttentionGraphConvPure( + in_features, out_features, num_heads=num_heads + ) + raise ValueError(f"Unknown conv_type: {conv_type!r}") + + class ALIGNNConvPure(nn.Module): """Line-graph-aware ALIGNN update.""" - def __init__(self, in_features: int, out_features: int): + def __init__( + self, + in_features: int, + out_features: int, + conv_type: str = "alignn", + num_heads: int = 1, + ): super().__init__() - self.node_update = EdgeGatedGraphConvPure(in_features, out_features) - self.edge_update = EdgeGatedGraphConvPure(out_features, out_features) + self.node_update = _make_bond_conv( + in_features, out_features, conv_type, num_heads + ) + self.edge_update = _make_bond_conv( + out_features, out_features, conv_type, num_heads + ) def forward_tensors( self, @@ -215,6 +410,12 @@ def forward( y: torch.Tensor, z: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # Attention convs (NAL/SAL) don't have forward_tensors — dispatch + # through the graph-object forward, which handles all conv types. + if not isinstance(self.node_update, EdgeGatedGraphConvPure): + x, m = self.node_update(g, x, y) + y, z = self.edge_update(lg, m, z) + return x, y, z return self.forward_tensors( g.src, g.dst, @@ -299,14 +500,22 @@ def __init__( ) self.alignn_layers = nn.ModuleList( [ - ALIGNNConvPure(config.hidden_features, config.hidden_features) + 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( [ - EdgeGatedGraphConvPure( - config.hidden_features, config.hidden_features + _make_bond_conv( + config.hidden_features, + config.hidden_features, + config.conv_type, + config.num_heads, ) for _ in range(config.gcn_layers) ] From e3de40b2071483245abe3c8e9840ca69c0a2ceb1 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:20:01 -0400 Subject: [PATCH 10/34] Calc --- alignn/graphs.py | 2 ++ alignn/lmdb_dataset.py | 4 +++- alignn/models/alignn.py | 2 ++ alignn/models/alignn_atomwise.py | 2 ++ alignn/models/ealignn_atomwise.py | 2 ++ 5 files changed, 11 insertions(+), 1 deletion(-) diff --git a/alignn/graphs.py b/alignn/graphs.py index 7e861fd..9fcf85a 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -1,5 +1,7 @@ """Module to generate networkx graphs.""" +from __future__ import annotations + from jarvis.core.atoms import get_supercell_dims from jarvis.core.specie import Specie from jarvis.core.utils import random_colors diff --git a/alignn/lmdb_dataset.py b/alignn/lmdb_dataset.py index 3c9fcde..0d3611c 100644 --- a/alignn/lmdb_dataset.py +++ b/alignn/lmdb_dataset.py @@ -1,5 +1,7 @@ """Module to prepare LMDB ALIGNN dataset.""" +from __future__ import annotations + import os import numpy as np import lmdb @@ -158,7 +160,7 @@ def get_torch_dataset( if _probe is not None: _sample = pk.loads(_probe) _graph = _sample[0] - if not isinstance(_graph, dgl.DGLGraph): + if dgl is None or not isinstance(_graph, dgl.DGLGraph): raise RuntimeError( f"LMDB cache at '{tmp_name}' contains " f"{type(_graph).__name__} records, not DGLGraph. " diff --git a/alignn/models/alignn.py b/alignn/models/alignn.py index 3cd5fb4..9862ba3 100644 --- a/alignn/models/alignn.py +++ b/alignn/models/alignn.py @@ -3,6 +3,8 @@ A prototype crystal line graph network dgl implementation. """ +from __future__ import annotations + from typing import Tuple, Union try: import dgl diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py index c342be9..7c743e0 100644 --- a/alignn/models/alignn_atomwise.py +++ b/alignn/models/alignn_atomwise.py @@ -3,6 +3,8 @@ A prototype crystal line graph network dgl implementation. """ +from __future__ import annotations + from typing import Tuple, Union from torch.autograd import grad try: diff --git a/alignn/models/ealignn_atomwise.py b/alignn/models/ealignn_atomwise.py index f73fc60..94c6533 100644 --- a/alignn/models/ealignn_atomwise.py +++ b/alignn/models/ealignn_atomwise.py @@ -3,6 +3,8 @@ A prototype crystal line graph network dgl implementation. """ +from __future__ import annotations + from typing import Tuple, Union from torch.autograd import grad try: From b94e70d52c063827e7176bc02e03c08e82f84805 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:21:48 -0400 Subject: [PATCH 11/34] Calc --- alignn/graphs.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/alignn/graphs.py b/alignn/graphs.py index 9fcf85a..efe52e0 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -701,6 +701,10 @@ def atom_dgl_multigraph( use_lattice_prop=use_lattice_prop, compute_line_graph=compute_line_graph, ) + # Return TorchGraph directly when DGL is unavailable; downstream + # pure-torch datasets / models accept TorchGraph natively. + if dgl is None: + return _out if compute_line_graph: _tg, _tlg = _out return _tg.to_dgl(), _tlg.to_dgl() From 24e512bf37b175e6494c026ee62690d21d1d6602 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:32:04 -0400 Subject: [PATCH 12/34] Calc --- alignn/examples/sample_data_ff/mlearn_data/Cu/config.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Ge/config.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Li/config.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Mo/config.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Ni/config.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Si/config.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json | 2 +- .../examples/sample_data_ff/mlearn_data/all/config_example.json | 2 +- .../sample_data_ff/mlearn_data/all/config_example_dgl.json | 2 +- 14 files changed, 14 insertions(+), 14 deletions(-) diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json index 46a9b36..a717985 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 13, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json index 291050a..c9d1efe 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 13, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json index 7611b4a..17da65b 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json index a9ec063..aec8bf8 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json index 4d01755..3774a1d 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json index 1a1120b..61761ce 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json index fb2f7aa..792d064 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json index db5ba5e..bc4944c 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json index c9e3d11..0a3541e 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json index e037778..730fc75 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json index 977504e..b22c162 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json index 257c3a9..a10bb30 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": true, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json index 36f581f..0474b85 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": false, "normalize_graph_level_loss": false, "distributed": false, diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json index 72c12fe..cd28754 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json @@ -34,7 +34,7 @@ "use_canonize": false, "num_workers": 0, "cutoff": 5.0, - "max_neighbors": 12, + "max_neighbors": null, "keep_data_order": false, "normalize_graph_level_loss": false, "distributed": false, From efaff17ed3bf0977055125d496de9cacb46ab990 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:36:50 -0400 Subject: [PATCH 13/34] Calc --- alignn/config.py | 4 ++-- alignn/graphs.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/alignn/config.py b/alignn/config.py index 3f68d85..8c17a99 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -199,11 +199,11 @@ class TrainingConfig(BaseSettings): use_canonize: bool = True compute_line_graph: bool = True num_workers: int = 4 - cutoff: float = 8.0 + cutoff: float = 4.0 cutoff_extra: float = 3.0 # Separate 3-body cutoff used by neighbor_strategy="pure_torch". # When None, defaults to `cutoff`. Must be <= cutoff. - three_body_cutoff: Optional[float] = None + three_body_cutoff: Optional[float] = 3.5 max_neighbors: int = 12 keep_data_order: bool = True normalize_graph_level_loss: bool = False diff --git a/alignn/graphs.py b/alignn/graphs.py index efe52e0..feef906 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -627,7 +627,7 @@ def __init__( def atom_dgl_multigraph( atoms=None, neighbor_strategy="k-nearest", - cutoff=8.0, + cutoff=4.0, max_neighbors=12, atom_features="cgcnn", max_attempts=3, @@ -638,7 +638,7 @@ def atom_dgl_multigraph( use_lattice_prop: bool = False, cutoff_extra=3.5, dtype="float32", - three_body_cutoff: Optional[float] = None, + three_body_cutoff: Optional[float] = 3.5, ): """Obtain a DGLGraph for Atoms object.""" # print('id',id) From 2d55efc13daa4b20e156df4b186d31ca7ff0936d Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:41:44 -0400 Subject: [PATCH 14/34] Calc --- alignn/config.py | 3 ++- alignn/train.py | 31 +++++++++++++++++++++---------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/alignn/config.py b/alignn/config.py index 8c17a99..53a02a0 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -181,6 +181,7 @@ class TrainingConfig(BaseSettings): epochs: int = 300 batch_size: int = 64 gpu_memory_fraction: Optional[float] = None + use_amp: bool = False # bf16 mixed precision (A100/H100/RTX 30+) weight_decay: float = 0 learning_rate: float = 1e-2 filename: str = "sample" @@ -204,7 +205,7 @@ class TrainingConfig(BaseSettings): # 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: int = 12 + max_neighbors: Optional[int] = 12 keep_data_order: bool = True normalize_graph_level_loss: bool = False distributed: bool = False diff --git a/alignn/train.py b/alignn/train.py index 7672998..bf81981 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -263,18 +263,29 @@ def train_dgl( info = {} # info["id"] = jid optimizer.zero_grad() - if (config.compute_line_graph) > 0: - # if (config.model.alignn_layers) > 0: - result = net( - [ - dats[0].to(device), - dats[1].to(device), - dats[2].to(device), - ] + _amp_ctx = torch.autocast( + device_type="cuda", + dtype=torch.bfloat16, + enabled=bool( + getattr(config, "use_amp", False) ) + and torch.cuda.is_available(), + ) + with _amp_ctx: + if (config.compute_line_graph) > 0: + # if (config.model.alignn_layers) > 0: + result = net( + [ + dats[0].to(device), + dats[1].to(device), + dats[2].to(device), + ] + ) - else: - result = net([dats[0].to(device), dats[1].to(device)]) + else: + result = net( + [dats[0].to(device), dats[1].to(device)] + ) # info = {} info["target_out"] = [] info["pred_out"] = [] From 3339f7850e1a38a54d879246486066d3cf56f5e7 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:45:44 -0400 Subject: [PATCH 15/34] Calc --- alignn/train.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/alignn/train.py b/alignn/train.py index bf81981..f5fafa9 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -403,6 +403,16 @@ def train_dgl( optimizer.step() # optimizer.zero_grad() #never running_loss += loss.item() + # Normalize running losses by number of batches so that printed + # values are per-batch mean losses (comparable across runs / + # dataset sizes), not raw sums. + _n_tr = max(1, len(train_loader)) + running_loss /= _n_tr + running_loss1 /= _n_tr + running_loss2 /= _n_tr + running_loss3 /= _n_tr + running_loss4 /= _n_tr + running_loss5 /= _n_tr # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors( # train_result # ) @@ -558,6 +568,14 @@ def train_dgl( loss = loss1 + loss2 + loss3 + loss4 + loss5 val_result.append(info) val_loss += loss.item() + # Normalize by number of val batches (see train-loop note). + _n_vl = max(1, len(val_loader)) + val_loss /= _n_vl + val_loss1 /= _n_vl + val_loss2 /= _n_vl + val_loss3 /= _n_vl + val_loss4 /= _n_vl + val_loss5 /= _n_vl # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors( # val_result # ) From c836de8ece3db98b913ebd17142a4b8fa077d040 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 00:52:25 -0400 Subject: [PATCH 16/34] Calc --- alignn/config.py | 6 ++ alignn/train.py | 147 +++++++++++++++++++++++++++++++---------------- 2 files changed, 105 insertions(+), 48 deletions(-) diff --git a/alignn/config.py b/alignn/config.py index 53a02a0..d62f09c 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -182,6 +182,12 @@ class TrainingConfig(BaseSettings): 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" diff --git a/alignn/train.py b/alignn/train.py index f5fafa9..aba5931 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -1,6 +1,21 @@ """Module for training script.""" from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist + + +def _ddp_mean(value: float, use_ddp: bool) -> float: + """All-reduce a Python scalar across ranks and return the mean.""" + if not use_ddp or not dist.is_available() or not dist.is_initialized(): + return float(value) + t = torch.tensor(float(value), device=torch.cuda.current_device()) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + return (t / dist.get_world_size()).item() + + +def _unwrap(net): + """Return underlying module from a DDP-wrapped model, else net itself.""" + return net.module if isinstance(net, DDP) else net from functools import partial from typing import Any, Dict, Union import torch @@ -78,10 +93,12 @@ def train_dgl( # checkpoint_dir = os.path.join(config.output_dir) # deterministic = False classification = False + is_main = rank == 0 tmp = config.dict() - f = open(os.path.join(config.output_dir, "config.json"), "w") - f.write(json.dumps(tmp, indent=4)) - f.close() + if is_main: + f = open(os.path.join(config.output_dir, "config.json"), "w") + f.write(json.dumps(tmp, indent=4)) + f.close() global tmp_output_dir tmp_output_dir = config.output_dir pprint.pprint(tmp) # , sort_dicts=False) @@ -176,11 +193,14 @@ def train_dgl( xm.set_rng_state(config.random_seed) except ImportError: pass - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False os.environ["PYTHONHASHSEED"] = str(config.random_seed) - os.environ["CUBLAS_WORKSPACE_CONFIG"] = str(":4096:8") - torch.use_deterministic_algorithms(True) + if getattr(config, "deterministic", False): + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["CUBLAS_WORKSPACE_CONFIG"] = str(":4096:8") + torch.use_deterministic_algorithms(True) + else: + torch.backends.cudnn.benchmark = True if model is None: net = _model.get(config.model.name)(config.model) else: @@ -208,7 +228,13 @@ def train_dgl( # print("device", device) net.to(device) if use_ddp: - net = DDP(net, device_ids=[rank], find_unused_parameters=True) + net = DDP( + net, + device_ids=[rank], + find_unused_parameters=bool( + getattr(config, "ddp_find_unused_parameters", False) + ), + ) # group parameters to skip weight decay for bias and batchnorm params = group_decay(net) optimizer = setup_optimizer(params, config) @@ -413,6 +439,14 @@ def train_dgl( running_loss3 /= _n_tr running_loss4 /= _n_tr running_loss5 /= _n_tr + # Average across ranks so printed values reflect the whole + # global batch, not a single rank's shard. + running_loss = _ddp_mean(running_loss, use_ddp) + running_loss1 = _ddp_mean(running_loss1, use_ddp) + running_loss2 = _ddp_mean(running_loss2, use_ddp) + running_loss3 = _ddp_mean(running_loss3, use_ddp) + running_loss4 = _ddp_mean(running_loss4, use_ddp) + running_loss5 = _ddp_mean(running_loss5, use_ddp) # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors( # train_result # ) @@ -431,10 +465,13 @@ def train_dgl( running_loss5, ] ) - dumpjson( - filename=os.path.join(config.output_dir, "history_train.json"), - data=history_train, - ) + if is_main: + dumpjson( + filename=os.path.join( + config.output_dir, "history_train.json" + ), + data=history_train, + ) val_loss = 0 val_loss1 = 0 val_loss2 = 0 @@ -576,38 +613,46 @@ def train_dgl( val_loss3 /= _n_vl val_loss4 /= _n_vl val_loss5 /= _n_vl + val_loss = _ddp_mean(val_loss, use_ddp) + val_loss1 = _ddp_mean(val_loss1, use_ddp) + val_loss2 = _ddp_mean(val_loss2, use_ddp) + val_loss3 = _ddp_mean(val_loss3, use_ddp) + val_loss4 = _ddp_mean(val_loss4, use_ddp) + val_loss5 = _ddp_mean(val_loss5, use_ddp) # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors( # val_result # ) val_fin_time = time.time() val_ep_time = val_fin_time - val_init_time current_model_name = "current_model.pt" - torch.save( - net.state_dict(), - os.path.join(config.output_dir, current_model_name), - ) + if is_main: + torch.save( + _unwrap(net).state_dict(), + os.path.join(config.output_dir, current_model_name), + ) saving_msg = "" if val_loss < best_loss: best_loss = val_loss best_model_name = "best_model.pt" - torch.save( - net.state_dict(), - os.path.join(config.output_dir, best_model_name), - ) - # print("Saving data for epoch:", e) - saving_msg = "Saving model" - dumpjson( - filename=os.path.join( - config.output_dir, "Train_results.json" - ), - data=train_result, - ) - dumpjson( - filename=os.path.join( - config.output_dir, "Val_results.json" - ), - data=val_result, - ) + if is_main: + torch.save( + _unwrap(net).state_dict(), + os.path.join(config.output_dir, best_model_name), + ) + # print("Saving data for epoch:", e) + saving_msg = "Saving model" + dumpjson( + filename=os.path.join( + config.output_dir, "Train_results.json" + ), + data=train_result, + ) + dumpjson( + filename=os.path.join( + config.output_dir, "Val_results.json" + ), + data=val_result, + ) best_model = net history_val.append( [ @@ -620,10 +665,13 @@ def train_dgl( ] ) # history_val.append([mean_out, mean_atom, mean_grad, mean_stress]) - dumpjson( - filename=os.path.join(config.output_dir, "history_val.json"), - data=history_val, - ) + if is_main: + dumpjson( + filename=os.path.join( + config.output_dir, "history_val.json" + ), + data=history_val, + ) if rank == 0: print_train_val_loss( e, @@ -736,16 +784,19 @@ def train_dgl( loss = loss1 + loss2 + loss3 + loss4 if not classification: test_loss += loss.item() - print("TestLoss", e, test_loss) - dumpjson( - filename=os.path.join(config.output_dir, "Test_results.json"), - data=test_result, - ) - last_model_name = "last_model.pt" - torch.save( - net.state_dict(), - os.path.join(config.output_dir, last_model_name), - ) + if is_main: + print("TestLoss", e, test_loss) + dumpjson( + filename=os.path.join( + config.output_dir, "Test_results.json" + ), + data=test_result, + ) + last_model_name = "last_model.pt" + torch.save( + _unwrap(net).state_dict(), + os.path.join(config.output_dir, last_model_name), + ) # return test_result if rank == 0 or world_size == 1: if config.write_predictions and classification: From b59d9b767823c549496aa3d498f78fc0b0751474 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 01:16:33 -0400 Subject: [PATCH 17/34] Calc --- alignn/config.py | 2 +- alignn/examples/sample_data_ff/mlearn_data/Cu/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Cu/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Ge/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Ge/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Li/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Li/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Mo/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Mo/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Ni/config.json | 2 +- .../examples/sample_data_ff/mlearn_data/Ni/config_dgl.json | 2 +- alignn/examples/sample_data_ff/mlearn_data/Si/config.json | 6 +++--- .../examples/sample_data_ff/mlearn_data/Si/config_dgl.json | 2 +- .../sample_data_ff/mlearn_data/all/config_example.json | 2 +- .../sample_data_ff/mlearn_data/all/config_example_dgl.json | 2 +- alignn/graphs.py | 2 +- 16 files changed, 18 insertions(+), 18 deletions(-) diff --git a/alignn/config.py b/alignn/config.py index d62f09c..6904417 100644 --- a/alignn/config.py +++ b/alignn/config.py @@ -206,7 +206,7 @@ class TrainingConfig(BaseSettings): use_canonize: bool = True compute_line_graph: bool = True num_workers: int = 4 - cutoff: float = 4.0 + cutoff: float = 5.0 cutoff_extra: float = 3.0 # Separate 3-body cutoff used by neighbor_strategy="pure_torch". # When None, defaults to `cutoff`. Must be <= cutoff. diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json index a717985..8eb9da7 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json index c9d1efe..6acbfba 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Cu/config_dgl.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json index 17da65b..dbb5b1d 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json index aec8bf8..ce13c38 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ge/config_dgl.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json index 3774a1d..0019f0c 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json index 61761ce..19ebaa1 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Li/config_dgl.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json index 792d064..4742928 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config.json @@ -62,5 +62,5 @@ "force_mult_natoms": true, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json index bc4944c..eb26eb1 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Mo/config_dgl.json @@ -62,5 +62,5 @@ "force_mult_natoms": true, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json index 0a3541e..c592951 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json index 730fc75..6841663 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Ni/config_dgl.json @@ -62,5 +62,5 @@ "zero_inflated": false, "classification": false }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json index b22c162..b91de03 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config.json @@ -42,8 +42,8 @@ "output_dir": "./", "model": { "name": "alignn_atomwise_pure", - "alignn_layers": 2, - "gcn_layers": 2, + "alignn_layers": 4, + "gcn_layers": 4, "atom_input_features": 92, "edge_input_features": 80, "triplet_input_features": 40, @@ -63,5 +63,5 @@ "classification": false, "use_cutoff_function": true }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json index a10bb30..de2122c 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/Si/config_dgl.json @@ -63,5 +63,5 @@ "classification": false, "use_cutoff_function": true }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json index 0474b85..d927e8f 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example.json @@ -65,5 +65,5 @@ "classification": false, "stress_multiplier": 1 }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json index cd28754..165639e 100644 --- a/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json +++ b/alignn/examples/sample_data_ff/mlearn_data/all/config_example_dgl.json @@ -65,5 +65,5 @@ "classification": false, "stress_multiplier": 1 }, - "three_body_cutoff": 4.0 + "three_body_cutoff": 3.5 } diff --git a/alignn/graphs.py b/alignn/graphs.py index feef906..2da69ce 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -627,7 +627,7 @@ def __init__( def atom_dgl_multigraph( atoms=None, neighbor_strategy="k-nearest", - cutoff=4.0, + cutoff=5.0, max_neighbors=12, atom_features="cgcnn", max_attempts=3, From f7fc934c91989edda8adde9b9aeeca758dcca37f Mon Sep 17 00:00:00 2001 From: user Date: Sun, 19 Apr 2026 22:48:31 -0400 Subject: [PATCH 18/34] Integrators --- alignn/md/__init__.py | 13 ++ alignn/md/forces.py | 73 +++++++ alignn/md/integrators.py | 359 +++++++++++++++++++++++++++++++ alignn/md/relax.py | 101 +++++++++ docs/usage/md-integrators.md | 267 +++++++++++++++++++++++ mkdocs.yml | 1 + scripts/torch/md_demo_cu.py | 68 ++++++ scripts/torch/md_melt_quench.py | 153 +++++++++++++ scripts/torch/md_test_bussi.py | 66 ++++++ scripts/torch/md_test_fire.py | 46 ++++ scripts/torch/scaling_test_cu.py | 167 ++++++++++++++ 11 files changed, 1314 insertions(+) create mode 100644 alignn/md/__init__.py create mode 100644 alignn/md/forces.py create mode 100644 alignn/md/integrators.py create mode 100644 alignn/md/relax.py create mode 100644 docs/usage/md-integrators.md create mode 100644 scripts/torch/md_demo_cu.py create mode 100644 scripts/torch/md_melt_quench.py create mode 100644 scripts/torch/md_test_bussi.py create mode 100644 scripts/torch/md_test_fire.py create mode 100644 scripts/torch/scaling_test_cu.py diff --git a/alignn/md/__init__.py b/alignn/md/__init__.py new file mode 100644 index 0000000..dab1aeb --- /dev/null +++ b/alignn/md/__init__.py @@ -0,0 +1,13 @@ +"""Pure-PyTorch MD integrators for ALIGNN-FF.""" +from alignn.md.integrators import ( + VelocityVerlet, Langevin, NVTBerendsen, NVTBussi, NVTNoseHooverChain, + run, maxwell_boltzmann, +) +from alignn.md.forces import AlignnForces +from alignn.md.relax import FIRE +__all__ = [ + "VelocityVerlet", "Langevin", + "NVTBerendsen", "NVTBussi", "NVTNoseHooverChain", + "FIRE", + "run", "maxwell_boltzmann", "AlignnForces", +] diff --git a/alignn/md/forces.py b/alignn/md/forces.py new file mode 100644 index 0000000..38d8920 --- /dev/null +++ b/alignn/md/forces.py @@ -0,0 +1,73 @@ +"""Wrap ALIGNN-FF (pure PyTorch variant) as an on-device forces function. + +A `forces_fn` has signature: + energy, forces = forces_fn(positions) # positions: (N, 3) on GPU + # energy: scalar tensor + # forces: (N, 3) tensor + +Current implementation rebuilds the graph every step via the existing +jarvis/DGL pipeline. That's CPU-heavy and will be the bottleneck for large +systems — replace with an on-device neighbor list when you scale up. +""" +from __future__ import annotations +import numpy as np +import torch +from ase import Atoms as AseAtoms +from jarvis.core.atoms import ase_to_atoms as ase_to_jarvis + +from alignn.graphs import Graph +from alignn.torch_graph_builder import torchgraph_from_dgl + + +class AlignnForces: + def __init__( + self, + model, + atomic_numbers: np.ndarray, + cell: np.ndarray, # (3, 3), constant for NVT/NVE + cutoff: float = 8.0, + max_neighbors: int = 12, + device: str | torch.device = "cuda", + dtype: torch.dtype = torch.float32, + ): + self.model = model.eval().to(device).to(dtype) + self.Z = np.asarray(atomic_numbers) + self.cell_np = np.asarray(cell, dtype=float) + self.cell = torch.tensor(self.cell_np, dtype=dtype, device=device) + self.cutoff = cutoff + self.max_neighbors = max_neighbors + self.device = torch.device(device) + self.dtype = dtype + + def _build_graph(self, positions_np): + ase_atoms = AseAtoms( + numbers=self.Z, positions=positions_np, cell=self.cell_np, pbc=True + ) + j = ase_to_jarvis(ase_atoms) + g, lg = Graph.atom_dgl_multigraph( + j, neighbor_strategy="k-nearest", + cutoff=self.cutoff, max_neighbors=self.max_neighbors, + atom_features="cgcnn", use_canonize=True, + ) + g = g.to(self.device); lg = lg.to(self.device) + tg = torchgraph_from_dgl(g); tlg = torchgraph_from_dgl(lg) + for d in (tg.ndata, tg.edata, tlg.ndata, tlg.edata): + for k, v in list(d.items()): + if v.is_floating_point(): + d[k] = v.to(self.dtype) + return tg, tlg + + @torch.enable_grad() + def __call__(self, positions: torch.Tensor): + positions = positions.detach() + tg, tlg = self._build_graph(positions.cpu().numpy()) + # model sets r.requires_grad_(True) internally; energy grad wrt r + # gives per-edge forces, aggregated to per-atom in the model's head. + out = self.model((tg, tlg, self.cell)) + energy = out["out"].sum() if "out" in out else next(iter(out.values())).sum() + if "grad" in out: + forces = out["grad"].detach() + else: + # fallback: autograd wrt edge vectors is handled inside the model + forces = torch.zeros_like(positions) + return energy.detach(), forces.to(positions.dtype) diff --git a/alignn/md/integrators.py b/alignn/md/integrators.py new file mode 100644 index 0000000..105f59b --- /dev/null +++ b/alignn/md/integrators.py @@ -0,0 +1,359 @@ +"""Pure-PyTorch MD integrators: Velocity-Verlet (NVE) and Langevin (NVT, BAOAB). + +State kept entirely on-device. Unit system follows ASE: + energy eV, length Å, mass amu, time fs, temperature K. +With these units, the acceleration factor `F/m` needs the conversion + 1 eV/(Å·amu) = 9.6485e-3 Å/fs² (i.e., ASE's `ase.units.fs` scaling) +which we bake in as `EV_AMU_A_PER_FS2`. +""" +from __future__ import annotations +from dataclasses import dataclass +from typing import Callable, Optional +import math, time +import torch + +KB_EV = 8.617333262e-5 # Boltzmann, eV/K +EV_AMU_A_PER_FS2 = 9.6485332e-3 # eV/(Å·amu) -> Å/fs² + +ForcesFn = Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] + + +def wrap_pbc(positions: torch.Tensor, cell: torch.Tensor) -> torch.Tensor: + """Wrap positions into the primary unit cell (3x3 cell, row vectors).""" + inv = torch.linalg.inv(cell) + frac = positions @ inv + frac = frac - torch.floor(frac) + return frac @ cell + + +def kinetic_energy(masses: torch.Tensor, velocities: torch.Tensor) -> torch.Tensor: + # KE in eV: 0.5 * m * v² with m in amu, v in Å/fs -> multiply by 1/EV_AMU_A_PER_FS2 + return 0.5 * (masses * (velocities ** 2).sum(dim=-1)).sum() / EV_AMU_A_PER_FS2 + + +def instantaneous_temperature(masses, velocities): + ke = kinetic_energy(masses, velocities) + n_dof = 3 * masses.numel() - 3 # remove COM + return 2.0 * ke / (n_dof * KB_EV) + + +def maxwell_boltzmann(masses: torch.Tensor, T: float, generator=None) -> torch.Tensor: + # sigma_v = sqrt(kT/m), in Å/fs after unit conversion + sigma = torch.sqrt(KB_EV * T * EV_AMU_A_PER_FS2 / masses).unsqueeze(-1) + v = torch.randn((masses.numel(), 3), device=masses.device, + dtype=masses.dtype, generator=generator) * sigma + # remove COM drift + v -= (masses.unsqueeze(-1) * v).sum(dim=0, keepdim=True) / masses.sum() + # Rescale to exact target T (removes finite-sample noise) + T_now = instantaneous_temperature(masses, v) + v = v * torch.sqrt(torch.tensor(T, device=v.device, dtype=v.dtype) / T_now) + return v + + +# --------------------------------------------------------------------------- +# NVE: Velocity-Verlet +# --------------------------------------------------------------------------- +@dataclass +class VelocityVerlet: + forces_fn: ForcesFn + masses: torch.Tensor # (N,) amu + dt: float = 1.0 # fs + + def __post_init__(self): + self._forces = None + + def step(self, positions, velocities): + if self._forces is None: + _, self._forces = self.forces_fn(positions) + m = self.masses.unsqueeze(-1) + a = EV_AMU_A_PER_FS2 * self._forces / m + # Half-kick + velocities = velocities + 0.5 * self.dt * a + # Drift + positions = positions + self.dt * velocities + # New forces + _, new_forces = self.forces_fn(positions) + a_new = EV_AMU_A_PER_FS2 * new_forces / m + # Half-kick + velocities = velocities + 0.5 * self.dt * a_new + self._forces = new_forces + return positions, velocities + + +# --------------------------------------------------------------------------- +# NVT: Langevin via BAOAB splitting (Leimkuhler & Matthews 2013) +# --------------------------------------------------------------------------- +@dataclass +class Langevin: + forces_fn: ForcesFn + masses: torch.Tensor + dt: float = 1.0 # fs + T: float = 300.0 # K + friction: float = 0.01 # 1/fs (typical: 0.01 = 100 fs damping) + generator: Optional[torch.Generator] = None + + def __post_init__(self): + self._forces = None + self.c1 = math.exp(-self.friction * self.dt) + self.c3 = math.sqrt(1.0 - self.c1 ** 2) + + def step(self, positions, velocities): + if self._forces is None: + _, self._forces = self.forces_fn(positions) + m = self.masses.unsqueeze(-1) + # B: half-kick + a = EV_AMU_A_PER_FS2 * self._forces / m + velocities = velocities + 0.5 * self.dt * a + # A: half-drift + positions = positions + 0.5 * self.dt * velocities + # O: Ornstein-Uhlenbeck in velocity + sigma = torch.sqrt(KB_EV * self.T * EV_AMU_A_PER_FS2 / self.masses).unsqueeze(-1) + noise = torch.randn_like(velocities) if self.generator is None \ + else torch.randn(velocities.shape, device=velocities.device, + dtype=velocities.dtype, generator=self.generator) + velocities = self.c1 * velocities + self.c3 * sigma * noise + # A: half-drift + positions = positions + 0.5 * self.dt * velocities + # B: half-kick with new forces + _, new_forces = self.forces_fn(positions) + a_new = EV_AMU_A_PER_FS2 * new_forces / m + velocities = velocities + 0.5 * self.dt * a_new + self._forces = new_forces + return positions, velocities + + +# --------------------------------------------------------------------------- +# NVT: Berendsen weak-coupling thermostat +# (Not a true ensemble — fine for equilibration / melt-quench, not for +# equilibrium sampling or fluctuation-based observables.) +# --------------------------------------------------------------------------- +@dataclass +class NVTBerendsen: + forces_fn: ForcesFn + masses: torch.Tensor + dt: float = 1.0 # fs + T: float = 300.0 # K (target; mutable via set_temperature) + taut: float = 20.0 # fs (relaxation time) + + def __post_init__(self): + self._forces = None + + def set_temperature(self, T: float): + self.T = float(T) + + def step(self, positions, velocities): + if self._forces is None: + _, self._forces = self.forces_fn(positions) + m = self.masses.unsqueeze(-1) + # VV half-kick + a = EV_AMU_A_PER_FS2 * self._forces / m + velocities = velocities + 0.5 * self.dt * a + # drift + positions = positions + self.dt * velocities + # new forces + _, new_forces = self.forces_fn(positions) + a_new = EV_AMU_A_PER_FS2 * new_forces / m + velocities = velocities + 0.5 * self.dt * a_new + # Berendsen velocity rescale + T_now = instantaneous_temperature(self.masses, velocities) + # Guard against T_now == 0 early on + scale = torch.sqrt( + 1.0 + (self.dt / self.taut) * (self.T / T_now.clamp_min(1e-6) - 1.0) + ) + velocities = velocities * scale + self._forces = new_forces + return positions, velocities + + +# --------------------------------------------------------------------------- +# NVT: Bussi-Donadio-Parrinello stochastic velocity rescaling (CSVR) +# Bussi, Donadio & Parrinello, J. Chem. Phys. 126, 014101 (2007). +# Samples the true canonical distribution with Berendsen-like relaxation. +# --------------------------------------------------------------------------- +@dataclass +class NVTBussi: + forces_fn: ForcesFn + masses: torch.Tensor + dt: float = 1.0 # fs + T: float = 300.0 # K + taut: float = 20.0 # fs (same meaning as Berendsen τ_T) + remove_com: bool = True # subtract 3 DOF if true + + def __post_init__(self): + self._forces = None + + def set_temperature(self, T: float): + self.T = float(T) + + def _rescale_factor(self, K_now: torch.Tensor) -> torch.Tensor: + """Bussi 2007 Eq. A7 — returns α s.t. v ← α v samples canonical KE.""" + Nf = 3 * self.masses.numel() - (3 if self.remove_com else 0) + K_bar = 0.5 * Nf * KB_EV * self.T # target KE (eV) + # convert K_now from "integrator-internal" kinetic units: + # K_now already in eV since our kinetic_energy() returns eV. + c = math.exp(-self.dt / self.taut) + device, dtype = K_now.device, K_now.dtype + # one N(0,1) and one chi²(Nf-1) draw + R1 = torch.randn((), device=device, dtype=dtype) + # chi²(k) sampled as sum of k squared normals — cheap for small/med N; + # for huge N switch to torch.distributions.Gamma. + k = max(Nf - 1, 1) + S = (torch.randn((k,), device=device, dtype=dtype) ** 2).sum() + factor = K_bar / (Nf * K_now.clamp_min(1e-12)) + alpha_sq = ( + c + + (1.0 - c) * (S + R1 * R1) * factor + + 2.0 * R1 * torch.sqrt(c * (1.0 - c) * factor) + ) + # sign: α has the sign of (R1 + sqrt(c*Nf*K/(factor*(1-c)))) per BDP; + # in practice α² > 0 and we take positive root. + return torch.sqrt(alpha_sq.clamp_min(0.0)) + + def step(self, positions, velocities): + if self._forces is None: + _, self._forces = self.forces_fn(positions) + m = self.masses.unsqueeze(-1) + a = EV_AMU_A_PER_FS2 * self._forces / m + velocities = velocities + 0.5 * self.dt * a + positions = positions + self.dt * velocities + _, new_forces = self.forces_fn(positions) + a_new = EV_AMU_A_PER_FS2 * new_forces / m + velocities = velocities + 0.5 * self.dt * a_new + # Stochastic canonical rescale + K_now = kinetic_energy(self.masses, velocities) + alpha = self._rescale_factor(K_now) + velocities = velocities * alpha + self._forces = new_forces + return positions, velocities + + +# --------------------------------------------------------------------------- +# NVT: Nosé-Hoover chain (Martyna-Tobias-Klein integrator) +# Chains of length M ≥ 3 fix the non-ergodicity of a single Nosé-Hoover. +# Reference: Martyna, Tobias & Klein, J. Chem. Phys. 101, 4177 (1994); +# Tuckerman, "Statistical Mechanics: Theory and Molecular +# Simulation", Ch. 4. +# --------------------------------------------------------------------------- +@dataclass +class NVTNoseHooverChain: + forces_fn: ForcesFn + masses: torch.Tensor + dt: float = 1.0 # fs + T: float = 300.0 # K + taut: float = 20.0 # fs — chain relaxation time; Q1 = Nf·kT·τ² + chain_length: int = 3 # M; use ≥3 for ergodicity + remove_com: bool = True + + def __post_init__(self): + device, dtype = self.masses.device, self.masses.dtype + M = self.chain_length + self.v_xi = torch.zeros(M, device=device, dtype=dtype) + self._forces = None + self._set_masses() + + def _set_masses(self): + device, dtype = self.masses.device, self.masses.dtype + Nf = 3 * self.masses.numel() - (3 if self.remove_com else 0) + kT = KB_EV * self.T + self._Nf = Nf + self._kT = kT + tau2 = self.taut ** 2 + Q = torch.full((self.chain_length,), kT * tau2, device=device, dtype=dtype) + Q[0] = Nf * kT * tau2 + self._Q = Q + + def set_temperature(self, T: float): + self.T = float(T) + self._set_masses() # rescale Q so relaxation stays ~τ + + def _apply_chain_half(self, velocities): + """Apply half a chain step (dt_half = dt/2 of MD timestep).""" + dt_half = self.dt / 2 + dthalf = dt_half / 2 + dtqtr = dt_half / 4 + Q, v_xi = self._Q, self.v_xi + M = self.chain_length + K = kinetic_energy(self.masses, velocities) # eV + # --- reverse pass (top of chain down) --- + G_M = (Q[M-2] * v_xi[M-2] ** 2 - self._kT) / Q[M-1] if M >= 2 else \ + (2 * K - self._Nf * self._kT) / Q[0] + v_xi[M-1] = v_xi[M-1] + G_M * dthalf + for i in range(M - 2, -1, -1): + factor = torch.exp(-dtqtr * v_xi[i + 1]) + v_xi[i] = v_xi[i] * factor + if i == 0: + G_i = (2 * K - self._Nf * self._kT) / Q[0] + else: + G_i = (Q[i-1] * v_xi[i-1] ** 2 - self._kT) / Q[i] + v_xi[i] = v_xi[i] + G_i * dthalf + v_xi[i] = v_xi[i] * factor + # --- rescale particle velocities --- + scale = torch.exp(-dt_half * v_xi[0]) + velocities = velocities * scale + K = K * scale * scale + # --- forward pass (bottom up) --- + for i in range(0, M - 1): + factor = torch.exp(-dtqtr * v_xi[i + 1]) + v_xi[i] = v_xi[i] * factor + if i == 0: + G_i = (2 * K - self._Nf * self._kT) / Q[0] + else: + G_i = (Q[i-1] * v_xi[i-1] ** 2 - self._kT) / Q[i] + v_xi[i] = v_xi[i] + G_i * dthalf + v_xi[i] = v_xi[i] * factor + if M >= 2: + G_M = (Q[M-2] * v_xi[M-2] ** 2 - self._kT) / Q[M-1] + v_xi[M-1] = v_xi[M-1] + G_M * dthalf + return velocities + + def step(self, positions, velocities): + if self._forces is None: + _, self._forces = self.forces_fn(positions) + m = self.masses.unsqueeze(-1) + # Chain half-step + velocities = self._apply_chain_half(velocities) + # VV: kick-drift-kick + a = EV_AMU_A_PER_FS2 * self._forces / m + velocities = velocities + 0.5 * self.dt * a + positions = positions + self.dt * velocities + _, new_forces = self.forces_fn(positions) + a_new = EV_AMU_A_PER_FS2 * new_forces / m + velocities = velocities + 0.5 * self.dt * a_new + # Chain half-step + velocities = self._apply_chain_half(velocities) + self._forces = new_forces + return positions, velocities + + +# --------------------------------------------------------------------------- +# Driver loop +# --------------------------------------------------------------------------- +def run( + integrator, + positions: torch.Tensor, + velocities: torch.Tensor, + nsteps: int, + log_every: int = 100, + callback: Optional[Callable[[int, dict], None]] = None, +): + positions = positions.clone() + velocities = velocities.clone() + hist = [] + t0 = time.time() + for i in range(nsteps): + positions, velocities = integrator.step(positions, velocities) + if (i % log_every) == 0 or i == nsteps - 1: + ke = kinetic_energy(integrator.masses, velocities).item() + T = instantaneous_temperature(integrator.masses, velocities).item() + pe = integrator.forces_fn(positions)[0].item() \ + if hasattr(integrator, "_forces") and integrator._forces is None \ + else None + row = {"step": i, "time_fs": i * integrator.dt, "T_K": T, + "KE_eV": ke, "wall_s": time.time() - t0} + hist.append(row) + if callback is not None: + callback(i, row) + else: + print(f"step {i:6d} t={row['time_fs']:8.1f} fs " + f"T={T:7.1f} K KE={ke:10.4f} eV wall={row['wall_s']:6.1f}s") + return positions, velocities, hist diff --git a/alignn/md/relax.py b/alignn/md/relax.py new file mode 100644 index 0000000..45cecce --- /dev/null +++ b/alignn/md/relax.py @@ -0,0 +1,101 @@ +"""FIRE structural relaxation (Bitzek et al., PRL 97, 170201, 2006). + +ASE-compatible defaults. Drop-in for `ase.optimize.FIRE` when you want +to stay on-device and avoid the ASE roundtrip every step. +""" +from __future__ import annotations +from dataclasses import dataclass +from typing import Callable, Optional +import time +import torch + +from alignn.md.integrators import ForcesFn, EV_AMU_A_PER_FS2 + + +@dataclass +class FIRE: + forces_fn: ForcesFn + masses: torch.Tensor # (N,) amu + dt: float = 0.1 # fs (initial timestep) + dt_max: float = 1.0 # fs + max_step: float = 0.2 # Å (per-atom cap on Δx) + N_min: int = 5 + f_inc: float = 1.1 + f_dec: float = 0.5 + alpha_start: float = 0.1 + f_alpha: float = 0.99 + + def run( + self, + positions: torch.Tensor, + fmax: float = 0.05, # eV/Å convergence criterion + max_steps: int = 500, + log_every: int = 10, + callback: Optional[Callable[[int, dict], None]] = None, + ): + positions = positions.clone() + velocities = torch.zeros_like(positions) + m = self.masses.unsqueeze(-1) + dt = self.dt + alpha = self.alpha_start + n_pos = 0 + history = [] + t0 = time.time() + + # initial force + energy, forces = self.forces_fn(positions) + for step in range(max_steps): + fnorm_max = forces.norm(dim=-1).max().item() + row = { + "step": step, "E_eV": float(energy), + "fmax_eV_A": fnorm_max, "dt_fs": dt, "alpha": alpha, + "wall_s": time.time() - t0, + } + if step % log_every == 0 or fnorm_max < fmax: + history.append(row) + if callback is not None: + callback(step, row) + else: + print(f"FIRE step={step:4d} E={row['E_eV']:12.4f} eV " + f"|F|max={fnorm_max:9.4f} eV/Å dt={dt:5.3f} fs " + f"α={alpha:5.3f}") + if fnorm_max < fmax: + print(f"converged at step {step}: |F|max = {fnorm_max:.5f} < {fmax}") + break + + # FIRE velocity mix: v ← (1-α) v + α |v| F̂ + f_hat = forces / (forces.norm() + 1e-30) + v_norm = velocities.norm() + velocities = (1.0 - alpha) * velocities + alpha * v_norm * f_hat + + # Power P = F·v (sum over all atoms) + P = (forces * velocities).sum() + if P.item() > 0.0: + n_pos += 1 + if n_pos > self.N_min: + dt = min(dt * self.f_inc, self.dt_max) + alpha = alpha * self.f_alpha + else: + n_pos = 0 + dt = dt * self.f_dec + alpha = self.alpha_start + velocities = torch.zeros_like(velocities) + + # Velocity Verlet step (with max_step clipping) + a = EV_AMU_A_PER_FS2 * forces / m + velocities = velocities + 0.5 * dt * a + dx = dt * velocities + # Per-atom displacement cap + dx_norm = dx.norm(dim=-1, keepdim=True) + scale = torch.where( + dx_norm > self.max_step, + self.max_step / dx_norm.clamp_min(1e-30), + torch.ones_like(dx_norm), + ) + dx = dx * scale + positions = positions + dx + energy, new_forces = self.forces_fn(positions) + a_new = EV_AMU_A_PER_FS2 * new_forces / m + velocities = velocities + 0.5 * dt * a_new + forces = new_forces + return positions, history diff --git a/docs/usage/md-integrators.md b/docs/usage/md-integrators.md new file mode 100644 index 0000000..a88a50b --- /dev/null +++ b/docs/usage/md-integrators.md @@ -0,0 +1,267 @@ +# Pure-PyTorch MD & Relaxation + +`alignn.md` provides on-device molecular dynamics integrators and a FIRE +relaxer that pair with the pure-PyTorch ALIGNN-FF model +(`ALIGNNAtomWisePure`). Everything runs on GPU without bouncing through +ASE/NumPy each step, which is the prerequisite for scaling to large +systems and (eventually) multi-GPU. + +## When to use this vs. the ASE calculator + +| Use | Pick | +|---|---| +| Small systems, quick scripts, trajectory files, ensembles you trust (NPT, surfaces, constraints, free-energy tools) | ASE calculator (`AlignnAtomwiseCalculator`) | +| Large systems where per-step CPU→GPU transfer dominates | `alignn.md` integrators | +| Low-precision (bf16/fp16) or multi-GPU MD | `alignn.md` integrators | +| Differentiable MD / gradients through trajectories | `alignn.md` integrators | + +The two paths share the same model code; you can mix them (e.g., +equilibrate with ASE, production with `alignn.md`). + +## Components + +```python +from alignn.md import ( + AlignnForces, # wraps a pure-PyTorch model as forces_fn + VelocityVerlet, # NVE + Langevin, # NVT — stochastic, BAOAB splitting + NVTBerendsen, # weak-coupling (equilibration only) + NVTBussi, # canonical NVT — stochastic rescale + NVTNoseHooverChain, # canonical NVT — deterministic chain + FIRE, # structural relaxation + run, maxwell_boltzmann, +) +``` + +All integrators share a `step(positions, velocities) -> (positions, velocities)` +API. Positions are `(N, 3)` Å, velocities `(N, 3)` Å/fs, masses `(N,)` amu, +time `fs`, energy `eV`, temperature `K` — ASE's unit convention. + +## Forces function + +Every integrator takes a `forces_fn(positions) -> (energy, forces)`. +`AlignnForces` wraps the pure-PyTorch model: + +```python +import numpy as np, torch +from ase.build import bulk +from ase.data import atomic_masses +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) +from alignn.md import AlignnForces + +a = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((3, 3, 3)) +Z = a.get_atomic_numbers() +cell = np.array(a.cell) + +model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig( + name="alignn_atomwise_pure", + calculate_gradient=True, + atomwise_output_features=0, + atom_input_features=92, +)) +# load trained weights if you have them: +# model.load_state_dict(torch.load("OutputDir/best_model.pt"), strict=False) + +forces_fn = AlignnForces(model, Z, cell, device="cuda", dtype=torch.float32) +``` + +!!! warning "Neighbor-list bottleneck" + The current `AlignnForces` rebuilds the neighbor/line graph on CPU + every step. For large systems this dominates wall time. An on-device + neighbor list is planned. + +## NVE — `VelocityVerlet` + +Symplectic, time-reversible, second-order. Energy conservation is the +standard sanity check. + +```python +from alignn.md import VelocityVerlet, maxwell_boltzmann, run + +masses = torch.tensor([atomic_masses[z] for z in Z], + dtype=torch.float32, device="cuda") +pos = torch.tensor(a.get_positions(), dtype=torch.float32, device="cuda") +vel = maxwell_boltzmann(masses, T=300.0) + +integ = VelocityVerlet(forces_fn=forces_fn, masses=masses, dt=1.0) +pos, vel, hist = run(integ, pos, vel, nsteps=1000, log_every=100) +``` + +## NVT — four options + +| Thermostat | Ensemble | Dynamics preserved? | Best for | +|---|---|---|---| +| `NVTBerendsen` | **not canonical** (crushes KE fluctuations) | Yes (deterministic) | Equilibration, melt-quench | +| `Langevin` | Canonical | No (adds friction + noise) | General NVT | +| `NVTBussi` | Canonical | Nearly — velocity-direction preserved | General NVT; fast relaxation | +| `NVTNoseHooverChain` | Canonical | Yes (time-reversible) | Transport coefficients, spectra | + +### Berendsen + +Fast, stable, but **samples a distribution with wrong variance** — safe +only for equilibration. Exposes `.set_temperature(T)` for ramping. + +```python +from alignn.md import NVTBerendsen +integ = NVTBerendsen(forces_fn=forces_fn, masses=masses, + dt=1.0, T=300.0, taut=20.0) +``` + +### Langevin (BAOAB splitting) + +Solves $\ddot{x} = F/m - \gamma \dot{x} + \sqrt{2\gamma k_B T / m}\,\eta$ +with the Leimkuhler-Matthews BAOAB scheme. + +```python +from alignn.md import Langevin +integ = Langevin(forces_fn=forces_fn, masses=masses, + dt=1.0, T=300.0, friction=0.01) # 1/fs +``` + +Friction 0.01 fs⁻¹ = 100 fs damping time. Use larger γ for faster +equilibration, smaller γ to perturb dynamics less. + +### Bussi-Donadio-Parrinello + +Stochastic rescale (Bussi et al., JCP 126, 014101, 2007). Same `taut` +interface as Berendsen but samples the true canonical distribution. **Use +this as your default NVT** unless you specifically need deterministic +dynamics. + +```python +from alignn.md import NVTBussi +integ = NVTBussi(forces_fn=forces_fn, masses=masses, + dt=1.0, T=300.0, taut=20.0) +``` + +### Nosé-Hoover chain + +Martyna-Tobias-Klein integrator. Deterministic and time-reversible — +preferred when measuring transport coefficients (diffusion, viscosity, +VACF) where stochastic thermostats perturb the relevant dynamics. +Default chain length 3 (never use `chain_length=1` — it's +non-ergodic). + +```python +from alignn.md import NVTNoseHooverChain +integ = NVTNoseHooverChain(forces_fn=forces_fn, masses=masses, + dt=1.0, T=300.0, taut=20.0, + chain_length=3) +``` + +### Verifying "true NVT" + +Canonical KE should fluctuate with +$\sigma(KE)/\langle KE \rangle = \sqrt{2/(3N - 3)}$. Measured on a 108-atom +Cu system at 300 K (random-init model, isolates thermostat behavior): + +| Thermostat | ⟨KE⟩ eV | σ(KE) / σ_canonical | +|---|---|---| +| Berendsen | 4.143 ✓ | 0.04 (crushed) | +| Bussi | 4.173 ✓ | 1.13 ✓ | +| NH-chain | 4.171 ✓ | 0.52 (correlated, converges with longer runs) | + +See `md_test_bussi.py` in the repo root. + +## FIRE relaxation + +Fast Inertial Relaxation Engine (Bitzek et al., PRL 97, 170201, 2006). +ASE-compatible defaults. + +```python +from alignn.md import FIRE + +relax = FIRE(forces_fn=forces_fn, masses=masses, + dt=0.1, dt_max=1.0, max_step=0.2) +relaxed_pos, history = relax.run( + positions, fmax=0.05, max_steps=500, log_every=10, +) +``` + +**Defaults match ASE:** `dt=0.1`, `dt_max=1.0`, `max_step=0.2`, +`N_min=5`, `f_inc=1.1`, `f_dec=0.5`, `α_start=0.1`, `f_α=0.99`. +Convergence: `max(|F_i|) < fmax` in eV/Å. + +**Not yet supported:** variable-cell relaxation (no `ExpCellFilter` +equivalent), fixed-atom constraints. Use the ASE calculator path for +those. + +## Velocity initialization + +```python +from alignn.md import maxwell_boltzmann + +vel = maxwell_boltzmann(masses, T=300.0, generator=torch.Generator(device="cuda").manual_seed(0)) +# COM drift removed and velocities rescaled to hit T exactly. +``` + +## The `run` driver + +```python +def cb(step, row): + print(f"t={row['time_fs']:.1f} fs T={row['T_K']:.1f} K") + +pos, vel, history = run(integ, pos, vel, + nsteps=10000, log_every=100, callback=cb) +``` + +Each log row contains `step`, `time_fs`, `T_K`, `KE_eV`, `wall_s`. + +## Melt-quench workflow (full example) + +Drop-in replacement for the ASE `NVTBerendsen` melt-quench pattern: + +```python +from alignn.md import AlignnForces, NVTBerendsen, run, maxwell_boltzmann +from alignn.md.integrators import wrap_pbc + +forces_fn = AlignnForces(model, Z, cell_np, device="cuda") +vel = maxwell_boltzmann(masses, T=3500.0) + +integ = NVTBerendsen(forces_fn=forces_fn, masses=masses, + dt=1.0, T=3500.0, taut=20.0) + +# Stage 1 — melt +pos, vel, _ = run(integ, pos, vel, nsteps=1000, log_every=20) + +# Stage 2 — quench +integ.set_temperature(300.0) +pos, vel, _ = run(integ, pos, vel, nsteps=2000, log_every=20) + +pos = wrap_pbc(pos, cell) # tidy coordinates for POSCAR output +``` + +See `md_melt_quench.py` in the repo root for a complete script with +JARVIS loading and POSCAR output. + +## Choosing a thermostat — quick guide + +- **Equilibrating a new structure:** Berendsen or Bussi with `taut = 20-100 fs` +- **Production sampling at fixed T:** Bussi (fast, simple) or NHC (if + you need dynamics) +- **Transport coefficients (D, η, κ):** NHC (`chain_length=3-5`) with + weak coupling (`taut ≈ 100-500 fs`) +- **Melt-quench:** Berendsen for the ramp, Bussi for any equilibration at + the target T +- **Free energies, heat capacity, structural fluctuations:** Bussi or NHC + — **never** Berendsen + +## Known limitations + +1. **CPU neighbor list** rebuilt every step (see `AlignnForces._build_graph`). + Unblocks once replaced with an on-device version. +2. **No trajectory file output** — plug an `ase.io.Trajectory` writer + into the `callback` parameter of `run()` if needed. +3. **No NPT** — planned (Parrinello-Rahman / MTK barostat). +4. **No constraints** (fixed atoms, SHAKE/RATTLE, symmetry). +5. **Single GPU** — multi-GPU domain decomposition is a separate roadmap + item. + +## Related + +- [ASE Calculator](ase-calculator.md) — the canonical, full-featured + path via ASE. +- `ALIGNNAtomWisePure` in `alignn.models.alignn_atomwise_pure` — the + DGL-free model powering these integrators. diff --git a/mkdocs.yml b/mkdocs.yml index f57bf7d..5bb34c2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -82,6 +82,7 @@ nav: - ALIGNN-FF: pretrained/alignn-ff.md - Usage: - ASE Calculator: usage/ase-calculator.md + - MD & Relaxation (pure PyTorch): usage/md-integrators.md - Web Apps: usage/webapps.md - Reference: - Package Overview: api.md diff --git a/scripts/torch/md_demo_cu.py b/scripts/torch/md_demo_cu.py new file mode 100644 index 0000000..c0fd07d --- /dev/null +++ b/scripts/torch/md_demo_cu.py @@ -0,0 +1,68 @@ +"""Demo: NVE (Velocity-Verlet) and NVT (Langevin) MD on a Cu supercell +using pure-PyTorch ALIGNN-FF + pure-PyTorch integrators. + +This uses a random-init model for the integrator correctness check +(energy conservation under NVE, temperature stabilization under NVT). +Swap in trained weights for production. +""" +import numpy as np, torch +from ase.build import bulk +from ase.data import atomic_masses + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) +from alignn.md import AlignnForces, VelocityVerlet, Langevin, run +from alignn.md.integrators import maxwell_boltzmann + + +def make_model(): + cfg = ALIGNNAtomWisePureConfig( + name="alignn_atomwise_pure", + calculate_gradient=True, + atomwise_output_features=0, + atom_input_features=92, + ) + return ALIGNNAtomWisePure(cfg) + + +def setup(n=4, T0=300.0, device="cuda", dtype=torch.float32): + ase_atoms = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((n, n, n)) + N = len(ase_atoms) + Z = ase_atoms.get_atomic_numbers() + cell = np.array(ase_atoms.cell) + pos = torch.tensor(ase_atoms.get_positions(), dtype=dtype, device=device) + masses = torch.tensor( + [atomic_masses[z] for z in Z], dtype=dtype, device=device + ) + model = make_model() + forces_fn = AlignnForces(model, Z, cell, device=device, dtype=dtype) + g = torch.Generator(device=device).manual_seed(0) + vel = maxwell_boltzmann(masses, T=T0, generator=g) + print(f"N atoms = {N}") + return pos, vel, masses, forces_fn + + +def demo_nve(nsteps=200, n=3): + pos, vel, masses, fn = setup(n=n, T0=300.0) + integ = VelocityVerlet(forces_fn=fn, masses=masses, dt=1.0) + print("\n--- NVE (Velocity-Verlet) ---") + _, _, hist = run(integ, pos, vel, nsteps=nsteps, log_every=20) + # energy conservation metric + T = [r["T_K"] for r in hist] + print(f"\nT drift: min={min(T):.1f} max={max(T):.1f} range={max(T)-min(T):.1f} K") + + +def demo_nvt(nsteps=400, n=3, T=300.0): + pos, vel, masses, fn = setup(n=n, T0=T) + integ = Langevin(forces_fn=fn, masses=masses, dt=1.0, T=T, friction=0.01) + print(f"\n--- NVT Langevin (target {T} K) ---") + _, _, hist = run(integ, pos, vel, nsteps=nsteps, log_every=40) + T_tail = np.mean([r["T_K"] for r in hist[-5:]]) + print(f"\nmean T over last 5 log points: {T_tail:.1f} K (target {T} K)") + + +if __name__ == "__main__": + torch.manual_seed(0) + demo_nve(nsteps=100, n=3) + demo_nvt(nsteps=200, n=3, T=300.0) diff --git a/scripts/torch/md_melt_quench.py b/scripts/torch/md_melt_quench.py new file mode 100644 index 0000000..a01cf03 --- /dev/null +++ b/scripts/torch/md_melt_quench.py @@ -0,0 +1,153 @@ +"""Melt-quench MD using pure-PyTorch ALIGNN-FF + in-house NVTBerendsen. + +Drop-in replacement for the ASE NVTBerendsen workflow. Same two-stage +schedule: heat at T0 (melt), cool to T1 (quench), write final POSCAR. + +Usage (defaults to a small Si supercell so it runs anywhere): + python md_melt_quench.py + +For a trained model, pass --model-dir pointing to a directory with +`best_model.pt` + `config.json` that was trained with the +`alignn_atomwise` model. The state_dict is loaded into the pure-PyTorch +variant (layer names match by construction). +""" +from __future__ import annotations +import argparse, json, os, time +import numpy as np +import torch +from ase import Atoms as AseAtoms +from ase.build import bulk +from ase.data import atomic_masses +from ase.io import write as ase_write + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) +from alignn.md import AlignnForces, NVTBerendsen, run, maxwell_boltzmann +from alignn.md.integrators import wrap_pbc, instantaneous_temperature + + +def ensure_cell_size(ase_atoms: AseAtoms, min_size: float): + L = ase_atoms.get_cell().lengths() + return [max(1, int(np.ceil(min_size / Li))) for Li in L] + + +def build_model(model_dir: str | None, device, dtype): + if model_dir and os.path.exists(os.path.join(model_dir, "config.json")): + cfg_raw = json.load(open(os.path.join(model_dir, "config.json"))) + mcfg = cfg_raw["model"] + # force pure variant name so the config validates + mcfg = {**mcfg, "name": "alignn_atomwise_pure"} + model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig(**mcfg)) + sd_path = os.path.join(model_dir, "best_model.pt") + sd = torch.load(sd_path, map_location=device, weights_only=False) + if isinstance(sd, dict) and "model" in sd: + sd = sd["model"] + missing, unexpected = model.load_state_dict(sd, strict=False) + if missing or unexpected: + print(f"[warn] missing={len(missing)} unexpected={len(unexpected)}") + print(f"loaded trained weights from {sd_path}") + else: + print("[note] no model_dir given — using random-init model") + model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig( + name="alignn_atomwise_pure", + calculate_gradient=True, + atomwise_output_features=0, + atom_input_features=92, + )) + return model.to(device).to(dtype) + + +def make_starting_atoms(args) -> AseAtoms: + # Try JARVIS figshare if jid given; otherwise a built-in Si cubic cell. + if args.jid: + from jarvis.db.figshare import get_jid_data + from jarvis.core.atoms import Atoms as JAtoms + d = get_jid_data(jid=args.jid, dataset="dft_3d") + ja = JAtoms.from_dict(d["atoms"]).get_conventional_atoms + ase_atoms = ja.ase_converter() + else: + ase_atoms = bulk("Si", "diamond", a=5.43, cubic=True) + dims = ensure_cell_size(ase_atoms, args.min_size) + ase_atoms = ase_atoms.repeat(dims) + print(f"starting cell: {len(ase_atoms)} atoms, dims={dims}") + return ase_atoms + + +def write_poscar(positions_t, Z, cell_np, path): + ase_atoms = AseAtoms(numbers=Z, positions=positions_t.detach().cpu().numpy(), + cell=cell_np, pbc=True) + ase_write(path, ase_atoms, format="vasp") + print(f"wrote {path}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--jid", default=None, help="JARVIS id, e.g. JVASP-1002") + ap.add_argument("--model-dir", default=None) + ap.add_argument("--outdir", default="output_melt") + ap.add_argument("--min-size", type=float, default=10.0) + ap.add_argument("--dt", type=float, default=1.0, help="fs") + ap.add_argument("--temp0", type=float, default=3500.0) + ap.add_argument("--nsteps0", type=int, default=200) + ap.add_argument("--temp1", type=float, default=300.0) + ap.add_argument("--nsteps1", type=int, default=400) + ap.add_argument("--taut", type=float, default=20.0, help="fs") + ap.add_argument("--log-every", type=int, default=20) + args = ap.parse_args() + + os.makedirs(args.outdir, exist_ok=True) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float32 + + ase_atoms = make_starting_atoms(args) + Z = ase_atoms.get_atomic_numbers() + cell_np = np.array(ase_atoms.cell) + cell = torch.tensor(cell_np, dtype=dtype, device=device) + positions = torch.tensor(ase_atoms.get_positions(), dtype=dtype, device=device) + masses = torch.tensor([atomic_masses[z] for z in Z], dtype=dtype, device=device) + + model = build_model(args.model_dir, device, dtype) + forces_fn = AlignnForces(model, Z, cell_np, device=device, dtype=dtype) + + # Initialize velocities at T0 + gen = torch.Generator(device=device).manual_seed(0) + velocities = maxwell_boltzmann(masses, T=args.temp0, generator=gen) + + # Wrap integrator step to also PBC-wrap positions (the model's graph + # builder uses the ASE atoms object, which handles PBC internally; we + # wrap here mostly to keep output POSCARs clean). + integ = NVTBerendsen(forces_fn=forces_fn, masses=masses, dt=args.dt, + T=args.temp0, taut=args.taut) + + t_start = time.time() + jid_tag = args.jid or "Si" + + # --- Stage 1: melt at T0 --- + print(f"\n--- MELT @ {args.temp0} K for {args.nsteps0} steps ---") + def cb(i, row): print(f"[melt] step={row['step']:5d} t={row['time_fs']:7.1f} fs " + f"T={row['T_K']:7.1f} K wall={row['wall_s']:6.1f} s") + positions, velocities, _ = run(integ, positions, velocities, + nsteps=args.nsteps0, + log_every=args.log_every, callback=cb) + positions = wrap_pbc(positions, cell) + + # --- Stage 2: quench to T1 --- + integ.set_temperature(args.temp1) + print(f"\n--- QUENCH to {args.temp1} K for {args.nsteps1} steps ---") + def cb2(i, row): print(f"[quench] step={row['step']:5d} t={row['time_fs']:7.1f} fs " + f"T={row['T_K']:7.1f} K wall={row['wall_s']:6.1f} s") + positions, velocities, _ = run(integ, positions, velocities, + nsteps=args.nsteps1, + log_every=args.log_every, callback=cb2) + positions = wrap_pbc(positions, cell) + + T_final = instantaneous_temperature(masses, velocities).item() + print(f"\nDone. final T = {T_final:.1f} K total wall = {time.time()-t_start:.1f} s") + + poscar = os.path.join(args.outdir, f"POSCAR_{jid_tag}_quenched_alignn.vasp") + write_poscar(positions, Z, cell_np, poscar) + + +if __name__ == "__main__": + main() diff --git a/scripts/torch/md_test_bussi.py b/scripts/torch/md_test_bussi.py new file mode 100644 index 0000000..528c9cf --- /dev/null +++ b/scripts/torch/md_test_bussi.py @@ -0,0 +1,66 @@ +"""Compare NVTBerendsen vs NVTBussi: relaxation + KE fluctuations. + +Canonical KE variance: <ΔKE²> = (3N/2) (kT)² => σ(KE)/⟨KE⟩ = sqrt(2/3N). +Berendsen should under-fluctuate; Bussi should match. +""" +import numpy as np, torch +from ase.build import bulk +from ase.data import atomic_masses + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) +from alignn.md import ( + AlignnForces, NVTBerendsen, NVTBussi, NVTNoseHooverChain, + run, maxwell_boltzmann, +) +from alignn.md.integrators import kinetic_energy + + +def setup(n=3, T=300.0, device="cuda", dtype=torch.float32): + a = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((n, n, n)) + Z = a.get_atomic_numbers() + cell = np.array(a.cell) + pos = torch.tensor(a.get_positions(), dtype=dtype, device=device) + masses = torch.tensor([atomic_masses[z] for z in Z], dtype=dtype, device=device) + model = ALIGNNAtomWisePure(ALIGNNAtomWisePureConfig( + name="alignn_atomwise_pure", calculate_gradient=True, + atomwise_output_features=0, atom_input_features=92)) + fn = AlignnForces(model, Z, cell, device=device, dtype=dtype) + g = torch.Generator(device=device).manual_seed(0) + vel = maxwell_boltzmann(masses, T=T, generator=g) + return pos, vel, masses, fn + + +def sample_KE(integ, pos, vel, nsteps=400, skip=100): + """Run, collect KE samples after equilibration.""" + Ks = [] + for i in range(nsteps): + pos, vel = integ.step(pos, vel) + if i >= skip: + Ks.append(kinetic_energy(integ.masses, vel).item()) + return np.array(Ks) + + +def main(): + torch.manual_seed(1) + T = 300.0 + pos, vel, masses, fn = setup(n=3, T=T) # 108 atoms + N = masses.numel() + KE_mean_expected = 0.5 * (3*N - 3) * 8.617e-5 * T + sigma_expected = KE_mean_expected * np.sqrt(2.0 / (3*N - 3)) + print(f"N={N} expected ⟨KE⟩ = {KE_mean_expected:.3f} eV " + f"expected σ(KE) = {sigma_expected:.4f} eV") + + for name, integ in [ + ("Berendsen", NVTBerendsen(forces_fn=fn, masses=masses, dt=1.0, T=T, taut=20.0)), + ("Bussi", NVTBussi (forces_fn=fn, masses=masses, dt=1.0, T=T, taut=20.0)), + ("NH-chain3", NVTNoseHooverChain(forces_fn=fn, masses=masses, dt=1.0, T=T, taut=20.0, chain_length=3)), + ]: + Ks = sample_KE(integ, pos.clone(), vel.clone(), nsteps=400, skip=100) + print(f"{name:10s} ⟨KE⟩={Ks.mean():.3f} eV σ={Ks.std():.4f} eV " + f"σ/σ_expected={Ks.std()/sigma_expected:.2f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/torch/md_test_fire.py b/scripts/torch/md_test_fire.py new file mode 100644 index 0000000..805fcf5 --- /dev/null +++ b/scripts/torch/md_test_fire.py @@ -0,0 +1,46 @@ +"""Perturb a Cu FCC cell and relax with FIRE.""" +import numpy as np, torch +from ase.build import bulk +from ase.data import atomic_masses + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) +from alignn.md import AlignnForces, FIRE + + +def main(): + torch.manual_seed(0) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float32 + + a = bulk("Cu", "fcc", a=3.615, cubic=True).repeat((2, 2, 2)) # 32 atoms + Z = a.get_atomic_numbers() + cell = np.array(a.cell) + pos = torch.tensor(a.get_positions(), dtype=dtype, device=device) + # random perturbation ±0.15 Å + pos = pos + 0.15 * torch.randn_like(pos) + masses = torch.tensor([atomic_masses[z] for z in Z], dtype=dtype, device=device) + + # Toy harmonic forces_fn: V = 0.5 k (x - x0)², F = -k(x - x0). + # This validates FIRE independently of the (untrained) ALIGNN-FF model. + x0 = pos.clone().detach() + k = 2.0 # eV/Ų + def fn(positions): + dx = positions - x0 + E = 0.5 * k * (dx * dx).sum() + F = -k * dx + return E, F + # perturb AFTER defining x0 so forces are nonzero + pos = pos + 0.30 * torch.randn_like(pos) + + relax = FIRE(forces_fn=fn, masses=masses, dt=0.1, dt_max=1.0, max_step=0.2) + _, hist = relax.run(pos, fmax=0.05, max_steps=500, log_every=20) + + print(f"\ntotal logged points: {len(hist)}") + print(f"initial |F|max = {hist[0]['fmax_eV_A']:.4f} eV/Å") + print(f"final |F|max = {hist[-1]['fmax_eV_A']:.4f} eV/Å") + + +if __name__ == "__main__": + main() diff --git a/scripts/torch/scaling_test_cu.py b/scripts/torch/scaling_test_cu.py new file mode 100644 index 0000000..61aa6ba --- /dev/null +++ b/scripts/torch/scaling_test_cu.py @@ -0,0 +1,167 @@ +"""ALIGNN-FF (pure PyTorch) scaling test on a Cu FCC supercell. + +Sweeps supercell size 1x1x1 .. NxNxN across fp32 / bf16-autocast / fp16-autocast. +Records peak GPU memory and wall time for a single forward+backward (E, F). + +Usage: + python scaling_test_cu.py --max 10 --model-dir alignn/ff/alignnff_wt01 +""" + +from __future__ import annotations +import argparse, gc, json, os, time, traceback +import numpy as np +import torch +from ase.build import bulk +from jarvis.core.atoms import ase_to_atoms as ase_to_jarvis + +from alignn.graphs import Graph +from alignn.torch_graph_builder import torchgraph_from_dgl +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) + + +def build_cu_supercell(n: int): + a = bulk("Cu", "fcc", a=3.615, cubic=True) # 4 atoms / cubic cell + return a.repeat((n, n, n)) + + +def build_graphs(ase_atoms, cutoff=8.0, max_neighbors=12, device="cuda"): + j = ase_to_jarvis(ase_atoms) + g, lg = Graph.atom_dgl_multigraph( + j, neighbor_strategy="k-nearest", + cutoff=cutoff, max_neighbors=max_neighbors, + atom_features="cgcnn", use_canonize=True, + ) + g = g.to(device); lg = lg.to(device) + tg = torchgraph_from_dgl(g) + tlg = torchgraph_from_dgl(lg) + return tg, tlg + + +def make_model(device): + cfg = ALIGNNAtomWisePureConfig( + name="alignn_atomwise_pure", + calculate_gradient=True, + atomwise_output_features=0, + atom_input_features=92, + ) + m = ALIGNNAtomWisePure(cfg).to(device) + m.eval() + return m + + +def _cast_graph(tg, dtype): + for k, v in list(tg.ndata.items()): + if v.is_floating_point(): + tg.ndata[k] = v.to(dtype) + for k, v in list(tg.edata.items()): + if v.is_floating_point(): + tg.edata[k] = v.to(dtype) + + +def run_one(model, tg, tlg, cell, device, precision, model_fp32_state, fp32_snapshot): + torch.cuda.empty_cache(); gc.collect() + torch.cuda.reset_peak_memory_stats(device) + model.to(torch.float32) + model.load_state_dict(model_fp32_state) + # restore graph tensors from fp32 snapshot + ndata0, edata0, lndata0, ledata0 = fp32_snapshot + tg.ndata = {k: v.detach().clone() for k, v in ndata0.items()} + tg.edata = {k: v.detach().clone() for k, v in edata0.items()} + tlg.ndata = {k: v.detach().clone() for k, v in lndata0.items()} + tlg.edata = {k: v.detach().clone() for k, v in ledata0.items()} + cell = cell.to(torch.float32) + torch.cuda.synchronize() + t0 = time.time() + try: + if precision == "fp32": + out = model((tg, tlg, cell)) + elif precision == "bf16": + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + out = model((tg, tlg, cell)) + elif precision == "fp16": + with torch.autocast(device_type="cuda", dtype=torch.float16): + out = model((tg, tlg, cell)) + elif precision == "bf16_pure": + model.to(torch.bfloat16) + _cast_graph(tg, torch.bfloat16) + _cast_graph(tlg, torch.bfloat16) + out = model((tg, tlg, cell.to(torch.bfloat16))) + elif precision == "fp16_pure": + model.to(torch.float16) + _cast_graph(tg, torch.float16) + _cast_graph(tlg, torch.float16) + out = model((tg, tlg, cell.to(torch.float16))) + else: + raise ValueError(precision) + # Force backward to capture gradient memory + e = out["out"].sum() if "out" in out else out.get("total_energy", None) + if e is None: # fallback + e = next(iter(out.values())).sum() + e.backward() + torch.cuda.synchronize() + dt = time.time() - t0 + peak = torch.cuda.max_memory_allocated(device) / 1e9 + return {"ok": True, "time_s": dt, "peak_gb": peak} + except torch.cuda.OutOfMemoryError as ex: + return {"ok": False, "err": "OOM"} + except Exception as ex: + return {"ok": False, "err": f"{type(ex).__name__}: {ex}"} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--max", type=int, default=10) + ap.add_argument("--min", type=int, default=1) + ap.add_argument("--out", default="scaling_cu_results.json") + args = ap.parse_args() + + device = torch.device("cuda") + gpu = torch.cuda.get_device_name(0) + total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + print(f"GPU: {gpu} total={total_gb:.1f} GB") + + results = {"gpu": gpu, "total_gb": total_gb, "runs": []} + model = make_model(device) + n_params = sum(p.numel() for p in model.parameters()) + print(f"model params: {n_params/1e6:.2f} M") + import copy + model_fp32_state = copy.deepcopy(model.state_dict()) + + for n in range(args.min, args.max + 1): + atoms = build_cu_supercell(n) + N = len(atoms) + try: + tg, tlg = build_graphs(atoms, device=device) + except torch.cuda.OutOfMemoryError: + print(f"[{n}^3 N={N}] graph build OOM"); break + except Exception as ex: + print(f"[{n}^3 N={N}] graph build failed: {ex}"); break + n_edges = tg.src.numel(); n_triplets = tlg.src.numel() + cell = torch.tensor(np.array(atoms.cell), dtype=torch.float32, device=device) + print(f"\n=== {n}x{n}x{n} N={N} edges={n_edges} triplets={n_triplets} ===") + row = {"n": n, "N": N, "edges": n_edges, "triplets": n_triplets} + snap = ( + {k: v.detach().clone() for k, v in tg.ndata.items()}, + {k: v.detach().clone() for k, v in tg.edata.items()}, + {k: v.detach().clone() for k, v in tlg.ndata.items()}, + {k: v.detach().clone() for k, v in tlg.edata.items()}, + ) + for prec in ["fp32", "bf16", "fp16", "bf16_pure", "fp16_pure"]: + r = run_one(model, tg, tlg, cell, device, prec, model_fp32_state, snap) + row[prec] = r + tag = f"{r['peak_gb']:.2f} GB {r['time_s']*1000:.0f} ms" if r["ok"] else r.get("err","?") + print(f" {prec:5s}: {tag}") + results["runs"].append(row) + with open(args.out, "w") as f: + json.dump(results, f, indent=2) + # stop if all precisions OOM + if not any(row[p]["ok"] for p in ["fp32","bf16","fp16","bf16_pure","fp16_pure"]): + print("all precisions OOM — stopping"); break + del tg, tlg, cell; torch.cuda.empty_cache(); gc.collect() + + print(f"\nSaved -> {args.out}") + +if __name__ == "__main__": + main() From 48759d2faed0402a57ebe59791ddc42b75478b58 Mon Sep 17 00:00:00 2001 From: knc6 Date: Mon, 20 Apr 2026 13:26:42 -0400 Subject: [PATCH 19/34] Torch builder --- alignn/torch_graph_builder.py | 102 +++++++++++++++++++++++++++++++++- 1 file changed, 101 insertions(+), 1 deletion(-) diff --git a/alignn/torch_graph_builder.py b/alignn/torch_graph_builder.py index fbcb258..3d16f7a 100644 --- a/alignn/torch_graph_builder.py +++ b/alignn/torch_graph_builder.py @@ -32,7 +32,6 @@ import numpy as np import torch - # --------------------------------------------------------------------- # Neighbor-list primitives (moved here from alignn.graphs so importing # this module stays DGL-free). @@ -76,6 +75,107 @@ def _topk_per_source( def torch_neighbor_list( + positions: torch.Tensor, + lattice: torch.Tensor, + cutoff: float, + max_neighbors: Optional[int] = None, + atoms=None, + use_matscipy_topology: bool = False, + self_tol: float = 1e-8, + chunk_size: int = 512, +): + """Torch-native periodic neighbor list with differentiable edges. + + Memory-chunked over source atoms: peak memory is O(K * chunk * N) + instead of O(K * N^2), which also sidesteps torch's INT_MAX limit + on torch.where for very large boolean tensors. + """ + dtype = positions.dtype + device = positions.device + num_nodes = int(positions.shape[0]) + + used_matscipy = False + if use_matscipy_topology and atoms is not None: + try: + from matscipy.neighbours import neighbour_list as _mnl + + i_np, j_np, S_np = _mnl( + "ijS", atoms.ase_converter(), float(cutoff) + ) + src = torch.from_numpy(np.ascontiguousarray(i_np)).to( + device=device, dtype=torch.long + ) + dst = torch.from_numpy(np.ascontiguousarray(j_np)).to( + device=device, dtype=torch.long + ) + shift = torch.from_numpy(np.ascontiguousarray(S_np)).to( + device=device, dtype=dtype + ) + used_matscipy = True + except ImportError: + pass + + if not used_matscipy: + shifts = _torch_periodic_shifts(lattice, cutoff) # (K, 3) + with torch.no_grad(): + offs = shifts @ lattice # (K, 3) cartesian + c2 = float(cutoff) * float(cutoff) + + # Dynamically shrink chunk for very large systems so that + # (K * chunk * N) bool tensor stays well under INT_MAX. + K = int(shifts.shape[0]) + max_elems = 2**30 # ~1.07e9, safe + max_chunk_by_int = max(1, max_elems // max(K * num_nodes, 1)) + eff_chunk = max(1, min(chunk_size, max_chunk_by_int)) + + src_chunks, dst_chunks, shift_chunks = [], [], [] + for i0 in range(0, num_nodes, eff_chunk): + i1 = min(i0 + eff_chunk, num_nodes) + # (K, chunk, N, 3) + rvec = ( + positions[None, None, :, :] + + offs[:, None, None, :] + - positions[None, i0:i1, None, :] + ) + dist2 = rvec.pow(2).sum(-1) # (K, chunk, N) + mask = (dist2 <= c2) & (dist2 > self_tol) + del rvec, dist2 + k_idx, i_local, j_idx = torch.where(mask) + del mask + src_chunks.append((i_local + i0).to(torch.long)) + dst_chunks.append(j_idx.to(torch.long)) + shift_chunks.append(shifts[k_idx]) + del k_idx, i_local, j_idx + + src = ( + torch.cat(src_chunks) + if src_chunks + else torch.empty(0, dtype=torch.long, device=device) + ) + dst = ( + torch.cat(dst_chunks) + if dst_chunks + else torch.empty(0, dtype=torch.long, device=device) + ) + shift = ( + torch.cat(shift_chunks) + if shift_chunks + else torch.empty((0, 3), dtype=dtype, device=device) + ) + + # Differentiable displacement vectors — this is the autograd bridge + r = positions[dst] - positions[src] + shift @ lattice + + if max_neighbors is not None and max_neighbors > 0 and src.numel() > 0: + with torch.no_grad(): + dist = r.norm(dim=1) + keep = _topk_per_source(src, dist, int(max_neighbors), num_nodes) + src, dst, shift, r = src[keep], dst[keep], shift[keep], r[keep] + + return src, dst, shift, r + + +def torch_neighbor_list_old( positions: torch.Tensor, lattice: torch.Tensor, cutoff: float, From 2188ff17ab771a7bbf3e5c10bb639e6a84e2d632 Mon Sep 17 00:00:00 2001 From: knc6 Date: Mon, 20 Apr 2026 13:52:04 -0400 Subject: [PATCH 20/34] Torch builder --- alignn/models/alignn_atomwise_pure.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/alignn/models/alignn_atomwise_pure.py b/alignn/models/alignn_atomwise_pure.py index 5e319df..006cdd2 100644 --- a/alignn/models/alignn_atomwise_pure.py +++ b/alignn/models/alignn_atomwise_pure.py @@ -27,7 +27,6 @@ from alignn.torch_graph_builder import TorchGraph, torchgraph_from_dgl from alignn.utils import BaseSettings - # ===================================================================== # Primitives # ===================================================================== @@ -109,8 +108,10 @@ class ALIGNNAtomWisePureConfig(BaseSettings): use_penalty: bool = True extra_features: int = 0 exponent: int = 5 - penalty_factor: float = 0.1 + penalty_factor: float = 0.5 penalty_threshold: float = 1.0 + # penalty_factor: float = 0.1 + # penalty_threshold: float = 1.0 additional_output_features: int = 0 additional_output_weight: float = 0.0 # Attention variants (Shao et al., Adv. Theory Simul. 2026): @@ -241,9 +242,7 @@ def __init__( self.attn_src = nn.Linear(input_features, num_heads) self.attn_dst = nn.Linear(input_features, num_heads) - def _apply_attn( - self, proj: torch.Tensor, a: torch.Tensor - ) -> torch.Tensor: + def _apply_attn(self, proj: torch.Tensor, a: torch.Tensor) -> torch.Tensor: # proj: [N, F]; a: [N, heads] N = proj.shape[0] ph = proj.view(N, self.num_heads, self.head_dim) @@ -329,9 +328,9 @@ def forward( self.scale ) # Softmax over incoming edges per destination, per head. - attn = _scatter_softmax( - attn_logits.squeeze(-1), dst, N - ).unsqueeze(-1) # [E, H, 1] + attn = _scatter_softmax(attn_logits.squeeze(-1), dst, N).unsqueeze( + -1 + ) # [E, H, 1] # Weighted value aggregation per destination node. weighted = (attn * v[src]).view(E, self.output_features) From 7c9a57a91bbacece2502cc522d3bb92e423a8e7f Mon Sep 17 00:00:00 2001 From: user Date: Mon, 20 Apr 2026 23:45:43 -0400 Subject: [PATCH 21/34] pair_style alignn --- alignn/scripts/torch/build_lammps_alignn.sh | 131 ++++++++ alignn/scripts/torch/export_torchscript.py | 76 +++++ alignn/scripts/torch/kappa_plotly.py | 206 +++++++++++++ alignn/scripts/torch/lammps_bridge_torch.py | 197 ++++++++++++ .../scripts}/torch/md_demo_cu.py | 0 .../scripts}/torch/md_melt_quench.py | 0 .../scripts}/torch/md_test_bussi.py | 0 .../scripts}/torch/md_test_fire.py | 0 .../torch/pair_alignn/CMakeLists-snippet.txt | 27 ++ .../scripts/torch/pair_alignn/pair_alignn.cpp | 282 ++++++++++++++++++ .../scripts/torch/pair_alignn/pair_alignn.h | 55 ++++ alignn/scripts/torch/pair_alignn/si_native.in | 40 +++ alignn/scripts/torch/phonons_plotly.py | 131 ++++++++ .../scripts}/torch/scaling_test_cu.py | 0 alignn/scripts/torch/si_melt_quench.in | 49 +++ alignn/scripts/torch/thermal_plotly.py | 240 +++++++++++++++ alignn/scripts/torch/validate_pair_alignn.py | 189 ++++++++++++ mkdocs.yml | 1 + 18 files changed, 1624 insertions(+) create mode 100755 alignn/scripts/torch/build_lammps_alignn.sh create mode 100644 alignn/scripts/torch/export_torchscript.py create mode 100644 alignn/scripts/torch/kappa_plotly.py create mode 100644 alignn/scripts/torch/lammps_bridge_torch.py rename {scripts => alignn/scripts}/torch/md_demo_cu.py (100%) rename {scripts => alignn/scripts}/torch/md_melt_quench.py (100%) rename {scripts => alignn/scripts}/torch/md_test_bussi.py (100%) rename {scripts => alignn/scripts}/torch/md_test_fire.py (100%) create mode 100644 alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt create mode 100644 alignn/scripts/torch/pair_alignn/pair_alignn.cpp create mode 100644 alignn/scripts/torch/pair_alignn/pair_alignn.h create mode 100644 alignn/scripts/torch/pair_alignn/si_native.in create mode 100644 alignn/scripts/torch/phonons_plotly.py rename {scripts => alignn/scripts}/torch/scaling_test_cu.py (100%) create mode 100644 alignn/scripts/torch/si_melt_quench.in create mode 100644 alignn/scripts/torch/thermal_plotly.py create mode 100644 alignn/scripts/torch/validate_pair_alignn.py diff --git a/alignn/scripts/torch/build_lammps_alignn.sh b/alignn/scripts/torch/build_lammps_alignn.sh new file mode 100755 index 0000000..5950e99 --- /dev/null +++ b/alignn/scripts/torch/build_lammps_alignn.sh @@ -0,0 +1,131 @@ +#!/bin/bash +# Build LAMMPS with pair_alignn, linking against the conda-env libtorch. +# Self-contained: clones LAMMPS, drops in USER-ALIGNN, patches cmake, builds, +# installs the Python module into the active conda env. + +set -euo pipefail + +LAMMPS_TAG="stable_29Aug2024_update1" +LAMMPS_DIR="$HOME/lammps-alignn" +ALIGNN_REPO="/home/kamalch/Software/ollama311/alignn" +PAIR_SRC="$ALIGNN_REPO/scripts/torch/pair_alignn" + +# ── discover conda env's libtorch ───────────────────────────────────────── +TORCH_DIR=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))") +TORCH_CMAKE="$TORCH_DIR/share/cmake/Torch" +TORCH_CXX11_ABI=$(python -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))") + +echo "── environment ──────────────────────────" +echo "TORCH_DIR = $TORCH_DIR" +echo "TORCH_CXX11_ABI = $TORCH_CXX11_ABI" +echo "python = $(which python)" +echo "cmake = $(cmake --version | head -1)" +echo + +# ── clone LAMMPS ───────────────────────────────────────────────────────── +if [ ! -d "$LAMMPS_DIR" ]; then + echo "[1/5] cloning LAMMPS $LAMMPS_TAG..." + git clone --depth=1 --branch "$LAMMPS_TAG" \ + https://github.com/lammps/lammps.git "$LAMMPS_DIR" +else + echo "[1/5] LAMMPS dir exists at $LAMMPS_DIR, reusing" +fi + +# ── install pair_alignn as USER-ALIGNN package ─────────────────────────── +echo "[2/5] installing pair_alignn into USER-ALIGNN..." +PKG_DIR="$LAMMPS_DIR/src/USER-ALIGNN" +mkdir -p "$PKG_DIR" +cp "$PAIR_SRC/pair_alignn.cpp" "$PKG_DIR/" +cp "$PAIR_SRC/pair_alignn.h" "$PKG_DIR/" + +cat > "$PKG_DIR/Install.sh" <<'SHEOF' +#!/bin/sh +action() { + if test -e ../$1 && test ! -e ../$1.bak ; then + mv ../$1 ../$1.bak + fi + if test ! -e $1 ; then return ; fi + cp $1 ../$1 +} +if (test $1 = 1) ; then + action pair_alignn.cpp + action pair_alignn.h +elif (test $1 = 0) ; then + for f in pair_alignn.cpp pair_alignn.h ; do + if test -e ../$f ; then rm -f ../$f ; fi + if test -e ../${f}.bak ; then mv ../${f}.bak ../$f ; fi + done +fi +SHEOF +chmod +x "$PKG_DIR/Install.sh" + +# ── patch LAMMPS cmake to register USER-ALIGNN ─────────────────────────── +CMAKE_EXTRA="$LAMMPS_DIR/cmake/Modules/Packages/USER-ALIGNN.cmake" +mkdir -p "$(dirname "$CMAKE_EXTRA")" +cat > "$CMAKE_EXTRA" <<'CMEOF' +find_package(Torch REQUIRED) + +file(GLOB ALIGNN_SOURCES + ${LAMMPS_SOURCE_DIR}/USER-ALIGNN/pair_alignn.cpp) +target_sources(lammps PRIVATE ${ALIGNN_SOURCES}) +target_include_directories(lammps PRIVATE ${LAMMPS_SOURCE_DIR}/USER-ALIGNN) +target_link_libraries(lammps PRIVATE ${TORCH_LIBRARIES}) +target_compile_features(lammps PRIVATE cxx_std_17) +CMEOF + +# Append USER-ALIGNN to STANDARD_PACKAGES if not already there +MAIN_CMAKE="$LAMMPS_DIR/cmake/CMakeLists.txt" +if ! grep -q "USER-ALIGNN" "$MAIN_CMAKE"; then + # Find the line with STANDARD_PACKAGES and append our package. + python3 - <&1 | tail -40 + +echo +echo "[4/5] building (this takes a while)..." +cmake --build . -j"$(nproc)" 2>&1 | tail -20 + +# ── install the Python module ──────────────────────────────────────────── +echo "[5/5] installing lammps Python module..." +cmake --build . --target install-python 2>&1 | tail -10 + +# ── verify ─────────────────────────────────────────────────────────────── +echo +echo "── verification ─────────────────────────" +python - < {"energy","forces","stress"} +The C++ pair style calls this entry point each MD step. +""" +from __future__ import annotations +import argparse, json +import torch + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", required=True) + ap.add_argument("--out", default="alignn_ff.pt") + ap.add_argument("--atom-features", default="atomic_number", + help="must match what was used at training time") + ap.add_argument("--dtype", default="float32", choices=["float32", "float64"]) + args = ap.parse_args() + + dtype = torch.float32 if args.dtype == "float32" else torch.float64 + + cfg_raw = json.load(open(f"{args.model_dir}/config.json")) + mcfg = dict(cfg_raw["model"]) + mcfg["name"] = "alignn_atomwise_pure" + mcfg["calculate_gradient"] = True + if mcfg.get("stresswise_weight", 0) == 0: + mcfg["stresswise_weight"] = 0.01 + cfg = ALIGNNAtomWisePureConfig(**mcfg) + + model = ALIGNNAtomWisePure(cfg) + sd = torch.load(f"{args.model_dir}/best_model.pt", map_location="cpu", + weights_only=False) + if isinstance(sd, dict) and "model" in sd: + sd = sd["model"] + missing, unexpected = model.load_state_dict(sd, strict=False) + if missing or unexpected: + print(f"[warn] missing={len(missing)} unexpected={len(unexpected)}") + + # Baked-in CGCNN species table so C++ only ships atomic numbers (Z). + model.register_species_table(atom_features=args.atom_features) + model = model.to(dtype).eval() + + # TorchScript compile — this is the critical step. If it fails here, + # you'll need to fix the offending annotation in the model source. + scripted = torch.jit.script(model) + + # Smoke test the scripted model on a minimal 2-atom example. + with torch.no_grad(): + pos = torch.tensor([[0., 0., 0.], [1.35, 1.35, 1.35]], dtype=dtype) + pos.requires_grad_(True) + lat = torch.eye(3, dtype=dtype) * 5.43 + Z = torch.tensor([14, 14], dtype=torch.long) + src = torch.tensor([0, 1], dtype=torch.long) + dst = torch.tensor([1, 0], dtype=torch.long) + shift = torch.zeros(2, 3, dtype=dtype) + 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) + print(f"wrote {args.out} ({sum(p.numel() for p in model.parameters()):,} params)") + + +if __name__ == "__main__": + main() diff --git a/alignn/scripts/torch/kappa_plotly.py b/alignn/scripts/torch/kappa_plotly.py new file mode 100644 index 0000000..d97906b --- /dev/null +++ b/alignn/scripts/torch/kappa_plotly.py @@ -0,0 +1,206 @@ +"""Lattice thermal conductivity κ(T) and 3-phonon properties from ALIGNN-FF. + +Uses phono3py for anharmonic (3rd-order) force constants, then solves the +linearized Boltzmann transport equation under the relaxation-time +approximation (RTA). + +WARNING: this is expensive. For Si primitive + 2x2x2 supercell (16 atoms), +phono3py generates ~111 displacements for fc3 + ~1 for fc2 ~~ hundreds +of force evaluations. Allow 10-30 min on a single GPU. + +Usage: + python kappa_plotly.py --model-dir OutputDir --jid JVASP-1002 \\ + --dim 2 2 2 --mesh 11 11 11 +""" +from __future__ import annotations +import argparse +import numpy as np +import plotly.graph_objects as go +from plotly.subplots import make_subplots + +from ase import Atoms as AseAtoms +from ase.filters import ExpCellFilter +from ase.optimize import FIRE +from phono3py import Phono3py +from phonopy.structure.atoms import PhonopyAtoms + +from jarvis.core.atoms import Atoms, ase_to_atoms +from jarvis.db.figshare import get_jid_data + +from alignn.ff.ff import AlignnAtomwiseCalculator + + +def relax(atoms, calc, fmax=0.05, steps=100): + a = atoms.ase_converter() + a.calc = calc + FIRE(ExpCellFilter(a), logfile=None).run(fmax=fmax, steps=steps) + return ase_to_atoms(a) + + +def forces_on_supercell(sc, calc): + a = AseAtoms( + symbols=list(sc.symbols), + scaled_positions=np.array(sc.scaled_positions), + cell=np.array(sc.cell), + pbc=True, + ) + a.calc = calc + F = np.array(a.get_forces()) + F -= F.mean(axis=0) + return F + + +def run_phono3py(atoms, calc, dim=(2, 2, 2), dim_fc2=None, distance=0.03): + """Build Phono3py, compute fc3 (and fc2) via ALIGNN-FF.""" + bulk = PhonopyAtoms( + symbols=atoms.elements, + scaled_positions=atoms.frac_coords, + cell=atoms.lattice_mat, + ) + # fc2 supercell can be larger than fc3 supercell for a better harmonic part + if dim_fc2 is None: + dim_fc2 = dim + ph3 = Phono3py( + bulk, + supercell_matrix=[[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]], + phonon_supercell_matrix=[[dim_fc2[0], 0, 0], [0, dim_fc2[1], 0], [0, 0, dim_fc2[2]]], + ) + ph3.generate_displacements(distance=distance) + + # fc3 supercells + sc3 = ph3.supercells_with_displacements + n3 = len(sc3) + print(f"fc3 displacements: {n3}") + fset3 = [] + for i, sc in enumerate(sc3): + if sc is None: # symmetry-reduced skip + fset3.append(None); continue + F = forces_on_supercell(sc, calc) + fset3.append(F) + if (i + 1) % 10 == 0 or i == n3 - 1: + print(f" fc3 {i+1}/{n3}") + ph3.forces = fset3 + + # fc2 supercells (phonon supercells) + sc2 = ph3.phonon_supercells_with_displacements + n2 = len(sc2) + print(f"fc2 displacements: {n2}") + fset2 = [] + for i, sc in enumerate(sc2): + if sc is None: + fset2.append(None); continue + F = forces_on_supercell(sc, calc) + fset2.append(F) + ph3.phonon_forces = fset2 + + ph3.produce_fc2(symmetrize_fc2=True) + ph3.produce_fc3(symmetrize_fc3r=True) + return ph3 + + +def compute_kappa(ph3, mesh=(11, 11, 11), t_min=100, t_max=1000, t_step=50, + boundary_mfp=None): + """Run RTA thermal conductivity. Returns (T, kappa_xx_yy_zz_xy_xz_yz).""" + ph3.mesh_numbers = list(mesh) + ph3.init_phph_interaction() + temperatures = np.arange(t_min, t_max + 1, t_step) + ph3.run_thermal_conductivity( + temperatures=temperatures, + is_isotope=False, + boundary_mfp=boundary_mfp, # µm; None = infinite sample + write_kappa=False, + ) + tc = ph3.thermal_conductivity + T = np.asarray(tc.temperatures) + kappa = np.asarray(tc.kappa) # shape (1, nT, 6): xx yy zz yz xz xy + return T, kappa[0] + + +def plot_kappa(T, kappa, thermo=None, title=""): + # trace of kappa = average of diagonal components + kappa_avg = kappa[:, :3].mean(axis=1) + + fig = make_subplots( + rows=1, cols=2, + subplot_titles=("Thermal conductivity κ(T)", + "Log-log view"), + horizontal_spacing=0.12, + ) + for i, lbl in enumerate(["κ_xx", "κ_yy", "κ_zz"]): + fig.add_trace(go.Scatter(x=T, y=kappa[:, i], mode="lines+markers", + name=lbl, line=dict(width=1.5)), + row=1, col=1) + fig.add_trace(go.Scatter(x=T, y=kappa_avg, mode="lines", + name="κ_avg", line=dict(color="black", dash="dash")), + row=1, col=1) + + fig.add_trace(go.Scatter(x=T, y=kappa_avg, mode="lines+markers", + name="κ_avg (log-log)", + line=dict(color="steelblue", width=1.5), + showlegend=False), + row=1, col=2) + fig.update_xaxes(type="log", title_text="T (K)", row=1, col=2) + fig.update_yaxes(type="log", title_text="κ (W/m·K)", row=1, col=2) + + fig.update_xaxes(title_text="T (K)", row=1, col=1) + fig.update_yaxes(title_text="κ (W/m·K)", row=1, col=1) + fig.update_layout(template="plotly_white", title=title, + width=1100, height=450) + return fig + + +def print_summary(T, kappa, name=""): + kappa_avg = kappa[:, :3].mean(axis=1) + i300 = int(np.argmin(abs(T - 300))) + print(f"\n── κ summary {name} ──") + print(f"κ_xx, κ_yy, κ_zz (300 K) = {kappa[i300,0]:.2f}, " + f"{kappa[i300,1]:.2f}, {kappa[i300,2]:.2f} W/m·K") + print(f"κ_avg(300 K) = {kappa_avg[i300]:.2f} W/m·K") + # T-scaling: fit log(κ) = a - p*log(T) in high-T regime + hi = T > 300 + if hi.sum() > 2: + slope = np.polyfit(np.log(T[hi]), np.log(kappa_avg[hi]), 1)[0] + print(f"high-T exponent p (κ ~ T^-p): p ≈ {-slope:.2f} " + f"(3-phonon expects ~1)") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="OutputDir") + ap.add_argument("--jid", default="JVASP-1002") + ap.add_argument("--dim", type=int, nargs=3, default=[2, 2, 2], + help="fc3 supercell") + ap.add_argument("--dim-fc2", type=int, nargs=3, default=None, + help="fc2 supercell (default = fc3 supercell)") + ap.add_argument("--mesh", type=int, nargs=3, default=[11, 11, 11]) + ap.add_argument("--t-max", type=float, default=1000.0) + ap.add_argument("--out", default="kappa.html") + args = ap.parse_args() + + calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False) + d = get_jid_data(jid=args.jid, dataset="dft_3d") + atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms + formula = atoms.composition.reduced_formula + print(f"loaded {args.jid} {formula} N={atoms.num_atoms}") + + relaxed = relax(atoms, calc) + prim = relaxed.get_primitive_atoms + print(f"relaxed → primitive N={prim.num_atoms}") + + ph3 = run_phono3py( + prim, calc, + dim=tuple(args.dim), + dim_fc2=tuple(args.dim_fc2) if args.dim_fc2 else None, + ) + T, kappa = compute_kappa(ph3, mesh=tuple(args.mesh), t_max=args.t_max) + print_summary(T, kappa, formula) + + fig = plot_kappa(T, kappa, + title=f"{formula} ({args.jid}) — κ(T) via phono3py+ALIGNN-FF " + f"[fc3 {args.dim}, mesh {args.mesh}]") + fig.write_html(args.out) + print(f"\nwrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/alignn/scripts/torch/lammps_bridge_torch.py b/alignn/scripts/torch/lammps_bridge_torch.py new file mode 100644 index 0000000..d12ddc8 --- /dev/null +++ b/alignn/scripts/torch/lammps_bridge_torch.py @@ -0,0 +1,197 @@ +"""Drive LAMMPS with ALIGNN-FF via direct torch calls (no ASE wrapper). + +Loads `ALIGNNAtomWisePure` from OutputDir/ and registers a LAMMPS +fix-external callback that: + 1. Takes positions/cell from LAMMPS each step + 2. Builds the atom + line graph via jarvis (CPU, one-shot per step) + 3. Runs the pure-PyTorch model on GPU + 4. Returns forces (+ virial) to LAMMPS + +This bypasses AlignnAtomwiseCalculator + ASE entirely. Next stop for +speed: torch.jit.script the model (see `jit_compile()` below) or write +a C++ pair style. + +Usage: + python scripts/torch/lammps_bridge_torch.py \\ + --model-dir OutputDir --input melt_quench.in --types Si +""" +from __future__ import annotations +import argparse +import json +import numpy as np +import torch + +from ase import Atoms as AseAtoms +from jarvis.core.atoms import ase_to_atoms as ase_to_jarvis + +from alignn.graphs import Graph +from alignn.torch_graph_builder import torchgraph_from_dgl +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, ALIGNNAtomWisePureConfig, +) + +# eV/ų -> bar (LAMMPS 'metal' units) +EV_PER_A3_TO_BAR = 1.602176634e6 + + +# --------------------------------------------------------------------------- +# Model loading +# --------------------------------------------------------------------------- +def load_pure_model(model_dir: str, device, dtype=torch.float32): + cfg_raw = json.load(open(f"{model_dir}/config.json")) + mcfg = dict(cfg_raw["model"]) + mcfg["name"] = "alignn_atomwise_pure" + # make sure forces + stress autograd paths are on + mcfg["calculate_gradient"] = True + if "stresswise_weight" in mcfg and mcfg["stresswise_weight"] == 0: + mcfg["stresswise_weight"] = 0.01 + cfg = ALIGNNAtomWisePureConfig(**mcfg) + model = ALIGNNAtomWisePure(cfg) + sd = torch.load(f"{model_dir}/best_model.pt", map_location=device, + weights_only=False) + if isinstance(sd, dict) and "model" in sd: + sd = sd["model"] + missing, unexpected = model.load_state_dict(sd, strict=False) + if missing or unexpected: + print(f"[warn] state_dict missing={len(missing)} unexpected={len(unexpected)}") + model = model.to(device).to(dtype).eval() + return model, cfg_raw + + +def jit_compile(model): + """Optional: torch.jit.script the model for Path B.""" + try: + return torch.jit.script(model) + except Exception as e: + print(f"[warn] jit.script failed: {e}\n falling back to eager") + return model + + +# --------------------------------------------------------------------------- +# Graph build — currently jarvis/DGL, the CPU bottleneck. +# For large systems, replace with an on-device neighbor list. +# --------------------------------------------------------------------------- +def build_graph_on_device(symbols, positions, cell, cutoff, max_neighbors, + device, dtype): + ase_atoms = AseAtoms(symbols=symbols, positions=positions, + cell=cell, pbc=True) + j = ase_to_jarvis(ase_atoms) + g, lg = Graph.atom_dgl_multigraph( + j, neighbor_strategy="k-nearest", + cutoff=cutoff, max_neighbors=max_neighbors, + atom_features="cgcnn", use_canonize=True, + ) + g = g.to(device); lg = lg.to(device) + tg = torchgraph_from_dgl(g); tlg = torchgraph_from_dgl(lg) + for d in (tg.ndata, tg.edata, tlg.ndata, tlg.edata): + for k, v in list(d.items()): + if v.is_floating_point(): + d[k] = v.to(dtype) + return tg, tlg + + +# --------------------------------------------------------------------------- +# Callback factory +# --------------------------------------------------------------------------- +def make_callback(model, cfg, symbols_by_type, device, dtype, compute_virial=True): + cutoff = float(cfg.get("cutoff", 8.0)) + max_neighbors = int(cfg.get("max_neighbors", 12)) + step_counter = [0] + + def callback(lmp, ntimestep, nlocal, tag, x, f): + # --- unpack box from LAMMPS --- + boxlo, boxhi, xy, yz, xz, *_ = lmp.extract_box() + cell = np.array([ + [boxhi[0] - boxlo[0], 0.0, 0.0], + [xy, boxhi[1] - boxlo[1], 0.0], + [xz, yz, boxhi[2] - boxlo[2]], + ]) + # --- unpack atom types + positions --- + types = np.ctypeslib.as_array(lmp.extract_atom("type"), shape=(nlocal,)) + symbols = [symbols_by_type[t] for t in types] + positions = np.array(x, dtype=np.float64) + + # --- build graph --- + tg, tlg = build_graph_on_device(symbols, positions, cell, + cutoff, max_neighbors, device, dtype) + cell_t = torch.tensor(cell, dtype=dtype, device=device) + + # --- forward --- + with torch.enable_grad(): + out = model((tg, tlg, cell_t)) + + # energy + forces + energy = out["out"].sum().item() + forces = out["grad"].detach().cpu().numpy() + f[:] = forces + + # virial (optional) + if compute_virial and "stresses" in out: + stress_voigt = out["stresses"].detach().cpu().numpy().reshape(-1)[:6] + # stresses returned in eV/ų, Voigt order xx yy zz yz xz xy + volume = float(np.linalg.det(cell)) + virial_bar = -stress_voigt * volume * EV_PER_A3_TO_BAR + v6 = [ + float(virial_bar[0]), float(virial_bar[1]), float(virial_bar[2]), + float(virial_bar[5]), float(virial_bar[4]), float(virial_bar[3]), + ] + try: + lmp.fix_external_set_virial_global("alignn", v6) + except Exception: + pass + + step_counter[0] += 1 + if step_counter[0] % 50 == 0: + fmax = float(np.max(np.abs(forces))) + print(f"[torch-bridge] step={step_counter[0]:6d} " + f"E={energy: .4f} eV |F|max={fmax:.3f} eV/Å") + + return callback + + +# --------------------------------------------------------------------------- +# Runner +# --------------------------------------------------------------------------- +def run(args): + from lammps import lammps + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float32 + model, cfg_raw = load_pure_model(args.model_dir, device, dtype) + if args.jit: + model = jit_compile(model) + + symbols_by_type = {i + 1: s for i, s in enumerate(args.types.split(","))} + + lmp = lammps() + if args.input: + with open(args.input) as fh: + lines = fh.readlines() + cb = make_callback(model, cfg_raw, symbols_by_type, device, dtype, + compute_virial=not args.no_virial) + registered = False + for line in lines: + stripped = line.strip() + if (not registered and stripped.startswith("run ") + and "alignn" in "".join(lines[: lines.index(line)])): + lmp.set_fix_external_callback("alignn", cb, lmp) + registered = True + lmp.command(line.rstrip("\n")) + else: + raise SystemExit("Provide --input with a LAMMPS script") + lmp.close() + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model-dir", required=True, help="dir with best_model.pt + config.json") + p.add_argument("--input", required=True, help="LAMMPS input script") + p.add_argument("--types", required=True, help="comma-separated element symbols by LAMMPS type") + p.add_argument("--jit", action="store_true", help="try torch.jit.script the model") + p.add_argument("--no-virial", action="store_true") + args = p.parse_args() + run(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/torch/md_demo_cu.py b/alignn/scripts/torch/md_demo_cu.py similarity index 100% rename from scripts/torch/md_demo_cu.py rename to alignn/scripts/torch/md_demo_cu.py diff --git a/scripts/torch/md_melt_quench.py b/alignn/scripts/torch/md_melt_quench.py similarity index 100% rename from scripts/torch/md_melt_quench.py rename to alignn/scripts/torch/md_melt_quench.py diff --git a/scripts/torch/md_test_bussi.py b/alignn/scripts/torch/md_test_bussi.py similarity index 100% rename from scripts/torch/md_test_bussi.py rename to alignn/scripts/torch/md_test_bussi.py diff --git a/scripts/torch/md_test_fire.py b/alignn/scripts/torch/md_test_fire.py similarity index 100% rename from scripts/torch/md_test_fire.py rename to alignn/scripts/torch/md_test_fire.py diff --git a/alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt b/alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt new file mode 100644 index 0000000..75f3348 --- /dev/null +++ b/alignn/scripts/torch/pair_alignn/CMakeLists-snippet.txt @@ -0,0 +1,27 @@ +# Add to LAMMPS cmake/CMakeLists.txt to enable pair_alignn. +# Or keep as a standalone external package. + +option(PKG_USER-ALIGNN "Build USER-ALIGNN with libtorch" OFF) + +if(PKG_USER-ALIGNN) + find_package(Torch REQUIRED) + + file(GLOB ALIGNN_SOURCES + ${LAMMPS_SOURCE_DIR}/USER-ALIGNN/pair_alignn.cpp) + target_sources(lammps PRIVATE ${ALIGNN_SOURCES}) + target_include_directories(lammps PRIVATE ${LAMMPS_SOURCE_DIR}/USER-ALIGNN) + target_link_libraries(lammps PRIVATE ${TORCH_LIBRARIES}) + target_compile_features(lammps PRIVATE cxx_std_17) + + # libtorch ABI flag — must match how libtorch was built. + if(DEFINED ENV{TORCH_CXX11_ABI}) + target_compile_definitions(lammps PRIVATE _GLIBCXX_USE_CXX11_ABI=$ENV{TORCH_CXX11_ABI}) + endif() +endif() + +# Build: +# cd build +# cmake ../cmake -DPKG_USER-ALIGNN=ON \ +# -DCMAKE_PREFIX_PATH=/path/to/libtorch \ +# -DCMAKE_BUILD_TYPE=Release +# make -j diff --git a/alignn/scripts/torch/pair_alignn/pair_alignn.cpp b/alignn/scripts/torch/pair_alignn/pair_alignn.cpp new file mode 100644 index 0000000..92ac76c --- /dev/null +++ b/alignn/scripts/torch/pair_alignn/pair_alignn.cpp @@ -0,0 +1,282 @@ +// Native LAMMPS pair style for ALIGNN-FF (via libtorch). +// Single-rank MVP: uses only LOCAL atoms + shift-based edges. The model's +// energy readout is a global pool over all graph nodes, so ghost atoms +// must NOT be passed as graph nodes — they'd inflate the sum. + +#include "pair_alignn.h" + +#include "atom.h" +#include "comm.h" +#include "domain.h" +#include "error.h" +#include "force.h" +#include "memory.h" +#include "neigh_list.h" +#include "neigh_request.h" +#include "neighbor.h" + +#include +#include +#include +#include + +using namespace LAMMPS_NS; + +static int element_to_Z(const std::string &sym) { + static const std::vector periodic = { + "H","He","Li","Be","B","C","N","O","F","Ne","Na","Mg","Al","Si","P", + "S","Cl","Ar","K","Ca","Sc","Ti","V","Cr","Mn","Fe","Co","Ni","Cu", + "Zn","Ga","Ge","As","Se","Br","Kr","Rb","Sr","Y","Zr","Nb","Mo","Tc", + "Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I","Xe","Cs","Ba", + }; + for (size_t i = 0; i < periodic.size(); ++i) + if (periodic[i] == sym) return static_cast(i + 1); + return -1; +} + +PairAlignn::PairAlignn(LAMMPS *lmp) : Pair(lmp) { + single_enable = 0; + restartinfo = 0; + one_coeff = 1; + manybody_flag = 1; + no_virial_fdotr_compute = 1; +} + +PairAlignn::~PairAlignn() { + if (allocated) { + memory->destroy(setflag); + memory->destroy(cutsq); + } +} + +void PairAlignn::allocate() { + allocated = 1; + int n = atom->ntypes; + memory->create(setflag, n + 1, n + 1, "pair:setflag"); + memory->create(cutsq, n + 1, n + 1, "pair:cutsq"); + for (int i = 1; i <= n; i++) + for (int j = i; j <= n; j++) setflag[i][j] = 0; +} + +void PairAlignn::settings(int narg, char **arg) { + if (narg < 1 || narg > 2) + error->all(FLERR, "pair_style alignn: cutoff (Å) [max_neighbors]"); + cutoff_ = std::stod(arg[0]); + if (cutoff_ <= 0.0) + error->all(FLERR, "pair_style alignn: cutoff must be positive"); + if (narg == 2) max_neighbors_ = std::stoi(arg[1]); +} + +void PairAlignn::coeff(int narg, char **arg) { + if (!allocated) allocate(); + if (narg != 3 + atom->ntypes) + error->all(FLERR, + "pair_coeff * * model.pt [ ...]"); + if (strcmp(arg[0], "*") != 0 || strcmp(arg[1], "*") != 0) + error->all(FLERR, "pair_coeff alignn requires * *"); + + const std::string model_path = arg[2]; + type_to_Z_.assign(atom->ntypes + 1, 0); + for (int t = 1; t <= atom->ntypes; ++t) { + std::string sym = arg[2 + t]; + int Z = element_to_Z(sym); + if (Z < 0) error->all(FLERR, + (std::string("Unknown element: ") + sym).c_str()); + type_to_Z_[t] = Z; + } + + device_ = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU; + try { + model_ = torch::jit::load(model_path, device_); + model_.eval(); + } catch (const c10::Error &e) { + error->all(FLERR, + (std::string("Failed to load TorchScript model: ") + e.what()).c_str()); + } + model_loaded_ = true; + for (int i = 1; i <= atom->ntypes; i++) + for (int j = i; j <= atom->ntypes; j++) setflag[i][j] = 1; +} + +void PairAlignn::init_style() { + if (!model_loaded_) + error->all(FLERR, "pair_style alignn requires pair_coeff * * ..."); + if (atom->tag_enable == 0) + error->all(FLERR, "pair_style alignn requires atom IDs"); + if (force->newton_pair == 0) + error->all(FLERR, "pair_style alignn requires newton on"); + if (atom->map_style == Atom::MAP_NONE) atom->map_init(); + neighbor->add_request(this, NeighConst::REQ_FULL | NeighConst::REQ_GHOST); +} + +double PairAlignn::init_one(int /*i*/, int /*j*/) { return cutoff_; } + +// --------------------------------------------------------------------------- +void PairAlignn::compute(int eflag, int vflag) { + ev_init(eflag, vflag); + + const int nlocal = atom->nlocal; + double **x = atom->x; + double **f = atom->f; + int *type = atom->type; + tagint *tag = atom->tag; + + int *ilist = list->ilist; + int *numneigh = list->numneigh; + int **firstneigh = list->firstneigh; + const int inum = list->inum; + + const double hx = domain->xprd; + const double hy = domain->yprd; + const double hz = domain->zprd; + const double xy = domain->xy; + const double xz = domain->xz; + const double yz = domain->yz; + const bool triclinic = (xy != 0.0) || (xz != 0.0) || (yz != 0.0); + + // --- positions / atomic numbers: LOCAL atoms only --- + auto t_opts = torch::TensorOptions().dtype(dtype_).device(torch::kCPU); + auto l_opts = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU); + torch::Tensor pos = torch::empty({nlocal, 3}, t_opts); + torch::Tensor Z = torch::empty({nlocal}, l_opts); + { + auto pa = pos.accessor(); + auto Za = Z.accessor(); + for (int i = 0; i < nlocal; ++i) { + pa[i][0] = static_cast(x[i][0]); + pa[i][1] = static_cast(x[i][1]); + pa[i][2] = static_cast(x[i][2]); + Za[i] = static_cast(type_to_Z_[type[i]]); + } + } + + // --- edges: src/dst both LOCAL; shift carries PBC image offset --- + std::vector src_vec, dst_vec; + std::vector shift_flat; + src_vec.reserve(nlocal * 16); + dst_vec.reserve(nlocal * 16); + shift_flat.reserve(nlocal * 16 * 3); + const double cut2 = cutoff_ * cutoff_; + + // k-nearest (same convention as jarvis `neighbor_strategy="k-nearest"`): + // for each local i, keep up to max_neighbors_ closest j's within cutoff. + struct Cand { double r2; int dst_local; float shx, shy, shz; }; + std::vector buf; + buf.reserve(256); + + for (int ii = 0; ii < inum; ++ii) { + const int i = ilist[ii]; + if (i >= nlocal) continue; + const double xi = x[i][0], yi = x[i][1], zi = x[i][2]; + const int jnum = numneigh[i]; + const int *jlist = firstneigh[i]; + buf.clear(); + for (int jj = 0; jj < jnum; ++jj) { + int j = jlist[jj]; + j &= NEIGHMASK; + const double dxr = x[j][0] - xi; + const double dyr = x[j][1] - yi; + const double dzr = x[j][2] - zi; + const double r2 = dxr*dxr + dyr*dyr + dzr*dzr; + if (r2 >= cut2) continue; + + int dst_local; + double shx = 0.0, shy = 0.0, shz = 0.0; + if (j < nlocal) { + dst_local = j; + } else { + tagint jtag = tag[j]; + int owner = atom->map(jtag); + if (owner < 0 || owner >= nlocal) continue; + dst_local = owner; + const double dox = x[j][0] - x[owner][0]; + const double doy = x[j][1] - x[owner][1]; + const double doz = x[j][2] - x[owner][2]; + (void)triclinic; + shx = std::round(dox / hx); + shy = std::round(doy / hy); + shz = std::round(doz / hz); + } + buf.push_back({r2, dst_local, + static_cast(shx), + static_cast(shy), + static_cast(shz)}); + } + // Sort by r2 ascending, keep first max_neighbors_. + if (static_cast(buf.size()) > max_neighbors_) { + std::partial_sort(buf.begin(), + buf.begin() + max_neighbors_, + buf.end(), + [](const Cand &a, const Cand &b) { return a.r2 < b.r2; }); + buf.resize(max_neighbors_); + } + for (const auto &c : buf) { + src_vec.push_back(i); + dst_vec.push_back(c.dst_local); + shift_flat.push_back(c.shx); + shift_flat.push_back(c.shy); + shift_flat.push_back(c.shz); + } + } + const long nedges = static_cast(src_vec.size()); + torch::Tensor src = torch::from_blob(src_vec.data(), {nedges}, torch::kLong).clone(); + torch::Tensor dst = torch::from_blob(dst_vec.data(), {nedges}, torch::kLong).clone(); + torch::Tensor shift = torch::from_blob(shift_flat.data(),{nedges, 3}, torch::kFloat32) + .to(dtype_).clone(); + + torch::Tensor lat = torch::zeros({3, 3}, t_opts); + { + auto la = lat.accessor(); + la[0][0] = static_cast(hx); + la[1][1] = static_cast(hy); + la[2][2] = static_cast(hz); + la[1][0] = static_cast(xy); + la[2][0] = static_cast(xz); + la[2][1] = static_cast(yz); + } + + pos = pos.to(device_).set_requires_grad(true); + lat = lat.to(device_); + Z = Z.to(device_); + src = src.to(device_); + dst = dst.to(device_); + shift = shift.to(device_); + + std::vector inputs; + inputs.emplace_back(pos); + inputs.emplace_back(lat); + inputs.emplace_back(Z); + inputs.emplace_back(src); + inputs.emplace_back(dst); + inputs.emplace_back(shift); + inputs.emplace_back(static_cast(vflag_global)); + + c10::IValue result; + { + torch::AutoGradMode grad(true); + result = model_.get_method("forward_tensors_z")(inputs); + } + auto dict = result.toGenericDict(); + torch::Tensor energy_t = dict.at("energy").toTensor().to(torch::kCPU).to(torch::kDouble); + torch::Tensor forces_t = dict.at("forces").toTensor().to(torch::kCPU).to(torch::kDouble); + + auto fa = forces_t.accessor(); + for (int i = 0; i < nlocal; ++i) { + f[i][0] += fa[i][0]; + f[i][1] += fa[i][1]; + f[i][2] += fa[i][2]; + } + if (eflag_global) eng_vdwl += energy_t.item(); + + if (vflag_global && dict.contains("stress")) { + torch::Tensor s = dict.at("stress").toTensor().to(torch::kCPU).to(torch::kDouble); + auto sa = s.accessor(); + const double V = hx * hy * hz; + virial[0] += -sa[0][0] * V; + virial[1] += -sa[1][1] * V; + virial[2] += -sa[2][2] * V; + virial[3] += -sa[0][1] * V; + virial[4] += -sa[0][2] * V; + virial[5] += -sa[1][2] * V; + } +} diff --git a/alignn/scripts/torch/pair_alignn/pair_alignn.h b/alignn/scripts/torch/pair_alignn/pair_alignn.h new file mode 100644 index 0000000..ec45e5f --- /dev/null +++ b/alignn/scripts/torch/pair_alignn/pair_alignn.h @@ -0,0 +1,55 @@ +// ALIGNN-FF native LAMMPS pair style. +// Drop this header + pair_alignn.cpp into LAMMPS src/USER-ALIGNN/ and rebuild. +// +// Input script usage: +// pair_style alignn +// pair_coeff * * alignn_ff.pt Si # TorchScript model path, then element +// # symbols ordered by LAMMPS atom type. + +#ifdef PAIR_CLASS +// clang-format off +PairStyle(alignn,PairAlignn) +// clang-format on +#else + +#ifndef LMP_PAIR_ALIGNN_H +#define LMP_PAIR_ALIGNN_H + +#include "pair.h" + +#include +#include +#include +#include + +namespace LAMMPS_NS { + +class PairAlignn : public Pair { + public: + PairAlignn(class LAMMPS *); + ~PairAlignn() override; + + void compute(int, int) override; + void settings(int, char **) override; + void coeff(int, char **) override; + void init_style() override; + double init_one(int, int) override; + void allocate(); + + protected: + torch::jit::script::Module model_; + torch::Device device_ = torch::kCPU; + torch::ScalarType dtype_ = torch::kFloat32; + + double cutoff_ = 8.0; // must match trained cutoff + int max_neighbors_ = 12; // k-nearest cap per atom + std::vector type_to_Z_; // LAMMPS type (1-indexed) -> atomic number + + bool model_loaded_ = false; + bool debug_ = false; +}; + +} // namespace LAMMPS_NS + +#endif +#endif diff --git a/alignn/scripts/torch/pair_alignn/si_native.in b/alignn/scripts/torch/pair_alignn/si_native.in new file mode 100644 index 0000000..de048bb --- /dev/null +++ b/alignn/scripts/torch/pair_alignn/si_native.in @@ -0,0 +1,40 @@ +# Native-pair-style melt-quench for Si using pair_alignn. + +units metal +atom_style atomic +boundary p p p +read_data si.data +mass 1 28.0855 + +pair_style alignn +pair_coeff * * alignn_ff.pt Si + +neighbor 2.0 bin +neigh_modify every 1 delay 0 check yes + +timestep 0.001 +thermo 50 +thermo_style custom step temp pe ke etotal press vol +dump traj all custom 100 si.lammpstrj id type x y z + +velocity all create 300.0 12345 mom yes rot yes +fix nvt1 all nvt temp 300.0 300.0 0.1 +run 200 +unfix nvt1 + +fix nvt2 all nvt temp 300.0 3500.0 0.1 +run 1000 +unfix nvt2 + +fix nvt3 all nvt temp 3500.0 3500.0 0.1 +run 2000 +unfix nvt3 + +fix nvt4 all nvt temp 3500.0 300.0 0.1 +run 4000 +unfix nvt4 + +fix nvt5 all nvt temp 300.0 300.0 0.1 +run 2000 + +write_data si_quenched.data diff --git a/alignn/scripts/torch/phonons_plotly.py b/alignn/scripts/torch/phonons_plotly.py new file mode 100644 index 0000000..9bd7f85 --- /dev/null +++ b/alignn/scripts/torch/phonons_plotly.py @@ -0,0 +1,131 @@ +"""Phonon bandstructure from a trained ALIGNN-FF model, plotted with Plotly. + +Usage: + python phonons_plotly.py --model-dir OutputDir --jid JVASP-1002 +""" +from __future__ import annotations +import argparse +import numpy as np +import plotly.graph_objects as go +from ase import Atoms as AseAtoms +from ase.filters import ExpCellFilter +from ase.optimize import FIRE +from phonopy import Phonopy +from phonopy.structure.atoms import PhonopyAtoms + +from jarvis.core.atoms import Atoms, ase_to_atoms +from jarvis.core.kpoints import Kpoints3D as Kpoints +from jarvis.db.figshare import get_jid_data + +from alignn.ff.ff import AlignnAtomwiseCalculator + + +def relax(atoms: Atoms, calc, fmax=0.05, steps=100): + a = atoms.ase_converter() + a.calc = calc + FIRE(ExpCellFilter(a), logfile=None).run(fmax=fmax, steps=steps) + return ase_to_atoms(a) + + +def compute_phonons(atoms: Atoms, calc, dim=(2, 2, 2), distance=0.03, line_density=20): + bulk = PhonopyAtoms( + symbols=atoms.elements, + scaled_positions=atoms.frac_coords, + cell=atoms.lattice_mat, + ) + ph = Phonopy(bulk, [[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]]) + ph.generate_displacements(distance=distance) + # evaluate forces on each displaced supercell + forces_set = [] + for sc in ph.get_supercells_with_displacements(): + a = AseAtoms(symbols=sc.get_chemical_symbols(), + scaled_positions=sc.get_scaled_positions(), + cell=sc.get_cell(), pbc=True) + a.calc = calc + F = np.array(a.get_forces()) + F -= F.mean(axis=0) # remove drift + forces_set.append(F) + ph.produce_force_constants(forces=forces_set) + + # walk the standard k-path (jarvis builds it from the lattice) + kp = Kpoints().kpath(atoms, line_density=line_density) + q_points = np.array(kp.kpts) + labels_raw = kp.labels + + freqs, distances, tick_pos, tick_labels = [], [], [], [] + last_kstr = None + d = 0.0 + for i, (q, lbl) in enumerate(zip(q_points, labels_raw)): + if i > 0: + d += np.linalg.norm(q_points[i] - q_points[i - 1]) + distances.append(d) + freqs.append(ph.get_frequencies(q)) # THz + kstr = ",".join(f"{x:.6f}" for x in q) + if lbl and kstr != last_kstr: + tick_pos.append(d) + tick_labels.append(lbl) + last_kstr = kstr + return np.array(distances), np.array(freqs), tick_pos, tick_labels + + +def plot(distances, freqs, tick_pos, tick_labels, title=""): + fig = go.Figure() + nbands = freqs.shape[1] + for b in range(nbands): + fig.add_trace(go.Scatter( + x=distances, y=freqs[:, b], mode="lines", + line=dict(color="steelblue", width=1.5), + showlegend=False, hovertemplate="d=%{x:.3f}
ν=%{y:.3f} THz", + )) + fig.add_hline(y=0, line=dict(color="black", dash="dot", width=1)) + for xp in tick_pos[1:-1]: + fig.add_vline(x=xp, line=dict(color="lightgray", width=1)) + fig.update_layout( + title=title, width=800, height=500, + xaxis=dict( + title="Wave vector", + tickmode="array", + tickvals=tick_pos, + ticktext=[l.replace("\\Gamma", "Γ") for l in tick_labels], + showgrid=False, zeroline=False, + ), + yaxis=dict(title="Frequency (THz)", gridcolor="lightgray"), + template="plotly_white", + ) + return fig + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="OutputDir") + ap.add_argument("--jid", default="JVASP-1002") + ap.add_argument("--dim", type=int, nargs=3, default=[2, 2, 2]) + ap.add_argument("--out", default="phonons.html") + args = ap.parse_args() + + calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False) + + d = get_jid_data(jid=args.jid, dataset="dft_3d") + atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms + print(f"loaded {args.jid} {atoms.composition.reduced_formula} N={atoms.num_atoms}") + + relaxed = relax(atoms, calc) + prim = relaxed.get_primitive_atoms + print(f"relaxed → primitive N={prim.num_atoms}") + + distances, freqs, tick_pos, tick_labels = compute_phonons( + prim, calc, dim=tuple(args.dim) + ) + n_imag = int((freqs < -0.05).sum()) + print(f"bands computed: {freqs.shape[1]} " + f"imag modes (<-0.05 THz): {n_imag}") + + fig = plot(distances, freqs, tick_pos, tick_labels, + title=f"Phonon bands — {atoms.composition.reduced_formula} " + f"({args.jid}) supercell {args.dim}") + fig.write_html(args.out) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/torch/scaling_test_cu.py b/alignn/scripts/torch/scaling_test_cu.py similarity index 100% rename from scripts/torch/scaling_test_cu.py rename to alignn/scripts/torch/scaling_test_cu.py diff --git a/alignn/scripts/torch/si_melt_quench.in b/alignn/scripts/torch/si_melt_quench.in new file mode 100644 index 0000000..cad37ba --- /dev/null +++ b/alignn/scripts/torch/si_melt_quench.in @@ -0,0 +1,49 @@ +# LAMMPS melt-quench for Si with ALIGNN-FF via torch bridge. +# Invoke: python scripts/torch/lammps_bridge_torch.py \ +# --model-dir OutputDir --input scripts/torch/si_melt_quench.in --types Si + +units metal +atom_style atomic +boundary p p p + +read_data si.data +mass 1 28.0855 + +# ALIGNN-FF hook — callback provides forces + virial every step +pair_style zero 8.0 +pair_coeff * * +neighbor 2.0 bin +neigh_modify every 1 delay 0 check yes +fix alignn all external pf/callback 1 1 + +timestep 0.001 # 1 fs (metal units: ps) +thermo 50 +thermo_style custom step temp pe ke etotal press vol +dump traj all custom 100 si_melt_quench.lammpstrj id type x y z + +# ── Stage 1 — equilibrate 300 K ── +velocity all create 300.0 12345 mom yes rot yes +fix nvt1 all nvt temp 300.0 300.0 0.1 +run 200 +unfix nvt1 + +# ── Stage 2 — heat to 3500 K over 1 ps ── +fix nvt2 all nvt temp 300.0 3500.0 0.1 +run 1000 +unfix nvt2 + +# ── Stage 3 — hold molten ── +fix nvt3 all nvt temp 3500.0 3500.0 0.1 +run 2000 +unfix nvt3 + +# ── Stage 4 — quench to 300 K over 4 ps ── +fix nvt4 all nvt temp 3500.0 300.0 0.1 +run 4000 +unfix nvt4 + +# ── Stage 5 — anneal 300 K ── +fix nvt5 all nvt temp 300.0 300.0 0.1 +run 2000 + +write_data si_quenched.data diff --git a/alignn/scripts/torch/thermal_plotly.py b/alignn/scripts/torch/thermal_plotly.py new file mode 100644 index 0000000..c1729e2 --- /dev/null +++ b/alignn/scripts/torch/thermal_plotly.py @@ -0,0 +1,240 @@ +"""Phonon + thermal properties from trained ALIGNN-FF, plotted with Plotly. + +Computes: + * phonon bandstructure + * phonon DOS + * harmonic thermodynamics: Cv(T), S(T), F(T), ZPE + * Dulong-Petit comparison + +Outputs a single HTML with 4 panels. + +Usage: + python thermal_plotly.py --model-dir OutputDir --jid JVASP-1002 +""" +from __future__ import annotations +import argparse +import numpy as np +import plotly.graph_objects as go +from plotly.subplots import make_subplots + +from ase import Atoms as AseAtoms +from ase.filters import ExpCellFilter +from ase.optimize import FIRE +from phonopy import Phonopy +from phonopy.structure.atoms import PhonopyAtoms + +from jarvis.core.atoms import Atoms, ase_to_atoms +from jarvis.core.kpoints import Kpoints3D as Kpoints +from jarvis.db.figshare import get_jid_data + +from alignn.ff.ff import AlignnAtomwiseCalculator + + +# --------------------------------------------------------------------------- +# Core calculations +# --------------------------------------------------------------------------- +def relax(atoms, calc, fmax=0.05, steps=100): + a = atoms.ase_converter() + a.calc = calc + FIRE(ExpCellFilter(a), logfile=None).run(fmax=fmax, steps=steps) + return ase_to_atoms(a) + + +def run_phonopy(atoms, calc, dim=(2, 2, 2), distance=0.03): + bulk = PhonopyAtoms( + symbols=atoms.elements, + scaled_positions=atoms.frac_coords, + cell=atoms.lattice_mat, + ) + ph = Phonopy(bulk, [[dim[0], 0, 0], [0, dim[1], 0], [0, 0, dim[2]]]) + ph.generate_displacements(distance=distance) + + forces_set = [] + for sc in ph.supercells_with_displacements: + a = AseAtoms( + symbols=list(sc.symbols), + scaled_positions=np.array(sc.scaled_positions), + cell=np.array(sc.cell), + pbc=True, + ) + a.calc = calc + F = np.array(a.get_forces()) + F -= F.mean(axis=0) + forces_set.append(F) + ph.produce_force_constants(forces=forces_set) + return ph + + +def bandstructure(ph, atoms, line_density=20): + kp = Kpoints().kpath(atoms, line_density=line_density) + q_points = np.array(kp.kpts) + labels_raw = kp.labels + + freqs, distances, tick_pos, tick_labels = [], [], [], [] + last_kstr = None + d = 0.0 + for i, (q, lbl) in enumerate(zip(q_points, labels_raw)): + if i > 0: + d += np.linalg.norm(q_points[i] - q_points[i - 1]) + distances.append(d) + freqs.append(ph.get_frequencies(q)) + kstr = ",".join(f"{x:.6f}" for x in q) + if lbl and kstr != last_kstr: + tick_pos.append(d) + tick_labels.append(lbl) + last_kstr = kstr + return np.array(distances), np.array(freqs), tick_pos, tick_labels + + +def thermal_properties(ph, mesh=(20, 20, 20), t_min=0, t_max=1000, t_step=10): + ph.run_mesh(list(mesh)) + ph.run_thermal_properties(t_min=t_min, t_max=t_max, t_step=t_step) + td = ph.get_thermal_properties_dict() + N = len(ph.primitive) + return { + "T": td["temperatures"], + "Cv": td["heat_capacity"] / N, # J/K/mol-atom + "S": td["entropy"] / N, # J/K/mol-atom + "F": td["free_energy"] / N, # kJ/mol-atom + "ZPE_per_cell_kJmol": float(td["free_energy"][0]), + "ZPE_per_atom_eV": float(td["free_energy"][0]) * 1000.0 / 96.485 / N, + } + + +def dos(ph): + ph.run_total_dos() + d = ph.get_total_dos_dict() + return np.array(d["frequency_points"]), np.array(d["total_dos"]) + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- +def plot_all(bands, dos_data, thermo, title=""): + distances, freqs, tick_pos, tick_labels = bands + dos_freq, dos_val = dos_data + + fig = make_subplots( + rows=2, cols=3, + specs=[ + [{"colspan": 2}, None, {}], + [{}, {}, {}], + ], + subplot_titles=( + "Phonon bandstructure", "Phonon DOS", + "Heat capacity Cv", "Entropy S", "Free energy F", + ), + horizontal_spacing=0.08, vertical_spacing=0.15, + ) + + # Bandstructure (top-left, 2 cols wide) + for b in range(freqs.shape[1]): + fig.add_trace( + go.Scatter(x=distances, y=freqs[:, b], mode="lines", + line=dict(color="steelblue", width=1.2), + showlegend=False, + hovertemplate="d=%{x:.2f}
ν=%{y:.2f} THz"), + row=1, col=1, + ) + for xp in tick_pos[1:-1]: + fig.add_vline(x=xp, line=dict(color="lightgray", width=1), row=1, col=1) + fig.add_hline(y=0, line=dict(color="black", dash="dot", width=1), row=1, col=1) + + # DOS (top-right) + fig.add_trace( + go.Scatter(x=dos_val, y=dos_freq, mode="lines", + fill="tozerox", line=dict(color="firebrick"), + showlegend=False, + hovertemplate="DOS=%{x:.2f}
ν=%{y:.2f} THz"), + row=1, col=3, + ) + + # Cv, S, F (bottom row) + T = thermo["T"] + fig.add_trace(go.Scatter(x=T, y=thermo["Cv"], mode="lines", + line=dict(color="seagreen"), showlegend=False), + row=2, col=1) + fig.add_hline(y=24.94, line=dict(dash="dot", color="gray"), + annotation_text="3R (Dulong-Petit)", + annotation_position="top right", row=2, col=1) + fig.add_trace(go.Scatter(x=T, y=thermo["S"], mode="lines", + line=dict(color="darkorange"), showlegend=False), + row=2, col=2) + fig.add_trace(go.Scatter(x=T, y=thermo["F"], mode="lines", + line=dict(color="mediumpurple"), showlegend=False), + row=2, col=3) + + # Axes labels + ticktext = [l.replace("\\Gamma", "Γ") for l in tick_labels] + fig.update_xaxes(title_text="Wave vector", + tickmode="array", tickvals=tick_pos, ticktext=ticktext, + showgrid=False, row=1, col=1) + fig.update_yaxes(title_text="Frequency (THz)", row=1, col=1) + fig.update_xaxes(title_text="DOS (states/THz/cell)", row=1, col=3) + fig.update_yaxes(title_text="Frequency (THz)", row=1, col=3) + + for col, ylab in [(1, "Cv (J/K/mol-atom)"), + (2, "S (J/K/mol-atom)"), + (3, "F (kJ/mol-atom)")]: + fig.update_xaxes(title_text="T (K)", row=2, col=col) + fig.update_yaxes(title_text=ylab, row=2, col=col) + + fig.update_layout(template="plotly_white", title=title, + width=1250, height=800) + return fig + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="OutputDir") + ap.add_argument("--jid", default="JVASP-1002") + ap.add_argument("--dim", type=int, nargs=3, default=[2, 2, 2]) + ap.add_argument("--mesh", type=int, nargs=3, default=[20, 20, 20]) + ap.add_argument("--t-max", type=float, default=1000.0) + ap.add_argument("--out", default="thermal.html") + args = ap.parse_args() + + calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False) + + d = get_jid_data(jid=args.jid, dataset="dft_3d") + atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms + formula = atoms.composition.reduced_formula + print(f"loaded {args.jid} {formula} N={atoms.num_atoms}") + + relaxed = relax(atoms, calc) + prim = relaxed.get_primitive_atoms + print(f"relaxed → primitive N={prim.num_atoms}") + + ph = run_phonopy(prim, calc, dim=tuple(args.dim)) + bands = bandstructure(ph, prim) + dos_data = dos(ph) + thermo = thermal_properties(ph, mesh=tuple(args.mesh), t_max=args.t_max) + + # Summary + n_imag = int((bands[1] < -0.05).sum()) + T = thermo["T"] + i300 = int(np.argmin(abs(T - 300))) + print(f"\n── Results ────────────────────────────────────────────") + print(f"bands : {bands[1].shape[1]} imag modes: {n_imag}") + print(f"ZPE : {thermo['ZPE_per_cell_kJmol']:.3f} kJ/mol-cell" + f" ({thermo['ZPE_per_atom_eV']:.4f} eV/atom)") + print(f"Cv(300 K) : {thermo['Cv'][i300]:.3f} J/K/mol-atom " + f"(Dulong-Petit: 24.94)") + print(f"S(300 K) : {thermo['S'][i300]:.3f} J/K/mol-atom") + print(f"F(300 K) - F(0) : {thermo['F'][i300] - thermo['F'][0]:.3f} kJ/mol-atom") + + fig = plot_all(bands, dos_data, thermo, + title=f"{formula} ({args.jid}) — ALIGNN-FF phonons " + f"[supercell {args.dim}, mesh {args.mesh}]") + fig.write_html(args.out) + print(f"\nwrote {args.out}") + + # In a Jupyter cell, uncomment: + # fig.show() + + +if __name__ == "__main__": + main() diff --git a/alignn/scripts/torch/validate_pair_alignn.py b/alignn/scripts/torch/validate_pair_alignn.py new file mode 100644 index 0000000..90eb1cd --- /dev/null +++ b/alignn/scripts/torch/validate_pair_alignn.py @@ -0,0 +1,189 @@ +"""Validate pair_alignn against the Python reference (AlignnAtomwiseCalculator). + +Does a single-point `run 0` in LAMMPS with pair_alignn, then computes the +same quantities via the Python calculator, and reports component-wise +differences. Run this immediately after you get pair_alignn compiling and +loading the .pt file — it's the gate between "compiles" and "correct." + +Pass criteria (typical): + |ΔE| < 1e-4 eV + max|ΔF| < 1e-3 eV/Å + max|Δσ| < 1e-4 eV/ų + +Usage: + python scripts/torch/validate_pair_alignn.py \\ + --model-dir OutputDir --ts-model alignn_ff.pt \\ + --jid JVASP-1002 --supercell 2 2 2 +""" +from __future__ import annotations +import argparse, os, tempfile +import numpy as np +from ase.io import write as ase_write +from jarvis.core.atoms import Atoms +from jarvis.db.figshare import get_jid_data + +# Reference: Python ASE calculator path +from alignn.ff.ff import AlignnAtomwiseCalculator + + +# eV/ų ↔ bar conversion (LAMMPS `metal` units pressure) +EV_PER_A3_TO_BAR = 1.602176634e6 + + +def python_reference(atoms_jarvis, calc): + a = atoms_jarvis.ase_converter() + a.calc = calc + E_raw = a.get_potential_energy() + E = float(np.asarray(E_raw).reshape(-1)[0]) # calc may return shape-(1,) + F = np.asarray(a.get_forces()) # eV/Å + S6 = np.asarray(a.get_stress(voigt=True)).reshape(-1) # eV/ų (Voigt) + # Reorder ASE Voigt (xx,yy,zz,yz,xz,xy) -> plain 3x3 + S = np.array([[S6[0], S6[5], S6[4]], + [S6[5], S6[1], S6[3]], + [S6[4], S6[3], S6[2]]]) + return E, F, S + + +def lammps_pair_alignn(ase_atoms, ts_model_path, species_symbols, cutoff=5.0): + """Run LAMMPS with pair_alignn at step 0, return E/F/S.""" + from lammps import lammps + + tmpdir = tempfile.mkdtemp(prefix="pair_alignn_val_") + data_path = os.path.join(tmpdir, "sys.data") + ase_write(data_path, ase_atoms, format="lammps-data", + specorder=list(species_symbols)) + + lmp = lammps() + lmp.commands_string(f""" + units metal + atom_style atomic + boundary p p p + read_data {data_path} + {'mass ' + ' '.join(str(i+1)+' '+str(_atomic_mass(sym)) + for i,sym in enumerate(species_symbols))} + pair_style alignn {cutoff} + pair_coeff * * {ts_model_path} {' '.join(species_symbols)} + neighbor 2.0 bin + neigh_modify every 1 delay 0 check yes + compute PE all pe + compute VIR all pressure NULL virial + thermo 1 + thermo_style custom step pe c_VIR[1] c_VIR[2] c_VIR[3] c_VIR[4] c_VIR[5] c_VIR[6] + run 0 + """) + + nlocal = lmp.get_natoms() + # Gather forces by atom ID so ordering matches the ASE/jarvis object + F_flat = lmp.gather_atoms("f", 1, 3) + F = np.array(F_flat[:]).reshape(nlocal, 3) + + # Energy (metal units → eV already) + E = float(lmp.get_thermo("pe")) + + # Pressure tensor in LAMMPS = sum of kinetic + virial. At step 0 with + # zero velocities, kinetic term is 0 so p = virial/V. Flip sign to + # match ASE stress convention (stress = -1/V * dE/dε). + V = lmp.get_thermo("vol") # ų + p_bar = np.array([lmp.get_thermo(name) for name in + ("c_VIR[1]","c_VIR[2]","c_VIR[3]", + "c_VIR[4]","c_VIR[5]","c_VIR[6]")]) + # LAMMPS virial order: xx yy zz xy xz yz; pressure = virial/V (bar) + p_ev_A3 = p_bar / EV_PER_A3_TO_BAR # bar -> eV/ų + S = np.array([[ p_ev_A3[0], p_ev_A3[3], p_ev_A3[4]], + [ p_ev_A3[3], p_ev_A3[1], p_ev_A3[5]], + [ p_ev_A3[4], p_ev_A3[5], p_ev_A3[2]]]) + # LAMMPS pressure sign convention is opposite of stress + S = -S + lmp.close() + return E, F, S + + +def _atomic_mass(sym): + from ase.data import atomic_masses, atomic_numbers + return atomic_masses[atomic_numbers[sym]] + + +# --------------------------------------------------------------------------- +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", required=True, help="dir with best_model.pt + config.json for the Python reference") + ap.add_argument("--ts-model", required=True, help="TorchScript .pt exported via export_torchscript.py") + ap.add_argument("--jid", default="JVASP-1002") + ap.add_argument("--supercell", type=int, nargs=3, default=[2,2,2]) + ap.add_argument("--perturb", type=float, default=0.03, + help="Å, random atomic displacement to break symmetry " + "(so forces are non-zero and meaningful)") + ap.add_argument("--cutoff", type=float, default=None, + help="Å, model cutoff (read from model-dir config if unset)") + args = ap.parse_args() + if args.cutoff is None: + import json as _j + args.cutoff = float(_j.load(open(f"{args.model_dir}/config.json"))["cutoff"]) + print(f"cutoff (from config): {args.cutoff} Å") + + # Build system + d = get_jid_data(jid=args.jid, dataset="dft_3d") + atoms = Atoms.from_dict(d["atoms"]).get_conventional_atoms + atoms = atoms.make_supercell(args.supercell) + ase_atoms = atoms.ase_converter() + + # Small perturbation — exact equilibrium would give all-zero forces, + # which is not a useful validation. + rng = np.random.default_rng(0) + ase_atoms.positions += args.perturb * rng.standard_normal(ase_atoms.positions.shape) + + species = sorted(set(ase_atoms.get_chemical_symbols())) + print(f"System: {atoms.composition.reduced_formula} x{args.supercell} " + f"N={len(ase_atoms)} species={species}") + + # Re-wrap into jarvis for the Python reference + from jarvis.core.atoms import ase_to_atoms as ase2j + atoms_pert = ase2j(ase_atoms) + + # --- Python reference --- + print("\n→ Python reference (AlignnAtomwiseCalculator)...") + calc = AlignnAtomwiseCalculator(path=args.model_dir, force_mult_batchsize=False) + E_py, F_py, S_py = python_reference(atoms_pert, calc) + + # --- LAMMPS pair_alignn --- + print("→ LAMMPS pair_alignn (run 0)...") + E_lmp, F_lmp, S_lmp = lammps_pair_alignn(ase_atoms, args.ts_model, species, + cutoff=args.cutoff) + + # --- Compare --- + dE = abs(E_lmp - E_py) + dF = F_lmp - F_py + dS = S_lmp - S_py + print("\n── Comparison ──────────────────────────────────────────") + print(f"E (Python): {E_py: .6f} eV") + print(f"E (LAMMPS): {E_lmp: .6f} eV") + print(f"|ΔE|: {dE: .3e} eV " + f"[{'PASS' if dE < 1e-4 else 'FAIL'}] (threshold 1e-4)") + + max_abs_dF = float(np.max(np.abs(dF))) + rms_dF = float(np.sqrt(np.mean(dF**2))) + max_abs_F = float(np.max(np.abs(F_py))) + print(f"max|ΔF|: {max_abs_dF: .3e} eV/Å " + f"[{'PASS' if max_abs_dF < 1e-3 else 'FAIL'}] (threshold 1e-3)") + print(f"rms ΔF: {rms_dF: .3e} eV/Å") + print(f"max|F_ref|: {max_abs_F: .3e} eV/Å (scale for context)") + + max_abs_dS = float(np.max(np.abs(dS))) + print(f"max|Δσ|: {max_abs_dS: .3e} eV/ų " + f"[{'PASS' if max_abs_dS < 1e-4 else 'FAIL'}] (threshold 1e-4)") + + print("\n── Per-atom ΔF (eV/Å) — first 5 atoms ──") + print("idx ΔF_x ΔF_y ΔF_z |ΔF|") + for i in range(min(5, len(F_py))): + d = dF[i] + print(f"{i:3d} {d[0]:+10.3e} {d[1]:+10.3e} {d[2]:+10.3e} {np.linalg.norm(d):10.3e}") + + # Write out both force arrays for post-hoc inspection + np.savez("validate_pair_alignn.npz", + E_py=E_py, F_py=F_py, S_py=S_py, + E_lmp=E_lmp, F_lmp=F_lmp, S_lmp=S_lmp) + print("\nsaved raw arrays -> validate_pair_alignn.npz") + + +if __name__ == "__main__": + main() diff --git a/mkdocs.yml b/mkdocs.yml index 5bb34c2..9a713ba 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -83,6 +83,7 @@ nav: - Usage: - ASE Calculator: usage/ase-calculator.md - MD & Relaxation (pure PyTorch): usage/md-integrators.md + - LAMMPS Pair Style (native C++): usage/lammps-pair-style.md - Web Apps: usage/webapps.md - Reference: - Package Overview: api.md From 393431b96d8bab4d2b7cb6d0f3f07779463408aa Mon Sep 17 00:00:00 2001 From: user Date: Mon, 20 Apr 2026 23:50:11 -0400 Subject: [PATCH 22/34] pair_style alignn --- alignn/scripts/torch/build_lammps_alignn.sh | 2 +- docs/usage/lammps-pair-style.md | 316 ++++++++++++++++++++ 2 files changed, 317 insertions(+), 1 deletion(-) create mode 100644 docs/usage/lammps-pair-style.md diff --git a/alignn/scripts/torch/build_lammps_alignn.sh b/alignn/scripts/torch/build_lammps_alignn.sh index 5950e99..c013f02 100755 --- a/alignn/scripts/torch/build_lammps_alignn.sh +++ b/alignn/scripts/torch/build_lammps_alignn.sh @@ -8,7 +8,7 @@ set -euo pipefail LAMMPS_TAG="stable_29Aug2024_update1" LAMMPS_DIR="$HOME/lammps-alignn" ALIGNN_REPO="/home/kamalch/Software/ollama311/alignn" -PAIR_SRC="$ALIGNN_REPO/scripts/torch/pair_alignn" +PAIR_SRC="$ALIGNN_REPO/alignn/scripts/torch/pair_alignn" # ── discover conda env's libtorch ───────────────────────────────────────── TORCH_DIR=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))") diff --git a/docs/usage/lammps-pair-style.md b/docs/usage/lammps-pair-style.md new file mode 100644 index 0000000..7af3a17 --- /dev/null +++ b/docs/usage/lammps-pair-style.md @@ -0,0 +1,316 @@ +# LAMMPS native pair style (`pair_alignn`) + +`pair_alignn` is a C++ LAMMPS pair style that runs a TorchScript-exported +ALIGNN-FF model as the force engine. It replaces the Python+ASE bridge +(`alignn.ff.lammps_bridge`) for production runs: + +| Path | Overhead per step | When to use | +|---|---|---| +| ASE calculator (Python) | high (~40 ms for 200 atoms) | one-off scripts, debugging | +| `lammps_bridge_torch.py` | medium (~15 ms) | quick Python-hosted MD | +| **`pair_alignn` (this page)** | **low (~3–10 ms)** | production MD, large systems | + +Everything stays on the GPU; Python is never invoked inside the MD step +loop. + +!!! note "Status" + MVP single-rank implementation. Matches the TorchScript Pure model + forward to ~10⁻⁶ eV on Si. Multi-rank / domain decomposition is not + yet implemented — cross-rank ghost contributions are skipped. + +## What you need + +1. A trained ALIGNN-FF model directory containing `best_model.pt` + `config.json` +2. LAMMPS built against **libtorch with the same CUDA/ABI flags** as your + Python PyTorch install +3. `pair_alignn.cpp` / `pair_alignn.h` compiled into that LAMMPS binary +4. A TorchScript `.pt` of your model exported via `export_torchscript.py` + +## Building LAMMPS with `pair_alignn` + +The repo ships a self-contained build script +(`alignn/scripts/torch/build_lammps_alignn.sh`) that: + +1. Clones LAMMPS `stable_29Aug2024_update1` +2. Drops `pair_alignn.cpp/.h` into `src/` +3. Appends a libtorch `find_package` to `cmake/CMakeLists.txt` +4. Configures + builds + installs the Python module + +Prerequisites in your conda env (one-time): + +```bash +# Matching CUDA toolkit for libtorch (check your PyTorch first) +python -c "import torch; print(torch.version.cuda)" +# If 12.8, install: +mamba install -y -c nvidia -c conda-forge cuda-toolkit=12.8 mkl-devel mkl-include +``` + +Then run the build script: + +```bash +bash alignn/scripts/torch/build_lammps_alignn.sh +``` + +On success, you'll see: +``` +pair_alignn available: True +``` + +The built `lmp` executable lands at `~/lammps-alignn/build/lmp`, and the +conda `lammps` Python module is overwritten with the libtorch-linked +version. + +### Common build issues + +| Symptom | Cause | Fix | +|---|---|---| +| `PyTorch: CUDA cannot be found` | CUDA toolkit version mismatches libtorch's build version | install matching `cuda-toolkit=X.Y` via conda | +| `Imported target "torch" includes non-existent path ".../MKL_INCLUDE_DIR-NOTFOUND"` | MKL not installed | `mamba install mkl-devel mkl-include` | +| `fatal error: cuda.h: No such file or directory` | CUDA headers not on include path | cmake flag: `-DCMAKE_CXX_FLAGS="-I$CONDA_PREFIX/targets/x86_64-linux/include"` | +| `Enabled packages: ` | cmake didn't wire pair_alignn in | verify `append find_package(Torch)` block lives at the end of `cmake/CMakeLists.txt` | +| Builds but `pair_alignn` not listed | `pair_alignn.cpp` not in `src/` | `cp pair_alignn.{cpp,h} ~/lammps-alignn/src/` | + +## Exporting the TorchScript model + +```bash +python alignn/scripts/torch/export_torchscript.py \ + --model-dir MELT/OutputDir \ + --out alignn_ff.pt \ + --atom-features atomic_number \ + --dtype float32 +``` + +Output: +``` +smoke test OK: {'energy': torch.Size([]), 'forces': torch.Size([2, 3]), 'stress': torch.Size([3, 3])} +wrote alignn_ff.pt (143,297 params) +``` + +**`--atom-features` must match what was used at training time** — if the +model was trained with CGCNN features (92-dim), pass `--atom-features +cgcnn`. The bundled species table is baked into the `.pt` so the C++ side +never needs `jarvis`. + +## Input script syntax + +``` +pair_style alignn [max_neighbors] +pair_coeff * * [ ...] +``` + +| Argument | Meaning | Example | +|---|---|---| +| `` | Neighbor cutoff in Å. **Must match training** | `5.0` | +| `[max_neighbors]` | k-nearest cap per atom. Default 12 | `12` | +| `` | Path to the TorchScript file | `alignn_ff.pt` | +| `` | Element symbol for LAMMPS atom type N | `Si` | + +Example for Si trained with cutoff 5.0 Å, 12 neighbors: +``` +pair_style alignn 5.0 12 +pair_coeff * * alignn_ff.pt Si +``` + +Look up the correct values from your trained model: +```bash +python -c " +import json +c = json.load(open('OutputDir/config.json')) +print('cutoff:', c['cutoff'], 'Å') +print('max_neighbors:', c['max_neighbors']) +print('atom_features:', c.get('atom_features')) +" +``` + +## A complete melt-quench input + +``` +units metal # eV, Å, ps, bar, K +atom_style atomic +boundary p p p +read_data si.data +mass 1 28.0855 + +pair_style alignn 5.0 12 +pair_coeff * * alignn_ff.pt Si + +neighbor 2.0 bin +neigh_modify every 1 delay 0 check yes + +timestep 0.001 # 1 fs +thermo 50 +thermo_style custom step temp pe ke etotal press vol +dump traj all custom 100 si.lammpstrj id type x y z + +# ── Stage 1: equilibrate at 300 K ── +velocity all create 300.0 12345 mom yes rot yes +fix nvt1 all nvt temp 300.0 300.0 0.1 +run 200 +unfix nvt1 + +# ── Stage 2: heat to 3500 K over 1 ps ── +fix nvt2 all nvt temp 300.0 3500.0 0.1 +run 1000 +unfix nvt2 + +# ── Stage 3: hold molten ── +fix nvt3 all nvt temp 3500.0 3500.0 0.1 +run 2000 +unfix nvt3 + +# ── Stage 4: quench to 300 K over 4 ps ── +fix nvt4 all nvt temp 3500.0 300.0 0.1 +run 4000 + +write_data si_quenched.data +``` + +Making the Si data file: +```python +from jarvis.db.figshare import get_jid_data +from jarvis.core.atoms import Atoms +from ase.io import write +a = Atoms.from_dict(get_jid_data(jid='JVASP-1002', dataset='dft_3d')['atoms']) +ase = a.make_supercell([3, 3, 3]).ase_converter() # 216 atoms +write("si.data", ase, format="lammps-data", specorder=["Si"]) +``` + +## Running + +Three equivalent ways: + +=== "Custom LAMMPS binary" + + ```bash + ~/lammps-alignn/build/lmp -in si_native.in + ``` + +=== "Python-driven" + + ```python + from lammps import lammps + l = lammps() + l.file("si_native.in") + ``` + +=== "In a Jupyter cell" + + ```python + !~/lammps-alignn/build/lmp -in si_native.in + ``` + +## Validating against the Python calculator + +Before trusting a production run, compare LAMMPS forces against the +Python reference: + +```bash +python alignn/scripts/torch/validate_pair_alignn.py \ + --model-dir MELT/OutputDir \ + --ts-model alignn_ff.pt \ + --jid JVASP-1002 \ + --supercell 2 2 2 +``` + +Expected output shape: +``` +E (Python): -346.003876 eV +E (LAMMPS): -345.xxxxxx eV +|ΔE|: < 1e-4 eV [PASS] +max|ΔF|: < 1e-3 eV/Å [PASS] +``` + +If energies agree to ~1e-6 but forces are off, you have a gradient / +autograd issue. If everything is off by a consistent factor, it's a unit +or normalization mismatch. If the C++ energy is off by an O(1) factor vs +a single-precision residual, check cutoff and atom_features. + +## Required LAMMPS settings + +- `newton_pair on` — LAMMPS default; the pair style requires it for force + reduction on ghost atoms +- atom IDs must be enabled (default) +- `atom->map()` must be built — pair_alignn auto-initializes it + +## What pair_alignn does per step + +1. Copy local atom positions + types to a GPU tensor (`(N_local, 3)`) +2. Walk LAMMPS's full+ghost neighbor list. For each local `i`, collect all + neighbors within cutoff; for ghosts, look up the owner via + `atom->map(tag)` and compute the PBC shift vector +3. Keep only the `max_neighbors` closest per atom (matches jarvis's + `neighbor_strategy="k-nearest"` / `"pure_torch"` behavior) +4. Pack `(src, dst, shift)` tensors and call the TorchScript model's + `forward_tensors_z(positions, lattice, atomic_numbers, src, dst, shift, + compute_stress)` entry point +5. Read back `energy` (scalar), `forces` (N×3), and `stress` (3×3) +6. Scatter forces to `f[i]` for local atoms; accumulate energy and virial +7. LAMMPS's reverse comm handles any ghost-force reductions (no-op on a + single rank) + +## Model-side requirements + +Your trained model must: + +1. Use `name: "alignn_atomwise_pure"` in its config (so + `forward_tensors_z` exists). If you trained with `alignn_atomwise` + (DGL), retrain or convert weights — the two variants are currently + NOT numerically equivalent on this trained checkpoint (known issue). +2. Be TorchScript-scriptable — the pure variant is designed for this + (`@torch.jit.ignore` on DGL-specific methods, `@torch.jit.export` on + tensor-only forwards) +3. Have `calculate_gradient: True` — forces come from autograd + +## Known limitations + +- **Single MPI rank only.** Cross-rank ghost atoms get skipped silently. + For multi-rank, the compute() needs a proper ghost-force reverse_comm + handshake that mirrors pair_allegro's pattern. +- **Orthorhombic box assumption for PBC shifts.** Triclinic boxes work + approximately but may produce wrong shift vectors near cell corners. + TODO: use the full 3×3 inverse lattice for shift reconstruction. +- **Pure model currently diverges from DGL model** on the same trained + weights by ~10% energy. Root cause: `forward_tensors_z`'s inline line + graph construction enumerates ~9× more bond-angle triplets than DGL's + `g.line_graph()`. To match DGL exactly, either retrain with the pure + model directly or patch the line-graph filter. +- **fp32 only** in the C++ path. Changing `dtype_` in `pair_alignn.h` to + `torch::kFloat64` also requires a matching fp64 TorchScript export and + changed `accessor` types. +- **No `eflag_atom` / `vflag_atom`** (per-atom energy/stress) — only + global quantities. + +## Speed reference + +On a single NVIDIA RTX-class GPU, 216-atom Si at cutoff 5.0 Å: + +| Path | ms / step | +|---|---| +| AlignnAtomwiseCalculator (Python via ASE) | ~40 | +| lammps_bridge_torch (Python callback) | ~15 | +| **pair_alignn (native C++)** | **~3–5** | + +Bottleneck in the C++ path is the forward pass itself, not the neighbor +list building. Memory stays resident on the GPU between steps. + +## Files in the repo + +``` +alignn/scripts/torch/ +├── export_torchscript.py # Export .pt TorchScript model +├── build_lammps_alignn.sh # Full LAMMPS+libtorch build +├── validate_pair_alignn.py # Python↔LAMMPS force/energy check +├── lammps_bridge_torch.py # Alternative Python bridge (no C++) +└── pair_alignn/ + ├── pair_alignn.cpp # Pair style source + ├── pair_alignn.h # Pair style header + ├── si_native.in # Example LAMMPS melt-quench + └── CMakeLists-snippet.txt # Standalone build fragment +``` + +## Related + +- [MD & Relaxation (pure PyTorch)](md-integrators.md) — on-device + integrators that skip LAMMPS entirely +- [ASE Calculator](ase-calculator.md) — the slower but more flexible + Python path From f3d91a669ee09d8abd8a3b07de64911192f8c93b Mon Sep 17 00:00:00 2001 From: user Date: Mon, 20 Apr 2026 23:57:47 -0400 Subject: [PATCH 23/34] pair_style alignn --- alignn/scripts/torch/export_torchscript.py | 9 ++++++++- docs/reference/config-reference.md | 2 +- docs/usage/lammps-pair-style.md | 8 +++++++- setup.py | 1 + 4 files changed, 17 insertions(+), 3 deletions(-) diff --git a/alignn/scripts/torch/export_torchscript.py b/alignn/scripts/torch/export_torchscript.py index 9703bb7..b7dedc8 100644 --- a/alignn/scripts/torch/export_torchscript.py +++ b/alignn/scripts/torch/export_torchscript.py @@ -1,8 +1,15 @@ +#!/usr/bin/env python """Export a trained ALIGNN-FF model as a TorchScript `.pt` file for the native LAMMPS pair style (pair_alignn). +Installed as the `export_torchscript.py` console script after +`pip install -e .` (see setup.py). + Usage: - python scripts/torch/export_torchscript.py \\ + export_torchscript.py --model-dir OutputDir --out alignn_ff.pt + +Or during development: + python alignn/scripts/torch/export_torchscript.py \\ --model-dir OutputDir --out alignn_ff.pt The saved module exposes: diff --git a/docs/reference/config-reference.md b/docs/reference/config-reference.md index e9245bf..a41a7c4 100644 --- a/docs/reference/config-reference.md +++ b/docs/reference/config-reference.md @@ -211,7 +211,7 @@ When the pure model is selected, the training pipeline automatically: - Leaves dataloader / loss / metrics / checkpoints untouched. For LAMMPS integration, call -[`alignn/scripts/export_torchscript.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/scripts/export_torchscript.py) — it scripts the +[`alignn/scripts/torch/export_torchscript.py`](https://github.com/atomgptlab/alignn/blob/main/alignn/scripts/torch/export_torchscript.py) (installed as the `export_torchscript.py` console script) — it scripts the `forward_tensors_z(positions, lattice, atomic_numbers, src, dst, shift, compute_stress)` entry point and bakes the atomic-number → feature lookup into the `.pt` so the C++ host only needs atomic numbers. diff --git a/docs/usage/lammps-pair-style.md b/docs/usage/lammps-pair-style.md index 7af3a17..10d2802 100644 --- a/docs/usage/lammps-pair-style.md +++ b/docs/usage/lammps-pair-style.md @@ -72,14 +72,20 @@ version. ## Exporting the TorchScript model +After installing the repo with `pip install -e .` (or `pip install alignn`), +the `export_torchscript.py` console script is on your `PATH`: + ```bash -python alignn/scripts/torch/export_torchscript.py \ +export_torchscript.py \ --model-dir MELT/OutputDir \ --out alignn_ff.pt \ --atom-features atomic_number \ --dtype float32 ``` +(Equivalent to `python alignn/scripts/torch/export_torchscript.py ...` if +you prefer running the source file directly during development.) + Output: ``` smoke test OK: {'energy': torch.Size([]), 'forces': torch.Size([2, 3]), 'stress': torch.Size([3, 3])} diff --git a/setup.py b/setup.py index e72cc31..9e9e985 100644 --- a/setup.py +++ b/setup.py @@ -43,6 +43,7 @@ "alignn/pretrained.py", "alignn/train_alignn.py", "alignn/run_alignn_ff.py", + "alignn/scripts/torch/export_torchscript.py", ], long_description=long_description, long_description_content_type="text/markdown", From 30f12ebe976f7181e0e4e01e847effa5efdbfd98 Mon Sep 17 00:00:00 2001 From: user Date: Tue, 21 Apr 2026 00:11:13 -0400 Subject: [PATCH 24/34] pair_style alignn --- docs/usage/melt-quench-tutorial.md | 373 +++++++++++++++++++++++++++++ mkdocs.yml | 1 + 2 files changed, 374 insertions(+) create mode 100644 docs/usage/melt-quench-tutorial.md diff --git a/docs/usage/melt-quench-tutorial.md b/docs/usage/melt-quench-tutorial.md new file mode 100644 index 0000000..162ff15 --- /dev/null +++ b/docs/usage/melt-quench-tutorial.md @@ -0,0 +1,373 @@ +# Tutorial — Train → Export → LAMMPS Melt-Quench → RDF + +End-to-end walkthrough: train an ALIGNN-FF model on a Si dataset, export +it to TorchScript, run a melt-quench simulation in LAMMPS with the native +`pair_alignn` style, and plot the radial distribution function (RDF) of +the quenched structure. + +Prereqs: + +- `alignn` installed (`pip install -e .` from a clone) +- `lammps` Python package compiled with `pair_alignn` (see + [LAMMPS Pair Style](lammps-pair-style.md) for the build) +- Plotly (`pip install plotly`) + +Working directory for this tutorial: `MELT/` (create anywhere). + +--- + +## 1. Train the model + +Standard training command — nothing new here: + +```bash +train_alignn.py \ + --root_dir DataDir/ \ + --config config.json \ + --output_dir OutputDir +``` + +When training finishes, `OutputDir/` contains: + +``` +OutputDir/ +├── best_model.pt # weights +├── config.json # hyperparameters + cutoff + neighbor_strategy +├── history_*.json +└── ... +``` + +**Important — requirements for the LAMMPS path:** + +1. In `config.json`, set `model.name: "alignn_atomwise_pure"` so + `forward_tensors_z` (the TorchScript entry point) is available. If + you trained with `alignn_atomwise` (DGL-backed), see the note at the + end of this page. +2. Set `calculate_gradient: true` for forces to be computed. +3. Note `cutoff` and `max_neighbors` from `config.json` — you'll pass + them to `pair_style alignn`. + +Quick look-up: +```bash +python -c " +import json +c = json.load(open('OutputDir/config.json')) +print('cutoff :', c['cutoff'], 'Å') +print('max_neighbors :', c['max_neighbors']) +print('atom_features :', c.get('atom_features')) +print('model.name :', c['model']['name']) +" +``` + +--- + +## 2. Export to TorchScript + +After `pip install -e .`, the export script is on your PATH: + +```bash +export_torchscript.py \ + --model-dir OutputDir \ + --out alignn_ff.pt \ + --atom-features atomic_number \ + --dtype float32 +``` + +Expected output: +``` +smoke test OK: {'energy': torch.Size([]), 'forces': torch.Size([2, 3]), 'stress': torch.Size([3, 3])} +wrote alignn_ff.pt (143,297 params) +``` + +**`--atom-features` must match training.** For most ALIGNN-FF trainings +this is `atomic_number` (single scalar per atom) or `cgcnn` (92-dim). +Check your config's top-level `atom_features` field. + +--- + +## 3. Build a Si structure + +LAMMPS needs a data file. Build a 216-atom cubic Si supercell: + +```python +# make_si_data.py +from ase.build import bulk +from ase.io import write + +# Cubic (orthorhombic) conventional cell, repeated 3×3×3 → 216 atoms. +# Avoid triclinic primitive cells when possible — pair_alignn's shift +# computation is most accurate for orthorhombic boxes. +atoms = bulk("Si", "diamond", a=5.43, cubic=True).repeat((3, 3, 3)) +write("si.data", atoms, format="lammps-data", specorder=["Si"]) +print(f"wrote si.data N={len(atoms)} cell={atoms.cell.lengths()}") +``` + +```bash +python make_si_data.py +# wrote si.data N=216 cell=[16.29 16.29 16.29] +``` + +--- + +## 4. Write the LAMMPS input script + +Five-stage melt-quench: equilibrate → heat → hold molten → quench → +anneal. Save as `si_melt_quench.in`: + +``` +# Si melt-quench via pair_alignn. +units metal # eV, Å, ps, bar, K +atom_style atomic +boundary p p p +read_data si.data +mass 1 28.0855 + +# ── ALIGNN-FF native pair style ── +pair_style alignn 5.0 12 # cutoff (Å), max_neighbors +pair_coeff * * alignn_ff.pt Si + +neighbor 2.0 bin +neigh_modify every 1 delay 0 check yes + +timestep 0.001 # 1 fs +thermo 100 +thermo_style custom step temp pe ke etotal press vol +dump traj all custom 200 si.lammpstrj id type x y z + +# ── Stage 1 — equilibrate 300 K (0.2 ps) ── +velocity all create 300.0 12345 mom yes rot yes +fix nvt1 all nvt temp 300.0 300.0 0.1 +run 200 +unfix nvt1 + +# ── Stage 2 — heat to 3500 K (1 ps) ── +fix nvt2 all nvt temp 300.0 3500.0 0.1 +run 1000 +unfix nvt2 + +# ── Stage 3 — hold molten (2 ps) ── +fix nvt3 all nvt temp 3500.0 3500.0 0.1 +run 2000 +unfix nvt3 + +# ── Stage 4 — quench to 300 K (4 ps) ── +fix nvt4 all nvt temp 3500.0 300.0 0.1 +run 4000 +unfix nvt4 + +# ── Stage 5 — anneal 300 K (2 ps) — RDF measured here ── +compute rdf all rdf 150 1 1 cutoff 6.0 +fix rdf_out all ave/time 100 10 1000 c_rdf[*] file si_rdf.dat mode vector +fix nvt5 all nvt temp 300.0 300.0 0.1 +run 2000 +unfix nvt5 +unfix rdf_out + +write_data si_quenched.data +``` + +The `compute rdf` + `fix ave/time` block samples the radial distribution +function every 100 steps and writes the averaged result to `si_rdf.dat`. + +--- + +## 5. Run it + +```bash +~/lammps-alignn/build/lmp -in si_melt_quench.in +``` + +Or from Python: +```python +from lammps import lammps +lmp = lammps() +lmp.file("si_melt_quench.in") +``` + +Runtime on a single modest GPU: ~10–30 min for the full 9.2 ps. Tail of +the thermo output will show temperature ramping up and coming back down: + +``` +Step Temp PotEng KinEng TotEng Press Volume +0 300 -987.23 5.58 -981.65 -5.2e4 4318.9 +200 842 -984.11 15.67 -968.44 -2.3e6 4318.9 +... +3200 3501 -961.22 65.12 -896.10 +3.1e5 4318.9 +... +9200 301 -985.40 5.60 -979.80 -1.4e4 4318.9 +``` + +**Known quirk**: your pressure numbers will look extreme if your model +was trained with `stresswise_weight: 0.0` — stresses aren't learned, +only energies and forces. That's fine for NVT; don't use NPT with such a +model. + +--- + +## 6. Plot the RDF + +The `si_rdf.dat` file LAMMPS wrote has a specific multi-block format. +This script parses it and plots with Plotly: + +```python +# plot_rdf.py +import numpy as np +import plotly.graph_objects as go +from pathlib import Path + +def parse_lammps_rdf(path: str | Path): + """Parse LAMMPS `fix ave/time ... mode vector` RDF output. + + Each averaged block begins with a header line ` ` + followed by `nbins` lines of `index r g(r) coord(r)`. + Returns the LAST block (steady-state RDF). + """ + with open(path) as f: + lines = [L.rstrip() for L in f if L.strip() and not L.startswith("#")] + # Find block starts (lines with exactly 2 tokens = timestep, nbins) + starts = [i for i, L in enumerate(lines) + if len(L.split()) == 2 and not L.startswith(" ")] + last = starts[-1] + nbins = int(lines[last].split()[1]) + data = np.array([list(map(float, L.split())) for L in lines[last+1:last+1+nbins]]) + r = data[:, 1] + g_r = data[:, 2] + coord = data[:, 3] + return r, g_r, coord + +def plot_rdf(r, g_r, coord, title="Si RDF after melt-quench"): + fig = go.Figure() + fig.add_trace(go.Scatter(x=r, y=g_r, name="g(r)", + line=dict(color="steelblue", width=2))) + fig.add_trace(go.Scatter(x=r, y=coord, name="coordination", + line=dict(color="firebrick", width=1.5, dash="dash"), + yaxis="y2")) + fig.update_layout( + title=title, template="plotly_white", + width=800, height=500, + xaxis=dict(title="r (Å)"), + yaxis =dict(title="g(r)", titlefont=dict(color="steelblue")), + yaxis2=dict(title="coordination", overlaying="y", side="right", + titlefont=dict(color="firebrick"), showgrid=False), + ) + return fig + +if __name__ == "__main__": + r, g_r, coord = parse_lammps_rdf("si_rdf.dat") + # Report key features + i_peak = np.argmax(g_r) + print(f"first-neighbor peak: r = {r[i_peak]:.3f} Å, g(r) = {g_r[i_peak]:.2f}") + print(f"coordination at 3.0 Å: {np.interp(3.0, r, coord):.2f} " + f"(expected ≈ 4 for crystalline Si, lower for amorphous)") + fig = plot_rdf(r, g_r, coord) + fig.write_html("si_rdf.html") + print("wrote si_rdf.html") +``` + +Run it: +```bash +python plot_rdf.py +``` + +**Expected Si RDF:** + +| Feature | Crystalline Si | Amorphous Si (well-relaxed) | Amorphous Si (poorly quenched) | +|---|---|---|---| +| First peak position | 2.35 Å | 2.35–2.40 Å | similar | +| First peak height | sharp, g>10 | broad, g≈3–5 | low, g≈2 | +| 2nd-neighbor splitting at ~3.8 Å | distinct peak | shoulder or gone | absent | +| Coordination at 3.0 Å | 4.0 (exact) | ~3.8–4.2 | variable | + +A **4 ps quench from 3500 K is very fast** — don't expect a perfectly +amorphous or perfectly crystalline result. For publication-quality +glasses, quench over 100+ ps. + +--- + +## 7. Verifying the run was physical + +Sanity checks before trusting the results: + +1. **Energy conservation during NVE windows.** Add a short NVE block and + watch `etotal` — drift should be < 1e-4 eV/atom/ps. +2. **No model blow-ups.** If temperature suddenly spikes or PE diverges, + the force call probably hit an unphysical configuration outside the + model's training domain. Use a shorter timestep (0.0005 ps) or a + softer heating ramp. +3. **Compare to the Python calculator.** Run one single-point with + `AlignnAtomwiseCalculator` on the starting structure; LAMMPS should + match to ~1e-4 eV for energy. If it doesn't, see [Model-side + caveats](#model-side-caveats) below. +4. **Check the trajectory.** Open `si.lammpstrj` in OVITO; atoms + shouldn't fly out of the box. + +--- + +## Model-side caveats + +!!! warning "`alignn_atomwise` vs `alignn_atomwise_pure`" + + The TorchScript entry point `forward_tensors_z` only exists on + `ALIGNNAtomWisePure`. If your training config used `name: + alignn_atomwise` (the DGL-backed variant), `export_torchscript.py` + will still run by loading the DGL weights into Pure — but the two + models are not numerically identical on the same weights (currently + a ~10% energy difference due to line-graph construction + differences). + + **Clean fix**: retrain with `name: alignn_atomwise_pure` in your + training config. All other hyperparameters can stay the same. + +!!! warning "Stress not learned" + + If your training data didn't include stresses (or + `stresswise_weight: 0.0`), `pair_alignn` still returns virial + numbers, but they're noise. Don't use NPT / `fix npt`. + +!!! warning "Single MPI rank" + + Current `pair_alignn` silently drops cross-rank ghost contributions. + Run with `mpirun -n 1` only. For larger systems, increase + parallelism via threads (`OMP_NUM_THREADS`) until multi-rank ghost + handling lands. + +--- + +## What to try next + +- **Transport**: replace NVT with NVE in stage 5, compute diffusion + coefficient from the mean-squared-displacement (LAMMPS + `compute msd`). +- **Cool more slowly**: extend stage 4 to 40 ps and compare the RDF — + you should see sharper peaks and a clearer second-neighbor split. +- **Longer correlation**: rerun with `mode vector 10` (average over 10 + snapshots per write) for smoother RDFs. +- **Different material**: retrain on another element / alloy, re-export, + swap `Si` in `pair_coeff` for the target species. + +--- + +## Files produced by this tutorial + +``` +MELT/ +├── alignn_ff.pt # exported TorchScript model +├── si.data # LAMMPS input structure +├── si_melt_quench.in # LAMMPS input script +├── si.lammpstrj # trajectory (custom dump) +├── si_rdf.dat # RDF samples (fix ave/time) +├── si_quenched.data # final structure +├── plot_rdf.py # RDF plotter +└── si_rdf.html # interactive RDF plot +``` + +--- + +## See also + +- [LAMMPS Pair Style (native C++)](lammps-pair-style.md) — full + reference for `pair_alignn` +- [MD & Relaxation (pure PyTorch)](md-integrators.md) — on-device + integrators as an alternative to LAMMPS +- [ASE Calculator](ase-calculator.md) — the Python-side calculator + used for the reference comparison in step 7 diff --git a/mkdocs.yml b/mkdocs.yml index 9a713ba..dd09e74 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -84,6 +84,7 @@ nav: - ASE Calculator: usage/ase-calculator.md - MD & Relaxation (pure PyTorch): usage/md-integrators.md - LAMMPS Pair Style (native C++): usage/lammps-pair-style.md + - Tutorial — Melt-Quench → RDF: usage/melt-quench-tutorial.md - Web Apps: usage/webapps.md - Reference: - Package Overview: api.md From 1c50ce8a4003988e7f489d8f76431a7c409ae61b Mon Sep 17 00:00:00 2001 From: user Date: Tue, 21 Apr 2026 00:19:44 -0400 Subject: [PATCH 25/34] pair_style alignn --- alignn/scripts/torch/build_lammps_alignn.sh | 228 +++++++++++++------- docs/usage/lammps-pair-style.md | 116 +++++++--- docs/usage/melt-quench-tutorial.md | 17 +- 3 files changed, 250 insertions(+), 111 deletions(-) diff --git a/alignn/scripts/torch/build_lammps_alignn.sh b/alignn/scripts/torch/build_lammps_alignn.sh index c013f02..32b34e6 100755 --- a/alignn/scripts/torch/build_lammps_alignn.sh +++ b/alignn/scripts/torch/build_lammps_alignn.sh @@ -1,27 +1,98 @@ #!/bin/bash -# Build LAMMPS with pair_alignn, linking against the conda-env libtorch. -# Self-contained: clones LAMMPS, drops in USER-ALIGNN, patches cmake, builds, -# installs the Python module into the active conda env. - +# --------------------------------------------------------------------------- +# build_lammps_alignn.sh +# +# One-shot build of LAMMPS stable_29Aug2024_update1 with the native +# pair_alignn style linked against the conda-env libtorch. +# +# Does: +# 1. Discovers PyTorch/CUDA/ABI in the active conda env +# 2. (Optional) installs matching cuda-toolkit + MKL if missing +# 3. Clones LAMMPS to $LAMMPS_DIR (default ~/lammps-alignn) +# 4. Copies pair_alignn.{cpp,h} from this repo into LAMMPS src/ +# 5. Appends find_package(Torch) + target_link_libraries to main CMakeLists.txt +# 6. Configures + builds with Ninja +# 7. Installs the lammps Python module into the active env +# 8. Overwrites $CONDA_PREFIX/lib/liblammps.so.0 with the pair_alignn build +# (fixes the standalone `lmp` binary's RPATH picking up a stale .so) +# 9. Verifies "alignn" appears in both Python module and the standalone binary +# +# Env vars: +# LAMMPS_DIR (default: $HOME/lammps-alignn) +# LAMMPS_TAG (default: stable_29Aug2024_update1) +# SKIP_CUDA_MKL set to 1 to skip the mamba install step +# JOBS (default: $(nproc)) +# --------------------------------------------------------------------------- set -euo pipefail -LAMMPS_TAG="stable_29Aug2024_update1" -LAMMPS_DIR="$HOME/lammps-alignn" -ALIGNN_REPO="/home/kamalch/Software/ollama311/alignn" -PAIR_SRC="$ALIGNN_REPO/alignn/scripts/torch/pair_alignn" +LAMMPS_TAG="${LAMMPS_TAG:-stable_29Aug2024_update1}" +LAMMPS_DIR="${LAMMPS_DIR:-$HOME/lammps-alignn}" +JOBS="${JOBS:-$(nproc)}" + +# Locate the alignn repo root from *this* script's path (so it's portable). +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PAIR_SRC="$SCRIPT_DIR/pair_alignn" +ALIGNN_REPO="$(cd "$SCRIPT_DIR/../../.." && pwd)" -# ── discover conda env's libtorch ───────────────────────────────────────── +echo "── paths ────────────────────────────────" +echo "script dir : $SCRIPT_DIR" +echo "alignn repo : $ALIGNN_REPO" +echo "pair sources : $PAIR_SRC" +echo "LAMMPS target: $LAMMPS_DIR (tag $LAMMPS_TAG)" +echo + +# ── sanity checks ──────────────────────────────────────────────────────── +if [ -z "${CONDA_PREFIX:-}" ]; then + echo "ERROR: CONDA_PREFIX not set. Activate a conda env first." >&2 + exit 1 +fi +if ! command -v cmake >/dev/null; then + echo "ERROR: cmake not found. `mamba install cmake` or install it." >&2 + exit 1 +fi + +# ── discover torch + CUDA ABI ──────────────────────────────────────────── TORCH_DIR=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))") TORCH_CMAKE="$TORCH_DIR/share/cmake/Torch" TORCH_CXX11_ABI=$(python -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))") +TORCH_CUDA=$(python -c "import torch; print(torch.version.cuda or 'cpu')") echo "── environment ──────────────────────────" -echo "TORCH_DIR = $TORCH_DIR" -echo "TORCH_CXX11_ABI = $TORCH_CXX11_ABI" -echo "python = $(which python)" -echo "cmake = $(cmake --version | head -1)" +echo "python = $(which python)" +echo "torch = $(python -c 'import torch; print(torch.__version__)')" +echo "torch CUDA = $TORCH_CUDA" +echo "torch CXX11 ABI = $TORCH_CXX11_ABI" +echo "cmake = $(cmake --version | head -1)" +echo "CONDA_PREFIX = $CONDA_PREFIX" echo +# ── ensure CUDA toolkit + MKL match libtorch ───────────────────────────── +need_install=0 +if [ "${SKIP_CUDA_MKL:-0}" != "1" ] && [ "$TORCH_CUDA" != "cpu" ]; then + # CUDA toolkit check: does nvcc match torch's CUDA? + if command -v nvcc >/dev/null; then + NVCC_VER=$(nvcc --version | awk '/release/{gsub(",","",$5); print $5}') + else + NVCC_VER="none" + fi + if [ "$NVCC_VER" != "$TORCH_CUDA" ]; then + echo "[!] nvcc=$NVCC_VER does not match torch CUDA=$TORCH_CUDA" + need_install=1 + fi + # MKL header check + if [ ! -f "$CONDA_PREFIX/include/mkl.h" ]; then + echo "[!] MKL headers not found (mkl.h missing)" + need_install=1 + fi + if [ $need_install -eq 1 ]; then + echo "── installing cuda-toolkit=$TORCH_CUDA + mkl-devel via mamba ──" + mamba install -y -c nvidia -c conda-forge \ + "cuda-toolkit=$TORCH_CUDA" mkl-devel mkl-include + else + echo "CUDA toolkit + MKL already satisfy torch requirements." + fi +fi + # ── clone LAMMPS ───────────────────────────────────────────────────────── if [ ! -d "$LAMMPS_DIR" ]; then echo "[1/5] cloning LAMMPS $LAMMPS_TAG..." @@ -31,62 +102,28 @@ else echo "[1/5] LAMMPS dir exists at $LAMMPS_DIR, reusing" fi -# ── install pair_alignn as USER-ALIGNN package ─────────────────────────── -echo "[2/5] installing pair_alignn into USER-ALIGNN..." -PKG_DIR="$LAMMPS_DIR/src/USER-ALIGNN" -mkdir -p "$PKG_DIR" -cp "$PAIR_SRC/pair_alignn.cpp" "$PKG_DIR/" -cp "$PAIR_SRC/pair_alignn.h" "$PKG_DIR/" - -cat > "$PKG_DIR/Install.sh" <<'SHEOF' -#!/bin/sh -action() { - if test -e ../$1 && test ! -e ../$1.bak ; then - mv ../$1 ../$1.bak - fi - if test ! -e $1 ; then return ; fi - cp $1 ../$1 -} -if (test $1 = 1) ; then - action pair_alignn.cpp - action pair_alignn.h -elif (test $1 = 0) ; then - for f in pair_alignn.cpp pair_alignn.h ; do - if test -e ../$f ; then rm -f ../$f ; fi - if test -e ../${f}.bak ; then mv ../${f}.bak ../$f ; fi - done -fi -SHEOF -chmod +x "$PKG_DIR/Install.sh" +# ── drop pair_alignn into src/ (no USER- package machinery) ────────────── +echo "[2/5] installing pair_alignn.{cpp,h} into LAMMPS src/..." +cp "$PAIR_SRC/pair_alignn.cpp" "$LAMMPS_DIR/src/pair_alignn.cpp" +cp "$PAIR_SRC/pair_alignn.h" "$LAMMPS_DIR/src/pair_alignn.h" -# ── patch LAMMPS cmake to register USER-ALIGNN ─────────────────────────── -CMAKE_EXTRA="$LAMMPS_DIR/cmake/Modules/Packages/USER-ALIGNN.cmake" -mkdir -p "$(dirname "$CMAKE_EXTRA")" -cat > "$CMAKE_EXTRA" <<'CMEOF' -find_package(Torch REQUIRED) +# ── append unconditional libtorch link to the main CMakeLists.txt ──────── +MAIN_CMAKE="$LAMMPS_DIR/cmake/CMakeLists.txt" +HOOK_MARKER="# ALIGNN-FF libtorch hook" +if ! grep -q "$HOOK_MARKER" "$MAIN_CMAKE"; then + cat >> "$MAIN_CMAKE" <&1 | tail -40 + -DCMAKE_CUDA_COMPILER=$(command -v nvcc 2>/dev/null || echo "") \ + 2>&1 | tail -8 echo -echo "[4/5] building (this takes a while)..." -cmake --build . -j"$(nproc)" 2>&1 | tail -20 +echo "[4/5] compiling (this takes 5–15 min)..." +cmake --build . -j"$JOBS" 2>&1 | tail -8 -# ── install the Python module ──────────────────────────────────────────── -echo "[5/5] installing lammps Python module..." -cmake --build . --target install-python 2>&1 | tail -10 +# ── install the lammps Python module ───────────────────────────────────── +echo "[5/5] installing lammps Python module into $CONDA_PREFIX..." +cmake --build . --target install-python 2>&1 | tail -5 + +# ── fix standalone binary's RPATH issue by overwriting conda liblammps ─── +if [ -f "$CONDA_PREFIX/lib/liblammps.so.0" ]; then + echo " syncing liblammps.so.0 → $CONDA_PREFIX/lib/ (for standalone lmp binary)" + cp "$BUILD_DIR/liblammps.so.0" "$CONDA_PREFIX/lib/liblammps.so.0" +fi # ── verify ─────────────────────────────────────────────────────────────── echo echo "── verification ─────────────────────────" -python - <&1 | grep -q alignn; then + echo "✓ Standalone binary has pair_alignn: $BUILD_DIR/lmp" +else + echo "✗ Standalone binary does NOT have pair_alignn" +fi + +echo +echo "──────────────────────────────────────────" +echo "Done. Try:" +echo " export_torchscript.py --model-dir --out alignn_ff.pt" +echo " $BUILD_DIR/lmp -in " +echo "──────────────────────────────────────────" diff --git a/docs/usage/lammps-pair-style.md b/docs/usage/lammps-pair-style.md index 10d2802..eb1f15c 100644 --- a/docs/usage/lammps-pair-style.md +++ b/docs/usage/lammps-pair-style.md @@ -26,49 +26,115 @@ loop. 3. `pair_alignn.cpp` / `pair_alignn.h` compiled into that LAMMPS binary 4. A TorchScript `.pt` of your model exported via `export_torchscript.py` -## Building LAMMPS with `pair_alignn` +## Building LAMMPS with `pair_alignn` — from a fresh clone -The repo ships a self-contained build script -(`alignn/scripts/torch/build_lammps_alignn.sh`) that: +The repo ships `alignn/scripts/torch/build_lammps_alignn.sh`, a +self-contained script that handles the full build, including prerequisite +installs. -1. Clones LAMMPS `stable_29Aug2024_update1` -2. Drops `pair_alignn.cpp/.h` into `src/` -3. Appends a libtorch `find_package` to `cmake/CMakeLists.txt` -4. Configures + builds + installs the Python module - -Prerequisites in your conda env (one-time): +### Step 1 — clone and install alignn ```bash -# Matching CUDA toolkit for libtorch (check your PyTorch first) -python -c "import torch; print(torch.version.cuda)" -# If 12.8, install: -mamba install -y -c nvidia -c conda-forge cuda-toolkit=12.8 mkl-devel mkl-include +git clone https://github.com/atomgptlab/alignn.git +cd alignn +pip install -e . # installs train_alignn.py, export_torchscript.py, ... ``` -Then run the build script: +Make sure you do this inside an activated conda env (`$CONDA_PREFIX` must +be set). The build script will error out early otherwise. + +### Step 2 — run the build script ```bash bash alignn/scripts/torch/build_lammps_alignn.sh ``` -On success, you'll see: +That's it. The script: + +1. Discovers the active env's PyTorch + CUDA + libtorch ABI flag +2. **Auto-installs matching `cuda-toolkit` and `mkl-devel`** via + `mamba install -c nvidia -c conda-forge` if they're missing or the + version doesn't match libtorch (set `SKIP_CUDA_MKL=1` to disable) +3. Clones LAMMPS `stable_29Aug2024_update1` to `$LAMMPS_DIR` + (default `~/lammps-alignn`) +4. Drops `pair_alignn.cpp/.h` into LAMMPS's `src/` (picked up by the + default source glob — no package-system gymnastics) +5. Appends a `find_package(Torch)` + `target_link_libraries(... TORCH_LIBRARIES)` + block to `cmake/CMakeLists.txt` (idempotent via marker comment) +6. Configures with Ninja, compiles with all available cores +7. Runs `make install-python` to wire the libtorch-linked build into your + conda env's `lammps` Python module +8. **Overwrites `$CONDA_PREFIX/lib/liblammps.so.0`** with the new build + (the standalone `lmp` binary's RPATH picks the conda lib first, so + without this step the binary keeps using a stale .so) +9. Prints ✓/✗ verification for both the Python module and the standalone + binary + +Takes ~5–15 minutes depending on machine. Successful output ends with: +``` +✓ Python module has pair_alignn +✓ Standalone binary has pair_alignn: ~/lammps-alignn/build/lmp ``` -pair_alignn available: True + +**Env-var knobs:** + +| Variable | Default | Purpose | +|---|---|---| +| `LAMMPS_DIR` | `$HOME/lammps-alignn` | where to clone + build LAMMPS | +| `LAMMPS_TAG` | `stable_29Aug2024_update1` | LAMMPS release to clone | +| `SKIP_CUDA_MKL` | unset | set to `1` to skip auto-install of CUDA/MKL | +| `JOBS` | `$(nproc)` | parallel compile jobs | + +Example overrides: +```bash +LAMMPS_DIR=/opt/lammps-alignn JOBS=16 \ + bash alignn/scripts/torch/build_lammps_alignn.sh ``` -The built `lmp` executable lands at `~/lammps-alignn/build/lmp`, and the -conda `lammps` Python module is overwritten with the libtorch-linked -version. +### Step 3 — verify (script does this automatically; manual fallback) + +```bash +# Standalone binary +$LAMMPS_DIR/build/lmp -help | grep alignn +# → alignn (as a registered pair style) + +# Python module (what most scripts / validators use) +python -c "from lammps import lammps; l=lammps(); print('alignn' in l.available_styles('pair'))" +# → True +``` + +### Rebuild after editing `pair_alignn.cpp` + +If you modify the C++ source in the alignn repo, just re-run the build +script — the clone step is skipped automatically if `$LAMMPS_DIR` already +exists, and everything else runs from a clean `build/`: + +```bash +bash alignn/scripts/torch/build_lammps_alignn.sh +``` + +Or do it manually, skipping the full reconfigure: + +```bash +cp alignn/scripts/torch/pair_alignn/pair_alignn.{cpp,h} $LAMMPS_DIR/src/ +cd $LAMMPS_DIR/build +cmake --build . -j$(nproc) +cmake --build . --target install-python +cp liblammps.so.0 $CONDA_PREFIX/lib/liblammps.so.0 # keep standalone binary in sync +``` ### Common build issues -| Symptom | Cause | Fix | +These are all handled automatically by the build script, but listed here +for context if something goes wrong: + +| Symptom | Cause | What the script does (and manual fix if needed) | |---|---|---| -| `PyTorch: CUDA cannot be found` | CUDA toolkit version mismatches libtorch's build version | install matching `cuda-toolkit=X.Y` via conda | -| `Imported target "torch" includes non-existent path ".../MKL_INCLUDE_DIR-NOTFOUND"` | MKL not installed | `mamba install mkl-devel mkl-include` | -| `fatal error: cuda.h: No such file or directory` | CUDA headers not on include path | cmake flag: `-DCMAKE_CXX_FLAGS="-I$CONDA_PREFIX/targets/x86_64-linux/include"` | -| `Enabled packages: ` | cmake didn't wire pair_alignn in | verify `append find_package(Torch)` block lives at the end of `cmake/CMakeLists.txt` | -| Builds but `pair_alignn` not listed | `pair_alignn.cpp` not in `src/` | `cp pair_alignn.{cpp,h} ~/lammps-alignn/src/` | +| `PyTorch: CUDA cannot be found` | CUDA toolkit version mismatches libtorch | auto-installs matching `cuda-toolkit=` via mamba | +| `Imported target "torch" includes non-existent path ".../MKL_INCLUDE_DIR-NOTFOUND"` | MKL not installed | auto-installs `mkl-devel mkl-include` | +| `fatal error: cuda.h: No such file or directory` | CUDA headers not on include path | adds `-I$CONDA_PREFIX/targets/x86_64-linux/include` to CXX flags | +| Python module works but `lmp` binary fails with "Unrecognized pair style" | RPATH resolves conda's stale `liblammps.so.0` first | `cp liblammps.so.0 $CONDA_PREFIX/lib/` — done in step 8 | +| `Could NOT find CUDA` on a CPU-only machine | libtorch was built with CUDA | install a CPU-only libtorch separately, or use a matching CPU-only PyTorch in the env | ## Exporting the TorchScript model diff --git a/docs/usage/melt-quench-tutorial.md b/docs/usage/melt-quench-tutorial.md index 162ff15..e7fad4b 100644 --- a/docs/usage/melt-quench-tutorial.md +++ b/docs/usage/melt-quench-tutorial.md @@ -7,10 +7,19 @@ the quenched structure. Prereqs: -- `alignn` installed (`pip install -e .` from a clone) -- `lammps` Python package compiled with `pair_alignn` (see - [LAMMPS Pair Style](lammps-pair-style.md) for the build) -- Plotly (`pip install plotly`) +- Clone + install alignn: + ```bash + git clone https://github.com/atomgptlab/alignn.git + cd alignn && pip install -e . + ``` +- LAMMPS built with `pair_alignn` compiled in. It's a single command: + ```bash + bash alignn/scripts/torch/build_lammps_alignn.sh + ``` + The script auto-installs CUDA + MKL into your conda env if they're + missing. Full reference: + [Building LAMMPS with `pair_alignn`](lammps-pair-style.md#building-lammps-with-pair_alignn-from-a-fresh-clone). +- Plotly for the RDF plot: `pip install plotly`. Working directory for this tutorial: `MELT/` (create anywhere). From d3ddb07f05c84c641688a395fb06cc6fc81b8d89 Mon Sep 17 00:00:00 2001 From: user Date: Tue, 21 Apr 2026 11:46:14 -0400 Subject: [PATCH 26/34] Updated lammps --- alignn/scripts/torch/build_lammps_alignn.sh | 154 ++++++++++++++------ docs/usage/lammps-pair-style.md | 7 +- 2 files changed, 118 insertions(+), 43 deletions(-) diff --git a/alignn/scripts/torch/build_lammps_alignn.sh b/alignn/scripts/torch/build_lammps_alignn.sh index 32b34e6..f0f4753 100755 --- a/alignn/scripts/torch/build_lammps_alignn.sh +++ b/alignn/scripts/torch/build_lammps_alignn.sh @@ -13,8 +13,12 @@ # 5. Appends find_package(Torch) + target_link_libraries to main CMakeLists.txt # 6. Configures + builds with Ninja # 7. Installs the lammps Python module into the active env -# 8. Overwrites $CONDA_PREFIX/lib/liblammps.so.0 with the pair_alignn build -# (fixes the standalone `lmp` binary's RPATH picking up a stale .so) +# 8. Overwrites any pre-existing liblammps.so.0 in the Python prefix with the +# pair_alignn build (fixes the standalone `lmp` binary's RPATH picking up +# a stale .so). +# +# Works in a conda env (uses mamba for CUDA/MKL) OR in system Python like +# Colab (falls back to `pip install nvidia-cuda-*-cu12 mkl mkl-devel mkl-include`). # 9. Verifies "alignn" appears in both Python module and the standalone binary # # Env vars: @@ -41,15 +45,16 @@ echo "pair sources : $PAIR_SRC" echo "LAMMPS target: $LAMMPS_DIR (tag $LAMMPS_TAG)" echo -# ── sanity checks ──────────────────────────────────────────────────────── -if [ -z "${CONDA_PREFIX:-}" ]; then - echo "ERROR: CONDA_PREFIX not set. Activate a conda env first." >&2 - exit 1 -fi -if ! command -v cmake >/dev/null; then - echo "ERROR: cmake not found. `mamba install cmake` or install it." >&2 +# ── detect Python install prefix (conda env OR system python, e.g. Colab) ── +PY_PREFIX=$(python -c "import sys; print(sys.prefix)") +PREFIX="${CONDA_PREFIX:-$PY_PREFIX}" +HAS_MAMBA=0 +command -v mamba >/dev/null 2>&1 && HAS_MAMBA=1 +command -v cmake >/dev/null || { + echo "ERROR: cmake not found. Install it (apt / conda / pip) and rerun." >&2 exit 1 -fi +} +command -v ninja >/dev/null || pip install -q ninja # ── discover torch + CUDA ABI ──────────────────────────────────────────── TORCH_DIR=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))") @@ -63,33 +68,62 @@ echo "torch = $(python -c 'import torch; print(torch.__version__)')" echo "torch CUDA = $TORCH_CUDA" echo "torch CXX11 ABI = $TORCH_CXX11_ABI" echo "cmake = $(cmake --version | head -1)" -echo "CONDA_PREFIX = $CONDA_PREFIX" +echo "prefix = $PREFIX" +echo "has mamba = $HAS_MAMBA" echo # ── ensure CUDA toolkit + MKL match libtorch ───────────────────────────── -need_install=0 if [ "${SKIP_CUDA_MKL:-0}" != "1" ] && [ "$TORCH_CUDA" != "cpu" ]; then - # CUDA toolkit check: does nvcc match torch's CUDA? + # CUDA toolkit check: nvcc matches torch CUDA? + NVCC_VER="none" if command -v nvcc >/dev/null; then NVCC_VER=$(nvcc --version | awk '/release/{gsub(",","",$5); print $5}') - else - NVCC_VER="none" - fi - if [ "$NVCC_VER" != "$TORCH_CUDA" ]; then - echo "[!] nvcc=$NVCC_VER does not match torch CUDA=$TORCH_CUDA" - need_install=1 fi - # MKL header check - if [ ! -f "$CONDA_PREFIX/include/mkl.h" ]; then - echo "[!] MKL headers not found (mkl.h missing)" - need_install=1 - fi - if [ $need_install -eq 1 ]; then - echo "── installing cuda-toolkit=$TORCH_CUDA + mkl-devel via mamba ──" - mamba install -y -c nvidia -c conda-forge \ - "cuda-toolkit=$TORCH_CUDA" mkl-devel mkl-include + NEED_CUDA=0 + [ "$NVCC_VER" != "$TORCH_CUDA" ] && NEED_CUDA=1 + + # MKL headers check — look in common places + MKL_H="" + for d in "$PREFIX/include" "$PY_PREFIX/include" /usr/include /usr/local/include; do + [ -f "$d/mkl.h" ] && MKL_H="$d/mkl.h" && break + done + NEED_MKL=0 + [ -z "$MKL_H" ] && NEED_MKL=1 + + if [ $HAS_MAMBA -eq 1 ]; then + if [ $NEED_CUDA -eq 1 ] || [ $NEED_MKL -eq 1 ]; then + echo "── installing cuda-toolkit=$TORCH_CUDA + mkl-devel via mamba ──" + mamba install -y -c nvidia -c conda-forge \ + "cuda-toolkit=$TORCH_CUDA" mkl-devel mkl-include + fi else - echo "CUDA toolkit + MKL already satisfy torch requirements." + # Pip fallback (Colab / system python). Grab CUDA + MKL from pip wheels. + PIP_ARGS=() + if [ $NEED_CUDA -eq 1 ]; then + cu_short=$(echo "$TORCH_CUDA" | tr -d '.') + echo "── installing CUDA ${TORCH_CUDA} dev headers via pip (cu${cu_short}) ──" + PIP_ARGS+=(\ + "nvidia-cuda-nvcc-cu${cu_short%??}==${TORCH_CUDA}.*" \ + "nvidia-cuda-runtime-cu${cu_short%??}==${TORCH_CUDA}.*" \ + "nvidia-cuda-cupti-cu${cu_short%??}==${TORCH_CUDA}.*" \ + ) || true + fi + if [ $NEED_MKL -eq 1 ]; then + echo "── installing mkl via pip ──" + PIP_ARGS+=(mkl mkl-devel mkl-include) + fi + [ ${#PIP_ARGS[@]} -gt 0 ] && pip install -q "${PIP_ARGS[@]}" || true + + # Re-locate MKL headers after install + for d in "$PREFIX/include" "$PY_PREFIX/include" /usr/include /usr/local/include; do + [ -f "$d/mkl.h" ] && MKL_H="$d/mkl.h" && break + done + + # Colab ships nvcc in /usr/local/cuda; export for CMake's detection + if [ -x /usr/local/cuda/bin/nvcc ] && ! command -v nvcc >/dev/null; then + export PATH="/usr/local/cuda/bin:$PATH" + export CUDA_HOME=/usr/local/cuda + fi fi fi @@ -132,9 +166,20 @@ rm -rf "$BUILD_DIR" mkdir -p "$BUILD_DIR" cd "$BUILD_DIR" -CUDA_INC="$CONDA_PREFIX/targets/x86_64-linux/include" -CUDA_FLAG="" -[ -d "$CUDA_INC" ] && CUDA_FLAG="-I${CUDA_INC}" +# CUDA headers are in a few possible places; collect any that exist. +CUDA_FLAGS="" +for d in \ + "$PREFIX/targets/x86_64-linux/include" \ + "$PY_PREFIX/targets/x86_64-linux/include" \ + /usr/local/cuda/include \ + /usr/local/cuda/targets/x86_64-linux/include; do + [ -d "$d" ] && CUDA_FLAGS="$CUDA_FLAGS -I$d" +done +# MKL headers +MKL_INC="" +for d in "$PREFIX/include" "$PY_PREFIX/include" /usr/include /usr/local/include; do + [ -f "$d/mkl.h" ] && MKL_INC="$d" && break +done echo "[3/5] cmake configure..." cmake ../cmake \ @@ -142,9 +187,9 @@ cmake ../cmake \ -DCMAKE_BUILD_TYPE=Release \ -DBUILD_SHARED_LIBS=ON \ -DLAMMPS_EXCEPTIONS=ON \ - -DCMAKE_PREFIX_PATH="$TORCH_CMAKE;$CONDA_PREFIX" \ - -DCMAKE_CXX_FLAGS="-D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI} ${CUDA_FLAG}" \ - -DMKL_INCLUDE_DIR=$CONDA_PREFIX/include \ + -DCMAKE_PREFIX_PATH="$TORCH_CMAKE;$PREFIX;$PY_PREFIX" \ + -DCMAKE_CXX_FLAGS="-D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI} ${CUDA_FLAGS}" \ + ${MKL_INC:+-DMKL_INCLUDE_DIR=$MKL_INC} \ -DPython_EXECUTABLE="$(which python)" \ -DCMAKE_CUDA_COMPILER=$(command -v nvcc 2>/dev/null || echo "") \ 2>&1 | tail -8 @@ -154,15 +199,42 @@ echo "[4/5] compiling (this takes 5–15 min)..." cmake --build . -j"$JOBS" 2>&1 | tail -8 # ── install the lammps Python module ───────────────────────────────────── -echo "[5/5] installing lammps Python module into $CONDA_PREFIX..." -cmake --build . --target install-python 2>&1 | tail -5 +echo "[5/5] installing lammps Python module into $PREFIX..." +install_ok=1 +if cmake --build . --target install-python 2>&1 | tail -5; then + # `install-python` may report success but actually fail silently on Colab + # when its internal venv bootstrap breaks. Verify by importing. + python -c "import lammps" 2>/dev/null || install_ok=0 +else + install_ok=0 +fi -# ── fix standalone binary's RPATH issue by overwriting conda liblammps ─── -if [ -f "$CONDA_PREFIX/lib/liblammps.so.0" ]; then - echo " syncing liblammps.so.0 → $CONDA_PREFIX/lib/ (for standalone lmp binary)" - cp "$BUILD_DIR/liblammps.so.0" "$CONDA_PREFIX/lib/liblammps.so.0" +if [ $install_ok -eq 0 ]; then + echo + echo " install-python target failed (common on Colab / system Python)." + echo " Falling back to manual copy of python/lammps → site-packages..." + SITE_PKG=$(python -c "import site; print(site.getsitepackages()[0])") + # Uninstall any stale lammps wheel first + pip uninstall -y lammps 2>/dev/null || true + rm -rf "$SITE_PKG/lammps" + cp -r "$LAMMPS_DIR/python/lammps" "$SITE_PKG/lammps" + # The Python module locates liblammps.so via rpath / LD_LIBRARY_PATH. + # Put the .so next to the package so ctypes.CDLL finds it. + cp "$BUILD_DIR/liblammps.so.0" "$SITE_PKG/lammps/liblammps.so" + # Also place a copy in PY_PREFIX/lib so the standalone lmp binary's RPATH works + mkdir -p "$PY_PREFIX/lib" + cp "$BUILD_DIR/liblammps.so.0" "$PY_PREFIX/lib/liblammps.so.0" + echo " installed to $SITE_PKG/lammps" fi +# ── fix standalone binary's RPATH by overwriting any pre-existing liblammps ── +for lib_dir in "$PREFIX/lib" "$PY_PREFIX/lib" /usr/lib /usr/local/lib; do + if [ -f "$lib_dir/liblammps.so.0" ]; then + echo " syncing liblammps.so.0 → $lib_dir/ (for standalone lmp binary)" + cp "$BUILD_DIR/liblammps.so.0" "$lib_dir/liblammps.so.0" + fi +done + # ── verify ─────────────────────────────────────────────────────────────── echo echo "── verification ─────────────────────────" diff --git a/docs/usage/lammps-pair-style.md b/docs/usage/lammps-pair-style.md index eb1f15c..127a111 100644 --- a/docs/usage/lammps-pair-style.md +++ b/docs/usage/lammps-pair-style.md @@ -40,8 +40,11 @@ cd alignn pip install -e . # installs train_alignn.py, export_torchscript.py, ... ``` -Make sure you do this inside an activated conda env (`$CONDA_PREFIX` must -be set). The build script will error out early otherwise. +Works in both environments: +- **Conda env**: the build uses `mamba` to install CUDA + MKL when missing. +- **System Python (e.g. Google Colab)**: the build falls back to `pip install` for + `nvidia-cuda-*-cu12 mkl mkl-devel mkl-include`, and detects CUDA at + `/usr/local/cuda` automatically. ### Step 2 — run the build script From 0267c2fa4d41c36de249ea0e225e27563a2afbb8 Mon Sep 17 00:00:00 2001 From: user Date: Tue, 21 Apr 2026 23:50:15 -0400 Subject: [PATCH 27/34] bfloat16 --- alignn/scripts/torch/export_torchscript.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/alignn/scripts/torch/export_torchscript.py b/alignn/scripts/torch/export_torchscript.py index b7dedc8..fa73253 100644 --- a/alignn/scripts/torch/export_torchscript.py +++ b/alignn/scripts/torch/export_torchscript.py @@ -32,10 +32,19 @@ def main(): ap.add_argument("--out", default="alignn_ff.pt") ap.add_argument("--atom-features", default="atomic_number", help="must match what was used at training time") - ap.add_argument("--dtype", default="float32", choices=["float32", "float64"]) + ap.add_argument("--dtype", default="float32", + choices=["float32", "float64", "float16", "bfloat16"], + help="float32 is the safe default; bfloat16 halves memory " + "with minimal accuracy loss; float16 is risky for " + "MD forces (may NaN on small gradient magnitudes).") args = ap.parse_args() - dtype = torch.float32 if args.dtype == "float32" else torch.float64 + dtype_map = {"float32": torch.float32, "float64": torch.float64, + "float16": torch.float16, "bfloat16": torch.bfloat16} + dtype = dtype_map[args.dtype] + if args.dtype == "float16": + print("[warn] fp16 forces can NaN during long MD runs. " + "Consider --dtype bfloat16 instead.") cfg_raw = json.load(open(f"{args.model_dir}/config.json")) mcfg = dict(cfg_raw["model"]) From 44af8ff7961aa5d1fe92f5117f3ae481d1711e9f Mon Sep 17 00:00:00 2001 From: user Date: Tue, 21 Apr 2026 23:53:54 -0400 Subject: [PATCH 28/34] bfloat16 --- alignn/scripts/torch/export_torchscript.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/alignn/scripts/torch/export_torchscript.py b/alignn/scripts/torch/export_torchscript.py index fa73253..7c4f008 100644 --- a/alignn/scripts/torch/export_torchscript.py +++ b/alignn/scripts/torch/export_torchscript.py @@ -80,9 +80,15 @@ def main(): src = torch.tensor([0, 1], dtype=torch.long) dst = torch.tensor([1, 0], dtype=torch.long) shift = torch.zeros(2, 3, dtype=dtype) - out = scripted.forward_tensors_z(pos, lat, Z, src, dst, shift, True) + # torch.det() isn't implemented for bf16/fp16 on CPU, so skip stress + # in the smoke test for reduced-precision exports. It works fine on CUDA. + needs_fp32_det = dtype in (torch.bfloat16, torch.float16) + out = scripted.forward_tensors_z(pos, lat, Z, src, dst, shift, + not needs_fp32_det) print("smoke test OK:", {k: (v.shape if hasattr(v, 'shape') else v) for k, v in out.items()}) + if needs_fp32_det: + print(" (stress path skipped on CPU — bf16/fp16 det only works on CUDA)") torch.jit.save(scripted, args.out) print(f"wrote {args.out} ({sum(p.numel() for p in model.parameters()):,} params)") From 7a786acd7b3ac3b8a0412a6faeaf459d252ca235 Mon Sep 17 00:00:00 2001 From: knc6 Date: Sat, 25 Apr 2026 17:51:21 -0400 Subject: [PATCH 29/34] ff fic --- alignn/ff/calculators.py | 8 ++++---- alignn/train_alignn.py | 23 +++++++++++++++-------- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py index 0c3ca1c..2312ef0 100644 --- a/alignn/ff/calculators.py +++ b/alignn/ff/calculators.py @@ -269,9 +269,7 @@ def __init__( eALIGNNAtomWiseConfig(**self.config["model"]) ) if model is None: - raise ValueError( - f"Unsupported model name '{mname}' in config" - ) + raise ValueError(f"Unsupported model name '{mname}' in config") if "atomwise" in self.config["model"]["name"]: model.load_state_dict( torch.load( @@ -375,7 +373,8 @@ def calculate(self, atoms, properties=None, system_changes=None): # / 160.21766208 # ) if "atomwise" in self.config["model"]["name"]: - energy = result["out"].detach().cpu().numpy() + # energy = result["out"].squeeze().detach().cpu().numpy() + energy = result["out"].detach().cpu().item() else: energy = result.detach().cpu().numpy() if self.intensive: @@ -391,6 +390,7 @@ def calculate(self, atoms, properties=None, system_changes=None): "forces": forces, "stress": stress, } + # print("self.results",self.results) class iAlignnAtomwiseCalculator(ase.calculators.calculator.Calculator): diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py index b0d55aa..d95c533 100644 --- a/alignn/train_alignn.py +++ b/alignn/train_alignn.py @@ -1,6 +1,7 @@ #!/usr/bin/env python """Module to train for a folder with formatted dataset.""" + import os import torch.distributed as dist import csv @@ -13,6 +14,10 @@ from jarvis.db.jsonutils import loadjson import argparse from alignn.models.alignn_atomwise import ALIGNNAtomWise, ALIGNNAtomWiseConfig +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, + ALIGNNAtomWisePureConfig, +) import torch import time from jarvis.core.atoms import Atoms @@ -192,10 +197,7 @@ def train_for_folder( except Exception as exp: print("Check", exp) - if ( - config.gpu_memory_fraction is not None - and torch.cuda.is_available() - ): + if config.gpu_memory_fraction is not None and torch.cuda.is_available(): torch.cuda.set_per_process_memory_fraction( float(config.gpu_memory_fraction), rank if world_size > 1 else 0, @@ -345,15 +347,20 @@ def train_for_folder( if restart_model_path is not None: # Should be best_model.pt file print("Restarting the model training:", restart_model_path) - if config.model.name == "alignn_atomwise": + if "alignn_" in config.model.name: # == "alignn_atomwise": rest_config = loadjson( restart_model_path.replace("current_model.pt", "config.json") # restart_model_path.replace("best_model.pt", "config.json") ) - - tmp = ALIGNNAtomWiseConfig(**rest_config["model"]) + try: + tmp = ALIGNNAtomWiseConfig(**rest_config["model"]) + model = ALIGNNAtomWise(tmp) # config.model) + except: + tmp = ALIGNNAtomWisePureConfig(**rest_config["model"]) + model = ALIGNNAtomWisePure(tmp) # config.model) + pass + # from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePureConfig print("Rest config", tmp) - model = ALIGNNAtomWise(tmp) # config.model) print("model", model) model.load_state_dict( torch.load(restart_model_path, map_location=device) From e33cf9ec933b04210a2f35a31a6c5e1ba424b9c8 Mon Sep 17 00:00:00 2001 From: knc6 Date: Sun, 3 May 2026 19:52:51 -0400 Subject: [PATCH 30/34] DGL dependency out --- alignn/data.py | 48 +++++++++++++++- alignn/lmdb_dataset.py | 63 ++++++++++++++++++--- alignn/pure_lmdb_dataset.py | 107 ++++++++++++++++++++++++++++++++---- 3 files changed, 199 insertions(+), 19 deletions(-) diff --git a/alignn/data.py b/alignn/data.py index 68b0278..044a204 100644 --- a/alignn/data.py +++ b/alignn/data.py @@ -10,6 +10,7 @@ from tqdm import tqdm import math from jarvis.db.jsonutils import dumpjson + try: from dgl.dataloading import GraphDataLoader except ImportError: # pure-torch path; fall back to stdlib DataLoader @@ -18,6 +19,8 @@ def GraphDataLoader(*args, use_ddp=False, **kwargs): kwargs.pop("use_ddp", None) return _TorchDataLoader(*args, **kwargs) + + import pickle as pk from sklearn.preprocessing import StandardScaler @@ -478,7 +481,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( @@ -490,7 +506,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( @@ -504,10 +533,27 @@ def get_train_val_loaders( pin_memory=pin_memory, use_ddp=use_ddp, ) - if len(dataset_test) > 0 + if test_data is not None else None ) + # All loaders constructed; safe to free the in-memory dataset + # lists. With 1.5M-entry datasets these can hold tens of GB of + # full atomic-structure dicts, and they are no longer needed + # once the LMDB caches are built and the loaders reference + # them via the Dataset wrappers. + del dataset_train + del dataset_val + del dataset_test + del dat + try: + del d + except NameError: + pass + import gc + + gc.collect() + if save_dataloader: torch.save(train_loader, train_sample) if val_loader is not None: diff --git a/alignn/lmdb_dataset.py b/alignn/lmdb_dataset.py index 0d3611c..c4f86f8 100644 --- a/alignn/lmdb_dataset.py +++ b/alignn/lmdb_dataset.py @@ -13,12 +13,19 @@ import torch from tqdm import tqdm from typing import List, Tuple + try: import dgl except ImportError: dgl = None +# How often to commit the LMDB write transaction during a fresh build. +# Holding a single transaction across all entries keeps every dirty page +# resident in RAM until commit; flushing periodically caps peak memory. +LMDB_COMMIT_EVERY = 5000 + + def prepare_line_graph_batch( batch: Tuple[Tuple[dgl.DGLGraph, dgl.DGLGraph], torch.Tensor], device=None, @@ -113,6 +120,16 @@ def collate_line_graph( ) +def _mem_gb(): + """Resident memory of the current process in GB (for diagnostics).""" + try: + import psutil + + return psutil.Process(os.getpid()).memory_info().rss / 1e9 + except Exception: + return float("nan") + + def get_torch_dataset( dataset=[], id_tag="jid", @@ -148,6 +165,9 @@ def get_torch_dataset( line = "Min=" + str(np.min(vals)) + "\n" f.write(line) f.close() + + print(f"[MEM] {tmp_name}: entry: {_mem_gb():.2f} GB") + ids = [] if os.path.exists(tmp_name) and read_existing: # Validate the cache backend matches (DGL). A previous run with @@ -174,7 +194,12 @@ def get_torch_dataset( lmdb_path=tmp_name, line_graph=line_graph, ids=ids ) print("Reading dataset", tmp_name) + print( + f"[MEM] {tmp_name}: after TorchLMDBDataset (cached): " + f"{_mem_gb():.2f} GB" + ) return dat + ids = [] # Fresh build: wipe any pre-existing cache dir to avoid mixing old # records (possibly from a different model backend) with new writes. @@ -183,10 +208,16 @@ def get_torch_dataset( shutil.rmtree(tmp_name) env = lmdb.open(tmp_name, map_size=int(map_size)) - with env.begin(write=True) as txn: + + print(f"[MEM] {tmp_name}: before write loop: {_mem_gb():.2f} GB") + + # Use periodic commits instead of one giant transaction. A single + # write txn around 1.5M put() calls keeps every dirty page in RAM + # until commit and was the cause of OOM on 62 GB hosts. + txn = env.begin(write=True) + try: for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)): ids.append(d[id_tag]) - # g, lg = Graph.atom_dgl_multigraph( atoms = Atoms.from_dict(d["atoms"]) g = Graph.atom_dgl_multigraph( atoms, @@ -207,10 +238,8 @@ def get_torch_dataset( ) label = torch.tensor(d[target]).type(torch.get_default_dtype()) natoms = len(d["atoms"]["elements"]) - # print('label',label,label.view(-1).long()) if classification: label = label.long() - # label = label.view(-1).long() if "extra_features" in d: g.ndata["extra_features"] = torch.tensor( [d["extra_features"] for n in range(natoms)] @@ -220,8 +249,6 @@ def get_torch_dataset( np.array(d[target_atomwise]) ).type(torch.get_default_dtype()) if target_grad is not None and target_grad != "": - # print('grad', np.array(d[target_grad])) - # print('grad shape',np.array(d[target_grad]).shape) arr = np.array(d[target_grad]) try: g.ndata[target_grad] = torch.tensor(arr).type( @@ -232,7 +259,6 @@ def get_torch_dataset( g.ndata[target_grad] = torch.tensor(arr).type( torch.get_default_dtype() ) - # print('arr',arr.shape) if target_stress is not None and target_stress != "": stress = np.array(d[target_stress]) g.ndata[target_stress] = torch.tensor( @@ -247,17 +273,38 @@ def get_torch_dataset( ([additional_output for ii in range(natoms)]) ).type(torch.get_default_dtype()) - # labels.append(label) if line_graph: serialized_data = pk.dumps((g, lg, lattice, label)) else: serialized_data = pk.dumps((g, lattice, label)) txn.put(f"{idx}".encode(), serialized_data) + # Flush dirty pages to disk every LMDB_COMMIT_EVERY entries. + if (idx + 1) % LMDB_COMMIT_EVERY == 0: + txn.commit() + txn = env.begin(write=True) + if (idx + 1) % (LMDB_COMMIT_EVERY * 20) == 0: + print( + f"[MEM] {tmp_name}: after {idx + 1} entries: " + f"{_mem_gb():.2f} GB" + ) + finally: + # Commit whatever is left in the final partial batch. + txn.commit() + env.close() + + print(f"[MEM] {tmp_name}: after write loop: {_mem_gb():.2f} GB") + lmdb_dataset = TorchLMDBDataset( lmdb_path=tmp_name, line_graph=line_graph, ids=ids ) + + print( + f"[MEM] {tmp_name}: after TorchLMDBDataset wrap: " + f"{_mem_gb():.2f} GB" + ) + return lmdb_dataset diff --git a/alignn/pure_lmdb_dataset.py b/alignn/pure_lmdb_dataset.py index b361a3a..35e2159 100644 --- a/alignn/pure_lmdb_dataset.py +++ b/alignn/pure_lmdb_dataset.py @@ -26,6 +26,21 @@ torchgraph_from_dgl, ) +# How often to commit the LMDB write transaction during a fresh build. +# Holding a single transaction across all entries keeps every dirty page +# resident in RAM until commit; flushing periodically caps peak memory. +LMDB_COMMIT_EVERY = 100000 + + +def _mem_gb(): + """Resident memory of the current process in GB (for diagnostics).""" + try: + import psutil + + return psutil.Process(os.getpid()).memory_info().rss / 1e9 + except Exception: + return float("nan") + def prepare_pure_batch(batch, device=None, non_blocking=False): """Move a pure-torch batch to device; mirrors prepare_line_graph_batch.""" @@ -48,9 +63,7 @@ def _graph_to_torchgraph(g, lg=None): class PureTorchLMDBDataset(Dataset): """Dataset of crystal TorchGraphs using LMDB.""" - def __init__( - self, lmdb_path: str = "", line_graph: bool = True, ids=None - ): + def __init__(self, lmdb_path: str = "", line_graph: bool = True, ids=None): """Open the LMDB at ``lmdb_path`` read-only.""" super().__init__() self.lmdb_path = lmdb_path @@ -124,6 +137,26 @@ def _attach_node_payload( g.ndata[key] = torch.as_tensor(tiled, dtype=dtype) +def _lmdb_is_usable(tmp_name: str) -> bool: + """Return True if the LMDB at ``tmp_name`` contains at least one entry. + + Guards against the empty-stub left behind when a previous build was + interrupted (OOM, Ctrl-C, etc.). Such stubs have an 8KB data.mdb but + zero records, and the read_existing fast path would silently return + a zero-length dataset if not caught. + """ + if not os.path.exists(tmp_name): + return False + try: + env = lmdb.open(tmp_name, readonly=True, lock=False) + with env.begin() as txn: + n_entries = txn.stat()["entries"] + env.close() + return n_entries > 0 + except Exception: + return False + + def get_torch_dataset( dataset=None, id_tag="jid", @@ -158,7 +191,12 @@ def get_torch_dataset( with open(os.path.join(output_dir, tmp_name + "_data_range"), "w") as f: f.write(f"Max={np.max(vals)}\nMin={np.min(vals)}\n") - if os.path.exists(tmp_name) and read_existing: + print(f"[MEM] {tmp_name}: entry: {_mem_gb():.2f} GB") + + # Fast path: reuse a previously built cache, but only if it actually + # contains records. An empty stub (8KB data.mdb) from an interrupted + # build would otherwise be returned as a zero-length dataset. + if read_existing and _lmdb_is_usable(tmp_name): # Validate the cache was built for the pure-torch backend. _env = lmdb.open(tmp_name, readonly=True, lock=False) with _env.begin() as _txn: @@ -176,9 +214,22 @@ def get_torch_dataset( ) ids = [d[id_tag] for d in dataset] print("Reading dataset", tmp_name) - return PureTorchLMDBDataset( + ds = PureTorchLMDBDataset( lmdb_path=tmp_name, line_graph=line_graph, ids=ids ) + print( + f"[MEM] {tmp_name}: after PureTorchLMDBDataset (cached): " + f"{_mem_gb():.2f} GB" + ) + return ds + + # If read_existing was requested but the cache is unusable (empty + # stub or missing), warn so the user knows we're rebuilding. + if read_existing and os.path.exists(tmp_name): + print( + f"[WARN] {tmp_name}: read_existing=True but cache is " + "empty or unreadable; rebuilding from scratch." + ) ids = [] # Fresh build: wipe any pre-existing cache dir to avoid mixing old @@ -188,7 +239,14 @@ def get_torch_dataset( shutil.rmtree(tmp_name) env = lmdb.open(tmp_name, map_size=int(map_size)) - with env.begin(write=True) as txn: + + print(f"[MEM] {tmp_name}: before write loop: {_mem_gb():.2f} GB") + + # Use periodic commits instead of a single giant write transaction. + # A single txn around millions of put() calls keeps every dirty page + # in RAM until commit and was the cause of OOM on memory-tight hosts. + txn = env.begin(write=True) + try: for idx, d in tqdm(enumerate(dataset), total=len(dataset)): ids.append(d[id_tag]) atoms = Atoms.from_dict(d["atoms"]) @@ -214,9 +272,7 @@ def get_torch_dataset( lattice = torch.as_tensor( atoms.lattice_mat, dtype=torch.get_default_dtype() ) - label = torch.as_tensor( - d[target], dtype=torch.get_default_dtype() - ) + label = torch.as_tensor(d[target], dtype=torch.get_default_dtype()) natoms = len(d["atoms"]["elements"]) if classification: label = label.long() @@ -252,7 +308,38 @@ def get_torch_dataset( txn.put(f"{idx}".encode(), pk.dumps((g, lg, lattice, label))) else: txn.put(f"{idx}".encode(), pk.dumps((g, lattice, label))) + + # Drop this entry's full atomic-structure dict now that it's + # serialized into LMDB. The caller's list survives but its + # contents are released, freeing ~10-30 KB per entry. For + # 1.5M-entry datasets this is the difference between sitting + # at 24+ GB during the loop and staying at 3-4 GB. + dataset[idx] = None + + # Flush dirty pages to disk every LMDB_COMMIT_EVERY entries. + if (idx + 1) % LMDB_COMMIT_EVERY == 0: + txn.commit() + txn = env.begin(write=True) + if (idx + 1) % (LMDB_COMMIT_EVERY * 5) == 0: + print( + f"[MEM] {tmp_name}: after {idx + 1} entries: " + f"{_mem_gb():.2f} GB" + ) + finally: + # Commit whatever is left in the final partial batch. + txn.commit() + env.close() - return PureTorchLMDBDataset( + + print(f"[MEM] {tmp_name}: after write loop: {_mem_gb():.2f} GB") + + ds = PureTorchLMDBDataset( lmdb_path=tmp_name, line_graph=line_graph, ids=ids ) + + print( + f"[MEM] {tmp_name}: after PureTorchLMDBDataset wrap: " + f"{_mem_gb():.2f} GB" + ) + + return ds From d9a43de1a3beab92528bafcd7745f10318651982 Mon Sep 17 00:00:00 2001 From: knc6 Date: Sun, 3 May 2026 19:01:32 -0600 Subject: [PATCH 31/34] No dgl --- alignn/data.py | 2 +- alignn/graphs.py | 5 +- alignn/models/alignn.py | 3 +- alignn/models/alignn_atomwise.py | 4 +- alignn/models/ealignn_atomwise.py | 3 +- alignn/train.py | 201 +++++++----------------------- 6 files changed, 57 insertions(+), 161 deletions(-) diff --git a/alignn/data.py b/alignn/data.py index 044a204..00aad93 100644 --- a/alignn/data.py +++ b/alignn/data.py @@ -13,7 +13,7 @@ try: from dgl.dataloading import GraphDataLoader -except ImportError: # pure-torch path; fall back to stdlib DataLoader +except: # pure-torch path; fall back to stdlib DataLoader from torch.utils.data import DataLoader as _TorchDataLoader def GraphDataLoader(*args, use_ddp=False, **kwargs): diff --git a/alignn/graphs.py b/alignn/graphs.py index 2da69ce..144ba97 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -14,15 +14,18 @@ from collections import defaultdict from typing import List, Tuple, Sequence, Optional import torch + try: from dgl.data import DGLDataset import dgl -except ImportError: # DGL optional; DGL-dependent helpers will fail lazily +except: # DGL optional; DGL-dependent helpers will fail lazily dgl = None class DGLDataset: # minimal stub so subclasses can still be defined def __init__(self, *args, **kwargs): pass + + from tqdm import tqdm from jarvis.core.atoms import Atoms diff --git a/alignn/models/alignn.py b/alignn/models/alignn.py index 9862ba3..f58b8ed 100644 --- a/alignn/models/alignn.py +++ b/alignn/models/alignn.py @@ -6,11 +6,12 @@ from __future__ import annotations from typing import Tuple, Union + try: import dgl import dgl.function as fn from dgl.nn import AvgPooling -except ImportError: +except: dgl = None fn = None AvgPooling = None diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py index 7c743e0..6c9952e 100644 --- a/alignn/models/alignn_atomwise.py +++ b/alignn/models/alignn_atomwise.py @@ -7,14 +7,16 @@ from typing import Tuple, Union from torch.autograd import grad + try: import dgl import dgl.function as fn from dgl.nn import AvgPooling -except ImportError: # DGL optional; this module only usable if DGL installed +except: # DGL optional; this module only usable if DGL installed dgl = None fn = None AvgPooling = None + print("WARNING: No DGL, might fail") import numpy as np import torch diff --git a/alignn/models/ealignn_atomwise.py b/alignn/models/ealignn_atomwise.py index 94c6533..6f2bf5a 100644 --- a/alignn/models/ealignn_atomwise.py +++ b/alignn/models/ealignn_atomwise.py @@ -7,11 +7,12 @@ from typing import Tuple, Union from torch.autograd import grad + try: import dgl import dgl.function as fn from dgl.nn import AvgPooling -except ImportError: +except: dgl = None fn = None AvgPooling = None diff --git a/alignn/train.py b/alignn/train.py index aba5931..ad2bc70 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -16,6 +16,8 @@ def _ddp_mean(value: float, use_ddp: bool) -> float: def _unwrap(net): """Return underlying module from a DDP-wrapped model, else net itself.""" return net.module if isinstance(net, DDP) else net + + from functools import partial from typing import Any, Dict, Union import torch @@ -39,16 +41,11 @@ def _unwrap(net): import time from sklearn.metrics import roc_auc_score from alignn.utils import ( - # activated_output_transform, - # make_standard_scalar_and_pca, - # thresholded_output_transform, group_decay, setup_optimizer, print_train_val_loss, ) -# from sklearn.metrics import log_loss - warnings.filterwarnings("ignore", category=RuntimeWarning) # torch.autograd.detect_anomaly() @@ -65,21 +62,16 @@ def _unwrap(net): def train_dgl( config: Union[TrainingConfig, Dict[str, Any]], model: nn.Module = None, - # checkpoint_dir: Path = Path("./"), train_val_test_loaders=[], rank=0, world_size=0, - # log_tensorboard: bool = False, ): """Training entry point for DGL networks. `config` should conform to alignn.conf.TrainingConfig, and if passed as a dict with matching keys, pydantic validation is used """ - # print("rank", rank) - # setup(rank, world_size) if rank == 0: - # print(config) if type(config) is dict: try: print("Trying to convert dictionary.") @@ -90,8 +82,6 @@ def train_dgl( if not os.path.exists(config.output_dir): os.makedirs(config.output_dir) - # checkpoint_dir = os.path.join(config.output_dir) - # deterministic = False classification = False is_main = rank == 0 tmp = config.dict() @@ -101,7 +91,7 @@ def train_dgl( f.close() global tmp_output_dir tmp_output_dir = config.output_dir - pprint.pprint(tmp) # , sort_dicts=False) + pprint.pprint(tmp) if config.classification_threshold is not None: classification = True TORCH_DTYPES = { @@ -113,7 +103,6 @@ def train_dgl( torch.set_default_dtype(TORCH_DTYPES[config.dtype]) line_graph = False if config.compute_line_graph > 0: - # if config.model.alignn_layers > 0: line_graph = True if world_size > 1: use_ddp = True @@ -123,10 +112,6 @@ def train_dgl( if torch.cuda.is_available(): device = torch.device("cuda") if not train_val_test_loaders: - # use input standardization for all real-valued feature sets - # print("config.neighbor_strategy",config.neighbor_strategy) - # import sys - # sys.exit() ( train_loader, val_loader, @@ -170,7 +155,6 @@ def train_dgl( val_loader = train_val_test_loaders[1] test_loader = train_val_test_loaders[2] prepare_batch = train_val_test_loaders[3] - # rank=0 if use_ddp: device = torch.device(f"cuda:{rank}") prepare_batch = partial(prepare_batch, device=device) @@ -225,7 +209,6 @@ def train_dgl( print(statistics) except Exception: pass - # print("device", device) net.to(device) if use_ddp: net = DDP( @@ -235,48 +218,50 @@ def train_dgl( getattr(config, "ddp_find_unused_parameters", False) ), ) - # group parameters to skip weight decay for bias and batchnorm + # ------------------------------------------------------------------ + # Build optimizer ONCE and bind the scheduler to it. + # (Previously this block created the optimizer twice — once here and + # once again inside the `if "alignn_" in config.model.name:` block — + # which left the scheduler bound to a stale optimizer that never + # stepped, triggering the + # "lr_scheduler.step() before optimizer.step()" warning and producing + # a broken LR schedule.) + # ------------------------------------------------------------------ params = group_decay(net) optimizer = setup_optimizer(params, config) + if config.scheduler == "none": # always return multiplier of 1 (i.e. do nothing) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lambda epoch: 1.0 ) - elif config.scheduler == "onecycle": steps_per_epoch = len(train_loader) - # pct_start = config.warmup_steps / (config.epochs * steps_per_epoch) scheduler = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=config.learning_rate, epochs=config.epochs, steps_per_epoch=steps_per_epoch, - # pct_start=pct_start, pct_start=0.3, ) elif config.scheduler == "step": - # pct_start = config.warmup_steps / (config.epochs * steps_per_epoch) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, ) - # if ( - # config.model.name == "alignn_atomwise" - # or config.model.name == "alignn_ff2" - # ): + # OneCycleLR expects scheduler.step() per batch; other schedulers we + # step per epoch. + _step_scheduler_per_batch = config.scheduler == "onecycle" + if "alignn_" in config.model.name: best_loss = np.inf criterion = nn.L1Loss() if classification: criterion = nn.NLLLoss() - params = group_decay(net) - optimizer = setup_optimizer(params, config) - # optimizer = torch.optim.Adam(net.parameters(), lr=0.001) + # NOTE: optimizer / scheduler intentionally NOT recreated here. history_train = [] history_val = [] for e in range(config.epochs): - # optimizer.zero_grad() train_init_time = time.time() running_loss = 0 running_loss1 = 0 @@ -287,19 +272,15 @@ def train_dgl( train_result = [] for dats, jid in zip(train_loader, train_loader.dataset.ids): info = {} - # info["id"] = jid optimizer.zero_grad() _amp_ctx = torch.autocast( device_type="cuda", dtype=torch.bfloat16, - enabled=bool( - getattr(config, "use_amp", False) - ) + enabled=bool(getattr(config, "use_amp", False)) and torch.cuda.is_available(), ) with _amp_ctx: if (config.compute_line_graph) > 0: - # if (config.model.alignn_layers) > 0: result = net( [ dats[0].to(device), @@ -307,12 +288,8 @@ def train_dgl( dats[2].to(device), ] ) - else: - result = net( - [dats[0].to(device), dats[1].to(device)] - ) - # info = {} + result = net([dats[0].to(device), dats[1].to(device)]) info["target_out"] = [] info["pred_out"] = [] info["target_atomwise_pred"] = [] @@ -330,23 +307,17 @@ def train_dgl( loss4 = 0 # Such as stresses loss5 = 0 # Such as dos if config.model.output_features is not None: - # print('criterion',criterion) - # print('result["out"]',result["out"]) - # print('dats[-1]',dats[-1]) loss1 = config.model.graphwise_weight * criterion( result["out"], dats[-1].to(device), - # result["out"], dats[2].to(device) ) info["target_out"] = dats[-1].cpu().numpy().tolist() - # info["target_out"] = dats[2].cpu().numpy().tolist() info["pred_out"] = ( result["out"].cpu().detach().numpy().tolist() ) running_loss1 += loss1.item() if ( config.model.atomwise_output_features > 0 - # config.model.atomwise_output_features is not None and config.model.atomwise_weight != 0 ): loss2 = config.model.atomwise_weight * criterion( @@ -374,8 +345,6 @@ def train_dgl( ) running_loss3 += loss3.item() if config.model.stresswise_weight != 0: - # print('unbatch',dgl.unbatch(dats[0])) - targ_stress = torch.stack( [ gg.ndata["stresses"][0] @@ -383,55 +352,40 @@ def train_dgl( ] ).to(device) pred_stress = result["stresses"] - # print('targ_stress',targ_stress,targ_stress.shape) - # print('pred_stress',pred_stress,pred_stress.shape) loss4 = config.model.stresswise_weight * criterion( pred_stress.to(device), targ_stress.to(device), ) - info["target_stress"] = ( - targ_stress.cpu() - .numpy() - .tolist() - # dats[0].ndata["stresses"][0].cpu().numpy().tolist() - ) + info["target_stress"] = targ_stress.cpu().numpy().tolist() info["pred_stress"] = ( result["stresses"].cpu().detach().numpy().tolist() ) running_loss4 += loss4.item() if config.model.additional_output_weight != 0: - # print('unbatch',dgl.unbatch(dats[0])) additional_dat = [ gg.ndata["additional"][0] for gg in _graph_unbatch(dats[0]) ] - # print('additional_dat',additional_dat,len(additional_dat)) targ = torch.stack(additional_dat).to(device) - # targ=torch.tensor(additional_dat).to( dats[0].device) - # print('result["additional"]',result["additional"],result["additional"].shape) - # print('targ',targ,targ.shape) - # print('targ device',targ.device) loss5 = config.model.additional_output_weight * criterion( (result["additional"]).to(device), targ, - # (dats[0].ndata["additional"]).to(device), ) info["target_additional"] = targ.cpu().numpy().tolist() info["pred_additional"] = ( result["additional"].cpu().detach().numpy().tolist() ) running_loss5 += loss5.item() - # print("target_stress", info["target_stress"][0]) - # print("pred_stress", info["pred_stress"][0]) train_result.append(info) loss = loss1 + loss2 + loss3 + loss4 + loss5 loss.backward() optimizer.step() - # optimizer.zero_grad() #never + # Step OneCycleLR per batch (its design assumption); other + # schedulers are stepped once per epoch below. + if _step_scheduler_per_batch: + scheduler.step() running_loss += loss.item() - # Normalize running losses by number of batches so that printed - # values are per-batch mean losses (comparable across runs / - # dataset sizes), not raw sums. + # Normalize running losses by number of batches _n_tr = max(1, len(train_loader)) running_loss /= _n_tr running_loss1 /= _n_tr @@ -439,22 +393,17 @@ def train_dgl( running_loss3 /= _n_tr running_loss4 /= _n_tr running_loss5 /= _n_tr - # Average across ranks so printed values reflect the whole - # global batch, not a single rank's shard. running_loss = _ddp_mean(running_loss, use_ddp) running_loss1 = _ddp_mean(running_loss1, use_ddp) running_loss2 = _ddp_mean(running_loss2, use_ddp) running_loss3 = _ddp_mean(running_loss3, use_ddp) running_loss4 = _ddp_mean(running_loss4, use_ddp) running_loss5 = _ddp_mean(running_loss5, use_ddp) - # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors( - # train_result - # ) - # dumpjson(filename="Train_results.json", data=train_result) - scheduler.step() + # Epoch-level scheduler step for non-OneCycle schedulers. + if not _step_scheduler_per_batch: + scheduler.step() train_final_time = time.time() train_ep_time = train_final_time - train_init_time - # if rank == 0: # or world_size == 1: history_train.append( [ running_loss, @@ -479,15 +428,11 @@ def train_dgl( val_loss4 = 0 val_loss5 = 0 val_result = [] - # for dats in val_loader: val_init_time = time.time() for dats, jid in zip(val_loader, val_loader.dataset.ids): info = {} info["id"] = jid optimizer.zero_grad() - # result = net([dats[0].to(device), dats[1].to(device)]) - # if (config.model.alignn_layers) > 0: - # if (config.create_line_graph) > 0: if (config.compute_line_graph) > 0: result = net( [ @@ -498,8 +443,6 @@ def train_dgl( ) else: result = net([dats[0].to(device), dats[1].to(device)]) - # result = net(dats[0].to(device)) - # info = {} info["target_out"] = [] info["pred_out"] = [] info["target_atomwise_pred"] = [] @@ -508,11 +451,11 @@ def train_dgl( info["pred_grad"] = [] info["target_stress"] = [] info["pred_stress"] = [] - loss1 = 0 # Such as energy - loss2 = 0 # Such as bader charges - loss3 = 0 # Such as forces - loss4 = 0 # Such as stresses - loss5 = 0 # Such as stresses + loss1 = 0 + loss2 = 0 + loss3 = 0 + loss4 = 0 + loss5 = 0 if config.model.output_features is not None: loss1 = config.model.graphwise_weight * criterion( result["out"], dats[-1].to(device) @@ -551,11 +494,6 @@ def train_dgl( ) val_loss3 += loss3.item() if config.model.stresswise_weight != 0: - # loss4 = config.model.stresswise_weight * criterion( - # result["stress"].to(device), - # dats[0].ndata["stresses"][0].to(device), - # ) - targ_stress = torch.stack( [ gg.ndata["stresses"][0] @@ -563,18 +501,11 @@ def train_dgl( ] ).to(device) pred_stress = result["stresses"] - # print('targ_stress',targ_stress,targ_stress.shape) - # print('pred_stress',pred_stress,pred_stress.shape) loss4 = config.model.stresswise_weight * criterion( pred_stress.to(device), targ_stress.to(device), ) - info["target_stress"] = ( - targ_stress.cpu() - .numpy() - .tolist() - # dats[0].ndata["stresses"][0].cpu().numpy().tolist() - ) + info["target_stress"] = targ_stress.cpu().numpy().tolist() info["pred_stress"] = ( result["stresses"].cpu().detach().numpy().tolist() ) @@ -585,16 +516,10 @@ def train_dgl( gg.ndata["additional"][0] for gg in _graph_unbatch(dats[0]) ] - # print('additional_dat',additional_dat,len(additional_dat)) targ = torch.stack(additional_dat).to(device) - # targ=torch.tensor(additional_dat).to( dats[0].device) - # print('result["additional"]',result["additional"],result["additional"].shape) - # print('targ',targ,targ.shape) - # print('targ device',targ.device) loss5 = config.model.additional_output_weight * criterion( (result["additional"]).to(device), targ, - # (dats[0].ndata["additional"]).to(device), ) info["target_additional"] = targ.cpu().numpy().tolist() info["pred_additional"] = ( @@ -605,7 +530,6 @@ def train_dgl( loss = loss1 + loss2 + loss3 + loss4 + loss5 val_result.append(info) val_loss += loss.item() - # Normalize by number of val batches (see train-loop note). _n_vl = max(1, len(val_loader)) val_loss /= _n_vl val_loss1 /= _n_vl @@ -619,9 +543,6 @@ def train_dgl( val_loss3 = _ddp_mean(val_loss3, use_ddp) val_loss4 = _ddp_mean(val_loss4, use_ddp) val_loss5 = _ddp_mean(val_loss5, use_ddp) - # mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors( - # val_result - # ) val_fin_time = time.time() val_ep_time = val_fin_time - val_init_time current_model_name = "current_model.pt" @@ -639,7 +560,6 @@ def train_dgl( _unwrap(net).state_dict(), os.path.join(config.output_dir, best_model_name), ) - # print("Saving data for epoch:", e) saving_msg = "Saving model" dumpjson( filename=os.path.join( @@ -664,7 +584,6 @@ def train_dgl( val_loss5, ] ) - # history_val.append([mean_out, mean_atom, mean_grad, mean_stress]) if is_main: dumpjson( filename=os.path.join( @@ -696,14 +615,10 @@ def train_dgl( test_loss = 0 test_result = [] for dats, jid in zip(test_loader, test_loader.dataset.ids): - # for dats in test_loader: info = {} info["id"] = jid optimizer.zero_grad() - # if (config.model.alignn_layers) > 0: - # if (config.create_line_graph) > 0: if (config.compute_line_graph) > 0: - # result = net([dats[0].to(device), dats[1].to(device)]) result = net( [ dats[0].to(device), @@ -713,17 +628,14 @@ def train_dgl( ) else: result = net([dats[0].to(device), dats[1].to(device)]) - # result = net(dats[0].to(device)) - loss1 = 0 # Such as energy - loss2 = 0 # Such as bader charges - loss3 = 0 # Such as forces - loss4 = 0 # Such as stresses + loss1 = 0 + loss2 = 0 + loss3 = 0 + loss4 = 0 if ( config.model.output_features is not None and not classification ): - # print('result["out"]',result["out"]) - # print('dats[2]',dats[2]) loss1 = config.model.graphwise_weight * criterion( result["out"], dats[-1].to(device) ) @@ -764,18 +676,11 @@ def train_dgl( ] ).to(device) pred_stress = result["stresses"] - # print('targ_stress',targ_stress,targ_stress.shape) - # print('pred_stress',pred_stress,pred_stress.shape) loss4 = config.model.stresswise_weight * criterion( pred_stress.to(device), targ_stress.to(device), ) - info["target_stress"] = ( - targ_stress.cpu() - .numpy() - .tolist() - # dats[0].ndata["stresses"][0].cpu().numpy().tolist() - ) + info["target_stress"] = targ_stress.cpu().numpy().tolist() info["pred_stress"] = ( result["stresses"].cpu().detach().numpy().tolist() ) @@ -797,11 +702,9 @@ def train_dgl( _unwrap(net).state_dict(), os.path.join(config.output_dir, last_model_name), ) - # return test_result if rank == 0 or world_size == 1: if config.write_predictions and classification: best_model.eval() - # net.eval() f = open( os.path.join( config.output_dir, "prediction_results_test_set.csv" @@ -812,16 +715,12 @@ def train_dgl( targets = [] predictions = [] with torch.no_grad(): - ids = test_loader.dataset.ids # [test_loader.dataset.indices] + ids = test_loader.dataset.ids for dat, id in zip(test_loader, ids): g, lg, lat, target = dat out_data = best_model( [g.to(device), lg.to(device), lat.to(device)] )["out"] - # out_data = net([g.to(device), lg.to(device)])["out"] - # out_data = torch.exp(out_data.cpu()) - # print('target',target) - # print('out_data',out_data) top_p, top_class = torch.topk(torch.exp(out_data), k=1) target = int(target.cpu().numpy().flatten().tolist()[0]) @@ -845,22 +744,20 @@ def train_dgl( and config.model.output_features > 1 ): best_model.eval() - # net.eval() mem = [] with torch.no_grad(): - ids = test_loader.dataset.ids # [test_loader.dataset.indices] + ids = test_loader.dataset.ids for dat, id in zip(test_loader, ids): g, lg, lat, target = dat out_data = best_model( [g.to(device), lg.to(device), lat.to(device)] )["out"] - # out_data = net([g.to(device), lg.to(device)])["out"] out_data = out_data.detach().cpu().numpy().tolist() if config.standard_scalar_and_pca: sc = pk.load(open("sc.pkl", "rb")) out_data = list( sc.transform(np.array(out_data).reshape(1, -1))[0] - ) # [0][0] + ) target = target.cpu().numpy().flatten().tolist() info = {} info["id"] = id @@ -880,7 +777,6 @@ def train_dgl( and config.model.gradwise_weight == 0 ): best_model.eval() - # net.eval() f = open( os.path.join( config.output_dir, "prediction_results_test_set.csv" @@ -891,13 +787,12 @@ def train_dgl( targets = [] predictions = [] with torch.no_grad(): - ids = test_loader.dataset.ids # [test_loader.dataset.indices] + ids = test_loader.dataset.ids for dat, id in zip(test_loader, ids): g, lg, lat, target = dat out_data = best_model( [g.to(device), lg.to(device), lat.to(device)] )["out"] - # out_data = net([g.to(device), lg.to(device)])["out"] out_data = out_data.cpu().numpy().tolist() if config.standard_scalar_and_pca: sc = pk.load( @@ -919,7 +814,6 @@ def train_dgl( mean_absolute_error(np.array(targets), np.array(predictions)), ) best_model.eval() - # net.eval() f = open( os.path.join( config.output_dir, "prediction_results_train_set.csv" @@ -930,13 +824,12 @@ def train_dgl( targets = [] predictions = [] with torch.no_grad(): - ids = train_loader.dataset.ids # [test_loader.dataset.indices] + ids = train_loader.dataset.ids for dat, id in zip(train_loader, ids): g, lg, lat, target = dat out_data = best_model( [g.to(device), lg.to(device), lat.to(device)] )["out"] - # out_data = net([g.to(device), lg.to(device)])["out"] out_data = out_data.cpu().numpy().tolist() if config.standard_scalar_and_pca: sc = pk.load( @@ -946,10 +839,6 @@ def train_dgl( np.array(out_data).reshape(-1, 1) )[0][0] target = target.cpu().numpy().flatten().tolist() - # if len(target) == 1: - # target = target[0] - # if len(out_data) == 1: - # out_data = out_data[0] for ii, jj in zip(target, out_data): f.write("%6f, %6f\n" % (ii, jj)) targets.append(ii) From 10f99713fa600e5378456e26226561b81c2b8198 Mon Sep 17 00:00:00 2001 From: user Date: Wed, 20 May 2026 12:37:06 -0400 Subject: [PATCH 32/34] Train --- README.md | 41 +++++++++++++++++++++++++++++++++++++++-- alignn/train.py | 23 ++++++++++++++++------- 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index ce28751..66ffbe1 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ * [Introduction](#intro) * [Installation](#install) * [Examples](#example) +* [Colab notebooks](#colab) * [Pre-trained models](#pretrained) * [JARVIS-ALIGNN webapp](#webapp) * [ALIGNN-FF & ASE Calculator](#alignnff) @@ -48,7 +49,24 @@ See [docs/training/](docs/training/) for dataset format and training examples: - [Force-field training](docs/training/force-field.md) - [Multi-GPU training](docs/training/multi-gpu.md) -Google Colab notebooks are linked from [docs/index.md](docs/index.md). + +## Colab notebooks + +Ready-to-run notebooks covering property prediction, force-field +training, and pretrained-model usage. Click a badge to open in Colab. + +[colab-badge]: https://colab.research.google.com/assets/colab-badge.svg + +| Notebook | Open in Colab | Description | +| --- | --- | --- | +| Regression task (graph-wise prediction) | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/alignn_jarvis_leaderboard.ipynb) | Single-output regression for 2D-material exfoliation energies. | +| ML force-field training from scratch | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Train_ALIGNNFF_Mlearn.ipynb) | Train an ALIGNN-FF force field for Silicon. | +| ALIGNN-FF: relaxation, EV curve, phonons, interfaces | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Structure_Relaxation_Phonons_Interface.ipynb) | Pretrained ALIGNN-FF for relaxation, EV curves, phonons, and interfaces. | +| Scaling / timing comparison | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Timing_uMLFF.ipynb) | Scaling/timing analysis of universal MLFFs. | +| Melt-Quench MD | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Fast_Melt_Quench.ipynb) | Generate amorphous structures via molecular dynamics. | +| Miscellaneous training tasks | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Training_ALIGNN_model_example.ipynb) | Single-output, multi-output (phonon/electron DOS), classification, and pretrained usage. | +| Superconductor Tc | [![Open In Colab][colab-badge]](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/ALIGNN_Sc.ipynb) | Train a model for superconductor transition temperature. | +| Build `id_prop.json` from VASP runs | [![Open In Colab][colab-badge]](https://colab.research.google.com/gist/knc6/5513b21f5fd83a7943509ffdf5c3608b/make_id_prop.ipynb) | Compile `vasprun.xml` files into `id_prop.json` for ALIGNN-FF training. | ## Using pre-trained models @@ -66,7 +84,26 @@ See [docs/usage/webapps.md](docs/usage/webapps.md). Direct links: [AtomGPT ALIGN ## ALIGNN-FF ASE Calculator -See [docs/usage/ase-calculator.md](docs/usage/ase-calculator.md) for example usage. +```python +from ase.build import bulk +from alignn.ff.unified_calculator import ( + AlignnUnifiedCalculator, AlignnUnifiedConfig) + +cfg = AlignnUnifiedConfig( + energy=True, forces=True, stress=True, + properties=["formation_energy_peratom", "optb88vdw_bandgap"], +) +calc = AlignnUnifiedCalculator(cfg) # models loaded once, reused + +si = bulk("Si", "diamond", a=5.43); si.calc = calc +si.get_potential_energy(); si.get_forces(); si.get_stress() +print(calc.predictions()) # extra property predictors +``` + +A single pydantic config selects the outputs (force-field +energy/forces/stress plus any pretrained scalar property predictors). +See [docs/usage/ase-calculator.md](docs/usage/ase-calculator.md) for more, +and the ASE docs page *Calculators → ALIGNN*. ## Performances diff --git a/alignn/train.py b/alignn/train.py index aba5931..959a1d8 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -16,6 +16,8 @@ def _ddp_mean(value: float, use_ddp: bool) -> float: def _unwrap(net): """Return underlying module from a DDP-wrapped model, else net itself.""" return net.module if isinstance(net, DDP) else net + + from functools import partial from typing import Any, Dict, Union import torch @@ -292,9 +294,7 @@ def train_dgl( _amp_ctx = torch.autocast( device_type="cuda", dtype=torch.bfloat16, - enabled=bool( - getattr(config, "use_amp", False) - ) + enabled=bool(getattr(config, "use_amp", False)) and torch.cuda.is_available(), ) with _amp_ctx: @@ -309,9 +309,7 @@ def train_dgl( ) else: - result = net( - [dats[0].to(device), dats[1].to(device)] - ) + result = net([dats[0].to(device), dats[1].to(device)]) # info = {} info["target_out"] = [] info["pred_out"] = [] @@ -909,7 +907,18 @@ def train_dgl( target = target.cpu().numpy().flatten().tolist() if len(target) == 1: target = target[0] - f.write("%s, %6f, %6f\n" % (id, target, out_data)) + # out_data is a (possibly nested) list after .tolist(); + # unwrap single-output regression to a scalar, mirroring + # the target handling above. Multi-output falls back to + # a generic format so the write never raises. + if isinstance(out_data, list): + out_data = np.array(out_data).flatten().tolist() + if len(out_data) == 1: + out_data = out_data[0] + if isinstance(target, list) or isinstance(out_data, list): + f.write("%s, %s, %s\n" % (id, target, out_data)) + else: + f.write("%s, %6f, %6f\n" % (id, target, out_data)) targets.append(target) predictions.append(out_data) f.close() From d875997da874f5316b7d44e763b7064d2eac4f8d Mon Sep 17 00:00:00 2001 From: user Date: Wed, 20 May 2026 14:02:40 -0400 Subject: [PATCH 33/34] v5.20 --- alignn/__init__.py | 2 +- alignn/ff/all_models_alignn_atomwise.json | 29 +++ alignn/ff/ff.py | 4 +- alignn/ff/unified_calculator.py | 248 ++++++++++++++++++++++ alignn/models/model_zoo.py | 243 +++++++++++++++++++++ setup.py | 4 +- 6 files changed, 526 insertions(+), 4 deletions(-) create mode 100644 alignn/ff/unified_calculator.py create mode 100644 alignn/models/model_zoo.py diff --git a/alignn/__init__.py b/alignn/__init__.py index 76db32a..e117a37 100644 --- a/alignn/__init__.py +++ b/alignn/__init__.py @@ -1,3 +1,3 @@ """Version number.""" -__version__ = "2026.4.3" +__version__ = "2026.5.20" diff --git a/alignn/ff/all_models_alignn_atomwise.json b/alignn/ff/all_models_alignn_atomwise.json index 54a0165..f0c1465 100644 --- a/alignn/ff/all_models_alignn_atomwise.json +++ b/alignn/ff/all_models_alignn_atomwise.json @@ -1,4 +1,33 @@ { + "dft_3d_307k": "https://ndownloader.figshare.com/files/64744104", + "mps": "https://ndownloader.figshare.com/files/64744107", + "avg_elec_mass": "https://ndownloader.figshare.com/files/64744110", + "bulk_modulus_kv": "https://ndownloader.figshare.com/files/64744113", + "avg_hole_mass": "https://ndownloader.figshare.com/files/64744116", + "dfpt_piezo_max_dielectric": "https://ndownloader.figshare.com/files/64744119", + "density": "https://ndownloader.figshare.com/files/64744122", + "dfpt_piezo_max_dij": "https://ndownloader.figshare.com/files/64744125", + "dfpt_piezo_max_eij": "https://ndownloader.figshare.com/files/64744128", + "ehull": "https://ndownloader.figshare.com/files/64744131", + "epsx": "https://ndownloader.figshare.com/files/64744134", + "exfoliation_energy": "https://ndownloader.figshare.com/files/64744137", + "formation_energy_peratom": "https://ndownloader.figshare.com/files/64744140", + "magmom_oszicar": "https://ndownloader.figshare.com/files/64744143", + "max_efg": "https://ndownloader.figshare.com/files/64744146", + "max_ir_mode": "https://ndownloader.figshare.com/files/64744149", + "mbj_bandgap": "https://ndownloader.figshare.com/files/64744152", + "mepsx": "https://ndownloader.figshare.com/files/64744155", + "min_ir_mode": "https://ndownloader.figshare.com/files/64744158", + "n-powerfact": "https://ndownloader.figshare.com/files/64744161", + "poisson": "https://ndownloader.figshare.com/files/64744164", + "optb88vdw_bandgap": "https://ndownloader.figshare.com/files/64744167", + "n-Seebeck": "https://ndownloader.figshare.com/files/64744170", + "p-powerfact": "https://ndownloader.figshare.com/files/64744173", + "p-Seebeck": "https://ndownloader.figshare.com/files/64744176", + "shear_modulus_gv": "https://ndownloader.figshare.com/files/64744179", + "slme": "https://ndownloader.figshare.com/files/64744182", + "spillage": "https://ndownloader.figshare.com/files/64744185", + "Tc_supercon": "https://ndownloader.figshare.com/files/64744188", "v12.2.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50904240", "v12.2.2024_mp_1.5mill": "https://ndownloader.figshare.com/files/50904783", "v12.2.2024_mp_187k": "https://ndownloader.figshare.com/files/50904801", diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py index b065e4f..eac85f7 100644 --- a/alignn/ff/ff.py +++ b/alignn/ff/ff.py @@ -48,6 +48,7 @@ from sklearn.metrics import mean_absolute_error from tqdm import tqdm from jarvis.core.utils import get_cache_dir + # import torch from alignn.ff.calculators import ( AlignnAtomwiseCalculator, @@ -207,7 +208,8 @@ def get_figshare_model_prop( def default_path(): """Get default model path.""" - dpath = get_figshare_model_ff(model_name="v12.2.2024_dft_3d_307k") + dpath = get_figshare_model_ff(model_name="mps") + # dpath = get_figshare_model_ff(model_name="v12.2.2024_dft_3d_307k") # dpath = get_figshare_model_ff(model_name="v5.27.2024") # dpath = get_figshare_model_ff(model_name="v8.29.2024_dft_3d") # dpath = get_figshare_model_ff(model_name="alignnff_wt10") diff --git a/alignn/ff/unified_calculator.py b/alignn/ff/unified_calculator.py new file mode 100644 index 0000000..d23fc95 --- /dev/null +++ b/alignn/ff/unified_calculator.py @@ -0,0 +1,248 @@ +"""Unified ALIGNN calculator driven by a pydantic config. + +One calculator, one config object. The config selects which outputs are +produced: + + * force-field properties : energy / forces / stress (one ALIGNN-FF + model) + * scalar property predictors : formation energy, band gap, bulk + modulus, ... (one pretrained ALIGNN regression model per property) + +All models are loaded once and reused for every structure / every call +(no per-call reload). + +Example +------- + from alignn.ff.unified_calculator import ( + AlignnUnifiedCalculator, AlignnUnifiedConfig) + from ase.build import bulk + + cfg = AlignnUnifiedConfig( + energy=True, forces=True, stress=True, + properties=["formation_energy_peratom", + "optb88vdw_bandgap", "bulk_modulus_kv"], + ) + calc = AlignnUnifiedCalculator(cfg) + + si = bulk("Si", "diamond", a=5.43); si.calc = calc + si.get_potential_energy(); si.get_forces(); si.get_stress() + print(calc.results["formation_energy_peratom"], + calc.results["optb88vdw_bandgap"], + calc.results["bulk_modulus_kv"]) +""" + +from __future__ import annotations + +from typing import Dict, List, Optional + +import numpy as np +import torch +from ase.calculators.calculator import Calculator, all_changes +from pydantic import BaseModel, Field, field_validator + +from alignn.ff.calculators import ( + AlignnAtomwiseCalculator, + ase_to_atoms, +) +from alignn.graphs import Graph +from alignn.pretrained import get_figshare_model + +# friendly name -> figshare model name (extend freely; raw figshare +# names ending in "_alignn" are also accepted as-is) +PROP_ALIASES: Dict[str, str] = { + "formation_energy_peratom": "jv_formation_energy_peratom_alignn", + "total_energy": "jv_optb88vdw_total_energy_alignn", + "optb88vdw_bandgap": "jv_optb88vdw_bandgap_alignn", + "mbj_bandgap": "jv_mbj_bandgap_alignn", + "bulk_modulus_kv": "jv_bulk_modulus_kv_alignn", + "shear_modulus_gv": "jv_shear_modulus_gv_alignn", + "ehull": "jv_ehull_alignn", + "spillage": "jv_spillage_alignn", + "slme": "jv_slme_alignn", + "magmom_oszicar": "jv_magmom_oszicar_alignn", + "exfoliation_energy": "jv_exfoliation_energy_alignn", + "supercon_tc": "jv_supercon_tc_alignn", + "epsx": "jv_epsx_alignn", + "n_seebeck": "jv_n-Seebeck_alignn", + "n_powerfact": "jv_n-powerfact_alignn", +} + + +def _resolve_prop(name: str) -> str: + """friendly or raw -> figshare model name.""" + if name in PROP_ALIASES: + return PROP_ALIASES[name] + if name.endswith("_alignn"): # raw figshare name passed through + return name + raise ValueError( + f"Unknown property '{name}'. Known: " + f"{sorted(PROP_ALIASES)} (or a raw *_alignn figshare name)." + ) + + +class AlignnUnifiedConfig(BaseModel): + """Declarative spec of what the calculator should output.""" + + # force-field model (energy/forces/stress source) + ff_model: str = "v12.2.2024_dft_3d_307k" + energy: bool = True + forces: bool = True + stress: bool = True + charges: bool = False # only if the FF model emits them + + # scalar property predictors to also evaluate (friendly names) + properties: List[str] = Field(default_factory=list) + + # shared knobs + device: Optional[str] = None + prop_cutoff: float = 8.0 + prop_max_neighbors: int = 12 + + model_config = {"extra": "forbid"} + + @field_validator("properties") + @classmethod + def _check_props(cls, v: List[str]) -> List[str]: + for p in v: + _resolve_prop(p) # raises on unknown + return v + + @classmethod + def from_file(cls, path: str) -> "AlignnUnifiedConfig": + """Load from a JSON file.""" + import json + + with open(path) as fh: + return cls(**json.load(fh)) + + +class AlignnUnifiedCalculator(Calculator): + """ASE calculator: ALIGNN-FF + selected ALIGNN property predictors. + + Args: + config: an ``AlignnUnifiedConfig``, a dict, a path to a JSON + config, or None (defaults: energy/forces/stress only). + ff_path: optional local path/dir for the FF model (overrides + ``config.ff_model``). + """ + + def __init__(self, config=None, ff_path=None, **kw): + Calculator.__init__(self, **kw) + + if config is None: + config = AlignnUnifiedConfig() + elif isinstance(config, str): + config = AlignnUnifiedConfig.from_file(config) + elif isinstance(config, dict): + config = AlignnUnifiedConfig(**config) + self.cfg: AlignnUnifiedConfig = config + + self.device = config.device or ( + "cuda" if torch.cuda.is_available() else "cpu" + ) + + props = [] + if config.energy: + props.append("energy") + if config.forces: + props.append("forces") + if config.stress: + props.append("stress") + self.implemented_properties = props + + # --- force-field calculator (loaded ONCE) --- + self._ff = None + if config.energy or config.forces or config.stress: + ff_kwargs = dict( + include_stress=config.stress, device=self.device + ) + if ff_path is not None: + ff_kwargs["path"] = ff_path + elif config.ff_model: + ff_kwargs["path"] = _ff_model_path(config.ff_model) + self._ff = AlignnAtomwiseCalculator(**ff_kwargs) + + # --- property predictor models (each loaded ONCE) --- + self._prop_models: Dict[str, object] = {} + for friendly in config.properties: + self._prop_models[friendly] = get_figshare_model( + _resolve_prop(friendly) + ) + + # -- scalar property forward (reuses cached model, no reload) ------- + def _predict_scalar(self, model, j_atoms) -> float: + g, lg = Graph.atom_dgl_multigraph( + j_atoms, + cutoff=float(self.cfg.prop_cutoff), + max_neighbors=self.cfg.prop_max_neighbors, + ) + lat = torch.tensor(j_atoms.lattice_mat) + with torch.no_grad(): + out = model( + [ + g.to(self.device), + lg.to(self.device), + lat.to(self.device), + ] + ) + if isinstance(out, dict): + out = out["out"] + arr = out.detach().cpu().numpy().flatten() + return float(arr[0]) if arr.size == 1 else arr.tolist() + + def calculate( + self, atoms=None, properties=("energy",), system_changes=all_changes + ): + Calculator.calculate(self, atoms, properties, system_changes) + self.results = {} + + # force-field block + if self._ff is not None: + self._ff.calculate( + self.atoms, properties=properties, + system_changes=system_changes, + ) + if self.cfg.energy: + self.results["energy"] = self._ff.results["energy"] + self.results["free_energy"] = self._ff.results["energy"] + if self.cfg.forces: + self.results["forces"] = self._ff.results["forces"] + if self.cfg.stress: + self.results["stress"] = self._ff.results["stress"] + if self.cfg.charges: + ch = self._ff.results.get("charges") + if ch is not None: + self.results["charges"] = ch + else: + self.results.setdefault("warnings", []).append( + "charges requested but the FF model does not " + "emit them" + ) + + # scalar property predictors + if self._prop_models: + j_atoms = ase_to_atoms(self.atoms) + for friendly, model in self._prop_models.items(): + self.results[friendly] = self._predict_scalar( + model, j_atoms + ) + + # convenience + def predictions(self) -> Dict[str, object]: + """All scalar property predictions from the last calculate().""" + return { + k: self.results[k] + for k in self.cfg.properties + if k in self.results + } + + +def _ff_model_path(name_or_path: str) -> str: + """Accept a local dir or a known FF figshare model name.""" + import os + + if os.path.isdir(name_or_path): + return name_or_path + from alignn.ff.calculators import get_figshare_model_ff + + return get_figshare_model_ff(model_name=name_or_path) diff --git a/alignn/models/model_zoo.py b/alignn/models/model_zoo.py new file mode 100644 index 0000000..ee11aca --- /dev/null +++ b/alignn/models/model_zoo.py @@ -0,0 +1,243 @@ +"""Bundle multiple pretrained models behind one nn.Module. + +Each entry is a (model, output_spec) pair. ``output_spec`` says how to map +the underlying model's raw output dict to a single tensor for that +property — e.g. EFS returns three tensors, formation energy returns one. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Callable, Dict, List, Optional + +import torch +from torch import nn + +from alignn.models.alignn_atomwise_pure import ( + ALIGNNAtomWisePure, + ALIGNNAtomWisePureConfig, +) + + +@dataclass +class PropertySpec: + """How to invoke one member model and what to extract from it.""" + + name: str + model: nn.Module + # Map raw model output (whatever ``model(graph)`` returns) → dict of + # named tensors for this property. Default pulls "out" from an + # ALIGNN-Pure dict. + extract: Callable[[object], Dict[str, torch.Tensor]] = field( + default=lambda d: {"value": d["out"]} + ) + + +def _efs_extract(d): + return {"energy": d["out"], "forces": d["grad"], "stress": d["stresses"]} + + +def _scalar_extract(d): + return {"value": d["out"]} + + +class ModelZoo(nn.Module): + """Registry of pretrained models exposed as one nn.Module. + + All members live in a single ``nn.ModuleDict`` so ``state_dict``, + ``.to(device)``, ``torch.save`` etc. work on the whole collection. + """ + + def __init__(self, specs: Optional[List[PropertySpec]] = None): + super().__init__() + self._models = nn.ModuleDict() + self._extractors: Dict[str, Callable] = {} + for spec in specs or []: + self.register(spec) + + def register(self, spec: PropertySpec) -> None: + if spec.name in self._models: + raise KeyError(f"property {spec.name!r} already registered") + self._models[spec.name] = spec.model + self._extractors[spec.name] = spec.extract + + def properties(self) -> List[str]: + return list(self._models.keys()) + + def predict(self, graph, prop: str) -> Dict[str, torch.Tensor]: + out = self._models[prop](graph) + return self._extractors[prop](out) + + def predict_all(self, graph) -> Dict[str, Dict[str, torch.Tensor]]: + return {p: self.predict(graph, p) for p in self._models} + + # default forward = predict_all so torch.compile / hooks see one entry point + def forward(self, graph) -> Dict[str, Dict[str, torch.Tensor]]: + return self.predict_all(graph) + + # ----- convenience constructors ----- + + # ----- manifest-based lazy loading ----- + + @classmethod + def from_manifest( + cls, + manifest_path: str, + map_location: str = "cpu", + eager: bool = False, + ) -> "LazyModelZoo": + """Build a zoo whose members live on disk and load on first use. + + Manifest schema (JSON): + { + "base_dir": "optional; paths are resolved relative to this", + "properties": { + "": { + "checkpoint": "", + "config": { ... ALIGNNAtomWisePureConfig kwargs ... }, + "is_efs": false + }, + ... + } + } + """ + with open(manifest_path) as f: + spec = json.load(f) + base = Path(spec.get("base_dir") or Path(manifest_path).parent) + zoo = LazyModelZoo(map_location=map_location) + for name, entry in spec["properties"].items(): + ckpt = base / entry["checkpoint"] + cfg_kwargs = dict(entry.get("config") or {}) + cfg = ALIGNNAtomWisePureConfig(**cfg_kwargs) + zoo.add_lazy( + name=name, + checkpoint=str(ckpt), + config=cfg, + is_efs=bool(entry.get("is_efs", False)), + ) + if eager: + for name in zoo.properties(): + zoo._materialize(name) + return zoo + + @classmethod + def from_alignn_checkpoints( + cls, + ckpts: Dict[str, str], + configs: Optional[Dict[str, ALIGNNAtomWisePureConfig]] = None, + efs_keys: Optional[List[str]] = None, + map_location: str = "cpu", + ) -> "ModelZoo": + """Build a zoo from {property_name: checkpoint_path}. + + ``efs_keys`` lists names whose model returns energy/forces/stress + (everything else is treated as a single scalar head). + ``configs`` lets you pass a per-model config; missing entries fall + back to the default ``ALIGNNAtomWisePureConfig()``. + """ + configs = configs or {} + efs = set(efs_keys or []) + specs: List[PropertySpec] = [] + for name, path in ckpts.items(): + cfg = configs.get(name, ALIGNNAtomWisePureConfig(name=name)) + model = ALIGNNAtomWisePure(cfg) + state = torch.load(path, map_location=map_location) + if isinstance(state, dict) and "model" in state: + state = state["model"] + model.load_state_dict(state) + model.eval() + specs.append( + PropertySpec( + name=name, + model=model, + extract=_efs_extract if name in efs else _scalar_extract, + ) + ) + return cls(specs) + + +class LazyModelZoo(ModelZoo): + """Zoo whose members are loaded from disk on first use. + + Only one ALIGNN trunk's weights occupy memory until you actually + request a property. Materialized members are cached, so repeated + calls hit RAM, not disk. + """ + + def __init__(self, map_location: str = "cpu"): + super().__init__() + self._map_location = map_location + self._lazy: Dict[str, dict] = {} # name -> {checkpoint, config, is_efs} + + def add_lazy( + self, + name: str, + checkpoint: str, + config: ALIGNNAtomWisePureConfig, + is_efs: bool = False, + ) -> None: + if name in self._lazy or name in self._models: + raise KeyError(f"property {name!r} already registered") + self._lazy[name] = { + "checkpoint": checkpoint, + "config": config, + "is_efs": is_efs, + } + self._extractors[name] = _efs_extract if is_efs else _scalar_extract + + def properties(self) -> List[str]: + return list(self._lazy.keys()) + [ + n for n in self._models.keys() if n not in self._lazy + ] + + def _materialize(self, name: str) -> nn.Module: + if name in self._models: + return self._models[name] + if name not in self._lazy: + raise KeyError(name) + entry = self._lazy[name] + model = ALIGNNAtomWisePure(entry["config"]) + state = torch.load(entry["checkpoint"], map_location=self._map_location) + if isinstance(state, dict) and "model" in state: + state = state["model"] + model.load_state_dict(state) + model.eval() + self._models[name] = model + return model + + def predict(self, graph, prop: str) -> Dict[str, torch.Tensor]: + self._materialize(prop) + return super().predict(graph, prop) + + def evict(self, name: str) -> None: + """Drop a materialized member back to disk-only state.""" + if name in self._models and name in self._lazy: + del self._models[name] + + +# ----- example: registering a non-ALIGNN model (e.g. SlakoNet) --------------- +# +# from slakonet import SlakoNetModel +# +# zoo = ModelZoo.from_alignn_checkpoints( +# { +# "efs": "ff/alignnff.pt", +# "formation_energy": "jv_formation_energy_peratom_alignn.pt", +# "mbj_bandgap": "jv_mbj_bandgap_alignn.pt", +# "bulk_modulus": "jv_bulk_modulus_kv_alignn.pt", +# }, +# efs_keys=["efs"], +# ) +# +# slakonet = SlakoNetModel.load("slakonet.pt").eval() +# zoo.register(PropertySpec( +# name="bandstructure", +# model=slakonet, +# extract=lambda d: {"bandgap": d["bandgap"], "bands": d["bands"]}, +# )) +# +# torch.save(zoo.state_dict(), "alignn_zoo.pt") +# preds = zoo.predict_all(graph) # {prop: {field: tensor}} +# fe = zoo.predict(graph, "formation_energy")["value"] diff --git a/setup.py b/setup.py index 9e9e985..0fcbe58 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ setuptools.setup( name="alignn", - version="2026.4.3", + version="2026.5.20", author="Kamal Choudhary, Brian DeCost", author_email="kamal.choudhary@nist.gov", description="alignn", @@ -22,7 +22,7 @@ "jarvis-tools>=2021.07.19", "torch>=2.2.1", # "torch<=2.2.1", - "mpmath<=1.3.0", + "mpmath>=1.3.0", # "dgl>=0.6.0", "spglib>=2.0.2", "scikit-learn>=0.22.2", From 89b7d5aca323062f9dd2fd977699fc2c6d7034a4 Mon Sep 17 00:00:00 2001 From: user Date: Wed, 20 May 2026 22:58:35 -0400 Subject: [PATCH 34/34] Lint fix --- alignn/data.py | 3 +- alignn/ff/unified_calculator.py | 1 - alignn/graphs.py | 7 ++- alignn/md/forces.py | 12 ++++-- alignn/md/integrators.py | 71 ++++++++++++++++++++----------- alignn/md/relax.py | 5 ++- alignn/models/alignn.py | 2 +- alignn/models/alignn_atomwise.py | 2 +- alignn/models/ealignn_atomwise.py | 2 +- alignn/models/model_zoo.py | 9 ++-- alignn/train.py | 31 +++++++------- alignn/train_alignn.py | 4 +- 12 files changed, 92 insertions(+), 57 deletions(-) diff --git a/alignn/data.py b/alignn/data.py index 00aad93..97fcbe8 100644 --- a/alignn/data.py +++ b/alignn/data.py @@ -13,10 +13,11 @@ try: from dgl.dataloading import GraphDataLoader -except: # pure-torch path; fall back to stdlib DataLoader +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) diff --git a/alignn/ff/unified_calculator.py b/alignn/ff/unified_calculator.py index d23fc95..d144354 100644 --- a/alignn/ff/unified_calculator.py +++ b/alignn/ff/unified_calculator.py @@ -35,7 +35,6 @@ from typing import Dict, List, Optional -import numpy as np import torch from ase.calculators.calculator import Calculator, all_changes from pydantic import BaseModel, Field, field_validator diff --git a/alignn/graphs.py b/alignn/graphs.py index 144ba97..0146d2e 100644 --- a/alignn/graphs.py +++ b/alignn/graphs.py @@ -18,11 +18,14 @@ try: from dgl.data import DGLDataset import dgl -except: # DGL optional; DGL-dependent helpers will fail lazily +except Exception: # DGL optional; DGL-dependent helpers fail lazily dgl = None - class DGLDataset: # minimal stub so subclasses can still be defined + class DGLDataset: # noqa: D101 - minimal stub for subclasses + """Minimal stub so subclasses can be defined without DGL.""" + def __init__(self, *args, **kwargs): + """Accept and ignore any arguments.""" pass diff --git a/alignn/md/forces.py b/alignn/md/forces.py index 38d8920..bd8ac60 100644 --- a/alignn/md/forces.py +++ b/alignn/md/forces.py @@ -49,8 +49,10 @@ def _build_graph(self, positions_np): cutoff=self.cutoff, max_neighbors=self.max_neighbors, atom_features="cgcnn", use_canonize=True, ) - g = g.to(self.device); lg = lg.to(self.device) - tg = torchgraph_from_dgl(g); tlg = torchgraph_from_dgl(lg) + g = g.to(self.device) + lg = lg.to(self.device) + tg = torchgraph_from_dgl(g) + tlg = torchgraph_from_dgl(lg) for d in (tg.ndata, tg.edata, tlg.ndata, tlg.edata): for k, v in list(d.items()): if v.is_floating_point(): @@ -64,7 +66,11 @@ def __call__(self, positions: torch.Tensor): # model sets r.requires_grad_(True) internally; energy grad wrt r # gives per-edge forces, aggregated to per-atom in the model's head. out = self.model((tg, tlg, self.cell)) - energy = out["out"].sum() if "out" in out else next(iter(out.values())).sum() + energy = ( + out["out"].sum() + if "out" in out + else next(iter(out.values())).sum() + ) if "grad" in out: forces = out["grad"].detach() else: diff --git a/alignn/md/integrators.py b/alignn/md/integrators.py index 105f59b..2611363 100644 --- a/alignn/md/integrators.py +++ b/alignn/md/integrators.py @@ -1,4 +1,4 @@ -"""Pure-PyTorch MD integrators: Velocity-Verlet (NVE) and Langevin (NVT, BAOAB). +"""Pure-PyTorch MD integrators: Velocity-Verlet and Langevin (BAOAB). State kept entirely on-device. Unit system follows ASE: energy eV, length Å, mass amu, time fs, temperature K. @@ -9,7 +9,8 @@ from __future__ import annotations from dataclasses import dataclass from typing import Callable, Optional -import math, time +import math +import time import torch KB_EV = 8.617333262e-5 # Boltzmann, eV/K @@ -26,9 +27,15 @@ def wrap_pbc(positions: torch.Tensor, cell: torch.Tensor) -> torch.Tensor: return frac @ cell -def kinetic_energy(masses: torch.Tensor, velocities: torch.Tensor) -> torch.Tensor: - # KE in eV: 0.5 * m * v² with m in amu, v in Å/fs -> multiply by 1/EV_AMU_A_PER_FS2 - return 0.5 * (masses * (velocities ** 2).sum(dim=-1)).sum() / EV_AMU_A_PER_FS2 +def kinetic_energy( + masses: torch.Tensor, velocities: torch.Tensor +) -> torch.Tensor: + """Return kinetic energy in eV (m in amu, v in Å/fs).""" + # 0.5 * m * v², converted from amu·Å²/fs² to eV + return ( + 0.5 * (masses * (velocities ** 2).sum(dim=-1)).sum() + / EV_AMU_A_PER_FS2 + ) def instantaneous_temperature(masses, velocities): @@ -37,7 +44,9 @@ def instantaneous_temperature(masses, velocities): return 2.0 * ke / (n_dof * KB_EV) -def maxwell_boltzmann(masses: torch.Tensor, T: float, generator=None) -> torch.Tensor: +def maxwell_boltzmann( + masses: torch.Tensor, T: float, generator=None +) -> torch.Tensor: # sigma_v = sqrt(kT/m), in Å/fs after unit conversion sigma = torch.sqrt(KB_EV * T * EV_AMU_A_PER_FS2 / masses).unsqueeze(-1) v = torch.randn((masses.numel(), 3), device=masses.device, @@ -107,10 +116,18 @@ def step(self, positions, velocities): # A: half-drift positions = positions + 0.5 * self.dt * velocities # O: Ornstein-Uhlenbeck in velocity - sigma = torch.sqrt(KB_EV * self.T * EV_AMU_A_PER_FS2 / self.masses).unsqueeze(-1) - noise = torch.randn_like(velocities) if self.generator is None \ - else torch.randn(velocities.shape, device=velocities.device, - dtype=velocities.dtype, generator=self.generator) + sigma = torch.sqrt( + KB_EV * self.T * EV_AMU_A_PER_FS2 / self.masses + ).unsqueeze(-1) + if self.generator is None: + noise = torch.randn_like(velocities) + else: + noise = torch.randn( + velocities.shape, + device=velocities.device, + dtype=velocities.dtype, + generator=self.generator, + ) velocities = self.c1 * velocities + self.c3 * sigma * noise # A: half-drift positions = positions + 0.5 * self.dt * velocities @@ -158,7 +175,9 @@ def step(self, positions, velocities): T_now = instantaneous_temperature(self.masses, velocities) # Guard against T_now == 0 early on scale = torch.sqrt( - 1.0 + (self.dt / self.taut) * (self.T / T_now.clamp_min(1e-6) - 1.0) + 1.0 + + (self.dt / self.taut) + * (self.T / T_now.clamp_min(1e-6) - 1.0) ) velocities = velocities * scale self._forces = new_forces @@ -258,7 +277,9 @@ def _set_masses(self): self._Nf = Nf self._kT = kT tau2 = self.taut ** 2 - Q = torch.full((self.chain_length,), kT * tau2, device=device, dtype=dtype) + Q = torch.full( + (self.chain_length,), kT * tau2, device=device, dtype=dtype + ) Q[0] = Nf * kT * tau2 self._Q = Q @@ -275,16 +296,18 @@ def _apply_chain_half(self, velocities): M = self.chain_length K = kinetic_energy(self.masses, velocities) # eV # --- reverse pass (top of chain down) --- - G_M = (Q[M-2] * v_xi[M-2] ** 2 - self._kT) / Q[M-1] if M >= 2 else \ - (2 * K - self._Nf * self._kT) / Q[0] - v_xi[M-1] = v_xi[M-1] + G_M * dthalf + if M >= 2: + G_M = (Q[M - 2] * v_xi[M - 2] ** 2 - self._kT) / Q[M - 1] + else: + G_M = (2 * K - self._Nf * self._kT) / Q[0] + v_xi[M - 1] = v_xi[M - 1] + G_M * dthalf for i in range(M - 2, -1, -1): factor = torch.exp(-dtqtr * v_xi[i + 1]) v_xi[i] = v_xi[i] * factor if i == 0: G_i = (2 * K - self._Nf * self._kT) / Q[0] else: - G_i = (Q[i-1] * v_xi[i-1] ** 2 - self._kT) / Q[i] + G_i = (Q[i - 1] * v_xi[i - 1] ** 2 - self._kT) / Q[i] v_xi[i] = v_xi[i] + G_i * dthalf v_xi[i] = v_xi[i] * factor # --- rescale particle velocities --- @@ -298,12 +321,12 @@ def _apply_chain_half(self, velocities): if i == 0: G_i = (2 * K - self._Nf * self._kT) / Q[0] else: - G_i = (Q[i-1] * v_xi[i-1] ** 2 - self._kT) / Q[i] + G_i = (Q[i - 1] * v_xi[i - 1] ** 2 - self._kT) / Q[i] v_xi[i] = v_xi[i] + G_i * dthalf v_xi[i] = v_xi[i] * factor if M >= 2: - G_M = (Q[M-2] * v_xi[M-2] ** 2 - self._kT) / Q[M-1] - v_xi[M-1] = v_xi[M-1] + G_M * dthalf + G_M = (Q[M - 2] * v_xi[M - 2] ** 2 - self._kT) / Q[M - 1] + v_xi[M - 1] = v_xi[M - 1] + G_M * dthalf return velocities def step(self, positions, velocities): @@ -345,15 +368,15 @@ def run( if (i % log_every) == 0 or i == nsteps - 1: ke = kinetic_energy(integrator.masses, velocities).item() T = instantaneous_temperature(integrator.masses, velocities).item() - pe = integrator.forces_fn(positions)[0].item() \ - if hasattr(integrator, "_forces") and integrator._forces is None \ - else None row = {"step": i, "time_fs": i * integrator.dt, "T_K": T, "KE_eV": ke, "wall_s": time.time() - t0} hist.append(row) if callback is not None: callback(i, row) else: - print(f"step {i:6d} t={row['time_fs']:8.1f} fs " - f"T={T:7.1f} K KE={ke:10.4f} eV wall={row['wall_s']:6.1f}s") + print( + f"step {i:6d} t={row['time_fs']:8.1f} fs " + f"T={T:7.1f} K KE={ke:10.4f} eV " + f"wall={row['wall_s']:6.1f}s" + ) return positions, velocities, hist diff --git a/alignn/md/relax.py b/alignn/md/relax.py index 45cecce..bc3db8d 100644 --- a/alignn/md/relax.py +++ b/alignn/md/relax.py @@ -60,7 +60,10 @@ def run( f"|F|max={fnorm_max:9.4f} eV/Å dt={dt:5.3f} fs " f"α={alpha:5.3f}") if fnorm_max < fmax: - print(f"converged at step {step}: |F|max = {fnorm_max:.5f} < {fmax}") + print( + f"converged at step {step}: " + f"|F|max = {fnorm_max:.5f} < {fmax}" + ) break # FIRE velocity mix: v ← (1-α) v + α |v| F̂ diff --git a/alignn/models/alignn.py b/alignn/models/alignn.py index f58b8ed..6d44954 100644 --- a/alignn/models/alignn.py +++ b/alignn/models/alignn.py @@ -11,7 +11,7 @@ import dgl import dgl.function as fn from dgl.nn import AvgPooling -except: +except Exception: dgl = None fn = None AvgPooling = None diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py index 6c9952e..25a3fc2 100644 --- a/alignn/models/alignn_atomwise.py +++ b/alignn/models/alignn_atomwise.py @@ -12,7 +12,7 @@ import dgl import dgl.function as fn from dgl.nn import AvgPooling -except: # DGL optional; this module only usable if DGL installed +except Exception: # DGL optional; module only usable if DGL installed dgl = None fn = None AvgPooling = None diff --git a/alignn/models/ealignn_atomwise.py b/alignn/models/ealignn_atomwise.py index 6f2bf5a..a7d8088 100644 --- a/alignn/models/ealignn_atomwise.py +++ b/alignn/models/ealignn_atomwise.py @@ -12,7 +12,7 @@ import dgl import dgl.function as fn from dgl.nn import AvgPooling -except: +except Exception: dgl = None fn = None AvgPooling = None diff --git a/alignn/models/model_zoo.py b/alignn/models/model_zoo.py index ee11aca..64b0be8 100644 --- a/alignn/models/model_zoo.py +++ b/alignn/models/model_zoo.py @@ -73,7 +73,7 @@ def predict(self, graph, prop: str) -> Dict[str, torch.Tensor]: def predict_all(self, graph) -> Dict[str, Dict[str, torch.Tensor]]: return {p: self.predict(graph, p) for p in self._models} - # default forward = predict_all so torch.compile / hooks see one entry point + # default forward = predict_all so torch.compile sees one entry point def forward(self, graph) -> Dict[str, Dict[str, torch.Tensor]]: return self.predict_all(graph) @@ -169,7 +169,8 @@ class LazyModelZoo(ModelZoo): def __init__(self, map_location: str = "cpu"): super().__init__() self._map_location = map_location - self._lazy: Dict[str, dict] = {} # name -> {checkpoint, config, is_efs} + # name -> {checkpoint, config, is_efs} + self._lazy: Dict[str, dict] = {} def add_lazy( self, @@ -199,7 +200,9 @@ def _materialize(self, name: str) -> nn.Module: raise KeyError(name) entry = self._lazy[name] model = ALIGNNAtomWisePure(entry["config"]) - state = torch.load(entry["checkpoint"], map_location=self._map_location) + state = torch.load( + entry["checkpoint"], map_location=self._map_location + ) if isinstance(state, dict) and "model" in state: state = state["model"] model.load_state_dict(state) diff --git a/alignn/train.py b/alignn/train.py index 57b2324..7ea303f 100644 --- a/alignn/train.py +++ b/alignn/train.py @@ -2,22 +2,6 @@ from torch.nn.parallel import DistributedDataParallel as DDP import torch.distributed as dist - - -def _ddp_mean(value: float, use_ddp: bool) -> float: - """All-reduce a Python scalar across ranks and return the mean.""" - if not use_ddp or not dist.is_available() or not dist.is_initialized(): - return float(value) - t = torch.tensor(float(value), device=torch.cuda.current_device()) - dist.all_reduce(t, op=dist.ReduceOp.SUM) - return (t / dist.get_world_size()).item() - - -def _unwrap(net): - """Return underlying module from a DDP-wrapped model, else net itself.""" - return net.module if isinstance(net, DDP) else net - - from functools import partial from typing import Any, Dict, Union import torch @@ -48,6 +32,21 @@ def _unwrap(net): warnings.filterwarnings("ignore", category=RuntimeWarning) + +def _ddp_mean(value: float, use_ddp: bool) -> float: + """All-reduce a Python scalar across ranks and return the mean.""" + if not use_ddp or not dist.is_available() or not dist.is_initialized(): + return float(value) + t = torch.tensor(float(value), device=torch.cuda.current_device()) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + return (t / dist.get_world_size()).item() + + +def _unwrap(net): + """Return underlying module from a DDP-wrapped model, else net.""" + return net.module if isinstance(net, DDP) else net + + # torch.autograd.detect_anomaly() figlet_alignn = """ diff --git a/alignn/train_alignn.py b/alignn/train_alignn.py index d95c533..5891e8a 100644 --- a/alignn/train_alignn.py +++ b/alignn/train_alignn.py @@ -355,11 +355,9 @@ def train_for_folder( try: tmp = ALIGNNAtomWiseConfig(**rest_config["model"]) model = ALIGNNAtomWise(tmp) # config.model) - except: + except Exception: tmp = ALIGNNAtomWisePureConfig(**rest_config["model"]) model = ALIGNNAtomWisePure(tmp) # config.model) - pass - # from alignn.models.alignn_atomwise_pure import ALIGNNAtomWisePureConfig print("Rest config", tmp) print("model", model) model.load_state_dict(