Scandium-Labs/Scandium-Dataset

1"""Parallel BVSE processor — processes Li/Na entries using multiprocessing.2Launched from nohup; checkpoints every N entries per worker.3"""4import json, os, sys, time, argparse, warnings5from pathlib import Path6from collections import defaultdict7from multiprocessing import Pool, cpu_count8import numpy as np9 10BASE_DIR = Path(__file__).resolve().parent.parent11sys.path.insert(0, str(BASE_DIR / "dataset"))12os.environ["OPENBLAS_NUM_THREADS"] = "1"13os.environ["OMP_NUM_THREADS"] = "1"14os.environ["MKL_NUM_THREADS"] = "1"15warnings.filterwarnings("ignore")16 17PARQUET_PATH = BASE_DIR / "dataset" / "entries_v3.parquet"18INDEX_PATH = BASE_DIR / "dataset" / "entries_v3.index.json"19CHECKPOINT_PATH = BASE_DIR / "dataset" / "bvse_checkpoint.json"20 21MOBILITY_THRESHOLDS = [22    ("superionic", 0.25),23    ("good", 0.50),24    ("moderate", 0.75),25    ("poor", float('inf')),26]27 28 29def process_one(args):30    sid, struct_json, mobile_ion = args31    from pymatgen.core import Structure32    from pymatgen.analysis.bond_valence import BVAnalyzer33    from bvlain import Lain34    result = {"sid": sid, "barrier": None, "cls": "none", "dim": "none", "skip": ""}35    try:36        struct = Structure.from_dict(json.loads(struct_json))37        bva = BVAnalyzer()38        st_oxi = bva.get_oxi_state_decorated_structure(struct)39        lain = Lain(verbose=False)40        lain.read_structure(st_oxi, oxi_check=False)41        lain.bvse_distribution(mobile_ion=f"{mobile_ion}1+", resolution=0.5)42        barriers = lain.percolation_barriers(encut=10.0, n_jobs=1)43        min_barrier = min(v for v in barriers.values())44        if min_barrier == float('inf') or min_barrier < 0:45            result["skip"] = "no connected pathway"46            return result47        eps = 0.0548        dim_map = {1: "1D", 2: "2D", 3: "3D"}49        dims = []50        for d in [1, 2, 3]:51            if barriers[f"E_{d}D"] <= min_barrier + eps:52                dims.append(d)53        percolation_dim = dim_map.get(max(dims), "3D") if dims else "none"54        mobility_class = "poor"55        for cls_name, threshold in MOBILITY_THRESHOLDS:56            if min_barrier < threshold:57                mobility_class = cls_name58                break59        result["barrier"] = round(float(min_barrier), 4)60        result["cls"] = mobility_class61        result["dim"] = percolation_dim62    except Exception as exc:63        err = str(exc)64        if "No BVSE data" in err:65            result["skip"] = "no BV params"66        elif "oxidation state" in err.lower():67            result["skip"] = "oxidation failed"68        elif "min() iterable" in err:69            result["skip"] = "no path"70        else:71            result["skip"] = f"error: {err[:80]}"72    return result73 74 75def _write_checkpoint(table, updated_rows, all_ssb, sid_to_idx, all_table_rows):76    """Write updated ssb_screening values back to the Parquet table."""77    import pyarrow as pa78    import pyarrow.parquet as pq79    from dataset_store import _encode_value, _decode_value80    ssb_col = table.column("ssb_screening").to_pylist()81    for table_row in updated_rows:82        # Find the corresponding index in all_* arrays83        sid = _decode_value(table.column("source_id")[table_row].as_py())84        if sid in sid_to_idx:85            ssb_col[table_row] = _encode_value(all_ssb[sid_to_idx[sid]])86    table = table.set_column(87        table.schema.get_field_index("ssb_screening"), "ssb_screening",88        pa.chunked_array([pa.array(ssb_col)])89    )90    pq.write_table(table, PARQUET_PATH, compression="zstd")91 92 93def main():94    parser = argparse.ArgumentParser()95    parser.add_argument("--workers", type=int, default=max(1, cpu_count() - 1))96    parser.add_argument("--batch-size", type=int, default=5000)97    parser.add_argument("--max-sites", type=int, default=60)98    args = parser.parse_args()99 100    print(f"Parallel BVSE processor | workers={args.workers} batch={args.batch_size} max_sites={args.max_sites}")101    print(f"Loading Parquet...")102    import pyarrow.parquet as pq103    t0 = time.time()104    table = pq.read_table(PARQUET_PATH)105    # Decode columns106    from dataset_store import _decode_value, _encode_value107    all_sids = []108    all_structs = []109    all_mobiles = []110    all_ssb = []111    all_rows = []112    for i in range(table.num_rows):113        sid = _decode_value(table.column("source_id")[i].as_py())114        mobile = _decode_value(table.column("mobile_ion")[i].as_py())115        nsites_raw = table.column("nsites")[i].as_py()116        nsites = int(_decode_value(nsites_raw)) if nsites_raw else 999117        struct_raw = table.column("structure_json")[i].as_py()118        struct = _decode_value(struct_raw) if struct_raw else None119        ssb_raw = table.column("ssb_screening")[i].as_py()120        ssb = _decode_value(ssb_raw) if ssb_raw else {}121        if mobile in ("Li", "Na") and struct and nsites <= args.max_sites:122            all_sids.append(sid)123            all_structs.append(struct)124            all_mobiles.append(mobile)125            all_ssb.append(ssb)126            all_rows.append(i)127    print(f"  {len(all_sids):,} Li/Na entries (≤{args.max_sites} sites) loaded in {time.time()-t0:.1f}s")128 129    # Build lookup: source_id -> index in all_* arrays130    sid_to_idx = {sid: i for i, sid in enumerate(all_sids)}131    all_table_rows = all_rows  # table row indices for each entry132    updated_rows = set()133 134    # Load checkpoint135    completed_sids = set()136    if CHECKPOINT_PATH.exists():137        with open(CHECKPOINT_PATH) as f:138            completed_sids = set(json.load(f))139        print(f"  Resuming from checkpoint: {len(completed_sids):,} already computed")140 141    # Build worklist (skip already done)142    worklist = []143    for idx in range(len(all_sids)):144        if all_sids[idx] not in completed_sids:145            worklist.append((all_sids[idx], all_structs[idx], all_mobiles[idx]))146    print(f"  Remaining: {len(worklist):,} entries")147 148    if not worklist:149        print("  All entries already processed.")150        return151 152    # Process in batches153    total_processed = len(completed_sids)154    total_ok = 0155    total_skip = 0156    total_err = 0157    barriers = []158    class_counts = defaultdict(int)159 160    batch_num = 0161    checkpoint_counter = 0162    CHECKPOINT_EVERY_N_BATCHES = 50  # checkpoint every 50 batches (every ~2500 entries at batch=50)163    pool = Pool(processes=args.workers)164    t_start = time.time()165    while worklist:166        batch = worklist[:args.batch_size]167        worklist = worklist[args.batch_size:]168        print(f"  Batch {batch_num}: {len(batch)} entries, {len(worklist)} remaining...", flush=True)169        results = pool.map(process_one, batch)170        # Apply results171        for r in results:172            sid = r["sid"]173            if sid in sid_to_idx:174                idx = sid_to_idx[sid]175                ssb = all_ssb[idx]176                if r["barrier"] is not None:177                    ssb["bvse_migration_barrier_eV"] = r["barrier"]178                    ssb["bvse_mobility_class"] = r["cls"]179                    ssb["bvse_percolation_dimensionality"] = r["dim"]180                    barriers.append(r["barrier"])181                    class_counts[r["cls"]] += 1182                    total_ok += 1183                elif r["skip"]:184                    total_skip += 1185                else:186                    total_err += 1187                completed_sids.add(sid)188                updated_rows.add(all_table_rows[idx])189        total_processed += len(batch)190        batch_num += 1191 192        # Periodic checkpoint to Parquet193        checkpoint_counter += 1194        if checkpoint_counter >= CHECKPOINT_EVERY_N_BATCHES or not worklist:195            _write_checkpoint(table, updated_rows, all_ssb, sid_to_idx, all_table_rows)196            # Save completed SIDs for resume197            with open(CHECKPOINT_PATH, "w") as f:198                json.dump(list(completed_sids), f)199            rate = total_processed / (time.time() - t_start)200            eta = len(worklist) / rate if rate > 0 and worklist else 0201            print(f"    CHECKPOINT: {total_processed:,} done, {len(worklist):,} remain, "202                  f"{rate:.1f} ent/s, ETA {eta/3600:.1f}h", flush=True)203            checkpoint_counter = 0204 205    pool.close()206    pool.join()207 208    elapsed = time.time() - t_start209    print(f"\n{'=' * 60}")210    print(f"  BVSE COMPLETE")211    print(f"  Processed: {total_processed:,}")212    print(f"  Barriers computed: {total_ok}")213    print(f"  Skipped: {total_skip}")214    print(f"  Errors: {total_err}")215    print(f"  Time: {elapsed/60:.1f} min ({total_processed/elapsed:.1f} ent/s)")216    print(f"{'=' * 60}")217 218 219if __name__ == "__main__":220    main()221