cpu_generation_benchmark/run_cpu_generation_benchmark.py

#!/usr/bin/env python3
"""CPU-only parallel benchmark for PXRDNet XRD-to-structure generation."""

from __future__ import annotations

import argparse
import json
import math
import multiprocessing as mp
import os
import platform
import sys
import threading
import time
from pathlib import Path
from types import SimpleNamespace

import hydra
import numpy as np
import pandas as pd
import psutil
import torch
import torch.nn.functional as F
from pymatgen.analysis.structure_matcher import StructureMatcher
from torch.distributions import MultivariateNormal
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from tqdm import tqdm

ROOT = Path(__file__).resolve().parents[1]
SCRIPTS = ROOT / "scripts"
for path in (ROOT, SCRIPTS):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from compute_metrics import Crystal  # noqa: E402
from eval_utils import load_config  # noqa: E402
from evaluate_pxrdgen_match_only import (  # noqa: E402
    EPS,
    PXRDGEN_MATCHER_KWARGS,
    choose_checkpoint,
    crystal_match,
    json_safe,
    make_structures,
    save_crystal_cif,
)

WORKER_STATE = {}
PREPARED_BATCHES = []


def monitor_cpu_usage(parent_pid: int, output_dir: Path, stop_event: threading.Event, interval: float) -> None:
    timing_dir = output_dir / "timing"
    timing_dir.mkdir(parents=True, exist_ok=True)
    usage_path = timing_dir / "cpu_usage.tsv"
    parent = psutil.Process(parent_pid)

    def live_processes():
        processes = [parent]
        try:
            processes.extend(parent.children(recursive=True))
        except psutil.Error:
            pass
        return processes

    with usage_path.open("w") as f:
        f.write("elapsed_seconds\tworker_count\ttotal_process_cpu_seconds_delta\tinterval_seconds\ttotal_process_cpu_percent\tequivalent_logical_cores\tsystem_cpu_percent\n")
        start = time.perf_counter()
        last_wall = start
        last_cpu_by_pid = {}
        while not stop_event.wait(interval):
            now = time.perf_counter()
            interval_seconds = max(now - last_wall, 1e-9)
            current_cpu_by_pid = {}
            worker_count = 0
            for proc in live_processes():
                try:
                    times = proc.cpu_times()
                    current_cpu_by_pid[proc.pid] = times.user + times.system
                except psutil.Error:
                    continue
                if proc.pid != parent_pid:
                    worker_count += 1
            cpu_delta = 0.0
            for pid, current_cpu in current_cpu_by_pid.items():
                previous_cpu = last_cpu_by_pid.get(pid)
                if previous_cpu is not None:
                    cpu_delta += max(0.0, current_cpu - previous_cpu)
            total_percent = 100.0 * cpu_delta / interval_seconds
            elapsed = time.perf_counter() - start
            f.write(
                f"{elapsed:.6f}\t{worker_count}\t{cpu_delta:.6f}\t{interval_seconds:.6f}\t{total_percent:.6f}\t"
                f"{total_percent / 100.0:.6f}\t{psutil.cpu_percent(interval=None):.6f}\n"
            )
            f.flush()
            last_wall = now
            last_cpu_by_pid = current_cpu_by_pid


def summarize_cpu_usage(output_dir: Path) -> dict:
    usage_path = output_dir / "timing" / "cpu_usage.tsv"
    if not usage_path.exists():
        return {}
    rows = []
    with usage_path.open() as f:
        header = f.readline().strip().split("\t")
        for line in f:
            if line.strip():
                rows.append(dict(zip(header, line.strip().split("\t"))))
    if not rows:
        return {}
    total_cpu = np.array([float(row["total_process_cpu_percent"]) for row in rows], dtype=float)
    worker_counts = np.array([int(row["worker_count"]) for row in rows], dtype=int)
    return {
        "cpu_usage_samples": int(len(rows)),
        "cpu_usage_path": str(usage_path),
        "mean_process_cpu_percent": float(np.mean(total_cpu)),
        "max_process_cpu_percent": float(np.max(total_cpu)),
        "mean_equivalent_logical_cores": float(np.mean(total_cpu / 100.0)),
        "max_equivalent_logical_cores": float(np.max(total_cpu / 100.0)),
        "max_observed_worker_count": int(np.max(worker_counts)),
    }


def _set_cpu_env(num_threads: int) -> None:
    value = str(max(1, int(num_threads)))
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    os.environ["OMP_NUM_THREADS"] = value
    os.environ["MKL_NUM_THREADS"] = value
    os.environ["OPENBLAS_NUM_THREADS"] = value
    os.environ["NUMEXPR_NUM_THREADS"] = value
    torch.set_num_threads(int(value))
    torch.set_num_interop_threads(1)


def load_model_and_loader_cpu(args):
    model = load_model_cpu(args)
    _, test_loader = load_limited_test_loader(args)
    return model, test_loader


def normalize_cfg_paths(cfg, args):
    if args.data_root_override:
        cfg.data.root_path = args.data_root_override
    elif not Path(str(cfg.data.root_path)).exists():
        cfg.data.root_path = str(ROOT / "data" / Path(str(cfg.data.root_path)).name)
    cfg.data.datamodule.batch_size.test = 1
    if "decoder" in cfg.model and "scale_file" in cfg.model.decoder:
        cfg.model.decoder.scale_file = str(ROOT / "cdvae" / "pl_modules" / "gemnet" / "gemnet-dT.json")
    return cfg


def load_model_cpu(args):
    model_path = Path(args.model_path)
    cfg = normalize_cfg_paths(load_config(model_path), args)

    model = hydra.utils.instantiate(
        cfg.model,
        optim=cfg.optim,
        data=cfg.data,
        logging=cfg.logging,
        _recursive_=False,
    )
    ckpt = choose_checkpoint(model_path)
    checkpoint = torch.load(str(ckpt), map_location=torch.device("cpu"))
    model.load_state_dict(checkpoint["state_dict"], strict=True)
    model.lattice_scaler = torch.load(model_path / "lattice_scaler.pt", map_location=torch.device("cpu"))
    model.to("cpu")
    model.eval()
    model.freeze()
    return model


def load_limited_test_loader(args):
    model_path = Path(args.model_path)
    cfg = normalize_cfg_paths(load_config(model_path), args)
    limited_root = Path(args.output_dir) / "limited_data"
    limited_root.mkdir(parents=True, exist_ok=True)
    source_test_path = Path(str(cfg.data.datamodule.datasets.test[0].path))
    if not source_test_path.exists():
        source_test_path = Path(str(cfg.data.root_path)) / "test.csv"
    limited_test_path = limited_root / f"test_first{args.first_idx}_n{args.num_materials}.pkl"

    if not limited_test_path.exists() or args.rebuild_limited_data:
        df = pd.read_pickle(source_test_path)
        subset = df.iloc[args.first_idx: args.first_idx + args.num_materials].copy()
        if len(subset) != args.num_materials:
            raise ValueError(
                f"Requested {args.num_materials} materials from {source_test_path}, got {len(subset)}"
            )
        subset.to_pickle(limited_test_path)

    cfg.data.root_path = str(limited_root)
    cfg.data.datamodule.datasets.test[0].path = str(limited_test_path)
    cfg.data.datamodule.datasets.test[0].preprocess_workers = min(
        int(cfg.data.datamodule.datasets.test[0].preprocess_workers),
        max(1, args.num_materials),
    )

    datamodule = hydra.utils.instantiate(
        cfg.data.datamodule, _recursive_=False, scaler_path=model_path
    )
    datamodule.setup("test")
    return cfg, datamodule.test_dataloader()[0]


def calculate_accuracy(probabilities, labels):
    _, predicted_classes = torch.max(probabilities, dim=1)
    correct_predictions = (predicted_classes == labels).float()
    return (correct_predictions.sum() / labels.size(0)).item()


def optimize_latent_code_cpu(args, model, batch, target_xrd, material_index: int):
    m = MultivariateNormal(
        torch.zeros(model.hparams.hidden_dim, device="cpu"),
        torch.eye(model.hparams.hidden_dim, device="cpu"),
    )
    z = torch.randn(args.num_starting_points, model.hparams.hidden_dim, device="cpu")
    z.requires_grad = True
    opt = Adam([z], lr=args.lr)
    total_gradient_steps = args.num_gradient_steps * (1 + 2 + 4) - 1
    scheduler = CosineAnnealingWarmRestarts(opt, args.num_gradient_steps, T_mult=2, eta_min=args.min_lr)

    disable_bar = args.disable_bar or args.workers > 1
    with tqdm(
        total=total_gradient_steps,
        desc=f"material{material_index} latent opt",
        unit="steps",
        mininterval=args.progress_mininterval,
        disable=disable_bar,
    ) as pbar:
        for step_idx in range(total_gradient_steps):
            opt.zero_grad()
            pred_xrd = model.fc_property(z)
            target = target_xrd.broadcast_to(z.shape[0], target_xrd.shape[-1])
            xrd_loss = F.l1_loss(pred_xrd, target) if args.l1_loss else F.mse_loss(pred_xrd, target)
            prob = m.log_prob(z).mean()
            pred_num_atoms, _, _, _, pred_composition_per_atom = model.decode_stats(
                z, batch.num_atoms, batch.lengths, batch.angles, teacher_forcing=False
            )
            repeated_num_atoms = batch.num_atoms.repeat(args.num_starting_points)
            repeated_atom_types = (batch.atom_types - 1).repeat(args.num_starting_points)
            num_atom_loss = F.cross_entropy(pred_num_atoms, repeated_num_atoms)
            composition_loss = F.cross_entropy(pred_composition_per_atom, repeated_atom_types)
            if step_idx % args.progress_log_interval == 0 or step_idx == total_gradient_steps - 1:
                num_atom_accuracy = calculate_accuracy(pred_num_atoms, repeated_num_atoms)
                composition_accuracy = calculate_accuracy(pred_composition_per_atom, repeated_atom_types)
                pbar.set_postfix_str(
                    f"xrd={xrd_loss.item():.3e}; logp={prob.item():.3e}; "
                    f"num={num_atom_loss.item():.3e}/{num_atom_accuracy:.3f}; "
                    f"comp={composition_loss.item():.3e}/{composition_accuracy:.3f}",
                    refresh=True,
                )
            pbar.update(1)
            total_loss = (
                xrd_loss
                - args.l2_penalty * prob
                + args.num_atom_lambda * num_atom_loss
                + args.composition_lambda * composition_loss
            )
            total_loss.backward()
            opt.step()
            scheduler.step()
    return z


def init_worker(args_dict: dict, torch_threads: int):
    _set_cpu_env(torch_threads)
    args = SimpleNamespace(**args_dict)
    model = load_model_cpu(args)
    WORKER_STATE["args"] = args
    WORKER_STATE["model"] = model
    WORKER_STATE["matcher"] = StructureMatcher(**PXRDGEN_MATCHER_KWARGS)


def generate_material(local_index: int) -> dict:
    args = WORKER_STATE["args"]
    model = WORKER_STATE["model"]
    matcher = WORKER_STATE["matcher"]
    material_index = args.first_idx + local_index
    batch = PREPARED_BATCHES[local_index]

    start = time.perf_counter()
    batch = batch.to("cpu")
    mpid = batch.mpid[0]
    formula = batch.pretty_formula[0]
    material_dir = Path(args.output_dir) / "cpu_sample" / f"material{material_index}_{mpid}_{formula}"
    material_dir.mkdir(parents=True, exist_ok=True)

    xrd_dim = args.n_postsubsample
    target_xrd = batch.y.reshape(1, xrd_dim)
    z = optimize_latent_code_cpu(args, model, batch, target_xrd, material_index)

    init_num_atoms = batch.num_atoms.repeat(args.num_starting_points) if args.num_atom_lambda > EPS else None
    init_atom_types = batch.atom_types.repeat(args.num_starting_points) if args.composition_lambda > EPS else None
    dynamics_start = time.perf_counter()
    crystals = model.langevin_dynamics(
        z,
        SimpleNamespace(
            n_step_each=args.n_step_each,
            step_lr=args.step_lr,
            min_sigma=args.min_sigma,
            save_traj=False,
            disable_bar=True,
        ),
        gt_num_atoms=init_num_atoms,
        gt_atom_types=init_atom_types,
    )
    dynamics_seconds = time.perf_counter() - dynamics_start
    crystals = {k: crystals[k] for k in ["frac_coords", "atom_types", "num_atoms", "lengths", "angles"]}

    _pred_coords, _pred_atom_types, pred_crystal_dicts = make_structures(
        args,
        crystals["frac_coords"],
        crystals["num_atoms"],
        crystals["atom_types"],
        crystals["lengths"],
        crystals["angles"],
    )
    _gt_coords, _gt_atom_types, gt_crystal_dicts = make_structures(
        args,
        batch.frac_coords,
        batch.num_atoms,
        batch.atom_types,
        batch.lengths,
        batch.angles,
    )
    gt_crystal = Crystal(gt_crystal_dicts[0])
    save_crystal_cif(
        gt_crystal,
        material_dir / "gt" / "cif" / f"noSpacegroup_material{material_index}_{mpid}_{formula}.cif",
        material_dir / "gt" / "cif" / f"material{material_index}_{mpid}_{formula}.cif",
    )

    candidate_results = []
    target_match = False
    target_rms_values = []
    for i, pred_dict in enumerate(pred_crystal_dicts):
        pred_crystal = Crystal(pred_dict)
        cand_dir = material_dir / "pred" / f"candidate{i}" / "cif"
        save_crystal_cif(
            pred_crystal,
            cand_dir / f"noSpacegroup_material{material_index}_candidate{i}.cif",
            cand_dir / f"material{material_index}_candidate{i}.cif",
        )
        match_info = crystal_match(pred_crystal, gt_crystal, matcher)
        if match_info["match"]:
            target_match = True
            target_rms_values.append(match_info["rms_dist"])
        candidate_results.append({"candidate_index": i, **match_info})

    elapsed_seconds = time.perf_counter() - start
    best_rms = min(target_rms_values) if target_rms_values else None
    material_metrics = {
        "material_index": material_index,
        "mpid": mpid,
        "formula": formula,
        "num_candidates": len(candidate_results),
        "target_match": bool(target_match),
        "best_rms_dist": best_rms,
        "candidate_match_status": candidate_results,
        "matcher": PXRDGEN_MATCHER_KWARGS,
        "timing": {
            "elapsed_seconds": elapsed_seconds,
            "langevin_seconds": dynamics_seconds,
            "candidates_per_second": len(candidate_results) / elapsed_seconds if elapsed_seconds > 0 else None,
        },
        "cpu": {
            "pid": os.getpid(),
            "torch_threads": torch.get_num_threads(),
        },
    }
    with (material_dir / "metrics.json").open("w") as f:
        json.dump(json_safe(material_metrics), f, indent=2)
    return json_safe(material_metrics)


def positive_or_auto(value: int | None) -> int | None:
    if value is None or value <= 0:
        return None
    return value


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--data-root-override", default="")
    parser.add_argument("--first-idx", type=int, default=0)
    parser.add_argument("--num-materials", type=int, default=20)
    parser.add_argument("--num-starting-points", type=int, default=1)
    parser.add_argument("--workers", type=int, default=0, help="0 means auto.")
    parser.add_argument("--torch-threads-per-worker", type=int, default=0, help="0 means auto.")
    parser.add_argument("--num-gradient-steps", type=int, default=5000)
    parser.add_argument("--n-step-each", type=int, default=100)
    parser.add_argument("--step-lr", type=float, default=1e-4)
    parser.add_argument("--min-sigma", type=float, default=0)
    parser.add_argument("--disable-bar", action="store_true")
    parser.add_argument("--lr", type=float, default=0.1)
    parser.add_argument("--min-lr", type=float, default=1e-4)
    parser.add_argument("--l2-penalty", type=float, default=2e-4)
    parser.add_argument("--num-atom-lambda", type=float, default=0.1)
    parser.add_argument("--composition-lambda", type=float, default=0.1)
    parser.add_argument("--l1-loss", action="store_true", default=True)
    parser.add_argument("--wave-source", default="CuKa")
    parser.add_argument("--xrd-vector-dim", type=int, default=900)
    parser.add_argument("--min-theta", type=float, default=0)
    parser.add_argument("--max-theta", type=float, default=90)
    parser.add_argument("--progress-log-interval", type=int, default=500)
    parser.add_argument("--progress-mininterval", type=float, default=10.0)
    parser.add_argument("--rebuild-limited-data", action="store_true")
    parser.add_argument("--cpu-sample-interval", type=float, default=5.0)
    args = parser.parse_args()

    total_cpus = os.cpu_count() or 1
    workers = positive_or_auto(args.workers)
    if workers is None:
        workers = min(args.num_materials, total_cpus)
    workers = max(1, min(workers, args.num_materials))
    torch_threads = positive_or_auto(args.torch_threads_per_worker)
    if torch_threads is None:
        torch_threads = max(1, math.ceil(total_cpus / workers))

    args.workers = workers
    args.torch_threads_per_worker = torch_threads
    args.output_dir = str(Path(args.output_dir).resolve())
    Path(args.output_dir, "cpu_sample").mkdir(parents=True, exist_ok=True)
    Path(args.output_dir, "timing").mkdir(parents=True, exist_ok=True)

    print("Preparing limited test data in the main process...")
    _setup_start = time.perf_counter()
    _cfg, test_loader = load_limited_test_loader(args)
    batches = list(test_loader)
    if len(batches) != args.num_materials:
        raise ValueError(f"Expected {args.num_materials} prepared batches, got {len(batches)}")
    args.n_postsubsample = int(test_loader.dataset.n_postsubsample)
    data_setup_seconds = time.perf_counter() - _setup_start

    material_indices = list(range(args.first_idx, args.first_idx + args.num_materials))
    global PREPARED_BATCHES
    PREPARED_BATCHES = batches
    config = {
        "model_path": str(Path(args.model_path).resolve()),
        "output_dir": args.output_dir,
        "first_idx": args.first_idx,
        "num_materials": args.num_materials,
        "num_starting_points": args.num_starting_points,
        "cpu_count_logical": total_cpus,
        "workers": workers,
        "torch_threads_per_worker": torch_threads,
        "estimated_torch_threads_total": workers * torch_threads,
        "data_setup_seconds": data_setup_seconds,
        "limited_data_note": "Only first_idx:num_materials rows are written to a temporary pickle before CrystDataset preprocessing.",
        "python": sys.version,
        "platform": platform.platform(),
        "args": vars(args),
    }
    with Path(args.output_dir, "parameters.json").open("w") as f:
        json.dump(json_safe(config), f, indent=2)

    print(json.dumps(json_safe(config), indent=2))
    wall_start = time.perf_counter()
    stop_cpu_monitor = threading.Event()
    cpu_monitor = threading.Thread(
        target=monitor_cpu_usage,
        args=(os.getpid(), Path(args.output_dir), stop_cpu_monitor, args.cpu_sample_interval),
        daemon=True,
    )
    cpu_monitor.start()
    ctx = mp.get_context("fork")
    try:
        with ctx.Pool(
            processes=workers,
            initializer=init_worker,
            initargs=(vars(args), torch_threads),
        ) as pool:
            results = list(tqdm(pool.imap_unordered(generate_material, range(len(batches))), total=len(batches), desc="CPU materials"))
    finally:
        stop_cpu_monitor.set()
        cpu_monitor.join(timeout=max(1.0, args.cpu_sample_interval + 1.0))
    wall_seconds = time.perf_counter() - wall_start
    results = sorted(results, key=lambda item: item["material_index"])

    elapsed = [float(item["timing"]["elapsed_seconds"]) for item in results]
    matched_rms = [item["best_rms_dist"] for item in results if item["best_rms_dist"] is not None]
    summary = {
        "number_of_targets": len(results),
        "num_candidates_per_target": args.num_starting_points,
        "wall_seconds": wall_seconds,
        "mean_material_seconds": float(np.mean(elapsed)) if elapsed else None,
        "median_material_seconds": float(np.median(elapsed)) if elapsed else None,
        "min_material_seconds": float(np.min(elapsed)) if elapsed else None,
        "max_material_seconds": float(np.max(elapsed)) if elapsed else None,
        "materials_per_wall_second": len(results) / wall_seconds if wall_seconds > 0 else None,
        "mean_candidates_per_material_second": float(np.mean([
            item["timing"]["candidates_per_second"] for item in results
            if item["timing"]["candidates_per_second"] is not None
        ])) if results else None,
        "match_rate": float(sum(item["target_match"] for item in results) / len(results)) if results else None,
        "rms_dist": float(np.mean(matched_rms)) if matched_rms else None,
        "cpu_count_logical": total_cpus,
        "workers": workers,
        "torch_threads_per_worker": torch_threads,
        "estimated_torch_threads_total": workers * torch_threads,
        "data_setup_seconds": data_setup_seconds,
        "mean_material_seconds_excludes_data_setup": True,
    }
    summary.update(summarize_cpu_usage(Path(args.output_dir)))

    timing_path = Path(args.output_dir, "timing", "per_material_timing.tsv")
    with timing_path.open("w") as f:
        f.write("material_index\tmpid\tformula\tnum_candidates\telapsed_seconds\tlangevin_seconds\tcandidates_per_second\tpid\ttorch_threads\ttarget_match\tbest_rms_dist\n")
        for item in results:
            f.write(
                f"{item['material_index']}\t{item['mpid']}\t{item['formula']}\t{item['num_candidates']}\t"
                f"{item['timing']['elapsed_seconds']:.6f}\t{item['timing']['langevin_seconds']:.6f}\t"
                f"{item['timing']['candidates_per_second']:.6f}\t{item['cpu']['pid']}\t{item['cpu']['torch_threads']}\t"
                f"{item['target_match']}\t{item['best_rms_dist']}\n"
            )

    with Path(args.output_dir, "timing", "summary.json").open("w") as f:
        json.dump(json_safe(summary), f, indent=2)
    with Path(args.output_dir, "cpu_sample", "metrics").mkdir(parents=True, exist_ok=True) or Path(args.output_dir, "cpu_sample", "metrics", "aggregate_metrics.json").open("w") as f:
        aggregate = {
            "sample_label": "cpu_sample",
            "number_of_targets": len(results),
            "num_candidates_per_target": args.num_starting_points,
            "match_rate": summary["match_rate"],
            "rms_dist": summary["rms_dist"],
            "matcher": PXRDGEN_MATCHER_KWARGS,
            "benchmark_summary": summary,
            "materials": [
                {
                    "material_index": item["material_index"],
                    "mpid": item["mpid"],
                    "formula": item["formula"],
                    "target_match": item["target_match"],
                    "best_rms_dist": item["best_rms_dist"],
                    "elapsed_seconds": item["timing"]["elapsed_seconds"],
                }
                for item in results
            ],
        }
        json.dump(json_safe(aggregate), f, indent=2)

    print(json.dumps(json_safe(summary), indent=2))
    print(f"wrote timing: {timing_path}")
    print(f"wrote summary: {Path(args.output_dir, 'timing', 'summary.json')}")


if __name__ == "__main__":
    main()