| |
| """Build the released MolGenBench paper baselines from three evaluation rounds.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import ast |
| from collections import defaultdict |
| from pathlib import Path |
| import sys |
| from typing import Any |
|
|
| import numpy as np |
| import pandas as pd |
| import yaml |
|
|
| SPACE_ROOT = Path(__file__).resolve().parents[1] |
| if str(SPACE_ROOT) not in sys.path: |
| sys.path.insert(0, str(SPACE_ROOT)) |
|
|
| from src.leaderboard_schema import ( |
| HIT_DENOVO_COLUMNS, |
| HIT_H2L_COLUMNS, |
| LEADERBOARD_COLUMNS, |
| METRIC_COLUMNS, |
| ) |
|
|
| ROUNDS = ("Round1", "Round2", "Round3") |
| SPLITS = ("all", "seen", "unseen") |
| MODEL_SPECS = { |
| "denovo": [ |
| ("MolCraft_generated_molecules", "MolCraft"), |
| ("DecompDiff_generated_molecules", "DecompDiff"), |
| ("FLAG_generated_molecules", "FLAG"), |
| ("PocketFlow_generated_molecules", "PocketFlow"), |
| ("SurfGen_generated_molecules", "SurfGen"), |
| ("TargetDiff_generated_molecules", "TargetDiff"), |
| ("TamGen_generated_molecules", "TamGen"), |
| ("Pocket2Mol_generated_molecules", "Pocket2Mol"), |
| ("diffSBDD_cond_crossdocked_generated_molecules", "DiffSBDD-C"), |
| ("diffSBDD_cond_moad_generated_molecules", "DiffSBDD-M"), |
| ], |
| "h2l": [ |
| ("DeleteHit2Lead(CrossDock)_Hit_to_Lead", "Delete"), |
| ("DiffDec_Hit_to_Lead", "DiffDec"), |
| ("diffSBDD_cond_crossdocked_Hit_to_Lead", "DiffSBDD-C"), |
| ("diffSBDD_cond_moad_Hit_to_Lead", "DiffSBDD-M"), |
| ("PGMG_Hit_to_Lead", "PGMG"), |
| ("ShapeMol_Hit_to_Lead", "ShapeMol"), |
| ("shepherd_x1x3x4_mosesaq_submission_Hit_to_Lead", "ShEPhERD"), |
| ], |
| } |
| MODE_DIRS = {"denovo": "De_novo_Results", "h2l": "Hit_to_Lead_Results"} |
|
|
| SCALAR_METRICS = { |
| "Validity": ("Validity", "Validity"), |
| "QED": ("QED", "QED"), |
| "SA": ("SA", "SA"), |
| "Uniqueness": ("Uniqueness", "Uniqueness"), |
| "Diversity": ("Diversity", "Diversity"), |
| "MotifDist_AtomType": ("MotifDist", "Atom Type_Score"), |
| "MotifDist_RingType": ("MotifDist", "Ring Type_Score"), |
| "MotifDist_FunctionalGroup": ("MotifDist", "Functional Group_Score"), |
| "ChemFilter_smiles_pass_rate": ("ChemFilter", "smiles_pass_rate"), |
| "ChemFilter_scaffold_pass_rate": ("ChemFilter", "scaffold_pass_rate"), |
| "PoseBuster_pass_all": ("PoseBuster", "pass_all"), |
| "StrainEnergy_Gen_25": ("StrainEnergy", "StrainEnergy25%"), |
| "StrainEnergy_Gen_50": ("StrainEnergy", "StrainEnergy50%"), |
| "StrainEnergy_Gen_75": ("StrainEnergy", "StrainEnergy75%"), |
| "StrainEnergy_Vina_25": ("StrainEnergy_docked", "StrainEnergy25%"), |
| "StrainEnergy_Vina_50": ("StrainEnergy_docked", "StrainEnergy50%"), |
| "StrainEnergy_Vina_75": ("StrainEnergy_docked", "StrainEnergy75%"), |
| "RMSD_lt_2": ("RMSD", "RMSD_<2"), |
| "InteractionScore_Vina_25": ("InteractionScore_docked", "InteractionScore>p25"), |
| "InteractionScore_Vina_50": ("InteractionScore_docked", "InteractionScore>p50"), |
| "InteractionScore_Vina_75": ("InteractionScore_docked", "InteractionScore>p75"), |
| "InteractionScore_Gen_25": ("InteractionScore", "InteractionScore>p25"), |
| "InteractionScore_Gen_50": ("InteractionScore", "InteractionScore>p50"), |
| "InteractionScore_Gen_75": ("InteractionScore", "InteractionScore>p75"), |
| "ClashScore_Gen": ("ClashScore", "ClashScore"), |
| "ClashScore_Vina": ("ClashScore_docked", "ClashScore"), |
| "HitRediscover_smiles_hit_recovery": ("HitRediscover", "smiles_hit_recovery"), |
| "HitRediscover_smiles_hit_rate": ("HitRediscover", "smiles_hit_rate"), |
| "HitRediscover_smiles_hit_fraction": ("HitRediscover", "smiles_hit_fraction"), |
| "HitRediscover_scaffold_hit_recovery": ("HitRediscover", "scaffold_hit_recovery"), |
| "HitRediscover_scaffold_hit_rate": ("HitRediscover", "scaffold_hit_rate"), |
| "HitRediscover_scaffold_hit_fraction": ("HitRediscover", "scaffold_hit_fraction"), |
| "HitRediscover_Mean_normalized_affinity": ("HitRediscover", "Mean_normalized_affinity"), |
| "HitRediscover_max_normalized_affinity_ge_0_5": ( |
| "HitRediscover", "sum_max_normalized_affinity_count_>=_0.5" |
| ), |
| "HitRediscover_max_normalized_affinity_ge_0_8": ( |
| "HitRediscover", "sum_max_normalized_affinity_count_>=_0.8" |
| ), |
| "HitRediscover_max_normalized_affinity_ge_0_9": ( |
| "HitRediscover", "sum_max_normalized_affinity_count_>=_0.9" |
| ), |
| "HitRediscover_max_normalized_affinity_ge_0_95": ( |
| "HitRediscover", "sum_max_normalized_affinity_count_>=_0.95" |
| ), |
| "HitRediscover_max_normalized_affinity_ge_0_99": ( |
| "HitRediscover", "sum_max_normalized_affinity_count_>=_0.99" |
| ), |
| } |
|
|
| TASCORE_COLUMNS = { |
| "smiles_TAScore": [ |
| "HitRediscover_smiles_TAScore_0_1_count", |
| "HitRediscover_smiles_TAScore_1_10_count", |
| "HitRediscover_smiles_TAScore_10_100_count", |
| "HitRediscover_smiles_TAScore_gt_100_count", |
| ], |
| "scaffold_TAScore": [ |
| "HitRediscover_scaffold_TAScore_0_1_count", |
| "HitRediscover_scaffold_TAScore_1_10_count", |
| "HitRediscover_scaffold_TAScore_10_100_count", |
| "HitRediscover_scaffold_TAScore_gt_100_count", |
| ], |
| } |
|
|
|
|
| def _parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("paper_results", type=Path, help="Path to MolGenBench paper_results") |
| parser.add_argument("output", type=Path, help="Output paper_baselines.parquet") |
| parser.add_argument( |
| "--seen-targets-file", |
| type=Path, |
| help="sup_info/utils.py containing the authoritative uniprot_in_trainset list", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def _read_seen_targets(path: Path) -> set[str]: |
| tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) |
| for node in tree.body: |
| if isinstance(node, (ast.Assign, ast.AnnAssign)): |
| targets = node.targets if isinstance(node, ast.Assign) else [node.target] |
| if any( |
| isinstance(target, ast.Name) and target.id == "uniprot_in_trainset" |
| for target in targets |
| ): |
| return set(ast.literal_eval(node.value)) |
| raise ValueError(f"uniprot_in_trainset was not found in {path}") |
|
|
|
|
| def _read_round_documents(paper_results: Path) -> tuple[dict[tuple[str, str], list[dict[str, Any]]], int]: |
| documents: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) |
| expected_paths = { |
| paper_results / round_name / MODE_DIRS[task] / source_name / "results.yaml" |
| for task, models in MODEL_SPECS.items() |
| for source_name, _ in models |
| for round_name in ROUNDS |
| } |
| discovered_paths = set(paper_results.glob("Round*/**/results.yaml")) |
| if discovered_paths != expected_paths: |
| missing = sorted(str(path) for path in expected_paths - discovered_paths) |
| extra = sorted(str(path) for path in discovered_paths - expected_paths) |
| raise ValueError( |
| f"Expected exactly the 51 declared YAML files; missing={missing}, extra={extra}" |
| ) |
| yaml_count = 0 |
| for task, models in MODEL_SPECS.items(): |
| mode_dir = MODE_DIRS[task] |
| for source_name, _ in models: |
| for round_name in ROUNDS: |
| path = paper_results / round_name / mode_dir / source_name / "results.yaml" |
| if not path.is_file(): |
| raise FileNotFoundError(f"Missing paper result: {path}") |
| document = yaml.safe_load(path.read_text(encoding="utf-8")) |
| if not isinstance(document, dict): |
| raise ValueError(f"Expected a mapping in {path}") |
| if document.get("round") != round_name: |
| raise ValueError(f"Unexpected round in {path}: {document.get('round')!r}") |
| documents[(task, source_name)].append(document) |
| yaml_count += 1 |
| if yaml_count != 51: |
| raise ValueError(f"Expected exactly 51 YAML files, read {yaml_count}") |
| return documents, yaml_count |
|
|
|
|
| def _round_metric( |
| document: dict[str, Any], split: str, family: str, key: str |
| ) -> float: |
| """Read one metric from one paper round, preserving complete-round sparsity.""" |
|
|
| value = document.get("metrics", {}).get(family, {}).get(split, {}).get(key) |
| return np.nan if value is None else float(value) |
|
|
|
|
| def _load_tascore_details(paper_results: Path) -> pd.DataFrame: |
| frames: list[pd.DataFrame] = [] |
| for source_name, model_name in MODEL_SPECS["denovo"]: |
| for round_name in ROUNDS: |
| path = ( |
| paper_results |
| / round_name |
| / MODE_DIRS["denovo"] |
| / source_name |
| / "hitrediscovery_detail.csv" |
| ) |
| if not path.is_file(): |
| raise FileNotFoundError(f"Missing TAScore detail: {path}") |
| frame = pd.read_csv(path) |
| required = {"uniprot", *TASCORE_COLUMNS} |
| missing = required.difference(frame.columns) |
| if missing: |
| raise ValueError(f"Missing TAScore columns in {path}: {sorted(missing)}") |
| frame = frame[["uniprot", *TASCORE_COLUMNS]].copy() |
| frame["model_name"] = model_name |
| frame["round"] = round_name |
| frames.append(frame) |
| return pd.concat(frames, ignore_index=True) |
|
|
|
|
| def _tascore_counts( |
| details: pd.DataFrame, seen_targets: set[str] |
| ) -> dict[tuple[str, str, str], dict[str, float]]: |
| """Apply the paper-figure TAScore binning to each displayed round. |
| |
| The per-target detail files only contain observed targets. The paper figure |
| uses the complete round/split target universe as its denominator and assigns |
| targets missing for an individual model to the 0-1 bin. |
| """ |
|
|
| output: dict[tuple[str, str, str], dict[str, float]] = {} |
| for split in SPLITS: |
| selected = details |
| if split == "seen": |
| selected = details[details["uniprot"].isin(seen_targets)] |
| elif split == "unseen": |
| selected = details[~details["uniprot"].isin(seen_targets)] |
|
|
| for round_name in ROUNDS: |
| round_scores = selected[selected["round"] == round_name] |
| target_count = int(round_scores["uniprot"].nunique()) |
| for _, model_name in MODEL_SPECS["denovo"]: |
| model_scores = round_scores[round_scores["model_name"] == model_name] |
| row: dict[str, float] = {} |
| for source_column, output_columns in TASCORE_COLUMNS.items(): |
| scores = pd.to_numeric(model_scores[source_column], errors="coerce") |
| counts = [ |
| int(((scores >= 0) & (scores < 1)).sum()), |
| int(((scores >= 1) & (scores < 10)).sum()), |
| int(((scores >= 10) & (scores < 100)).sum()), |
| int((scores >= 100).sum()), |
| ] |
| counts[0] += max(target_count - sum(counts), 0) |
| row.update( |
| {column: float(count) for column, count in zip(output_columns, counts)} |
| ) |
| output[(model_name, round_name, split)] = row |
| return output |
|
|
|
|
| def build_paper_baselines( |
| paper_results: Path, seen_targets_file: Path |
| ) -> pd.DataFrame: |
| """Return one row per task, model, round, and All/Seen/Unseen split.""" |
|
|
| documents, _ = _read_round_documents(paper_results) |
| seen_targets = _read_seen_targets(seen_targets_file) |
| tascores = _tascore_counts(_load_tascore_details(paper_results), seen_targets) |
| rows: list[dict[str, Any]] = [] |
|
|
| for task, models in MODEL_SPECS.items(): |
| for source_name, model_name in models: |
| model_documents = documents[(task, source_name)] |
| for document in model_documents: |
| round_name = str(document["round"]) |
| for split in SPLITS: |
| row: dict[str, Any] = {column: np.nan for column in LEADERBOARD_COLUMNS} |
| slug = model_name.lower().replace(" ", "-") |
| row.update( |
| { |
| "source_type": "paper_baseline", |
| "target_split": split, |
| "submission_id": f"paper-baseline-{task}-{slug}-{round_name.lower()}", |
| "model_name": model_name, |
| "round": round_name, |
| "task": task, |
| "evaluation_level": "2d+3d+hitrediscovery", |
| "submitted_by": "MolGenBench authors", |
| "publish_time": None, |
| } |
| ) |
| for column, (family, key) in SCALAR_METRICS.items(): |
| row[column] = _round_metric(document, split, family, key) |
| if task == "denovo": |
| row.update(tascores[(model_name, round_name, split)]) |
| rows.append(row) |
|
|
| frame = pd.DataFrame(rows).reindex(columns=LEADERBOARD_COLUMNS) |
| if len(frame) != 153 or frame["model_name"].nunique() != 15: |
| |
| raise ValueError( |
| f"Expected 153 rows and 15 unique display labels, got {len(frame)} and " |
| f"{frame['model_name'].nunique()}" |
| ) |
| if frame.groupby(["task", "model_name", "round"]).size().ne(3).any(): |
| raise ValueError("Each paper task/model/round must contain all three target splits") |
| if frame.groupby(["task", "model_name"])["round"].nunique().ne(3).any(): |
| raise ValueError("Each paper task/model pair must contain Round1, Round2, and Round3") |
| missing_required = frame[[*METRIC_COLUMNS[:10], *METRIC_COLUMNS[26:32]]].isna().any(axis=1) |
| if missing_required.any(): |
| raise ValueError("2D and common HitRediscovery metrics must be complete for every paper row") |
| if frame.loc[frame["task"] == "denovo", HIT_DENOVO_COLUMNS].isna().any(axis=None): |
| raise ValueError("De novo TAScore bins must be complete for every paper row") |
| if frame.loc[frame["task"] == "h2l", HIT_H2L_COLUMNS].isna().any(axis=None): |
| raise ValueError("Hit-to-lead affinity metrics must be complete for every paper row") |
| return frame |
|
|
|
|
| def main() -> None: |
| args = _parse_args() |
| seen_file = args.seen_targets_file or args.paper_results.parent / "sup_info" / "utils.py" |
| frame = build_paper_baselines(args.paper_results, seen_file) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| frame.to_parquet(args.output, index=False) |
| print( |
| f"Wrote {len(frame)} round-level rows for " |
| f"{frame.groupby(['task', 'model_name']).ngroups} task/model pairs to {args.output}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
