alignn 2026.8.6__py3-none-any.whl

alignn/__init__.py

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

alignn/cli.py

@@ -0,0 +1,73 @@
+"""Ignite training cli."""
+
+import json
+import os
+import shutil
+from pathlib import Path
+
+# from typing import Any, Dict, Optional, Union
+from typing import Optional
+
+import torch
+import typer
+
+from alignn.config import TrainingConfig
+from alignn.profiler import profile_dgl
+from alignn.train import train_dgl
+
+
+def cli(
+    config: Optional[Path] = typer.Argument(None),
+    progress: bool = False,
+    checkpoint_dir: Path = Path("/tmp/models"),
+    store_outputs: bool = False,
+    tensorboard: bool = False,
+    profile: bool = False,
+):
+    """ALIGNN training cli.
+
+    config: path to json config file (conform to TrainingConfig)
+    progress: enable tqdm console logging
+    tensorboard: enable tensorboard logging
+    profile: run profiling script for one epoch instead of training
+    """
+    model_dir = config.parent
+
+    if config is None:
+        model_dir = os.getcwd()
+        config = TrainingConfig(epochs=10, n_train=32, n_val=32, batch_size=16)
+
+    elif config.is_file():
+        model_dir = config.parent
+        with open(config, "r") as f:
+            config = json.load(f)
+            config = TrainingConfig(**config)
+
+    if profile:
+        profile_dgl(config)
+        return
+
+    hist = train_dgl(
+        config,
+        progress=progress,
+        checkpoint_dir=checkpoint_dir,
+        store_outputs=store_outputs,
+        log_tensorboard=tensorboard,
+    )
+
+    # print(model_dir)
+    # with open(model_dir / "metrics.json", "w") as f:
+    #     json.dump(hist, f)
+
+    torch.save(hist, model_dir / "metrics.pt")
+
+    with open(model_dir / "fullconfig.json", "w") as f:
+        json.dump(json.loads(config.json()), f, indent=2)
+
+    # move temporary checkpoint data into model_dir
+    for checkpoint in checkpoint_dir.glob("*.pt"):
+        shutil.copy(checkpoint, model_dir / checkpoint.name)
+
+
+if __name__ == "__main__":
+    typer.run(cli)

alignn/config.py

@@ -0,0 +1,253 @@
+"""Pydantic model for default configuration and validation."""
+
+import subprocess
+from typing import Optional, Union
+import os
+from typing import Literal
+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.alignn_atomwise_pure_smooth import (
+    ALIGNNAtomWisePureSmoothConfig,
+)
+from alignn.models.ealignn_atomwise import eALIGNNAtomWiseConfig
+
+# import torch
+
+try:
+    VERSION = (
+        subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()
+    )
+except Exception:
+    VERSION = "NA"
+    pass
+
+
+FEATURESET_SIZE = {"basic": 11, "atomic_number": 1, "cfid": 438, "cgcnn": 92}
+
+
+TARGET_ENUM = Literal[
+    "formation_energy_peratom",
+    "optb88vdw_bandgap",
+    "bulk_modulus_kv",
+    "shear_modulus_gv",
+    "mbj_bandgap",
+    "slme",
+    "magmom_oszicar",
+    "spillage",
+    "kpoint_length_unit",
+    "encut",
+    "optb88vdw_total_energy",
+    "epsx",
+    "epsy",
+    "epsz",
+    "mepsx",
+    "mepsy",
+    "mepsz",
+    "max_ir_mode",
+    "min_ir_mode",
+    "n-Seebeck",
+    "p-Seebeck",
+    "n-powerfact",
+    "p-powerfact",
+    "ncond",
+    "pcond",
+    "nkappa",
+    "pkappa",
+    "ehull",
+    "exfoliation_energy",
+    "dfpt_piezo_max_dielectric",
+    "dfpt_piezo_max_eij",
+    "dfpt_piezo_max_dij",
+    "gap pbe",
+    "e_form",
+    "e_hull",
+    "energy_per_atom",
+    "formation_energy_per_atom",
+    "band_gap",
+    "e_above_hull",
+    "mu_b",
+    "bulk modulus",
+    "shear modulus",
+    "elastic anisotropy",
+    "U0",
+    "HOMO",
+    "LUMO",
+    "R2",
+    "ZPVE",
+    "omega1",
+    "mu",
+    "alpha",
+    "homo",
+    "lumo",
+    "gap",
+    "r2",
+    "zpve",
+    "U",
+    "H",
+    "G",
+    "Cv",
+    "A",
+    "B",
+    "C",
+    "all",
+    "target",
+    "max_efg",
+    "avg_elec_mass",
+    "avg_hole_mass",
+    "_oqmd_band_gap",
+    "_oqmd_delta_e",
+    "_oqmd_stability",
+    "edos_up",
+    "pdos_elast",
+    "bandgap",
+    "energy_total",
+    "net_magmom",
+    "b3lyp_homo",
+    "b3lyp_lumo",
+    "b3lyp_gap",
+    "b3lyp_scharber_pce",
+    "b3lyp_scharber_voc",
+    "b3lyp_scharber_jsc",
+    "log_kd_ki",
+    "max_co2_adsp",
+    "min_co2_adsp",
+    "lcd",
+    "pld",
+    "void_fraction",
+    "surface_area_m2g",
+    "surface_area_m2cm3",
+    "indir_gap",
+    "f_enp",
+    "final_energy",
+    "energy_per_atom",
+    "ead",
+]
+
+
+class TrainingConfig(BaseSettings):
+    """Training config defaults and validation."""
+
+    version: str = VERSION
+
+    # dataset configuration
+    dataset: Literal[
+        "dft_3d",
+        "jdft_3d-8-18-2021",
+        "dft_2d",
+        "megnet",
+        "megnet2",
+        "mp_3d_2020",
+        "qm9",
+        "qm9_dgl",
+        "qm9_std_jctc",
+        "user_data",
+        "oqmd_3d_no_cfid",
+        "edos_up",
+        "edos_pdos",
+        "qmof",
+        "qe_tb",
+        "hmof",
+        "hpov",
+        "pdbbind",
+        "pdbbind_core",
+        "tinnet_OH",
+        "tinnet_O",
+        "tinnet_N",
+    ] = "dft_3d"
+    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",
+        "fast_graph",
+        "torch_graph",
+        "pure_torch",
+    ] = "k-nearest"
+    id_tag: Literal["jid", "id", "_oqmd_entry_id"] = "jid"
+
+    # training configuration
+    dtype: str = "float32"
+    random_seed: Optional[int] = 123
+    classification_threshold: Optional[float] = None
+    # target_range: Optional[List] = None
+    n_val: Optional[int] = None
+    n_test: Optional[int] = None
+    n_train: Optional[int] = None
+    train_ratio: Optional[float] = 0.8
+    val_ratio: Optional[float] = 0.1
+    test_ratio: Optional[float] = 0.1
+    target_multiplication_factor: Optional[float] = None
+    epochs: int = 300
+    batch_size: int = 64
+    gpu_memory_fraction: Optional[float] = None
+    use_amp: bool = False  # bf16 mixed precision (A100/H100/RTX 30+)
+    # DDP tuning. find_unused_parameters=True is slow; enable only if
+    # your model has conditional branches whose gradients vary per step.
+    ddp_find_unused_parameters: bool = False
+    # Resume optimizer/scheduler/epoch from <output_dir>/current_state.pt if
+    # present. Lets a run that hit the walltime continue with a single
+    # continuous LR schedule instead of restarting it (weights alone are
+    # restored via --restart_model_path, which does not restore these).
+    resume_checkpoint: bool = False
+    # LR-schedule horizon in epochs, independent of how many epochs this job
+    # runs (`epochs`). Set it to the FINAL target when a run is split into
+    # resumed segments so OneCycle spans the whole run and stays continuous
+    # across restarts. None => use `epochs` (the single-job default).
+    lr_total_epochs: Optional[int] = None
+    # When True, forces cuDNN determinism (slower). Decoupled from seed so
+    # you can seed for reproducibility without paying the speed cost.
+    deterministic: bool = False
+    weight_decay: float = 0
+    learning_rate: float = 1e-2
+    filename: str = "sample"
+    warmup_steps: int = 2000
+    criterion: Literal["mse", "l1", "poisson", "zig"] = "mse"
+    optimizer: Literal["adamw", "sgd"] = "adamw"
+    scheduler: Literal["onecycle", "none"] = "onecycle"
+    pin_memory: bool = False
+    save_dataloader: bool = False
+    write_checkpoint: bool = True
+    write_predictions: bool = True
+    store_outputs: bool = True
+    progress: bool = True
+    log_tensorboard: bool = False
+    standard_scalar_and_pca: bool = False
+    use_canonize: bool = True
+    compute_line_graph: bool = True
+    num_workers: int = 4
+    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.
+    three_body_cutoff: Optional[float] = 3.5
+    max_neighbors: Optional[int] = 12
+    keep_data_order: bool = True
+    normalize_graph_level_loss: bool = False
+    distributed: bool = False
+    data_parallel: bool = False
+    n_early_stopping: Optional[int] = None  # typically 50
+    output_dir: str = os.path.abspath(".")
+    use_lmdb: bool = True
+    # When True, reuse an existing LMDB cache on disk (fast, but the
+    # cache must match the current model backend + graph config).
+    # Default False: always rebuild from scratch for safety.
+    read_existing: bool = False
+    # alignn_layers: int = 4
+    # gcn_layers: int =4
+    # edge_input_features: int= 80
+    # hidden_features: int= 256
+    # triplet_input_features: int=40
+    # embedding_features: int=64
+
+    # model configuration
+    model: Union[
+        ALIGNNConfig,
+        ALIGNNAtomWiseConfig,
+        ALIGNNAtomWisePureConfig,
+        ALIGNNAtomWisePureSmoothConfig,
+        eALIGNNAtomWiseConfig,
+    ] = ALIGNNAtomWiseConfig(name="alignn_atomwise")