| """Produce the magnet_benchmark results table by running MagNET and MagNET-x over the four test sets |
| and write it to performance_results.csv. |
| |
| Needs the checkpoints and datasets, so it does not run in CI; they ship in the repo (Git LFS via the |
| Hugging Face checkout) and resolve automatically. One row per (model, test set, nucleus) with median |
| AE, MAE, RMSE, and atom count. The three vibrated/stationary sets are sampled; isolated-chloroform is |
| run in full. The shipped CSV used --sample 500 --seed 0; re-running reproduces it (sampled rows within |
| tolerance, isolated rows exactly). |
| |
| Usage: |
| python magnet_benchmark_run.py --out performance_results.csv --n-passes 1 --sample 500 --seed 0 |
| """ |
| import argparse |
| import csv |
| import os |
| import sys |
| import time |
|
|
| from paths import ensure_on_path, repo_root, dataset_file |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| GH = repo_root(__file__) |
| if HERE not in sys.path: |
| sys.path.insert(0, HERE) |
| ensure_on_path("analysis", "code", "shared", file=__file__) |
| ensure_on_path("magnet", file=__file__) |
| for _d in ("sigma-shake", "gdb_qcd", "dft8k", "sigma-fresh"): |
| ensure_on_path("data", _d, file=__file__) |
|
|
| import magnet_benchmark as M |
| from run_magnet import compute_MagNET_foundation_shieldings, _predict_with |
|
|
|
|
| def magnet_x_chloroform_shieldings(ans, geos, n_passes, device): |
| """Absolute shieldings from the MagNET-x chloroform model (run as a plain predictor, no solvent), |
| used to score MagNET-x on the same test sets as the foundation model.""" |
| return _predict_with("MagNET-x-chloroform_H", "MagNET-x-chloroform_C", ans, geos, |
| n_passes=n_passes, device=device) |
|
|
|
|
| MODEL_RUNNERS = { |
| "MagNET": compute_MagNET_foundation_shieldings, |
| "MagNET-x": magnet_x_chloroform_shieldings, |
| } |
|
|
|
|
| def assemble_all(paths, sample, seed): |
| """Assemble the four test sets, returning {test_set: (ans, geos, dfts)}.""" |
| from decode_sigma_shake import SigmaShake |
| from decode_gdb_qcd import GDBQCD |
| from dft8k_reader import DFT8k |
| from decode_sigma_fresh import SigmaFresh |
|
|
| raw = {} |
| with SigmaShake(paths["sigma_shake"]) as ss: |
| raw["stationary_internal"] = M.cases_stationary_internal(ss, n_sample=sample, seed=seed) |
| with GDBQCD(paths["gdb_qcd"]) as gq: |
| raw["vibrated_internal"] = M.cases_vibrated_internal(gq, n_sample=sample // 3, seed=seed) |
| with DFT8k(paths["dft8k"]) as d8: |
| raw["vibrated_external"] = M.cases_vibrated_external(d8, n_sample=sample, seed=seed) |
| with SigmaFresh(paths["sigma_fresh"]) as sf: |
| |
| |
| raw["isolated_chloroform"] = M.cases_isolated_chloroform(sf, scrub=False) |
| |
| |
| cases = {} |
| for ts, (ans, geos, dfts) in raw.items(): |
| a, g, d, dropped = M.filter_supported(ans, geos, dfts) |
| if dropped: |
| print(f" {ts}: dropped {dropped} structure(s) with unsupported elements", flush=True) |
| cases[ts] = (a, g, d) |
| return cases |
|
|
|
|
| def main(): |
| """Assemble the four test sets, run both models over them, and write performance_results.csv.""" |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--out", default=os.path.join(HERE, "..", "..", "data", "magnet_benchmark", "performance_results.csv")) |
| ap.add_argument("--n-passes", type=int, default=1) |
| ap.add_argument("--sample", type=int, default=500) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--device", default="cpu") |
| ap.add_argument("--sigma-shake", default=dataset_file("sigma-shake", root=GH)) |
| ap.add_argument("--gdb-qcd", default=dataset_file("gdb_qcd", root=GH)) |
| ap.add_argument("--dft8k", default=dataset_file("dft8k", root=GH)) |
| ap.add_argument("--sigma-fresh", default=dataset_file("sigma-fresh", root=GH)) |
| args = ap.parse_args() |
|
|
| paths = {"sigma_shake": args.sigma_shake, "gdb_qcd": args.gdb_qcd, |
| "dft8k": args.dft8k, "sigma_fresh": args.sigma_fresh} |
| print(f"assembling test cases (sample={args.sample}, seed={args.seed})...", flush=True) |
| cases = assemble_all(paths, args.sample, args.seed) |
| for ts, (ans, _, _) in cases.items(): |
| print(f" {ts}: {len(ans)} structures", flush=True) |
|
|
| rows = [] |
| for model, runner in MODEL_RUNNERS.items(): |
| for ts in M.TEST_SETS: |
| ans, geos, dfts = cases[ts] |
| t = time.time() |
| pred = runner(ans, geos, args.n_passes, args.device) |
| tab = M.error_table(pred, dfts, ans) |
| dt = time.time() - t |
| for nucleus in ("1H", "13C"): |
| s = tab[nucleus] |
| rows.append(dict(model=model, test_set=ts, nucleus=nucleus, |
| median_ae=s["median_ae"], mae=s["mae"], rmse=s["rmse"], n=s["n"])) |
| print(f"{model:9s} {ts:20s} ({dt:.0f}s) " |
| f"1H mae {tab['1H']['mae']:.4f} 13C mae {tab['13C']['mae']:.4f}", flush=True) |
|
|
| with open(args.out, "w", newline="") as f: |
| w = csv.DictWriter(f, fieldnames=["model", "test_set", "nucleus", "median_ae", "mae", "rmse", "n"]) |
| w.writeheader() |
| for r in rows: |
| w.writerow(r) |
| print(f"wrote {args.out} ({len(rows)} rows)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
