"""Compute retrieval metrics from a result pickle.""" import argparse import os import pickle import sys from pathlib import Path def _add_local_dependency_paths() -> None: """Ensure ``flare`` and ``massspecgym`` are importable when run as a script.""" flare_repo_root = Path(__file__).resolve().parents[1] candidate_roots = [flare_repo_root, flare_repo_root / "massspecgym"] env_value = os.environ.get("MASSSPECGYM_ROOT") if env_value: candidate_roots.append(Path(env_value).expanduser()) for root in candidate_roots: if root.exists(): root_str = str(root) if root_str not in sys.path: sys.path.insert(0, root_str) _add_local_dependency_paths() import numpy as np from tqdm import tqdm def rank_at_k(ranks, k: int) -> float: return (ranks <= k).mean() def confidence_interval( values, confidence_level: float = 0.999, n_resamples: int = 20_000, seed: int = 0, ) -> str: """Return a bootstrap confidence interval formatted as in msgym_result.ipynb.""" from scipy.stats import bootstrap ci = bootstrap( (np.asarray(values),), np.mean, confidence_level=confidence_level, n_resamples=n_resamples, random_state=seed, ).confidence_interval return f"{ci.low:.2f}-{ci.high:.2f}" def compute_mces(result, n_processes: int): """Compute the per-row MCES between the target and top-ranked candidate.""" # Lazy import: needs massspecgym / torchmetrics, only required for --mces. print("Importing MCES helpers (flare.utils.eval)...") from flare.utils import eval as eval_utils required = {"candidates", "labels", "scores"} missing = required.difference(result.columns) if missing: raise SystemExit( f"MCES requires columns {sorted(required)}; missing: {sorted(missing)}" ) print("Extracting targets and top-1 candidates...") targets = [ eval_utils.get_target(candidates, labels) for candidates, labels in tqdm( zip(result["candidates"], result["labels"]), total=len(result), desc="targets", ) ] top_candidates = [ eval_utils.get_top_cand(candidates, scores) for candidates, scores in tqdm( zip(result["candidates"], result["scores"]), total=len(result), desc="cand@1", ) ] pairs = list(zip(targets, top_candidates)) unique_pairs = list(dict.fromkeys(pairs)) print( f"Computing MCES for {len(unique_pairs)} unique target/candidate pairs " f"({n_processes} processes)..." ) computed = eval_utils.Compute_Myopic_MCES.compute_mces_parallel( unique_pairs, n_processes=n_processes ) print("Assigning MCES values to rows...") mces_by_pair = {tuple(pair): mces for pair, mces in computed} return np.asarray([mces_by_pair[pair] for pair in pairs], dtype=float) def main(): parser = argparse.ArgumentParser( description="Compute rank@k and optional CI/MCES metrics from a result pickle." ) parser.add_argument("result_file", type=Path, help="Path to result_*.pkl") parser.add_argument( "--rank_col", default="rank", help="Column with ranks (default: rank)", ) parser.add_argument( "--ks", nargs="+", type=int, default=[1, 5, 20], help="k values for rank@k (default: 1 5 20)", ) parser.add_argument( "--ci", "--compute_ci", dest="compute_ci", action="store_true", help="Compute bootstrap confidence intervals (hit rates are percentages)", ) parser.add_argument( "--mces", "--compute_mces", dest="compute_mces", action="store_true", help="Compute the mean MCES for each target/top-candidate pair", ) parser.add_argument( "--confidence_level", type=float, default=0.999, help="Bootstrap confidence level (default: 0.999)", ) parser.add_argument( "--n_resamples", type=int, default=20_000, help="Number of bootstrap resamples (default: 20000)", ) parser.add_argument( "--seed", type=int, default=0, help="Bootstrap random seed (default: 0)", ) parser.add_argument( "--mces_processes", type=int, default=25, help="Number of MCES worker processes (default: 25)", ) args = parser.parse_args() result_path = args.result_file.resolve() print(f"Loading result file: {result_path}") with open(result_path, "rb") as f: result = pickle.load(f) print(f"Loaded {len(result)} rows; columns={list(result.columns)}") if args.rank_col not in result.columns: raise SystemExit( f"Column '{args.rank_col}' not found. Available: {list(result.columns)}" ) ranks = result[args.rank_col] if len(ranks) == 0: raise SystemExit("The result file contains no rows.") print(f"Computing rank@k for ks={args.ks} (rank_col={args.rank_col})...") lines = [f"n={len(ranks)}", f"rank_col={args.rank_col}"] for k in args.ks: val = rank_at_k(ranks, k) line = f"rank@{k}: {val:.6f}" lines.append(line) print(line) mces_values = None if args.compute_mces: print("MCES flag set; starting MCES computation...") mces_values = compute_mces(result, args.mces_processes) line = f"mces: {mces_values.mean():.6f}" lines.append(line) print(line) elif "mces" in result.columns: print("Using precomputed 'mces' column from result file.") mces_values = result["mces"].to_numpy(dtype=float) if args.compute_ci: ci_label = f"{args.confidence_level * 100:g}%_ci" print( f"Computing bootstrap CIs " f"(level={args.confidence_level}, n_resamples={args.n_resamples}, seed={args.seed})..." ) ci_jobs = [] for k in args.ks: # The notebook reports hit-rate confidence intervals on a 0-100 scale. hit_rate_percent = (ranks <= k).to_numpy(dtype=float) * 100 ci_jobs.append((f"rank@{k}_{ci_label}_percent", hit_rate_percent)) if mces_values is not None: ci_jobs.append((f"mces_{ci_label}", mces_values)) else: print("No MCES values available; skipping MCES CI.") for name, values in tqdm(ci_jobs, desc="bootstrap CI"): print(f" Bootstrapping {name}...") interval = confidence_interval( values, confidence_level=args.confidence_level, n_resamples=args.n_resamples, seed=args.seed, ) line = f"{name}: {interval}" lines.append(line) print(line) out_path = result_path.with_name(result_path.stem + "_rank_at_k.txt") print(f"Writing metrics to {out_path}...") out_path.write_text("\n".join(lines) + "\n") print(f"Done. Wrote {out_path}") if __name__ == "__main__": main()