Scandium-Labs/Scandium-Dataset

1"""Set up MLIP infrastructure for high-throughput migration barrier computation.2 3Installs and validates MLIP tools for nudged elastic band (NEB) calculations:4  - CHGNet: universal crystal Hamiltonian Graph neural Network5  - MACE-MP-0: MACE architecture trained on Materials Project trajectories6  - M3GNet: universal potential from Materials Project7  - Orb-v3: Orbital-based MLIP8 9This script:10  1. Checks what's installed11  2. Attempts installation of missing packages12  3. Validates each potential on a test structure13  4. Generates a configuration file for the NEB pipeline14 15Usage:16    python scripts/setup_mlip_infrastructure.py17    python scripts/setup_mlip_infrastructure.py --check-only18    python scripts/setup_mlip_infrastructure.py --install19"""20import argparse, os, sys, subprocess, json, warnings21from pathlib import Path22 23MLIP_PACKAGES = {24    "chgnet": "chgnet",25    "mace": "mace-torch",26    "matgl": "matgl",27    "orb": "orb-models",28}29 30TEST_STRUCTURE = """31{32    "@module": "pymatgen.core.structure",33    "@class": "Structure",34    "lattice": {"matrix": [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 3.0]], "pbc": [true, true, true]},35    "sites": [36        {"species": [{"element": "Li", "occu": 1}], "abc": [0.0, 0.0, 0.0]},37        {"species": [{"element": "Cl", "occu": 1}], "abc": [0.5, 0.5, 0.5]}38    ]39}40"""41 42 43def check_installed():44    """Check which MLIP packages are installed."""45    results = {}46    for name, pkg in MLIP_PACKAGES.items():47        try:48            __import__(name.replace("-", "_"))49            results[name] = "installed"50        except ImportError:51            try:52                __import__(pkg.replace("-", "_"))53                results[name] = "installed"54            except ImportError:55                results[name] = "not found"56    return results57 58 59def install_packages(packages):60    """Install MLIP packages via pip."""61    for name, pkg in packages.items():62        print(f"  Installing {pkg}...")63        result = subprocess.run(64            [sys.executable, "-m", "pip", "install", pkg],65            capture_output=True, text=True66        )67        if result.returncode == 0:68            print(f"    {name}: installed")69        else:70            print(f"    {name}: failed — {result.stderr[-200:]}")71 72 73def validate_chgnet(structure_dict):74    """Validate CHGNet can predict on test structure."""75    import json76    from pymatgen.core import Structure77    from chgnet.model import CHGNet78    from chgnet.utils import write_structures_to_POSCAR79    80    struct = Structure.from_dict(structure_dict)81    model = CHGNet.load()82    prediction = model.predict_structure(struct)83    return {84        "energy": float(prediction["e"]),85        "forces_shape": list(prediction["f"].shape),86    }87 88 89def validate_mace(structure_dict):90    """Validate MACE can predict on test structure."""91    import torch92    from mace.calculators import MACECalculator93    from ase.io import read94    from pymatgen.core import Structure95    from pymatgen.io.ase import AseAtomsAdaptor96    97    struct = Structure.from_dict(structure_dict)98    atoms = AseAtomsAdaptor.get_atoms(struct)99    100    calc = MACECalculator(model_path="medium", device="cpu")101    atoms.set_calculator(calc)102    energy = atoms.get_potential_energy()103    forces = atoms.get_forces()104    105    return {106        "energy": float(energy),107        "forces_shape": list(forces.shape),108    }109 110 111def main():112    parser = argparse.ArgumentParser(description="MLIP infrastructure setup")113    parser.add_argument("--check-only", action="store_true",114                        help="Check installed packages only")115    parser.add_argument("--install", action="store_true",116                        help="Install missing MLIP packages")117    parser.add_argument("--validate", action="store_true",118                        help="Validate installed potentials on test structure")119    args = parser.parse_args()120    121    BASE_DIR = Path(__file__).resolve().parent.parent122    123    print("=" * 60)124    print("  MLIP INFRASTRUCTURE SETUP")125    print("  High-throughput migration barrier computation pipeline")126    print("=" * 60)127    128    # Check installed packages129    print("\n  Checking installed MLIP packages...")130    installed = check_installed()131    for name, status in installed.items():132        print(f"    {name:12s}: {status}")133    134    if args.install:135        to_install = {k: v for k, v in MLIP_PACKAGES.items() if installed[k] == "not found"}136        if to_install:137            print(f"\n  Installing {len(to_install)} packages...")138            install_packages(to_install)139        else:140            print("\n  All packages already installed.")141    142    if args.validate:143        print("\n  Validating potentials...")144        struct_dict = json.loads(TEST_STRUCTURE)145        146        if installed.get("chgnet") == "installed":147            try:148                result = validate_chgnet(struct_dict)149                print(f"    CHGNet: OK (energy={result['energy']:.3f} eV)")150            except Exception as e:151                print(f"    CHGNet: validation failed — {str(e)[:80]}")152        153        if installed.get("mace") == "installed":154            try:155                result = validate_mace(struct_dict)156                print(f"    MACE: OK (energy={result['energy']:.3f} eV)")157            except Exception as e:158                print(f"    MACE: validation failed — {str(e)[:80]}")159    160    # Generate config file161    if not args.check_only:162        config = {163            "potentials": installed,164            "pipeline": {165                "bvse_barrier_threshold": 0.5,166                "mlip_neb_grid": [5, 5, 5],167                "mlip_neb_spring_constant": 5.0,168                "mlip_neb_fmax": 0.05,169                "mlip_neb_steps": 500,170            },171            "target_subset": "gold_battery_li",172            "description": "Li-containing Gold-tier battery-family entries",173        }174        175        config_path = BASE_DIR / "configs" / "mlip_pipeline.json"176        print(f"\n  Writing config to {config_path}...")177        config_path.parent.mkdir(parents=True, exist_ok=True)178        with open(config_path, "w") as f:179            json.dump(config, f, indent=2)180    181    # Print next steps182    print(f"\n{'─' * 60}")183    print("  NEXT STEPS")184    print(f"  {'─' * 60}")185    print("""186  1. Install MLIP packages:187     pip install chgnet mace-torch matgl orb-models188     189  2. Run BVSE pre-filter on Li/Na entries:190     python scripts/compute_bvse_barriers.py --subset gold --limit 50000191     192  3. Run MLIP-NEB on BVSE-filtered subset:193     python scripts/run_mlip_neb_pipeline.py --input dataset/bvse_filtered.json194     195  4. Update sse_candidate_score with full 5 gates:196     python scripts/compute_sse_candidate_score.py197  """)198    print("=" * 60)199 200 201if __name__ == "__main__":202    main()203