#!/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()