From 1c20c28b961078676fd0b0a7ec5760e8b3cdb90b Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:39:43 +0100 Subject: [PATCH 01/26] chore: port analysis scripts --- scripts/analyze_fdr_overlap.py | 1261 ++++++++++++ scripts/analyze_features.py | 802 ++++++++ scripts/analyze_novelty.py | 1814 +++++++++++++++++ scripts/analyze_upscored_fps.py | 705 +++++++ scripts/benchmark_runtime.py | 694 +++++++ scripts/benchmark_scaling.py | 421 ++++ scripts/calibrator_generalisation_utils.py | 114 ++ scripts/evaluate_calibrator_generalisation.py | 406 ++++ scripts/fdr_tool_comparison_preprocess.py | 774 +++++++ scripts/fdr_tool_comparison_summaries.py | 533 +++++ scripts/feature_subsets.py | 74 + scripts/plot_ablation_summary.py | 641 ++++++ scripts/plot_acfm_minus_lcfm_fdr.py | 441 ++++ scripts/plot_analysis.py | 1133 ++++++++++ .../plot_calibrator_generalisation_heatmap.py | 265 +++ scripts/plot_eval_results.py | 984 +++++++++ scripts/plot_fdr_method_comparison.py | 1097 ++++++++++ scripts/plot_feature_investigation.py | 1323 ++++++++++++ .../run_external_peptide_holdout_benchmark.py | 1115 ++++++++++ scripts/run_feature_ablations.py | 1532 ++++++++++++++ 20 files changed, 16129 insertions(+) create mode 100644 scripts/analyze_fdr_overlap.py create mode 100644 scripts/analyze_features.py create mode 100644 scripts/analyze_novelty.py create mode 100644 scripts/analyze_upscored_fps.py create mode 100644 scripts/benchmark_runtime.py create mode 100644 scripts/benchmark_scaling.py create mode 100644 scripts/calibrator_generalisation_utils.py create mode 100644 scripts/evaluate_calibrator_generalisation.py create mode 100644 scripts/fdr_tool_comparison_preprocess.py create mode 100644 scripts/fdr_tool_comparison_summaries.py create mode 100644 scripts/feature_subsets.py create mode 100644 scripts/plot_ablation_summary.py create mode 100644 scripts/plot_acfm_minus_lcfm_fdr.py create mode 100644 scripts/plot_analysis.py create mode 100644 scripts/plot_calibrator_generalisation_heatmap.py create mode 100644 scripts/plot_eval_results.py create mode 100644 scripts/plot_fdr_method_comparison.py create mode 100644 scripts/plot_feature_investigation.py create mode 100644 scripts/run_external_peptide_holdout_benchmark.py create mode 100644 scripts/run_feature_ablations.py diff --git a/scripts/analyze_fdr_overlap.py b/scripts/analyze_fdr_overlap.py new file mode 100644 index 00000000..cbe44ef2 --- /dev/null +++ b/scripts/analyze_fdr_overlap.py @@ -0,0 +1,1261 @@ +#!/usr/bin/env python3 +"""Post-FDR overlap analysis: Winnow-filtered identifications vs database search. + +For each project (see ``plot_eval_results.py`` CLI pattern), at 1 %, 5 %, and 10 % +nominal FDR: + + * Count retained PSMs / unique peptides vs database-search reference peptides at + the same nominal FDR (Winnow: non-parametric on calibrated confidence; + database search: database-grounded on raw confidence). + * Match rule: exact ProForma sequence after I/L equivalence; PTM differences + are not a match. + * Categorise discordant calls (partial match, PTM candidate, single-AA variant, + near-miss edit distance 2-3, fully discordant). + * Full-search Venns: Winnow (non-parametric calibrated confidence) vs database + search unique peptides at the same nominal FDR (database-grounded raw confidence). + * Violin plots comparing database-matched vs fully novel retained PSMs at + matched FDR (same retention rules as overlap summaries). + +Inputs are ``winnow predict`` output folders arranged as subdirectories under two +roots: an **unlabelled** tree (full-search Winnow predictions) and a **labelled** +tree (database-search reference with ``sequence``). Each project folder (flat or +``PXD*//`` nested) must contain ``preds_and_fdr_metrics.csv``; +``metadata.csv`` is merged when present for violin plots. +""" + +from __future__ import annotations + +import json +import logging +import re +from collections import defaultdict +from pathlib import Path +from typing import Annotated, cast + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer +import yaml +from matplotlib_venn import venn2 +from rich.logging import RichHandler + +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl +from winnow.fdr.nonparametric import NonParametricFDRControl + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_NOVEL_COLOUR = _PALETTE[2] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)-?") +_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") + +_PXD_ACCESSION_PREFIX = "PXD" + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] +_DB_GROUNDED_DROP = 10 +RAW_CONFIDENCE_COL = "confidence" +DB_Q_VALUE_COL = "db_psm_q_value" + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", + "PXD004732": "ProteomeTools-1", +} + +_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") +_UNLABELLED_FOLDER_SUFFIXES = ("_raw", "_unlabelled") +_LABELLED_FOLDER_SUFFIXES = ("_annotated", "_labelled") + +NOVEL_FEATURE_COLUMNS: list[tuple[str, str]] = [ + ("spectral_angle", "Spectral angle"), + ("ion_matches", "Ion match rate"), + ("ion_match_intensity", "Ion match intensity"), + ("precursor_charge", "Precursor charge"), + ("mass_error_da", "Precursor mass error (Da)"), + ("irt_error", "iRT error"), + ("confidence", "Raw confidence"), + ("margin", "Beam margin"), +] + +_DISCORDANCE_COUNT_COLS = [ + "n_partial_match", + "n_ptm_candidate", + "n_single_aa_variant", + "n_near_miss_edit_dist", + "n_fully_discordant", +] + +_PTM_DELTAS = { + "oxidation": 15.995, + "phosphorylation": 79.966, + "deamidation": 0.984, + "acetylation": 42.011, + "methylation": 14.016, + "carbamidomethyl": 57.021, +} +_PTM_TOLERANCE_DA = 0.02 + +_MIN_VIOLIN_GROUP_SIZE = 5 +_MAX_VIOLIN_PSMs_PER_GROUP = 5000 +_VIOLIN_SUBSAMPLE_SEED = 42 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _display_name(key: str) -> str: + return DATASET_DISPLAY_NAMES.get(key, key) + + +def _project_key_from_folder(folder_name: str) -> str: + """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" + for suffix in _FOLDER_SUFFIXES: + if folder_name.endswith(suffix): + return folder_name[: -len(suffix)] + return folder_name + + +def _search_space_tag_from_folder(folder_name: str) -> str: + """Infer eval-type label for tables/plots from the unlabelled subfolder name.""" + for suffix in _UNLABELLED_FOLDER_SUFFIXES: + if folder_name.endswith(suffix): + return suffix[1:] # raw | unlabelled + return "full_search" + + +def _eval_type_display(eval_type: str) -> str: + """Human-readable eval-type label for plot titles.""" + return { + "full_search": "full search space", + "raw": "raw", + "unlabelled": "unlabelled", + }.get(eval_type, eval_type) + + +def _get_residue_masses() -> dict[str, float]: + config_path = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" + with open(config_path) as f: + cfg = yaml.safe_load(f) + return cfg["residue_masses"] + + +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _sequence_match_key(seq: str) -> str: + """Exact match key: ProForma with mods preserved, I/L equivalent.""" + if not seq or not isinstance(seq, str): + return "" + return seq.replace("I", "L") + + +def _strip_mods(seq: str) -> str: + """Strip PTM annotations and normalise I -> L (discordance subtyping only).""" + if not seq or not isinstance(seq, str): + return "" + s = _MOD_PLUS.sub("", seq) + s = _MOD_UNIMOD.sub("", s) + return s.replace("I", "L") + + +def _levenshtein(s: str, t: str) -> int: + n, m = len(s), len(t) + if n == 0: + return m + if m == 0: + return n + prev = list(range(m + 1)) + for i in range(1, n + 1): + curr = [i] + [0] * m + for j in range(1, m + 1): + cost = 0 if s[i - 1] == t[j - 1] else 1 + curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost) + prev = curr + return prev[m] + + +def _mass_from_sequence(seq_str: str, residue_masses: dict[str, float]) -> float | None: + key = _strip_mods(seq_str) + if not key: + return None + total = 18.010565 + for aa in key: + m = residue_masses.get(aa) + if m is None: + return None + total += m + return total + + +def _db_stripped_by_length(db_stripped_list: list[str]) -> dict[int, list[str]]: + by_len: dict[int, list[str]] = defaultdict(list) + for s in db_stripped_list: + by_len[len(s)].append(s) + return by_len + + +def _min_edit_distance_to_db( + pred_stripped: str, db_stripped_by_len: dict[int, list[str]] +) -> int: + if not pred_stripped: + return 999 + plen = len(pred_stripped) + best = 999 + for length in range(max(0, plen - 3), plen + 4): + for db in db_stripped_by_len.get(length, []): + d = _levenshtein(pred_stripped, db) + if d < best: + best = d + if best == 0: + return 0 + return best + + +def _build_db_reference_sets( + db_df: pd.DataFrame, +) -> tuple[set[str], set[str], list[str], dict[int, list[str]]]: + sequences = db_df["sequence"].dropna().astype(str) + db_keys = {_sequence_match_key(s) for s in sequences if _sequence_match_key(s)} + db_stripped = [_strip_mods(s) for s in sequences if _strip_mods(s)] + db_stripped_set = set(db_stripped) + db_stripped_unique = list(dict.fromkeys(db_stripped)) + return ( + db_keys, + db_stripped_set, + db_stripped_unique, + _db_stripped_by_length(db_stripped_unique), + ) + + +def _is_db_match(pred: str, db_keys: set[str]) -> bool: + return _sequence_match_key(pred) in db_keys + + +# --------------------------------------------------------------------------- +# Discordance classification +# --------------------------------------------------------------------------- +LabelledDiscordanceKey = tuple[str, str, int] +LabelledDiscordanceCache = dict[LabelledDiscordanceKey, str] +DbDiscordanceCache = dict[str, str] + + +def _discordance_from_edit_distance(seq_norm: str, pred_norm: str) -> str | None: + if not (seq_norm and pred_norm): + return None + ed = _levenshtein(seq_norm, pred_norm) + if ed == 1: + return "single_aa_variant" + if ed in (2, 3): + return "near_miss_edit_dist" + return None + + +def _discordance_from_ptm_mass( + seq_str: str, + pred_str: str, + residue_masses: dict[str, float], +) -> str | None: + seq_mass = _mass_from_sequence(seq_str, residue_masses) + pred_mass = _mass_from_sequence(pred_str, residue_masses) + if seq_mass is None or pred_mass is None: + return None + delta = abs(seq_mass - pred_mass) + for ptm_delta in _PTM_DELTAS.values(): + if abs(delta - ptm_delta) < _PTM_TOLERANCE_DA: + return "ptm_candidate" + return None + + +def _labelled_discordance_key(row: pd.Series) -> LabelledDiscordanceKey: + return ( + str(row.get("sequence", "")), + str(row.get("prediction", "")), + int(row.get("num_matches", 0)), + ) + + +def _lookup_labelled_discordance( + row: pd.Series, cache: LabelledDiscordanceCache +) -> str: + return cache[_labelled_discordance_key(row)] + + +def classify_discordance_labelled( + seq_str: str, + pred_str: str, + num_matches: int, + residue_masses: dict[str, float], +) -> str: + """Classify a non-matching PSM on a labelled spectrum.""" + if num_matches > 0: + return "partial_match" + + seq_norm = _strip_mods(seq_str) + pred_norm = _strip_mods(pred_str) + if seq_norm == pred_norm and _sequence_match_key(seq_str) != _sequence_match_key( + pred_str + ): + return "ptm_candidate" + + edit_label = _discordance_from_edit_distance(seq_norm, pred_norm) + if edit_label is not None: + return edit_label + + ptm_label = _discordance_from_ptm_mass(seq_str, pred_str, residue_masses) + if ptm_label is not None: + return ptm_label + + return "fully_discordant" + + +def classify_discordance_vs_db( + pred_str: str, + db_keys: set[str], + db_stripped_set: set[str], + db_stripped_by_len: dict[int, list[str]], +) -> str: + """Classify a discordant full-search PSM vs the database peptide reference set.""" + pred_key = _sequence_match_key(pred_str) + if pred_key in db_keys: + return "db_match" + + pred_stripped = _strip_mods(pred_str) + if pred_stripped in db_stripped_set: + return "ptm_candidate" + + ed = _min_edit_distance_to_db(pred_stripped, db_stripped_by_len) + if ed == 1: + return "single_aa_variant" + if ed in (2, 3): + return "near_miss_edit_dist" + return "fully_discordant" + + +def _build_discordance_cache( + predictions: pd.Series, + db_keys: set[str], + db_stripped_set: set[str], + db_stripped_by_len: dict[int, list[str]], +) -> dict[str, str]: + """Classify each unique discordant prediction string once.""" + cache: dict[str, str] = {} + for pred in predictions.dropna().unique(): + key = str(pred) + if _is_db_match(key, db_keys): + cache[key] = "db_match" + elif key not in cache: + cache[key] = classify_discordance_vs_db( + key, db_keys, db_stripped_set, db_stripped_by_len + ) + return cache + + +def _classify_predictions_vs_db( + predictions: pd.Series, + cache: dict[str, str], +) -> pd.Series: + return predictions.map(lambda p: cache.get(str(p), "fully_discordant")) + + +# --------------------------------------------------------------------------- +# Data loading +# --------------------------------------------------------------------------- +def _preds_header(folder: Path) -> set[str] | None: + preds_path = folder / "preds_and_fdr_metrics.csv" + if not preds_path.is_file(): + return None + return set(pd.read_csv(preds_path, nrows=0).columns.tolist()) + + +def _is_unlabelled_preds_folder(folder: Path) -> bool: + header = _preds_header(folder) + return header is not None and "sequence" not in header + + +def _is_labelled_preds_folder(folder: Path) -> bool: + header = _preds_header(folder) + return header is not None and "sequence" in header + + +def _load_from_folder(folder: Path) -> pd.DataFrame: + """Load preds_and_fdr_metrics.csv merged with metadata.csv from a project folder.""" + preds_path = folder / "preds_and_fdr_metrics.csv" + if not preds_path.is_file(): + raise FileNotFoundError(f"Missing predictions file: {preds_path}") + + preds_df = pd.read_csv(preds_path) + meta_path = folder / "metadata.csv" + if meta_path.is_file(): + meta_df = pd.read_csv(meta_path) + overlap_cols = [ + c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" + ] + if overlap_cols: + meta_df = meta_df.drop(columns=overlap_cols) + return preds_df.merge(meta_df, on="spectrum_id", how="left") + return preds_df + + +def _discover_unlabelled_folders(root: Path) -> dict[str, Path]: + """Map project key -> full-search folder under ``root``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts (``{root}/PXD004452//``) used by new eval sets. + """ + projects: dict[str, Path] = {} + if not root.is_dir(): + return projects + + def _register(key: str, folder: Path) -> None: + if key in projects: + logger.warning( + "Duplicate unlabelled project key %r: %s and %s", + key, + projects[key], + folder, + ) + return + projects[key] = folder + + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if _is_unlabelled_preds_folder(child): + _register(_project_key_from_folder(child.name), child) + continue + if not child.name.startswith(_PXD_ACCESSION_PREFIX): + continue + for run_dir in sorted(child.iterdir()): + if run_dir.is_dir() and _is_unlabelled_preds_folder(run_dir): + _register(run_dir.name, run_dir) + return projects + + +def _discover_labelled_folders(root: Path) -> dict[str, Path]: + """Map project key -> database-reference folder under ``root``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts (``{root}/PXD004452//``) used by new eval sets. + """ + projects: dict[str, Path] = {} + if not root.is_dir(): + return projects + + def _register(key: str, folder: Path) -> None: + if key in projects: + logger.warning( + "Duplicate labelled project key %r: %s and %s", + key, + projects[key], + folder, + ) + return + projects[key] = folder + + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if _is_labelled_preds_folder(child): + _register(_project_key_from_folder(child.name), child) + continue + if not child.name.startswith(_PXD_ACCESSION_PREFIX): + continue + for run_dir in sorted(child.iterdir()): + if run_dir.is_dir() and _is_labelled_preds_folder(run_dir): + _register(run_dir.name, run_dir) + return projects + + +def discover_project_pairs( + unlabelled_dir: Path, + labelled_dir: Path, + *, + projects_filter: set[str] | None = None, +) -> list[tuple[str, Path, Path, str]]: + """Return ``(project, unlabelled_folder, labelled_folder, search_space_tag)`` pairs.""" + unlabelled = _discover_unlabelled_folders(unlabelled_dir) + labelled = _discover_labelled_folders(labelled_dir) + + keys = sorted(unlabelled.keys() & labelled.keys()) + if projects_filter is not None: + keys = [k for k in keys if k in projects_filter] + + pairs: list[tuple[str, Path, Path, str]] = [] + for key in keys: + pairs.append( + ( + key, + unlabelled[key], + labelled[key], + _search_space_tag_from_folder(unlabelled[key].name), + ) + ) + + for key in sorted(unlabelled.keys() - labelled.keys()): + if projects_filter is None or key in projects_filter: + logger.warning("No labelled folder for unlabelled project %r", key) + for key in sorted(labelled.keys() - unlabelled.keys()): + if projects_filter is None or key in projects_filter: + logger.warning("No unlabelled folder for labelled project %r", key) + + return pairs + + +def _effective_db_grounded_drop(n_rows: int, drop: int = _DB_GROUNDED_DROP) -> int: + return min(drop, max(0, n_rows - 1)) + + +def _add_q_values( + df: pd.DataFrame, + conf_col: str = "calibrated_confidence", + *, + q_col: str = "psm_q_value", +) -> pd.DataFrame: + """Attach PSM q-values from a non-parametric FDR fit on *conf_col*.""" + if q_col in df.columns: + return df + if conf_col not in df.columns: + raise ValueError(f"Missing confidence column {conf_col!r}") + + existing_q = df["psm_q_value"] if "psm_q_value" in df.columns else None + work = df.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") + + fdr = NonParametricFDRControl() + fdr.fit(dataset=work[conf_col]) + out = fdr.add_psm_q_value(work, confidence_col=conf_col) + if q_col != "psm_q_value": + out = out.rename(columns={"psm_q_value": q_col}) + if existing_q is not None and q_col != "psm_q_value": + out["psm_q_value"] = existing_q + return out + + +def _add_database_grounded_q_values( + df: pd.DataFrame, + residue_masses: dict[str, float], + confidence_col: str = RAW_CONFIDENCE_COL, + *, + q_col: str = DB_Q_VALUE_COL, + correct_col: str = "correct", +) -> pd.DataFrame: + """Attach PSM q-values from database-grounded FDR on *confidence_col*.""" + if q_col in df.columns: + return df + if confidence_col not in df.columns: + raise ValueError(f"Missing confidence column {confidence_col!r}") + if "sequence" not in df.columns or "prediction" not in df.columns: + raise ValueError( + "Database-grounded FDR requires 'sequence' and 'prediction' columns" + ) + + # Drop Winnow NP q-values before fitting; merge-based add_psm_q_value can + # leave duplicate psm_q_value_* columns and skip the rename to db_psm_q_value. + work = df.drop( + columns=[q_col, "psm_q_value", "psm_fdr", "fdr"], + errors="ignore", + ).copy() + ctrl = DatabaseGroundedFDRControl( + confidence_feature=confidence_col, + residue_masses=residue_masses, + drop=_effective_db_grounded_drop(len(work)), + ) + ctrl.fit(dataset=work.copy(), correct_column=correct_col) + q_df = ctrl.add_psm_q_value( + work[[confidence_col]].copy(), confidence_col=confidence_col + ) + work[q_col] = q_df["psm_q_value"].values + return work + + +def _unique_peptides_at_fdr( + df: pd.DataFrame, + sequence_col: str, + q_col: str, + fdr_t: float, +) -> set[str]: + retained = df[df[q_col] <= fdr_t] + return set(retained[sequence_col].dropna().map(_sequence_match_key)) - {""} + + +def _empty_overlap_row( + project: str, + eval_type: str, + fdr_t: float, + n_db_peptides: int, + labelled_subset: bool, +) -> dict: + row: dict = { + "project": project, + "eval_type": eval_type, + "fdr_threshold": fdr_t, + "n_psms_retained": 0, + "n_unique_peptides_retained": 0, + "n_db_search_peptides": n_db_peptides, + "n_matching": 0, + "pct_matching": 0.0, + "n_discordant": 0, + "pct_discordant": 0.0, + } + for col in _DISCORDANCE_COUNT_COLS: + if col == "n_partial_match" and not labelled_subset: + continue + row[col] = 0 + return row + + +def _discordance_cache_for_fdr_retained( + df: pd.DataFrame, + db_keys: set[str], + db_stripped_set: set[str], + db_stripped_by_len: dict[int, list[str]], + *, + labelled_subset: bool, + residue_masses: dict[str, float] | None, +) -> LabelledDiscordanceCache | DbDiscordanceCache: + """Build discordance lookup for all predictions retained at any FDR threshold.""" + df = _add_q_values(df) + retained = df[df["psm_q_value"] <= max(FDR_THRESHOLDS)] + if labelled_subset and residue_masses is not None: + cache: LabelledDiscordanceCache = {} + disc = retained[ + retained["prediction"].map(_sequence_match_key) + != retained["sequence"].map(_sequence_match_key) + ] + for _, row in disc.drop_duplicates( + subset=["sequence", "prediction"] + ).iterrows(): + trip = ( + str(row.get("sequence", "")), + str(row.get("prediction", "")), + int(row.get("num_matches", 0)), + ) + if trip not in cache: + cache[trip] = classify_discordance_labelled( + trip[0], trip[1], trip[2], residue_masses + ) + return cache + + return _build_discordance_cache( + retained["prediction"], db_keys, db_stripped_set, db_stripped_by_len + ) + + +def compute_overlap_table( + df: pd.DataFrame, + project: str, + eval_type: str, + db_df: pd.DataFrame, + discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, + residue_masses: dict[str, float], + *, + labelled_subset: bool = False, +) -> pd.DataFrame: + """Overlap summary at each FDR threshold. + + Winnow uses non-parametric FDR on calibrated confidence (full-search run). + Database reference peptides use database-grounded FDR on raw confidence in the + database-labelled run (see ``plot_full_search_venn``). + """ + df = _add_q_values(df.copy()) + db_scored = _add_database_grounded_q_values( + db_df.copy(), + residue_masses, + confidence_col=RAW_CONFIDENCE_COL, + q_col=DB_Q_VALUE_COL, + ) + + rows: list[dict] = [] + for fdr_t in FDR_THRESHOLDS: + db_keys_at_fdr = _unique_peptides_at_fdr( + db_scored, "sequence", DB_Q_VALUE_COL, fdr_t + ) + n_db_peptides = len(db_keys_at_fdr) + + retained = df[df["psm_q_value"] <= fdr_t].copy() + n_retained = len(retained) + if n_retained == 0: + rows.append( + _empty_overlap_row( + project, eval_type, fdr_t, n_db_peptides, labelled_subset + ) + ) + continue + + if labelled_subset: + retained["pred_key"] = retained["prediction"].map(_sequence_match_key) + retained["seq_key"] = retained["sequence"].map(_sequence_match_key) + is_match = retained["pred_key"] == retained["seq_key"] + else: + is_match = retained["prediction"].map( + lambda p: _is_db_match(str(p), db_keys_at_fdr) + ) + + n_matching = int(is_match.sum()) + n_discordant = n_retained - n_matching + n_unique_peptides = int( + retained["prediction"].map(_sequence_match_key).replace("", pd.NA).nunique() + ) + + disc = retained[~is_match] + cat_counts: dict[str, int] = {} + if len(disc) > 0: + if labelled_subset: + labelled_cache = cast(LabelledDiscordanceCache, discordance_cache) + cats = disc.apply( + _lookup_labelled_discordance, axis=1, cache=labelled_cache + ) + else: + cats = _classify_predictions_vs_db( + disc["prediction"], + cast(DbDiscordanceCache, discordance_cache), + ) + cat_counts = cats.value_counts().to_dict() + + row: dict = { + "project": project, + "eval_type": eval_type, + "fdr_threshold": fdr_t, + "n_psms_retained": n_retained, + "n_unique_peptides_retained": n_unique_peptides, + "n_db_search_peptides": n_db_peptides, + "n_matching": n_matching, + "pct_matching": round(n_matching / n_retained * 100, 2), + "n_discordant": n_discordant, + "pct_discordant": round(n_discordant / n_retained * 100, 2), + "n_ptm_candidate": cat_counts.get("ptm_candidate", 0), + "n_single_aa_variant": cat_counts.get("single_aa_variant", 0), + "n_near_miss_edit_dist": cat_counts.get("near_miss_edit_dist", 0), + "n_fully_discordant": cat_counts.get("fully_discordant", 0), + } + if labelled_subset: + row["n_partial_match"] = cat_counts.get("partial_match", 0) + rows.append(row) + + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Plots +# --------------------------------------------------------------------------- +def _plot_venn_panels( + winnow_df: pd.DataFrame, + project: str, + output_path: Path, + *, + winnow_label: str, + suptitle: str, + residue_masses: dict[str, float] | None = None, + db_df: pd.DataFrame | None = None, + db_peptides_static: set[str] | None = None, +) -> None: + if db_df is None and db_peptides_static is None: + raise ValueError("Provide db_df or db_peptides_static for Venn panels") + + winnow_scored = _add_q_values(winnow_df.copy()) + db_scored: pd.DataFrame | None = None + if db_df is not None: + if residue_masses is None: + raise ValueError("residue_masses required when db_df is provided") + db_scored = _add_database_grounded_q_values( + db_df.copy(), + residue_masses, + confidence_col=RAW_CONFIDENCE_COL, + q_col=DB_Q_VALUE_COL, + ) + + n_thresholds = len(FDR_THRESHOLDS) + fig, axes = plt.subplots(1, n_thresholds, figsize=(5 * n_thresholds, 5)) + if n_thresholds == 1: + axes = [axes] + + for ax, fdr_t in zip(axes, FDR_THRESHOLDS): + winnow_peptides = _unique_peptides_at_fdr( + winnow_scored, "prediction", "psm_q_value", fdr_t + ) + if db_scored is not None: + db_peptides = _unique_peptides_at_fdr( + db_scored, "sequence", DB_Q_VALUE_COL, fdr_t + ) + else: + assert db_peptides_static is not None + db_peptides = db_peptides_static + + pct = int(fdr_t * 100) + if not winnow_peptides and not db_peptides: + ax.set_title(f"No peptides retained at {pct}% FDR") + ax.axis("off") + continue + + if not winnow_peptides or not db_peptides: + missing = "Winnow" if not winnow_peptides else "Database search" + ax.text( + 0.5, + 0.5, + f"No {missing} peptides at {pct}% FDR", + ha="center", + va="center", + transform=ax.transAxes, + ) + ax.set_title(f"{pct}% FDR") + ax.axis("off") + continue + + venn2( + [db_peptides, winnow_peptides], + set_labels=("Database search", winnow_label), + set_colors=(_CORRECT_COLOUR, _INCORRECT_COLOUR), + alpha=0.6, + ax=ax, + ) + ax.set_title(f"Unique peptides at {pct}% FDR") + _style_ax(ax) + + fig.suptitle(suptitle, fontsize=12) + fig.tight_layout() + _save_fig(fig, output_path) + + +def plot_full_search_venn( + df: pd.DataFrame, + db_df: pd.DataFrame, + project: str, + plots_dir: Path, + residue_masses: dict[str, float], +) -> None: + """Venn diagrams of FDR-filtered DB vs Winnow full-search unique peptides. + + Database peptides use database-grounded q-values on raw ``confidence`` in the + database-labelled run. Winnow peptides use non-parametric q-values on + ``calibrated_confidence`` in the full-search run. + """ + display = _display_name(project) + _plot_venn_panels( + df, + project, + plots_dir / f"venn_{project}_full_search", + winnow_label="Winnow", + suptitle=f"Database search vs Winnow full search at matched FDR for {display}", + residue_masses=residue_masses, + db_df=db_df, + ) + + +def plot_labelled_subset_venn( + df: pd.DataFrame, + project: str, + plots_dir: Path, +) -> None: + """Venn diagrams of labelled reference peptides vs Winnow predictions per FDR.""" + db_peptides = { + _sequence_match_key(s) + for s in df["sequence"].dropna() + if _sequence_match_key(s) + } + display = _display_name(project) + _plot_venn_panels( + df, + project, + plots_dir / f"venn_{project}_labelled_subset", + winnow_label="Winnow", + suptitle=f"Peptide overlap on labelled spectra for {display}", + db_peptides_static=db_peptides, + ) + + +def _assign_retained_groups( + retained: pd.DataFrame, + db_keys: set[str] | None, + discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, + *, + labelled_subset: bool, +) -> pd.Series: + if labelled_subset and "sequence" in retained.columns: + labelled_cache = cast(LabelledDiscordanceCache, discordance_cache) + is_match = retained["prediction"].map(_sequence_match_key) == retained[ + "sequence" + ].map(_sequence_match_key) + groups = pd.Series("Database match", index=retained.index) + disc_mask = ~is_match + if disc_mask.any(): + disc = retained.loc[disc_mask] + groups.loc[disc_mask] = disc.apply( + _lookup_labelled_discordance, axis=1, cache=labelled_cache + ).values + return groups + + db_cache = cast(DbDiscordanceCache, discordance_cache) + if db_keys is None: + raise ValueError("db_keys required for full-search retained-group assignment") + is_match = retained["prediction"].map(lambda p: _is_db_match(str(p), db_keys)) + groups = pd.Series("Database match", index=retained.index) + disc_mask = ~is_match + if disc_mask.any(): + groups.loc[disc_mask] = _classify_predictions_vs_db( + retained.loc[disc_mask, "prediction"], + db_cache, + ).values + return groups + + +def _subsample_violin_groups(df: pd.DataFrame, category_col: str) -> pd.DataFrame: + """Limit points per category so violin plots stay responsive.""" + parts: list[pd.DataFrame] = [] + for _cat, group in df.groupby(category_col, observed=True): + if len(group) > _MAX_VIOLIN_PSMs_PER_GROUP: + group = group.sample( + n=_MAX_VIOLIN_PSMs_PER_GROUP, + random_state=_VIOLIN_SUBSAMPLE_SEED, + ) + parts.append(group) + return pd.concat(parts, ignore_index=True) if parts else df + + +def plot_novel_feature_violins( + df: pd.DataFrame, + db_df: pd.DataFrame | None, + residue_masses: dict[str, float] | None, + discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, + project: str, + eval_type: str, + plots_dir: Path, + *, + labelled_subset: bool = False, +) -> None: + """Violin plots: database-matched vs fully discordant retained PSMs. + + Winnow retention uses non-parametric FDR on calibrated confidence. For + full-search comparisons, database match uses peptides retained at the same + nominal FDR via database-grounded FDR on raw confidence (see overlap table). + """ + available = [ + (col, label) for col, label in NOVEL_FEATURE_COLUMNS if col in df.columns + ] + if not available: + logger.warning( + "%s: no feature columns for violin plots (metadata merge missing?)", + project, + ) + return + + df = _add_q_values(df.copy()) + display = _display_name(project) + + db_scored: pd.DataFrame | None = None + if not labelled_subset: + if db_df is None or residue_masses is None: + raise ValueError( + "Full-search violin plots require db_df and residue_masses" + ) + db_scored = _add_database_grounded_q_values( + db_df.copy(), + residue_masses, + confidence_col=RAW_CONFIDENCE_COL, + q_col=DB_Q_VALUE_COL, + ) + + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t].copy() + if len(retained) < _MIN_VIOLIN_GROUP_SIZE: + logger.info( + "%s: skip violins at %d%% FDR (n=%d retained)", + project, + int(fdr_t * 100), + len(retained), + ) + continue + + if labelled_subset: + match_keys: set[str] | None = None + else: + assert db_scored is not None + match_keys = _unique_peptides_at_fdr( + db_scored, "sequence", DB_Q_VALUE_COL, fdr_t + ) + + groups = _assign_retained_groups( + retained, + match_keys, + discordance_cache, + labelled_subset=labelled_subset, + ) + retained = retained.assign(_overlap_group=groups) + plot_df = retained[ + retained["_overlap_group"].isin(["Database match", "fully_discordant"]) + ].copy() + plot_df["Category"] = plot_df["_overlap_group"].map( + { + "Database match": "Database match", + "fully_discordant": "Novel", + } + ) + plot_df = _subsample_violin_groups(plot_df, "Category") + + n_match = (plot_df["Category"] == "Database match").sum() + n_novel = (plot_df["Category"] == "Novel").sum() + if n_match < _MIN_VIOLIN_GROUP_SIZE or n_novel < _MIN_VIOLIN_GROUP_SIZE: + logger.info( + "%s: skip violins at %d%% FDR (match=%d, novel=%d)", + project, + int(fdr_t * 100), + n_match, + n_novel, + ) + continue + + n_feats = len(available) + n_cols = 4 + n_rows = int(np.ceil(n_feats / n_cols)) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows)) + axes_flat = np.atleast_1d(axes).flatten() + + palette = {"Database match": _CORRECT_COLOUR, "Novel": _NOVEL_COLOUR} + cat_order = ["Database match", "Novel"] + + for ax, (col, label) in zip(axes_flat, available): + sub = plot_df[[col, "Category"]].dropna() + if sub["Category"].nunique() < 2: + ax.set_visible(False) + continue + sns.violinplot( + data=sub, + x="Category", + y=col, + order=cat_order, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + ax.set_ylabel(label) + ax.tick_params(axis="x", rotation=15) + _style_ax(ax) + + for ax in axes_flat[len(available) :]: + ax.set_visible(False) + + pct = int(fdr_t * 100) + fig.suptitle( + f"{display} ({_eval_type_display(eval_type)}): " + f"database-matched vs novel features at {pct}% FDR", + fontsize=12, + ) + fig.tight_layout() + _save_fig(fig, plots_dir / f"novel_feature_violins_{project}_fdr{pct}") + + +# --------------------------------------------------------------------------- +# Per-project orchestration +# --------------------------------------------------------------------------- +def generate_all_analyses( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, + db_df: pd.DataFrame | None, + residue_masses: dict[str, float], +) -> pd.DataFrame: + """Tables and plots for one project.""" + output_dir.mkdir(parents=True, exist_ok=True) + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + if db_df is None: + raise ValueError(f"Full-search analysis requires DB reference for {project}") + db_keys, db_stripped_set, _, db_by_len = _build_db_reference_sets(db_df) + disc_cache = _discordance_cache_for_fdr_retained( + df, + db_keys, + db_stripped_set, + db_by_len, + labelled_subset=False, + residue_masses=None, + ) + overlap = compute_overlap_table( + df, + project, + eval_type, + db_df, + disc_cache, + residue_masses, + labelled_subset=False, + ) + plot_full_search_venn(df, db_df, project, plots_dir, residue_masses) + plot_novel_feature_violins( + df, + db_df, + residue_masses, + disc_cache, + project, + eval_type, + plots_dir, + labelled_subset=False, + ) + + overlap.to_csv(output_dir / f"{project}_overlap_summary.csv", index=False) + with open(output_dir / f"{project}_overlap_summary.json", "w") as f: + json.dump(overlap.to_dict(orient="records"), f, indent=2) + + logger.info("\n%s", overlap.to_string(index=False)) + return overlap + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +@app.command() +def main( + unlabelled_dir: Annotated[ + Path, + typer.Option( + "--unlabelled-dir", + help="Root with per-project full-search folders (e.g. gluc_raw/, PXD014877_unlabelled/).", + ), + ], + labelled_dir: Annotated[ + Path, + typer.Option( + "--labelled-dir", + help="Root with per-project database-search folders (e.g. gluc_annotated/, PXD014877_labelled/).", + ), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for tables and plots."), + ], + projects: Annotated[ + str | None, + typer.Option( + "--projects", + help="Optional space- or comma-separated project keys to restrict analysis.", + ), + ] = None, +) -> None: + """Post-FDR overlap: full-search Winnow identifications vs database search.""" + logging.basicConfig(level=logging.INFO, format="%(message)s", datefmt="%H:%M:%S") + + projects_filter: set[str] | None = None + if projects is not None: + project_list = [ + p.strip() for p in projects.replace(",", " ").split() if p.strip() + ] + if not project_list: + raise typer.BadParameter("No projects specified in --projects.") + projects_filter = set(project_list) + + pairs = discover_project_pairs( + unlabelled_dir, labelled_dir, projects_filter=projects_filter + ) + if not pairs: + logger.error( + "No paired projects found under unlabelled-dir=%s and labelled-dir=%s", + unlabelled_dir, + labelled_dir, + ) + raise typer.Exit(code=1) + + output_dir.mkdir(parents=True, exist_ok=True) + residue_masses = _get_residue_masses() + + all_tables: list[pd.DataFrame] = [] + for project, unlabelled_folder, labelled_folder, search_tag in pairs: + display = _display_name(project) + logger.info( + "Processing %s (%s): %s vs %s", + project, + display, + unlabelled_folder.name, + labelled_folder.name, + ) + + try: + df = _load_from_folder(unlabelled_folder) + db_df = _load_from_folder(labelled_folder) + except FileNotFoundError as exc: + logger.warning("Skipping %s: %s", project, exc) + continue + + logger.info( + " Full search: %d rows; DB reference: %d rows", len(df), len(db_df) + ) + + try: + table = generate_all_analyses( + df, + project, + search_tag, + output_dir, + db_df, + residue_masses, + ) + all_tables.append(table) + except ValueError as exc: + logger.warning("Skipping %s: %s", project, exc) + + if not all_tables: + logger.error("No projects produced overlap output under %s", output_dir) + raise typer.Exit(code=1) + + combined = pd.concat(all_tables, ignore_index=True) + combined.to_csv(output_dir / "all_projects_overlap_summary.csv", index=False) + with open(output_dir / "all_projects_overlap_summary.json", "w") as f: + json.dump(combined.to_dict(orient="records"), f, indent=2) + + logger.info("FDR overlap analysis complete. Output in %s", output_dir) + + +if __name__ == "__main__": + app() diff --git a/scripts/analyze_features.py b/scripts/analyze_features.py new file mode 100644 index 00000000..e8a6ec51 --- /dev/null +++ b/scripts/analyze_features.py @@ -0,0 +1,802 @@ +"""Analyze feature importance and correlations for a pretrained calibrator. + +This script provides comprehensive analysis of feature importance: + - Permutation importance on test set + - SHAP values with training background on test set + - Feature correlation analysis on training data + - Optional visualization of results +""" + +import logging +import pickle +from pathlib import Path +from typing import Annotated, Any, Dict, List, Optional + +import matplotlib.pyplot as plt +from matplotlib.colors import LinearSegmentedColormap +import numpy as np +import pandas as pd +import seaborn as sns +import shap +import torch +import typer +import yaml +from rich.console import Console +from rich.theme import Theme +from sklearn.inspection import permutation_importance + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.datasets.data_loaders import InstaNovoDatasetLoader + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette + "sunset" diverging colourmap +# --------------------------------------------------------------------------- +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +_SUNSET_COLORS = [ + "#364B9A", + "#4A7BB7", + "#6EA6CD", + "#98CAE1", + "#C2E4EF", + "#EAECCC", + "#FEDA8B", + "#FDB366", + "#F67E4B", + "#DD3D2D", + "#A50026", +] +_BAD_COLOUR = "#FFFFFF" + + +def _sunset_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) + cmap.set_bad(color=_BAD_COLOUR) + return cmap + + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + + +def _style_axes(ax: plt.Axes) -> None: + """Apply standard axes formatting: no grid, black spines.""" + ax.set_axisbelow(True) + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +logger = logging.getLogger("winnow") +logger.setLevel(logging.INFO) + +logging.getLogger("shap").setLevel(logging.WARNING) + +# --------------------------------------------------------------------------- +# Constants — loaded from the canonical Winnow YAML configs +# --------------------------------------------------------------------------- +SEED = 42 + +_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "winnow" / "configs" + +with open(_CONFIGS_DIR / "residues.yaml") as _f: + RESIDUE_MASSES: dict[str, float] = yaml.safe_load(_f)["residue_masses"] + +with open(_CONFIGS_DIR / "data_loader" / "instanovo.yaml") as _f: + _instanovo_cfg = yaml.safe_load(_f) + RESIDUE_REMAPPING: dict[str, str] = _instanovo_cfg.get("residue_remapping", {}) + BEAM_COLUMNS: dict[str, str] | None = _instanovo_cfg.get("beam_columns") + +COLUMN_DISPLAY_NAMES = { + "confidence": "Raw confidence", + "mass_error": "Mass error", + "mass_error_ppm": "Mass error (ppm)", + "mass_error_da": "Mass error (Da)", + "spectral_angle": "Spectral angle", + "ion_matches": "Ion matches", + "ion_match_intensity": "Ion match intensity", + "chimeric_ion_matches": "Chimeric ion matches", + "chimeric_ion_match_intensity": "Chimeric ion match intensity", + "irt_error": "iRT error", + "margin": "Margin", + "median_margin": "Median margin", + "entropy": "Entropy", + "z-score": "Z-score", + "edit_distance": "Edit distance", + "xcorr": "XCorr", + "chimeric_xcorr": "Chimeric XCorr", + "longest_ion_series": "Longest ion series", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Max ion gap", + "chimeric_longest_ion_series": "Chimeric longest ion series", + "chimeric_complementary_ion_count": "Chimeric complementary ion count", + "chimeric_max_ion_gap": "Chimeric max ion gap", + "is_missing_fragment_match_features": "Missing fragment match", + "is_missing_chimeric_features": "Missing chimeric", + "is_missing_irt_error": "Missing iRT", + "sequence_length": "Sequence length", + "precursor_charge": "Precursor charge", + "min_token_probability": "Min token probability", + "std_token_probability": "Std token probability", +} + +error_theme = Theme({"error": "red bold", "error_highlight": "red bold underline"}) +console = Console(theme=error_theme) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +def feature_display_name(column: str) -> str: + """Human-readable label for a feature column.""" + if column in COLUMN_DISPLAY_NAMES: + return COLUMN_DISPLAY_NAMES[column] + return column.replace("_", " ").replace("-", " ").strip().title() + + +def to_sentence_case(name: str) -> str: + """Convert a feature display name to sentence case.""" + return name.lower() + + +_SUPPORTED_EXTENSIONS = {".parquet", ".ipc", ".mgf"} + + +def _load_and_compute_features( + spectrum_path: Path, + predictions_path: Path, + loader: InstaNovoDatasetLoader, + calibrator: ProbabilityCalibrator, +) -> CalibrationDataset: + """Load spectra (single file or directory) and compute calibration features. + + When *spectrum_path* is a directory the contained spectrum files are + processed one at a time and the resulting metadata frames are + concatenated, mirroring the batched logic in + ``winnow.scripts.main._compute_features_directory``. + """ + if spectrum_path.is_dir(): + files = sorted( + f for f in spectrum_path.iterdir() if f.suffix in _SUPPORTED_EXTENSIONS + ) + if not files: + raise FileNotFoundError( + f"No spectrum files found in {spectrum_path}. " + f"Supported extensions: {', '.join(sorted(_SUPPORTED_EXTENSIONS))}" + ) + all_metadata: list[pd.DataFrame] = [] + for file_path in files: + logger.info(" Processing experiment file: %s", file_path.name) + ds = loader.load(data_path=file_path, predictions_path=predictions_path) + calibrator.compute_features(ds) + all_metadata.append(ds.metadata) + combined = pd.concat(all_metadata, ignore_index=True) + return CalibrationDataset(metadata=combined, predictions=None) + + dataset = loader.load(data_path=spectrum_path, predictions_path=predictions_path) + calibrator.compute_features(dataset) + return dataset + + +# --------------------------------------------------------------------------- +# Model wrapper for sklearn-compatible predict_proba +# --------------------------------------------------------------------------- +class _CalibratorPredictor: + """Wraps a fitted ProbabilityCalibrator as an sklearn-style estimator. + + Provides ``predict_proba`` and ``predict`` on *pre-normalised* feature + arrays so that permutation importance and SHAP can treat it like a + classifier. The ``classes_`` attribute is set to ``[0, 1]``. + """ + + def __init__(self, calibrator: ProbabilityCalibrator) -> None: + assert calibrator.network is not None + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + + self.network = calibrator.network + self.feature_mean = calibrator.feature_mean.cpu() + self.feature_std = calibrator.feature_std.cpu() + self.classes_ = np.array([0, 1]) + + def fit(self, x_input: np.ndarray, y: np.ndarray) -> "_CalibratorPredictor": + """No-op fit to satisfy sklearn estimator interface.""" + return self + + def score(self, x_input: np.ndarray, y: np.ndarray) -> float: + """Return accuracy to satisfy sklearn estimator interface.""" + return float(np.mean(self.predict(x_input) == y)) + + def predict_proba(self, x_input: np.ndarray) -> np.ndarray: # noqa: N803 + """Return class probabilities for each sample.""" + x = torch.as_tensor(x_input, dtype=torch.float32) + self.network.eval() + with torch.no_grad(): + logits = self.network(x) + probs = torch.sigmoid(logits).numpy().ravel() + return np.column_stack([1 - probs, probs]) + + def predict(self, x_input: np.ndarray) -> np.ndarray: # noqa: N803 + """Return binary predictions for each sample.""" + return (self.predict_proba(x_input)[:, 1] >= 0.5).astype(int) + + +# --------------------------------------------------------------------------- +# Plotting functions +# --------------------------------------------------------------------------- +def plot_feature_importance( + importance_scores: Dict[str, float], + title: str, + output_path_base: Path, +) -> None: + """Plot horizontal bar chart of permutation feature importance scores.""" + plt.figure(figsize=(8, 6)) + features = list(importance_scores.keys()) + scores = list(importance_scores.values()) + + sorted_idx = np.argsort(scores) + features = [features[i] for i in sorted_idx] + scores = [scores[i] for i in sorted_idx] + + plt.barh(range(len(features)), scores, color=_PALETTE[0]) + plt.yticks(range(len(features)), features) + plt.xlabel("Importance score") + plt.title(title) + _style_axes(plt.gca()) + + plt.savefig(f"{output_path_base}.pdf", bbox_inches="tight", dpi=300) + plt.savefig(f"{output_path_base}.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_feature_correlations(features: pd.DataFrame, output_path_base: Path) -> None: + """Plot lower-triangle feature correlation heatmap.""" + plt.figure(figsize=(12, 10)) + corr_matrix = features.corr() + mask = np.triu(np.ones_like(corr_matrix, dtype=bool)) + + sns.heatmap( + corr_matrix, + mask=mask, + cmap=_sunset_cmap(), + vmin=-1, + vmax=1, + center=0, + square=True, + annot=True, + fmt=".2f", + cbar_kws={"shrink": 0.5}, + ) + plt.title("Feature correlation matrix") + _style_axes(plt.gca()) + + plt.savefig(f"{output_path_base}.pdf", bbox_inches="tight", dpi=300) + plt.savefig(f"{output_path_base}.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_shap_summary(shap_values, correct_class_idx: int, output_dir: Path) -> None: + """Plot SHAP beeswarm summary for the correct class.""" + plt.figure(figsize=(8, 6)) + shap.plots.beeswarm( + shap_values[:, :, correct_class_idx], + show=False, + max_display=12, + color=_sunset_cmap(), + ) + plt.title(r"SHAP feature impact on $P(\text{correct})$") + _style_axes(plt.gca()) + + plt.savefig(output_dir / "shap_summary.pdf", bbox_inches="tight", dpi=300) + plt.savefig(output_dir / "shap_summary.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_shap_bar( + shap_values, + test_features_scaled, + test_labels, + correct_class_idx: int, + output_dir: Path, +) -> None: + """Plot SHAP bar chart with hierarchical clustering.""" + plt.figure(figsize=(8, 6)) + clustering = shap.utils.hclust(test_features_scaled, test_labels) + shap.plots.bar( + shap_values[:, :, correct_class_idx], + clustering=clustering, + show=False, + clustering_cutoff=0.5, + max_display=12, + ) + ax = plt.gca() + for patch in ax.patches: + patch.set_facecolor(_PALETTE[1]) + plt.title(r"SHAP feature importance for $P(\text{correct})$") + _style_axes(plt.gca()) + + plt.savefig(output_dir / "shap_importance.pdf", bbox_inches="tight", dpi=300) + plt.savefig(output_dir / "shap_importance.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_shap_dependence( + shap_values, + feature_names: list, + display_feature_names: list, + correct_class_idx: int, + output_dir: Path, + top_n: int = 3, +) -> None: + """Plot SHAP dependence scatter for the top-N most important features.""" + mean_abs_shap = np.abs(shap_values.values[:, :, correct_class_idx]).mean(axis=0) + top_features_idx = np.argsort(mean_abs_shap)[-top_n:][::-1] + + for idx in top_features_idx: + feature_name = display_feature_names[idx] + original_feature_name = feature_names[idx] + plt.figure(figsize=(8, 6)) + shap.plots.scatter( + shap_values[:, idx, correct_class_idx], + show=False, + color=_PALETTE[2], + ) + plt.title( + "SHAP dependence plot for " + + to_sentence_case(feature_name) + + "\n" + + r"(impact on $P(\text{correct})$)" + ) + + ax = plt.gca() + ylabel = ax.get_ylabel() + if "SHAP value for" in ylabel: + ax.set_ylabel(ylabel.replace(feature_name, to_sentence_case(feature_name))) + + _style_axes(ax) + plt.savefig( + output_dir / f"shap_dependence_{original_feature_name}.pdf", + bbox_inches="tight", + dpi=300, + ) + plt.savefig( + output_dir / f"shap_dependence_{original_feature_name}.png", + bbox_inches="tight", + dpi=300, + ) + plt.close() + + +def plot_shap_interactions( + shap_values, + feature_names: list, + display_feature_names: list, + correct_class_idx: int, + output_dir: Path, + top_n: int = 3, +) -> None: + """Plot pairwise SHAP interaction scatter plots for top-N features.""" + mean_abs_shap = np.abs(shap_values.values[:, :, correct_class_idx]).mean(axis=0) + top_features_idx = np.argsort(mean_abs_shap)[-top_n:][::-1] + + for i, idx1 in enumerate(top_features_idx): + f1_display = display_feature_names[idx1] + f1_orig = feature_names[idx1] + for j, idx2 in enumerate(top_features_idx): + if i == j: + continue + f2_display = display_feature_names[idx2] + f2_orig = feature_names[idx2] + + plt.figure(figsize=(8, 6)) + shap.plots.scatter( + shap_values[:, f1_display, correct_class_idx], + color=shap_values[:, f2_display, correct_class_idx], + show=False, + cmap=_sunset_cmap(), + ) + plt.title( + "SHAP interaction plot for " + + to_sentence_case(f1_display) + + " vs " + + to_sentence_case(f2_display) + + "\n" + + r" (impact on $P(\text{correct})$)" + ) + + ax = plt.gca() + ylabel = ax.get_ylabel() + if "SHAP value for" in ylabel: + ax.set_ylabel(ylabel.replace(f1_display, to_sentence_case(f1_display))) + + _style_axes(ax) + plt.savefig( + output_dir / f"shap_interaction_{f1_orig}_vs_{f2_orig}.pdf", + bbox_inches="tight", + dpi=300, + ) + plt.savefig( + output_dir / f"shap_interaction_{f1_orig}_vs_{f2_orig}.png", + bbox_inches="tight", + dpi=300, + ) + plt.close() + + +def plot_shap_heatmap(shap_values, correct_class_idx: int, output_dir: Path) -> None: + """Plot SHAP heatmap showing per-sample feature contributions.""" + plt.figure(figsize=(8, 6)) + shap.plots.heatmap( + shap_values[:, :, correct_class_idx], + max_display=12, + show=False, + cmap=_sunset_cmap(), + ) + plt.title("SHAP feature impact heatmap\n" + r"(impact on $P(\text{correct})$)") + _style_axes(plt.gca()) + + plt.savefig(output_dir / "shap_heatmap.pdf", bbox_inches="tight", dpi=300) + plt.savefig(output_dir / "shap_heatmap.png", bbox_inches="tight", dpi=300) + plt.close() + + +def create_all_plots( + perm_importance_dict: Dict[str, float], + shap_values, + train_features_scaled_df: pd.DataFrame, + test_features_scaled: np.ndarray, + test_labels: np.ndarray, + feature_names: list, + display_feature_names: list, + correct_class_idx: int, + output_dir: Path, +) -> None: + """Generate all analysis plots (importance, correlations, SHAP).""" + logger.info("Creating plots...") + + plot_feature_importance( + perm_importance_dict, + "Permutation feature importance", + output_dir / "permutation_importance", + ) + plot_shap_summary(shap_values, correct_class_idx, output_dir) + plot_shap_bar( + shap_values, test_features_scaled, test_labels, correct_class_idx, output_dir + ) + plot_shap_dependence( + shap_values, feature_names, display_feature_names, correct_class_idx, output_dir + ) + plot_shap_interactions( + shap_values, feature_names, display_feature_names, correct_class_idx, output_dir + ) + plot_shap_heatmap(shap_values, correct_class_idx, output_dir) + plot_feature_correlations( + train_features_scaled_df, output_dir / "feature_correlations" + ) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _load_features_from_parquet( + path: Path, + feature_columns: list[str], +) -> tuple[np.ndarray, np.ndarray]: + """Load feature matrix and labels from a Parquet file or directory.""" + import polars as pl + + p = Path(path) + if p.is_dir(): + parquet_files = sorted(p.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError(f"No .parquet files found in {p}") + df = pl.concat([pl.read_parquet(f) for f in parquet_files]) + else: + df = pl.read_parquet(p) + + if "correct" not in df.columns: + raise ValueError(f"Parquet at {path} must contain a 'correct' column") + missing = [c for c in feature_columns if c not in df.columns] + if missing: + raise ValueError(f"Missing feature columns in Parquet: {missing}") + + features = df.select(feature_columns).to_numpy().astype(np.float32) + labels = df["correct"].to_numpy().astype(np.float32) + return features, labels + + +def _parse_koina_constants(raw: Optional[List[str]]) -> Optional[Dict[str, Any]]: + """Parse ``KEY=VALUE`` pairs into a dict, casting numeric strings.""" + if not raw: + return None + out: Dict[str, Any] = {} + for item in raw: + if "=" not in item: + raise typer.BadParameter( + f"Invalid --koina-input-constant format: '{item}'. Expected KEY=VALUE." + ) + key, value = item.split("=", 1) + try: + out[key] = int(value) + except ValueError: + try: + out[key] = float(value) + except ValueError: + out[key] = value + return out + + +@app.command() +def main( + model_path: Annotated[ + Path, typer.Option(help="Path to pretrained calibrator model directory.") + ], + output_dir: Annotated[ + Path, typer.Option(help="Directory to save analysis results and plots.") + ], + data_dir: Annotated[ + Optional[Path], + typer.Option( + help="Directory containing train and test data files (raw spectra path)." + ), + ] = None, + train_features_path: Annotated[ + Optional[Path], + typer.Option( + help="Path to pre-computed training feature Parquet (alternative to --data-dir)." + ), + ] = None, + test_features_path: Annotated[ + Optional[Path], + typer.Option( + help="Path to pre-computed test feature Parquet (alternative to --data-dir)." + ), + ] = None, + train_spectra: Annotated[ + str, typer.Option(help="Filename of training spectra parquet inside data-dir.") + ] = "general_train.parquet", + train_preds: Annotated[ + str, typer.Option(help="Filename of training predictions CSV inside data-dir.") + ] = "general_train_beams.csv", + test_spectra: Annotated[ + str, typer.Option(help="Filename of test spectra parquet inside data-dir.") + ] = "general_test.parquet", + test_preds: Annotated[ + str, typer.Option(help="Filename of test predictions CSV inside data-dir.") + ] = "general_test_beams.csv", + koina_url: Annotated[ + Optional[str], + typer.Option( + "--koina-url", + help="Override Koina server URL on loaded features (e.g. localhost:8500).", + ), + ] = None, + koina_ssl: Annotated[ + bool, + typer.Option( + "--koina-ssl/--no-koina-ssl", + help="Use SSL for Koina (disable for in-pod Triton).", + ), + ] = True, + koina_input_constant: Annotated[ + Optional[List[str]], + typer.Option( + help="Koina model input constant as KEY=VALUE (repeatable). " + "E.g. --koina-input-constant collision_energies=27 " + "--koina-input-constant fragmentation_types=HCD", + ), + ] = None, + n_background_samples: Annotated[ + int, typer.Option(help="Background samples for SHAP.", min=1, max=10000) + ] = 500, + n_test_samples: Annotated[ + int, typer.Option(help="Test samples for SHAP.", min=1, max=10000) + ] = 1000, + create_plots: Annotated[ + bool, typer.Option("--create-plots/--no-plots", help="Whether to create plots.") + ] = True, +) -> None: + """Analyze feature importance and correlations for a pretrained calibrator.""" + use_parquet = train_features_path is not None or test_features_path is not None + if use_parquet and (train_features_path is None or test_features_path is None): + raise typer.BadParameter( + "--train-features-path and --test-features-path must both be provided." + ) + if not use_parquet and data_dir is None: + raise typer.BadParameter( + "Either --data-dir or --train-features-path/--test-features-path must be provided." + ) + + output_dir.mkdir(parents=True, exist_ok=True) + + # Load calibrator + logger.info("Loading pretrained calibrator from %s", model_path) + calibrator = ProbabilityCalibrator.load(model_path) + + if koina_url is not None or not koina_ssl: + calibrator.apply_koina_server_overrides( + server_url=koina_url, + ssl=koina_ssl, + ) + + koina_constants = _parse_koina_constants(koina_input_constant) + if koina_constants: + logger.info("Applying Koina input constant overrides: %s", koina_constants) + calibrator.apply_koina_model_input_overrides( + model_input_constants=koina_constants, + ) + + # Build sklearn-compatible predictor wrapper + predictor = _CalibratorPredictor(calibrator) + + if use_parquet: + assert train_features_path is not None + assert test_features_path is not None + + feature_columns = ["confidence"] + calibrator.columns + + logger.info("Loading training features from Parquet: %s", train_features_path) + train_features, train_labels = _load_features_from_parquet( + train_features_path, + feature_columns, + ) + logger.info( + " %d training samples, %d features", + len(train_labels), + train_features.shape[1], + ) + + logger.info("Loading test features from Parquet: %s", test_features_path) + test_features, test_labels = _load_features_from_parquet( + test_features_path, + feature_columns, + ) + logger.info( + " %d test samples, %d features", len(test_labels), test_features.shape[1] + ) + + feature_names = feature_columns + else: + assert data_dir is not None + loader = InstaNovoDatasetLoader( + residue_masses=RESIDUE_MASSES, + residue_remapping=RESIDUE_REMAPPING, + beam_columns=BEAM_COLUMNS, + add_index_cols=True, + ) + + logger.info("Loading and computing features for training set...") + train_dataset = _load_and_compute_features( + data_dir / train_spectra, + data_dir / train_preds, + loader, + calibrator, + ) + train_features, train_labels = calibrator._extract_feature_matrix( + train_dataset, labelled=True + ) + + logger.info("Loading and computing features for test set...") + test_dataset = _load_and_compute_features( + data_dir / test_spectra, + data_dir / test_preds, + loader, + calibrator, + ) + test_features, test_labels = calibrator._extract_feature_matrix( + test_dataset, labelled=True + ) + + feature_names = [train_dataset.confidence_column] + calibrator.columns + + display_feature_names = [feature_display_name(name) for name in feature_names] + + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + feature_mean = calibrator.feature_mean.cpu().numpy() + feature_std = calibrator.feature_std.cpu().numpy() + train_features_scaled = (train_features - feature_mean) / feature_std + test_features_scaled = (test_features - feature_mean) / feature_std + + correct_class_idx = 1 # class 1 = correct + + # 1. Permutation importance on test set + logger.info("Computing permutation importance on test set...") + perm_importance = permutation_importance( + predictor, + test_features_scaled, + test_labels, + n_repeats=10, + random_state=SEED, + n_jobs=-1, + ) + perm_importance_dict = dict( + zip(display_feature_names, perm_importance.importances_mean) + ) + + # 2. SHAP values + logger.info("Computing SHAP values...") + background = shap.sample( + train_features_scaled, + min(n_background_samples, len(train_features_scaled)), + random_state=SEED, + ) + + explainer = shap.KernelExplainer( + model=predictor.predict_proba, + data=background, + seed=SEED, + link="identity", + ) + + np.random.seed(SEED) + n_samples = min(n_test_samples, test_features_scaled.shape[0]) + indices = np.random.choice( + test_features_scaled.shape[0], size=n_samples, replace=False + ) + + shap_values = explainer(test_features_scaled[indices]) + + # Switch to original feature space for visualisation + shap_values.data = test_features[indices] + shap_values.feature_names = display_feature_names + + # 3. Feature correlations on training data + logger.info("Computing feature correlations on training data...") + train_features_scaled_df = pd.DataFrame( + train_features_scaled, columns=display_feature_names + ) + + if create_plots: + create_all_plots( + perm_importance_dict=perm_importance_dict, + shap_values=shap_values, + train_features_scaled_df=train_features_scaled_df, + test_features_scaled=test_features_scaled, + test_labels=test_labels, + feature_names=feature_names, + display_feature_names=display_feature_names, + correct_class_idx=correct_class_idx, + output_dir=output_dir, + ) + + # Save raw objects + logger.info("Saving raw analysis objects...") + + with open(output_dir / "perm_importance.pkl", "wb") as f: + pickle.dump(perm_importance, f) + + with open(output_dir / "shap_values.pkl", "wb") as f: + pickle.dump(shap_values, f) + + logger.info("Analysis complete!") + logger.info("Results saved to %s", output_dir) + logger.info( + "Permutation Feature Importance: computed on %d test samples", len(test_labels) + ) + logger.info( + "SHAP values: computed on %d test samples with %d training samples as background", + n_samples, + len(background), + ) + logger.info( + "Correlation matrix: computed on %d training samples", len(train_labels) + ) + + saved_files = ["perm_importance.pkl", "shap_values.pkl"] + if create_plots: + saved_files.append("All plots in PDF and PNG formats") + else: + logger.info("Plots were skipped (--no-plots flag used)") + + logger.info("Saved files: %s", ", ".join(saved_files)) + print(f"\nResults saved to {output_dir}") + + +if __name__ == "__main__": + app() diff --git a/scripts/analyze_novelty.py b/scripts/analyze_novelty.py new file mode 100644 index 00000000..8b87cdc4 --- /dev/null +++ b/scripts/analyze_novelty.py @@ -0,0 +1,1814 @@ +#!/usr/bin/env python3 +"""Analyse Winnow calibrator behaviour on out-of-distribution / novel peptides. + +Two analyses demonstrate that the calibrator does not penalise peptides absent +from the standard tryptic database-search training distribution: + +1. **Non-tryptic enzyme digest (``nontryptic_digest`` subcommand)** -- The model + was trained on tryptic data. Enzymes such as GluC, AspN, LysC or chymotrypsin cleave at + non-K/R sites, so retained peptides often have a C-terminus that is *not* + K or R. We classify predictions by whether their C-terminal residue + is tryptic (K/R) or non-tryptic, report terminus proportions before and + after FDR, compare raw InstaNovo versus Winnow calibrated scores, and + quantify calibration shifts (``calibrated_confidence - confidence``). + + Inputs are **full search space** Winnow predictions (acfm / unlabelled eval: + all candidate spectra, not the labelled database-search subset and not + acfm-minus-lcfm). The ``proteome_hit`` column flags predictions whose + stripped sequence occurs in the reference proteome FASTA; plots and tables + label this cohort **full search space**. + + *Non-tryptic* is defined solely by the C-terminal residue of the + mod-stripped, I/L-normalised prediction. N-terminal context is not + checked because positional information is lost in the substring proteome + match. + +2. **ProteomeTools-1 PXD004732 (``proteometools`` subcommand)** -- Synthetic + peptide library. The *lcfm* set contains database-search-confirmed + peptides; the *acfm* set contains all candidates. For each acfm + prediction we check whether it exactly matches, is a subsequence of, or + shares no overlap with any lcfm peptide. Subsequence matches are + validated novel identifications the search engine missed. + +3. **Summary (``summary`` subcommand)** -- Combines tables from both analyses + into a single grouped bar chart. +""" + +from __future__ import annotations + +import re +import warnings +from pathlib import Path +from typing import Annotated + +import ahocorasick +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import typer +from Bio import SeqIO +from scipy.stats import gaussian_kde + +from winnow.fdr.nonparametric import NonParametricFDRControl + +warnings.filterwarnings("ignore", module="winnow") + +# ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_NOVEL_COLOUR = _PALETTE[2] +_MAIN_LINE_COLOUR = _PALETTE[3] +_RAW_LINE_COLOUR = _PALETTE[5] +_IDEAL_LINE_COLOUR = _PALETTE[6] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)") +_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") +_PROTEOME_JOIN_SEP = "\x1f" + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] + +# Display label for the ``proteome_hit`` cohort (full-search / acfm evaluation). +FULL_SEARCH_SPACE_LABEL = "full search space" +COHORT_FULL_SEARCH = "full_search_space" +COHORT_FULL_SEARCH_AT_FDR = "full_search_space_at_fdr" + +_NONTRYPTIC_CALIBRATION_SCATTER_MAX_POINTS = 10_000 +_NONTRYPTIC_CALIBRATION_SCATTER_RANDOM_STATE = 42 +_C_TERMINUS_ORDER = ("Tryptic (K/R)", "Non-tryptic") +_NONTRYPTIC_GROUP_ALPHA = 0.6 +_TRYPTIC_GROUP_ALPHA = 0.7 + +FEATURE_COLUMNS = [ + "spectral_angle", + "xcorr", + "ion_matches", + "ion_match_intensity", + "irt_error", + "mass_error_ppm", +] + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "PXD004732": "ProteomeTools-1", +} + +app = typer.Typer( + add_completion=False, + no_args_is_help=True, + pretty_exceptions_show_locals=False, +) + + +# ── Shared helpers ──────────────────────────────────────────────────── + + +def _spine_fmt(ax: plt.Axes) -> None: + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _save(fig: plt.Figure, out_dir: Path, name: str) -> None: + base = out_dir / name + fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + print(f" saved {name}") + + +def _subsample_psms( + df: pd.DataFrame, + max_points: int, + random_state: int = _NONTRYPTIC_CALIBRATION_SCATTER_RANDOM_STATE, +) -> pd.DataFrame: + """Return up to ``max_points`` rows without replacement.""" + if len(df) <= max_points: + return df + return df.sample(n=max_points, random_state=random_state) + + +def _strip_mods(seq: str) -> str: + """Strip PTM annotations and normalise I -> L.""" + if not seq or not isinstance(seq, str): + return "" + s = _MOD_PLUS.sub("", seq) + s = _MOD_UNIMOD.sub("", s) + return s.replace("I", "L") + + +def _load_data(predictions_dir: Path) -> pl.DataFrame: + """Load and join ``preds_and_fdr_metrics.csv`` + ``metadata.csv``.""" + preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") + meta_path = predictions_dir / "metadata.csv" + if meta_path.exists(): + meta = pl.read_csv(meta_path) + join_cols = ["spectrum_id"] + [ + c for c in meta.columns if c != "spectrum_id" and c not in preds.columns + ] + if len(join_cols) > 1: + preds = preds.join(meta.select(join_cols), on="spectrum_id", how="inner") + return preds + + +def _add_q_values( + df: pd.DataFrame, + conf_col: str = "calibrated_confidence", +) -> pd.DataFrame: + """Fit non-parametric FDR and append ``psm_q_value`` if missing.""" + if "psm_q_value" in df.columns: + return df + fdr = NonParametricFDRControl() + fdr.fit(dataset=df[conf_col]) + return fdr.add_psm_q_value(df, confidence_col=conf_col) + + +def _load_proteome_haystack(fasta_file: Path) -> str: + """Load a FASTA proteome into a single string for substring matching.""" + parts: list[str] = [] + for record in SeqIO.parse(fasta_file, "fasta"): + s = str(record.seq).replace("I", "L") + if s: + parts.append(s) + return _PROTEOME_JOIN_SEP.join(parts) + + +def _batch_substring_hits( + needles: list[str], + haystack: str, +) -> list[bool]: + """Aho-Corasick batch substring matching.""" + n = len(needles) + out = [False] * n + if not haystack: + return out + + by_needle: dict[str, list[int]] = {} + for i, p in enumerate(needles): + if not p: + continue + by_needle.setdefault(p, []).append(i) + if not by_needle: + return out + + auto = ahocorasick.Automaton() + needle_for_pid: list[str] = [] + for pid, needle in enumerate(by_needle): + auto.add_word(needle, pid) + needle_for_pid.append(needle) + auto.make_automaton() + + matched_pids: set[int] = set() + for _end_idx, pid in auto.iter(haystack): + matched_pids.add(pid) + + for pid in matched_pids: + needle = needle_for_pid[pid] + for row_i in by_needle[needle]: + out[row_i] = True + return out + + +def _nice_feature_label(col: str) -> str: + labels = { + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation", + "ion_matches": "Ion match rate", + "ion_match_intensity": "Ion match intensity", + "irt_error": "iRT prediction error", + "mass_error_ppm": "Mass error (ppm)", + } + return labels.get(col, col.replace("_", " ").capitalize()) + + +# ── Non-tryptic digest analysis ───────────────────────────────────────── + + +def _is_tryptic_cterm(seq: str) -> bool: + """Return True if the mod-stripped C-terminal residue is K or R.""" + stripped = _strip_mods(seq) + if not stripped: + return False + return stripped[-1] in ("K", "R") + + +def _nontryptic_full_search_panel_title(fdr_t: float, n_psms: int) -> str: + """Subplot title for FDR-filtered full-search-space PSMs.""" + pct = int(fdr_t * 100) + return ( + f"{FULL_SEARCH_SPACE_LABEL.title()} identifications at {pct}% FDR\n" + f"(n={n_psms:,})" + ) + + +def _nontryptic_annotate( + df: pl.DataFrame, + fasta_path: Path, +) -> pl.DataFrame: + """Annotate full-search predictions with proteome and tryptic-terminus flags. + + ``proteome_hit`` is True when the mod-stripped prediction is a substring of + the reference proteome FASTA (same rule as ``annotate_preds_proteome_hits``). + """ + haystack = _load_proteome_haystack(fasta_path) + + processed = df["prediction"].map_elements( + lambda x: _strip_mods(x) if isinstance(x, str) else "", + return_dtype=pl.Utf8, + ) + hits = _batch_substring_hits(processed.to_list(), haystack) + tryptic = df["prediction"].map_elements( + lambda x: _is_tryptic_cterm(x) if isinstance(x, str) else False, + return_dtype=pl.Boolean, + ) + return df.with_columns( + pl.Series("proteome_hit", hits, dtype=pl.Boolean), + tryptic.alias("tryptic_cterm"), + ) + + +def _terminus_count_row( + sub: pd.DataFrame, + *, + cohort: str, + fdr_threshold: float | None, +) -> dict: + """Return one row of tryptic / non-tryptic counts for *sub*.""" + n = len(sub) + n_tryp = int(sub["tryptic_cterm"].sum()) if n > 0 else 0 + n_non = n - n_tryp + return { + "cohort": cohort, + "fdr_threshold": fdr_threshold, + "n": n, + "n_tryptic": n_tryp, + "n_non_tryptic": n_non, + "pct_tryptic": round(n_tryp / n * 100, 2) if n > 0 else 0.0, + "pct_non_tryptic": round(n_non / n * 100, 2) if n > 0 else 0.0, + } + + +def _nontryptic_terminus_proportions_table(df: pd.DataFrame) -> pd.DataFrame: + """Tryptic versus non-tryptic counts across cohorts and FDR cutoffs.""" + df = _add_q_values(df) + rows: list[dict] = [ + _terminus_count_row(df, cohort="all_predictions", fdr_threshold=None), + _terminus_count_row( + df[df["proteome_hit"]], + cohort=COHORT_FULL_SEARCH, + fdr_threshold=None, + ), + ] + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t] + rows.append( + _terminus_count_row( + retained, + cohort="retained_at_fdr", + fdr_threshold=fdr_t, + ) + ) + rows.append( + _terminus_count_row( + retained[retained["proteome_hit"]], + cohort=COHORT_FULL_SEARCH_AT_FDR, + fdr_threshold=fdr_t, + ) + ) + return pd.DataFrame(rows) + + +def _nontryptic_calibration_delta_rows( + sub: pd.DataFrame, + *, + cohort: str, + fdr_threshold: float | None, +) -> list[dict]: + """Build calibration-shift summary rows for tryptic and non-tryptic groups.""" + out: list[dict] = [] + for tryptic, label in ((True, "tryptic"), (False, "non_tryptic")): + grp = sub[sub["tryptic_cterm"] == tryptic] + n = len(grp) + out.append( + { + "cohort": cohort, + "fdr_threshold": fdr_threshold, + "terminus_group": label, + "n": n, + "mean_confidence": ( + round(float(grp["confidence"].mean()), 4) if n > 0 else float("nan") + ), + "mean_calibrated_confidence": ( + round(float(grp["calibrated_confidence"].mean()), 4) + if n > 0 + else float("nan") + ), + "mean_delta_confidence": ( + round(float(grp["delta_confidence"].mean()), 4) + if n > 0 + else float("nan") + ), + "median_delta_confidence": ( + round(float(grp["delta_confidence"].median()), 4) + if n > 0 + else float("nan") + ), + } + ) + return out + + +def _nontryptic_calibration_delta_table(df: pd.DataFrame) -> pd.DataFrame: + """Mean calibration shift by terminus and cohort.""" + if "confidence" not in df.columns: + return pd.DataFrame() + + work = _add_q_values(df.copy()) + work["delta_confidence"] = work["calibrated_confidence"] - work["confidence"] + work = work.dropna( + subset=["confidence", "calibrated_confidence", "delta_confidence"] + ) + + rows: list[dict] = [] + rows.extend( + _nontryptic_calibration_delta_rows( + work, cohort="all_predictions", fdr_threshold=None + ) + ) + rows.extend( + _nontryptic_calibration_delta_rows( + work[work["proteome_hit"]], + cohort=COHORT_FULL_SEARCH, + fdr_threshold=None, + ) + ) + for fdr_t in FDR_THRESHOLDS: + retained = work[(work["psm_q_value"] <= fdr_t) & work["proteome_hit"]] + rows.extend( + _nontryptic_calibration_delta_rows( + retained, + cohort=COHORT_FULL_SEARCH_AT_FDR, + fdr_threshold=fdr_t, + ) + ) + return pd.DataFrame(rows) + + +def _nontryptic_summary_table(df: pd.DataFrame) -> pd.DataFrame: + """Build the tryptic-summary table at each FDR threshold.""" + df = _add_q_values(df) + rows: list[dict] = [] + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t] + n_retained = len(retained) + hits = retained[retained["proteome_hit"]] + n_hit = len(hits) + tryptic_hits = hits[hits["tryptic_cterm"]] + non_tryptic_hits = hits[~hits["tryptic_cterm"]] + n_tryp = len(tryptic_hits) + n_non = len(non_tryptic_hits) + rows.append( + { + "fdr_threshold": fdr_t, + "n_retained": n_retained, + "n_full_search_space": n_hit, + "n_tryptic_hit": n_tryp, + "n_non_tryptic_hit": n_non, + "pct_non_tryptic_among_hits": ( + round(n_non / n_hit * 100, 2) if n_hit > 0 else 0.0 + ), + "mean_cal_conf_tryptic": ( + round(float(tryptic_hits["calibrated_confidence"].mean()), 4) + if n_tryp > 0 + else float("nan") + ), + "mean_cal_conf_non_tryptic": ( + round(float(non_tryptic_hits["calibrated_confidence"].mean()), 4) + if n_non > 0 + else float("nan") + ), + } + ) + return pd.DataFrame(rows) + + +def _nontryptic_score_label(score_col: str) -> str: + if score_col == "confidence": + return "Raw InstaNovo confidence" + return "Calibrated confidence" + + +def _finalize_nontryptic_violin_figure( + fig: plt.Figure, + axes: list[plt.Axes], + *, + y_label: str, + suptitle: str, +) -> None: + """Apply shared y-axis label and suptitle after multi-panel violin plots.""" + if len(axes) > 0: + axes[0].set_ylabel(y_label) + fig.tight_layout(rect=[0, 0, 1, 0.92]) + fig.suptitle(suptitle, fontsize=13, y=0.98) + + +def _plot_nontryptic_score_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + score_col: str, + save_name: str, + suptitle: str, +) -> None: + """Violin plot of a score column split by C-terminal residue at each FDR.""" + if score_col not in df.columns: + print(f" skipping {save_name} (missing {score_col})") + return + + df = _add_q_values(df) + y_label = _nontryptic_score_label(score_col) + palette = {"Tryptic (K/R)": _MAIN_LINE_COLOUR, "Non-tryptic": _NOVEL_COLOUR} + + panels: list[tuple[float, pd.DataFrame]] = [] + for fdr_t in FDR_THRESHOLDS: + retained = df[(df["psm_q_value"] <= fdr_t) & df["proteome_hit"]] + retained = retained.dropna(subset=[score_col]) + if len(retained) < 5: + print( + f" skipping {int(fdr_t * 100)}% FDR panel in {save_name} " + f"(n={len(retained):,})" + ) + continue + retained = retained.copy() + retained["C-terminus"] = retained["tryptic_cterm"].map( + {True: "Tryptic (K/R)", False: "Non-tryptic"}, + ) + panels.append((fdr_t, retained)) + + if not panels: + print( + f" skipping {save_name} (no FDR panels with enough " + f"{FULL_SEARCH_SPACE_LABEL} PSMs)" + ) + return + + n_cols = len(panels) + fig, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 5), sharey=True) + if n_cols == 1: + axes = [axes] + + for ax, (fdr_t, retained) in zip(axes, panels): + sns.violinplot( + data=retained, + x="C-terminus", + y=score_col, + order=list(_C_TERMINUS_ORDER), + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + pct = int(fdr_t * 100) + ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) + ax.grid(False) + _spine_fmt(ax) + + _finalize_nontryptic_violin_figure(fig, axes, y_label=y_label, suptitle=suptitle) + _save(fig, out_dir, save_name) + + +def _plot_nontryptic_conf_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Violin plot of calibrated confidence split by C-terminal residue.""" + _plot_nontryptic_score_by_terminus( + df, + out_dir, + score_col="calibrated_confidence", + save_name=f"{prefix}_conf_by_terminus", + suptitle=( + f"Calibrated confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" + "by C-terminal residue" + ), + ) + + +def _plot_nontryptic_raw_conf_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Violin plot of raw InstaNovo confidence split by C-terminal residue.""" + _plot_nontryptic_score_by_terminus( + df, + out_dir, + score_col="confidence", + save_name=f"{prefix}_raw_conf_by_terminus", + suptitle=( + f"Raw InstaNovo confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" + "by C-terminal residue" + ), + ) + + +def _plot_nontryptic_overlapping_score_histogram( + tryp: np.ndarray, + non_tryp: np.ndarray, + *, + score_col: str, + title: str, + out_dir: Path, + save_name: str, +) -> None: + """Overlapping tryptic / non-tryptic histogram with KDE overlays.""" + if len(tryp) + len(non_tryp) < 2: + print(f" skipping {save_name} (too few PSMs)") + return + + fig, ax = plt.subplots(figsize=(7, 5)) + bins = 50 + + ax.hist( + non_tryp, + bins=bins, + alpha=_NONTRYPTIC_GROUP_ALPHA, + label=f"Non-tryptic (n={len(non_tryp):,})", + density=False, + edgecolor="black", + color=_NOVEL_COLOUR, + ) + ax.hist( + tryp, + bins=bins, + alpha=_TRYPTIC_GROUP_ALPHA, + label=f"Tryptic (K/R, n={len(tryp):,})", + density=False, + edgecolor="black", + color=_MAIN_LINE_COLOUR, + ) + + all_vals = np.concatenate([tryp, non_tryp]) if len(non_tryp) else tryp + x_min, x_max = float(all_vals.min()), float(all_vals.max()) + if x_max <= x_min: + x_max = x_min + 1e-6 + x_grid = np.linspace(x_min, x_max, 300) + bin_width = (x_max - x_min) / bins if bins > 1 else 1.0 + + if len(non_tryp) > 1: + y_non = gaussian_kde(non_tryp)(x_grid) * len(non_tryp) * bin_width + ax.plot(x_grid, y_non, color=_NOVEL_COLOUR, lw=1.5) + if len(tryp) > 1: + y_tryp = gaussian_kde(tryp)(x_grid) * len(tryp) * bin_width + ax.plot(x_grid, y_tryp, color=_MAIN_LINE_COLOUR, lw=1.5) + + ax.set_xlabel(_nontryptic_score_label(score_col)) + ax.set_ylabel("Frequency") + ax.set_title(title) + ax.legend(loc="upper center") + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, save_name) + + +def _plot_nontryptic_score_histograms( + df: pd.DataFrame, + out_dir: Path, + *, + score_col: str, + save_name: str, + title: str, + subset: pd.DataFrame | None = None, + min_psms: int = 10, +) -> None: + """Overlapping tryptic / non-tryptic histograms for a score column.""" + if score_col not in df.columns: + print(f" skipping {save_name} (missing {score_col})") + return + + retained = df if subset is None else subset + retained = retained.dropna(subset=[score_col]) + if len(retained) < min_psms: + print(f" skipping {save_name} (too few PSMs)") + return + + tryp = retained.loc[retained["tryptic_cterm"], score_col].to_numpy() + non_tryp = retained.loc[~retained["tryptic_cterm"], score_col].to_numpy() + _plot_nontryptic_overlapping_score_histogram( + tryp, + non_tryp, + score_col=score_col, + title=title, + out_dir=out_dir, + save_name=save_name, + ) + + +def _plot_nontryptic_score_histograms_at_fdr( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Calibrated confidence histogram at 5% FDR for full search space PSMs.""" + df = _add_q_values(df) + retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="calibrated_confidence", + save_name=f"{prefix}_score_histograms", + title=( + f"Calibrated confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} " + "at 5% FDR" + ), + subset=retained, + ) + + +def _plot_nontryptic_raw_score_histograms_at_fdr( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Raw InstaNovo confidence histogram at 5% FDR for full search space PSMs.""" + df = _add_q_values(df) + retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="confidence", + save_name=f"{prefix}_raw_score_histograms", + title=( + f"Raw InstaNovo confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} " + "at 5% FDR" + ), + subset=retained, + ) + + +def _plot_nontryptic_full_score_histograms( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Full-dataset raw and calibrated histograms by C-terminal residue.""" + all_preds = df.dropna(subset=["calibrated_confidence"]) + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="calibrated_confidence", + save_name=f"{prefix}_calibrated_score_histogram_full", + title=( + f"Calibrated confidence for all {dataset_label} predictions\n" + "by C-terminal residue" + ), + subset=all_preds, + min_psms=2, + ) + if "confidence" not in df.columns: + print(f" skipping {prefix}_raw_score_histogram_full (missing confidence)") + return + raw_preds = df.dropna(subset=["confidence"]) + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="confidence", + save_name=f"{prefix}_raw_score_histogram_full", + title=( + f"Raw InstaNovo confidence for all {dataset_label} predictions\n" + "by C-terminal residue" + ), + subset=raw_preds, + min_psms=2, + ) + + has_raw = "confidence" in df.columns + panels: list[tuple[str, str, pd.DataFrame]] = [ + ( + "calibrated_confidence", + "Calibrated confidence", + all_preds, + ), + ] + if has_raw: + panels.append(("confidence", "Raw InstaNovo confidence", raw_preds)) + + n_cols = len(panels) + fig, axes = plt.subplots(1, n_cols, figsize=(7 * n_cols, 5), sharey=True) + if n_cols == 1: + axes = [axes] + + bins = 50 + for ax, (score_col, y_label, subset) in zip(axes, panels): + tryp = subset.loc[subset["tryptic_cterm"], score_col].to_numpy() + non_tryp = subset.loc[~subset["tryptic_cterm"], score_col].to_numpy() + ax.hist( + non_tryp, + bins=bins, + alpha=_NONTRYPTIC_GROUP_ALPHA, + label=f"Non-tryptic (n={len(non_tryp):,})", + density=False, + edgecolor="black", + color=_NOVEL_COLOUR, + ) + ax.hist( + tryp, + bins=bins, + alpha=_TRYPTIC_GROUP_ALPHA, + label=f"Tryptic (K/R, n={len(tryp):,})", + density=False, + edgecolor="black", + color=_MAIN_LINE_COLOUR, + ) + ax.set_xlabel(y_label) + ax.set_ylabel("Frequency") + ax.set_title(f"All predictions (n={len(subset):,})") + ax.legend(loc="upper center", fontsize=9) + ax.grid(False) + _spine_fmt(ax) + + fig.tight_layout(rect=[0, 0, 1, 0.92]) + fig.suptitle( + f"Score distributions for all {dataset_label} predictions by C-terminal residue", + fontsize=13, + y=0.98, + ) + _save(fig, out_dir, f"{prefix}_score_histogram_full_panel") + + +def _plot_nontryptic_calibration_scatter( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Subsampled scatter of raw versus calibrated confidence by C-terminus.""" + if "confidence" not in df.columns: + print(f" skipping {prefix}_calibration_scatter (missing confidence)") + return + + work = df.dropna(subset=["confidence", "calibrated_confidence"]) + n_total = len(work) + if n_total < 10: + print(f" skipping {prefix}_calibration_scatter (too few PSMs)") + return + + plot_df = _subsample_psms(work, _NONTRYPTIC_CALIBRATION_SCATTER_MAX_POINTS) + n_show = len(plot_df) + fig, ax = plt.subplots(figsize=(7.5, 7)) + panels = [ + ("Non-tryptic", _NOVEL_COLOUR, False, _NONTRYPTIC_GROUP_ALPHA), + ("Tryptic (K/R)", _MAIN_LINE_COLOUR, True, _TRYPTIC_GROUP_ALPHA), + ] + for label, colour, tryptic, alpha in panels: + sub = plot_df.loc[plot_df["tryptic_cterm"] == tryptic] + if len(sub) == 0: + continue + ax.scatter( + sub["confidence"], + sub["calibrated_confidence"], + c=colour, + s=12, + alpha=alpha, + rasterized=True, + label=f"{label}", + ) + + ax.plot( + [-0.01, 1.01], + [-0.01, 1.01], + ls="--", + color="black", + lw=1, + label="No recalibration", + zorder=5, + ) + ax.set_xlim(-0.01, 1.01) + ax.set_ylim(-0.01, 1.01) + ax.set_xlabel("Raw InstaNovo confidence") + ax.set_ylabel("Calibrated confidence") + if n_show < n_total: + ax.set_title( + f"Raw vs calibrated confidence for all {dataset_label} predictions" + ) + else: + ax.set_title(f"All {dataset_label} predictions") + ax.legend(loc="lower right", fontsize=9) + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, f"{prefix}_calibration_scatter") + + +def _plot_nontryptic_delta_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Calibration shift by C-terminal residue.""" + if "confidence" not in df.columns: + print(f" skipping {prefix}_delta_by_terminus (missing confidence)") + return + + df = _add_q_values(df) + y_label = "Calibration shift" + palette = {"Tryptic (K/R)": _MAIN_LINE_COLOUR, "Non-tryptic": _NOVEL_COLOUR} + work = df.copy() + work["delta_confidence"] = work["calibrated_confidence"] - work["confidence"] + + panels: list[tuple[float, pd.DataFrame]] = [] + for fdr_t in FDR_THRESHOLDS: + retained = work[(work["psm_q_value"] <= fdr_t) & work["proteome_hit"]] + retained = retained.dropna(subset=["delta_confidence"]) + if len(retained) < 5: + print( + f" skipping {int(fdr_t * 100)}% FDR panel in " + f"{prefix}_delta_by_terminus (n={len(retained):,})" + ) + continue + retained = retained.copy() + retained["C-terminus"] = retained["tryptic_cterm"].map( + {True: "Tryptic (K/R)", False: "Non-tryptic"}, + ) + panels.append((fdr_t, retained)) + + if not panels: + print(f" skipping {prefix}_delta_by_terminus (no FDR panels with enough PSMs)") + return + + n_cols = len(panels) + fig, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 5), sharey=True) + if n_cols == 1: + axes = [axes] + + for ax, (fdr_t, retained) in zip(axes, panels): + sns.violinplot( + data=retained, + x="C-terminus", + y="delta_confidence", + order=list(_C_TERMINUS_ORDER), + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.axhline(0.0, ls="--", color=_IDEAL_LINE_COLOUR, lw=1) + ax.set_xlabel("") + pct = int(fdr_t * 100) + ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) + ax.grid(False) + _spine_fmt(ax) + + _finalize_nontryptic_violin_figure( + fig, + axes, + y_label=y_label, + suptitle=( + f"Winnow calibration shift for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" + "by C-terminal residue" + ), + ) + _save(fig, out_dir, f"{prefix}_delta_by_terminus") + + +def _pooled_feature_mean_std(retained: pd.DataFrame, col: str) -> tuple[float, float]: + vals = retained[col].dropna() + if len(vals) == 0: + return float("nan"), float("nan") + if len(vals) == 1: + return float(vals.iloc[0]), float("nan") + return float(vals.mean()), float(vals.std()) + + +def _feature_group_median_z_row( + sub: pd.DataFrame, + available: list[str], + pooled: dict[str, tuple[float, float]], +) -> dict[str, float]: + row: dict[str, float] = {} + for col in available: + vals = sub[col].dropna() + if len(vals) == 0: + row[f"median_{col}"] = float("nan") + row[f"z_median_{col}"] = float("nan") + continue + med = float(vals.median()) + row[f"median_{col}"] = round(med, 4) + mu, std = pooled[col] + if np.isnan(std) or std == 0: + row[f"z_median_{col}"] = float("nan") + else: + row[f"z_median_{col}"] = round((med - mu) / std, 4) + return row + + +def _feature_median_z_score_table( + retained: pd.DataFrame, + available: list[str], + groups: list[tuple[str, str, pd.Series]], + *, + group_col: str, + reference: pd.DataFrame | None = None, +) -> pd.DataFrame: + """Per-group feature medians and z-scores relative to *reference* or *retained* PSMs.""" + pool_from = reference if reference is not None else retained + pooled = {col: _pooled_feature_mean_std(pool_from, col) for col in available} + + rows: list[dict] = [] + for group_key, _label, mask in groups: + sub = retained[mask] + row: dict = {group_col: group_key, "n": len(sub)} + row.update(_feature_group_median_z_row(sub, available, pooled)) + rows.append(row) + return pd.DataFrame(rows) + + +def _nontryptic_feature_table(df: pd.DataFrame) -> pd.DataFrame: + """Median feature values for tryptic vs non-tryptic full search space PSMs at 5% FDR.""" + df = _add_q_values(df) + retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] + available = [c for c in FEATURE_COLUMNS if c in retained.columns] + if not available: + return pd.DataFrame() + + return _feature_median_z_score_table( + retained, + available, + [ + ("tryptic", "Tryptic (K/R)", retained["tryptic_cterm"]), + ("non_tryptic", "Non-tryptic", ~retained["tryptic_cterm"]), + ], + group_col="group", + ) + + +def _plot_grouped_feature_z_scores( + feat_df: pd.DataFrame, + *, + group_col: str, + out_dir: Path, + save_name: str, + title: str, + group_style: list[tuple[str, str, str]], + z_score_ylabel: str = "Median z-score", +) -> None: + """Grouped bar chart of pooled z-scored feature medians.""" + if feat_df.empty: + print(f" skipping {save_name} (no feature data)") + return + + z_cols = [c for c in feat_df.columns if c.startswith("z_median_")] + if not z_cols: + print(f" skipping {save_name} (no z-scored feature columns)") + return + + plot_df = feat_df.set_index(group_col) + feature_labels = [_nice_feature_label(c.replace("z_median_", "")) for c in z_cols] + x = np.arange(len(z_cols)) + + present = [ + (key, label, colour) + for key, label, colour in group_style + if key in plot_df.index + ] + n_groups = len(present) + total_width = 0.7 + bar_w = total_width / max(n_groups, 1) + + fig, ax = plt.subplots(figsize=(9, 6.5)) + ax.axhline(0.0, color=_IDEAL_LINE_COLOUR, lw=0.8, zorder=0) + bar_groups = [] + for plot_i, (group_key, label, colour) in enumerate(present): + offset = (plot_i - (n_groups - 1) / 2) * bar_w + vals = plot_df.loc[group_key, z_cols].to_numpy(dtype=float) + bars = ax.bar( + x + offset, + vals, + bar_w, + label=label, + color=colour, + edgecolor="black", + linewidth=1, + ) + bar_groups.append(bars) + + if not bar_groups: + print(f" skipping {save_name} (no groups to plot)") + plt.close(fig) + return + + ax.set_xticks(x) + ax.set_xticklabels(feature_labels, rotation=30, ha="right") + ax.set_ylabel(z_score_ylabel) + ax.set_title(title) + ax.legend(loc="upper right", fontsize=9) + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, save_name) + + +def _plot_nontryptic_feature_comparison( + feat_df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Grouped bar chart of median features, tryptic vs non-tryptic.""" + _plot_grouped_feature_z_scores( + feat_df, + group_col="group", + out_dir=out_dir, + save_name=f"{prefix}_feature_comparison", + title=( + f"Median feature values for {dataset_label} tryptic versus " + f"non-tryptic {FULL_SEARCH_SPACE_LABEL} PSMs at 5% FDR" + ), + group_style=[ + ("tryptic", "Tryptic (K/R)", _MAIN_LINE_COLOUR), + ("non_tryptic", "Non-tryptic", _NOVEL_COLOUR), + ], + ) + + +def _nontryptic_digest_analysis( + predictions_dir: Path, + fasta: Path, + output_dir: Path, + *, + file_prefix: str, + dataset_label: str, +) -> None: + """Shared tryptic vs non-tryptic digest analysis (any non-tryptic enzyme dataset).""" + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"Loading predictions from {predictions_dir}") + df_pl = _load_data(predictions_dir) + print(f" {df_pl.height:,} rows loaded") + + print(f"Annotating {FULL_SEARCH_SPACE_LABEL} predictions against {fasta}") + df_pl = _nontryptic_annotate(df_pl, fasta) + n_hits = df_pl.filter(pl.col("proteome_hit")).height + print(f" {n_hits:,} PSMs in {FULL_SEARCH_SPACE_LABEL} (proteome substring match)") + + df = df_pl.to_pandas() + + print("Building tryptic summary table") + summary = _nontryptic_summary_table(df) + summary.to_csv(output_dir / f"{file_prefix}_tryptic_summary.csv", index=False) + print(summary.to_string(index=False)) + + print("Building terminus proportion table") + prop_df = _nontryptic_terminus_proportions_table(df) + prop_df.to_csv(output_dir / f"{file_prefix}_terminus_proportions.csv", index=False) + print(prop_df.to_string(index=False)) + + if "confidence" in df.columns: + print("Building calibration shift table") + delta_df = _nontryptic_calibration_delta_table(df) + delta_df.to_csv( + output_dir / f"{file_prefix}_calibration_delta_summary.csv", index=False + ) + print(delta_df.to_string(index=False)) + else: + print(" skipping calibration shift table (missing raw confidence)") + + print("Building feature comparison table") + feat_df = _nontryptic_feature_table(df) + if not feat_df.empty: + feat_df.to_csv( + output_dir / f"{file_prefix}_feature_comparison.csv", index=False + ) + + plot_kw = {"prefix": file_prefix, "dataset_label": dataset_label} + print("Plotting") + _plot_nontryptic_conf_by_terminus(df, output_dir, **plot_kw) + _plot_nontryptic_raw_conf_by_terminus(df, output_dir, **plot_kw) + _plot_nontryptic_score_histograms_at_fdr(df, output_dir, **plot_kw) + _plot_nontryptic_raw_score_histograms_at_fdr(df, output_dir, **plot_kw) + _plot_nontryptic_full_score_histograms(df, output_dir, **plot_kw) + _plot_nontryptic_calibration_scatter(df, output_dir, **plot_kw) + _plot_nontryptic_delta_by_terminus(df, output_dir, **plot_kw) + _plot_nontryptic_feature_comparison(feat_df, output_dir, **plot_kw) + + print(f"\n{dataset_label} analysis complete. Output in {output_dir}") + + +@app.command() +def nontryptic_digest( + predictions_dir: Annotated[ + Path, + typer.Option( + "--predictions-dir", + help=( + "winnow predict output folder for the non-tryptic enzyme digest " + "(full search space / acfm)." + ), + ), + ], + fasta: Annotated[ + Path, + typer.Option("--fasta", help="Proteome FASTA for substring matching."), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for output tables and plots."), + ], + file_prefix: Annotated[ + str, + typer.Option( + "--file-prefix", + help="Prefix for output CSV and plot filenames (e.g. chymotrypsin).", + ), + ] = "nontryptic_digest", + dataset_label: Annotated[ + str, + typer.Option( + "--dataset-label", + help="Human-readable dataset name used in plot titles.", + ), + ] = "Non-tryptic digest", +) -> None: + """Analyse calibrator behaviour on non-tryptic enzyme digest peptides.""" + _nontryptic_digest_analysis( + predictions_dir, + fasta, + output_dir, + file_prefix=file_prefix, + dataset_label=dataset_label, + ) + + +@app.command() +def chymotrypsin( + predictions_dir: Annotated[ + Path, + typer.Option( + "--predictions-dir", + help="winnow predict output folder for HeLa chymotrypsin (full search space / acfm).", + ), + ], + fasta: Annotated[ + Path, + typer.Option("--fasta", help="Human proteome FASTA for substring matching."), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for output tables and plots."), + ], +) -> None: + """Convenience wrapper for ``nontryptic_digest`` on HeLa chymotrypsin data.""" + _nontryptic_digest_analysis( + predictions_dir, + fasta, + output_dir, + file_prefix="chymotrypsin", + dataset_label="HeLa chymotrypsin", + ) + + +@app.command() +def gluc( + predictions_dir: Annotated[ + Path, + typer.Option( + "--predictions-dir", + help="[Deprecated] winnow predict output folder for GluC raw.", + ), + ], + fasta: Annotated[ + Path, + typer.Option("--fasta", help="Human proteome FASTA for substring matching."), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for output tables and plots."), + ], +) -> None: + """[Deprecated] Use ``nontryptic_digest`` with ``--file-prefix gluc``.""" + typer.secho( + "Warning: gluc is deprecated; use nontryptic_digest --file-prefix gluc", + fg=typer.colors.YELLOW, + err=True, + ) + _nontryptic_digest_analysis( + predictions_dir, + fasta, + output_dir, + file_prefix="gluc", + dataset_label="HeLa degradome", + ) + + +# ── ProteomeTools-1 analysis ──────────────────────────────────────────── + + +def _exact_match_category(fits_precursor: bool) -> str: + if fits_precursor: + return "exact_match_and_fits_precursor" + return "exact_match_and_no_precursor_fit" + + +def _subsequence_category(fits_precursor: bool) -> str: + if fits_precursor: + return "subsequence_and_fits_precursor" + return "subsequence_and_no_precursor_fit" + + +def _neither_category(fits_precursor: bool) -> str: + if fits_precursor: + return "neither_and_fits_precursor" + return "neither_and_no_precursor_fit" + + +def _classify_exact_matches( + categories: list[str], + predictions: list[str], + fits_precursor: list[bool], + lcfm_peptide_set: set[str], +) -> None: + for i, (peptide, fit) in enumerate(zip(predictions, fits_precursor)): + if peptide in lcfm_peptide_set: + categories[i] = _exact_match_category(fit) + + +def _classify_subsequence_matches( + categories: list[str], + predictions: list[str], + fits_precursor: list[bool], + lcfm_haystack: str, +) -> None: + remaining_indices = [ + i for i, category in enumerate(categories) if category == "neither" + ] + remaining_peps = [predictions[i] for i in remaining_indices] + remaining_fits = [fits_precursor[i] for i in remaining_indices] + if not remaining_peps: + return + + hits = _batch_substring_hits(remaining_peps, lcfm_haystack) + for idx, fit, hit in zip(remaining_indices, remaining_fits, hits): + if hit: + categories[idx] = _subsequence_category(fit) + + +def _classify_remaining_neither( + categories: list[str], + fits_precursor: list[bool], +) -> None: + for i, fit in enumerate(fits_precursor): + if categories[i] == "neither": + categories[i] = _neither_category(fit) + + +def _classify_predictions( + predictions: list[str], + fits_precursor: list[bool], + lcfm_peptide_set: set[str], + lcfm_haystack: str, +) -> list[str]: + """Classify each prediction by lcfm overlap and precursor mass fit (<20 ppm). + + Uses Aho-Corasick to find which predictions are substrings of at least + one lcfm peptide (the haystack is built by joining all lcfm peptides + with a separator). Unmatched predictions are split by precursor fit. + """ + categories = ["neither"] * len(predictions) + _classify_exact_matches(categories, predictions, fits_precursor, lcfm_peptide_set) + _classify_subsequence_matches( + categories, predictions, fits_precursor, lcfm_haystack + ) + _classify_remaining_neither(categories, fits_precursor) + return categories + + +def _proteometools_summary_table(df: pd.DataFrame) -> pd.DataFrame: + """Build novelty summary table at each FDR threshold.""" + df = _add_q_values(df) + rows: list[dict] = [] + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t] + n = len(retained) + if n == 0: + rows.append( + { + "fdr_threshold": fdr_t, + "n_retained": 0, + "n_exact_match_and_no_precursor_fit": 0, + "n_exact_match_and_fits_precursor": 0, + "n_subsequence_and_no_precursor_fit": 0, + "n_subsequence_and_fits_precursor": 0, + "n_neither_and_no_precursor_fit": 0, + "n_neither_and_fits_precursor": 0, + "pct_exact_or_sub_and_fit_among_retained": 0.0, + } + ) + continue + + cats = retained["novelty_category"] + n_exact_and_no_fit = int((cats == "exact_match_and_no_precursor_fit").sum()) + n_exact_and_fits_precursor = int( + (cats == "exact_match_and_fits_precursor").sum() + ) + n_sub_and_no_fit = int((cats == "subsequence_and_no_precursor_fit").sum()) + n_sub_and_fits_precursor = int((cats == "subsequence_and_fits_precursor").sum()) + n_neither_and_no_fit = int((cats == "neither_and_no_precursor_fit").sum()) + n_neither_and_fits_precursor = int((cats == "neither_and_fits_precursor").sum()) + + rows.append( + { + "fdr_threshold": fdr_t, + "n_retained": n, + "n_exact_match_and_no_precursor_fit": n_exact_and_no_fit, + "n_exact_match_and_fits_precursor": n_exact_and_fits_precursor, + "n_subsequence_and_no_precursor_fit": n_sub_and_no_fit, + "n_subsequence_and_fits_precursor": n_sub_and_fits_precursor, + "n_neither_and_no_precursor_fit": n_neither_and_no_fit, + "n_neither_and_fits_precursor": n_neither_and_fits_precursor, + "pct_exact_or_sub_and_fit_among_retained": round( + (n_exact_and_fits_precursor + n_sub_and_fits_precursor) / n * 100, 2 + ) + if n > 0 + else 0.0, + } + ) + return pd.DataFrame(rows) + + +_PROTEOMETOOLS_CONF_PLOT_LABELS: dict[str, str] = { + "exact_match_and_no_precursor_fit": "ID-", + "exact_match_and_fits_precursor": "ID+", + "subsequence_and_no_precursor_fit": "Sub-", + "subsequence_and_fits_precursor": "Sub+", + "neither_and_no_precursor_fit": "Novel-", + "neither_and_fits_precursor": "Novel+", +} + +_PROTEOMETOOLS_CONF_CATEGORY_ORDER = list(_PROTEOMETOOLS_CONF_PLOT_LABELS.keys()) + + +def _proteometools_conf_category_legend(ax: plt.Axes) -> None: + handles = [ + Line2D([], [], color="none", label="ID: Exact sequence match to labelled set."), + Line2D([], [], color="none", label="Sub: Subsequence of labelled set peptide."), + Line2D([], [], color="none", label="Novel: No sequence match to labelled set."), + Line2D( + [], + [], + color="none", + label="+: Matches precursor mass within 20 ppm.", + ), + ] + ( + Line2D( + [], + [], + color="none", + label="-: Does not match precursor mass within 20 ppm.", + ), + ) + ax.legend( + handles=handles, + loc="upper left", + bbox_to_anchor=(1.02, 1.0), + borderaxespad=0, + frameon=True, + fontsize=9, + ) + + +def _plot_proteometools_conf_by_category( + df: pd.DataFrame, + out_dir: Path, +) -> None: + """Violin plot of calibrated confidence by novelty category (all unlabelled PSMs).""" + plot_df = df.dropna(subset=["calibrated_confidence"]).copy() + if len(plot_df) < 5: + print(" skipping proteometools_conf_by_category (too few PSMs)") + return + + palette = { + _PROTEOMETOOLS_CONF_PLOT_LABELS[k]: _PALETTE[i] + for i, k in enumerate(_PROTEOMETOOLS_CONF_CATEGORY_ORDER) + } + + plot_df["Category"] = plot_df["novelty_category"].map( + _PROTEOMETOOLS_CONF_PLOT_LABELS + ) + present_cats = [ + _PROTEOMETOOLS_CONF_PLOT_LABELS[c] + for c in _PROTEOMETOOLS_CONF_CATEGORY_ORDER + if _PROTEOMETOOLS_CONF_PLOT_LABELS[c] in plot_df["Category"].values + ] + if not present_cats: + print(" skipping proteometools_conf_by_category (no categories present)") + return + + conf = plot_df["calibrated_confidence"] + y_min, y_max = float(conf.min()), float(conf.max()) + + fig, ax = plt.subplots(figsize=(10, 5)) + sns.violinplot( + data=plot_df, + x="Category", + y="calibrated_confidence", + order=present_cats, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_ylim(y_min, y_max) + ax.set_xlabel("") + ax.set_ylabel("Calibrated confidence") + ax.set_title( + "Calibrated confidence for ProteomeTools-1 predictions\nby novelty category" + ) + ax.grid(False) + _spine_fmt(ax) + _proteometools_conf_category_legend(ax) + + fig.tight_layout() + _save(fig, out_dir, "proteometools_conf_by_category") + + +def _plot_proteometools_hit_rate( + df: pd.DataFrame, + out_dir: Path, +) -> None: + """Line plot: validated hit rate (exact or subsequence match and fits precursor mass within 20ppm) by calibrated confidence decile.""" + if len(df) < 20: + print(" skipping proteometools_hit_rate_vs_conf (too few PSMs)") + return + + df = df.copy() + df["is_validated"] = df["novelty_category"].isin( + ["exact_match_and_fits_precursor", "subsequence_and_fits_precursor"] + ) + df["conf_decile"] = pd.qcut( + df["calibrated_confidence"], + q=10, + duplicates="drop", + ) + grouped = ( + df.groupby("conf_decile", observed=True) + .agg( + hit_rate=("is_validated", "mean"), + mid=("calibrated_confidence", "mean"), + ) + .sort_values("mid") + ) + + fig, ax = plt.subplots(figsize=(8, 6)) + ax.plot( + grouped["mid"], + grouped["hit_rate"], + color=_MAIN_LINE_COLOUR, + linewidth=1.5, + marker="o", + markersize=6, + label="Validated hit rate", + ) + overall = float(df["is_validated"].mean()) + ax.axhline( + overall, + color=_IDEAL_LINE_COLOUR, + lw=1, + linestyle="--", + label=f"Overall mean ({overall:.2%})", + ) + ax.set_xlabel("Mean calibrated confidence per decile") + ax.set_ylabel( + "Fraction validated\n(exact match or subsequence fitting precursor mass)" + ) + ax.set_title( + "Validated hit rate by calibrated confidence decile for ProteomeTools-1" + ) + ax.legend(loc="lower right") + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, "proteometools_hit_rate_vs_conf") + + +def _proteometools_feature_table( + df: pd.DataFrame, + labelled_df: pd.DataFrame, +) -> pd.DataFrame: + """Median features for exact / novel / neither at 5% FDR.""" + df = _add_q_values(df) + labelled_df = _add_q_values(labelled_df) + retained = df[df["psm_q_value"] <= 0.05] + labelled_ref = labelled_df[labelled_df["psm_q_value"] <= 0.05] + available = [c for c in FEATURE_COLUMNS if c in retained.columns] + if not available: + return pd.DataFrame() + + return _feature_median_z_score_table( + retained, + available, + [ + ( + "exact_match_and_fits_precursor", + "Exact match, fits precursor mass", + retained["novelty_category"] == "exact_match_and_fits_precursor", + ), + ( + "exact_match_and_no_precursor_fit", + "Exact match, no precursor mass fit", + retained["novelty_category"] == "exact_match_and_no_precursor_fit", + ), + ( + "subsequence_and_fits_precursor", + "Subsequence, precursor mass fit", + retained["novelty_category"] == "subsequence_and_fits_precursor", + ), + ( + "subsequence_and_no_precursor_fit", + "Subsequence, no precursor mass fit", + retained["novelty_category"] == "subsequence_and_no_precursor_fit", + ), + ( + "neither_and_fits_precursor", + "Novel, precursor mass fit", + retained["novelty_category"] == "neither_and_fits_precursor", + ), + ( + "neither_and_no_precursor_fit", + "Novel, no precursor mass fit", + retained["novelty_category"] == "neither_and_no_precursor_fit", + ), + ], + group_col="category", + reference=labelled_ref, + ) + + +def _plot_proteometools_feature_comparison( + feat_df: pd.DataFrame, + out_dir: Path, +) -> None: + """Grouped bar chart of median features by category.""" + _plot_grouped_feature_z_scores( + feat_df, + group_col="category", + out_dir=out_dir, + save_name="proteometools_feature_comparison", + title=( + "Median feature values for ProteomeTools-1 predictions by novelty " + "category at 5% FDR" + ), + group_style=[ + ( + "exact_match_and_fits_precursor", + "Exact match, precursor mass fit", + _PALETTE[1], + ), + ( + "exact_match_and_no_precursor_fit", + "Exact match, no precursor mass fit", + _PALETTE[0], + ), + ( + "subsequence_and_fits_precursor", + "Subsequence, precursor mass fit", + _PALETTE[3], + ), + ( + "subsequence_and_no_precursor_fit", + "Subsequence, no precursor mass fit", + _PALETTE[2], + ), + ("neither_and_fits_precursor", "Novel, precursor mass fit", _PALETTE[5]), + ( + "neither_and_no_precursor_fit", + "Novel, no precursor mass fit", + _PALETTE[4], + ), + ], + z_score_ylabel="Median z-score (vs labelled PSMs at 5% FDR)", + ) + + +@app.command() +def proteometools( + lcfm_predictions_dir: Annotated[ + Path, + typer.Option( + "--lcfm-predictions-dir", + help="winnow predict output for PXD004732 lcfm (labelled).", + ), + ], + acfm_predictions_dir: Annotated[ + Path, + typer.Option( + "--acfm-predictions-dir", + help="winnow predict output for PXD004732 acfm (unlabelled).", + ), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for output tables and plots."), + ], +) -> None: + """Analyse calibrator behaviour on ProteomeTools-1 novel identifications.""" + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"Loading lcfm predictions from {lcfm_predictions_dir}") + lcfm_pl = _load_data(lcfm_predictions_dir) + print(f" {lcfm_pl.height:,} lcfm rows") + + print(f"Loading acfm predictions from {acfm_predictions_dir}") + acfm_pl = _load_data(acfm_predictions_dir) + print(f" {acfm_pl.height:,} acfm rows") + + print("Building lcfm peptide set for subsequence matching") + lcfm_sequences = lcfm_pl["sequence"].drop_nulls().to_list() + lcfm_peptide_set: set[str] = set() + for seq in lcfm_sequences: + stripped = _strip_mods(seq) + if stripped: + lcfm_peptide_set.add(stripped) + print(f" {len(lcfm_peptide_set):,} unique lcfm peptides") + + lcfm_haystack = _PROTEOME_JOIN_SEP.join(sorted(lcfm_peptide_set)) + + print("Classifying unlabelled predictions") + unlabelled_pl = acfm_pl.join( + lcfm_pl.select("spectrum_id"), on="spectrum_id", how="anti" + ) + unlabelled_pl = unlabelled_pl.with_columns( + (pl.col("delta_mass_ppm").abs() < 20).alias("fits_precursor") + ) + unlabelled_preds_raw = unlabelled_pl["prediction"].to_list() + unlabelled_preds_stripped = [ + _strip_mods(p) if isinstance(p, str) else "" for p in unlabelled_preds_raw + ] + fits_precursor = unlabelled_pl["fits_precursor"].to_list() + + categories = _classify_predictions( + unlabelled_preds_stripped, + fits_precursor, + lcfm_peptide_set, + lcfm_haystack, + ) + unlabelled_pl = unlabelled_pl.with_columns( + pl.Series("novelty_category", categories, dtype=pl.Utf8), + ) + df = unlabelled_pl.to_pandas() + labelled_df = lcfm_pl.to_pandas() + + print("Building novelty summary table") + summary = _proteometools_summary_table(df) + summary.to_csv(output_dir / "proteometools_novelty_summary.csv", index=False) + print(summary.to_string(index=False)) + + print("Building feature comparison table") + feat_df = _proteometools_feature_table(df, labelled_df) + if not feat_df.empty: + feat_df.to_csv(output_dir / "proteometools_feature_comparison.csv", index=False) + + print("Plotting") + _plot_proteometools_conf_by_category(df, output_dir) + _plot_proteometools_hit_rate(df, output_dir) + _plot_proteometools_feature_comparison(feat_df, output_dir) + + print(f"\nProteomeTools-1 analysis complete. Output in {output_dir}") + + +# ── Summary figure ──────────────────────────────────────────────────── + + +@app.command() +def summary( + proteometools_dir: Annotated[ + Path, + typer.Option( + "--proteometools-dir", + help="Output directory from the proteometools subcommand.", + ), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for the combined summary figure."), + ], + nontryptic_digest_dir: Annotated[ + Path | None, + typer.Option( + "--nontryptic-digest-dir", + help="Output directory from the nontryptic_digest subcommand.", + ), + ] = None, + chymotrypsin_dir: Annotated[ + Path | None, + typer.Option( + "--chymotrypsin-dir", + help="Legacy alias for --nontryptic-digest-dir (chymotrypsin outputs).", + ), + ] = None, + gluc_dir: Annotated[ + Path | None, + typer.Option( + "--gluc-dir", + help="Legacy alias for --nontryptic-digest-dir (gluc outputs).", + ), + ] = None, +) -> None: + """Produce a combined summary bar chart from both analyses.""" + output_dir.mkdir(parents=True, exist_ok=True) + + digest_dir = nontryptic_digest_dir or chymotrypsin_dir or gluc_dir + if digest_dir is None: + raise typer.BadParameter( + "Provide --nontryptic-digest-dir (or --chymotrypsin-dir / --gluc-dir)." + ) + + digest_csvs = sorted(digest_dir.glob("*_tryptic_summary.csv")) + if len(digest_csvs) != 1: + raise typer.BadParameter( + f"Expected exactly one *_tryptic_summary.csv under {digest_dir}, " + f"found {len(digest_csvs)}" + ) + digest_csv = digest_csvs[0] + pt_csv = proteometools_dir / "proteometools_novelty_summary.csv" + if not digest_csv.is_file(): + raise typer.BadParameter(f"Missing tryptic summary under {digest_dir}") + if not pt_csv.is_file(): + raise typer.BadParameter(f"Missing {pt_csv}") + + +if __name__ == "__main__": + app() diff --git a/scripts/analyze_upscored_fps.py b/scripts/analyze_upscored_fps.py new file mode 100644 index 00000000..eaeef641 --- /dev/null +++ b/scripts/analyze_upscored_fps.py @@ -0,0 +1,705 @@ +#!/usr/bin/env python3 +"""Characterise up-scored false positives from Winnow calibration. + +For each labelled evaluation dataset (where both ``sequence`` and ``prediction`` +are available), this script quantifies the false positives that calibration +"rescues" into high-confidence regions and compares their feature profiles to +true positives. + +Inputs are ``winnow predict`` output folders, each containing +``preds_and_fdr_metrics.csv`` and ``metadata.csv``. +""" + +from __future__ import annotations + +import json +import logging +import re +from pathlib import Path +from typing import Annotated + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer +import yaml +from instanovo.utils.metrics import Metrics +from instanovo.utils.residues import ResidueSet +from rich.logging import RichHandler + +from winnow.fdr.nonparametric import NonParametricFDRControl + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] + +TP_COLOR = _CORRECT_COLOUR +FP_COLOR = _INCORRECT_COLOUR + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)-?") +_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] +_FEATURE_VIOLIN_FDR_THRESHOLDS = [0.05, 0.10] +_MIN_FEATURE_VIOLIN_PSMs = 20 + +# Max PSMs per correctness panel in the raw-vs-calibrated confidence scatter. +_CONFIDENCE_SCATTER_MAX_POINTS = 10_000 +_CONFIDENCE_SCATTER_RANDOM_STATE = 42 + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD004732": "ProteomeTools-1", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", +} + +_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") + +# new_eval_sets_results layout: lcfm/PXD004452//preds_and_fdr_metrics.csv +_PXD_ACCESSION_PREFIX = "PXD" + +FEATURE_COLUMNS_OF_INTEREST = [ + "spectral_angle", + "xcorr", + "ion_matches", + "ion_match_intensity", + "irt_error", + "mass_error_ppm", + "margin", + "entropy", + "confidence", +] + +_MASS_ERROR_COLUMNS = ("mass_error_da", "mass_error_ppm") + +_NICE_LABELS: dict[str, str] = { + "ion_matches": "Ion match rate", + "ion_match_intensity": "Ion match intensity", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Max ion gap", + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation (XCorr)", + "mass_error_ppm": "Precursor mass error (ppm)", + "mass_error_da": "Precursor mass error (Da)", + "irt_error": "iRT prediction error", + "confidence": "Model confidence", + "margin": "Beam margin", + "median_margin": "Beam median margin", + "entropy": "Beam entropy", + "z-score": "Beam z-score", + "edit_distance": "Runner-up edit distance", + "min_token_probability": "Min. token probability", + "std_token_probability": "Std. token probability", +} + + +def _nice_label(col: str) -> str: + return _NICE_LABELS.get(col, col.replace("_", " ").capitalize()) + + +def _mass_error_column(df: pd.DataFrame, *, min_count: int = 10) -> str | None: + """Return ``mass_error_da`` or ``mass_error_ppm`` when present with enough data.""" + for col in _MASS_ERROR_COLUMNS: + if col in df.columns and df[col].notna().sum() > min_count: + return col + return None + + +def _violin_feature_columns(df: pd.DataFrame) -> list[str]: + """Feature list for violin plots, resolving Da vs ppm mass error.""" + cols: list[str] = [] + for col in FEATURE_COLUMNS_OF_INTEREST: + if col == "mass_error_ppm": + mass_col = _mass_error_column(df) + if mass_col is not None: + cols.append(mass_col) + else: + cols.append(col) + return cols + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _get_residue_masses() -> dict[str, float]: + config_path = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" + with open(config_path) as f: + cfg = yaml.safe_load(f) + return cfg["residue_masses"] + + +def _save_fig(fig: plt.Figure, base_path: Path, fmt: str = "both") -> None: + if fmt in ("pdf", "both"): + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + if fmt in ("png", "both"): + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _folder_display_name(folder_name: str) -> str: + """Map an evaluation folder name to a publication-ready dataset label.""" + key = folder_name + for suffix in _FOLDER_SUFFIXES: + if key.endswith(suffix): + key = key[: -len(suffix)] + break + return DATASET_DISPLAY_NAMES.get(key, key) + + +def _subsample_psms( + df: pd.DataFrame, + max_points: int, + random_state: int = _CONFIDENCE_SCATTER_RANDOM_STATE, +) -> pd.DataFrame: + """Return up to ``max_points`` rows without replacement.""" + if len(df) <= max_points: + return df + return df.sample(n=max_points, random_state=random_state) + + +def _strip_mods(seq: str) -> str: + if not seq or not isinstance(seq, str): + return "" + s = _MOD_PLUS.sub("", seq) + s = _MOD_UNIMOD.sub("", s) + return s.replace("I", "L") + + +def _project_key_from_folder(folder_name: str) -> str: + """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" + for suffix in _FOLDER_SUFFIXES: + if folder_name.endswith(suffix): + return folder_name[: -len(suffix)] + return folder_name + + +def _is_labelled_preds_folder(folder: Path) -> bool: + preds_csv = folder / "preds_and_fdr_metrics.csv" + if not preds_csv.is_file(): + return False + header = pd.read_csv(preds_csv, nrows=0).columns.tolist() + required = {"sequence", "prediction", "calibrated_confidence"} + return required.issubset(header) + + +def _discover_labelled_folders(root: Path) -> dict[str, Path]: + """Find folders with labelled ``preds_and_fdr_metrics.csv``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts used by new eval sets (``{root}/PXD004452//``). + """ + results: dict[str, Path] = {} + if not root.is_dir(): + return results + + def _register(key: str, folder: Path) -> None: + if key in results: + logger.warning( + "Duplicate labelled project key %r: %s and %s", + key, + results[key], + folder, + ) + return + results[key] = folder + + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if _is_labelled_preds_folder(child): + _register(_project_key_from_folder(child.name), child) + continue + if not child.name.startswith(_PXD_ACCESSION_PREFIX): + continue + for run_dir in sorted(child.iterdir()): + if run_dir.is_dir() and _is_labelled_preds_folder(run_dir): + _register(run_dir.name, run_dir) + return results + + +def _load_dataset(folder: Path) -> pd.DataFrame: + """Load and merge preds + metadata CSVs for a single evaluation folder.""" + preds = pd.read_csv(folder / "preds_and_fdr_metrics.csv") + meta_path = folder / "metadata.csv" + if meta_path.is_file(): + meta = pd.read_csv(meta_path) + join_cols = ["spectrum_id"] + [ + c for c in meta.columns if c != "spectrum_id" and c not in preds.columns + ] + if len(join_cols) > 1: + preds = preds.merge( + meta[join_cols].drop_duplicates(subset=["spectrum_id"]), + on="spectrum_id", + how="left", + ) + if "correct" not in preds.columns and {"sequence", "prediction"}.issubset( + preds.columns + ): + preds = preds.copy() + preds["correct"] = preds["sequence"] == preds["prediction"] + return preds + + +def _add_q_values( + df: pd.DataFrame, conf_col: str = "calibrated_confidence" +) -> pd.DataFrame: + """Fit non-parametric FDR and append ``psm_q_value`` if missing.""" + if "psm_q_value" in df.columns: + return df + fdr = NonParametricFDRControl() + fdr.fit(dataset=df[conf_col]) + df = fdr.add_psm_q_value(df, confidence_col=conf_col) + return df + + +# --------------------------------------------------------------------------- +# Analysis +# --------------------------------------------------------------------------- +def _upscored_summary_table( + df: pd.DataFrame, + delta_threshold: float, + dataset_name: str, +) -> pd.DataFrame: + """Build per-FDR-threshold summary of up-scored TP / FP counts.""" + df = _add_q_values(df) + upscored = df["delta_confidence"] > delta_threshold + + rows = [] + for fdr_t in FDR_THRESHOLDS: + passing = df["psm_q_value"] <= fdr_t + for label, mask in [ + ("all", pd.Series(True, index=df.index)), + ("up-scored", upscored), + ("not up-scored", ~upscored), + ]: + sub = df[mask & passing] + n = len(sub) + n_correct = int(sub["correct"].sum()) if "correct" in sub.columns else 0 + n_incorrect = n - n_correct + rows.append( + { + "dataset": dataset_name, + "fdr_threshold": fdr_t, + "subset": label, + "n_passing": n, + "n_correct": n_correct, + "n_incorrect": n_incorrect, + "pct_correct": round(n_correct / n * 100, 2) if n > 0 else 0.0, + } + ) + return pd.DataFrame(rows) + + +def _plot_confidence_scatter( + df: pd.DataFrame, + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Subsampled scatter of raw vs calibrated confidence, colored by correctness.""" + display = _folder_display_name(dataset_name) + fig, ax = plt.subplots(figsize=(8, 6)) + + # Define panels as before + panels = [ + ("Correct", TP_COLOR, df["correct"].astype(bool)), + ("Incorrect", FP_COLOR, ~df["correct"].astype(bool)), + ] + + handles = [] + for label, colour, mask in panels: + sub = df.loc[mask, ["confidence", "calibrated_confidence"]].dropna() + n_total = len(sub) + if n_total < 2: + # Only skip plotting, no data for this class + continue + + plot_df = _subsample_psms(sub, _CONFIDENCE_SCATTER_MAX_POINTS) + handle = ax.scatter( + plot_df["confidence"], + plot_df["calibrated_confidence"], + c=colour, + s=10, + alpha=0.3, + rasterized=True, + label=f"{label}", + ) + handles.append(handle) + + ax.plot( + [-0.01, 1.01], + [-0.01, 1.01], + ls="--", + color="black", + lw=1, + label="No recalibration", + zorder=5, + ) + ax.set_xlim(-0.01, 1.01) + ax.set_ylim(-0.01, 1.01) + ax.set_xlabel("Raw confidence") + ax.set_ylabel("Calibrated confidence") + ax.legend(loc="lower right", fontsize=9) + _style_ax(ax) + + fig.suptitle( + f"Raw versus calibrated confidence for {display}", + fontsize=13, + ) + fig.tight_layout() + _save_fig(fig, output_dir / f"confidence_scatter_{dataset_name}", plot_format) + + +def _plot_feature_distributions_at_fdr( + upscored: pd.DataFrame, + *, + fdr_t: float, + delta_threshold: float, + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Violin plots of features for up-scored TPs vs FPs retained at one FDR cutoff.""" + n = len(upscored) + pct = int(fdr_t * 100) + if n < _MIN_FEATURE_VIOLIN_PSMs: + logger.info( + "Skipping feature violins for %s at %d%% FDR (n=%d up-scored retained)", + dataset_name, + pct, + n, + ) + return + + available = [ + c + for c in _violin_feature_columns(upscored) + if c in upscored.columns + and upscored[c].notna().sum() >= _MIN_FEATURE_VIOLIN_PSMs + ] + if not available: + logger.warning( + "No feature columns with >=%d values for %s at %d%% FDR", + _MIN_FEATURE_VIOLIN_PSMs, + dataset_name, + pct, + ) + return + + plot_df = upscored.copy() + plot_df["label"] = plot_df["correct"].map({True: "Correct", False: "Incorrect"}) + n_features = len(available) + n_cols = min(3, n_features) + n_rows = (n_features + n_cols - 1) // n_cols + fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows)) + axes = np.atleast_1d(axes).flatten() + + palette = {"Correct": TP_COLOR, "Incorrect": FP_COLOR} + n_plotted = 0 + + for i, col in enumerate(available): + ax = axes[i] + sub = plot_df[[col, "label"]].dropna(subset=[col]) + if len(sub) < _MIN_FEATURE_VIOLIN_PSMs: + ax.set_visible(False) + continue + sns.violinplot( + data=sub, + x="label", + y=col, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + ax.set_ylabel(_nice_label(col)) + ax.set_title(_nice_label(col)) + _style_ax(ax) + n_plotted += 1 + + for i in range(len(available), len(axes)): + axes[i].set_visible(False) + + if n_plotted == 0: + plt.close(fig) + logger.info( + "Skipping feature violins for %s at %d%% FDR (no feature panels with n>=%d)", + dataset_name, + pct, + _MIN_FEATURE_VIOLIN_PSMs, + ) + return + + display = _folder_display_name(dataset_name) + fig.suptitle( + f"Feature distributions for up-scored PSMs retained at {pct}% FDR " + f"(calibration increase > {delta_threshold:.2f}) on {display}\n" + f"(n={n:,})", + fontsize=13, + ) + fig.tight_layout() + _save_fig( + fig, + output_dir / f"upscored_features_{dataset_name}_fdr{pct}", + plot_format, + ) + + +def _plot_feature_distributions( + df: pd.DataFrame, + delta_threshold: float, + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Violin plots at 5% and 10% FDR for up-scored correct vs incorrect PSMs.""" + df = _add_q_values(df) + upscored_mask = df["delta_confidence"] > delta_threshold + for fdr_t in _FEATURE_VIOLIN_FDR_THRESHOLDS: + retained = df[upscored_mask & (df["psm_q_value"] <= fdr_t)].copy() + _plot_feature_distributions_at_fdr( + retained, + fdr_t=fdr_t, + delta_threshold=delta_threshold, + dataset_name=dataset_name, + output_dir=output_dir, + plot_format=plot_format, + ) + + +def _upscored_fp_detail( + df: pd.DataFrame, + delta_threshold: float, + dataset_name: str, + metrics: Metrics, +) -> pd.DataFrame: + """Detailed characterisation of up-scored FPs that pass FDR thresholds.""" + df = _add_q_values(df) + upscored_fps = df[ + (df["delta_confidence"] > delta_threshold) & (~df["correct"].astype(bool)) + ].copy() + + if len(upscored_fps) == 0: + return pd.DataFrame() + + def _match_fraction(row: pd.Series) -> float: + nm = row.get("num_matches", 0) + seq = row.get("sequence", "") + if isinstance(seq, str): + tokens = metrics._split_peptide(seq) + else: + tokens = seq if seq else [] + return nm / len(tokens) if tokens else 0.0 + + upscored_fps["match_fraction"] = upscored_fps.apply(_match_fraction, axis=1) + + def _edit_dist(row: pd.Series) -> int: + s = _strip_mods(str(row.get("sequence", ""))) + p = _strip_mods(str(row.get("prediction", ""))) + if not s or not p: + return -1 + return _levenshtein(s, p) + + upscored_fps["edit_distance_norm"] = upscored_fps.apply(_edit_dist, axis=1) + + rows = [] + for fdr_t in FDR_THRESHOLDS: + sub = upscored_fps[upscored_fps["psm_q_value"] <= fdr_t] + if len(sub) == 0: + rows.append( + {"dataset": dataset_name, "fdr_threshold": fdr_t, "n_upscored_fps": 0} + ) + continue + rows.append( + { + "dataset": dataset_name, + "fdr_threshold": fdr_t, + "n_upscored_fps": len(sub), + "mean_match_fraction": round(float(sub["match_fraction"].mean()), 4), + "median_edit_distance": int(sub["edit_distance_norm"].median()), + "n_edit_dist_le2": int((sub["edit_distance_norm"] <= 2).sum()), + "n_partial_match": int((sub["match_fraction"] > 0).sum()), + } + ) + return pd.DataFrame(rows) + + +def _levenshtein(s: str, t: str) -> int: + """Simple Levenshtein distance for short peptide strings.""" + n, m = len(s), len(t) + if n == 0: + return m + if m == 0: + return n + prev = list(range(m + 1)) + for i in range(1, n + 1): + curr = [i] + [0] * m + for j in range(1, m + 1): + cost = 0 if s[i - 1] == t[j - 1] else 1 + curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost) + prev = curr + return prev[m] + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +_DEFAULT_PREDICTIONS_ROOT = Path("predictions/general_model") +_DEFAULT_OUTPUT_DIR = Path("analysis/upscored_fps") + + +@app.command() +def main( + predictions_root: Annotated[ + Path, + typer.Option( + help="Root directory containing winnow predict output folders.", + ), + ] = _DEFAULT_PREDICTIONS_ROOT, + output_dir: Annotated[ + Path, + typer.Option(help="Directory for output tables and plots."), + ] = _DEFAULT_OUTPUT_DIR, + delta_threshold: Annotated[ + float, + typer.Option( + help="Minimum delta (calibrated - raw) to classify a PSM as up-scored. " + "Default 0.2 (20 percentage-point increase).", + ), + ] = 0.2, + plot_format: Annotated[ + str, + typer.Option(help="Plot format: 'pdf', 'png', or 'both'."), + ] = "both", + residues_config: Annotated[ + Path, + typer.Option(help="Path to residues.yaml for InstaNovo Metrics."), + ] = _REPO_ROOT / "winnow" / "configs" / "residues.yaml", +) -> None: + """Characterise false positives that calibration up-scores into high-confidence regions.""" + output_dir.mkdir(parents=True, exist_ok=True) + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + residue_masses = _get_residue_masses() + metrics = Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + folders = _discover_labelled_folders(predictions_root) + if not folders: + logger.error("No labelled output folders found under %s", predictions_root) + raise typer.Exit(code=1) + + logger.info("Found %d labelled folder(s): %s", len(folders), list(folders.keys())) + + all_summary: list[pd.DataFrame] = [] + all_detail: list[pd.DataFrame] = [] + + for name, folder in folders.items(): + logger.info("Processing %s ...", name) + df = _load_dataset(folder) + + missing_conf = [ + c for c in ("confidence", "calibrated_confidence") if c not in df.columns + ] + if missing_conf: + logger.warning( + "Skipping %s: missing confidence columns %s", + name, + missing_conf, + ) + continue + if "correct" not in df.columns: + logger.warning("Skipping %s: missing 'correct' column", name) + continue + + df["delta_confidence"] = df["calibrated_confidence"] - df["confidence"] + + logger.info( + " %s: %d PSMs, %d correct, delta stats: mean=%.3f, q75=%.3f", + name, + len(df), + int(df["correct"].sum()), + df["delta_confidence"].mean(), + df["delta_confidence"].quantile(0.75), + ) + + summary = _upscored_summary_table(df, delta_threshold, name) + all_summary.append(summary) + + _plot_confidence_scatter(df, name, plots_dir, plot_format) + _plot_feature_distributions(df, delta_threshold, name, plots_dir, plot_format) + + detail = _upscored_fp_detail(df, delta_threshold, name, metrics) + if len(detail) > 0: + all_detail.append(detail) + + if all_summary: + combined = pd.concat(all_summary, ignore_index=True) + combined.to_csv(output_dir / "upscored_summary.csv", index=False) + logger.info("Summary table:\n%s", combined.to_string(index=False)) + + with open(output_dir / "upscored_summary.json", "w") as f: + json.dump(combined.to_dict(orient="records"), f, indent=2) + + if all_detail: + detail_df = pd.concat(all_detail, ignore_index=True) + detail_df.to_csv(output_dir / "upscored_fp_detail.csv", index=False) + logger.info("FP detail table:\n%s", detail_df.to_string(index=False)) + + logger.info("Up-scored FP analysis complete. Output in %s", output_dir) + + +if __name__ == "__main__": + app() diff --git a/scripts/benchmark_runtime.py b/scripts/benchmark_runtime.py new file mode 100644 index 00000000..b70998cf --- /dev/null +++ b/scripts/benchmark_runtime.py @@ -0,0 +1,694 @@ +#!/usr/bin/env python3 +"""Benchmark wall-clock time and memory for the winnow prediction pipeline. + +Measures end-to-end processing time and peak memory, broken down by: + (i) data loading, + (ii) per-feature computation (individually timed), + (iii) MLP calibration inference, and + (iv) FDR / q-value computation. + +Two configurations are benchmarked by default: + 1. Full feature set (including Prosit/Koina-derived features). + 2. Without Prosit features (fragment match + iRT removed). + +This directly addresses the modularity claim: Prosit-based features can be +omitted without disrupting the pipeline. + +Usage examples: + # Both configurations on sample data (requires Koina server for full run) + python scripts/benchmark_runtime.py + + # Only the no-Prosit configuration (no Koina server required) + python scripts/benchmark_runtime.py --no-prosit + + # Custom dataset and locally-trained model + python scripts/benchmark_runtime.py \ + --spectrum-path data/spectra.ipc \ + --predictions-path data/predictions.csv \ + --model-path models/my_model \ + --data-loader instanovo + + # Save structured results to JSON + python scripts/benchmark_runtime.py --output-json results/benchmark.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import resource +import time +import tracemalloc +from contextlib import contextmanager +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.calibration.features.fragment_match import FragmentMatchFeatures +from winnow.calibration.features.retention_time import RetentionTimeFeature +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.fdr.nonparametric import NonParametricFDRControl + + +PROSIT_FEATURE_CLASSES = (FragmentMatchFeatures, RetentionTimeFeature) + + +# --------------------------------------------------------------------------- +# Measurement helpers +# --------------------------------------------------------------------------- + + +@dataclass +class StageResult: + """Timing and memory result for a single pipeline stage.""" + + name: str + device: str + wall_time_s: float + peak_mem_mb: float + is_feature: bool = False + is_prosit: bool = False + columns: List[str] = field(default_factory=list) + + +@contextmanager +def measure(): + """Context manager that yields a dict populated with wall_time_s and peak_mem_mb on exit.""" + result: Dict[str, float] = {} + tracemalloc.start() + # Reset the peak so we measure only this block + tracemalloc.reset_peak() + t0 = time.perf_counter() + try: + yield result + finally: + result["wall_time_s"] = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + result["peak_mem_mb"] = peak / (1024 * 1024) + + +# --------------------------------------------------------------------------- +# Hardware info +# --------------------------------------------------------------------------- + + +def get_hardware_info() -> Dict[str, str]: + """Collect CPU, RAM, and GPU identifiers for the benchmark report.""" + info: Dict[str, str] = {} + + # CPU + cpu_name = None + try: + with open("/proc/cpuinfo") as f: + for line in f: + if line.startswith("model name"): + cpu_name = line.split(":", 1)[1].strip() + break + except OSError: + pass + if not cpu_name: + cpu_name = platform.processor() or "unknown" + info["cpu"] = cpu_name + info["cpu_cores"] = str(os.cpu_count() or "unknown") + + # RAM + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal"): + kb = int(line.split()[1]) + info["ram_gb"] = f"{kb / (1024**2):.0f}" + break + except OSError: + info["ram_gb"] = "unknown" + + # GPU + if torch.cuda.is_available(): + info["gpu"] = torch.cuda.get_device_name(0) + else: + info["gpu"] = "none" + + return info + + +# --------------------------------------------------------------------------- +# Pipeline stages +# --------------------------------------------------------------------------- + + +def load_dataset( + spectrum_path: str, + predictions_path: Optional[str], + data_loader_name: str, +) -> CalibrationDataset: + """Load and filter the dataset, returning a CalibrationDataset.""" + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + overrides = [f"data_loader={data_loader_name}"] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name="benchmark", + ): + cfg = compose(config_name="predict", overrides=overrides) + + data_loader = instantiate(cfg.data_loader) + dataset = data_loader.load( + data_path=spectrum_path, + predictions_path=predictions_path, + ) + + from winnow.scripts.main import filter_dataset + + dataset = filter_dataset(dataset) + return dataset + + +def compute_features_individually( + calibrator: ProbabilityCalibrator, + dataset: CalibrationDataset, +) -> List[StageResult]: + """Run each feature's prepare+compute with individual timing.""" + results: List[StageResult] = [] + + # Dependencies (currently all features return [], but measure for completeness) + for dep in calibrator.dependencies.values(): + with measure() as m: + dep.compute(dataset=dataset) + results.append( + StageResult( + name=f"Dependency: {dep.name}", + device="CPU", + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + ) + + for name, feat in calibrator.feature_dict.items(): + is_prosit = isinstance(feat, PROSIT_FEATURE_CLASSES) + device = "CPU + network" if is_prosit else "CPU" + + with measure() as m: + feat.prepare(dataset=dataset) + feat.compute(dataset=dataset) + + results.append( + StageResult( + name=f"Feature: {name}", + device=device, + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + is_feature=True, + is_prosit=is_prosit, + columns=list(feat.columns), + ) + ) + + return results + + +def run_mlp_inference( + calibrator: ProbabilityCalibrator, + dataset: CalibrationDataset, +) -> StageResult: + """Run MLP calibration inference.""" + if calibrator.network is None: + raise RuntimeError("Calibrator network is not loaded") + device = str(next(calibrator.network.parameters()).device) + with measure() as m: + calibrator.predict(dataset) + return StageResult( + name="MLP calibration inference", + device=device.upper() if device == "cpu" else device, + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + + +def run_fdr( + dataset: CalibrationDataset, + confidence_column: str = "calibrated_confidence", + fdr_threshold: float = 0.05, +) -> StageResult: + """Run FDR / q-value computation.""" + fdr_control = NonParametricFDRControl() + + with measure() as m: + fdr_control.fit(dataset=dataset.metadata[confidence_column]) + dataset.metadata = fdr_control.add_psm_pep(dataset.metadata, confidence_column) + dataset.metadata = fdr_control.add_psm_fdr(dataset.metadata, confidence_column) + dataset.metadata = fdr_control.add_psm_q_value( + dataset.metadata, confidence_column + ) + confidence_cutoff = fdr_control.get_confidence_cutoff(threshold=fdr_threshold) + _ = dataset.metadata[dataset.metadata[confidence_column] >= confidence_cutoff] + + return StageResult( + name="FDR / q-value computation", + device="CPU", + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + + +# --------------------------------------------------------------------------- +# Single benchmark run +# --------------------------------------------------------------------------- + + +@dataclass +class BenchmarkRun: + """Results from a single pipeline configuration.""" + + config_label: str + n_spectra: int + n_features: int + n_columns: int + stages: List[StageResult] + + @property + def total_wall_time_s(self) -> float: + """Sum of wall times across all recorded stages.""" + return sum(s.wall_time_s for s in self.stages) + + @property + def feature_wall_time_s(self) -> float: + """Sum of wall times for feature computation stages only.""" + return sum(s.wall_time_s for s in self.stages if s.is_feature) + + @property + def peak_mem_mb(self) -> float: + """Maximum peak memory across stages, in megabytes.""" + return max(s.peak_mem_mb for s in self.stages) if self.stages else 0.0 + + +def _model_matches_features(calibrator: ProbabilityCalibrator) -> bool: + """Check whether the loaded MLP input dim matches the current feature set.""" + if calibrator.network is None or calibrator.feature_mean is None: + return False + expected_dim = calibrator.feature_mean.shape[0] + actual_dim = 1 + len(calibrator.columns) # confidence + feature columns + return expected_dim == actual_dim + + +def run_benchmark( + spectrum_path: str, + predictions_path: Optional[str], + model_path: str, + data_loader_name: str, + include_prosit: bool, + koina_url: Optional[str] = None, + koina_ssl: Optional[bool] = None, +) -> BenchmarkRun: + """Execute the full prediction pipeline with per-stage timing.""" + config_label = "Full feature set" if include_prosit else "Without Prosit features" + + # Load calibrator + calibrator = ProbabilityCalibrator.load(pretrained_model_name_or_path=model_path) + + # Apply Koina server overrides if provided + if koina_url is not None or koina_ssl is not None: + calibrator.apply_koina_server_overrides(server_url=koina_url, ssl=koina_ssl) + + # Remove Prosit features if requested + if not include_prosit: + to_remove = [ + name + for name, feat in calibrator.feature_dict.items() + if isinstance(feat, PROSIT_FEATURE_CLASSES) + ] + for name in to_remove: + calibrator.remove_feature(name) + + stages: List[StageResult] = [] + + # Stage 1: Data loading + with measure() as m: + dataset = load_dataset(spectrum_path, predictions_path, data_loader_name) + n_spectra = len(dataset.metadata) + stages.append( + StageResult( + name="Data loading", + device="CPU", + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + ) + + # Stage 2: Per-feature computation + feature_results = compute_features_individually(calibrator, dataset) + stages.extend(feature_results) + + n_features = len(calibrator.feature_dict) + n_columns = len(calibrator.columns) + + # Stage 3: MLP calibration inference + # The MLP input dimension must match the feature set. If features were + # removed (e.g. Prosit features dropped) but the model was trained with + # the full set, the dimensions won't match. In that case we skip MLP + + # FDR and note the mismatch -- these stages are sub-millisecond anyway + # and their cost is independent of the feature set used. + can_infer = _model_matches_features(calibrator) + if can_infer: + stages.append(run_mlp_inference(calibrator, dataset)) + # Stage 4: FDR / q-value + stages.append(run_fdr(dataset)) + else: + mean_dim = ( + calibrator.feature_mean.shape[0] + if calibrator.feature_mean is not None + else "unknown" + ) + print( + f" [note] MLP input dim ({mean_dim}) " + f"does not match current feature count " + f"({1 + len(calibrator.columns)}). " + f"Skipping MLP inference and FDR for this configuration.\n" + f" To benchmark these stages, supply a model trained with the " + f"matching feature set via --model-path-no-prosit." + ) + + return BenchmarkRun( + config_label=config_label, + n_spectra=n_spectra, + n_features=n_features, + n_columns=n_columns, + stages=stages, + ) + + +# --------------------------------------------------------------------------- +# Output formatting +# --------------------------------------------------------------------------- + + +def format_run(run: BenchmarkRun) -> str: + """Format a single benchmark run as a human-readable table.""" + lines: List[str] = [] + header = ( + f"=== Configuration: {run.config_label} " + f"({run.n_features} features, {run.n_columns} columns) ===" + ) + lines.append("") + lines.append(header) + lines.append("") + + col_w = [38, 16, 15, 15] + hdr = ( + f"{'Stage':<{col_w[0]}}| {'Device':<{col_w[1]}}| " + f"{'Wall time (s)':>{col_w[2]}}| {'Peak mem (MB)':>{col_w[3]}}" + ) + sep = ( + "-" * col_w[0] + + "|" + + "-" * (col_w[1] + 1) + + "|" + + "-" * (col_w[2] + 1) + + "|" + + "-" * (col_w[3] + 1) + ) + + lines.append(hdr) + lines.append(sep) + + feat_time = 0.0 + feat_mem = 0.0 + total_time = 0.0 + + def _row(label: str, device: str, t: float, mem: float) -> str: + return ( + f"{label:<{col_w[0]}}| {device:<{col_w[1]}}| " + f"{t:>{col_w[2]}.2f}| {mem:>{col_w[3]}.1f}" + ) + + for i, s in enumerate(run.stages): + lines.append(_row(s.name, s.device, s.wall_time_s, s.peak_mem_mb)) + total_time += s.wall_time_s + if s.is_feature: + feat_time += s.wall_time_s + feat_mem = max(feat_mem, s.peak_mem_mb) + + is_last_feature = s.is_feature and not any( + st.is_feature for st in run.stages[i + 1 :] + ) + if is_last_feature: + lines.append( + _row( + " Feature computation subtotal", + "", + feat_time, + feat_mem, + ) + ) + + lines.append(sep) + total_mem = max(s.peak_mem_mb for s in run.stages) if run.stages else 0.0 + lines.append(_row("End-to-end total", "", total_time, total_mem)) + + return "\n".join(lines) + + +def format_full_report( + hw: Dict[str, str], + mlp_device: str, + runs: List[BenchmarkRun], + n_spectra: int, + spectrum_path: str, +) -> str: + """Assemble the complete benchmark report.""" + lines: List[str] = [] + lines.append("=" * 88) + lines.append(" Winnow Pipeline Runtime Benchmark") + lines.append("=" * 88) + lines.append("") + lines.append( + f"Hardware: {hw['cpu']} ({hw['cpu_cores']} cores), " + f"{hw['ram_gb']} GB RAM, GPU: {hw['gpu']}" + ) + lines.append(f"Dataset: {n_spectra:,} spectra from {spectrum_path}") + lines.append(f"Calibrator MLP device: {mlp_device}") + + for run in runs: + lines.append(format_run(run)) + + lines.append("") + lines.append("Notes:") + lines.append( + '- "CPU + network" = gRPC calls to a Koina/Triton server for' + " Prosit-derived spectral predictions." + ) + lines.append(" No local GPU is used by winnow during prediction.") + lines.append( + "- The MLP runs on CPU after loading. GPU is only used during training" + " (not benchmarked here)." + ) + lines.append( + "- Koina predictions are not cached to disk between runs; batching is" + " handled internally by koinapy." + ) + peak_rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + lines.append(f"- Process peak RSS: {peak_rss_mb:.1f} MB") + lines.append("") + return "\n".join(lines) + + +def build_json_report( + hw: Dict[str, str], + mlp_device: str, + runs: List[BenchmarkRun], + n_spectra: int, + spectrum_path: str, +) -> Dict[str, Any]: + """Build a structured dict suitable for JSON serialisation.""" + report: Dict[str, Any] = { + "hardware": hw, + "dataset": { + "spectrum_path": spectrum_path, + "n_spectra": n_spectra, + }, + "mlp_device": mlp_device, + "configurations": [], + } + for run in runs: + cfg: Dict[str, Any] = { + "label": run.config_label, + "n_features": run.n_features, + "n_columns": run.n_columns, + "total_wall_time_s": run.total_wall_time_s, + "feature_wall_time_s": run.feature_wall_time_s, + "stages": [asdict(s) for s in run.stages], + } + report["configurations"].append(cfg) + report["process_peak_rss_mb"] = ( + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + ) + report["notes"] = { + "koina_caching": ( + "Koina predictions are not cached to disk; each invocation " + "re-queries the server." + ), + "koina_batching": ( + "Batching is handled internally by koinapy " + "(gRPC streaming to the Koina/Triton server)." + ), + "gpu_usage": ( + "No local GPU is used during prediction. GPU is only used " + "during training (not benchmarked here)." + ), + } + return report + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments for the runtime benchmark.""" + parser = argparse.ArgumentParser( + description="Benchmark winnow prediction pipeline runtime and memory.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--spectrum-path", + default="examples/example_data/spectra.ipc", + help="Path to the spectrum data file (default: example data).", + ) + parser.add_argument( + "--predictions-path", + default="examples/example_data/predictions.csv", + help="Path to predictions file (default: example data).", + ) + parser.add_argument( + "--model-path", + default="InstaDeepAI/winnow-general-model", + help=( + "Path to a local calibrator directory or HuggingFace model " + "identifier (default: InstaDeepAI/winnow-general-model)." + ), + ) + parser.add_argument( + "--model-path-no-prosit", + default=None, + help=( + "Path to a calibrator trained without Prosit features, used for " + "the no-Prosit benchmark. If omitted, --model-path is used and " + "MLP/FDR stages are skipped when the input dimension mismatches." + ), + ) + parser.add_argument( + "--data-loader", + default="instanovo", + help="Data loader to use (default: instanovo).", + ) + parser.add_argument( + "--no-prosit", + action="store_true", + help=( + "Only benchmark without Prosit/Koina features. When omitted, " + "both configurations (full and no-Prosit) are benchmarked." + ), + ) + parser.add_argument( + "--full-only", + action="store_true", + help="Only benchmark the full feature set (skip no-Prosit run).", + ) + parser.add_argument( + "--koina-url", + default=None, + help="Koina server URL (e.g. localhost:8500).", + ) + parser.add_argument( + "--koina-ssl", + default=None, + type=lambda x: x.lower() in ("true", "1", "yes"), + help="Use SSL for Koina (true/false).", + ) + parser.add_argument( + "--output-json", + default=None, + help="Save structured results to a JSON file.", + ) + return parser.parse_args() + + +def main() -> None: + """Run configured pipeline benchmarks and print (optionally save) results.""" + args = parse_args() + + hw = get_hardware_info() + + # Determine which configs to run + run_full = not args.no_prosit + run_no_prosit = not args.full_only + + # Detect MLP device from a probe load + probe_calibrator = ProbabilityCalibrator.load( + pretrained_model_name_or_path=args.model_path + ) + if probe_calibrator.network is None: + raise RuntimeError("Calibrator network is not loaded") + mlp_device = str(next(probe_calibrator.network.parameters()).device) + del probe_calibrator + + runs: List[BenchmarkRun] = [] + n_spectra = 0 + + if run_full: + print("\n>>> Benchmarking: Full feature set (including Prosit) ...") + result = run_benchmark( + spectrum_path=args.spectrum_path, + predictions_path=args.predictions_path, + model_path=args.model_path, + data_loader_name=args.data_loader, + include_prosit=True, + koina_url=args.koina_url, + koina_ssl=args.koina_ssl, + ) + runs.append(result) + n_spectra = result.n_spectra + + if run_no_prosit: + print("\n>>> Benchmarking: Without Prosit features ...") + no_prosit_model = args.model_path_no_prosit or args.model_path + result = run_benchmark( + spectrum_path=args.spectrum_path, + predictions_path=args.predictions_path, + model_path=no_prosit_model, + data_loader_name=args.data_loader, + include_prosit=False, + koina_url=args.koina_url, + koina_ssl=args.koina_ssl, + ) + runs.append(result) + n_spectra = n_spectra or result.n_spectra + + report = format_full_report(hw, mlp_device, runs, n_spectra, args.spectrum_path) + print(report) + + if args.output_json: + json_report = build_json_report( + hw, mlp_device, runs, n_spectra, args.spectrum_path + ) + out_path = Path(args.output_json) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(json_report, f, indent=2) + print(f"JSON results saved to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_scaling.py b/scripts/benchmark_scaling.py new file mode 100644 index 00000000..4fbf4596 --- /dev/null +++ b/scripts/benchmark_scaling.py @@ -0,0 +1,421 @@ +#!/usr/bin/env python3 +"""Measure pipeline scaling by running the no-Prosit benchmark at multiple dataset sizes. + +Writes subsampled spectrum and prediction files to a temporary directory, then +runs the full pipeline (data loading, feature computation, MLP inference, FDR) +from scratch at each size. Produces a JSON file with the raw measurements and +a matplotlib figure showing per-stage scaling. + +Usage: + python scripts/benchmark_scaling.py \ + --spectrum-path held_out_projects/.../dataset-helaqc-raw-0000-0001.parquet \ + --predictions-path held_out_projects/.../dataset-helaqc-raw-0000-0001.csv \ + --model-path models/benchmark_model_no_prosit \ + --output-dir analysis +""" + +from __future__ import annotations + +import argparse +import json +import random +import shutil +import tempfile +import time +import tracemalloc +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +import seaborn as sns + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.calibration.features.fragment_match import FragmentMatchFeatures +from winnow.calibration.features.retention_time import RetentionTimeFeature +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.fdr.nonparametric import NonParametricFDRControl + +plt.switch_backend("Agg") + +# Paul Tol "bright" palette (colorblind-safe) — same as plot_eval_results.py +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +PROSIT_FEATURE_CLASSES = (FragmentMatchFeatures, RetentionTimeFeature) + +DEFAULT_FRACTIONS = [0.1, 0.5, 1.0] + + +@contextmanager +def measure(): + """Context manager that yields a dict populated with wall_time_s and peak_mem_mb on exit.""" + result: Dict[str, float] = {} + tracemalloc.start() + tracemalloc.reset_peak() + t0 = time.perf_counter() + try: + yield result + finally: + result["wall_time_s"] = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + result["peak_mem_mb"] = peak / (1024 * 1024) + + +def write_subsampled_files( + spectrum_path: str, + predictions_path: str, + fraction: float, + output_dir: Path, + seed: int = 42, +) -> Tuple[Path, Path, int]: + """Write subsampled spectrum and prediction files, returning paths and row count. + + Both files are joined on spectrum_id, subsampled together, and written out + so the data loader sees consistent, smaller files. + """ + spectra = pl.read_parquet(spectrum_path) + preds = pl.read_csv(predictions_path) + + n = len(spectra) + if fraction >= 1.0: + k = n + indices = list(range(n)) + else: + k = max(1, int(n * fraction)) + rng = random.Random(seed) + indices = sorted(rng.sample(range(n), k)) + + spectra_sub = spectra[indices] + # Match predictions to the subsampled spectra by spectrum_id + keep_ids = set(spectra_sub["spectrum_id"].to_list()) + preds_sub = preds.filter(pl.col("spectrum_id").is_in(keep_ids)) + + spec_path = output_dir / f"spectra_{fraction:.2f}.parquet" + pred_path = output_dir / f"preds_{fraction:.2f}.csv" + spectra_sub.write_parquet(spec_path) + preds_sub.write_csv(pred_path) + + return spec_path, pred_path, len(spectra_sub) + + +def load_dataset_timed( + spectrum_path: str, + predictions_path: str, + data_loader_name: str, +) -> Tuple[CalibrationDataset, float]: + """Load a dataset through the full data loader and return it with timing.""" + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + overrides = [f"data_loader={data_loader_name}"] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name="benchmark_scaling", + ): + cfg = compose(config_name="predict", overrides=overrides) + + data_loader = instantiate(cfg.data_loader) + + with measure() as m: + dataset = data_loader.load( + data_path=spectrum_path, + predictions_path=predictions_path, + ) + from winnow.scripts.main import filter_dataset + + dataset = filter_dataset(dataset) + + return dataset, m["wall_time_s"] + + +@dataclass +class ScalingPoint: + """Timing measurements for one dataset-size fraction.""" + + fraction: float + n_spectra: int + stage_times: Dict[str, float] = field(default_factory=dict) + total_time: float = 0.0 + + +def run_at_size( + spec_path: Path, + pred_path: Path, + expected_n: int, + calibrator: ProbabilityCalibrator, + data_loader_name: str, + fraction: float, +) -> ScalingPoint: + """Run the full pipeline from disk at a given dataset size.""" + stage_times: Dict[str, float] = {} + + # Data loading (from subsampled files on disk) + dataset, load_time = load_dataset_timed( + str(spec_path), str(pred_path), data_loader_name + ) + n = len(dataset.metadata) + stage_times["Data loading"] = load_time + + # Feature computation (per feature) + feature_total = 0.0 + for name, feat in calibrator.feature_dict.items(): + with measure() as m: + feat.prepare(dataset=dataset) + feat.compute(dataset=dataset) + stage_times[f"Feature: {name}"] = m["wall_time_s"] + feature_total += m["wall_time_s"] + stage_times["Feature computation (total)"] = feature_total + + # MLP inference + with measure() as m: + calibrator.predict(dataset) + stage_times["MLP inference"] = m["wall_time_s"] + + # FDR / q-value + fdr = NonParametricFDRControl() + col = "calibrated_confidence" + with measure() as m: + fdr.fit(dataset=dataset.metadata[col]) + dataset.metadata = fdr.add_psm_pep(dataset.metadata, col) + dataset.metadata = fdr.add_psm_fdr(dataset.metadata, col) + dataset.metadata = fdr.add_psm_q_value(dataset.metadata, col) + cutoff = fdr.get_confidence_cutoff(threshold=0.05) + _ = dataset.metadata[dataset.metadata[col] >= cutoff] + stage_times["FDR / q-value"] = m["wall_time_s"] + + total = ( + load_time + + feature_total + + stage_times["MLP inference"] + + stage_times["FDR / q-value"] + ) + stage_times["End-to-end"] = total + + return ScalingPoint( + fraction=fraction, + n_spectra=n, + stage_times=stage_times, + total_time=total, + ) + + +def fit_exponent(sizes: List[int], times: List[float]) -> Tuple[float, float]: + """Fit t = c * n^alpha in log-log space; return (alpha, R²).""" + log_s = np.log(sizes) + log_t = np.log(times) + slope, intercept = np.polyfit(log_s, log_t, 1) + predicted = slope * log_s + intercept + ss_res = float(np.sum((log_t - predicted) ** 2)) + ss_tot = float(np.sum((log_t - np.mean(log_t)) ** 2)) + r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0 + return slope, r2 + + +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + """Save figure as both PNG and PDF.""" + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def plot_scaling(points: List[ScalingPoint], output_path: Path) -> None: + """Plot per-stage wall time vs. dataset size.""" + sizes = [p.n_spectra for p in points] + + stages_to_plot = [ + ("Data loading", "Data loading", "v", "-", _PALETTE[0]), + ("Feature computation", "Feature computation (total)", "o", "-", _PALETTE[2]), + ("MLP inference", "MLP inference", "^", "-", _PALETTE[1]), + ("FDR / q-value", "FDR / q-value", "D", "-", _PALETTE[3]), + ("End-to-end", "End-to-end", "s", "--", _PALETTE[5]), + ] + + fig, ax = plt.subplots(figsize=(6, 4)) + + for label, key, marker, linestyle, colour in stages_to_plot: + times = [p.stage_times[key] for p in points] + ax.plot( + sizes, + times, + marker=marker, + linestyle=linestyle, + label=label, + color=colour, + alpha=0.7, + ) + + max_size = max(sizes) + max_total = max(p.stage_times["End-to-end"] for p in points) + ref_sizes = np.linspace(0, max_size, 50) + ref_times = max_total * (ref_sizes / max_size) + ax.plot( + ref_sizes, + ref_times, + ls=":", + color=_PALETTE[6], + linewidth=1, + label="Linear reference", + ) + + ax.set_xlabel("Number of spectra") + ax.set_ylabel("Wall time (s)") + ax.set_ylim(top=300) + ax.set_title("Pipeline scaling\nexcluding Koina-dependent features") + ax.legend(loc="upper left", fontsize=9) + fig.tight_layout() + + base = output_path.with_suffix("") + _save_fig(fig, base) + print(f"Scaling plot saved to {base}.png and {base}.pdf") + + +def load_points_from_json(json_path: Path) -> List[ScalingPoint]: + """Load previously saved scaling measurements from JSON.""" + with open(json_path) as f: + data = json.load(f) + return [ + ScalingPoint( + fraction=p["fraction"], + n_spectra=p["n_spectra"], + stage_times=p["stage_times"], + total_time=p["total_time"], + ) + for p in data["points"] + ] + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments for the scaling benchmark.""" + parser = argparse.ArgumentParser( + description="Measure pipeline scaling at multiple dataset sizes.", + ) + parser.add_argument( + "--replot-json", + metavar="PATH", + help="Replot from a saved benchmark_scaling.json without rerunning benchmarks.", + ) + parser.add_argument( + "--spectrum-path", + help="Path to the spectrum data file.", + ) + parser.add_argument( + "--predictions-path", + help="Path to predictions file.", + ) + parser.add_argument( + "--model-path", + help="Path to a calibrator trained without Prosit features.", + ) + parser.add_argument( + "--data-loader", + default="instanovo", + help="Data loader to use (default: instanovo).", + ) + parser.add_argument( + "--fractions", + nargs="+", + type=float, + default=DEFAULT_FRACTIONS, + help="Dataset fractions to benchmark (default: 0.1 0.5 1.0).", + ) + parser.add_argument( + "--output-dir", + default="analysis", + help="Directory for output files.", + ) + args = parser.parse_args() + if args.replot_json is None: + for name in ("spectrum_path", "predictions_path", "model_path"): + if getattr(args, name) is None: + parser.error( + f"--{name.replace('_', '-')} is required unless --replot-json is set" + ) + return args + + +def main() -> None: + """Run scaling benchmarks or replot from a saved JSON file.""" + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + if args.replot_json: + replot_points = load_points_from_json(Path(args.replot_json)) + plot_scaling(replot_points, output_dir / "benchmark_scaling.png") + return + + # Load calibrator (without Prosit features) + calibrator = ProbabilityCalibrator.load( + pretrained_model_name_or_path=args.model_path + ) + to_remove = [ + name + for name, feat in calibrator.feature_dict.items() + if isinstance(feat, PROSIT_FEATURE_CLASSES) + ] + for name in to_remove: + calibrator.remove_feature(name) + + fractions = sorted(args.fractions) + tmpdir = Path(tempfile.mkdtemp(prefix="winnow_scaling_")) + points: List[ScalingPoint] = [] + + try: + # Pre-write all subsampled files + file_info: List[Tuple[float, Path, Path, int]] = [] + print("Preparing subsampled files ...") + for frac in fractions: + spec_p, pred_p, n = write_subsampled_files( + args.spectrum_path, args.predictions_path, frac, tmpdir + ) + file_info.append((frac, spec_p, pred_p, n)) + print(f" {frac:.0%}: {n:,} spectra -> {spec_p.name}, {pred_p.name}") + + for frac, spec_p, pred_p, expected_n in file_info: + print(f"\n>>> Running at {frac:.0%} ({expected_n:,} spectra) ...") + point = run_at_size( + spec_p, pred_p, expected_n, calibrator, args.data_loader, frac + ) + points.append(point) + print(f" {point.n_spectra:,} spectra -> {point.total_time:.2f} s total") + for stage, t in point.stage_times.items(): + if stage not in ("Feature computation (total)", "End-to-end"): + print(f" {stage}: {t:.3f} s") + + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + # Save JSON + json_path = output_dir / "benchmark_scaling.json" + json_data: Dict[str, Any] = { + "points": [ + { + "fraction": p.fraction, + "n_spectra": p.n_spectra, + "stage_times": p.stage_times, + "total_time": p.total_time, + } + for p in points + ], + } + with open(json_path, "w") as f: + json.dump(json_data, f, indent=2) + print(f"\nScaling data saved to {json_path}") + + # Plot + plot_path = output_dir / "benchmark_scaling.png" + plot_scaling(points, plot_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/calibrator_generalisation_utils.py b/scripts/calibrator_generalisation_utils.py new file mode 100644 index 00000000..278cc5ef --- /dev/null +++ b/scripts/calibrator_generalisation_utils.py @@ -0,0 +1,114 @@ +"""Shared helpers for calibrator generalisation analysis.""" + +from __future__ import annotations + +import logging +import re +from pathlib import Path + +import polars as pl + +logger = logging.getLogger(__name__) + +HEPG2_SOURCE = "PXD019483" + +SPECIES_NAME_MAPPING: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + HEPG2_SOURCE: "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD014877": "$\\it{C.\\;elegans}$", +} + + +def extract_project_name(parquet_path: Path) -> str: + """Extract project name from ``dataset-helaqc-annotated-0000-0001.parquet``.""" + match = re.match(r"dataset-(.+?)-annotated", parquet_path.stem) + if match: + return match.group(1) + return parquet_path.stem + + +def build_experiment_source_mapping(biological_validation_dir: Path) -> dict[str, str]: + """Map every experiment in biological validation parquets to its source label.""" + mapping: dict[str, str] = {} + parquet_files = sorted(biological_validation_dir.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError( + f"No parquet files found in biological validation directory: " + f"{biological_validation_dir}" + ) + + for parquet_path in parquet_files: + project = extract_project_name(parquet_path) + experiments = ( + pl.scan_parquet(parquet_path) + .select("experiment_name") + .unique() + .collect()["experiment_name"] + .to_list() + ) + for experiment_name in experiments: + mapping[experiment_name] = project + + logger.info( + "Built experiment->source mapping for %d experiments across %d projects", + len(mapping), + len(parquet_files), + ) + return mapping + + +def annotate_train_source_labels( + train_parquet: Path, + train_predictions: Path, + biological_validation_dir: Path, +) -> None: + """Add a ``source`` column to the train parquet and predictions CSV. + + Experiments found in ``biological_validation_dir`` inherit that project name. + All other experiments are labelled as HepG2 (``PXD019483``). + """ + experiment_to_source = build_experiment_source_mapping(biological_validation_dir) + lookup = pl.DataFrame( + { + "experiment_name": list(experiment_to_source.keys()), + "source": list(experiment_to_source.values()), + } + ) + + spectra = pl.read_parquet(train_parquet) + if "source" not in spectra.columns: + spectra = spectra.join(lookup, on="experiment_name", how="left").with_columns( + pl.col("source").fill_null(HEPG2_SOURCE) + ) + spectra.write_parquet(train_parquet) + logger.info("Wrote source labels to %s", train_parquet) + else: + logger.info( + "Parquet already has source column, leaving %s unchanged", train_parquet + ) + + predictions = pl.read_csv(train_predictions) + if "source" not in predictions.columns: + source_by_spectrum = spectra.select("spectrum_id", "source") + predictions = predictions.join(source_by_spectrum, on="spectrum_id", how="left") + missing = predictions.filter(pl.col("source").is_null()) + if len(missing) > 0: + raise ValueError( + f"{len(missing)} prediction rows in {train_predictions} have no matching " + "spectrum_id in the train parquet" + ) + predictions.write_csv(train_predictions) + logger.info("Wrote source labels to %s", train_predictions) + else: + logger.info( + "Predictions CSV already has source column, leaving %s unchanged", + train_predictions, + ) diff --git a/scripts/evaluate_calibrator_generalisation.py b/scripts/evaluate_calibrator_generalisation.py new file mode 100644 index 00000000..8947b51e --- /dev/null +++ b/scripts/evaluate_calibrator_generalisation.py @@ -0,0 +1,406 @@ +"""Evaluate calibrator generalisation by training on one source dataset and testing on all others. + +Uses the ``train_extra_small`` train parquet and predictions CSV, with a ``source`` +column derived from biological-validation experiment names (everything else is HepG2). +For each source, trains a fresh calibrator, evaluates it in-distribution (held-out +20 %) and out-of-distribution (every other source), then saves a combined results CSV. +""" + +import logging +import re +import sys +from pathlib import Path +from typing import Annotated, Dict, List, Optional + +import numpy as np +import pandas as pd +import yaml +from rich.logging import RichHandler +import typer + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.calibration.features import ( + BeamFeatures, + FragmentMatchFeatures, + MassErrorDaFeature, + RetentionTimeFeature, + TokenScoreFeatures, +) +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.datasets.data_loaders import InstaNovoDatasetLoader + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.calibrator_generalisation_utils import annotate_train_source_labels # noqa: E402 + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +logger = logging.getLogger("winnow.evaluate_generalization") +logger.setLevel(logging.INFO) +logger.propagate = False +logger.addHandler(RichHandler()) + +# --------------------------------------------------------------------------- +# Constants — loaded from the canonical Winnow YAML configs +# --------------------------------------------------------------------------- +SEED = 42 +TEST_SIZE = 0.2 + +_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "winnow" / "configs" + +with open(_CONFIGS_DIR / "residues.yaml") as _f: + RESIDUE_MASSES: dict[str, float] = yaml.safe_load(_f)["residue_masses"] + +with open(_CONFIGS_DIR / "data_loader" / "instanovo.yaml") as _f: + _instanovo_cfg = yaml.safe_load(_f) + RESIDUE_REMAPPING: dict[str, str] = _instanovo_cfg.get("residue_remapping", {}) + BEAM_COLUMNS: dict[str, str] | None = _instanovo_cfg.get("beam_columns") + +with open(_CONFIGS_DIR / "calibrator.yaml") as _f: + _calibrator_cfg = yaml.safe_load(_f) + _KOINA_CFG = _calibrator_cfg["koina"] + _KOINA_CONSTRAINTS = _KOINA_CFG["constraints"] + _KOINA_INPUT_CONSTANTS = _KOINA_CFG.get("input_constants") or { + "collision_energies": 27, + "fragmentation_types": "HCD", + } + _UNSUPPORTED_RESIDUES: list[str] = ( + _KOINA_CONSTRAINTS.get("unsupported_residues") or [] + ) + _MAX_PRECURSOR_CHARGE: int = _KOINA_CONSTRAINTS["max_precursor_charge"] + _MAX_PEPTIDE_LENGTH: int = _KOINA_CONSTRAINTS["max_peptide_length"] + _INTENSITY_MODEL: str = _KOINA_CFG["intensity_model"] + _IRT_MODEL: str = _KOINA_CFG["irt_model"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_IRT_TRAIN_FRACTION_OVERRIDES: Dict[str, float] = { + "herceptin": 0.15, +} + +# Mirrors Makefile train-extra-small-mass-error-da / EXTRA_SMALL_* overrides. +_EXTRA_SMALL_FRAGMENT_EXCLUDE = [ + "spectral_angle", + "xcorr", + "complementary_ion_count", + "max_ion_gap", +] +_EXTRA_SMALL_BEAM_EXCLUDE = ["edit_distance"] + + +def initialise_calibrator( + *, + koina_server_url: Optional[str] = None, + koina_ssl: bool = True, + train_project: Optional[str] = None, +) -> ProbabilityCalibrator: + """Create a fresh calibrator matching train-extra-small-mass-error-da.""" + koina_kwargs: Dict = {} + if koina_server_url is not None: + koina_kwargs["koina_server_url"] = koina_server_url + koina_kwargs["koina_ssl"] = koina_ssl + + irt_train_fraction = _IRT_TRAIN_FRACTION_OVERRIDES.get(train_project or "", 0.1) + + calibrator = ProbabilityCalibrator( + hidden_dims=(50, 50), + dropout=0.3, + learning_rate=0.0001, + weight_decay=0.001, + max_epochs=1000, + batch_size=1024, + n_iter_no_change=10, + tol=0.0001, + seed=SEED, + val_early_stopping_max_psms=None, + val_subsample_seed=None, + ) + calibrator.add_feature(MassErrorDaFeature(residue_masses=RESIDUE_MASSES)) + calibrator.add_feature( + FragmentMatchFeatures( + mz_tolerance_ppm=20, + learn_from_missing=False, + intensity_model_name=_INTENSITY_MODEL, + max_precursor_charge=_MAX_PRECURSOR_CHARGE, + max_peptide_length=_MAX_PEPTIDE_LENGTH, + unsupported_residues=_UNSUPPORTED_RESIDUES, + model_input_constants=_KOINA_INPUT_CONSTANTS, + excluded_columns=_EXTRA_SMALL_FRAGMENT_EXCLUDE, + **koina_kwargs, + ) + ) + calibrator.add_feature( + RetentionTimeFeature( + train_fraction=irt_train_fraction, + min_train_points=3, + learn_from_missing=False, + irt_model_name=_IRT_MODEL, + max_peptide_length=_MAX_PEPTIDE_LENGTH, + unsupported_residues=_UNSUPPORTED_RESIDUES, + **koina_kwargs, + ) + ) + calibrator.add_feature(BeamFeatures(excluded_columns=_EXTRA_SMALL_BEAM_EXCLUDE)) + calibrator.add_feature(TokenScoreFeatures()) + return calibrator + + +def load_dataset(data_path: Path, predictions_path: Path) -> CalibrationDataset: + """Load the combined train_extra_small dataset.""" + logger.info("Loading dataset from %s and %s", data_path, predictions_path) + loader = InstaNovoDatasetLoader( + residue_masses=RESIDUE_MASSES, + residue_remapping=RESIDUE_REMAPPING, + beam_columns=BEAM_COLUMNS, + ) + return loader.load(data_path=data_path, predictions_path=predictions_path) + + +def subset_dataset(dataset: CalibrationDataset, idx: np.ndarray) -> CalibrationDataset: + """Return a row subset of *dataset* with aligned beam predictions.""" + meta = dataset.metadata.iloc[idx].reset_index(drop=True) + preds = ( + [dataset.predictions[i] for i in idx.tolist()] + if dataset.predictions is not None + else None + ) + return CalibrationDataset(metadata=meta, predictions=preds) + + +def split_dataset_by_source( + dataset: CalibrationDataset, +) -> Dict[str, CalibrationDataset]: + """Split a combined dataset into one CalibrationDataset per ``source`` label.""" + if "source" not in dataset.metadata.columns: + raise ValueError( + "Expected a 'source' column in the train parquet metadata. " + "Run annotate_train_source_labels() first." + ) + + datasets: Dict[str, CalibrationDataset] = {} + for source in sorted(dataset.metadata["source"].unique()): + idx = np.where(dataset.metadata["source"].values == source)[0] + datasets[source] = subset_dataset(dataset, idx) + return datasets + + +_MOD_RE = re.compile(r"\[UNIMOD:\d+\]") + + +def _peptide_key(tokens: object) -> str: + """Normalise a tokenised peptide to a modification-free, I/L-collapsed key. + + Matches the strategy in ``scripts/split_annotated_raw_parquets.py``: + strip UNIMOD modifications, normalise I→L. + """ + if not isinstance(tokens, list): + return "__MISSING__" + stripped = [_MOD_RE.sub("", tok).replace("I", "L") for tok in tokens] + return "".join(stripped) + + +def create_train_test_split( + dataset: CalibrationDataset, +) -> tuple[CalibrationDataset, CalibrationDataset]: + """Split a dataset 80/20 by peptide so no peptide appears in both folds.""" + meta = dataset.metadata + n = len(meta) + if n <= 1: + return dataset, dataset + + pep_keys = meta["sequence"].apply(_peptide_key) + unique_peptides = pep_keys.unique() + + rng = np.random.default_rng(SEED) + perm = rng.permutation(len(unique_peptides)) + n_train = int(len(unique_peptides) * (1 - TEST_SIZE)) + + train_peptides = set(unique_peptides[perm[:n_train]]) + train_mask = pep_keys.isin(train_peptides).values + + train_idx = np.where(train_mask)[0] + test_idx = np.where(~train_mask)[0] + + return subset_dataset(dataset, train_idx), subset_dataset(dataset, test_idx) + + +def evaluate_model( + model: ProbabilityCalibrator, + test_dataset: CalibrationDataset, + train_project: str, + test_project: str, + evaluation_type: str, +) -> pd.DataFrame: + """Run prediction and tag the results.""" + model.compute_features(test_dataset) + model.predict(test_dataset) + + results = test_dataset.metadata.copy() + results["trained_on_dataset"] = train_project + results["test_dataset"] = test_project + results["evaluation_type"] = evaluation_type + return results + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +_DEFAULT_MODEL_OUTPUT_DIR = Path("models/generalisation") +_DEFAULT_RESULTS_OUTPUT_DIR = Path("results/generalisation") +_DEFAULT_TRAIN_PARQUET = Path("train_extra_small/train.parquet") +_DEFAULT_TRAIN_PREDS = Path("train_extra_small/train_preds.csv") +_DEFAULT_BIOLOGICAL_VALIDATION_DIR = Path( + "held_out_projects/biological_validation/annotated" +) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +@app.command() +def main( + train_parquet: Annotated[ + Path, typer.Option(help="Combined train_extra_small parquet file.") + ] = _DEFAULT_TRAIN_PARQUET, + train_predictions: Annotated[ + Path, typer.Option(help="Combined train_extra_small predictions CSV.") + ] = _DEFAULT_TRAIN_PREDS, + biological_validation_dir: Annotated[ + Path, + typer.Option( + help=( + "Directory of biological validation annotated parquets used to map " + "experiment_name values to source labels." + ) + ), + ] = _DEFAULT_BIOLOGICAL_VALIDATION_DIR, + model_output_dir: Annotated[ + Path, typer.Option(help="Directory to save trained models.") + ] = _DEFAULT_MODEL_OUTPUT_DIR, + results_output_dir: Annotated[ + Path, typer.Option(help="Directory to save evaluation results.") + ] = _DEFAULT_RESULTS_OUTPUT_DIR, + koina_server_url: Annotated[ + Optional[str], typer.Option(help="Koina server URL override.") + ] = None, + koina_ssl: Annotated[bool, typer.Option(help="Use SSL for Koina server.")] = True, +) -> None: + """Evaluate calibrator generalisation across train_extra_small source datasets.""" + model_output_dir.mkdir(parents=True, exist_ok=True) + results_output_dir.mkdir(parents=True, exist_ok=True) + + if not train_parquet.exists(): + logger.error("Train parquet not found: %s", train_parquet) + raise typer.Exit(1) + if not train_predictions.exists(): + logger.error("Train predictions CSV not found: %s", train_predictions) + raise typer.Exit(1) + if not biological_validation_dir.exists(): + logger.error( + "Biological validation directory not found: %s", biological_validation_dir + ) + raise typer.Exit(1) + + annotate_train_source_labels( + train_parquet, train_predictions, biological_validation_dir + ) + + full_dataset = load_dataset(train_parquet, train_predictions) + datasets = split_dataset_by_source(full_dataset) + logger.info("Found %d source datasets: %s", len(datasets), list(datasets.keys())) + for source, dataset in datasets.items(): + logger.info(" %s: %d samples", source, len(dataset.metadata)) + + # Train-on-each, evaluate-on-all + all_results: List[pd.DataFrame] = [] + for train_project in datasets: + logger.info("=== Training on %s ===", train_project) + + train_ds, in_dist_test_ds = create_train_test_split(datasets[train_project]) + logger.info( + " train: %d, in-dist test: %d", + len(train_ds.metadata), + len(in_dist_test_ds.metadata), + ) + + calibrator = initialise_calibrator( + koina_server_url=koina_server_url, + koina_ssl=koina_ssl, + train_project=train_project, + ) + calibrator.fit(train_ds) + + model_path = model_output_dir / f"trained_on_{train_project}" + ProbabilityCalibrator.save(calibrator, model_path) + + # In-distribution evaluation + logger.info( + " Evaluating in-distribution on %s (%d samples)", + train_project, + len(in_dist_test_ds.metadata), + ) + all_results.append( + evaluate_model( + calibrator, + in_dist_test_ds, + train_project, + train_project, + "in_distribution", + ) + ) + + # Out-of-distribution evaluation + for test_project in datasets: + if test_project == train_project: + continue + test_ds = datasets[test_project] + logger.info( + " Evaluating out-of-distribution on %s (%d samples)", + test_project, + len(test_ds.metadata), + ) + all_results.append( + evaluate_model( + calibrator, + test_ds, + train_project, + test_project, + "out_of_distribution", + ) + ) + + # Combine and save + combined = pd.concat(all_results, ignore_index=True) + + # Drop large array columns to save space + array_cols = [c for c in ["mz_array", "intensity_array"] if c in combined.columns] + if array_cols: + combined = combined.drop(columns=array_cols) + + results_path = results_output_dir / "calibrator_generalisation_results.csv" + combined.to_csv(results_path, index=False) + logger.info("Results saved to %s", results_path) + + # Summary + logger.info("Evaluation summary:") + summary = ( + combined.groupby(["trained_on_dataset", "test_dataset", "evaluation_type"]) + .size() + .reset_index(name="num_samples") + ) + for _, row in summary.iterrows(): + logger.info( + " Trained on %s, tested on %s (%s): %d samples", + row["trained_on_dataset"], + row["test_dataset"], + row["evaluation_type"], + row["num_samples"], + ) + + +if __name__ == "__main__": + app() diff --git a/scripts/fdr_tool_comparison_preprocess.py b/scripts/fdr_tool_comparison_preprocess.py new file mode 100644 index 00000000..2eed87de --- /dev/null +++ b/scripts/fdr_tool_comparison_preprocess.py @@ -0,0 +1,774 @@ +"""Shared preprocessing helpers for FDR tool comparisons. + +PSM comparison and the external peptide score-mixture share: +1. Method-specific load / NovoBoard mass-delta → ProForma conversion. +2. Pair-gated NovoBoard target-decoy filters (equal twin counts). +3. :func:`filter_prediction_table` / :func:`filter_novoboard_prediction_table`. + +Labelled correctness uses Novor token matching; proteome-hit proxies use +PTM-stripped I→L substring search against an I→L FASTA haystack. + +Peptide score-mixture only then: +4. :func:`max_score_per_peptide` (no re-filtering). +5. NovoBoard max-target → twin-decoy peptide TDC helpers. +""" + +from __future__ import annotations + +import logging +import re +from functools import lru_cache +from pathlib import Path +from typing import Iterable + +import numpy as np +import pandas as pd +import yaml +from instanovo.utils.metrics import Metrics +from instanovo.utils.residues import ResidueSet + +from scripts.annotate_preds_proteome_hits import ( + _batch_peptide_substring_hits, +) + +logger = logging.getLogger(__name__) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_RESIDUES_YAML = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" + +MIN_PEPTIDE_LENGTH = 8 +# Labelled / reference sets use Novor agreement, so short peptides are valid. +# Keep a non-empty-key floor only. Unlabelled sets keep MIN_PEPTIDE_LENGTH +# because correctness is proteome substring membership. +LABELLED_MIN_PEPTIDE_LENGTH = 1 +_UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") +_MOD_SQUARE = re.compile(r"\[.*?\]") +_MOD_PAREN = re.compile(r"\(.*?\)") +_NOVOBOARD_TO_PROFORMA = { + "C(+57.02)": "C[UNIMOD:4]", + "M(+15.99)": "M[UNIMOD:35]", + "N(+0.98)": "N[UNIMOD:7]", + "Q(+0.98)": "Q[UNIMOD:7]", + "S(+79.97)": "S[UNIMOD:21]", + "T(+79.97)": "T[UNIMOD:21]", + "Y(+79.97)": "Y[UNIMOD:21]", +} + + +def normalize_peptide_key(peptide: object) -> str: + """Normalise sequence-only peptide identity (strip PTMs, I→L).""" + if isinstance(peptide, list): + peptide = "".join(str(token) for token in peptide) + if pd.isna(peptide) or not isinstance(peptide, str): + return "" + if len(peptide) > 4 and peptide[1] == "." and peptide[-2] == ".": + peptide = peptide[2:-2] + seq = _MOD_SQUARE.sub("", peptide) + seq = _MOD_PAREN.sub("", seq) + seq = "".join(c for c in seq if c.isalpha()) + return seq.replace("I", "L") + + +def has_unsupported_unimod(peptide: object) -> bool: + """Return True when *peptide* still contains an unsupported ``[UNIMOD:n]`` token.""" + if pd.isna(peptide) or not isinstance(peptide, str): + return True + return bool(_UNIMOD_RE.search(peptide)) + + +def novoboard_to_proforma(peptide: object) -> object: + """Convert NovoBoard's supported mass-delta notation to ProForma.""" + if pd.isna(peptide) or not isinstance(peptide, str): + return peptide + converted = peptide + for novoboard_mod, proforma_mod in _NOVOBOARD_TO_PROFORMA.items(): + converted = converted.replace(novoboard_mod, proforma_mod) + return converted + + +def sequence_only_correct_prediction(sequence: object, prediction: object) -> bool: + """Full sequence equality after PTM stripping and I/L normalisation.""" + sequence_key = normalize_peptide_key(sequence) + prediction_key = normalize_peptide_key(prediction) + return bool(sequence_key) and sequence_key == prediction_key + + +def sequence_only_correctness_mask( + sequences: pd.Series, predictions: pd.Series +) -> np.ndarray: + """Vectorized strip-PTM I→L equality (not for labelled Novor eval).""" + return np.array( + [ + sequence_only_correct_prediction(sequence, prediction) + for sequence, prediction in zip(sequences, predictions) + ], + dtype=bool, + ) + + +def load_residue_masses(residues_yaml: Path | None = None) -> dict[str, float]: + """Load residue masses from Winnow's residues YAML.""" + path = residues_yaml if residues_yaml is not None else DEFAULT_RESIDUES_YAML + with path.open(encoding="utf-8") as handle: + return yaml.safe_load(handle)["residue_masses"] + + +@lru_cache(maxsize=4) +def _metrics_from_residue_masses_frozen( + residues_items: tuple[tuple[str, float], ...], +) -> Metrics: + residue_masses = dict(residues_items) + return Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + +def metrics_from_residue_masses(residue_masses: dict[str, float]) -> Metrics: + """Build an InstaNovo ``Metrics`` instance for Novor matching.""" + items = tuple(sorted((str(k), float(v)) for k, v in residue_masses.items())) + return _metrics_from_residue_masses_frozen(items) + + +def novor_correct_prediction( + sequence: object, + prediction: object, + metrics: Metrics, +) -> bool: + """Winnow/InstaNovo Novor correctness: full residue-token match.""" + if isinstance(sequence, list): + gt = sequence + elif pd.isna(sequence) or not isinstance(sequence, str) or not sequence: + return False + else: + gt = metrics._split_peptide(sequence) + + if isinstance(prediction, list): + pred = prediction + elif pd.isna(prediction) or not isinstance(prediction, str) or not prediction: + return False + else: + pred = metrics._split_peptide(prediction) + + if not gt or not pred: + return False + num_matches = metrics._novor_match(gt, pred) + return bool(num_matches == len(gt) == len(pred)) + + +def novor_correctness_mask( + sequences: pd.Series | Iterable[object], + predictions: pd.Series | Iterable[object], + *, + residue_masses: dict[str, float] | None = None, + metrics: Metrics | None = None, +) -> np.ndarray: + """Vectorized Novor correctness (same rule as ``DatabaseGroundedFDRControl.fit``).""" + if metrics is None: + masses = residue_masses if residue_masses is not None else load_residue_masses() + metrics = metrics_from_residue_masses(masses) + return np.array( + [ + novor_correct_prediction(sequence, prediction, metrics) + for sequence, prediction in zip(sequences, predictions) + ], + dtype=bool, + ) + + +def dedupe_best_score_per_peptide( + df: pd.DataFrame, peptide_col: str, score_col: str +) -> pd.DataFrame: + """Keep the highest-scoring row per peptide key.""" + return ( + df.sort_values(score_col, ascending=False) + .groupby(peptide_col, as_index=False) + .first() + ) + + +def compute_q_values(fdr: np.ndarray) -> np.ndarray: + """Convert ranked FDR estimates to q-values using suffix minima.""" + values = np.asarray(fdr, dtype=float) + q_values = np.empty_like(values) + fdr_min = np.inf + for i in range(len(values) - 1, -1, -1): + fdr_min = min(fdr_min, values[i]) + q_values[i] = fdr_min + return q_values + + +def monotonize_q_by_confidence( + confidence: np.ndarray, q_value: np.ndarray +) -> np.ndarray: + """Enforce non-increasing q-values when confidence increases.""" + order = np.argsort(-np.asarray(confidence, dtype=float)) + q_sorted = np.asarray(q_value, dtype=float)[order] + q_mono = np.empty_like(q_sorted) + q_min = np.inf + for i in range(len(q_sorted) - 1, -1, -1): + if q_sorted[i] > q_min: + q_mono[i] = q_min + else: + q_mono[i] = q_sorted[i] + q_min = q_sorted[i] + out = np.empty_like(q_mono) + out[order] = q_mono + return out + + +def add_peptide_key( + df: pd.DataFrame, + peptide_col: str, + *, + key_col: str = "peptide_key", +) -> pd.DataFrame: + """Append normalised peptide keys.""" + work = df.copy() + work[key_col] = work[peptide_col].map(normalize_peptide_key) + return work + + +def filter_prediction_table( + df: pd.DataFrame, + peptide_col: str, + *, + min_length: int = MIN_PEPTIDE_LENGTH, + key_col: str = "peptide_key", + drop_unsupported_mods: bool = True, + log: bool = True, +) -> pd.DataFrame: + """Drop unsupported mods and peptides shorter than *min_length* (normalised).""" + work = add_peptide_key(df, peptide_col, key_col=key_col) + before = len(work) + if drop_unsupported_mods: + work = work[~work[peptide_col].map(has_unsupported_unimod)].copy() + work = work[work[key_col].str.len() >= min_length].copy() + dropped = before - len(work) + if log and dropped: + logger.info( + "Filtered %d/%d rows (unsupported mods and/or length < %d) on %s", + dropped, + before, + min_length, + peptide_col, + ) + return work.reset_index(drop=True) + + +def filter_novoboard_prediction_table( + df: pd.DataFrame, + *, + peptide_col: str = "Peptide", + min_length: int = MIN_PEPTIDE_LENGTH, + key_col: str = "_peptide_key", + log: bool = True, +) -> pd.DataFrame: + """Convert NovoBoard modifications to ProForma, then apply shared filters.""" + work = df.copy() + work[peptide_col] = work[peptide_col].map(novoboard_to_proforma) + return filter_prediction_table( + work, + peptide_col, + min_length=min_length, + key_col=key_col, + log=log, + ) + + +def filter_novoboard_target_decoy_pairs( + target: pd.DataFrame, + decoy: pd.DataFrame, + *, + peptide_col: str = "Peptide", + min_length: int = MIN_PEPTIDE_LENGTH, + key_col: str = "_peptide_key", + log: bool = True, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Filter NovoBoard target/decoy as spectrum twins so pair counts stay equal. + + Both sides are converted to ProForma and passed through the shared + mod/length filter. Only ``_pair_key`` values present in **both** filtered + tables are kept, so dropping an unsupported-mod decoy also drops its + target (and vice versa). + """ + if "_pair_key" not in target.columns or "_pair_key" not in decoy.columns: + raise ValueError( + "filter_novoboard_target_decoy_pairs requires '_pair_key' on both tables" + ) + + target_f = filter_novoboard_prediction_table( + target, + peptide_col=peptide_col, + min_length=min_length, + key_col=key_col, + log=log, + ) + decoy_f = filter_novoboard_prediction_table( + decoy, + peptide_col=peptide_col, + min_length=min_length, + key_col=key_col, + log=log, + ) + + def _valid_pair_keys(series: pd.Series) -> set[str]: + keys = series.astype(str) + return {k for k in keys if k and k != "nan"} + + shared_keys = _valid_pair_keys(target_f["_pair_key"]) & _valid_pair_keys( + decoy_f["_pair_key"] + ) + target_out = target_f[target_f["_pair_key"].astype(str).isin(shared_keys)].copy() + decoy_out = decoy_f[decoy_f["_pair_key"].astype(str).isin(shared_keys)].copy() + n_target = target_out["_pair_key"].astype(str).nunique() + n_decoy = decoy_out["_pair_key"].astype(str).nunique() + if n_target != n_decoy: + raise AssertionError( + f"Pair filter left unequal twin counts: targets={n_target} decoys={n_decoy}" + ) + if log: + before_pairs = len( + _valid_pair_keys(target_f["_pair_key"]) + | _valid_pair_keys(decoy_f["_pair_key"]) + ) + logger.info( + "NovoBoard pair filter: %s → %s twin spectra " + "(target rows %s → %s, decoy rows %s → %s)", + before_pairs, + n_target, + len(target_f), + len(target_out), + len(decoy_f), + len(decoy_out), + ) + return target_out, decoy_out + + +def restrict_winnow_to_novoboard_spectra( + winnow: pd.DataFrame, novoboard: pd.DataFrame +) -> pd.DataFrame: + """Trim Winnow to NovoBoard twin-valid spectra under the subset invariant. + + After shared peptide filters and NovoBoard pair-gating, NovoBoard targets are + expected to be a subset of Winnow-filtered spectra (same InstaNovo + predictions; NovoBoard additionally drops pairs whose decoy fails). The + shared pool is therefore the NovoBoard spectrum set: only Winnow is trimmed. + + Raises: + AssertionError: If any NovoBoard spectrum is missing from Winnow. + """ + winnow_ids = set(winnow["spectrum_id"].astype(str)) + novoboard_ids = set(novoboard["spectrum_id"].astype(str)) + only_novoboard = novoboard_ids - winnow_ids + if only_novoboard: + examples = sorted(only_novoboard)[:5] + raise AssertionError( + "NovoBoard twin-valid spectra are not a subset of Winnow-filtered " + f"spectra ({len(only_novoboard)} missing); examples={examples}. " + "Expected identical InstaNovo predictions after ProForma remapping." + ) + winnow_shared = winnow[winnow["spectrum_id"].astype(str).isin(novoboard_ids)].copy() + logger.info( + "Shared spectrum pool: %d spectra (trimmed Winnow=%d; NovoBoard unchanged)", + len(novoboard_ids), + len(winnow) - len(winnow_shared), + ) + return winnow_shared + + +def assert_shared_prediction_keys( + winnow: pd.DataFrame, + novoboard: pd.DataFrame, + *, + winnow_peptide_col: str = "prediction", + novoboard_peptide_col: str = "Peptide", +) -> None: + """Require I/L-normalised prediction identity on the shared spectrum pool.""" + w_ids = winnow["spectrum_id"].astype(str) + nb_ids = novoboard["spectrum_id"].astype(str) + if w_ids.nunique() != len(winnow) or nb_ids.nunique() != len(novoboard): + raise AssertionError( + "Shared-pool tables must have one row per spectrum_id before " + f"prediction-key assert (winnow rows={len(winnow)} unique={w_ids.nunique()}, " + f"novoboard rows={len(novoboard)} unique={nb_ids.nunique()})" + ) + merged = ( + winnow[["spectrum_id", winnow_peptide_col]] + .assign(spectrum_id=w_ids) + .merge( + novoboard[["spectrum_id", novoboard_peptide_col]].assign( + spectrum_id=nb_ids + ), + on="spectrum_id", + how="inner", + validate="one_to_one", + ) + ) + if len(merged) != len(winnow) or len(merged) != len(novoboard): + raise AssertionError( + "Shared-pool spectrum_id join is not 1:1 " + f"(winnow={len(winnow)} novoboard={len(novoboard)} inner={len(merged)})" + ) + w_keys = merged[winnow_peptide_col].map(normalize_peptide_key) + nb_keys = merged[novoboard_peptide_col].map(normalize_peptide_key) + mismatch = w_keys != nb_keys + if bool(mismatch.any()): + bad = merged.loc[ + mismatch, ["spectrum_id", winnow_peptide_col, novoboard_peptide_col] + ] + examples = bad.head(5).to_dict(orient="records") + raise AssertionError( + "Winnow and NovoBoard predictions disagree after I/L-normalised " + f"peptide keys ({int(mismatch.sum())} spectra); examples={examples}" + ) + + +def label_series_by_spectrum_id(winnow: pd.DataFrame, label_col: str) -> pd.Series: + """Map ``spectrum_id`` → boolean label from a Winnow table.""" + if label_col not in winnow.columns: + raise KeyError(f"Missing label column {label_col!r}") + ids = winnow["spectrum_id"].astype(str) + if ids.duplicated().any(): + raise AssertionError( + f"Duplicate spectrum_id values when building {label_col} label map" + ) + return pd.Series( + winnow[label_col].astype(bool).to_numpy(), + index=ids, + name=label_col, + ) + + +def attach_labels_by_spectrum_id( + novoboard: pd.DataFrame, + label_by_id: pd.Series, + *, + label_col: str, +) -> pd.DataFrame: + """Attach a shared label column to NovoBoard rows by ``spectrum_id``.""" + out = novoboard.copy() + mapped = out["spectrum_id"].astype(str).map(label_by_id) + if mapped.isna().any(): + missing = out.loc[mapped.isna(), "spectrum_id"].astype(str).head(5).tolist() + raise AssertionError( + f"NovoBoard rows missing shared {label_col} labels; examples={missing}" + ) + out[label_col] = mapped.astype(bool) + return out + + +def _best_alc_per_pair_key(df: pd.DataFrame) -> pd.DataFrame: + """One highest-ALC row per ``_pair_key``.""" + work = df.dropna(subset=["ALC (%)", "_pair_key"]) + work = work[work["_pair_key"].astype(str) != "nan"] + return ( + work.sort_values("ALC (%)", ascending=False) + .groupby("_pair_key", as_index=False) + .first() + ) + + +def proteome_hit_mask( + peptides: pd.Series | Iterable[str], + haystack: str, + *, + min_length: int = MIN_PEPTIDE_LENGTH, +) -> np.ndarray: + """True when normalised peptide key (length ≥ *min_length*) hits the proteome.""" + keys = [normalize_peptide_key(p) for p in peptides] + eligible = [bool(k) and len(k) >= min_length for k in keys] + unique_keys = sorted({k for k, ok in zip(keys, eligible) if ok}) + hit_map: dict[str, bool] = {} + if unique_keys: + hits = _batch_peptide_substring_hits(unique_keys, haystack) + hit_map = dict(zip(unique_keys, hits)) + return np.array( + [bool(eligible[i] and hit_map.get(keys[i], False)) for i in range(len(keys))], + dtype=bool, + ) + + +def max_score_per_peptide( + df: pd.DataFrame, + key_col: str, + score_col: str, +) -> pd.DataFrame: + """Keep the max-scoring row per peptide key (no filtering). + + Call after :func:`filter_prediction_table` or + :func:`filter_novoboard_prediction_table` so all methods share the same + filter → max-dedupe sequence in the peptide score-mixture benchmark. + """ + work = df.dropna(subset=[score_col, key_col]) + work = work[work[key_col].astype(str) != ""] + return dedupe_best_score_per_peptide(work, key_col, score_col).reset_index( + drop=True + ) + + +def confidence_to_log_prob(confidence: pd.Series | np.ndarray) -> np.ndarray: + """Map raw InstaNovo confidence in (0, 1] to Glissade-style log probabilities.""" + conf = np.asarray(confidence, dtype=float) + return np.log(np.clip(conf, 1e-300, 1.0)) + + +def _load_mgf_title_to_scan(mgf_path: Path) -> dict[str, str]: + """Parse TITLE→SCANS mapping from an MGF file.""" + mapping: dict[str, str] = {} + title: str | None = None + scan: str | None = None + with open(mgf_path, encoding="utf-8", errors="replace") as handle: + for line in handle: + value = line.strip() + if value.startswith("TITLE="): + title = value.removeprefix("TITLE=") + elif value.startswith("SCANS="): + scan = value.removeprefix("SCANS=") + elif value == "END IONS" and title is not None and scan is not None: + mapping[title] = scan + title = None + scan = None + return mapping + + +def attach_novoboard_pair_keys( + target: pd.DataFrame, + decoy: pd.DataFrame, + *, + novoboard_dir: Path, + split_prefix: str, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Attach ``_pair_key`` using Scan identity or decoy-MGF TITLE→SCANS mapping.""" + target_out = target.copy() + decoy_out = decoy.copy() + if "Scan" not in target_out.columns or "Scan" not in decoy_out.columns: + raise ValueError("NovoBoard tables require a 'Scan' column for twin pairing") + + target_key = target_out["Scan"].astype(str) + decoy_key = decoy_out["Scan"].astype(str) + best_decoy_key = decoy_key + best_overlap = len(set(target_key) & set(decoy_key)) + + mgf_path = novoboard_dir.parent / f"{split_prefix}.mgf" + if mgf_path.is_file(): + title_to_scan = _load_mgf_title_to_scan(mgf_path) + if title_to_scan: + mapped = decoy_key.map(title_to_scan) + mapped_overlap = len(set(target_key) & set(mapped.dropna())) + if mapped_overlap > best_overlap: + best_decoy_key = mapped + best_overlap = mapped_overlap + + target_out["_pair_key"] = target_key + decoy_out["_pair_key"] = best_decoy_key + logger.info( + "NovoBoard %s pair-key overlap: target=%d decoy=%d overlap=%d", + split_prefix, + target_key.nunique(dropna=True), + pd.Series(best_decoy_key).nunique(dropna=True), + best_overlap, + ) + return target_out, decoy_out + + +def prepare_novoboard_decoy_by_pair( + decoy_df: pd.DataFrame, + *, + min_length: int = MIN_PEPTIDE_LENGTH, + already_filtered: bool = False, +) -> pd.DataFrame: + """Index the best decoy row per ``_pair_key``. + + Args: + decoy_df: Decoy table with ``_pair_key``. Prefer pair-gated output from + :func:`filter_novoboard_target_decoy_pairs`. + already_filtered: When True, skip ProForma/mod/length filtering (caller + already pair-filtered). + """ + if "_pair_key" not in decoy_df.columns: + raise ValueError("NovoBoard twin TDC requires '_pair_key' on decoy") + if already_filtered: + decoy = decoy_df.copy() + if "_peptide_key" not in decoy.columns: + if "peptide_key" in decoy.columns: + decoy["_peptide_key"] = decoy["peptide_key"] + else: + decoy = add_peptide_key(decoy, "Peptide", key_col="_peptide_key") + elif "_peptide_key" in decoy_df.columns and decoy_df["_peptide_key"].notna().all(): + decoy = decoy_df.copy() + else: + decoy = filter_novoboard_prediction_table( + decoy_df, min_length=min_length, key_col="_peptide_key" + ) + decoy = decoy.dropna(subset=["ALC (%)", "_pair_key"]) + decoy = decoy[ + (decoy["_pair_key"].astype(str) != "nan") & (decoy["_peptide_key"] != "") + ] + return _best_alc_per_pair_key(decoy).set_index("_pair_key", drop=False) + + +def novoboard_psm_tdc( + target_df: pd.DataFrame, + decoy_df: pd.DataFrame, + *, + min_length: int = MIN_PEPTIDE_LENGTH, +) -> pd.DataFrame: + """Recompute NovoBoard's pooled PSM TDC after pair-gated filtering. + + Target and decoy are filtered as spectrum twins so unsupported-mod drops + remove the pair. Competition uses one best-ALC row per twin on each side, + guaranteeing ``sum(is_target) == sum(~is_target)``. + """ + target, decoy = filter_novoboard_target_decoy_pairs( + target_df, decoy_df, min_length=min_length + ) + target = _best_alc_per_pair_key(target).assign(is_target=True) + decoy = _best_alc_per_pair_key(decoy).assign(is_target=False) + n_target = int(target["_pair_key"].nunique()) + n_decoy = int(decoy["_pair_key"].nunique()) + if n_target != n_decoy or len(target) != len(decoy): + raise AssertionError( + f"PSM TDC unbalanced after pair gate: " + f"target_rows={len(target)} decoy_rows={len(decoy)} " + f"target_pairs={n_target} decoy_pairs={n_decoy}" + ) + + combined = pd.concat([target, decoy], ignore_index=True, sort=False) + combined = combined.sort_values( + ["ALC (%)", "is_target"], ascending=[False, False] + ).reset_index(drop=True) + return _assign_cumulative_tdc_fdr(combined) + + +def _prepare_targets_for_twin_tdc( + target_df: pd.DataFrame, + decoy_df: pd.DataFrame, + *, + min_length: int, + decoy_by_pair: pd.DataFrame | None, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Pair-gate or twin-filter targets and return ``(target, decoy_by_pair)``.""" + if "_pair_key" not in target_df.columns: + raise ValueError("NovoBoard twin TDC requires '_pair_key' on target") + if decoy_by_pair is None and "_pair_key" not in decoy_df.columns: + raise ValueError("NovoBoard twin TDC requires '_pair_key' on decoy") + + if decoy_by_pair is None: + target, decoy = filter_novoboard_target_decoy_pairs( + target_df, decoy_df, min_length=min_length, log=False + ) + decoy_by_pair = prepare_novoboard_decoy_by_pair( + decoy, min_length=min_length, already_filtered=True + ) + return target, decoy_by_pair + + target = target_df.copy() + if "_peptide_key" not in target.columns: + if "peptide_key" in target.columns: + target["_peptide_key"] = target["peptide_key"] + else: + target = filter_novoboard_prediction_table( + target, + min_length=min_length, + key_col="_peptide_key", + log=False, + ) + twin_keys = set(decoy_by_pair.index.astype(str)) + target = target[target["_pair_key"].astype(str).isin(twin_keys)] + return target, decoy_by_pair + + +def _assign_cumulative_tdc_fdr(combined: pd.DataFrame) -> pd.DataFrame: + """Add estimated FDR / q-value columns for a balanced target-decoy table.""" + out = combined.copy() + n_target = out["is_target"].astype(int).cumsum() + n_decoy = (~out["is_target"]).astype(int).cumsum() + out["estimated_fdr"] = np.divide( + n_decoy, + n_target, + out=np.ones(len(out), dtype=float), + where=n_target > 0, + ) + out["estimated_q_value"] = np.nan + target_mask = out["is_target"].to_numpy() + out.loc[target_mask, "estimated_q_value"] = compute_q_values( + out.loc[target_mask, "estimated_fdr"].to_numpy() + ) + return out + + +def novoboard_max_target_twin_decoy_tdc( + target_df: pd.DataFrame, + decoy_df: pd.DataFrame, + *, + min_length: int = MIN_PEPTIDE_LENGTH, + target_peptide_keys: set[str] | None = None, + decoy_by_pair: pd.DataFrame | None = None, + log_missing_twins: bool = True, +) -> pd.DataFrame: + """Peptide TDC: pair-gate, max ALC per target peptide, twin decoy by ``_pair_key``. + + Returns the combined ranked table with ``is_target``, ``estimated_fdr``, and + ``estimated_q_value`` (targets only). Targets without a twin-valid decoy are + dropped; the competition table is always 1:1. + + Args: + decoy_by_pair: Optional precomputed output of + :func:`prepare_novoboard_decoy_by_pair` from an already pair-gated + decoy table. When omitted, target/decoy are pair-filtered together. + """ + target, decoy_by_pair = _prepare_targets_for_twin_tdc( + target_df, + decoy_df, + min_length=min_length, + decoy_by_pair=decoy_by_pair, + ) + target = target.dropna(subset=["ALC (%)", "_pair_key"]) + target = target[ + (target["_pair_key"].astype(str) != "nan") & (target["_peptide_key"] != "") + ] + + if target_peptide_keys is not None: + target = target[target["_peptide_key"].isin(target_peptide_keys)] + + # Max-score only among twin-valid targets. + target_best = max_score_per_peptide(target, "_peptide_key", "ALC (%)") + pair_keys = target_best["_pair_key"].astype(str) + has_twin = pair_keys.isin(decoy_by_pair.index.astype(str)) + n_missing_twin = int((~has_twin).sum()) + if log_missing_twins and n_missing_twin: + logger.warning( + "NovoBoard twin-decoy TDC dropped %d/%d max-target peptides without twin", + n_missing_twin, + len(target_best), + ) + target_keep = target_best.loc[has_twin].copy() + if target_keep.empty: + return pd.DataFrame( + columns=[ + "spectrum_id", + "Peptide", + "ALC (%)", + "_peptide_key", + "_pair_key", + "is_target", + "estimated_fdr", + "estimated_q_value", + ] + ) + + decoy_keep = decoy_by_pair.loc[target_keep["_pair_key"].astype(str)].copy() + decoy_keep = decoy_keep.reset_index(drop=True) + target_keep = target_keep.assign(is_target=True).reset_index(drop=True) + decoy_keep = decoy_keep.assign(is_target=False) + if len(target_keep) != len(decoy_keep): + raise AssertionError( + f"Peptide TDC unbalanced: targets={len(target_keep)} decoys={len(decoy_keep)}" + ) + # Preserve 1:1 balance: one decoy row per retained target (do not dedupe decoys). + combined = pd.concat([target_keep, decoy_keep], ignore_index=True, sort=False) + combined = combined.sort_values( + ["ALC (%)", "is_target"], ascending=[False, False] + ).reset_index(drop=True) + return _assign_cumulative_tdc_fdr(combined) diff --git a/scripts/fdr_tool_comparison_summaries.py b/scripts/fdr_tool_comparison_summaries.py new file mode 100644 index 00000000..b48d1dcc --- /dev/null +++ b/scripts/fdr_tool_comparison_summaries.py @@ -0,0 +1,533 @@ +"""Summary tables for Winnow / NovoBoard / Glissade FDR tool comparisons. + +Produces two long-form CSVs: + +- ``*_acceptance.csv``: accepted counts and recovery at q-value thresholds. +- ``*_error_gain.csv``: observed FDP, excess over nominal FDR, optional mean + absolute q-value deviation vs a database-grounded reference (calibrated-score + DBG for Winnow; raw-score / ALC DBG for NovoBoard and Glissade), and relative + gain/loss of a primary method vs each comparator. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Iterable, Sequence, cast + +import numpy as np +import pandas as pd + +from scripts.fdr_tool_comparison_preprocess import compute_q_values + +logger = logging.getLogger(__name__) + +SUMMARY_THRESHOLDS: list[float] = [0.01, 0.05, 0.10] +# Match ``DatabaseGroundedFDRControl`` / PSM comparison default. +_DB_GROUNDED_DROP = 10 + +_KEY_COLS = ["dataset", "panel", "level", "method", "q_value_threshold"] + + +def _slug_comparator(name: str) -> str: + """Map a method label to a filesystem-/column-safe slug.""" + return ( + name.lower() + .replace(" ", "_") + .replace("(", "") + .replace(")", "") + .replace("-", "_") + ) + + +def acceptance_rows_from_q( + *, + dataset: str, + panel: str, + level: str, + method: str, + q_value: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, + label_mask: np.ndarray | None = None, + recovery_denom: int | None = None, +) -> list[dict[str, object]]: + """Build acceptance/yield rows for one method at each q-value threshold. + + Args: + dataset: Dataset key (e.g. ``helaqc``). + panel: Evaluation panel (e.g. ``labelled_test``, ``unlabelled``, ``external``). + level: ``psm`` or ``peptide``. + method: Method display label. + q_value: Per-row estimated q-values. + thresholds: Nominal FDR thresholds. + label_mask: Optional boolean correctness / proteome-hit labels aligned with + ``q_value``. When provided, ``n_correct`` is filled. + recovery_denom: Denominator for recovery percentage. Defaults to the number + of True labels when ``label_mask`` is given. + + Returns: + One dict per threshold. + """ + q = np.asarray(q_value, dtype=float) + labels = None if label_mask is None else np.asarray(label_mask, dtype=bool) + if labels is not None and len(labels) != len(q): + raise ValueError( + f"label_mask length {len(labels)} does not match q_value length {len(q)}" + ) + if recovery_denom is None and labels is not None: + recovery_denom = int(labels.sum()) + + rows: list[dict[str, object]] = [] + for threshold in thresholds: + valid = ~np.isnan(q) + accepted = valid & (q <= threshold) + n_accepted = int(accepted.sum()) + n_correct: float | int = np.nan + recovery_pct: float = np.nan + if labels is not None: + n_correct = int((accepted & labels).sum()) + if recovery_denom and recovery_denom > 0: + recovery_pct = 100.0 * float(n_correct) / float(recovery_denom) + rows.append( + { + "dataset": dataset, + "panel": panel, + "level": level, + "method": method, + "q_value_threshold": float(threshold), + "n_accepted": n_accepted, + "n_correct": n_correct, + "recovery_pct": recovery_pct, + } + ) + return rows + + +def observed_fdp_at_thresholds( + q_value: np.ndarray, + label_mask: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, +) -> list[float]: + """Return observed false-discovery proportion among accepted rows at each threshold. + + ``label_mask`` is True for correct (or proteome-hit) rows. Observed FDP is + ``1 - n_correct / n_accepted`` when any rows are accepted, else NaN. + """ + q = np.asarray(q_value, dtype=float) + labels = np.asarray(label_mask, dtype=bool) + if len(q) != len(labels): + raise ValueError( + f"label_mask length {len(labels)} does not match q_value length {len(q)}" + ) + out: list[float] = [] + for threshold in thresholds: + valid = ~np.isnan(q) + accepted = valid & (q <= threshold) + n_accepted = int(accepted.sum()) + if n_accepted == 0: + out.append(float("nan")) + continue + n_correct = int((accepted & labels).sum()) + out.append(1.0 - n_correct / n_accepted) + return out + + +def database_grounded_q_from_labels( + scores: np.ndarray, + labels: np.ndarray, + *, + drop: int = _DB_GROUNDED_DROP, +) -> np.ndarray: + """In-sample database-grounded q-values from ranked scores and boolean labels. + + Builds the empirical precision curve ``1 - cumsum(correct) / rank`` on scores + sorted descending (same construction as the proteome-hit shortcut in the PSM + comparison), drops the first *drop* ranks from the FDR map, assigns FDR by + score lookup, then converts to q-values. + + Args: + scores: Ranking scores (higher = more confident). + labels: Boolean correctness / hit labels aligned with *scores*. + drop: Leading ranks excluded from the FDR map (default 10). + + Returns: + q-value array aligned with *scores*. + """ + scores_a = np.asarray(scores, dtype=float) + labels_a = np.asarray(labels, dtype=bool) + n = len(scores_a) + if n == 0: + return np.asarray([], dtype=float) + if len(labels_a) != n: + raise ValueError( + f"labels length {len(labels_a)} does not match scores length {n}" + ) + + order = np.argsort(-scores_a, kind="mergesort") + precision = np.cumsum(labels_a[order].astype(float)) / np.arange(1, n + 1) + fdr_ranked = 1.0 - precision + drop_eff = min(drop, max(0, n - 1)) + fit_scores = scores_a[order][drop_eff:] + fit_fdr = fdr_ranked[drop_eff:] + n_fit = len(fit_scores) + + idx = np.searchsorted(-fit_scores, -scores_a, side="left") + fdr = np.empty(n, dtype=float) + below = (idx == n_fit) & (scores_a < fit_scores[-1]) + above = (idx == 0) & (scores_a > fit_scores[0]) + normal = ~(below | above) + fdr[below] = 1.0 + fdr[above] = float(fit_fdr[0]) + fdr[normal] = fit_fdr[np.clip(idx[normal], 0, n_fit - 1)] + + q_sorted = compute_q_values(fdr[order]) + q = np.empty(n, dtype=float) + q[order] = q_sorted + return q + + +def mean_abs_q_dev_vs_reference( + q_method: np.ndarray, + q_ref: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, +) -> list[float]: + """Mean absolute q-value deviation vs a row-aligned reference at each threshold. + + For each threshold, restrict to rows accepted by either method + (``q_method <= t`` or ``q_ref <= t``) with finite q for both, then report + ``mean(|q_method - q_ref|)``. + """ + q_m = np.asarray(q_method, dtype=float) + q_r = np.asarray(q_ref, dtype=float) + if len(q_m) != len(q_r): + raise ValueError( + f"q_ref length {len(q_r)} does not match q_method length {len(q_m)}" + ) + out: list[float] = [] + for threshold in thresholds: + both_finite = ~np.isnan(q_m) & ~np.isnan(q_r) + either_accepted = both_finite & ((q_m <= threshold) | (q_r <= threshold)) + if not np.any(either_accepted): + out.append(float("nan")) + continue + out.append(float(np.mean(np.abs(q_m[either_accepted] - q_r[either_accepted])))) + return out + + +def error_rows_from_q( + *, + dataset: str, + panel: str, + level: str, + method: str, + q_value: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, + label_mask: np.ndarray | None = None, + q_ref: np.ndarray | None = None, + observed_fdp: Sequence[float] | None = None, +) -> list[dict[str, object]]: + """Build error-metric rows for one method (without relative-gain columns). + + Args: + observed_fdp: Optional precomputed FDP values (e.g. from a mixture + benchmark). When omitted, FDP is derived from ``label_mask`` if given. + """ + if observed_fdp is not None and len(observed_fdp) != len(thresholds): + raise ValueError("observed_fdp length must match thresholds") + if observed_fdp is None and label_mask is not None: + fdp_values = observed_fdp_at_thresholds(q_value, label_mask, thresholds) + elif observed_fdp is not None: + fdp_values = [float(x) for x in observed_fdp] + else: + fdp_values = [float("nan")] * len(thresholds) + + if q_ref is not None: + q_dev = mean_abs_q_dev_vs_reference(q_value, q_ref, thresholds) + else: + q_dev = [float("nan")] * len(thresholds) + + rows: list[dict[str, object]] = [] + for threshold, fdp, dev in zip(thresholds, fdp_values, q_dev): + fdp_f = float(fdp) + excess = fdp_f - float(threshold) if np.isfinite(fdp_f) else float("nan") + rows.append( + { + "dataset": dataset, + "panel": panel, + "level": level, + "method": method, + "q_value_threshold": float(threshold), + "observed_fdp": fdp_f, + "fdp_excess": excess, + "mean_abs_q_dev_vs_db": float(dev), + } + ) + return rows + + +def _relative_gain_column(value_col: str, comparator: str) -> str: + """Return the plan-specified relative-gain column name for *value_col*.""" + slug = _slug_comparator(comparator) + if value_col == "n_accepted": + return f"accepted_pct_vs_{slug}" + if value_col == "recovery_pct": + return f"recovery_pct_vs_{slug}" + if value_col == "observed_fdp": + return f"fdp_delta_vs_{slug}" + if "fdp" in value_col: + return f"{value_col}_delta_vs_{slug}" + return f"{value_col}_pct_vs_{slug}" + + +def _relative_gain_value( + value_col: str, primary_val: object, comparator_val: object +) -> float: + """Compute primary-vs-comparator gain for one metric cell.""" + if pd.isna(primary_val) or pd.isna(comparator_val): + return float("nan") + primary = float(cast("float | int | str", primary_val)) + comparator = float(cast("float | int | str", comparator_val)) + if value_col == "observed_fdp" or "fdp" in value_col: + return primary - comparator + if comparator == 0: + return float("nan") + return 100.0 * (primary - comparator) / comparator + + +def _fill_primary_relative_gains( + work: pd.DataFrame, + group: pd.DataFrame, + *, + primary_method: str, + comparators: Sequence[str], + value_cols: Sequence[str], + group_cols: Sequence[str], + group_keys: tuple[object, ...], +) -> None: + """Write relative-gain columns onto the primary-method row for one group.""" + primary_rows = group[group["method"] == primary_method] + if primary_rows.empty: + return + primary = primary_rows.iloc[0] + mask = pd.Series(True, index=work.index) + for col, val in zip(group_cols, group_keys): + mask &= work[col] == val + primary_idx = work.index[mask & (work["method"] == primary_method)] + if len(primary_idx) == 0: + return + idx = primary_idx[0] + + for comparator in comparators: + comp_rows = group[group["method"] == comparator] + if comp_rows.empty: + continue + comp = comp_rows.iloc[0] + for col in value_cols: + out_col = _relative_gain_column(col, comparator) + gain = _relative_gain_value(col, primary[col], comp[col]) + if np.isfinite(gain): + work.at[idx, out_col] = gain + + +def add_relative_gain_columns( + df: pd.DataFrame, + *, + primary_method: str, + comparators: Iterable[str], + value_cols: Sequence[str], + group_cols: Sequence[str] = ("dataset", "panel", "level", "q_value_threshold"), +) -> pd.DataFrame: + """Attach primary-vs-comparator relative columns onto a long-form metrics table. + + For ``n_accepted`` / ``recovery_pct``, writes + ``100 * (primary - comparator) / comparator``. + For ``observed_fdp``, writes the signed difference ``primary - comparator``. + Relative columns are filled only on primary-method rows. + """ + if df.empty: + return df.copy() + + work = df.copy() + comparator_list = list(comparators) + for col in value_cols: + for comparator in comparator_list: + work[_relative_gain_column(col, comparator)] = np.nan + + group_list = list(group_cols) + for keys, group in work.groupby(group_list, dropna=False, sort=False): + if not isinstance(keys, tuple): + keys = (keys,) + _fill_primary_relative_gains( + work, + group, + primary_method=primary_method, + comparators=comparator_list, + value_cols=value_cols, + group_cols=group_list, + group_keys=keys, + ) + return work + + +def merge_acceptance_and_error( + acceptance: pd.DataFrame, + error: pd.DataFrame, + *, + key_cols: Sequence[str] | None = None, +) -> pd.DataFrame: + """Join acceptance counts onto error rows for relative-gain construction.""" + if acceptance.empty or error.empty: + return error.copy() + keys = list(key_cols) if key_cols is not None else list(_KEY_COLS) + keys = [c for c in keys if c in acceptance.columns and c in error.columns] + cols = [ + c + for c in ("n_accepted", "n_correct", "recovery_pct") + if c in acceptance.columns + ] + return error.merge( + acceptance[keys + cols], + on=keys, + how="left", + ) + + +def finalise_error_gain_table( + acceptance: pd.DataFrame, + error: pd.DataFrame, + *, + primary_method: str, + comparators: Sequence[str], + key_cols: Sequence[str] | None = None, + group_cols: Sequence[str] | None = None, +) -> pd.DataFrame: + """Merge counts into error rows and add primary-vs-comparator relative columns.""" + merged = merge_acceptance_and_error(acceptance, error, key_cols=key_cols) + value_cols = [ + c for c in ("n_accepted", "recovery_pct", "observed_fdp") if c in merged.columns + ] + gain_groups = ( + tuple(group_cols) + if group_cols is not None + else ("dataset", "panel", "level", "q_value_threshold") + ) + with_gain = add_relative_gain_columns( + merged, + primary_method=primary_method, + comparators=comparators, + value_cols=value_cols, + group_cols=gain_groups, + ) + # Keep error-table identity columns first; drop helper count cols that duplicate + # the acceptance table except when used only for gain calculation. + drop_helpers = [c for c in ("n_accepted", "n_correct") if c in with_gain.columns] + return with_gain.drop(columns=drop_helpers, errors="ignore") + + +def write_summary_tables( + acceptance_df: pd.DataFrame, + error_df: pd.DataFrame, + output_dir: Path, + stem: str, +) -> tuple[Path, Path]: + """Write acceptance and error/gain CSVs under *output_dir*.""" + output_dir.mkdir(parents=True, exist_ok=True) + acceptance_path = output_dir / f"{stem}_acceptance.csv" + error_path = output_dir / f"{stem}_error_gain.csv" + acceptance_df.to_csv(acceptance_path, index=False) + error_df.to_csv(error_path, index=False) + logger.info("Wrote %s", acceptance_path) + logger.info("Wrote %s", error_path) + return acceptance_path, error_path + + +def summarise_holdout_results( + raw: pd.DataFrame, + *, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, + primary_method: str = "Winnow", + comparators: Sequence[str] = ("NovoBoard", "Glissade"), + panel: str = "score_mixture", + level: str = "peptide", + group_extra: Sequence[str] = (), +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Aggregate mixture-benchmark iterations into the two summary tables. + + Args: + raw: Per-iteration rows from ``external_peptide_holdout_results.csv``. + thresholds: Nominal FDR thresholds to retain. + primary_method: Method used for relative gain/loss columns. + comparators: Comparator method labels. + panel: Panel name written into the summary tables. + level: Identification level written into the summary tables. + group_extra: Extra columns to group by (e.g. ``pi0_target``). + + Returns: + ``(acceptance_df, error_gain_df)``. + """ + if raw.empty: + return pd.DataFrame(), pd.DataFrame() + + filtered = raw[raw["q_value_threshold"].isin(thresholds)].copy() + if filtered.empty: + return pd.DataFrame(), pd.DataFrame() + + extra = [c for c in group_extra if c in filtered.columns] + group_cols = ["dataset", *extra, "method", "q_value_threshold"] + gain_group_cols = ("dataset", *extra, "panel", "level", "q_value_threshold") + agg_kwargs: dict[str, tuple[str, str]] = { + "n_accepted": ("accepted_peptides", "mean"), + "n_correct": ("true_correct_peptides", "mean"), + "recovery_pct": ("correct_discovery_pct", "mean"), + "observed_fdp": ("observed_fdp", "mean"), + "n_accepted_std": ("accepted_peptides", "std"), + "observed_fdp_std": ("observed_fdp", "std"), + "recovery_pct_std": ("correct_discovery_pct", "std"), + } + if "mean_abs_q_dev_vs_db" in filtered.columns: + agg_kwargs["mean_abs_q_dev_vs_db"] = ("mean_abs_q_dev_vs_db", "mean") + agg = ( + filtered.groupby(group_cols, as_index=False) + .agg(**agg_kwargs) + .sort_values(group_cols) + .reset_index(drop=True) + ) + agg["panel"] = panel + agg["level"] = level + agg["fdp_excess"] = agg["observed_fdp"] - agg["q_value_threshold"] + if "mean_abs_q_dev_vs_db" not in agg.columns: + agg["mean_abs_q_dev_vs_db"] = np.nan + + id_cols = ["dataset", *extra, "panel", "level", "method", "q_value_threshold"] + acceptance = agg[ + id_cols + + [ + "n_accepted", + "n_correct", + "recovery_pct", + "n_accepted_std", + "recovery_pct_std", + ] + ].copy() + + error = agg[ + id_cols + + [ + "observed_fdp", + "fdp_excess", + "mean_abs_q_dev_vs_db", + "observed_fdp_std", + ] + ].copy() + + error_gain = finalise_error_gain_table( + acceptance.drop( + columns=["n_accepted_std", "recovery_pct_std"], errors="ignore" + ), + error, + primary_method=primary_method, + comparators=comparators, + key_cols=id_cols, + group_cols=gain_group_cols, + ) + return acceptance, error_gain diff --git a/scripts/feature_subsets.py b/scripts/feature_subsets.py new file mode 100644 index 00000000..34712b4d --- /dev/null +++ b/scripts/feature_subsets.py @@ -0,0 +1,74 @@ +"""Feature column sets for subset calibrator training and evaluation.""" + +from __future__ import annotations + +from typing import TypedDict + + +class FeatureSubsetSpec(TypedDict): + """Metadata and column list for one feature-subset experiment.""" + + description: str + from_parquet: bool + columns: list[str] + + +# Full feature matrix columns (``confidence`` + features + ``correct`` label). +FULL_FEATURE_COLUMNS: list[str] = [ + "confidence", + "mass_error_ppm", + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", + "irt_error", + "margin", + "median_margin", + "entropy", + "z-score", + "edit_distance", + "min_token_probability", + "std_token_probability", +] + +_NO_XCORR_SPECTRAL = {"spectral_angle", "xcorr"} +_NO_FRAGMENT_SIMILARITY = _NO_XCORR_SPECTRAL | { + "complementary_ion_count", + "max_ion_gap", + "edit_distance", +} + + +def _columns_excluding(*, drop: set[str]) -> list[str]: + return [c for c in FULL_FEATURE_COLUMNS if c not in drop] + + +FEATURE_SUBSETS: dict[str, FeatureSubsetSpec] = { + "no_xcorr_spectral": { + "description": "Exclude spectral_angle and xcorr only.", + "from_parquet": True, + "columns": _columns_excluding(drop=_NO_XCORR_SPECTRAL), + }, + "no_fragment_similarity": { + "description": ( + "Exclude spectral_angle, xcorr, complementary_ion_count, " + "max_ion_gap, and edit_distance." + ), + "from_parquet": True, + "columns": _columns_excluding(drop=_NO_FRAGMENT_SIMILARITY), + }, + "mass_error_da_no_similarity": { + "description": ( + "Exclude mass_error_ppm and fragment-similarity features; " + "use mass_error_da (Daltons) instead. Requires full winnow train " + "(recomputes features from raw spectra)." + ), + "from_parquet": False, + "columns": list( + _columns_excluding(drop=_NO_FRAGMENT_SIMILARITY | {"mass_error_ppm"}) + ) + + ["mass_error_da"], + }, +} diff --git a/scripts/plot_ablation_summary.py b/scripts/plot_ablation_summary.py new file mode 100644 index 00000000..862d1829 --- /dev/null +++ b/scripts/plot_ablation_summary.py @@ -0,0 +1,641 @@ +#!/usr/bin/env python3 +"""Bar charts of ablation calibration metrics from ``ablation_summary.csv``. + +Designed for publication main text: tail ECE at FDR operating points (and optionally +Brier) per feature-group config, with a reference line at the full ``All features`` model. +""" + +from __future__ import annotations + +import json +import logging +import sys +from pathlib import Path +from typing import Annotated, Literal + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.plot_eval_results import ( # noqa: E402 + _display_name, + _fit_database_grounded_fdr, + _save_fig, + _style_ax, +) +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +# Paul Tol qualitative palette (colour-blind safe) — canonical ablation colours. +_ABLATION_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#EE7733", + "#0077BB", + "#33BBEE", + "#CC3311", +] + +# Ablation summary keys → ``plot_eval_results.DATASET_DISPLAY_NAMES`` keys. +_ABLATION_DATASET_KEYS: dict[str, str] = { + "Arabidopsis": "01747_C01_P018218_S00_I00_N03_R1", + "Astral": "astral", + "HCT116": "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46", +} + +logger = logging.getLogger(__name__) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +ABLATION_CONFIG_ORDER: list[str] = [ + "Confidence only", + "Confidence + mass error", + "Confidence + iRT error", + "Confidence + token-level", + "Confidence + beam search", + "Confidence + fragment matching", + "All features", +] + + +def assign_ablation_colors(config_names: list[str]) -> dict[str, str]: + """Assign a unique colour per ablation config (no palette wrap).""" + if len(config_names) > len(_ABLATION_PALETTE): + raise ValueError( + f"Need {len(config_names)} ablation colours but only " + f"{len(_ABLATION_PALETTE)} defined." + ) + return {name: _ABLATION_PALETTE[i] for i, name in enumerate(config_names)} + + +def ordered_ablation_configs(present: set[str]) -> list[str]: + """Canonical config order for ablation plots and colour assignment.""" + ordered = [c for c in ABLATION_CONFIG_ORDER if c in present] + extra = sorted(present - set(ordered)) + return ordered + extra + + +_CONFIG_SHORT_LABELS: dict[str, str] = { + "Confidence only": "Confidence", + "Confidence + mass error": "+ Mass", + "Confidence + iRT error": "+ iRT", + "Confidence + token-level": "+ Token", + "Confidence + beam search": "+ Beam", + "Confidence + fragment matching": "+ Fragment", + "All features": "All features", +} + +MetricName = Literal[ + "tail_ECE@5%FDR", + "tail_ECE@10%FDR", + "ECE", + "Brier", + "PR_AUC", + "fdr_bias@5%FDR", + "fdr_bias@10%FDR", + "q_dev@5%FDR", + "q_dev@10%FDR", +] + +FDR_TAIL_THRESHOLDS: tuple[float, ...] = (0.05, 0.10) +TAIL_ECE_COLUMN_BY_THRESHOLD: dict[float, str] = { + 0.05: "tail_ECE@5%FDR", + 0.10: "tail_ECE@10%FDR", +} +Q_DEV_COLUMN_BY_THRESHOLD: dict[float, str] = { + 0.05: "q_dev@5%FDR", + 0.10: "q_dev@10%FDR", +} +FDR_BIAS_COLUMN_BY_THRESHOLD: dict[float, str] = { + 0.05: "fdr_bias@5%FDR", + 0.10: "fdr_bias@10%FDR", +} + +_DEFAULT_SUMMARY = ( + Path.home() / "Documents/winnow/new_eval_sets_plots/ablations/ablation_summary.csv" +) + + +def load_ablation_summary(path: Path) -> pd.DataFrame: + """Load ``ablation_summary.csv`` or ``.json``.""" + if not path.is_file(): + raise FileNotFoundError(path) + if path.suffix == ".json": + with open(path) as f: + return pd.DataFrame(json.load(f)) + return pd.read_csv(path) + + +def compute_ece( + pred: np.ndarray, + labels: np.ndarray, + n_bins: int = 10, +) -> float: + """Expected calibration error.""" + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_indices = np.digitize(pred, bins) - 1 + bin_indices = np.clip(bin_indices, 0, n_bins - 1) + ece = 0.0 + for b in range(n_bins): + mask = bin_indices == b + if mask.sum() == 0: + continue + avg_conf = pred[mask].mean() + avg_acc = labels[mask].mean() + ece += mask.sum() / len(pred) * abs(avg_conf - avg_acc) + return float(ece) + + +def compute_tail_ece_at_fdr( + pred: np.ndarray, + labels: np.ndarray, + fdr_threshold: float, + *, + fdr_ctrl: NonParametricFDRControl | None = None, + n_bins: int = 10, +) -> float: + """ECE among PSMs accepted at a non-parametric FDR threshold.""" + if len(pred) == 0: + return float("nan") + + if fdr_ctrl is None: + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=pd.Series(pred, name="score")) + + cutoff = fdr_ctrl.get_confidence_cutoff(threshold=fdr_threshold) + if np.isnan(cutoff): + return float("nan") + + mask = pred >= cutoff + if not mask.any(): + return float("nan") + + return compute_ece(pred[mask], labels[mask], n_bins=n_bins) + + +def compute_tail_ece_at_fdr_thresholds( + df: pd.DataFrame, + *, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", + fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, +) -> dict[float, float]: + """Tail ECE at each non-parametric FDR operating point.""" + work = df[[confidence_col, label_col]].dropna() + if work.empty: + return {threshold: float("nan") for threshold in fdr_thresholds} + + pred = work[confidence_col].to_numpy(dtype=float) + labels = work[label_col].to_numpy(dtype=float) + + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=work[confidence_col]) + + return { + threshold: compute_tail_ece_at_fdr( + pred, + labels, + threshold, + fdr_ctrl=fdr_ctrl, + ) + for threshold in fdr_thresholds + } + + +def compute_fdr_bias_at_fdr_thresholds( + df: pd.DataFrame, + *, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", + fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, +) -> dict[float, float]: + """Signed FDR bias at each NP-FDR cutoff; equal to empirical sTECE.""" + work = df[[confidence_col, label_col]].dropna() + if work.empty: + return {threshold: float("nan") for threshold in fdr_thresholds} + + scores = work[confidence_col].to_numpy(dtype=float) + labels = work[label_col].to_numpy(dtype=float) + + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=work[confidence_col]) + + results: dict[float, float] = {} + for threshold in fdr_thresholds: + cutoff = fdr_ctrl.get_confidence_cutoff(threshold=threshold) + if np.isnan(cutoff): + results[threshold] = float("nan") + continue + mask = scores >= cutoff + if not mask.any(): + results[threshold] = float("nan") + continue + + # E[1-S | S>=tau] - E[1-Y | S>=tau] = E[Y-S | S>=tau]. + results[threshold] = float(np.mean(labels[mask] - scores[mask])) + return results + + +def compute_pr_auc( + df: pd.DataFrame, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", +) -> float: + """Area under the ablation PR curve (matches ``run_feature_ablations`` plots).""" + work = df[[confidence_col, label_col]].dropna() + if work.empty: + return float("nan") + + sorted_data = work.sort_values(by=confidence_col, ascending=False) + cum_correct = np.cumsum(sorted_data[label_col].values) + precision = cum_correct / np.arange(1, len(sorted_data) + 1) + total_correct = cum_correct[-1] if len(cum_correct) else 0 + if total_correct <= 0 or len(precision) < 2: + return 0.0 + + recall = cum_correct / total_correct + from sklearn.metrics import auc + + return float(auc(recall, precision)) + + +def _vectorized_psm_fdr( + scores: np.ndarray, + ctrl: NonParametricFDRControl, +) -> np.ndarray: + """Map confidence scores to PSM FDR using a fitted controller.""" + conf = np.asarray(ctrl._confidence_scores, dtype=float) + fdr = np.asarray(ctrl._fdr_values, dtype=float) + scores = np.asarray(scores, dtype=float) + idx = np.searchsorted(-conf, -scores, side="left") + idx = np.clip(idx, 0, max(len(fdr) - 1, 0)) + if len(fdr) == 0: + return np.ones_like(scores) + + out = fdr[idx] + below = (idx == len(conf)) & (scores < conf[-1]) + above = (idx == 0) & (scores > conf[0]) + out[below] = 1.0 + out[above] = fdr[0] + return out + + +def _vectorized_psm_q_values( + scores: np.ndarray, + ctrl: NonParametricFDRControl, +) -> np.ndarray: + """Assign PSM q-values without per-row ``compute_fdr`` calls.""" + row_fdr = _vectorized_psm_fdr(scores, ctrl) + order = np.argsort(-scores) + sorted_fdr = row_fdr[order] + q_sorted = np.empty_like(sorted_fdr) + fdr_min = np.inf + for i in range(len(sorted_fdr) - 1, -1, -1): + current = sorted_fdr[i] + if current > fdr_min: + q_sorted[i] = fdr_min + else: + q_sorted[i] = current + fdr_min = current + q_values = np.empty_like(q_sorted) + q_values[order] = q_sorted + return q_values + + +def compute_q_value_deviations( + df: pd.DataFrame, + *, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", + fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, +) -> dict[float, float]: + """Mean absolute q-value deviation among NP-accepted PSMs at each FDR level.""" + work = df[[confidence_col, label_col]].dropna().copy() + if work.empty or label_col not in work.columns: + return {threshold: float("nan") for threshold in fdr_thresholds} + + np_fdr = NonParametricFDRControl() + np_fdr.fit(dataset=work[confidence_col]) + + dbg_ctrl = _fit_database_grounded_fdr( + work, + confidence_col=confidence_col, + correct_col=label_col, + drop=0 if len(work) <= 10 else 10, + ) + + scores = work[confidence_col].to_numpy(dtype=float) + est_q = _vectorized_psm_q_values(scores, np_fdr) + true_q = _vectorized_psm_q_values(scores, dbg_ctrl) + deviations = np.abs(est_q - true_q) + + results: dict[float, float] = {} + for threshold in fdr_thresholds: + mask = est_q <= threshold + if not mask.any(): + results[threshold] = float("nan") + else: + results[threshold] = float(np.mean(deviations[mask])) + return results + + +def metrics_from_eval_parquet(path: Path) -> dict[str, float | str]: + """Compute tail ECE, PR-AUC, and q-value metrics from one eval-results Parquet.""" + df = pd.read_parquet(path) + config_name = str(df["config_name"].iloc[0]) + dataset_name = str(df["dataset_name"].iloc[0]) + meta = df.drop(columns=["config_name", "dataset_name"], errors="ignore") + + tail_ece = compute_tail_ece_at_fdr_thresholds(meta) + fdr_bias = compute_fdr_bias_at_fdr_thresholds(meta) + pr_auc = compute_pr_auc(meta) + q_dev = compute_q_value_deviations(meta) + + return { + "config": config_name, + "dataset": dataset_name, + TAIL_ECE_COLUMN_BY_THRESHOLD[0.05]: round(tail_ece[0.05], 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.10]: round(tail_ece[0.10], 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.05]: round(fdr_bias[0.05], 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.10]: round(fdr_bias[0.10], 5), + "PR_AUC": round(pr_auc, 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.05]: round(q_dev[0.05], 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.10]: round(q_dev[0.10], 5), + } + + +def enrich_summary_from_eval_results( + summary: pd.DataFrame, + eval_results_dir: Path, + *, + datasets: list[str] | None = None, +) -> pd.DataFrame: + """Add PR-AUC and q-value deviation columns using saved eval Parquets.""" + if not eval_results_dir.is_dir(): + raise FileNotFoundError(eval_results_dir) + + metric_rows: list[dict[str, float | str]] = [] + for path in sorted(eval_results_dir.glob("*.parquet")): + dataset_name = path.name.split("_", 1)[0] + if datasets is not None and dataset_name not in datasets: + continue + metric_rows.append(metrics_from_eval_parquet(path)) + + if not metric_rows: + raise FileNotFoundError( + f"No eval Parquets found under {eval_results_dir}" + + (f" for datasets {datasets!r}" if datasets else "") + ) + + metrics_df = pd.DataFrame(metric_rows) + merge_cols = ["config", "dataset"] + extra_cols = [ + *TAIL_ECE_COLUMN_BY_THRESHOLD.values(), + *FDR_BIAS_COLUMN_BY_THRESHOLD.values(), + "PR_AUC", + *Q_DEV_COLUMN_BY_THRESHOLD.values(), + "tail_ECE", + ] + summary = summary.drop(columns=extra_cols, errors="ignore") + return summary.merge(metrics_df, on=merge_cols, how="left") + + +def _ablation_dataset_display(dataset: str) -> str: + """Publication label via ``plot_eval_results._display_name``.""" + return _display_name(_ABLATION_DATASET_KEYS.get(dataset, dataset)) + + +def _wrap_title_before_dataset(title: str, *, max_line: int = 52) -> str: + """Break before ``on `` when the title would be too wide.""" + marker = " on " + if marker not in title or len(title) <= max_line: + return title + split = title.index(marker) + return f"{title[:split]}\n{title[split + 1 :]}" + + +def _metric_axis_label(metric: MetricName) -> str: + if metric == "tail_ECE@5%FDR": + return "Tail ECE at 5% FDR" + if metric == "tail_ECE@10%FDR": + return "Tail ECE at 10% FDR" + if metric == "ECE": + return "ECE" + if metric == "Brier": + return "Brier score" + if metric == "PR_AUC": + return "PR-AUC" + if metric == "fdr_bias@5%FDR": + return "FDR bias (= sTECE) at 5% FDR" + if metric == "fdr_bias@10%FDR": + return "FDR bias (= sTECE) at 10% FDR" + if metric == "q_dev@5%FDR": + return "Mean |q-value deviation| at 5% FDR" + return "Mean |q-value deviation| at 10% FDR" + + +def _metric_plot_title(metric: MetricName, dataset_display: str) -> str: + """Publication title: full sentence, ECE capitalised.""" + if metric == "tail_ECE@5%FDR": + title = ( + f"Tail expected calibration error among PSMs accepted at 5% FDR " + f"on {dataset_display}" + ) + elif metric == "tail_ECE@10%FDR": + title = ( + f"Tail expected calibration error among PSMs accepted at 10% FDR " + f"on {dataset_display}" + ) + elif metric == "ECE": + title = f"Expected calibration error (ECE) on {dataset_display}" + elif metric == "Brier": + title = f"Brier score on {dataset_display}." + elif metric == "PR_AUC": + title = f"Precision-recall AUC on {dataset_display}" + elif metric == "fdr_bias@5%FDR": + title = ( + f"FDR bias, equal to signed tail calibration error, " + f"at 5% FDR on {dataset_display}" + ) + elif metric == "fdr_bias@10%FDR": + title = ( + f"FDR bias, equal to signed tail calibration error, " + f"at 10% FDR on {dataset_display}" + ) + elif metric == "q_dev@5%FDR": + title = ( + f"Non-parametric q-value deviation from database-grounded q-values " + f"at 5% FDR on {dataset_display}" + ) + else: + title = ( + f"Non-parametric q-value deviation from database-grounded q-values " + f"at 10% FDR on {dataset_display}" + ) + return _wrap_title_before_dataset(title) + + +def plot_ablation_calibration_bars( + summary: pd.DataFrame, + dataset: str, + *, + metric: MetricName = "tail_ECE@5%FDR", + output_path: Path, + figsize: tuple[float, float] = (7.5, 4), +) -> pd.DataFrame: + """Bar chart of *metric* for one dataset; returns the plotted slice.""" + ds = summary.loc[summary["dataset"] == dataset].copy() + if ds.empty: + available = sorted(summary["dataset"].unique()) + raise ValueError(f"No rows for dataset {dataset!r}. Available: {available}") + + configs = ordered_ablation_configs(set(ds["config"])) + ds = ds.set_index("config").loc[configs].reset_index() + if metric not in ds.columns: + raise ValueError(f"Metric {metric!r} not in summary columns: {ds.columns}") + + values = ds[metric].to_numpy(dtype=float) + all_features_value = float(ds.loc[ds["config"] == "All features", metric].iloc[0]) + + colors = assign_ablation_colors(configs) + short_labels = [_CONFIG_SHORT_LABELS.get(c, c) for c in configs] + + fig, ax = plt.subplots(figsize=figsize) + x = np.arange(len(configs)) + bar_colors = [colors[c] for c in configs] + ax.bar(x, values, color=bar_colors, edgecolor="black", linewidth=0.6, zorder=2) + ax.axhline( + all_features_value, + color="#333333", + linestyle="--", + linewidth=1.2, + zorder=1, + label="All features", + ) + + display = _ablation_dataset_display(dataset) + ax.set_ylabel(_metric_axis_label(metric)) + ax.set_xlabel("Calibrator feature groups") + ax.set_title(_metric_plot_title(metric, display)) + ax.set_xticks(x) + ax.set_xticklabels(short_labels, rotation=35, ha="right") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s.png and %s.pdf", output_path, output_path) + return ds[["config", metric]] + + +@app.command() +def main( + summary: Annotated[ + Path, + typer.Option("--summary", help="ablation_summary.csv or .json"), + ] = _DEFAULT_SUMMARY, + dataset: Annotated[ + str, + typer.Option("--dataset", help="Dataset key in the summary table"), + ] = "Arabidopsis", + metric: Annotated[ + MetricName, + typer.Option("--metric", help="Calibration metric to plot"), + ] = "tail_ECE@5%FDR", + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for figure outputs"), + ] = _DEFAULT_SUMMARY.parent / "plots", + eval_results_dir: Annotated[ + Path | None, + typer.Option( + "--eval-results-dir", + help="Optional eval_results/ directory to enrich summary before plotting", + ), + ] = None, +) -> None: + """Plot ablation calibration bars for one dataset.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + sns.set_theme(style="white", context="paper", font_scale=1.5) + + summary_df = load_ablation_summary(summary) + if eval_results_dir is not None: + summary_df = enrich_summary_from_eval_results( + summary_df, + eval_results_dir, + datasets=[dataset], + ) + output_dir.mkdir(parents=True, exist_ok=True) + slug = dataset.lower().replace(" ", "_") + metric_slug = metric.lower().replace("%", "pct").replace("@", "_at_") + out_base = output_dir / f"ablation_{metric_slug}_{slug}" + table = plot_ablation_calibration_bars( + summary_df, dataset, metric=metric, output_path=out_base + ) + print(table.to_string(index=False)) + + +@app.command("recompute-summary") +def recompute_summary( + eval_results_dir: Annotated[ + Path, + typer.Option("--eval-results-dir", help="Directory of eval_results Parquets"), + ], + summary: Annotated[ + Path | None, + typer.Option( + "--summary", + help="Existing ablation_summary.csv to merge with (optional)", + ), + ] = None, + datasets: Annotated[ + list[str] | None, + typer.Option( + "--datasets", + help="Restrict to these dataset keys (repeatable)", + ), + ] = None, + output: Annotated[ + Path, + typer.Option("--output", help="Output CSV path"), + ] = _DEFAULT_SUMMARY, +) -> None: + """Recompute PR-AUC and q-value deviation columns from eval Parquets.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + if summary is not None: + base = load_ablation_summary(summary) + if datasets is not None: + base = base.loc[base["dataset"].isin(datasets)].copy() + enriched = enrich_summary_from_eval_results( + base, + eval_results_dir, + datasets=datasets, + ) + else: + rows = [] + for path in sorted(eval_results_dir.glob("*.parquet")): + dataset_name = path.name.split("_", 1)[0] + if datasets is not None and dataset_name not in datasets: + continue + rows.append(metrics_from_eval_parquet(path)) + if not rows: + raise typer.BadParameter(f"No eval Parquets found under {eval_results_dir}") + enriched = pd.DataFrame(rows) + enriched = enriched.sort_values(["dataset", "config"]).reset_index(drop=True) + + output.parent.mkdir(parents=True, exist_ok=True) + enriched.to_csv(output, index=False) + logger.info("Wrote %s", output) + print(enriched.to_string(index=False)) + + +if __name__ == "__main__": + app() diff --git a/scripts/plot_acfm_minus_lcfm_fdr.py b/scripts/plot_acfm_minus_lcfm_fdr.py new file mode 100644 index 00000000..8381eacb --- /dev/null +++ b/scripts/plot_acfm_minus_lcfm_fdr.py @@ -0,0 +1,441 @@ +"""Refit FDR on acfm predictions restricted to spectra not present in lcfm. + +For each external project, spectra are matched on ``spectrum_id``. The acfm +(unlabelled) set is filtered to ``spectrum_id`` values absent from the paired +lcfm (labelled) predictions, FDR is re-estimated on ``calibrated_confidence`` +for that subset only, and evaluation plots are written. +""" + +from __future__ import annotations + +import logging +import sys +from pathlib import Path +from typing import Annotated + +import pandas as pd +import typer +from rich.logging import RichHandler + +_REPO_ROOT = Path(__file__).resolve().parent.parent +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.plot_eval_results import ( # noqa: E402 + _compute_diagnostics, + _display_name, + _fit_database_grounded_fdr, + generate_all_plots, +) +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +def _safe_basename(project: str) -> str: + """Flat basename for outputs; project keys may contain path separators.""" + return project.replace("/", "_") + + +# Per-run keys under new_eval_sets_results/{lcfm,acfm}/// +_PXD_RUN_PARENTS: tuple[str, ...] = ("PXD004452", "PXD006939", "PXD013868") + + +def _preds_csv_candidates(root: Path, project: str, *, role: str) -> list[Path]: + """Paths to try for ``preds_and_fdr_metrics.csv`` under *root*. + + Supports flat layouts (``{root}/{run}/``), S3-style nesting + (``{root}/PXD006939/{run}/``), and explicit ``PXD006939/run`` project keys. + """ + if role not in ("labelled", "unlabelled"): + raise ValueError(f"Unknown role {role!r}") + + role_suffix = "_labelled" if role == "labelled" else "_unlabelled" + fname = "preds_and_fdr_metrics.csv" + seen: set[Path] = set() + candidates: list[Path] = [] + + def add(*relative: str) -> None: + path = root.joinpath(*relative, fname) + if path not in seen: + seen.add(path) + candidates.append(path) + + add(project) + base = project.split("/")[-1] + add(f"{base}{role_suffix}") + if "/" in project: + add(*project.split("/")) + + # S3 and download-new-eval-results: {root}/PXD*/{run}/; legacy flat {root}/{run}/. + if "/" not in project: + pxd_seen: set[str] = set() + for pxd in _PXD_RUN_PARENTS: + add(pxd, project) + add(pxd, f"{project}{role_suffix}") + pxd_seen.add(pxd) + if root.is_dir(): + for child in sorted(root.iterdir()): + if child.is_dir() and child.name.startswith("PXD"): + if child.name not in pxd_seen: + add(child.name, project) + add(child.name, f"{project}{role_suffix}") + + return candidates + + +def _resolve_preds_csv( + root: Path, + project: str, + *, + role: str, + alt_roots: list[Path] | None = None, +) -> Path: + """Resolve ``preds_and_fdr_metrics.csv`` for lcfm (labelled) or acfm (unlabelled).""" + roots_to_try: list[Path] = [root] + if alt_roots: + for alt in alt_roots: + if alt.resolve() != root.resolve() and alt not in roots_to_try: + roots_to_try.append(alt) + + tried: list[Path] = [] + for base in roots_to_try: + for path in _preds_csv_candidates(base, project, role=role): + tried.append(path) + if path.is_file(): + if base.resolve() != root.resolve(): + logger.info( + "Using %s predictions at %s (not under %s)", + role, + path, + root, + ) + return path + + hint = "" + if root.is_dir(): + children = sorted(p.name for p in root.iterdir())[:12] + hint = f" Children of {root}: {children}" + raise FileNotFoundError( + f"Missing {role} predictions for {project!r} under {root} " + f"(tried: {', '.join(str(p) for p in tried)}){hint}" + ) + + +def _infer_predictions_root( + labelled_dir: Path | None, + unlabelled_dir: Path | None, +) -> Path: + """Infer a common parent when only per-tree dirs are passed.""" + if labelled_dir is not None and unlabelled_dir is not None: + if labelled_dir.parent == unlabelled_dir.parent: + return labelled_dir.parent + if labelled_dir is not None: + return labelled_dir.parent + if unlabelled_dir is not None: + return unlabelled_dir.parent + raise typer.BadParameter( + "Provide --predictions-root or at least one of --labelled-dir / --unlabelled-dir." + ) + + +def _resolve_tree_root( + predictions_root: Path | None, + explicit: Path | None, + *, + role: str, +) -> tuple[Path, list[Path]]: + """Pick labelled (lcfm) or unlabelled (acfm) root; return alternates to try.""" + sub = "lcfm" if role == "labelled" else "acfm" + legacy = "labelled" if role == "labelled" else "unlabelled" + alt_roots: list[Path] = [] + + if explicit is not None: + if predictions_root is not None and len(explicit.parts) == 1: + under_predictions = predictions_root / explicit + if under_predictions.is_dir(): + root = under_predictions + if explicit.is_dir() and explicit.resolve() != root.resolve(): + alt_roots.append(explicit) + nested = predictions_root / sub + if nested.is_dir() and nested.resolve() != root.resolve(): + alt_roots.append(nested) + return root, alt_roots + return explicit, alt_roots + + if predictions_root is None: + raise typer.BadParameter( + f"Missing --predictions-root and --{legacy}-dir for {role} tree." + ) + + for name in (sub, legacy): + candidate = predictions_root / name + if candidate.is_dir(): + return candidate, alt_roots + return predictions_root, alt_roots + + +def _lcfm_spectrum_ids( + labelled_root: Path, + project: str, + *, + labelled_alt_roots: list[Path] | None = None, +) -> set[str]: + labelled_path = _resolve_preds_csv( + labelled_root, + project, + role="labelled", + alt_roots=labelled_alt_roots, + ) + ids = pd.read_csv(labelled_path, usecols=["spectrum_id"])["spectrum_id"] + return set(ids.astype(str)) + + +def _load_acfm_unlabelled( + unlabelled_root: Path, + project: str, + *, + unlabelled_alt_roots: list[Path] | None = None, +) -> pd.DataFrame: + """Load acfm predict outputs with metadata merged (same as plot_eval_results).""" + preds_path = _resolve_preds_csv( + unlabelled_root, + project, + role="unlabelled", + alt_roots=unlabelled_alt_roots, + ) + folder = preds_path.parent + preds_df = pd.read_csv(preds_path) + meta_path = folder / "metadata.csv" + if meta_path.is_file(): + meta_df = pd.read_csv(meta_path) + overlap = [ + c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" + ] + if overlap: + meta_df = meta_df.drop(columns=overlap) + df = preds_df.merge(meta_df, on="spectrum_id", how="left") + else: + df = preds_df + + if "proteome_hit" not in df.columns: + raise ValueError( + f"Expected 'proteome_hit' column for unlabelled acfm in {preds_path}" + ) + df["correct"] = df["proteome_hit"].astype(float) + required = ["confidence", "calibrated_confidence", "correct"] + missing = [c for c in required if c not in df.columns] + if missing: + raise ValueError(f"Missing columns {missing} in {preds_path}") + return df + + +def filter_acfm_minus_lcfm(acfm_df: pd.DataFrame, lcfm_ids: set[str]) -> pd.DataFrame: + """Keep acfm rows whose ``spectrum_id`` is not in the lcfm set.""" + mask = ~acfm_df["spectrum_id"].astype(str).isin(lcfm_ids) + return acfm_df.loc[mask].copy() + + +def refit_fdr_on_confidence( + df: pd.DataFrame, + confidence_col: str = "calibrated_confidence", +) -> pd.DataFrame: + """Fit non-parametric FDR on *df* and attach PSM FDR / q-value / PEP columns.""" + out = df.copy() + for col in ("psm_fdr", "psm_q_value", "psm_pep"): + if col in out.columns: + out = out.drop(columns=[col]) + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=out[confidence_col]) + out = fdr_ctrl.add_psm_fdr(out, confidence_col=confidence_col) + out = fdr_ctrl.add_psm_q_value(out, confidence_col=confidence_col) + out = fdr_ctrl.add_psm_pep(out, confidence_col=confidence_col) + return out + + +def process_project( + labelled_root: Path, + unlabelled_root: Path, + project: str, + output_dir: Path, + *, + labelled_alt_roots: list[Path] | None = None, + unlabelled_alt_roots: list[Path] | None = None, +) -> dict[str, int]: + """Filter acfm less lcfm, refit FDR, plot, and write tables for one project.""" + lcfm_ids = _lcfm_spectrum_ids( + labelled_root, project, labelled_alt_roots=labelled_alt_roots + ) + acfm_df = _load_acfm_unlabelled( + unlabelled_root, project, unlabelled_alt_roots=unlabelled_alt_roots + ) + subset_df = filter_acfm_minus_lcfm(acfm_df, lcfm_ids) + + counts = { + "n_lcfm_spectrum_ids": len(lcfm_ids), + "n_acfm": len(acfm_df), + "n_acfm_minus_lcfm": len(subset_df), + } + if counts["n_acfm_minus_lcfm"] == 0: + raise ValueError( + f"{project}: no acfm spectra remain after excluding lcfm spectrum_id values" + ) + + logger.info( + "%s: acfm=%s, lcfm ids=%s, acfm\\lcfm=%s", + project, + f"{counts['n_acfm']:,}", + f"{counts['n_lcfm_spectrum_ids']:,}", + f"{counts['n_acfm_minus_lcfm']:,}", + ) + + subset_df = refit_fdr_on_confidence(subset_df) + + project_dir = output_dir / project + project_dir.mkdir(parents=True, exist_ok=True) + safe = _safe_basename(project) + subset_df.to_csv(project_dir / "preds_and_fdr_metrics.csv", index=False) + + true_fdr_ctrl = _fit_database_grounded_fdr(subset_df) + db_fdr = true_fdr_ctrl.add_psm_fdr( + subset_df[["calibrated_confidence"]].copy(), + confidence_col="calibrated_confidence", + ) + subset_df["db_grounded_psm_fdr"] = db_fdr["psm_fdr"] + db_qval = true_fdr_ctrl.add_psm_q_value( + subset_df[["calibrated_confidence"]].copy(), + confidence_col="calibrated_confidence", + ) + subset_df["db_grounded_psm_q_value"] = db_qval["psm_q_value"] + + summary_cols = [ + c + for c in [ + "spectrum_id", + "prediction", + "confidence", + "calibrated_confidence", + "correct", + "psm_fdr", + "psm_q_value", + "db_grounded_psm_fdr", + "db_grounded_psm_q_value", + "proteome_hit", + ] + if c in subset_df.columns + ] + subset_df[summary_cols].to_csv(project_dir / f"{safe}_summary.csv", index=False) + + diag = _compute_diagnostics(subset_df, "unlabelled") + diag.to_csv(project_dir / f"{safe}_diagnostics.csv", index=False) + + pd.DataFrame([counts]).to_csv(project_dir / f"{safe}_counts.csv", index=False) + + generate_all_plots(subset_df, safe, "unlabelled", project_dir) + return counts + + +@app.command() +def main( + projects: Annotated[ + str, + typer.Option( + "--projects", + help="Space- or comma-separated project keys (e.g. 'PXD009935 PXD014877').", + ), + ], + output_dir: Annotated[ + Path, + typer.Option( + "--output-dir", + help="Directory for per-project plots and refitted prediction tables.", + ), + ], + predictions_root: Annotated[ + Path | None, + typer.Option( + "--predictions-root", + help=( + "Parent of labelled/unlabelled (or lcfm/acfm) trees. Optional when " + "both --labelled-dir and --unlabelled-dir are set (parent is inferred)." + ), + ), + ] = None, + labelled_dir: Annotated[ + Path | None, + typer.Option( + "--labelled-dir", + help=( + "Root with per-project lcfm folders ({project}/ or {project}_labelled/). " + "Use e.g. new_eval_sets_results/lcfm when mirroring S3. " + "If omitted, uses --predictions-root/lcfm when present." + ), + ), + ] = None, + unlabelled_dir: Annotated[ + Path | None, + typer.Option( + "--unlabelled-dir", + help=( + "Root with per-project acfm folders ({project}/ or {project}_unlabelled/). " + "Use e.g. new_eval_sets_results/acfm when mirroring S3. " + "If omitted, uses --predictions-root/acfm when present." + ), + ), + ] = None, +) -> None: + """Refit FDR on acfm less lcfm spectra and generate evaluation plots.""" + project_list = [p.strip() for p in projects.replace(",", " ").split() if p.strip()] + if not project_list: + raise typer.BadParameter("No projects specified.") + + preds_root = predictions_root + if preds_root is None: + preds_root = _infer_predictions_root(labelled_dir, unlabelled_dir) + + labelled_root, labelled_alt = _resolve_tree_root( + preds_root, labelled_dir, role="labelled" + ) + unlabelled_root, unlabelled_alt = _resolve_tree_root( + preds_root, unlabelled_dir, role="unlabelled" + ) + logger.info("Labelled (lcfm) root: %s", labelled_root.resolve()) + logger.info("Unlabelled (acfm) root: %s", unlabelled_root.resolve()) + + output_dir.mkdir(parents=True, exist_ok=True) + all_counts: list[dict[str, int | str]] = [] + + for project in project_list: + display = _display_name(project) + logger.info("Processing %s (%s)...", project, display) + try: + counts = process_project( + labelled_root, + unlabelled_root, + project, + output_dir, + labelled_alt_roots=labelled_alt, + unlabelled_alt_roots=unlabelled_alt, + ) + except FileNotFoundError as exc: + logger.warning("Skipping %s: %s", project, exc) + continue + all_counts.append({"project": project, **counts}) + + if all_counts: + pd.DataFrame(all_counts).to_csv(output_dir / "counts_summary.csv", index=False) + logger.info("Wrote counts summary to %s", output_dir / "counts_summary.csv") + else: + raise typer.Exit(code=1) + + logger.info("Done. Outputs in %s", output_dir) + + +if __name__ == "__main__": + app() diff --git a/scripts/plot_analysis.py b/scripts/plot_analysis.py new file mode 100644 index 00000000..40748198 --- /dev/null +++ b/scripts/plot_analysis.py @@ -0,0 +1,1133 @@ +"""Generate analysis plots from Winnow predict outputs. + +Usage: + python scripts/plot_analysis.py \ + --predictions-dir results/instanovo_helaqc_predictions_test/ \ + --split test \ + --label-mode labelled \ + --fasta fasta/human.fasta \ + [--model-dir models/instanovo_helaqc] \ + [--output-dir results/instanovo_helaqc_predictions_test/plots/] +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +import seaborn as sns +import yaml +from matplotlib.patches import Patch +from scipy.stats import gaussian_kde +from sklearn.calibration import calibration_curve +from sklearn.decomposition import PCA +from sklearn.preprocessing import StandardScaler + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT)) + +from winnow.calibration.calibrator import TrainingHistory # noqa: E402 +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 +from scripts.annotate_preds_proteome_hits import ( # noqa: E402 + filter_and_annotate_preds, + load_proteome_haystack, +) + +# ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_MAIN_LINE_COLOUR = _PALETTE[0] +_RAW_LINE_COLOUR = _PALETTE[5] +_IDEAL_LINE_COLOUR = _PALETTE[6] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) +warnings.filterwarnings("ignore", module="winnow") + + +def _spine_fmt(ax: plt.Axes) -> None: + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _save(fig: plt.Figure, out_dir: Path, name: str) -> None: + base = out_dir / name + fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + print(f" saved {name}") + + +# ── Plot functions ──────────────────────────────────────────────────── + + +def _df_for_raw_confidence_plots(df: pl.DataFrame) -> pl.DataFrame: + """Drop PSMs with negative raw confidence (Casanovo mass-mismatch penalty scores).""" + if "confidence" not in df.columns: + return df + n_neg = int((df["confidence"] < 0).sum()) + if n_neg == 0: + return df + print( + f" excluding {n_neg} PSMs with negative raw confidence " + "from raw-confidence plots" + ) + return df.filter(pl.col("confidence") >= 0) + + +def plot_calibration_curves( + df: pl.DataFrame, + label_col: str, + title: str, + bins: int = 10, + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Plot reliability curves for calibrated and (optionally) raw confidence.""" + fig, ax = plt.subplots(figsize=(8, 6)) + + frac_pos, mean_pred = calibration_curve( + df[label_col].to_numpy(), + df["calibrated_confidence"].to_numpy(), + n_bins=bins, + strategy="uniform", + ) + ax.plot( + mean_pred, + frac_pos, + marker="o", + color=_MAIN_LINE_COLOUR, + label="Calibrated confidence", + linewidth=1.5, + markersize=6, + zorder=3, + ) + + if df_raw is not None and len(df_raw) > 0 and "confidence" in df_raw.columns: + frac_pos, mean_pred = calibration_curve( + df_raw[label_col].to_numpy(), + df_raw["confidence"].to_numpy(), + n_bins=bins, + strategy="uniform", + ) + ax.plot( + mean_pred, + frac_pos, + marker="D", + color=_RAW_LINE_COLOUR, + label="Raw confidence", + linewidth=1.5, + markersize=6, + zorder=3, + ) + + ax.plot( + [0, 1], + [0, 1], + "--", + color=_IDEAL_LINE_COLOUR, + label="Perfectly calibrated", + alpha=0.7, + zorder=2, + ) + ax.set_xlabel("Mean predicted probability") + ax.set_ylabel("Fraction of positives") + ax.set_title(title) + ax.legend(loc="lower right") + ax.set_xlim([0, 1.05]) + ax.set_ylim([0, 1.05]) + ax.grid(False) + _spine_fmt(ax) + return fig + + +def plot_pr_curves( + df: pl.DataFrame, + label_col: str, + title: str, + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Plot precision–recall curves for calibrated and (optionally) raw confidence.""" + fig, ax = plt.subplots(figsize=(8, 6)) + + sorted_cal = df.sort("calibrated_confidence", descending=True) + labels = sorted_cal[label_col].to_numpy() + cum = np.cumsum(labels) + precision = cum / np.arange(1, len(labels) + 1) + recall = cum / len(labels) + ax.plot( + recall, + precision, + color=_MAIN_LINE_COLOUR, + label="Calibrated confidence", + linewidth=1.5, + ) + + if df_raw is not None and len(df_raw) > 0 and "confidence" in df_raw.columns: + sorted_raw = df_raw.sort("confidence", descending=True) + labels = sorted_raw[label_col].to_numpy() + cum = np.cumsum(labels) + precision = cum / np.arange(1, len(labels) + 1) + recall = cum / len(labels) + ax.plot( + recall, + precision, + color=_RAW_LINE_COLOUR, + label="Raw confidence", + linewidth=1.5, + ) + + ax.set_xlabel("Recall") + ax.set_ylabel("Precision") + ax.set_title(title) + ax.set_xlim(0, 1.05) + ax.set_ylim(0, 1.05) + ax.legend(loc="lower left") + ax.grid(False) + _spine_fmt(ax) + return fig + + +def plot_confidence_histogram( + df: pl.DataFrame, + label_col: str, + conf_col: str, + col_label: str, + title: str, + bins: int = 50, +) -> plt.Figure: + """Plot confidence histograms with KDE overlays for correct vs incorrect PSMs.""" + fig, ax = plt.subplots(1, 1, figsize=(7, 5)) + pos = df.filter(pl.col(label_col)) + neg = df.filter(~pl.col(label_col)) + n_data = neg[conf_col].to_numpy() + p_data = pos[conf_col].to_numpy() + + ax.hist( + n_data, + bins=bins, + alpha=0.6, + label="Incorrect", + density=False, + edgecolor="black", + color=_INCORRECT_COLOUR, + ) + ax.hist( + p_data, + bins=bins, + alpha=0.6, + label="Correct", + density=False, + edgecolor="black", + color=_CORRECT_COLOUR, + ) + + x_min = min(n_data.min(), p_data.min()) + x_max = max(n_data.max(), p_data.max()) + x_grid = np.linspace(x_min, x_max, 300) + bin_width = (x_max - x_min) / bins if bins > 1 else 1.0 + + if len(n_data) > 1: + y_neg = gaussian_kde(n_data)(x_grid) * len(n_data) * bin_width + ax.plot(x_grid, y_neg, color=_INCORRECT_COLOUR, lw=1.5) + if len(p_data) > 1: + y_pos = gaussian_kde(p_data)(x_grid) * len(p_data) * bin_width + ax.plot(x_grid, y_pos, color=_CORRECT_COLOUR, lw=1.5) + + ax.set_xlabel(col_label) + ax.set_ylabel("Frequency") + ax.legend(loc="upper center") + ax.grid(False) + ax.set_title(title) + _spine_fmt(ax) + fig.tight_layout() + return fig + + +def _fit_db_fdr( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + confidence_feature: str = "calibrated_confidence", + drop: int = 10, + use_proteome_shortcut: bool = False, +) -> DatabaseGroundedFDRControl: + """Fit a DatabaseGroundedFDRControl, using a proteome shortcut if labels lack sequences.""" + ctrl = DatabaseGroundedFDRControl( + confidence_feature=confidence_feature, + residue_masses=residue_masses, + drop=drop, + ) + if use_proteome_shortcut: + sorted_df = df.sort(confidence_feature, descending=True) + correct_vals = sorted_df[correct_col].to_numpy().astype(float) + confidence_vals = sorted_df[confidence_feature].to_numpy() + precision = np.cumsum(correct_vals) / np.arange(1, len(sorted_df) + 1) + ctrl._fdr_values = np.array(1 - precision[drop:]) + ctrl._confidence_scores = confidence_vals[drop:] + else: + ctrl.fit(dataset=df.to_pandas(), correct_column=correct_col) + return ctrl + + +def plot_fdr_accuracy( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, + metric: str = "fdr", + use_proteome_shortcut: bool = False, +) -> plt.Figure: + """Compare non-parametric vs database-grounded FDR or q-value vs confidence.""" + fig, ax = plt.subplots(figsize=(8, 6)) + col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" + winnow_col = col_name + + ctrl = _fit_db_fdr( + df, correct_col, residue_masses, use_proteome_shortcut=use_proteome_shortcut + ) + + if metric == "fdr": + db_pd = ctrl.add_psm_fdr(df.to_pandas(), "calibrated_confidence") + else: + df_pd = df.to_pandas() + if "psm_q_value" in df_pd.columns: + df_pd = df_pd.drop(columns=["psm_q_value"]) + db_pd = ctrl.add_psm_q_value(df_pd, "calibrated_confidence") + db_df = pl.from_pandas(db_pd).select(["spectrum_id", col_name]) + + merged = ( + df.select(["spectrum_id", "calibrated_confidence", winnow_col]) + .join(db_df, on="spectrum_id", how="inner", suffix="_db") + .sort("calibrated_confidence") + ) + + conf = merged["calibrated_confidence"].to_numpy() + ax.plot( + conf, + merged[winnow_col].to_numpy(), + color=_MAIN_LINE_COLOUR, + label="Non-parametric", + linewidth=1.5, + ) + ax.plot( + conf, + merged[f"{col_name}_db"].to_numpy(), + color=_RAW_LINE_COLOUR, + label="Database-grounded", + linewidth=1.5, + ) + + ax.set_xlabel("Calibrated confidence") + ylabel = "FDR" if metric == "fdr" else "Q-value" + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.legend(loc="upper right") + ax.grid(False) + _spine_fmt(ax) + return fig + + +def plot_ranked_qvalue( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, +) -> plt.Figure: + """Ranked predictions vs q-value (non-parametric & database-grounded).""" + ctrl = _fit_db_fdr(df, correct_col, residue_masses) + test_pd = df.to_pandas() + test_pd_no_q = test_pd.drop(columns=["psm_q_value"], errors="ignore") + db_q = ctrl.add_psm_q_value(test_pd_no_q, "calibrated_confidence") + + sorted_np = test_pd.sort_values( + "calibrated_confidence", ascending=False + ).reset_index(drop=True) + sorted_db = db_q.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + ranks = np.arange(1, len(sorted_np) + 1) + + fig, ax = plt.subplots(figsize=(8, 6)) + ax.plot( + ranks, + sorted_np["psm_q_value"].values, + color=_MAIN_LINE_COLOUR, + label="Non-parametric", + linewidth=1.5, + ) + ax.plot( + ranks, + sorted_db["psm_q_value"].values, + color=_RAW_LINE_COLOUR, + label="Database-grounded", + linewidth=1.5, + ) + ax.set_xlabel("Ranked predictions") + ax.set_ylabel("Q-value") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_ranked_fdr_pep(df: pl.DataFrame, title: str) -> plt.Figure: + """Ranked predictions vs non-parametric FDR and PEP.""" + sorted_df = df.sort("calibrated_confidence", descending=True) + ranks = np.arange(1, len(sorted_df) + 1) + fig, ax = plt.subplots(figsize=(8, 6)) + ax.plot( + ranks, + sorted_df["psm_fdr"].to_numpy(), + color=_MAIN_LINE_COLOUR, + label="FDR", + linewidth=1.5, + ) + if "psm_pep" in sorted_df.columns: + ax.plot( + ranks, + sorted_df["psm_pep"].to_numpy(), + color=_PALETTE[3], + label="PEP", + linewidth=1.5, + ) + ax.set_xlabel("Ranked predictions") + ax.set_ylabel("Error rate") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_ranked_fdr_raw_vs_cal( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, + metric: str = "fdr", + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Ranked predictions vs FDR/q-value for non-parametric and database-grounded on raw+calibrated.""" + test_pd = df.to_pandas() + test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") + + np_cal = test_pd.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + + db_cal_ctrl = _fit_db_fdr( + df, correct_col, residue_masses, confidence_feature="calibrated_confidence" + ) + raw_df = df_raw if df_raw is not None else df + db_raw_ctrl = _fit_db_fdr( + raw_df, correct_col, residue_masses, confidence_feature="confidence" + ) + + col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" + add_fn = "add_psm_fdr" if metric == "fdr" else "add_psm_q_value" + sort_col_cal = "calibrated_confidence" + sort_col_raw = "confidence" + + db_cal = getattr(db_cal_ctrl, add_fn)(test_pd_no_q.copy(), sort_col_cal) + db_cal = db_cal.sort_values(sort_col_cal, ascending=False).reset_index(drop=True) + + raw_pd_no_q = raw_df.to_pandas().drop( + columns=["psm_q_value", "psm_fdr"], errors="ignore" + ) + db_raw = getattr(db_raw_ctrl, add_fn)(raw_pd_no_q.copy(), sort_col_raw) + db_raw = db_raw.sort_values(sort_col_raw, ascending=False).reset_index(drop=True) + + ranks_cal = np.arange(1, len(np_cal) + 1) + ranks_raw = np.arange(1, len(db_raw) + 1) + + fig, ax = plt.subplots(figsize=(8, 6)) + ylabel = "FDR" if metric == "fdr" else "Q-value" + np_col = "psm_fdr" if metric == "fdr" else "psm_q_value" + ax.plot( + ranks_cal, + np_cal[np_col].values, + color=_MAIN_LINE_COLOUR, + label="Non-parametric (calibrated)", + linewidth=1.5, + ) + ax.plot( + ranks_cal, + db_cal[col_name].values, + color=_RAW_LINE_COLOUR, + label="Database-grounded (calibrated)", + linewidth=1.5, + ) + ax.plot( + ranks_raw, + db_raw[col_name].values, + color=_PALETTE[3], + label="Database-grounded (raw)", + linewidth=1.5, + ) + ax.set_xlabel("Ranked predictions") + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_bar_psms_fdr( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Bar plot of PSMs at q-value thresholds (calibrated vs raw, database-grounded).""" + test_pd = df.to_pandas() + test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") + + db_cal_ctrl = _fit_db_fdr( + df, correct_col, residue_masses, confidence_feature="calibrated_confidence" + ) + raw_df = df_raw if df_raw is not None else df + db_raw_ctrl = _fit_db_fdr( + raw_df, correct_col, residue_masses, confidence_feature="confidence" + ) + + db_cal = db_cal_ctrl.add_psm_q_value(test_pd_no_q.copy(), "calibrated_confidence") + raw_pd_no_q = raw_df.to_pandas().drop( + columns=["psm_q_value", "psm_fdr"], errors="ignore" + ) + db_raw = db_raw_ctrl.add_psm_q_value(raw_pd_no_q.copy(), "confidence") + + thresholds = [0.001, 0.01, 0.05, 0.1] + counts_cal = [int((db_cal["psm_q_value"] <= t).sum()) for t in thresholds] + counts_raw = [int((db_raw["psm_q_value"] <= t).sum()) for t in thresholds] + + x = np.arange(len(thresholds)) + width, gap = 0.32, 0.04 + + fig, ax = plt.subplots(figsize=(8, 6)) + bars_cal = ax.bar( + x - width / 2 - gap / 2, + counts_cal, + width, + label="Calibrated confidence", + color=_MAIN_LINE_COLOUR, + edgecolor="black", + linewidth=1, + ) + bars_raw = ax.bar( + x + width / 2 + gap / 2, + counts_raw, + width, + label="Raw confidence", + color=_RAW_LINE_COLOUR, + edgecolor="black", + linewidth=1, + ) + ax.set_xlabel("FDR threshold") + ax.set_ylabel("Peptide-spectrum matches") + ax.set_title(title) + ax.set_xticks(x) + ax.set_xticklabels([str(t) for t in thresholds]) + ax.legend(loc="upper left") + + for bar_group in [bars_cal, bars_raw]: + for bar in bar_group: + h = bar.get_height() + ax.annotate( + f"{h:,}", + xy=(bar.get_x() + bar.get_width() / 2, h), + xytext=(0, 3), + textcoords="offset points", + ha="center", + va="bottom", + fontsize=10, + ) + _spine_fmt(ax) + return fig + + +def plot_raw_vs_cal_scatter( + df: pl.DataFrame, + label_col: str, + title: str, +) -> plt.Figure: + """Raw confidence vs calibrated confidence coloured by correctness.""" + fig, ax = plt.subplots(figsize=(8, 7)) + inc = df.filter(~pl.col(label_col)) + cor = df.filter(pl.col(label_col)) + ax.scatter( + inc["confidence"].to_numpy(), + inc["calibrated_confidence"].to_numpy(), + c=_INCORRECT_COLOUR, + label="Incorrect", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.scatter( + cor["confidence"].to_numpy(), + cor["calibrated_confidence"].to_numpy(), + c=_CORRECT_COLOUR, + label="Correct", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.plot( + [0, 1], + [0, 1], + color=_IDEAL_LINE_COLOUR, + linestyle="--", + linewidth=1, + label="Identity", + ) + ax.set_xlabel("Raw confidence") + ax.set_ylabel("Calibrated confidence") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_pca_features( + df: pl.DataFrame, + label_col: str, + title: str, +) -> tuple[plt.Figure, PCA, list[str]]: + """PCA of calibrator features coloured by correctness.""" + feature_cols = [ + "confidence", + "mass_error_ppm", + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", + "irt_error", + "margin", + "median_margin", + "entropy", + "z-score", + "edit_distance", + "min_token_probability", + "std_token_probability", + ] + available = [c for c in feature_cols if c in df.columns] + feat_df = df.select(available).to_pandas().dropna() + labels = df.filter(pl.all_horizontal([pl.col(c).is_not_null() for c in available]))[ + label_col + ].to_numpy() + + scaler = StandardScaler() + features_scaled = scaler.fit_transform(feat_df.values) + pca = PCA(n_components=2) + coords = pca.fit_transform(features_scaled) + + fig, ax = plt.subplots(figsize=(8, 7)) + mask_inc, mask_cor = ~labels, labels + ax.scatter( + coords[mask_inc, 0], + coords[mask_inc, 1], + c=_INCORRECT_COLOUR, + label="Incorrect", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.scatter( + coords[mask_cor, 0], + coords[mask_cor, 1], + c=_CORRECT_COLOUR, + label="Correct", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.set_xlabel(f"PC 1 ({pca.explained_variance_ratio_[0]:.1%} variance)") + ax.set_ylabel(f"PC 2 ({pca.explained_variance_ratio_[1]:.1%} variance)") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig, pca, available + + +def plot_pca_loadings( + pca: PCA, + feature_names: list[str], + title: str, +) -> plt.Figure: + """PCA loadings for PC1 and PC2, ordered by |PC1|.""" + pretty = { + "confidence": "Raw confidence", + "mass_error_ppm": "Log absolute mass error (ppm)", + "ion_matches": "Ion matches", + "ion_match_intensity": "Ion match intensity", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Maximum ion gap", + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation", + "irt_error": "Retention time error", + "margin": "Margin", + "median_margin": "Median margin", + "entropy": "Entropy", + "z-score": "Z-score", + "edit_distance": "Edit distance", + "min_token_probability": "Minimum token probability", + "std_token_probability": "Token probability std. dev.", + } + pc1 = pca.components_[0] + pc2 = pca.components_[1] + names = [pretty.get(c, c) for c in feature_names] + order = np.argsort(np.abs(pc1))[::-1] + + y = np.arange(len(names)) + fig, ax = plt.subplots(figsize=(10, 7)) + ax.barh(y, pc1[order], color=_MAIN_LINE_COLOUR, alpha=0.6, edgecolor="black") + ax.barh(y, pc2[order], color=_RAW_LINE_COLOUR, alpha=0.4, edgecolor="black") + ax.set_yticks(y) + ax.set_yticklabels([names[i] for i in order]) + ax.invert_yaxis() + ax.set_xlabel("Loading value") + ax.set_title(title) + ax.axvline(0, color="black", linewidth=0.5) + ax.legend( + handles=[ + Patch(facecolor=_MAIN_LINE_COLOUR, alpha=0.6, label="PC 1 loading"), + Patch(facecolor=_RAW_LINE_COLOUR, alpha=0.4, label="PC 2 loading"), + ], + loc="lower right", + ) + _spine_fmt(ax) + return fig + + +def plot_scatter_feature_vs_conf( + df: pl.DataFrame, + label_col: str, + x_col: str, + y_col: str, + title: str, + x_label: str | None = None, + y_label: str | None = None, +) -> plt.Figure: + """Scatter of x_col vs y_col coloured by correctness.""" + fig, ax = plt.subplots(figsize=(8, 7)) + inc = df.filter(~pl.col(label_col)) + cor = df.filter(pl.col(label_col)) + ax.scatter( + inc[x_col].to_numpy(), + inc[y_col].to_numpy(), + c=_INCORRECT_COLOUR, + label="Incorrect", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.scatter( + cor[x_col].to_numpy(), + cor[y_col].to_numpy(), + c=_CORRECT_COLOUR, + label="Correct", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.set_xlabel(x_label or x_col.replace("_", " ").title()) + ax.set_ylabel(y_label or y_col.replace("_", " ").title()) + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +# ── Main logic ──────────────────────────────────────────────────────── + + +def _load_residue_masses() -> dict: + cfg_path = REPO_ROOT / "winnow" / "configs" / "residues.yaml" + with open(cfg_path) as f: + return yaml.safe_load(f)["residue_masses"] + + +def _load_data(predictions_dir: Path) -> pl.DataFrame: + preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") + meta_path = predictions_dir / "metadata.csv" + if meta_path.exists(): + meta = pl.read_csv(meta_path) + preds = preds.join(meta, on="spectrum_id", how="inner") + return preds + + +_SPLIT_DISPLAY_NAMES = { + "test": "test set", + "unlabelled": "unlabelled space", + "raw_less_train": "full search space", +} + + +def _split_display(split: str) -> str: + key = split.strip().replace("-", "_") + return _SPLIT_DISPLAY_NAMES.get(key, split) + + +def _dns_model_tag(dns_model: str | None) -> str: + return f" ({dns_model})" if dns_model else "" + + +def _split_title(split_label: str, dns_model: str | None = None) -> str: + return f"{split_label}{_dns_model_tag(dns_model)}" + + +def _eval_title( + prefix: str, + split_label: str, + eval_kind: str, + dns_model: str | None = None, +) -> str: + return f"{prefix} {_split_title(split_label, dns_model)}\nusing {eval_kind}" + + +def _save_training_history_plot(model_dir: Path, out_dir: Path) -> None: + hist_path = model_dir / "training_history.json" + if not hist_path.exists(): + return + print("Plotting training history") + th = TrainingHistory.load(str(hist_path)) + th.plot(output_path=out_dir / "training_history.png", show=False) + print(" saved training_history") + + +def _plot_calibration_and_pr( + df: pl.DataFrame, + df_raw_conf: pl.DataFrame, + split: str, + labelled: bool, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + print("Plotting calibration curves") + if labelled: + fig = plot_calibration_curves( + df, + "correct", + _eval_title( + "Calibration curves for", split_label, "database search", dns_model + ), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"calibration_{split}_db_search") + + fig = plot_calibration_curves( + df, + "proteome_hit", + _eval_title( + "Calibration curves for", split_label, "proteome mapping", dns_model + ), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"calibration_{split}_proteome") + + print("Plotting PR curves") + if labelled: + fig = plot_pr_curves( + df, + "correct", + _eval_title("PR curves for", split_label, "database search", dns_model), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"pr_{split}_db_search") + + fig = plot_pr_curves( + df, + "proteome_hit", + _eval_title("PR curves for", split_label, "proteome mapping", dns_model), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"pr_{split}_proteome") + + +def _plot_confidence_histograms( + df: pl.DataFrame, + df_raw_conf: pl.DataFrame, + split: str, + labelled: bool, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + print("Plotting confidence histograms") + for conf_col, conf_label, tag in [ + ("confidence", "Raw confidence", "raw"), + ("calibrated_confidence", "Calibrated confidence", "cal"), + ]: + hist_df = df_raw_conf if conf_col == "confidence" else df + if labelled: + fig = plot_confidence_histogram( + hist_df, + "correct", + conf_col, + conf_label, + _eval_title( + f"{conf_label} for", split_label, "database search", dns_model + ), + ) + _save(fig, out_dir, f"hist_{tag}_{split}_db_search") + + fig = plot_confidence_histogram( + hist_df, + "proteome_hit", + conf_col, + conf_label, + _eval_title( + f"{conf_label} for", split_label, "proteome mapping", dns_model + ), + ) + _save(fig, out_dir, f"hist_{tag}_{split}_proteome") + + +def _plot_fdr_accuracy_plots( + df: pl.DataFrame, + split: str, + labelled: bool, + residue_masses: dict, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + print("Plotting FDR accuracy") + use_shortcut = not labelled + for metric, tag in [("fdr", "fdr"), ("q_value", "qvalue")]: + metric_name = "FDR" if metric == "fdr" else "Q-value" + if labelled: + fig = plot_fdr_accuracy( + df, + "correct", + residue_masses, + _eval_title( + f"{metric_name} accuracy for", + split_label, + "database search", + dns_model, + ), + metric="fdr" if metric == "fdr" else "q_value", + use_proteome_shortcut=False, + ) + _save(fig, out_dir, f"{tag}_{split}_db_search") + + fig = plot_fdr_accuracy( + df, + "proteome_hit", + residue_masses, + _eval_title( + f"{metric_name} accuracy for", + split_label, + "proteome mapping", + dns_model, + ), + metric="fdr" if metric == "fdr" else "q_value", + use_proteome_shortcut=use_shortcut, + ) + _save(fig, out_dir, f"{tag}_{split}_proteome") + + +def _plot_labelled_diagnostics( + df: pl.DataFrame, + df_raw_conf: pl.DataFrame, + split: str, + residue_masses: dict, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + title_split = _split_title(split_label, dns_model) + label_col = "correct" + print("Plotting labelled-only diagnostics") + + fig = plot_ranked_qvalue( + df, + label_col, + residue_masses, + _eval_title( + "Ranked predictions vs q-value for", + split_label, + "database search", + dns_model, + ), + ) + _save(fig, out_dir, f"ranked_qvalue_{split}_db_search") + + fig = plot_ranked_fdr_pep( + df, f"Ranked predictions vs FDR and PEP for {title_split}" + ) + _save(fig, out_dir, f"ranked_fdr_pep_{split}_nonparametric") + + for metric, file_tag, metric_name in [ + ("fdr", "fdr", "FDR"), + ("q_value", "qvalue", "q-value"), + ]: + fig = plot_ranked_fdr_raw_vs_cal( + df, + label_col, + residue_masses, + _eval_title( + f"Ranked predictions vs {metric_name} for", + split_label, + "database search", + dns_model, + ), + metric=metric, + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"ranked_{file_tag}_raw_vs_cal_{split}_db_search") + + fig = plot_bar_psms_fdr( + df, + label_col, + residue_masses, + f"PSMs at database-grounded FDR thresholds for {title_split}", + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"bar_psms_fdr_thresholds_{split}_db_search") + + if len(df_raw_conf) > 0: + fig = plot_raw_vs_cal_scatter( + df_raw_conf, + label_col, + f"Raw vs calibrated confidence for {title_split}", + ) + _save(fig, out_dir, f"scatter_raw_vs_cal_confidence_{split}") + + if "margin" in df.columns and len(df_raw_conf) > 0: + fig = plot_scatter_feature_vs_conf( + df_raw_conf, + label_col, + "margin", + "confidence", + f"Raw confidence vs margin for {title_split}", + x_label="Margin", + y_label="Raw confidence", + ) + _save(fig, out_dir, f"scatter_raw_confidence_vs_margin_{split}") + + fig = plot_scatter_feature_vs_conf( + df, + label_col, + "margin", + "calibrated_confidence", + f"Calibrated confidence vs margin for {title_split}", + x_label="Margin", + y_label="Calibrated confidence", + ) + _save(fig, out_dir, f"scatter_cal_confidence_vs_margin_{split}") + + fig, pca_model, feat_names = plot_pca_features( + df, label_col, f"PCA of calibrator features for {title_split}" + ) + _save(fig, out_dir, f"pca_features_{split}") + + fig = plot_pca_loadings( + pca_model, feat_names, "PCA loadings for first two principal components" + ) + _save(fig, out_dir, f"pca_loadings_pc1_pc2_{split}") + + +def main() -> None: + """Load predictions, annotate proteome hits, and write analysis plots.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--predictions-dir", + type=Path, + required=True, + help="Winnow predict output dir with preds_and_fdr_metrics.csv", + ) + parser.add_argument( + "--split", + type=str, + required=True, + help="Split id for filenames; titles use test set / unlabelled space / full search space", + ) + parser.add_argument( + "--label-mode", + choices=["labelled", "unlabelled"], + required=True, + help="labelled = has 'correct' column; unlabelled = proteome mapping only", + ) + parser.add_argument( + "--fasta", type=Path, required=True, help="FASTA file for proteome annotation" + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + help="Output directory for plots (defaults to predictions-dir/plots/)", + ) + parser.add_argument( + "--model-dir", + type=Path, + default=None, + help="Model directory for training history plot", + ) + parser.add_argument( + "--dns-model", + type=str, + default=None, + help="Upstream DNS model name for plot titles (e.g. InstaNovo, Casanovo, $\\pi$-PrimeNovo)", + ) + args = parser.parse_args() + + out_dir = args.output_dir or (args.predictions_dir / "plots") + out_dir.mkdir(parents=True, exist_ok=True) + + split = args.split + dns_model = args.dns_model + labelled = args.label_mode == "labelled" + residue_masses = _load_residue_masses() + + print(f"Loading predictions from {args.predictions_dir}") + df = _load_data(args.predictions_dir) + + from instanovo.utils.metrics import Metrics + from instanovo.utils.residues import ResidueSet + + metrics = Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + print(f"Annotating with proteome hits from {args.fasta}") + haystack = load_proteome_haystack(str(args.fasta)) + df = filter_and_annotate_preds(df, haystack, metrics, min_residue_length=7) + df_raw_conf = _df_for_raw_confidence_plots(df) + + if args.model_dir is not None: + _save_training_history_plot(args.model_dir, out_dir) + + _plot_calibration_and_pr( + df, df_raw_conf, split, labelled, out_dir, dns_model=dns_model + ) + _plot_confidence_histograms( + df, df_raw_conf, split, labelled, out_dir, dns_model=dns_model + ) + _plot_fdr_accuracy_plots( + df, split, labelled, residue_masses, out_dir, dns_model=dns_model + ) + + if labelled: + _plot_labelled_diagnostics( + df, df_raw_conf, split, residue_masses, out_dir, dns_model=dns_model + ) + + print(f"\nAll plots saved to {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_calibrator_generalisation_heatmap.py b/scripts/plot_calibrator_generalisation_heatmap.py new file mode 100644 index 00000000..e82837da --- /dev/null +++ b/scripts/plot_calibrator_generalisation_heatmap.py @@ -0,0 +1,265 @@ +"""Plot PR-AUC heatmaps for calibrator generalisation results. + +Reads the combined CSV produced by ``evaluate_calibrator_generalisation.py`` +and creates heatmaps comparing raw vs calibrated confidence PR-AUC values. +""" + +import logging +import sys +from pathlib import Path +from typing import Annotated + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +from matplotlib.colors import LinearSegmentedColormap +from rich.logging import RichHandler +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.calibrator_generalisation_utils import SPECIES_NAME_MAPPING # noqa: E402 + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +logger = logging.getLogger("winnow.plot_generalisation_heatmap") +logger.setLevel(logging.INFO) +logger.propagate = False +logger.addHandler(RichHandler()) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette + "sunset" diverging colourmap +# --------------------------------------------------------------------------- +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +_SUNSET_COLORS = [ + "#364B9A", + "#4A7BB7", + "#6EA6CD", + "#98CAE1", + "#C2E4EF", + "#EAECCC", + "#FEDA8B", + "#FDB366", + "#F67E4B", + "#DD3D2D", + "#A50026", +] +_BAD_COLOUR = "#FFFFFF" + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=2) + + +def _diverging_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) + cmap.set_bad(color=_BAD_COLOUR) + return cmap + + +def _sequential_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list( + "tol_sunset_seq", _SUNSET_COLORS[5:], N=256 + ) + cmap.set_bad(color=_BAD_COLOUR) + return cmap + + +# --------------------------------------------------------------------------- +# PR-AUC computation +# --------------------------------------------------------------------------- +def compute_pr_auc( + input_dataset: pd.DataFrame, + confidence_column: str, + label_column: str, +) -> float: + """Compute Area Under Curve for precision-recall curve.""" + if len(input_dataset) == 0: + return 0.0 + + sorted_data = input_dataset[[confidence_column, label_column]].sort_values( + by=confidence_column, ascending=False + ) + + cum_correct = np.cumsum(sorted_data[label_column]) + precision = cum_correct / np.arange(1, len(sorted_data) + 1) + recall = ( + cum_correct / cum_correct.iloc[-1] + if cum_correct.iloc[-1] > 0 + else np.zeros_like(cum_correct) + ) + + if len(precision) < 2: + return 0.0 + + from sklearn.metrics import auc + + return auc(recall, precision) + + +# --------------------------------------------------------------------------- +# Heatmap creation +# --------------------------------------------------------------------------- +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + """Save figure as both PNG and PDF.""" + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def create_auc_heatmap( + auc_df: pd.DataFrame, + output_path: Path, + title: str = "Calibrator generalisation PR-AUC heatmap", +) -> None: + """Create and save a heatmap of PR-AUC values.""" + fig, ax = plt.subplots(figsize=(12, 10)) + + sns.heatmap( + auc_df, + annot=True, + fmt=".3f", + cmap=_sequential_cmap(), + cbar_kws={"label": "PR-AUC"}, + square=True, + linewidths=0.5, + ax=ax, + ) + + ax.set_title(title) + ax.set_xlabel("Test dataset") + ax.set_ylabel("Train dataset") + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") + ax.set_yticklabels(ax.get_yticklabels(), rotation=0) + + base = str(output_path).removesuffix(".png") + _save_fig(fig, Path(base)) + logger.info("Heatmap saved to %s", output_path) + + +def create_comparison_heatmaps(results_path: Path, output_dir: Path) -> None: + """Create heatmaps comparing raw vs calibrated confidence PR-AUC values.""" + logger.info("Scanning results from %s", results_path) + results = pl.scan_csv(results_path) + + trained_datasets = sorted( + results.select(pl.col("trained_on_dataset")) + .unique() + .collect() + .to_series() + .to_list() + ) + test_datasets = sorted( + results.select(pl.col("test_dataset")).unique().collect().to_series().to_list() + ) + logger.info("Trained datasets: %s", trained_datasets) + logger.info("Test datasets: %s", test_datasets) + + trained_labels = [SPECIES_NAME_MAPPING.get(ds, ds) for ds in trained_datasets] + test_labels = [SPECIES_NAME_MAPPING.get(ds, ds) for ds in test_datasets] + + # Compute PR-AUC matrices for both confidence types + auc_matrices = {} + for conf_type in ["confidence", "calibrated_confidence"]: + auc_matrix = [] + for trained_dataset in trained_datasets: + auc_row = [] + for test_dataset in test_datasets: + logger.info( + "Computing PR-AUC (%s) for trained=%s, test=%s", + conf_type, + trained_dataset, + test_dataset, + ) + subset = ( + results.filter( + (pl.col("trained_on_dataset") == trained_dataset) + & (pl.col("test_dataset") == test_dataset) + ) + .collect() + .to_pandas() + ) + + if len(subset) > 0: + auc_row.append(compute_pr_auc(subset, conf_type, "correct")) + else: + auc_row.append(np.nan) + auc_matrix.append(auc_row) + + auc_matrices[conf_type] = pd.DataFrame( + auc_matrix, index=trained_labels, columns=test_labels + ) + + # Individual heatmaps + for conf_type, auc_df in auc_matrices.items(): + conf_name = conf_type.replace("_", " ") + output_path = ( + output_dir / f"calibrator_generalisation_{conf_type}_auc_heatmap.png" + ) + create_auc_heatmap( + auc_df, + output_path, + f"Calibrator generalisation {conf_name} PR-AUC", + ) + + # Difference heatmap (calibrated - raw) + diff_matrix = auc_matrices["calibrated_confidence"] - auc_matrices["confidence"] + + fig, ax = plt.subplots(figsize=(12, 10)) + sns.heatmap( + diff_matrix, + annot=True, + fmt=".3f", + cmap=_diverging_cmap(), + center=0, + cbar_kws={"label": r"PR-AUC difference $(\mathrm{calibrated} - \mathrm{raw})$"}, + square=True, + linewidths=0.5, + ax=ax, + ) + ax.set_title("Calibrator generalisation PR-AUC improvement") + ax.set_xlabel("Test dataset") + ax.set_ylabel("Train dataset") + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") + ax.set_yticklabels(ax.get_yticklabels(), rotation=0) + + diff_base = output_dir / "calibrator_generalisation_auc_difference_heatmap" + _save_fig(fig, diff_base) + logger.info("Difference heatmap saved to %s", diff_base) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +_DEFAULT_OUTPUT_DIR = Path("results/generalisation/plots") + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +@app.command() +def main( + results_path: Annotated[ + Path, typer.Option(help="Path to calibrator generalisation results CSV.") + ], + output_dir: Annotated[ + Path, typer.Option(help="Directory to save plots.") + ] = _DEFAULT_OUTPUT_DIR, +) -> None: + """Create PR-AUC heatmaps for calibrator generalisation results.""" + output_dir.mkdir(parents=True, exist_ok=True) + + if not results_path.exists(): + logger.error("Results file not found: %s", results_path) + raise typer.Exit(1) + + logger.info("Loading results from: %s", results_path) + logger.info("Saving plots to: %s", output_dir) + + create_comparison_heatmaps(results_path, output_dir) + + +if __name__ == "__main__": + app() diff --git a/scripts/plot_eval_results.py b/scripts/plot_eval_results.py new file mode 100644 index 00000000..32bccbfa --- /dev/null +++ b/scripts/plot_eval_results.py @@ -0,0 +1,984 @@ +"""Generate publication-quality evaluation plots from ``winnow predict`` outputs. + +Supports both annotated (database-grounded) and raw (proteome-hit) evaluation +modes, producing six plots per project: precision-recall, FDR run, true vs +estimated FDR (full + zoomed), probability calibration, and before/after score +histograms. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Annotated +import warnings + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer +from rich.logging import RichHandler + +from winnow.fdr.nonparametric import NonParametricFDRControl + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# Filter by the specific message or category +warnings.filterwarnings("ignore", message=".*range of fitted confidence scores.*") +warnings.filterwarnings("ignore", message=".*range of fitted FDR thresholds.*") + +# --------------------------------------------------------------------------- +# Dataset display names +# --------------------------------------------------------------------------- +DATASET_DISPLAY_NAMES: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD004732": "ProteomeTools-1", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", +} + +_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") +_DISPLAY_NAME_LOOKUP = {k.lower(): v for k, v in DATASET_DISPLAY_NAMES.items()} + +# Paul Tol "bright" palette (colour-blind safe) +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_MAIN_LINE_COLOUR = _PALETTE[0] +_RAW_LINE_COLOUR = _PALETTE[5] +_IDEAL_LINE_COLOUR = _PALETTE[6] +_BAND_COLOUR = _PALETTE[0] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_DIAGNOSTIC_ALPHAS = (0.01, 0.05, 0.10) +_HOEFFDING_DELTA = 0.05 + + +def _normalize_project_key(key: str) -> str: + """Strip eval suffixes and nested path segments for display-name lookup.""" + key = key.strip() + if "/" in key: + key = key.rsplit("/", 1)[-1] + for suffix in _FOLDER_SUFFIXES: + if key.endswith(suffix): + return key[: -len(suffix)] + return key + + +def _display_name(key: str) -> str: + """Look up the publication-ready display name for a dataset key.""" + normalized = _normalize_project_key(key) + if normalized in DATASET_DISPLAY_NAMES: + return DATASET_DISPLAY_NAMES[normalized] + return _DISPLAY_NAME_LOOKUP.get(normalized.lower(), normalized) + + +def _ground_truth_qualifier(eval_type: str) -> str: + """Return the title-friendly ground truth qualifier for plot titles.""" + if eval_type in ("annotated", "labelled"): + return "using database search" + return "using proteome mapping" + + +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + """Save figure as both PNG and PDF.""" + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +# --------------------------------------------------------------------------- +# PR curve (non-standard cumulative definition) +# --------------------------------------------------------------------------- +def _compute_precision_recall( + df: pd.DataFrame, confidence_col: str = "calibrated_confidence" +) -> pd.DataFrame: + """Non-standard cumulative PR curve matching the codebase convention.""" + sorted_df = df.sort_values(confidence_col, ascending=False) + labels = sorted_df["correct"].values + cum_correct = np.cumsum(labels) + n = len(labels) + precision = cum_correct / np.arange(1, n + 1) + recall = cum_correct / n + return pd.DataFrame({"precision": precision, "recall": recall}) + + +def plot_precision_recall( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot precision-recall curve.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + pr_cal = _compute_precision_recall(df, "calibrated_confidence") + pr_raw = _compute_precision_recall(df, "confidence") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + pr_raw["recall"], + pr_raw["precision"], + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Raw confidence", + ) + ax.plot( + pr_cal["recall"], + pr_cal["precision"], + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Calibrated confidence", + ) + ax.set_xlabel("Recall") + ax.set_ylabel("Precision") + ax.set_title(f"{display} precision-recall {qualifier}") + ax.set_xlim(0, 1) + ax.set_ylim(0, 1.02) + ax.legend(loc="lower right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"pr_curve_{project}") + + +# --------------------------------------------------------------------------- +# FDR run plot +# --------------------------------------------------------------------------- +def plot_fdr_run( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot calibrated confidence vs estimated and true PSM FDR.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + true_fdr_df = true_fdr_ctrl.add_psm_fdr( + df.copy(), confidence_col="calibrated_confidence" + ) + true_fdr_df = true_fdr_df.sort_values("calibrated_confidence") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + df["calibrated_confidence"].values, + df["psm_fdr"].values, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_fdr_df["calibrated_confidence"].values, + true_fdr_df["psm_fdr"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM FDR") + ax.set_title(f"{display} FDR run {qualifier}") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_run_{project}") + + +# --------------------------------------------------------------------------- +# Q-value run plot +# --------------------------------------------------------------------------- +def _fit_database_grounded_fdr( + df: pd.DataFrame, + confidence_col: str = "calibrated_confidence", + correct_col: str = "correct", + drop: int = 10, +) -> NonParametricFDRControl: + """Fit an FDR controller using ground-truth labels. + + Replicates the fitting logic of ``DatabaseGroundedFDRControl`` (computing + FDR as 1 − precision over sorted predictions, with the first *drop* entries + removed) without pulling in the instanovo dependency. + """ + sorted_desc = df.sort_values(confidence_col, ascending=False) + labels = sorted_desc[correct_col].values.astype(float) + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + confidence = sorted_desc[confidence_col].values + + ctrl = NonParametricFDRControl() + ctrl._fdr_values = (1.0 - precision)[drop:] + ctrl._confidence_scores = confidence[drop:] + return ctrl + + +def plot_q_value_run( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot calibrated confidence vs estimated and true PSM q-values.""" + if "psm_q_value" not in df.columns: + logger.warning( + "Skipping q-value run plot for %s: psm_q_value column missing", project + ) + return + + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + sorted_df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + qval_input = df[["calibrated_confidence"]].copy() + true_q_df = true_fdr_ctrl.add_psm_q_value( + qval_input, confidence_col="calibrated_confidence" + ) + true_q_df = true_q_df.sort_values("calibrated_confidence") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + sorted_df["calibrated_confidence"].values, + sorted_df["psm_q_value"].values, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_q_df["calibrated_confidence"].values, + true_q_df["psm_q_value"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM q-value") + ax.set_title(f"{display} q-value run {qualifier}") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"qvalue_run_{project}") + + +# --------------------------------------------------------------------------- +# FDR / q-value run plots with Hoeffding confidence bands +# --------------------------------------------------------------------------- +def _hoeffding_band_arrays(n: int) -> np.ndarray: + """Compute pointwise Hoeffding half-widths for ranks 1..n (descending confidence).""" + ranks = np.arange(1, n + 1) + return np.sqrt(np.log(2.0 / _HOEFFDING_DELTA) / (2.0 * ranks)) + + +def plot_fdr_run_with_bands( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """FDR run plot with Hoeffding 95% confidence band on the non-parametric curve.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + true_fdr_df = true_fdr_ctrl.add_psm_fdr( + df.copy(), confidence_col="calibrated_confidence" + ) + true_fdr_df = true_fdr_df.sort_values("calibrated_confidence") + + fdr_vals = df["psm_fdr"].values + conf_vals = df["calibrated_confidence"].values + n = len(fdr_vals) + hw = _hoeffding_band_arrays(n)[::-1] + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.fill_between( + conf_vals, + np.clip(fdr_vals - hw, 0, None), + np.clip(fdr_vals + hw, None, 1), + color=_BAND_COLOUR, + alpha=0.2, + label="95% Hoeffding bound", + ) + ax.plot( + conf_vals, + fdr_vals, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_fdr_df["calibrated_confidence"].values, + true_fdr_df["psm_fdr"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM FDR") + ax.set_title(f"{display} FDR run with sampling error bounds {qualifier}") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_run_bands_{project}") + + +def plot_q_value_run_with_bands( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Q-value run plot with Hoeffding 95% confidence band on the non-parametric curve.""" + if "psm_q_value" not in df.columns: + logger.warning( + "Skipping banded q-value run plot for %s: psm_q_value column missing", + project, + ) + return + + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + sorted_df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + qval_input = df[["calibrated_confidence"]].copy() + true_q_df = true_fdr_ctrl.add_psm_q_value( + qval_input, confidence_col="calibrated_confidence" + ) + true_q_df = true_q_df.sort_values("calibrated_confidence") + + qvals = sorted_df["psm_q_value"].values + conf_vals = sorted_df["calibrated_confidence"].values + n = len(qvals) + hw = _hoeffding_band_arrays(n)[::-1] + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.fill_between( + conf_vals, + np.clip(qvals - hw, 0, None), + np.clip(qvals + hw, None, 1), + color=_BAND_COLOUR, + alpha=0.2, + label="95% Hoeffding bound", + ) + ax.plot( + conf_vals, + qvals, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_q_df["calibrated_confidence"].values, + true_q_df["psm_q_value"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM q-value") + ax.set_title(f"{display} q-value run with sampling error bounds {qualifier}") + ax.legend(loc="upper center") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"qvalue_run_bands_{project}") + + +# --------------------------------------------------------------------------- +# True FDR vs estimated FDR +# --------------------------------------------------------------------------- +def _compute_true_vs_estimated_fdr(df: pd.DataFrame) -> pd.DataFrame: + """Compute true and estimated FDR arrays, sorted by confidence descending.""" + sorted_df = df.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + + true_fdr_ctrl = _fit_database_grounded_fdr(sorted_df) + with_true_fdr = true_fdr_ctrl.add_psm_fdr( + sorted_df, confidence_col="calibrated_confidence" + ) + + return pd.DataFrame( + { + "estimated_fdr": sorted_df["psm_fdr"].values, + "true_fdr": with_true_fdr["psm_fdr"].values, + } + ) + + +def plot_true_vs_estimated_fdr( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, + *, + zoomed: bool = False, +) -> None: + """Plot true FDR vs estimated FDR.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + fdr_data = _compute_true_vs_estimated_fdr(df) + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + fdr_data["estimated_fdr"], + fdr_data["true_fdr"], + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Observed", + ) + + # Only plot the ideal line up to the max extent of the observed line + max_x = float(fdr_data["estimated_fdr"].max()) + max_y = float(fdr_data["true_fdr"].max()) + lim = 0.1 if zoomed else 1.0 + ideal_end = min(lim, max(max_x, max_y)) + + ax.plot( + [0, ideal_end], + [0, ideal_end], + ls="--", + color=_IDEAL_LINE_COLOUR, + lw=1, + label="Perfectly calibrated", + ) + ax.set_xlabel("Non-parametric estimated FDR") + ax.set_ylabel("Database-grounded FDR") + zoom_suffix = " (0 to 0.1)" if zoomed else "" + ax.set_title(f"{display} true vs estimated FDR{zoom_suffix} {qualifier}") + if zoomed: + ax.set_xlim(0, 0.1) + ax.set_ylim(0, 0.1) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + tag = "fdr_true_vs_est_zoom" if zoomed else "fdr_true_vs_est" + _save_fig(fig, output_dir / f"{tag}_{project}") + + +# --------------------------------------------------------------------------- +# True q-values vs estimated q-values +# --------------------------------------------------------------------------- +def _compute_true_vs_estimated_q_values(df: pd.DataFrame) -> pd.DataFrame: + """Compute true and estimated q-value arrays, sorted by confidence descending.""" + sorted_df = df.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + + true_q_val_ctrl = _fit_database_grounded_fdr(sorted_df) + qval_input = sorted_df[["calibrated_confidence"]].copy() + with_true_q_df = true_q_val_ctrl.add_psm_q_value( + qval_input, confidence_col="calibrated_confidence" + ) + return pd.DataFrame( + { + "estimated_q_value": sorted_df["psm_q_value"].values, + "true_q_value": with_true_q_df["psm_q_value"].values, + } + ) + + +def plot_true_vs_estimated_q_values( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, + *, + zoomed: bool = False, +) -> None: + """Plot true q-values vs estimated q-values.""" + if "psm_q_value" not in df.columns: + logger.warning( + "Skipping true vs estimated q-value plot for %s: psm_q_value column missing", + project, + ) + return + + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + q_value_data = _compute_true_vs_estimated_q_values(df) + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + q_value_data["estimated_q_value"], + q_value_data["true_q_value"], + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Observed", + ) + + # Only plot the ideal line up to the max extent of the observed line + max_x = float(q_value_data["estimated_q_value"].max()) + max_y = float(q_value_data["true_q_value"].max()) + lim = 0.1 if zoomed else 1.0 + ideal_end = min(lim, max(max_x, max_y)) + + ax.plot( + [0, ideal_end], + [0, ideal_end], + ls="--", + color=_IDEAL_LINE_COLOUR, + lw=1, + label="Perfectly calibrated", + ) + ax.set_xlabel("Non-parametric estimated q-values") + ax.set_ylabel("Database-grounded q-values") + zoom_suffix = " (0 to 0.1)" if zoomed else "" + ax.set_title(f"{display} true vs estimated q-values{zoom_suffix} {qualifier}") + if zoomed: + ax.set_xlim(0, 0.1) + ax.set_ylim(0, 0.1) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + tag = "qvalue_true_vs_est_zoom" if zoomed else "qvalue_true_vs_est" + _save_fig(fig, output_dir / f"{tag}_{project}") + + +# --------------------------------------------------------------------------- +# Probability calibration (reliability diagram) +# --------------------------------------------------------------------------- +def _compute_calibration_curve( + df: pd.DataFrame, + pred_col: str, + label_col: str, + n_bins: int = 10, +) -> pd.DataFrame: + """Fixed-width bin calibration curve.""" + data = df[[pred_col, label_col]].dropna().copy() + data[pred_col] = data[pred_col].clip(0.0, 1.0) + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_cats = pd.cut(data[pred_col], bins=bins, include_lowest=True) + bin_cats.name = "bin" + grouped = ( + data.groupby(bin_cats, observed=True) + .agg( + pred_mean=(pred_col, "mean"), + empirical=(label_col, "mean"), + count=(label_col, "size"), + ) + .reset_index() + ) + grouped = grouped[grouped["count"] > 0] + grouped["bin_center"] = grouped["bin"].apply(lambda iv: (iv.left + iv.right) / 2) + return grouped[["pred_mean", "empirical", "count", "bin_center"]] + + +def _estimate_calibration_values( + df: pd.DataFrame, + pred_col: str, + label_col: str, + n_bins: int = 20, +) -> np.ndarray: + """Estimate c(s) for each PSM via binned calibration. + + Returns an array of the same length as *df* where each entry is the + empirical accuracy of the bin that PSM falls into. + """ + scores = df[pred_col].values.clip(0.0, 1.0) + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_idx = np.digitize(scores, bins) - 1 + bin_idx = np.clip(bin_idx, 0, n_bins - 1) + labels = df[label_col].values.astype(float) + bin_sums = np.bincount(bin_idx, weights=labels, minlength=n_bins) + bin_counts = np.bincount(bin_idx, minlength=n_bins).astype(float) + bin_counts[bin_counts == 0] = 1.0 + bin_means = bin_sums / bin_counts + return bin_means[bin_idx] + + +def _hoeffding_halfwidth(k: int, delta: float = _HOEFFDING_DELTA) -> float: + """Hoeffding 95% confidence half-width for a mean of *k* bounded [0,1] r.v.s.""" + if k <= 0: + return float("nan") + return float(np.sqrt(np.log(2.0 / delta) / (2.0 * k))) + + +def plot_calibration( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot probability calibration (reliability diagram).""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + cal_calibrated = _compute_calibration_curve(df, "calibrated_confidence", "correct") + cal_raw = _compute_calibration_curve(df, "confidence", "correct") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + cal_raw["pred_mean"], + cal_raw["empirical"], + marker="D", + color=_RAW_LINE_COLOUR, + label="Raw confidence", + ) + ax.plot( + cal_calibrated["pred_mean"], + cal_calibrated["empirical"], + marker="o", + color=_MAIN_LINE_COLOUR, + label="Calibrated confidence", + ) + ax.plot( + [0, 1], + [0, 1], + ls="--", + color=_IDEAL_LINE_COLOUR, + lw=1, + label="Perfectly calibrated", + ) + ax.set_xlabel("Mean predicted probability") + ax.set_ylabel("Empirical accuracy") + ax.set_title(f"{display} probability calibration {qualifier}") + ax.set_xlim(0, 1) + ax.set_ylim(0, 1) + ax.legend(loc="lower right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"calibration_{project}") + + +# --------------------------------------------------------------------------- +# Before/after score histograms +# --------------------------------------------------------------------------- +def plot_score_histograms( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot before/after score histograms with correct/incorrect overlays.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + correct_mask = df["correct"].astype(bool) + + fig, axes = plt.subplots(2, 1, figsize=(7, 6), sharex=False) + + # Before calibration + ax = axes[0] + bins_before = np.linspace(0, 1, 51) + ax.hist( + df.loc[correct_mask, "confidence"], + bins=bins_before, + alpha=0.5, + color=_CORRECT_COLOUR, + edgecolor="black", + label="Correct", + ) + ax.hist( + df.loc[~correct_mask, "confidence"], + bins=bins_before, + alpha=0.5, + color=_INCORRECT_COLOUR, + edgecolor="black", + label="Incorrect", + ) + ax.set_xlabel("Raw confidence") + ax.set_ylabel("Count") + ax.set_title("Before calibration") + ax.legend(loc="upper center") + _style_ax(ax) + + # After calibration + ax = axes[1] + bins_after = np.linspace(0, 1, 51) + ax.hist( + df.loc[correct_mask, "calibrated_confidence"], + bins=bins_after, + alpha=0.5, + color=_CORRECT_COLOUR, + edgecolor="black", + label="Correct", + ) + ax.hist( + df.loc[~correct_mask, "calibrated_confidence"], + bins=bins_after, + alpha=0.5, + color=_INCORRECT_COLOUR, + edgecolor="black", + label="Incorrect", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("Count") + ax.set_title("After calibration") + ax.legend(loc="upper center") + _style_ax(ax) + + fig.suptitle(f"{display} score distributions {qualifier}", fontsize=13) + fig.tight_layout() + _save_fig(fig, output_dir / f"score_histograms_{project}") + + +# --------------------------------------------------------------------------- +# Diagnostics CSV +# --------------------------------------------------------------------------- +def _is_labelled(eval_type: str) -> bool: + return eval_type in ("annotated", "labelled") + + +def _compute_diagnostics( + df: pd.DataFrame, + eval_type: str, + alphas: tuple[float, ...] = _DIAGNOSTIC_ALPHAS, +) -> pd.DataFrame: + """Compute FDR diagnostics at each target alpha. + + Label-dependent metrics (sTECE, TECE, realised FDR, etc.) are only + populated for annotated/labelled eval types. + """ + labelled = _is_labelled(eval_type) + + np_ctrl = NonParametricFDRControl() + np_ctrl.fit(dataset=df["calibrated_confidence"]) + + if labelled: + db_ctrl = _fit_database_grounded_fdr(df) + c_hat = _estimate_calibration_values(df, "calibrated_confidence", "correct") + scores = df["calibrated_confidence"].values.clip(0.0, 1.0) + + rows: list[dict] = [] + for alpha in alphas: + tau_hat = np_ctrl.get_confidence_cutoff(threshold=alpha) + if np.isnan(tau_hat): + rows.append({"alpha": alpha, "tau_hat": float("nan")}) + continue + + mask_hat = df["calibrated_confidence"].values >= tau_hat + k = int(mask_hat.sum()) + est_fdr = float(np_ctrl.compute_fdr(tau_hat)) + eps = _hoeffding_halfwidth(k) + + row: dict = { + "alpha": alpha, + "tau_hat": float(tau_hat), + "k_accepted": k, + "estimated_fdr": est_fdr, + "hoeffding_halfwidth": eps, + } + + if labelled: + residuals = c_hat[mask_hat] - scores[mask_hat] + row["stece"] = float(np.mean(residuals)) + row["tece"] = float(np.mean(np.abs(residuals))) + row["tece_2"] = float(np.sqrt(np.mean(residuals**2))) + + realised_fdr = float(db_ctrl.compute_fdr(tau_hat)) + row["realised_fdr"] = realised_fdr + row["fdr_bias"] = est_fdr - realised_fdr + + tau_star = db_ctrl.get_confidence_cutoff(threshold=alpha) + row["tau_star"] = float(tau_star) + if not np.isnan(tau_star): + k_star = int((df["calibrated_confidence"].values >= tau_star).sum()) + row["discovery_count_shift"] = k - k_star + else: + row["discovery_count_shift"] = float("nan") + + rows.append(row) + + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- +def _load_project_data( + predictions_root: Path, + project: str, + suffix: str, + eval_type: str, +) -> pd.DataFrame: + """Load and merge metadata.csv and preds_and_fdr_metrics.csv for a project.""" + # folder = predictions_root / f"{project}_{suffix}" + folder = predictions_root / f"{project}" + preds_path = folder / "preds_and_fdr_metrics.csv" + meta_path = folder / "metadata.csv" + if not preds_path.is_file(): + raise FileNotFoundError(f"Missing predictions file: {preds_path}") + + preds_df = pd.read_csv(preds_path) + + if meta_path.is_file(): + meta_df = pd.read_csv(meta_path) + # Drop columns already present in preds to avoid duplicates on merge + overlap = [ + c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" + ] + if overlap: + meta_df = meta_df.drop(columns=overlap) + df = preds_df.merge(meta_df, on="spectrum_id", how="left") + else: + df = preds_df + + if eval_type in ("raw", "unlabelled"): + if "proteome_hit" not in df.columns: + raise ValueError( + f"Expected 'proteome_hit' column for eval-type={eval_type} in {preds_path}" + ) + df["correct"] = df["proteome_hit"].astype(float) + + required = ["confidence", "calibrated_confidence", "correct"] + missing = [c for c in required if c not in df.columns] + if missing: + raise ValueError(f"Missing columns {missing} in {preds_path}") + + return df + + +def generate_all_plots( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Generate all plots for a single project.""" + output_dir.mkdir(parents=True, exist_ok=True) + + plot_precision_recall(df, project, eval_type, output_dir) + plot_fdr_run(df, project, eval_type, output_dir) + plot_fdr_run_with_bands(df, project, eval_type, output_dir) + plot_q_value_run(df, project, eval_type, output_dir) + plot_q_value_run_with_bands(df, project, eval_type, output_dir) + plot_true_vs_estimated_fdr(df, project, eval_type, output_dir, zoomed=False) + plot_true_vs_estimated_fdr(df, project, eval_type, output_dir, zoomed=True) + plot_true_vs_estimated_q_values(df, project, eval_type, output_dir, zoomed=False) + plot_true_vs_estimated_q_values(df, project, eval_type, output_dir, zoomed=True) + plot_calibration(df, project, eval_type, output_dir) + plot_score_histograms(df, project, eval_type, output_dir) + + +_EVAL_TYPE_SUFFIX: dict[str, str] = { + "annotated": "annotated", + "raw": "raw", + "labelled": "labelled", + "unlabelled": "unlabelled", +} + + +@app.command() +def main( + predictions_root: Annotated[ + Path, + typer.Option( + "--predictions-root", + help="Root directory containing per-project prediction folders.", + ), + ], + projects: Annotated[ + str, + typer.Option( + "--projects", + help="Space- or comma-separated project keys (e.g. 'helaqc,gluc' or 'helaqc gluc').", + ), + ], + eval_type: Annotated[ + str, + typer.Option( + "--eval-type", + help="Evaluation type: annotated, raw, labelled, or unlabelled.", + ), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory to save plots and summary CSVs."), + ], +) -> None: + """Generate evaluation plots from winnow predict outputs.""" + logging.basicConfig(level=logging.INFO, format="%(message)s", datefmt="%H:%M:%S") + + if eval_type not in _EVAL_TYPE_SUFFIX: + raise typer.BadParameter( + f"Unknown eval-type {eval_type!r}. Expected one of: {list(_EVAL_TYPE_SUFFIX)}" + ) + + project_list = [p.strip() for p in projects.replace(",", " ").split() if p.strip()] + if not project_list: + raise typer.BadParameter("No projects specified.") + + suffix = _EVAL_TYPE_SUFFIX[eval_type] + output_dir.mkdir(parents=True, exist_ok=True) + + for project in project_list: + display = _display_name(project) + logger.info("Processing %s (%s, eval-type=%s)...", project, display, eval_type) + + df = _load_project_data(predictions_root, project, suffix, eval_type) + logger.info(" Loaded %d rows", len(df)) + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + db_fdr = true_fdr_ctrl.add_psm_fdr( + df[["calibrated_confidence"]].copy(), confidence_col="calibrated_confidence" + ) + df["db_grounded_psm_fdr"] = db_fdr["psm_fdr"] + db_qval = true_fdr_ctrl.add_psm_q_value( + df[["calibrated_confidence"]].copy(), confidence_col="calibrated_confidence" + ) + df["db_grounded_psm_q_value"] = db_qval["psm_q_value"] + + summary_cols = ["confidence", "calibrated_confidence", "correct"] + if "psm_fdr" in df.columns: + summary_cols.append("psm_fdr") + summary_cols.append("db_grounded_psm_fdr") + if "psm_q_value" in df.columns: + summary_cols.append("psm_q_value") + summary_cols.append("db_grounded_psm_q_value") + df[summary_cols].to_csv(output_dir / f"{project}_summary.csv", index=False) + + diag = _compute_diagnostics(df, eval_type) + diag.to_csv(output_dir / f"{project}_diagnostics.csv", index=False) + logger.info(" Diagnostics saved (%d alpha levels)", len(diag)) + + generate_all_plots(df, project, eval_type, output_dir) + logger.info(" Plots saved to %s", output_dir) + + logger.info("Done. All plots saved to %s", output_dir) + + +if __name__ == "__main__": + app() diff --git a/scripts/plot_fdr_method_comparison.py b/scripts/plot_fdr_method_comparison.py new file mode 100644 index 00000000..bd83385b --- /dev/null +++ b/scripts/plot_fdr_method_comparison.py @@ -0,0 +1,1097 @@ +#!/usr/bin/env python3 +"""Compare PSM-level FDR estimates from Winnow and NovoBoard. + +Plots and summaries cover labelled-test Novor correctness and unlabelled +reference-proteome membership at 1 %, 5 %, and 10 % FDR on a shared filtered +spectrum pool (NovoBoard mass-deltas converted to ProForma; unsupported +modifications dropped; NovoBoard target-decoy pairs gated). Unlabelled panels +also drop normalised peptides shorter than 8 residues (proteome-substring +proxy); labelled panels keep short peptides because correctness is Novor +agreement. The shared pool is the twin-valid NovoBoard set; Winnow is trimmed +to match under the invariant that NovoBoard ⊆ Winnow after identical InstaNovo +filters. + +A long-form ``fdr_method_comparison_curves.csv`` (per spectrum x method) is +written so plots and summary tables can be regenerated with ``--summarise-only``. +""" + +from __future__ import annotations + +import logging +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Annotated, Literal, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.annotate_preds_proteome_hits import load_proteome_haystack # noqa: E402 +from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 + LABELLED_MIN_PEPTIDE_LENGTH, + MIN_PEPTIDE_LENGTH, + assert_shared_prediction_keys, + attach_labels_by_spectrum_id, + attach_novoboard_pair_keys, + compute_q_values, + filter_novoboard_target_decoy_pairs, + filter_prediction_table, + label_series_by_spectrum_id, + load_residue_masses, + novoboard_psm_tdc, + novor_correctness_mask, + proteome_hit_mask, + restrict_winnow_to_novoboard_spectra, +) +from scripts.fdr_tool_comparison_summaries import ( # noqa: E402 + SUMMARY_THRESHOLDS, + acceptance_rows_from_q, + error_rows_from_q, + finalise_error_gain_table, + write_summary_tables, +) +from scripts.plot_eval_results import ( # noqa: E402 + _MAIN_LINE_COLOUR, + _PALETTE, + _RAW_LINE_COLOUR, + _display_name, + _ground_truth_qualifier, + _save_fig, + _style_ax, +) +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +PRIMARY_METHOD = "Winnow (non-parametric)" +DB_CAL_METHOD = "Database-grounded (calibrated confidence)" +DB_RAW_METHOD = "Database-grounded (raw confidence)" +NOVOBOARD_METHOD = "NovoBoard" +CURVES_CSV_NAME = "fdr_method_comparison_curves.csv" +_WINNOW_METHODS = (PRIMARY_METHOD, DB_CAL_METHOD, DB_RAW_METHOD) +_METHOD_ORDER = (*_WINNOW_METHODS, NOVOBOARD_METHOD) + +logger = logging.getLogger(__name__) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] +_DB_GROUNDED_DROP = 10 + +DEFAULT_NOVOBOARD_ROOT = Path("/home/j-daniel/repos/NovoBoard/datasets") +DEFAULT_WINNOW_RESULTS = _REPO_ROOT / "results" +DEFAULT_MODEL_ROOT = _REPO_ROOT / "models" +DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/fdr_method_comparison_psm" +DEFAULT_DATASETS = ["helaqc", "celegans"] +_METHOD_COLOURS = { + PRIMARY_METHOD: _MAIN_LINE_COLOUR, + DB_CAL_METHOD: _RAW_LINE_COLOUR, + DB_RAW_METHOD: _PALETTE[3], + NOVOBOARD_METHOD: _PALETTE[2], +} + +EvalType = Literal["labelled", "unlabelled"] + +_DATASET_META = { + "helaqc": { + "fasta": "fasta/human.fasta", + "novoboard_decoy": "0.50", + "winnow_suffix": "helaqc", + }, + "celegans": { + "fasta": "fasta/Celegans.fasta", + "novoboard_decoy": "0.70", + "winnow_suffix": "celegans", + }, + "sbrodae": { + "fasta": "fasta/Sb_proteome.fasta", + "novoboard_decoy": "0.50", + "winnow_suffix": "sbrodae", + }, + "PXD019483": { + "fasta": "fasta/human.fasta", + "novoboard_decoy": "0.70", + "winnow_suffix": "pxd019483", + }, +} + + +@dataclass(frozen=True) +class DatasetConfig: + """Paths and metadata for one evaluation dataset.""" + + key: str + fasta: Path + winnow_unlabelled: Path + winnow_test: Path + novoboard_dir: Path + novoboard_decoy_rate: str + calibrator_train_metadata: Path + + +def build_dataset_configs( + winnow_results: Path = DEFAULT_WINNOW_RESULTS, + novoboard_root: Path = DEFAULT_NOVOBOARD_ROOT, + model_root: Path = DEFAULT_MODEL_ROOT, +) -> dict[str, DatasetConfig]: + """Build per-dataset path bundles from repo roots.""" + configs: dict[str, DatasetConfig] = {} + for key, meta in _DATASET_META.items(): + suffix = meta["winnow_suffix"] + configs[key] = DatasetConfig( + key=key, + fasta=_REPO_ROOT / meta["fasta"], + winnow_unlabelled=winnow_results + / f"instanovo_{suffix}_predictions_unlabelled", + winnow_test=winnow_results / f"instanovo_{suffix}_predictions_test", + novoboard_dir=novoboard_root / f"{key}/novoboard", + novoboard_decoy_rate=meta["novoboard_decoy"], + calibrator_train_metadata=model_root + / f"instanovo_{suffix}/metadata_train.parquet", + ) + return configs + + +@dataclass +class MethodCurve: + """One method's confidence and q-value arrays for curve plotting.""" + + label: str + color: str + confidence: np.ndarray + q_value: np.ndarray + + +@dataclass +class MethodCounts: + """PSM counts per q-value threshold for one method.""" + + label: str + color: str + counts: list[int] + + +@dataclass +class MethodRecovery: + """Correct labelled identifications recovered at q-value thresholds.""" + + label: str + color: str + q_value: np.ndarray + correct: np.ndarray + + +def _load_residue_masses() -> dict[str, float]: + return load_residue_masses() + + +def _fit_database_grounded_fdr( + df: pd.DataFrame, + correct_col: str, + confidence_col: str, + residue_masses: dict[str, float], + *, + drop: int = _DB_GROUNDED_DROP, +) -> DatabaseGroundedFDRControl: + """Fit ``DatabaseGroundedFDRControl`` (proteome shortcut or labelled sequence fit).""" + ctrl = DatabaseGroundedFDRControl( + confidence_feature=confidence_col, + residue_masses=residue_masses, + drop=drop, + ) + if correct_col == "proteome_hit": + sorted_df = df.sort_values(confidence_col, ascending=False) + labels = sorted_df[correct_col].astype(float).to_numpy() + conf = sorted_df[confidence_col].to_numpy() + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + ctrl._fdr_values = np.array(1.0 - precision)[drop:] + ctrl._confidence_scores = conf[drop:] + else: + fit_df = df.copy() + if "sequence" not in fit_df.columns or "prediction" not in fit_df.columns: + raise ValueError( + "Labelled database-grounded FDR requires 'sequence' and 'prediction' columns" + ) + ctrl.fit(dataset=fit_df, correct_column=correct_col) + return ctrl + + +def load_winnow( + predictions_dir: Path, fasta: Path, eval_type: EvalType +) -> pd.DataFrame: + """Load Winnow preds + metadata; annotate proteome hits or use labelled ``correct``. + + Always drops unsupported ``[UNIMOD:n]`` tokens. Unlabelled panels also require + normalised peptide length ≥ :data:`MIN_PEPTIDE_LENGTH` (proteome-substring + proxy). Labelled panels only require a non-empty normalised key + (:data:`LABELLED_MIN_PEPTIDE_LENGTH`), because correctness is Novor agreement. + + Labelled ``correct`` keeps Winnow predict-time Novor labels when present; + otherwise recomputes Novor from ``sequence`` / ``prediction``. + """ + preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") + meta_path = predictions_dir / "metadata.csv" + if meta_path.exists(): + meta = pl.read_csv(meta_path, columns=["spectrum_id", "confidence"]) + preds = preds.join(meta, on="spectrum_id", how="inner") + + df = preds.to_pandas() + min_length = ( + LABELLED_MIN_PEPTIDE_LENGTH if eval_type == "labelled" else MIN_PEPTIDE_LENGTH + ) + df = filter_prediction_table( + df, "prediction", min_length=min_length, key_col="peptide_key" + ) + + if eval_type == "labelled": + if "correct" in df.columns: + pass + elif {"sequence", "prediction"}.issubset(df.columns): + df["correct"] = novor_correctness_mask(df["sequence"], df["prediction"]) + else: + raise ValueError( + f"Missing 'correct' (and sequence/prediction) in " + f"{predictions_dir}/preds_and_fdr_metrics.csv" + ) + return df + + haystack = load_proteome_haystack(fasta) + df["proteome_hit"] = proteome_hit_mask( + df["prediction"], haystack, min_length=MIN_PEPTIDE_LENGTH + ) + return df + + +def _effective_db_grounded_drop(n_rows: int, drop: int = _DB_GROUNDED_DROP) -> int: + """Cap drop so FDR fit retains at least one score when *n_rows* is small.""" + return min(drop, max(0, n_rows - 1)) + + +def _vectorized_fdr_from_control( + confidence: np.ndarray, ctrl: DatabaseGroundedFDRControl | NonParametricFDRControl +) -> np.ndarray: + """Vectorized equivalent of ``FDRControl.compute_fdr`` for an array of scores.""" + if ctrl._confidence_scores is None or ctrl._fdr_values is None: + raise AttributeError("FDR method not fitted, please call `fit()` first") + conf = np.asarray(confidence, dtype=float) + scores = np.asarray(ctrl._confidence_scores, dtype=float) + fdr_values = np.asarray(ctrl._fdr_values, dtype=float) + n = len(scores) + idx = np.searchsorted(-scores, -conf, side="left") + fdr = np.empty(len(conf), dtype=float) + below = (idx == n) & (conf < scores[-1]) + above = (idx == 0) & (conf > scores[0]) + normal = ~(below | above) + fdr[below] = 1.0 + fdr[above] = float(fdr_values[0]) + clipped = np.clip(idx[normal], 0, n - 1) + fdr[normal] = fdr_values[clipped] + return fdr + + +def _assign_q_values_fast( + df: pd.DataFrame, + confidence_col: str, + ctrl: DatabaseGroundedFDRControl | NonParametricFDRControl, + out_col: str, +) -> pd.DataFrame: + """Assign q-values without per-row ``compute_fdr`` applies (needed for large tables).""" + work = df.copy() + conf = work[confidence_col].to_numpy(dtype=float) + fdr = _vectorized_fdr_from_control(conf, ctrl) + order = np.argsort(-conf, kind="mergesort") + q_sorted = compute_q_values(fdr[order]) + q = np.empty_like(q_sorted) + q[order] = q_sorted + work[out_col] = q + return work + + +def _add_database_grounded_qvalues( + df: pd.DataFrame, + correct_col: str, + confidence_col: str, + out_col: str, + residue_masses: dict[str, float], + *, + fit_df: pd.DataFrame | None = None, + drop: int = _DB_GROUNDED_DROP, +) -> pd.DataFrame: + """Append database-grounded PSM q-values; fit on *fit_df* (defaults to *df*).""" + reference = fit_df if fit_df is not None else df + work = df.drop(columns=[out_col], errors="ignore").copy() + ctrl = _fit_database_grounded_fdr( + reference, + correct_col, + confidence_col, + residue_masses, + drop=_effective_db_grounded_drop(len(reference), drop), + ) + return _assign_q_values_fast(work, confidence_col, ctrl, out_col) + + +def _prepare_winnow_psm_table( + df: pd.DataFrame, + correct_col: str, + residue_masses: dict[str, float], + *, + fit_df: pd.DataFrame | None = None, +) -> pd.DataFrame: + """Append Winnow PSM-level q-value columns while retaining labels.""" + reference = fit_df if fit_df is not None else df + db_cal = _add_database_grounded_qvalues( + df, + correct_col, + "calibrated_confidence", + "psm_q_value_db_cal", + residue_masses, + fit_df=reference, + ) + db_raw = _add_database_grounded_qvalues( + db_cal, + correct_col, + "confidence", + "psm_q_value_db_raw", + residue_masses, + fit_df=reference, + ) + return db_raw + + +def _curves_df_from_winnow_table( + table: pd.DataFrame, + *, + dataset: str, + panel: str, + label_col: str, +) -> pd.DataFrame: + """Long-form curve rows for the three Winnow PSM q-value methods.""" + if "spectrum_id" not in table.columns: + raise KeyError("Winnow curve export requires spectrum_id") + if label_col not in table.columns: + raise KeyError(f"Missing label column {label_col!r}") + label = table[label_col].astype(bool).to_numpy() + spectrum_id = table["spectrum_id"].astype(str) + specs = ( + (PRIMARY_METHOD, "calibrated_confidence", "psm_q_value"), + (DB_CAL_METHOD, "calibrated_confidence", "psm_q_value_db_cal"), + (DB_RAW_METHOD, "confidence", "psm_q_value_db_raw"), + ) + parts: list[pd.DataFrame] = [] + for method, score_col, q_col in specs: + if score_col not in table.columns or q_col not in table.columns: + raise KeyError(f"Missing {score_col!r} / {q_col!r} for {method}") + parts.append( + pd.DataFrame( + { + "dataset": dataset, + "panel": panel, + "method": method, + "spectrum_id": spectrum_id, + "score": table[score_col].to_numpy(dtype=float), + "q_value": table[q_col].to_numpy(dtype=float), + "label": label, + } + ) + ) + return pd.concat(parts, ignore_index=True) + + +def _curves_df_from_novoboard( + df: pd.DataFrame, + *, + dataset: str, + panel: str, + label_col: str, +) -> pd.DataFrame: + """Long-form curve rows for NovoBoard PSM TDC targets.""" + if "spectrum_id" not in df.columns: + raise KeyError("NovoBoard curve export requires spectrum_id") + if label_col not in df.columns: + raise KeyError(f"Missing label column {label_col!r}") + return pd.DataFrame( + { + "dataset": dataset, + "panel": panel, + "method": NOVOBOARD_METHOD, + "spectrum_id": df["spectrum_id"].astype(str), + "score": df["ALC (%)"].to_numpy(dtype=float), + "q_value": df["estimated_q_value"].to_numpy(dtype=float), + "label": df[label_col].astype(bool).to_numpy(), + } + ) + + +def load_novoboard_target_decoy( + novoboard_dir: Path, split: Literal["unlabelled", "test"], decoy_rate: str +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Load NovoBoard target/decoy tables and attach twin ``_pair_key`` values.""" + prefix = "raw_unlabelled" if split == "unlabelled" else "annotated_test" + target_path = novoboard_dir / f"{prefix}.csv" + decoy_path = novoboard_dir / f"{prefix}_decoy_{decoy_rate}.csv" + if not target_path.is_file(): + raise FileNotFoundError(target_path) + if not decoy_path.is_file(): + raise FileNotFoundError(decoy_path) + target = pd.read_csv(target_path) + decoy = pd.read_csv(decoy_path) + return attach_novoboard_pair_keys( + target, decoy, novoboard_dir=novoboard_dir, split_prefix=prefix + ) + + +def _restrict_winnow_to_novoboard_spectra( + winnow: pd.DataFrame, novoboard: pd.DataFrame +) -> pd.DataFrame: + """Trim Winnow to NovoBoard twin-valid spectra under the subset invariant.""" + return restrict_winnow_to_novoboard_spectra(winnow, novoboard) + + +def _assert_shared_prediction_keys( + winnow: pd.DataFrame, + novoboard: pd.DataFrame, + *, + winnow_peptide_col: str = "prediction", + novoboard_peptide_col: str = "Peptide", +) -> None: + """Require I/L-normalised prediction identity on the shared spectrum pool.""" + assert_shared_prediction_keys( + winnow, + novoboard, + winnow_peptide_col=winnow_peptide_col, + novoboard_peptide_col=novoboard_peptide_col, + ) + + +def _label_series_by_spectrum_id(winnow: pd.DataFrame, label_col: str) -> pd.Series: + """Map ``spectrum_id`` → boolean label from a Winnow table.""" + return label_series_by_spectrum_id(winnow, label_col) + + +def _attach_labels_by_spectrum_id( + novoboard: pd.DataFrame, + label_by_id: pd.Series, + *, + label_col: str, +) -> pd.DataFrame: + """Attach a shared label column to NovoBoard rows by ``spectrum_id``.""" + return attach_labels_by_spectrum_id(novoboard, label_by_id, label_col=label_col) + + +def _method_curves_from_panel(panel_df: pd.DataFrame) -> list[MethodCurve]: + """Rebuild plot curves from long-form curve rows for one panel.""" + curves: list[MethodCurve] = [] + for method in _METHOD_ORDER: + sub = panel_df.loc[panel_df["method"] == method] + if sub.empty: + continue + colour = _METHOD_COLOURS.get(str(method), _PALETTE[0]) + curves.append( + MethodCurve( + str(method), + colour, + sub["score"].to_numpy(dtype=float), + sub["q_value"].to_numpy(dtype=float), + ) + ) + return curves + + +def _recovery_series_from_panel(panel_df: pd.DataFrame) -> list[MethodRecovery]: + """Rebuild labelled recovery series from long-form curve rows.""" + series: list[MethodRecovery] = [] + for method in _METHOD_ORDER: + sub = panel_df.loc[panel_df["method"] == method] + if sub.empty: + continue + colour = _METHOD_COLOURS.get(str(method), _PALETTE[0]) + series.append( + MethodRecovery( + str(method), + colour, + sub["q_value"].to_numpy(dtype=float), + sub["label"].astype(bool).to_numpy(), + ) + ) + return series + + +def plot_dataset_from_curves( + curves: pd.DataFrame, dataset_key: str, output_dir: Path +) -> None: + """Write PSM comparison plots for one dataset from a curves table.""" + out = output_dir / dataset_key + out.mkdir(parents=True, exist_ok=True) + ds = curves.loc[curves["dataset"] == dataset_key] + if ds.empty: + raise ValueError(f"No curve rows for dataset {dataset_key!r}") + + panel_specs: tuple[tuple[str, EvalType, str], ...] = ( + ("unlabelled", "unlabelled", "unlabelled"), + ("labelled_test", "labelled", "test"), + ) + for panel, eval_type, stem in panel_specs: + panel_df = ds.loc[ds["panel"] == panel] + method_curves = _method_curves_from_panel(panel_df) + if not method_curves: + continue + plot_qvalue_by_rank( + method_curves, + dataset_key, + eval_type, + out / f"psm_qvalue_by_rank_{stem}_{dataset_key}", + ) + plot_threshold_barplot( + _bar_series_from_curves(method_curves), + dataset_key, + eval_type, + out / f"psm_counts_{stem}_{dataset_key}", + ) + + labelled = ds.loc[ds["panel"] == "labelled_test"] + recovery = _recovery_series_from_panel(labelled) + if recovery: + plot_recovery_curves( + recovery, + dataset_key, + out / f"psm_recovery_test_{dataset_key}", + ) + + +def write_curves_csv(curves: pd.DataFrame, output_dir: Path) -> Path: + """Write the long-form replot curves table.""" + output_dir.mkdir(parents=True, exist_ok=True) + path = output_dir / CURVES_CSV_NAME + # Preserve float64 q/score values so threshold edge cases survive round-trip. + curves.to_csv(path, index=False, float_format="%.17g") + logger.info("Wrote %s (%d rows)", path, len(curves)) + return path + + +def plot_qvalue_by_rank( + curves: list[MethodCurve], + dataset_key: str, + eval_type: EvalType, + output_path: Path, + *, + title_suffix: str = "", +) -> None: + """Plot q-value against native-score rank/accepted count.""" + display = _display_name(dataset_key) + qualifier = _ground_truth_qualifier( + "labelled" if eval_type == "labelled" else "unlabelled" + ) + title = f"{display} PSM q-value by rank {qualifier}{title_suffix}" + + fig, ax = plt.subplots(figsize=(8, 6)) + q_max = 0.0 + for curve in curves: + order = np.argsort(-np.asarray(curve.confidence, dtype=float)) + y = np.asarray(curve.q_value, dtype=float)[order] + rank = np.arange(1, len(y) + 1) + valid = ~np.isnan(y) + if not np.any(valid): + continue + q_max = max(q_max, float(np.nanmax(y[valid]))) + ax.plot( + rank[valid], + y[valid], + color=curve.color, + lw=1.5, + label=curve.label, + ) + + ax.set_xlabel("Accepted PSMs by native-score rank") + ax.set_ylabel("PSM q-value") + ax.set_title(title) + y_top = min(max(q_max * 1.15, 0.05), 1.0) + ax.set_ylim(0, y_top) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s", output_path) + + +def _count_at_thresholds( + q: np.ndarray, thresholds: list[float] = FDR_THRESHOLDS +) -> list[int]: + q = np.asarray(q, dtype=float) + valid = q[~np.isnan(q)] + return [int((valid <= t).sum()) for t in thresholds] + + +def _bar_series_from_curves(curves: list[MethodCurve]) -> list[MethodCounts]: + return [ + MethodCounts( + label=c.label, + color=c.color, + counts=_count_at_thresholds(c.q_value), + ) + for c in curves + ] + + +def _recovery_at_thresholds( + q: np.ndarray, + correct: np.ndarray, + thresholds: list[float] = FDR_THRESHOLDS, +) -> list[float]: + """Return correct-identification recovery percentage at each q-value threshold.""" + q = np.asarray(q, dtype=float) + correct = np.asarray(correct, dtype=bool) + denom = int(correct.sum()) + if denom == 0: + return [np.nan for _ in thresholds] + valid = ~np.isnan(q) + return [100.0 * int((valid & correct & (q <= t)).sum()) / denom for t in thresholds] + + +def plot_recovery_curves( + series: list[MethodRecovery], + dataset_key: str, + output_path: Path, +) -> None: + """Plot correct-identification recovery versus q-value threshold.""" + display = _display_name(dataset_key) + fig, ax = plt.subplots(figsize=(8, 6)) + for item in series: + y = _recovery_at_thresholds(item.q_value, item.correct) + ax.plot( + FDR_THRESHOLDS, + y, + marker="o", + lw=1.5, + label=item.label, + color=item.color, + ) + + ax.set_xlim(0, max(FDR_THRESHOLDS)) + ax.set_ylim(0, 100) + ax.set_xlabel("Estimated q-value threshold") + ax.set_ylabel("Correct PSM recovery\n(% of labelled correct PSMs)") + ax.set_title(f"{display} labelled PSM recovery by q-value threshold") + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s", output_path) + + +def plot_threshold_barplot( + series: list[MethodCounts], + dataset_key: str, + eval_type: EvalType, + output_path: Path, +) -> None: + """Bar chart of identifications retained at each q-value threshold.""" + display = _display_name(dataset_key) + if eval_type == "labelled": + split_label = "labelled test set" + else: + split_label = "unlabelled set" + title = f"{display}: accepted PSMs on the {split_label} at q-value thresholds" + ylabel = "Peptide-spectrum matches" + + n_methods = len(series) + n_thresh = len(FDR_THRESHOLDS) + group_spacing = 0.825 + cluster_width = min(0.75, group_spacing * 0.92) + width = cluster_width / n_methods + x = np.arange(n_thresh) * group_spacing + + fig_w = max(10.0, 2.2 * n_thresh * group_spacing) + fig, ax = plt.subplots(figsize=(fig_w, 7)) + for i, item in enumerate(series): + offset = (i - (n_methods - 1) / 2) * width + bars = ax.bar( + x + offset, + item.counts, + width, + label=item.label, + color=item.color, + edgecolor="black", + linewidth=1, + ) + for bar in bars: + h = bar.get_height() + ax.annotate( + f"{int(h):,}", + xy=(bar.get_x() + bar.get_width() / 2, h), + xytext=(0, 3), + textcoords="offset points", + ha="center", + va="bottom", + fontsize=9, + ) + + max_count = max((c for item in series for c in item.counts), default=1) + y_headroom = (1.55 + 0.06 * n_methods) * (2 / 3) + ax.set_ylim(0, max_count * y_headroom) + + half_cluster = cluster_width / 2 + ax.set_xlim( + -half_cluster - 0.25, + (n_thresh - 1) * group_spacing + half_cluster + 0.25, + ) + + ax.set_xlabel("Q-value threshold") + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.set_xticks(x) + ax.set_xticklabels([str(t) for t in FDR_THRESHOLDS]) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s", output_path) + + +def _append_method_summary_rows( + acceptance_rows: list[dict[str, object]], + error_rows: list[dict[str, object]], + *, + dataset: str, + panel: str, + level: str, + method: str, + q_value: np.ndarray, + label_mask: np.ndarray | None = None, + recovery_denom: int | None = None, + q_ref: np.ndarray | None = None, +) -> None: + """Append acceptance and error rows for one method.""" + acceptance_rows.extend( + acceptance_rows_from_q( + dataset=dataset, + panel=panel, + level=level, + method=method, + q_value=q_value, + thresholds=SUMMARY_THRESHOLDS, + label_mask=label_mask, + recovery_denom=recovery_denom, + ) + ) + error_rows.extend( + error_rows_from_q( + dataset=dataset, + panel=panel, + level=level, + method=method, + q_value=q_value, + thresholds=SUMMARY_THRESHOLDS, + label_mask=label_mask, + q_ref=q_ref, + ) + ) + + +def summary_rows_from_curves( + curves: pd.DataFrame, +) -> tuple[list[dict[str, object]], list[dict[str, object]]]: + """Build acceptance and error summary rows from a long-form curves table.""" + acceptance_rows: list[dict[str, object]] = [] + error_rows: list[dict[str, object]] = [] + required = {"dataset", "panel", "method", "spectrum_id", "q_value", "label"} + missing = required - set(curves.columns) + if missing: + raise ValueError(f"Curves table missing columns: {sorted(missing)}") + + for (dataset, panel), group in curves.groupby(["dataset", "panel"], sort=False): + cal = group.loc[group["method"] == DB_CAL_METHOD, ["spectrum_id", "q_value"]] + q_ref_cal_by_id = cal.drop_duplicates("spectrum_id").set_index("spectrum_id")[ + "q_value" + ] + raw = group.loc[group["method"] == DB_RAW_METHOD, ["spectrum_id", "q_value"]] + q_ref_raw_by_id = raw.drop_duplicates("spectrum_id").set_index("spectrum_id")[ + "q_value" + ] + for method, mdf in group.groupby("method", sort=False): + method_s = str(method) + q_value = mdf["q_value"].to_numpy(dtype=float) + label_mask = mdf["label"].astype(bool).to_numpy() + recovery_denom = int(label_mask.sum()) + q_ref: np.ndarray | None = None + # Winnow-family methods: deviation vs calibrated-confidence DBG. + # NovoBoard / Glissade: deviation vs raw-confidence DBG. + if method_s in _WINNOW_METHODS: + q_ref = ( + mdf["spectrum_id"] + .astype(str) + .map(q_ref_cal_by_id) + .to_numpy(dtype=float) + ) + elif method_s in (NOVOBOARD_METHOD, "Glissade"): + q_ref = ( + mdf["spectrum_id"] + .astype(str) + .map(q_ref_raw_by_id) + .to_numpy(dtype=float) + ) + _append_method_summary_rows( + acceptance_rows, + error_rows, + dataset=str(dataset), + panel=str(panel), + level="psm", + method=method_s, + q_value=q_value, + label_mask=label_mask, + recovery_denom=recovery_denom, + q_ref=q_ref, + ) + return acceptance_rows, error_rows + + +def _comparators_for_panel(panel: str, level: str) -> list[str]: + """Comparator method labels present in a given summary panel.""" + del panel, level + return ["NovoBoard"] + + +def _finalise_method_comparison_tables( + acceptance_rows: list[dict[str, object]], + error_rows: list[dict[str, object]], +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Build acceptance and error/gain DataFrames with per-panel relative columns.""" + acceptance = pd.DataFrame(acceptance_rows) + error = pd.DataFrame(error_rows) + if acceptance.empty: + return acceptance, error + + gain_parts: list[pd.DataFrame] = [] + group_cols = ["dataset", "panel", "level"] + for keys, acc_group in acceptance.groupby(group_cols, sort=False): + if not isinstance(keys, tuple): + keys = (keys,) + _, panel, level = keys + err_mask = True + for col, val in zip(group_cols, keys): + err_mask = err_mask & (error[col] == val) + err_group = error.loc[err_mask] + gain_parts.append( + finalise_error_gain_table( + acc_group, + err_group, + primary_method=PRIMARY_METHOD, + comparators=_comparators_for_panel(str(panel), str(level)), + ) + ) + error_gain = pd.concat(gain_parts, ignore_index=True) if gain_parts else error + return acceptance, error_gain + + +def process_dataset(cfg: DatasetConfig, output_dir: Path) -> pd.DataFrame: + """Generate comparison plots and return long-form curve rows for one dataset.""" + residue_masses = _load_residue_masses() + + winnow_unlabelled = load_winnow(cfg.winnow_unlabelled, cfg.fasta, "unlabelled") + winnow_test = load_winnow(cfg.winnow_test, cfg.fasta, "labelled") + nb_u_target, nb_u_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "unlabelled", cfg.novoboard_decoy_rate + ) + nb_t_target, nb_t_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "test", cfg.novoboard_decoy_rate + ) + nb_u_target, nb_u_decoy = filter_novoboard_target_decoy_pairs( + nb_u_target, nb_u_decoy, min_length=MIN_PEPTIDE_LENGTH + ) + nb_t_target, nb_t_decoy = filter_novoboard_target_decoy_pairs( + nb_t_target, nb_t_decoy, min_length=LABELLED_MIN_PEPTIDE_LENGTH + ) + # Pair-gated NovoBoard ⊆ Winnow after identical InstaNovo filters; only trim Winnow. + winnow_unlabelled = _restrict_winnow_to_novoboard_spectra( + winnow_unlabelled, nb_u_target + ) + winnow_test = _restrict_winnow_to_novoboard_spectra(winnow_test, nb_t_target) + _assert_shared_prediction_keys(winnow_test, nb_t_target) + _assert_shared_prediction_keys(winnow_unlabelled, nb_u_target) + + # Shared labels once on Winnow; NovoBoard reuses them by spectrum_id. + winnow_test = winnow_test.copy() + winnow_test["correct"] = novor_correctness_mask( + winnow_test["sequence"], + winnow_test["prediction"], + residue_masses=residue_masses, + ) + if "proteome_hit" not in winnow_unlabelled.columns: + raise KeyError("Expected proteome_hit on unlabelled Winnow table") + correct_by_id = _label_series_by_spectrum_id(winnow_test, "correct") + hit_by_id = _label_series_by_spectrum_id(winnow_unlabelled, "proteome_hit") + + novoboard_unlabelled = novoboard_psm_tdc( + nb_u_target, nb_u_decoy, min_length=MIN_PEPTIDE_LENGTH + ) + n_u_tgt = int(novoboard_unlabelled["is_target"].sum()) + n_u_dec = int((~novoboard_unlabelled["is_target"]).sum()) + if n_u_tgt != n_u_dec: + raise AssertionError( + f"Unlabelled PSM TDC unbalanced: targets={n_u_tgt} decoys={n_u_dec}" + ) + novoboard_unlabelled = novoboard_unlabelled[ + novoboard_unlabelled["is_target"] + ].copy() + novoboard_test = novoboard_psm_tdc( + nb_t_target, nb_t_decoy, min_length=LABELLED_MIN_PEPTIDE_LENGTH + ) + n_t_tgt = int(novoboard_test["is_target"].sum()) + n_t_dec = int((~novoboard_test["is_target"]).sum()) + if n_t_tgt != n_t_dec: + raise AssertionError( + f"Labelled PSM TDC unbalanced: targets={n_t_tgt} decoys={n_t_dec}" + ) + novoboard_test = novoboard_test[novoboard_test["is_target"]].copy() + novoboard_test = _attach_labels_by_spectrum_id( + novoboard_test, correct_by_id, label_col="correct" + ) + novoboard_unlabelled = _attach_labels_by_spectrum_id( + novoboard_unlabelled, hit_by_id, label_col="proteome_hit" + ) + + n_w_correct = int(winnow_test["correct"].sum()) + n_nb_correct = int(novoboard_test["correct"].sum()) + if n_w_correct != n_nb_correct: + raise AssertionError( + f"Shared labelled correct counts disagree: Winnow={n_w_correct} " + f"NovoBoard={n_nb_correct}" + ) + n_w_hit = int(winnow_unlabelled["proteome_hit"].sum()) + n_nb_hit = int(novoboard_unlabelled["proteome_hit"].sum()) + if n_w_hit != n_nb_hit: + raise AssertionError( + f"Shared proteome-hit counts disagree: Winnow={n_w_hit} NovoBoard={n_nb_hit}" + ) + + logger.info( + "%s shared PSM pools: unlabelled=%d labelled=%d " + "(NovoBoard twin-valid; Winnow trimmed; shared labels correct=%d hits=%d)", + cfg.key, + len(winnow_unlabelled), + len(winnow_test), + n_w_correct, + n_w_hit, + ) + + winnow_u_psm_table = _prepare_winnow_psm_table( + winnow_unlabelled, "proteome_hit", residue_masses + ) + winnow_t_psm_table = _prepare_winnow_psm_table( + winnow_test, "correct", residue_masses + ) + curves = pd.concat( + [ + _curves_df_from_winnow_table( + winnow_t_psm_table, + dataset=cfg.key, + panel="labelled_test", + label_col="correct", + ), + _curves_df_from_novoboard( + novoboard_test, + dataset=cfg.key, + panel="labelled_test", + label_col="correct", + ), + _curves_df_from_winnow_table( + winnow_u_psm_table, + dataset=cfg.key, + panel="unlabelled", + label_col="proteome_hit", + ), + _curves_df_from_novoboard( + novoboard_unlabelled, + dataset=cfg.key, + panel="unlabelled", + label_col="proteome_hit", + ), + ], + ignore_index=True, + ) + plot_dataset_from_curves(curves, cfg.key, output_dir) + return curves + + +@app.command() +def main( + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for PNG/PDF outputs."), + ] = DEFAULT_OUTPUT_DIR, + datasets: Annotated[ + Optional[list[str]], + typer.Option("--datasets", help="Dataset keys to plot."), + ] = None, + novoboard_root: Annotated[ + Path, + typer.Option("--novoboard-root", help="NovoBoard datasets root."), + ] = DEFAULT_NOVOBOARD_ROOT, + winnow_results: Annotated[ + Path, + typer.Option("--winnow-results", help="Winnow results directory."), + ] = DEFAULT_WINNOW_RESULTS, + summarise_only: Annotated[ + Optional[Path], + typer.Option( + "--summarise-only", + help="Only write plots and summary CSVs from an existing curves CSV.", + ), + ] = None, +) -> None: + """Generate FDR method comparison plots and summary CSVs.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + output_dir.mkdir(parents=True, exist_ok=True) + + if summarise_only is not None: + curves = pd.read_csv(summarise_only, float_precision="round_trip") + if "label" in curves.columns: + curves["label"] = curves["label"].astype(bool) + dataset_keys = ( + datasets + if datasets is not None + else sorted(curves["dataset"].astype(str).unique()) + ) + for key in dataset_keys: + logger.info("Replotting %s from curves CSV", key) + plot_dataset_from_curves(curves, str(key), output_dir) + curves_out = curves.loc[curves["dataset"].astype(str).isin(dataset_keys)] + acceptance_rows, error_rows = summary_rows_from_curves(curves_out) + acceptance, error_gain = _finalise_method_comparison_tables( + acceptance_rows, error_rows + ) + write_summary_tables( + acceptance, error_gain, output_dir, "fdr_method_comparison" + ) + if summarise_only.resolve() != (output_dir / CURVES_CSV_NAME).resolve(): + write_curves_csv(curves_out, output_dir) + return + + dataset_keys = datasets if datasets is not None else list(DEFAULT_DATASETS) + configs = build_dataset_configs(winnow_results, novoboard_root) + + curve_parts: list[pd.DataFrame] = [] + for key in dataset_keys: + if key not in configs: + raise typer.BadParameter(f"Unknown dataset {key!r}") + logger.info("Processing %s", key) + curve_parts.append(process_dataset(configs[key], output_dir)) + + curves = pd.concat(curve_parts, ignore_index=True) + write_curves_csv(curves, output_dir) + acceptance_rows, error_rows = summary_rows_from_curves(curves) + acceptance, error_gain = _finalise_method_comparison_tables( + acceptance_rows, error_rows + ) + write_summary_tables(acceptance, error_gain, output_dir, "fdr_method_comparison") + + +if __name__ == "__main__": + app() diff --git a/scripts/plot_feature_investigation.py b/scripts/plot_feature_investigation.py new file mode 100644 index 00000000..2fd14664 --- /dev/null +++ b/scripts/plot_feature_investigation.py @@ -0,0 +1,1323 @@ +"""Generate feature investigation plots from calibrator training feature matrices. + +Produces KDE, scatter, violin, correlation, discriminative-power, pairplot, +mirror-spectrum, retention-time, token-stem, and beam-stem figures matching the +style of ``analysis/feature_investigation_new.ipynb``. + +Usage: + python scripts/plot_feature_investigation.py \ + --features-train models/instanovo_helaqc/features_train.parquet \ + [--features-val models/instanovo_helaqc/features_val.parquet] \ + [--metadata-train models/instanovo_helaqc/metadata_train.parquet] \ + [--metadata-val models/instanovo_helaqc/metadata_val.parquet] \ + [--predictions-csv held_out_projects/.../predictions.csv] \ + [--output-dir models/instanovo_helaqc/feature_investigation_plots] +""" + +from __future__ import annotations + +import argparse +import ast +import warnings +from pathlib import Path +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +from matplotlib.colors import LinearSegmentedColormap +from matplotlib.patches import Patch +from scipy import stats as sp_stats +from scipy.stats import gaussian_kde +from sklearn.decomposition import PCA +from sklearn.metrics import roc_auc_score +from sklearn.preprocessing import StandardScaler + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_NEUTRAL_COLOUR = _PALETTE[6] + +_HIGH_CONF_BEAM_COLOUR = _PALETTE[2] # green +_LOW_CONF_BEAM_COLOUR = _PALETTE[5] # purple +_MED_CONF_BEAM_COLOUR = _PALETTE[3] # yellow + +_OBS_COLOUR = _PALETTE[3] # yellow (observed spectrum) +_THEO_COLOUR = _PALETTE[5] # purple (predicted spectrum) + +HUE_LABEL_CORRECT = "Correct" +HUE_LABEL_INCORRECT = "Incorrect" +HUE_ORDER = [HUE_LABEL_CORRECT, HUE_LABEL_INCORRECT] +HUE_PALETTE = { + HUE_LABEL_CORRECT: _CORRECT_COLOUR, + HUE_LABEL_INCORRECT: _INCORRECT_COLOUR, +} + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_SUNSET_COLORS = [ + "#364B9A", + "#4A7BB7", + "#6EA6CD", + "#98CAE1", + "#C2E4EF", + "#EAECCC", + "#FEDA8B", + "#FDB366", + "#F67E4B", + "#DD3D2D", + "#A50026", +] + + +def _diverging_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) + cmap.set_bad(color="#FFFFFF") + return cmap + + +# --------------------------------------------------------------------------- +# Feature definitions +# --------------------------------------------------------------------------- +FEATURE_COLUMNS = [ + "confidence", + "mass_error_ppm", + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", + "irt_error", + "margin", + "median_margin", + "entropy", + "z-score", + "edit_distance", + "min_token_probability", + "std_token_probability", +] + +FRAGMENT_FEATURES = [ + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", +] + +BEAM_FEATURES = ["margin", "median_margin", "entropy", "z-score", "edit_distance"] + +TOKEN_FEATURES = ["min_token_probability", "std_token_probability"] + +SKEWED_FEATURES = {"irt_error"} + +_NICE_LABELS: dict[str, str] = { + "ion_matches": "Ion match rate", + "ion_match_intensity": "Ion match intensity", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Max ion gap", + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation (XCorr)", + "mass_error_ppm": "Precursor mass error", + "log_abs_mass_error_ppm": "Log-absolute mass error", + "log_abs_mass_error_da": "Log-absolute mass error", + "mass_error_da": "Precursor mass error", + "irt_error": "iRT prediction error", + "confidence": "Model confidence", + "margin": "Beam margin", + "median_margin": "Beam median margin", + "entropy": "Beam entropy", + "z-score": "Beam z-score", + "edit_distance": "Runner-up edit distance", + "min_token_probability": "Min. token probability", + "std_token_probability": "Std. token probability", + "predicted_irt": "Regressor-predicted iRT", + "irt": "Koina-predicted iRT", + "retention_time": "Retention time (s)", +} + + +def _nice_label(col: str) -> str: + return _NICE_LABELS.get(col, col.replace("_", " ").capitalize()) + + +# --------------------------------------------------------------------------- +# Plotting helpers +# --------------------------------------------------------------------------- +def _save_fig(fig: plt.Figure, name: str, output_dir: Path) -> None: + base = output_dir / name + fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _auto_ylim(df: pd.DataFrame, feature: str): + if feature in SKEWED_FEATURES: + q99 = df[feature].quantile(0.99) + q01 = df[feature].quantile(0.01) + margin = (q99 - q01) * 0.1 + return (q01 - margin, q99 + margin) + return None + + +def plot_feature_vs_confidence( + df: pd.DataFrame, + feature: str, + title: str | None = None, + ylim: tuple[float, float] | None = None, +) -> tuple[plt.Figure, plt.Axes]: + """Scatter plot of a feature against model confidence, coloured by class.""" + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label] + ax.scatter( + subset["confidence"], + subset[feature], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("confidence")) + ax.set_ylabel(_nice_label(feature)) + if ylim is not None: + ax.set_ylim(ylim) + if title: + ax.set_title(title) + ax.legend() + _style_ax(ax) + fig.tight_layout() + return fig, ax + + +def _plot_peak_normalised_kde( + ax: plt.Axes, + subset: pd.Series, + colour: str, + label: str, + fill: bool, + clip: tuple[float, float] | None, +) -> None: + """Plot a peak-normalised KDE curve on *ax*.""" + kde = gaussian_kde(subset) + lo = subset.min() if clip is None else clip[0] + hi = subset.max() if clip is None else clip[1] + xs = np.linspace(lo, hi, 500) + ys = kde(xs) + ys /= ys.max() + ax.plot(xs, ys, color=colour, label=label, linewidth=1.5) + if fill: + ax.fill_between(xs, ys, alpha=0.3, color=colour) + + +def plot_kde_by_class( + df: pd.DataFrame, + feature: str, + title: str | None = None, + fill: bool = True, + clip: tuple[float, float] | None = None, + peak_normalise: bool = False, +) -> tuple[plt.Figure, plt.Axes]: + """KDE density plot of a feature split by correct/incorrect class.""" + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label][feature].dropna() + if len(subset) < 2: + continue + if peak_normalise: + _plot_peak_normalised_kde(ax, subset, HUE_PALETTE[label], label, fill, clip) + else: + kw: dict = {} + if clip is not None: + kw["clip"] = clip + sns.kdeplot( + subset, + ax=ax, + color=HUE_PALETTE[label], + label=label, + fill=fill, + alpha=0.3, + linewidth=1.5, + **kw, + ) + ax.set_xlabel(_nice_label(feature)) + ax.set_ylabel("Peak-normalised density" if peak_normalise else "Density") + if title: + ax.set_title(title) + ax.legend(loc="upper center") + _style_ax(ax) + fig.tight_layout() + return fig, ax + + +def plot_mirror_spectrum( + obs_mz, + obs_int, + theo_mz, + theo_int, + annotations, + title: str, + ax: plt.Axes | None = None, +) -> tuple[plt.Figure, plt.Axes]: + """Mirror plot comparing observed vs predicted spectra.""" + own_fig = ax is None + if own_fig: + fig, ax = plt.subplots(figsize=(10, 4)) + else: + assert ax is not None + fig = ax.get_figure() + + obs_int_norm = np.array(obs_int) / max(obs_int) * 100 + theo_int_norm = np.array(theo_int) / max(theo_int) * 100 + + ax.vlines(obs_mz, 0, obs_int_norm, color=_OBS_COLOUR, linewidth=1.8) + ax.vlines(theo_mz, 0, -theo_int_norm, color=_THEO_COLOUR, linewidth=1.8) + + if annotations is not None: + for mz_val, intensity_val, ann in zip(theo_mz, theo_int_norm, annotations): + if intensity_val > 10: + label_text = ann.decode() if isinstance(ann, bytes) else str(ann) + ax.annotate( + label_text, + (mz_val, -intensity_val), + fontsize=8, + ha="center", + va="top", + rotation=90, + color=_THEO_COLOUR, + ) + + ax.axhline(0, color="black", linewidth=0.5) + ax.set_xlabel("m/z") + ax.set_ylabel("Relative intensity (%)") + ax.set_title(title) + + ax.text( + 0.99, + 0.95, + "Observed", + transform=ax.transAxes, + ha="right", + va="top", + fontsize=9, + color=_OBS_COLOUR, + fontweight="bold", + ) + ax.text( + 0.99, + 0.05, + "Predicted", + transform=ax.transAxes, + ha="right", + va="bottom", + fontsize=9, + color=_THEO_COLOUR, + fontweight="bold", + ) + + _style_ax(ax) + if own_fig: + fig.tight_layout() + return fig, ax + + +def compute_discriminative_stats(df: pd.DataFrame, features: list[str]) -> pd.DataFrame: + """Compute AUROC, KS statistic, and Cohen's d for each feature.""" + results = [] + labels = df["correct"].astype(int) + for feat in features: + vals = df[feat].dropna() + valid_mask = df[feat].notna() + valid_labels = labels[valid_mask] + valid_vals = vals + if len(valid_vals) < 10 or valid_labels.nunique() < 2: + results.append( + { + "feature": feat, + "auroc": np.nan, + "ks_stat": np.nan, + "cohens_d": np.nan, + } + ) + continue + try: + auroc = roc_auc_score(valid_labels, valid_vals) + auroc = max(auroc, 1 - auroc) + except ValueError: + auroc = np.nan + correct_vals = valid_vals[valid_labels == 1] + incorrect_vals = valid_vals[valid_labels == 0] + ks_stat, _ = sp_stats.ks_2samp(correct_vals, incorrect_vals) + pooled_std = np.sqrt( + ( + (len(correct_vals) - 1) * correct_vals.std() ** 2 + + (len(incorrect_vals) - 1) * incorrect_vals.std() ** 2 + ) + / (len(correct_vals) + len(incorrect_vals) - 2) + ) + cohens_d = ( + abs(correct_vals.mean() - incorrect_vals.mean()) / pooled_std + if pooled_std > 0 + else np.nan + ) + results.append( + {"feature": feat, "auroc": auroc, "ks_stat": ks_stat, "cohens_d": cohens_d} + ) + return ( + pd.DataFrame(results) + .sort_values("auroc", ascending=False) + .reset_index(drop=True) + ) + + +# --------------------------------------------------------------------------- +# Token / beam stem helpers +# --------------------------------------------------------------------------- +_STEM_FIGSIZE = (10, 4.5) + + +def _parse_token_probs(row): + """Extract token probabilities and residue labels from a row.""" + try: + token_probs = np.exp(np.array(ast.literal_eval(row["token_log_probs"]))) + except (ValueError, SyntaxError): + return np.array([]), [] + seq_str = row["prediction"] + residues: list[str] = [] + j = 0 + while j < len(seq_str): + if j + 1 < len(seq_str) and seq_str[j + 1] == "[": + end = seq_str.index("]", j + 1) + 1 + residues.append(seq_str[j:end]) + j = end + else: + residues.append(seq_str[j]) + j += 1 + n_tokens = min(len(token_probs), len(residues)) + return token_probs[:n_tokens], residues[:n_tokens] + + +def _plot_token_stem(row, beam_colour: str, title_suffix: str = ""): + """Stem plot of per-residue token probabilities for one PSM.""" + token_probs, residues = _parse_token_probs(row) + if len(token_probs) == 0: + return None + + n = len(token_probs) + fig, ax = plt.subplots(figsize=_STEM_FIGSIZE) + markerline, stemlines, baseline = ax.stem( + range(n), + token_probs, + linefmt="-", + markerfmt="o", + basefmt="k-", + ) + plt.setp(stemlines, color=beam_colour, linewidth=2.5) + plt.setp(markerline, color=beam_colour, markersize=7, zorder=5) + + ax.set_xticks(range(n)) + ax.set_xticklabels( + residues, + fontsize=11, + rotation=45, + ha="right", + rotation_mode="anchor", + ) + ax.set_xlim(-0.5, n - 0.5) + ax.set_ylim(-0.03, 1.05) + ax.set_ylabel("Token probability") + ax.set_xlabel("Residue") + + charge = int(row["precursor_charge"]) if "precursor_charge" in row.index else "?" + ax.set_title( + f"Token probabilities for {row['prediction']}, " + f"+{charge}, confidence={row['confidence']:.3f}{title_suffix}", + ) + _style_ax(ax) + fig.tight_layout() + return fig + + +def _infer_charge(row) -> str: + """Best-effort charge extraction from a beam CSV row.""" + for col in ("precursor_charge", "charge"): + if col in row.index and pd.notna(row[col]): + return str(int(row[col])) + return "?" + + +def _plot_beam_stem(row, beam_log_prob_cols, beam_seq_cols, colour: str): + """Stem plot of per-beam confidence for one spectrum.""" + probs: list[float] = [] + labels: list[str] = [] + for i, (lp_col, seq_col) in enumerate(zip(beam_log_prob_cols, beam_seq_cols)): + lp = row[lp_col] + seq = row[seq_col] + if pd.isna(lp) or np.isinf(lp): + continue + probs.append(np.exp(float(lp))) + label = str(seq) if pd.notna(seq) else f"beam {i}" + labels.append(label) + + if len(probs) < 2: + return None + + n = len(probs) + fig, ax = plt.subplots(figsize=_STEM_FIGSIZE) + markerline, stemlines, baseline = ax.stem( + range(n), + probs, + linefmt="-", + markerfmt="o", + basefmt="k-", + ) + plt.setp(stemlines, color=colour, linewidth=2.5) + plt.setp(markerline, color=colour, markersize=7, zorder=5) + + ax.set_xticks(range(n)) + ax.set_xlim(-0.5, n - 0.5) + ax.set_ylim(-max(probs) * 0.03, max(probs) * 1.15) + ax.set_ylabel("Beam confidence") + ax.set_xlabel("Beam prediction index") + ax.set_title(f"Beam confidence for {labels[0]}, +{_infer_charge(row)}") + _style_ax(ax) + fig.tight_layout() + return fig + + +# --------------------------------------------------------------------------- +# Section generators — each mirrors a notebook section +# --------------------------------------------------------------------------- +def plot_confidence(df: pd.DataFrame, output_dir: Path) -> None: + """Section 1: confidence distribution.""" + fig, ax = plot_kde_by_class(df, "confidence", title="Model confidence distribution") + _save_fig(fig, "01a_confidence_kde", output_dir) + + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label] + ax.hist( + subset["confidence"], + bins=50, + alpha=0.5, + color=HUE_PALETTE[label], + edgecolor="black", + label=label, + density=True, + ) + ax.set_xlabel(_nice_label("confidence")) + ax.set_ylabel("Density") + ax.set_title("Model confidence histogram") + ax.legend(loc="upper center") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "01b_confidence_histogram", output_dir) + + +def _mass_error_log_column(df: pd.DataFrame) -> tuple[str, str]: + """Return (raw mass error column, log-absolute column) for plotting.""" + if "mass_error_da" in df.columns: + return "mass_error_da", "log_abs_mass_error_da" + if "mass_error_ppm" in df.columns: + return "mass_error_ppm", "log_abs_mass_error_ppm" + raise KeyError( + "Feature matrix must contain 'mass_error_da' or 'mass_error_ppm' for mass error plots" + ) + + +def plot_mass_error(df: pd.DataFrame, output_dir: Path) -> None: + """Section 2: mass error vs confidence (Da or ppm).""" + raw_col, log_col = _mass_error_log_column(df) + work = df.copy() + work[log_col] = np.log(work[raw_col].abs().clip(lower=1e-12)) + + fig, _ = plot_kde_by_class( + work, + raw_col, + title=f"{_nice_label(raw_col)} distribution", + ) + _save_fig(fig, f"02a_{raw_col}_kde", output_dir) + + fig, _ = plot_feature_vs_confidence( + work, + raw_col, + title=f"{_nice_label(raw_col)} vs model confidence", + ) + _save_fig(fig, f"02b_{raw_col}_vs_confidence", output_dir) + + fig, _ = plot_kde_by_class( + work, + log_col, + title="Log-absolute precursor mass error distribution", + ) + _save_fig(fig, "02c_mass_error_log_kde", output_dir) + + fig, _ = plot_feature_vs_confidence( + work, + log_col, + title="Log-absolute precursor mass error vs model confidence", + ) + _save_fig(fig, "02d_mass_error_log_vs_confidence", output_dir) + + if raw_col == "mass_error_da": + da_ylim = (-0.2, 0.2) + fig, _ = plot_kde_by_class( + work, + raw_col, + title=f"{_nice_label(raw_col)} distribution", + clip=da_ylim, + ) + _save_fig(fig, "02e_mass_error_da_kde_within_0.2da", output_dir) + + fig, _ = plot_feature_vs_confidence( + work, + raw_col, + title=f"{_nice_label(raw_col)} vs model confidence", + ylim=da_ylim, + ) + _save_fig(fig, "02f_mass_error_da_vs_confidence_within_0.2da", output_dir) + + +def plot_mirror_spectra(df_meta: pd.DataFrame, output_dir: Path) -> None: + """Section 3: mirror plots of observed vs predicted spectra.""" + required = { + "theoretical_mz", + "mz_array", + "intensity_array", + "theoretical_intensity", + } + if not required.issubset(df_meta.columns): + print(" Skipping mirror plots — missing spectrum columns in metadata.") + return + + valid_mirror = df_meta[ + df_meta["theoretical_mz"].apply(lambda x: x is not None and len(x) > 0) + & df_meta["mz_array"].apply(lambda x: x is not None and len(x) > 0) + ].copy() + + if len(valid_mirror) == 0: + print(" Skipping mirror plots — no rows with valid spectrum arrays.") + return + + has_annotations = "theoretical_annotation" in valid_mirror.columns + + def _add_mirror_margin(ax, y_frac=0.11): + ymin, ymax = ax.get_ylim() + y_pad = (ymax - ymin) * y_frac + ax.set_ylim(ymin - y_pad, ymax + y_pad) + + def _mirror_title(row) -> str: + pred = row.get("prediction", "?") + charge = ( + int(row["precursor_charge"]) if "precursor_charge" in row.index else "?" + ) + return f"Observed vs predicted spectrum for {pred}, +{charge}" + + correct_high = valid_mirror[valid_mirror["correct"]].nlargest(3, "confidence") + incorrect_low = valid_mirror[~valid_mirror["correct"]].nsmallest(3, "confidence") + + conf_middle_lo, conf_middle_hi = 0.45, 0.55 + middle_mask = valid_mirror["confidence"].between(conf_middle_lo, conf_middle_hi) + n_correct_mid = (valid_mirror["correct"] & middle_mask).sum() + n_incorrect_mid = (~valid_mirror["correct"] & middle_mask).sum() + correct_middle = valid_mirror[valid_mirror["correct"] & middle_mask].sample( + n=min(3, n_correct_mid), random_state=42 + ) + incorrect_middle = valid_mirror[~valid_mirror["correct"] & middle_mask].sample( + n=min(3, n_incorrect_mid), random_state=42 + ) + + groups = [ + (correct_high, "correct", "03a_mirror_high_conf"), + (incorrect_low, "incorrect", "03b_mirror_low_conf"), + (correct_middle, "correct", "03c_mirror_middle_conf_correct"), + (incorrect_middle, "incorrect", "03d_mirror_middle_conf_incorrect"), + ] + + for subset, _status, prefix in groups: + for i, (_, row) in enumerate(subset.iterrows()): + fig, ax = plt.subplots(figsize=(8, 5)) + annotations = row.get("theoretical_annotation") if has_annotations else None + plot_mirror_spectrum( + obs_mz=row["mz_array"], + obs_int=row["intensity_array"], + theo_mz=row["theoretical_mz"], + theo_int=row["theoretical_intensity"], + annotations=annotations, + title=_mirror_title(row), + ax=ax, + ) + _add_mirror_margin(ax) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, f"{prefix}_{i}", output_dir) + + +def plot_fragment_features(df: pd.DataFrame, output_dir: Path) -> None: + """Section 4: fragment ion match features vs confidence.""" + available = [f for f in FRAGMENT_FEATURES if f in df.columns] + for feat in available: + fig, _ = plot_feature_vs_confidence( + df, + feat, + title=f"{_nice_label(feat)} vs model confidence", + ylim=_auto_ylim(df, feat), + ) + _save_fig(fig, f"04a_fragment_{feat}_vs_confidence", output_dir) + + fig, _ = plot_kde_by_class(df, feat, title=f"{_nice_label(feat)} distribution") + _save_fig(fig, f"04b_fragment_{feat}_kde", output_dir) + + +def plot_irt(df: pd.DataFrame, df_meta: pd.DataFrame | None, output_dir: Path) -> None: + """Section 7: iRT error plots + RT scatter when metadata is available.""" + if df_meta is not None: + has_rt = ( + "retention_time" in df_meta.columns and "predicted_irt" in df_meta.columns + ) + has_koina_irt = "irt" in df_meta.columns + + if has_rt and has_koina_irt: + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df_meta[df_meta["hue"] == label] + ax.scatter( + subset["retention_time"], + subset["irt"], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("retention_time")) + ax.set_ylabel(_nice_label("irt")) + ax.set_title("Retention time vs Koina-predicted iRT") + ax.legend(markerscale=3, frameon=True) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "07a_rt_vs_koina_irt", output_dir) + + if has_koina_irt and has_rt: + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df_meta[df_meta["hue"] == label] + ax.scatter( + subset["predicted_irt"], + subset["irt"], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("predicted_irt")) + ax.set_ylabel(_nice_label("irt")) + ax.set_title("Koina-predicted iRT vs regressor-predicted iRT") + ax.legend(markerscale=3, frameon=True) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "07b_predicted_vs_koina_irt", output_dir) + + if "irt_error" not in df.columns: + return + + irt_ylim = _auto_ylim(df, "irt_error") + fig, _ = plot_feature_vs_confidence( + df, "irt_error", title="iRT prediction error vs model confidence", ylim=irt_ylim + ) + _save_fig(fig, "07c_irt_error_vs_confidence", output_dir) + + fig, _ = plot_kde_by_class( + df, + "irt_error", + title="iRT prediction error distribution", + clip=(0, df["irt_error"].quantile(0.99)), + ) + _save_fig(fig, "07d_irt_error_kde", output_dir) + + +def plot_token_stems(df_meta: pd.DataFrame, output_dir: Path) -> None: + """Section 8: token-level probability stem plots from metadata.""" + if "token_log_probs" not in df_meta.columns: + print(" Skipping token stem plots — no token_log_probs column in metadata.") + return + if "prediction" not in df_meta.columns: + print(" Skipping token stem plots — no prediction column in metadata.") + return + + # High confidence + high_conf_pool = df_meta[df_meta["confidence"] >= 0.9] + high_samples = ( + high_conf_pool.sample(3, random_state=42) + if len(high_conf_pool) >= 3 + else high_conf_pool + ) + for i, (_, row) in enumerate(high_samples.iterrows()): + fig = _plot_token_stem(row, _HIGH_CONF_BEAM_COLOUR) + if fig is not None: + _save_fig(fig, f"08a_token_stem_high_conf_{i}", output_dir) + + # Medium confidence + med_conf_pool = df_meta[ + (df_meta["confidence"] >= 0.4) & (df_meta["confidence"] <= 0.7) + ] + med_samples = ( + med_conf_pool.sample(3, random_state=42) + if len(med_conf_pool) >= 3 + else med_conf_pool + ) + for i, (_, row) in enumerate(med_samples.iterrows()): + fig = _plot_token_stem(row, _MED_CONF_BEAM_COLOUR) + if fig is not None: + _save_fig(fig, f"08b_token_stem_med_conf_{i}", output_dir) + + # Low confidence + low_conf_pool = df_meta[df_meta["confidence"] <= 0.2] + low_samples = ( + low_conf_pool.sample(3, random_state=42) + if len(low_conf_pool) >= 3 + else low_conf_pool + ) + for i, (_, row) in enumerate(low_samples.iterrows()): + fig = _plot_token_stem(row, _LOW_CONF_BEAM_COLOUR) + if fig is not None: + _save_fig(fig, f"08c_token_stem_low_conf_{i}", output_dir) + + +def plot_beam_stems(predictions_csv: Path, output_dir: Path) -> None: + """Section 8b: beam confidence stem plots from the predictions CSV.""" + beam_csv = pd.read_csv(predictions_csv) + + beam_log_prob_cols = sorted( + [ + c + for c in beam_csv.columns + if c.startswith("predictions_log_probability_beam_") + ], + key=lambda c: int(c.rsplit("_", 1)[1]), + ) + beam_seq_cols = sorted( + [ + c + for c in beam_csv.columns + if c.startswith("predictions_beam_") + and "log_probability" not in c + and "token" not in c + ], + key=lambda c: int(c.rsplit("_", 1)[1]), + ) + + if not beam_log_prob_cols or not beam_seq_cols: + print(" Skipping beam stem plots — no beam columns in predictions CSV.") + return + + beam_csv["top_confidence"] = np.exp(beam_csv[beam_log_prob_cols[0]].astype(float)) + + # Filter rows where all beams are -inf or NaN + valid_beams = beam_csv.dropna(subset=beam_log_prob_cols, how="all").copy() + for col in beam_log_prob_cols: + valid_beams[col] = pd.to_numeric(valid_beams[col], errors="coerce") + valid_beams = valid_beams[ + valid_beams[beam_log_prob_cols].apply( + lambda row: not all(np.isinf(row) | row.isna()), axis=1 + ) + ] + valid_beams = valid_beams[ + valid_beams[beam_log_prob_cols].apply( + lambda row: any(np.exp(row.dropna()) > 1e-15), axis=1 + ) + ] + + if len(valid_beams) == 0: + print(" Skipping beam stem plots — no valid beam rows after filtering.") + return + + # High confidence beams + high_beam = valid_beams[valid_beams["top_confidence"] >= 0.9] + high_beam_samples = ( + high_beam.sample(3, random_state=42) if len(high_beam) >= 3 else high_beam + ) + for i, (_, row) in enumerate(high_beam_samples.iterrows()): + fig = _plot_beam_stem( + row, beam_log_prob_cols, beam_seq_cols, _HIGH_CONF_BEAM_COLOUR + ) + if fig is not None: + _save_fig(fig, f"08d_beam_conf_high_{i}", output_dir) + + # Low confidence beams + low_beam = valid_beams[valid_beams["top_confidence"] <= 0.2] + low_beam_samples = ( + low_beam.sample(3, random_state=42) if len(low_beam) >= 3 else low_beam + ) + for i, (_, row) in enumerate(low_beam_samples.iterrows()): + fig = _plot_beam_stem( + row, beam_log_prob_cols, beam_seq_cols, _LOW_CONF_BEAM_COLOUR + ) + if fig is not None: + _save_fig(fig, f"08e_beam_conf_low_{i}", output_dir) + + +def plot_beam_features(df: pd.DataFrame, output_dir: Path) -> None: + """Section 9: beam search features vs confidence.""" + available = [f for f in BEAM_FEATURES if f in df.columns] + for feat in available: + fig, _ = plot_feature_vs_confidence( + df, feat, title=f"{_nice_label(feat)} vs model confidence" + ) + _save_fig(fig, f"09a_beam_{feat}_scatter", output_dir) + + fig, _ = plot_kde_by_class(df, feat, title=f"{_nice_label(feat)} distribution") + _save_fig(fig, f"09b_beam_{feat}_kde", output_dir) + + +def plot_token_features(df: pd.DataFrame, output_dir: Path) -> None: + """Section 11: token-level features.""" + if "min_token_probability" not in df.columns: + return + + fig, _ = plot_kde_by_class( + df, "min_token_probability", title="Min. token probability distribution" + ) + _save_fig(fig, "11a_min_token_prob_kde", output_dir) + + fig, _ = plot_kde_by_class( + df, "std_token_probability", title="Std. token probability distribution" + ) + _save_fig(fig, "11b_std_token_prob_kde", output_dir) + + fig, _ = plot_feature_vs_confidence( + df, "min_token_probability", title="Min. token probability vs confidence" + ) + _save_fig(fig, "11c_min_token_prob_scatter", output_dir) + + fig, _ = plot_feature_vs_confidence( + df, "std_token_probability", title="Std. token probability vs confidence" + ) + _save_fig(fig, "11d_std_token_prob_scatter", output_dir) + + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label] + ax.scatter( + subset["min_token_probability"], + subset["std_token_probability"], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("min_token_probability")) + ax.set_ylabel(_nice_label("std_token_probability")) + ax.set_title("Token-level feature space") + ax.legend(markerscale=3, frameon=True) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "11e_token_feature_2d", output_dir) + + +def _plot_pca(df: pd.DataFrame, available: list[str], output_dir: Path) -> None: + """12f/12g — PCA scatter and loadings for calibrator features.""" + feat_df = df[available].dropna() + if len(feat_df) < 10: + return + + hue_pca = df.loc[feat_df.index, "hue"].values + scaler = StandardScaler() + x_scaled = scaler.fit_transform(feat_df.values) + pca = PCA(n_components=2) + z_pca = pca.fit_transform(x_scaled) + + fig, ax = plt.subplots(figsize=(8, 7)) + for label, colour in zip(reversed(HUE_ORDER), [_INCORRECT_COLOUR, _CORRECT_COLOUR]): + mask = hue_pca == label + ax.scatter( + z_pca[mask, 0], + z_pca[mask, 1], + c=colour, + label=label, + s=10, + alpha=0.3, + rasterized=True, + ) + ax.set_xlabel(f"PC 1 ({pca.explained_variance_ratio_[0]:.1%} variance)") + ax.set_ylabel(f"PC 2 ({pca.explained_variance_ratio_[1]:.1%} variance)") + ax.set_title("PCA of calibrator features") + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "12f_pca_features", output_dir) + + pc1 = pca.components_[0] + pc2 = pca.components_[1] + names = [_nice_label(c) for c in available] + order = np.argsort(np.abs(pc1))[::-1] + + y = np.arange(len(names)) + fig, ax = plt.subplots(figsize=(10, 7)) + ax.barh( + y, + pc1[order], + color=_CORRECT_COLOUR, + alpha=0.6, + edgecolor="black", + linewidth=0.4, + ) + ax.barh( + y, pc2[order], color=_PALETTE[5], alpha=0.4, edgecolor="black", linewidth=0.4 + ) + ax.set_yticks(y) + ax.set_yticklabels([names[i] for i in order]) + ax.invert_yaxis() + ax.set_xlabel("Loading value") + ax.set_title("PCA loadings for first two principal components") + ax.axvline(0, color="black", linewidth=0.5) + ax.legend( + handles=[ + Patch(facecolor=_CORRECT_COLOUR, alpha=0.6, label="PC 1 loading"), + Patch(facecolor=_PALETTE[5], alpha=0.4, label="PC 2 loading"), + ], + loc="lower right", + ) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "12g_pca_loadings", output_dir) + + +def plot_discriminative_analysis(df: pd.DataFrame, output_dir: Path) -> None: + """Section 12: discriminative stats, correlation, violins, pairplot.""" + available = [f for f in FEATURE_COLUMNS if f in df.columns] + disc_stats = compute_discriminative_stats(df, available) + + # 12a — AUROC bar chart + fig, ax = plt.subplots(figsize=(8, 7)) + colours = [ + _PALETTE[0] if v >= 0.7 else _NEUTRAL_COLOUR for v in disc_stats["auroc"] + ] + ax.barh(range(len(disc_stats)), disc_stats["auroc"], color=colours) + ax.set_yticks(range(len(disc_stats))) + ax.set_yticklabels([_nice_label(f) for f in disc_stats["feature"]], fontsize=9) + ax.set_xlabel("AUROC") + ax.set_title("Per-feature AUROC for separating correct vs incorrect") + ax.axvline(0.5, color="grey", linestyle="--", linewidth=0.8) + ax.invert_yaxis() + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "12a_feature_auroc_ranking", output_dir) + + # 12b — correlation matrix + corr = df[available].corr() + fig, ax = plt.subplots(figsize=(14, 12)) + mask = np.triu(np.ones_like(corr, dtype=bool), k=1) + sns.heatmap( + corr, + mask=mask, + cmap=_diverging_cmap(), + center=0, + ax=ax, + xticklabels=[_nice_label(c) for c in available], + yticklabels=[_nice_label(c) for c in available], + annot=True, + fmt=".2f", + annot_kws={"size": 6}, + linewidths=0.5, + square=True, + vmin=-1, + vmax=1, + cbar_kws={"label": "Pearson r"}, + ) + ax.set_title("Feature correlation matrix") + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right", fontsize=7) + ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=7) + fig.tight_layout() + _save_fig(fig, "12b_correlation_matrix", output_dir) + + # 12c — violin plots per feature group + feature_groups = { + "Fragment match": [f for f in FRAGMENT_FEATURES if f in df.columns], + "Beam search": [f for f in BEAM_FEATURES if f in df.columns], + "Token-level": [f for f in TOKEN_FEATURES if f in df.columns], + } + for _group_name, group_feats in feature_groups.items(): + for feat in group_feats: + fig, ax = plt.subplots(figsize=(6, 5)) + sns.violinplot( + data=df, + x="hue", + y=feat, + hue="hue", + ax=ax, + palette=HUE_PALETTE, + order=HUE_ORDER, + hue_order=HUE_ORDER, + inner="quartile", + cut=0, + linewidth=0.8, + legend=False, + ) + ax.set_xlabel("") + ax.set_ylabel(_nice_label(feat)) + ax.set_title(f"{_nice_label(feat)} by identification status") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, f"12c_violin_{feat}", output_dir) + + # 12e — pairplot of top-5 features + top5 = disc_stats.head(5)["feature"].tolist() + sample_size = min(2000, len(df)) + df_sample = df[top5 + ["hue"]].sample(n=sample_size, random_state=42) + + g = sns.pairplot( + df_sample, + vars=top5, + hue="hue", + palette=HUE_PALETTE, + hue_order=HUE_ORDER, + diag_kind="kde", + plot_kws={"alpha": 0.25, "s": 8, "rasterized": True}, + diag_kws={"fill": True, "alpha": 0.3}, + height=2.2, + ) + g.figure.suptitle("Pairplot of top-5 discriminative features", y=1.01, fontsize=13) + g._legend.set_title("Identification") + for ax_row in g.axes: + for ax_item in ax_row: + xl = ax_item.get_xlabel() + yl = ax_item.get_ylabel() + if xl: + ax_item.set_xlabel(_nice_label(xl), fontsize=7) + if yl: + ax_item.set_ylabel(_nice_label(yl), fontsize=7) + _style_ax(ax_item) + _save_fig(g.figure, "12e_pairplot_top5", output_dir) + + _plot_pca(df, available, output_dir) + + # Save discriminative stats as CSV for reference + disc_stats.to_csv(output_dir / "discriminative_stats.csv", index=False) + + +def print_summary(df: pd.DataFrame) -> None: + """Section 13: summary statistics printed to stdout.""" + available = [f for f in FEATURE_COLUMNS if f in df.columns] + disc_stats = compute_discriminative_stats(df, available) + + print("=" * 80) + print("DISCRIMINATIVE STATISTICS SUMMARY") + print("=" * 80) + n_correct = df["correct"].sum() + n_incorrect = (~df["correct"]).sum() + print( + f"\nDataset: {len(df):,} spectra | {n_correct:,} correct | {n_incorrect:,} incorrect" + ) + print(f"Class balance: {df['correct'].mean():.1%} correct\n") + + print("Per-feature discriminative power (sorted by AUROC):") + print("-" * 80) + print(disc_stats.to_string(index=False, float_format="%.3f")) + + print("\nTop-5 features by AUROC:") + for _, row in disc_stats.head(5).iterrows(): + print( + f" {_nice_label(row['feature']):40s} AUROC={row['auroc']:.3f} " + f"KS={row['ks_stat']:.3f} d={row['cohens_d']:.3f}" + ) + + print("\nBottom-5 features by AUROC:") + for _, row in disc_stats.tail(5).iterrows(): + print( + f" {_nice_label(row['feature']):40s} AUROC={row['auroc']:.3f} " + f"KS={row['ks_stat']:.3f} d={row['cohens_d']:.3f}" + ) + + print("\nConfidence statistics:") + print(f" Overall mean confidence: {df['confidence'].mean():.3f}") + print(f" Correct mean confidence: {df[df['correct']]['confidence'].mean():.3f}") + print(f" Incorrect mean confidence: {df[~df['correct']]['confidence'].mean():.3f}") + print( + f" Confidence AUROC: " + f"{roc_auc_score(df['correct'].astype(int), df['confidence']):.3f}" + ) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse command-line arguments for feature investigation plots.""" + parser = argparse.ArgumentParser( + description="Feature investigation plots from calibrator training matrices.", + ) + parser.add_argument( + "--features-train", + type=Path, + required=True, + help="Path to the training features parquet produced by `winnow compute-features`.", + ) + parser.add_argument( + "--features-val", + type=Path, + default=None, + help="Optional path to validation features parquet. When provided the train " + "and val splits are concatenated for richer plots.", + ) + parser.add_argument( + "--metadata-train", + type=Path, + default=None, + help="Optional path to full training metadata parquet (produced by compute-features " + "with metadata_output_path set). Enables mirror plots, RT scatter, and token stems.", + ) + parser.add_argument( + "--metadata-val", + type=Path, + default=None, + help="Optional path to full validation metadata parquet.", + ) + parser.add_argument( + "--predictions-csv", + type=Path, + default=None, + help="Optional path to InstaNovo-style predictions CSV with beam columns. " + "Enables beam confidence stem plots.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + help="Directory for output plots. Defaults to a `feature_investigation_plots` " + "subdirectory next to --features-train.", + ) + return parser.parse_args(argv) + + +def _load_metadata(args: argparse.Namespace) -> pd.DataFrame | None: + """Load and concatenate metadata parquets when provided.""" + parts: list[pd.DataFrame] = [] + for path in (args.metadata_train, args.metadata_val): + if path is not None: + print(f"Loading metadata from {path}") + parts.append(pl.read_parquet(path).to_pandas()) + if not parts: + return None + df_meta = pd.concat(parts, ignore_index=True) + print(f"Combined metadata: {len(df_meta):,} rows, {len(df_meta.columns)} columns") + df_meta["hue"] = df_meta["correct"].map( + {True: HUE_LABEL_CORRECT, False: HUE_LABEL_INCORRECT} + ) + return df_meta + + +def main(argv: list[str] | None = None) -> None: + """Generate all feature investigation plots.""" + warnings.filterwarnings("ignore", category=FutureWarning) + args = parse_args(argv) + + # -- Load features (always required) -- + print(f"Loading training features from {args.features_train}") + df = pl.read_parquet(args.features_train).to_pandas() + + if args.features_val is not None: + print(f"Loading validation features from {args.features_val}") + df_val = pl.read_parquet(args.features_val).to_pandas() + df = pd.concat([df, df_val], ignore_index=True) + print(f"Combined dataset: {len(df):,} spectra") + + df["hue"] = df["correct"].map({True: HUE_LABEL_CORRECT, False: HUE_LABEL_INCORRECT}) + + # -- Load metadata (optional) -- + df_meta = _load_metadata(args) + + has_metadata = df_meta is not None + has_predictions_csv = args.predictions_csv is not None + + output_dir = args.output_dir + if output_dir is None: + output_dir = args.features_train.parent / "feature_investigation_plots" + output_dir.mkdir(parents=True, exist_ok=True) + print(f"Saving plots to {output_dir}/\n") + + print(f"Dataset shape: {df.shape}") + n_correct = df["correct"].sum() + n_incorrect = (~df["correct"]).sum() + print(f"Correct: {n_correct:,} | Incorrect: {n_incorrect:,} | Total: {len(df):,}") + print(f"Class balance: {df['correct'].mean():.1%} correct\n") + + n_steps = 7 + has_metadata * 2 + has_predictions_csv + step = 0 + + step += 1 + print(f"[{step}/{n_steps}] Confidence distribution...") + plot_confidence(df, output_dir) + + step += 1 + print(f"[{step}/{n_steps}] Mass error vs confidence...") + plot_mass_error(df, output_dir) + + if has_metadata: + step += 1 + print(f"[{step}/{n_steps}] Mirror spectrum plots...") + plot_mirror_spectra(df_meta, output_dir) + + step += 1 + print(f"[{step}/{n_steps}] Fragment ion match features...") + plot_fragment_features(df, output_dir) + + step += 1 + print(f"[{step}/{n_steps}] iRT error...") + plot_irt(df, df_meta, output_dir) + + if has_metadata: + step += 1 + print(f"[{step}/{n_steps}] Token-level stem plots...") + plot_token_stems(df_meta, output_dir) + + if has_predictions_csv: + step += 1 + print(f"[{step}/{n_steps}] Beam confidence stem plots...") + plot_beam_stems(args.predictions_csv, output_dir) + + step += 1 + print(f"[{step}/{n_steps}] Beam search features...") + plot_beam_features(df, output_dir) + + step += 1 + print(f"[{step}/{n_steps}] Token-level features...") + plot_token_features(df, output_dir) + + step += 1 + print( + f"[{step}/{n_steps}] Discriminative analysis (AUROC, correlation, violins, pairplot)..." + ) + plot_discriminative_analysis(df, output_dir) + + print() + print_summary(df) + + print(f"\nDone — plots saved to {output_dir}/") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_external_peptide_holdout_benchmark.py b/scripts/run_external_peptide_holdout_benchmark.py new file mode 100644 index 00000000..d535cac6 --- /dev/null +++ b/scripts/run_external_peptide_holdout_benchmark.py @@ -0,0 +1,1115 @@ +#!/usr/bin/env python3 +"""Glissade-style external peptide score-mixture benchmark. + +Builds a **shared** matched pool S_m (labelled-test Novor-correct peptides) and +external pool S_e (unlabelled proteome-external peptides) after filter → +max-score-per-peptide (NovoBoard mass-deltas converted to ProForma; unsupported +mods dropped; NovoBoard target-decoy pairs gated so every retained key has a +twin). Unlabelled / external peptides require normalised length ≥ 8 (proteome +substring proxy); labelled matched peptides and Glissade's training-split +reference keep short peptides (Novor agreement). Novor and proteome-hit labels +are computed once on Winnow and reused for NovoBoard by ``spectrum_id``. +Method-specific scores are attached to the same peptide keys. Mixtures control +π₀ explicitly. + +Mixtures are drawn **without replacement** so every peptide key is unique. All +three methods therefore score the identical mixture and realise the same π₀; +NovoBoard's max-score-per-peptide step is a no-op, which is asserted per +mixture. + +Each tool draws its null/reference information from the same place, the +annotated *training* split of its own organism: Winnow through the pretrained +per-dataset calibrator, Glissade through the training-split matched score +distribution, NovoBoard through its training-tuned decoy masking rate. No tool +fits on the evaluation labels. + +NovoBoard peptide FDR uses max-target → twin-decoy TDC. Winnow uses max +calibrated confidence then nonparametric FDR (PSM-calibrator proxy). Glissade +uses native bootstrap FDR with NumPy seeded from the benchmark RNG. +""" + +from __future__ import annotations + +import importlib +import logging +import sys +from pathlib import Path +from typing import Annotated, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 + LABELLED_MIN_PEPTIDE_LENGTH, + MIN_PEPTIDE_LENGTH, + assert_shared_prediction_keys, + attach_labels_by_spectrum_id, + confidence_to_log_prob, + filter_novoboard_target_decoy_pairs, + filter_prediction_table, + label_series_by_spectrum_id, + max_score_per_peptide, + novoboard_max_target_twin_decoy_tdc, + novor_correctness_mask, + prepare_novoboard_decoy_by_pair, + restrict_winnow_to_novoboard_spectra, +) +from scripts.fdr_tool_comparison_summaries import ( # noqa: E402 + SUMMARY_THRESHOLDS, + database_grounded_q_from_labels, + mean_abs_q_dev_vs_reference, + summarise_holdout_results, + write_summary_tables, +) +from scripts.plot_eval_results import _PALETTE, _display_name, _save_fig, _style_ax # noqa: E402 +from scripts.plot_fdr_method_comparison import ( # noqa: E402 + DEFAULT_MODEL_ROOT, + DEFAULT_NOVOBOARD_ROOT, + DEFAULT_WINNOW_RESULTS, + build_dataset_configs, + load_novoboard_target_decoy, + load_winnow, +) +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +logger = logging.getLogger(__name__) +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark_v2" +DEFAULT_GLISSADE_REPO = Path("/home/j-daniel/repos/glissade") +DEFAULT_DATASETS = ["helaqc", "celegans"] +DEFAULT_Q_THRESHOLDS = [round(float(x), 2) for x in np.linspace(0.0, 0.25, 26)] +DEFAULT_PI0_GRID = [0.5, 0.6, 0.7, 0.8, 0.9] +DEFAULT_N_ITERATIONS = 20 +# Matches the Glissade package default for run_bootstraps. +DEFAULT_N_BOOTSTRAPS = 10 +# Prefer the full matched pool; |S_c| is then capped so the highest π₀ remains +# drawable without replacement from S_e. +DEFAULT_HOLDOUT_FRAC = 1.0 +DEFAULT_SEED = 42 +METHODS = ("Winnow", "NovoBoard", "Glissade") + + +def max_correct_pool_for_pi0_grid(n_external_pool: int, pi0_grid: list[float]) -> int: + """Largest |S_c| such that every π₀ in ``pi0_grid`` fits without replacement. + + Requires ``round(π₀ / (1 - π₀) · |S_c|) ≤ |S_e|`` for each target π₀. + """ + if n_external_pool < 1: + return 0 + limit = n_external_pool + for pi0 in pi0_grid: + if not 0.0 < pi0 < 1.0: + continue + ratio = pi0 / (1.0 - pi0) + # Largest n with round(ratio * n) <= n_external_pool. + # For ratio = k integer (e.g. 0.9 → 9), this is floor(pool / k). + hi = int(n_external_pool / ratio) + 2 + n_ok = 0 + for n in range(1, hi + 1): + if int(round(ratio * n)) <= n_external_pool: + n_ok = n + else: + break + limit = min(limit, n_ok) + return max(0, limit) + + +def _load_glissade_functions(glissade_repo: Path): + repo = str(glissade_repo.resolve()) + if repo not in sys.path: + sys.path.insert(0, repo) + module = importlib.import_module("glissade.glissade") + return module.run_bootstraps, module.annotate_results, module.compute_fdr_transform + + +def _load_winnow_with_raw_confidence( + predictions_dir: Path, fasta: Path, eval_type: str +) -> pd.DataFrame: + """Load Winnow preds and ensure raw ``confidence`` is present.""" + df = load_winnow(predictions_dir, fasta, eval_type) # type: ignore[arg-type] + if "confidence" not in df.columns: + meta_path = predictions_dir / "metadata.csv" + if not meta_path.is_file(): + raise FileNotFoundError(meta_path) + meta = pd.read_csv(meta_path, usecols=["spectrum_id", "confidence"]) + df = df.merge(meta, on="spectrum_id", how="inner") + return df + + +def build_glissade_training_reference( + train_metadata: Path, + *, + min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, +) -> pd.DataFrame: + """Matched reference score distribution for Glissade, from the training split. + + Glissade anchors its null-fraction estimate on a database-matched score + distribution. Taking that anchor from the annotated training split puts it on + the same data the Winnow calibrator was trained on and the NovoBoard decoy + masking rate was tuned on, and keeps it disjoint from the evaluation spectra. + Short peptides are retained: the reference is labelled (Novor-correct). + + Args: + train_metadata: Calibrator training metadata parquet. + min_length: Minimum normalised peptide length (default: labelled floor). + + Returns: + One row per Novor-correct training peptide with ``raw_confidence`` and + ``score_glissade``. + """ + if not train_metadata.is_file(): + raise FileNotFoundError(train_metadata) + train = pl.read_parquet( + train_metadata, + columns=["spectrum_id", "prediction", "sequence", "confidence"], + ).to_pandas() + train = filter_prediction_table( + train, "prediction", min_length=min_length, key_col="peptide_key" + ) + train["correct"] = novor_correctness_mask(train["sequence"], train["prediction"]) + matched = max_score_per_peptide( + train.loc[train["correct"]], "peptide_key", "confidence" + ) + reference = matched[["peptide_key", "confidence"]].rename( + columns={"confidence": "raw_confidence"} + ) + reference["score_glissade"] = confidence_to_log_prob(reference["raw_confidence"]) + logger.info( + "Glissade training reference: %d matched peptides from %s", + len(reference), + train_metadata, + ) + return reference + + +def _namespace_pair_keys(df: pd.DataFrame, namespace: str) -> pd.DataFrame: + """Prefix ``_pair_key`` to avoid collisions across splits.""" + work = df.copy() + if "_pair_key" not in work.columns: + raise ValueError("Missing '_pair_key'") + work["_pair_key"] = namespace + ":" + work["_pair_key"].astype(str) + return work + + +def build_shared_score_tables( + *, + dataset: str, + winnow_results: Path, + novoboard_root: Path, + model_root: Path = DEFAULT_MODEL_ROOT, + unlabelled_min_length: int = MIN_PEPTIDE_LENGTH, + labelled_min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, +) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """Return shared matched/external keys with per-method scores. + + All methods follow filter → :func:`max_score_per_peptide`. NovoBoard tables + are pair-gated (equal target/decoy twins) before max-dedupe so mixture keys + always have a twin. Labelled S_m and Glissade's training reference keep short + peptides (``labelled_min_length``); unlabelled S_e uses + ``unlabelled_min_length`` for the proteome-substring proxy. Labelled S_m uses + Novor correctness computed once on Winnow and reused for NovoBoard by + ``spectrum_id``; S_e uses proteome-hit the same way. Glissade scores are + ``log`` raw InstaNovo confidence on the shared keys. + + Returns: + matched_scores: one row per peptide_key in S_m with method scores. + external_scores: same for S_e. + nb_decoy_combined: namespaced twin-valid decoys for twin TDC. + glissade_reference: training-split matched scores for Glissade's anchor. + """ + cfg = build_dataset_configs(winnow_results, novoboard_root, model_root)[dataset] + + winnow_test = _load_winnow_with_raw_confidence( + cfg.winnow_test, cfg.fasta, "labelled" + ) + winnow_unlab = _load_winnow_with_raw_confidence( + cfg.winnow_unlabelled, cfg.fasta, "unlabelled" + ) + # load_winnow already applied the labelled/unlabelled length floors; re-apply + # explicitly so callers can override without depending on that path. + winnow_test = filter_prediction_table( + winnow_test, + "prediction", + min_length=labelled_min_length, + key_col="peptide_key", + ) + w_unlab = filter_prediction_table( + winnow_unlab, + "prediction", + min_length=unlabelled_min_length, + key_col="peptide_key", + ) + if "proteome_hit" not in w_unlab.columns: + raise KeyError("Expected proteome_hit on unlabelled Winnow table") + + nb_test_target, nb_test_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "test", cfg.novoboard_decoy_rate + ) + nb_unlab_target, nb_unlab_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "unlabelled", cfg.novoboard_decoy_rate + ) + nb_test_target = _namespace_pair_keys(nb_test_target, "test") + nb_test_decoy = _namespace_pair_keys(nb_test_decoy, "test") + nb_unlab_target = _namespace_pair_keys(nb_unlab_target, "unlabelled") + nb_unlab_decoy = _namespace_pair_keys(nb_unlab_decoy, "unlabelled") + + nb_test_target, nb_test_decoy = filter_novoboard_target_decoy_pairs( + nb_test_target, + nb_test_decoy, + min_length=labelled_min_length, + key_col="peptide_key", + ) + nb_unlab_target, nb_unlab_decoy = filter_novoboard_target_decoy_pairs( + nb_unlab_target, + nb_unlab_decoy, + min_length=unlabelled_min_length, + key_col="peptide_key", + ) + nb_decoy_combined = pd.concat( + [nb_test_decoy, nb_unlab_decoy], ignore_index=True, sort=False + ) + + # Twin-valid NovoBoard ⊆ Winnow; shared labels once on Winnow. + winnow_test = restrict_winnow_to_novoboard_spectra(winnow_test, nb_test_target) + w_unlab = restrict_winnow_to_novoboard_spectra(w_unlab, nb_unlab_target) + assert_shared_prediction_keys(winnow_test, nb_test_target) + assert_shared_prediction_keys(w_unlab, nb_unlab_target) + + winnow_test = winnow_test.copy() + winnow_test["correct"] = novor_correctness_mask( + winnow_test["sequence"], winnow_test["prediction"] + ) + correct_by_id = label_series_by_spectrum_id(winnow_test, "correct") + hit_by_id = label_series_by_spectrum_id(w_unlab, "proteome_hit") + nb_test_target = attach_labels_by_spectrum_id( + nb_test_target, correct_by_id, label_col="correct" + ) + nb_unlab_f = attach_labels_by_spectrum_id( + nb_unlab_target, hit_by_id, label_col="proteome_hit" + ) + + w_matched = max_score_per_peptide( + winnow_test.loc[winnow_test["correct"]], + "peptide_key", + "calibrated_confidence", + ) + w_matched_raw = max_score_per_peptide( + winnow_test.loc[winnow_test["correct"]], + "peptide_key", + "confidence", + ) + nb_correct = nb_test_target.loc[nb_test_target["correct"]].copy() + nb_matched = max_score_per_peptide(nb_correct, "peptide_key", "ALC (%)") + + w_external = max_score_per_peptide( + w_unlab.loc[~w_unlab["proteome_hit"].astype(bool)], + "peptide_key", + "calibrated_confidence", + ) + w_external_raw = max_score_per_peptide( + w_unlab.loc[~w_unlab["proteome_hit"].astype(bool)], + "peptide_key", + "confidence", + ) + nb_external = max_score_per_peptide( + nb_unlab_f.loc[~nb_unlab_f["proteome_hit"].astype(bool)], + "peptide_key", + "ALC (%)", + ) + + # Shared membership = Winnow ∩ twin-valid NovoBoard keys. + sm_keys = ( + set(w_matched["peptide_key"]) + & set(nb_matched["peptide_key"]) + & set(w_matched_raw["peptide_key"]) + ) + se_keys = ( + set(w_external["peptide_key"]) + & set(nb_external["peptide_key"]) + & set(w_external_raw["peptide_key"]) + ) + if sm_keys != set(w_matched["peptide_key"]) or sm_keys != set( + nb_matched["peptide_key"] + ): + raise AssertionError( + f"{dataset} shared S_m keys disagree after shared Novor labels: " + f"winnow={len(w_matched)} novoboard={len(nb_matched)} " + f"intersection={len(sm_keys)}" + ) + if se_keys != set(w_external["peptide_key"]) or se_keys != set( + nb_external["peptide_key"] + ): + raise AssertionError( + f"{dataset} shared S_e keys disagree after shared proteome-hit labels: " + f"winnow={len(w_external)} novoboard={len(nb_external)} " + f"intersection={len(se_keys)}" + ) + if not sm_keys: + raise ValueError(f"Empty shared matched pool for {dataset}") + if not se_keys: + raise ValueError(f"Empty shared external pool for {dataset}") + + matched = pd.DataFrame({"peptide_key": sorted(sm_keys)}) + matched = matched.merge( + w_matched[["peptide_key", "calibrated_confidence"]].rename( + columns={"calibrated_confidence": "score_winnow"} + ), + on="peptide_key", + how="left", + ) + matched = matched.merge( + nb_matched[["peptide_key", "ALC (%)"]].rename( + columns={"ALC (%)": "score_novoboard"} + ), + on="peptide_key", + how="left", + ) + matched = matched.merge( + w_matched_raw[["peptide_key", "confidence"]].rename( + columns={"confidence": "raw_confidence"} + ), + on="peptide_key", + how="left", + ) + matched["score_glissade"] = confidence_to_log_prob(matched["raw_confidence"]) + + external = pd.DataFrame({"peptide_key": sorted(se_keys)}) + external = external.merge( + w_external[["peptide_key", "calibrated_confidence"]].rename( + columns={"calibrated_confidence": "score_winnow"} + ), + on="peptide_key", + how="left", + ) + external = external.merge( + nb_external[["peptide_key", "ALC (%)"]].rename( + columns={"ALC (%)": "score_novoboard"} + ), + on="peptide_key", + how="left", + ) + external = external.merge( + w_external_raw[["peptide_key", "confidence"]].rename( + columns={"confidence": "raw_confidence"} + ), + on="peptide_key", + how="left", + ) + external["score_glissade"] = confidence_to_log_prob(external["raw_confidence"]) + + # Attach NovoBoard pair metadata for twin TDC on mixture subsets. + nb_ext_meta = nb_unlab_f.loc[ + ~nb_unlab_f["proteome_hit"].astype(bool), + ["peptide_key", "Peptide", "ALC (%)", "spectrum_id", "_pair_key", "Scan"], + ].copy() + nb_ext_meta = ( + nb_ext_meta.sort_values("ALC (%)", ascending=False) + .groupby("peptide_key", as_index=False) + .first() + ) + external = external.merge( + nb_ext_meta.rename( + columns={ + "Peptide": "peptide_novoboard", + "spectrum_id": "spectrum_id_novoboard", + } + ), + on="peptide_key", + how="left", + ) + + nb_match_meta = ( + nb_correct[ + ["peptide_key", "Peptide", "ALC (%)", "spectrum_id", "_pair_key", "Scan"] + ] + .sort_values("ALC (%)", ascending=False) + .groupby("peptide_key", as_index=False) + .first() + .rename( + columns={ + "Peptide": "peptide_novoboard", + "spectrum_id": "spectrum_id_novoboard", + } + ) + ) + matched = matched.merge(nb_match_meta, on="peptide_key", how="left") + + twin_coverage_m = ( + float(matched["_pair_key"].notna().mean()) if len(matched) else 0.0 + ) + twin_coverage_e = ( + float(external["_pair_key"].notna().mean()) if len(external) else 0.0 + ) + if twin_coverage_m < 1.0 or twin_coverage_e < 1.0: + raise AssertionError( + f"{dataset} NovoBoard twin coverage incomplete: " + f"Sm={twin_coverage_m:.3f} Se={twin_coverage_e:.3f}" + ) + + glissade_reference = build_glissade_training_reference( + cfg.calibrator_train_metadata, min_length=labelled_min_length + ) + + logger.info( + "%s shared pools: matched=%d external=%d glissade_reference=%d " + "(shared Novor/proteome-hit labels; NB twin coverage 100%%)", + dataset, + len(matched), + len(external), + len(glissade_reference), + ) + return matched, external, nb_decoy_combined, glissade_reference + + +def _estimate_winnow_q_values(mixed: pd.DataFrame) -> pd.DataFrame: + work = mixed[["peptide_key", "score_winnow", "source"]].copy() + work = work.rename(columns={"score_winnow": "score"}) + work = work.dropna(subset=["score"]) + ctrl = NonParametricFDRControl() + ctrl.fit(work["score"]) + q_table = ctrl.add_psm_q_value(work.copy(), "score") + return pd.DataFrame( + { + "peptide_key": q_table["peptide_key"], + "score": q_table["score"], + "source": q_table["source"], + "q_value": q_table["psm_q_value"], + "method": "Winnow", + } + ) + + +def _estimate_novoboard_q_values( + mixed: pd.DataFrame, + decoy_by_pair: pd.DataFrame, +) -> pd.DataFrame: + target = pd.DataFrame( + { + "Peptide": mixed["peptide_novoboard"].fillna(mixed["peptide_key"]), + "ALC (%)": mixed["score_novoboard"], + "spectrum_id": mixed.get( + "spectrum_id_novoboard", + pd.Series(np.arange(len(mixed)), dtype=str), + ), + "_pair_key": mixed["_pair_key"], + "Scan": mixed.get("Scan", mixed["_pair_key"]), + "source": mixed["source"], + "peptide_key": mixed["peptide_key"], + } + ) + target = target.dropna(subset=["ALC (%)", "_pair_key"]) + # Restrict twin TDC to the mixture peptide keys only. + table = novoboard_max_target_twin_decoy_tdc( + target, + pd.DataFrame(), + target_peptide_keys=set(target["peptide_key"].astype(str)), + decoy_by_pair=decoy_by_pair, + log_missing_twins=False, + ) + targets = table[table["is_target"]].copy() + # Map back sources for FDP. + source_map = mixed.set_index("peptide_key")["source"].to_dict() + key_col = "_peptide_key" if "_peptide_key" in targets.columns else "peptide_key" + targets["peptide_key"] = targets[key_col] + targets["source"] = targets["peptide_key"].map(source_map) + return pd.DataFrame( + { + "peptide_key": targets["peptide_key"], + "score": targets["ALC (%)"], + "source": targets["source"], + "q_value": targets["estimated_q_value"], + "method": "NovoBoard", + } + ) + + +def _estimate_glissade_q_values( + mixed: pd.DataFrame, + matched_reference: pd.DataFrame, + *, + glissade_repo: Path, + n_bootstraps: int, + rng: np.random.Generator, +) -> pd.DataFrame: + run_bootstraps, annotate_results, compute_fdr_transform = _load_glissade_functions( + glissade_repo + ) + # Glissade FDR is defined on the mixture scores; the reference is the + # training-split matched distribution and is never scored itself. + mixed_scores = mixed["score_glissade"].astype(float).to_numpy() + matched_scores = matched_reference["score_glissade"].astype(float).to_numpy() + if len(mixed_scores) < 10 or len(matched_scores) < 10: + raise ValueError("Glissade FDR requires ≥10 reference and mixture scores") + + peptides = mixed["peptide_key"].astype(str).tolist() + np.random.seed(int(rng.integers(0, 2**32 - 1))) + fdrs, grid, _ = run_bootstraps( + matched_scores, + mixed_scores, + n_bootstraps=n_bootstraps, + ) + out_peptides, peptide_fdrs, scores = annotate_results( + peptides, mixed_scores, fdrs, grid + ) + peptide_fdrs = compute_fdr_transform(peptide_fdrs) + source_map = mixed.set_index("peptide_key")["source"].to_dict() + return pd.DataFrame( + { + "peptide_key": out_peptides, + "score": scores, + "source": [source_map.get(p) for p in out_peptides], + "q_value": peptide_fdrs, + "method": "Glissade", + } + ) + + +def _mixture_result_rows( + *, + dataset: str, + pi0: float, + true_pi0: float, + holdout_frac: float, + seed: int, + iteration: int, + method: str, + q_table: pd.DataFrame, + correct_keys: set[str], + n_external: int, + thresholds: list[float], + q_ref: np.ndarray | None = None, +) -> list[dict[str, object]]: + """Build per-threshold result rows for one method on one mixture.""" + work = q_table.dropna(subset=["q_value"]) + if q_ref is not None: + if len(q_ref) != len(q_table): + raise ValueError( + f"q_ref length {len(q_ref)} does not match q_table length {len(q_table)}" + ) + q_ref_aligned = pd.Series(np.asarray(q_ref, dtype=float), index=q_table.index) + q_ref_work = q_ref_aligned.loc[work.index].to_numpy(dtype=float) + q_devs = mean_abs_q_dev_vs_reference( + work["q_value"].to_numpy(dtype=float), q_ref_work, thresholds + ) + else: + q_devs = [float("nan")] * len(thresholds) + + rows: list[dict[str, object]] = [] + for threshold, q_dev in zip(thresholds, q_devs): + accepted = work[work["q_value"] <= threshold] + n_accepted = len(accepted) + n_true = int(accepted["peptide_key"].isin(correct_keys).sum()) + n_false = n_accepted - n_true + rows.append( + { + "dataset": dataset, + "pi0_target": float(pi0), + "true_pi0": float(true_pi0), + "holdout_frac": float(holdout_frac), + "seed": seed, + "iteration": iteration, + "method": method, + "q_value_threshold": float(threshold), + "mixed_external_peptides": n_external, + "correct_peptides": len(correct_keys), + "accepted_peptides": n_accepted, + "true_correct_peptides": n_true, + "false_external_peptides": n_false, + "observed_fdp": (n_false / n_accepted if n_accepted else np.nan), + "correct_discovery_pct": ( + 100.0 * n_true / len(correct_keys) if correct_keys else np.nan + ), + "mean_abs_q_dev_vs_db": float(q_dev), + } + ) + return rows + + +def _evaluate_mixture_methods( + *, + mixed: pd.DataFrame, + estimator_reference: pd.DataFrame, + decoy_by_pair: pd.DataFrame, + glissade_repo: Path, + n_bootstraps: int, + iter_rng: np.random.Generator, + dataset: str, + pi0: float, + true_pi0: float, + holdout_frac: float, + seed: int, + iteration: int, + correct_keys: set[str], + n_external: int, + thresholds: list[float], +) -> list[dict[str, object]]: + """Run Winnow / NovoBoard / Glissade FDR on one mixture and collect rows.""" + estimators: dict[str, object] = { + "Winnow": lambda m, _r: _estimate_winnow_q_values(m), + "NovoBoard": lambda m, _r: _estimate_novoboard_q_values(m, decoy_by_pair), + "Glissade": lambda m, r, rng=iter_rng: _estimate_glissade_q_values( + m, + r, + glissade_repo=glissade_repo, + n_bootstraps=n_bootstraps, + rng=rng, + ), + } + # Raw-confidence DBG on the mixture used for NovoBoard and Glissade q-deviation. Winnow keeps calibrated-score DBG unset here. + is_correct = mixed["peptide_key"].astype(str).isin(correct_keys).to_numpy() + q_db_raw = database_grounded_q_from_labels( + mixed["score_novoboard"].to_numpy(dtype=float), is_correct + ) + q_db_raw_by_key = dict(zip(mixed["peptide_key"].astype(str), q_db_raw, strict=True)) + + rows: list[dict[str, object]] = [] + mixture_keys = set(mixed["peptide_key"].astype(str)) + for method, estimator in estimators.items(): + try: + q_table = estimator(mixed, estimator_reference) # type: ignore[operator] + except Exception as exc: # noqa: BLE001 - boundary around external tool + logger.warning( + "%s FDR failed dataset=%s pi0=%.3g iter=%d: %s", + method, + dataset, + pi0, + iteration, + exc, + ) + continue + scored_keys = set(q_table["peptide_key"].astype(str)) + if scored_keys != mixture_keys: + raise AssertionError( + f"{method} scored a different mixture on dataset={dataset} " + f"pi0={pi0:.3g} iter={iteration}: mixture={len(mixture_keys)} " + f"scored={len(scored_keys)} " + f"missing={len(mixture_keys - scored_keys)} " + f"extra={len(scored_keys - mixture_keys)}" + ) + q_ref: np.ndarray | None = None + if method in ("NovoBoard", "Glissade"): + q_ref = ( + q_table["peptide_key"] + .astype(str) + .map(q_db_raw_by_key) + .to_numpy(dtype=float) + ) + rows.extend( + _mixture_result_rows( + dataset=dataset, + pi0=pi0, + true_pi0=true_pi0, + holdout_frac=holdout_frac, + seed=seed, + iteration=iteration, + method=method, + q_table=q_table, + correct_keys=correct_keys, + n_external=n_external, + thresholds=thresholds, + q_ref=q_ref, + ) + ) + return rows + + +def _sample_correct_component( + *, + dataset: str, + matched: pd.DataFrame, + external: pd.DataFrame, + pi0_grid: list[float], + holdout_frac: float, + rng: np.random.Generator, +) -> tuple[pd.DataFrame, set[str], float]: + """Sample S_c from matched, capped so the π₀ grid fits in S_e.""" + n_from_frac = max(1, int(round(len(matched) * holdout_frac))) + n_from_frac = min(n_from_frac, len(matched)) + n_cap = max_correct_pool_for_pi0_grid(len(external), pi0_grid) + if n_cap < 1: + raise ValueError( + f"{dataset}: external pool of {len(external)} cannot support any " + f"π₀ in {pi0_grid} without replacement" + ) + n_correct = min(n_from_frac, n_cap) + if n_correct < n_from_frac: + logger.info( + "%s capping |S_c| from %d to %d so π₀ grid %s fits in |S_e|=%d " + "without replacement", + dataset, + n_from_frac, + n_correct, + pi0_grid, + len(external), + ) + correct = matched.sample(n=n_correct, random_state=int(rng.integers(0, 2**32 - 1))) + correct = correct.copy() + correct["source"] = "correct" + correct_keys = set(correct["peptide_key"]) + effective_holdout_frac = n_correct / len(matched) if len(matched) else 0.0 + overlap = correct_keys & set(external["peptide_key"]) + if overlap: + raise AssertionError( + f"{dataset} S_c and S_e share {len(overlap)} peptide keys; " + "mixture sources would be ambiguous" + ) + return correct, correct_keys, effective_holdout_frac + + +def _external_draw_size( + *, + dataset: str, + pi0: float, + n_correct: int, + n_external_pool: int, +) -> int | None: + """Return |S_e'| for π₀, or None to skip; raise if the pool is too small.""" + if not 0.0 < pi0 < 1.0: + return None + n_external = int(round(pi0 / (1.0 - pi0) * n_correct)) + if n_external < 1: + logger.warning("Skipping pi0=%.3g: requested |S_e'|=%d", pi0, n_external) + return None + if n_external > n_external_pool: + raise AssertionError( + f"{dataset} pi0={pi0:.3g}: |S_e'|={n_external} exceeds pool " + f"{n_external_pool} after |S_c| cap {n_correct}" + ) + return n_external + + +def evaluate_controlled_mixtures( + *, + dataset: str, + matched: pd.DataFrame, + external: pd.DataFrame, + nb_decoy: pd.DataFrame, + glissade_reference: pd.DataFrame, + pi0_grid: list[float], + holdout_frac: float, + seed: int, + n_iterations: int, + thresholds: list[float], + glissade_repo: Path, + n_bootstraps: int, +) -> list[dict[str, object]]: + """Evaluate all methods on shared mixtures with controlled π₀. + + Uses as much of the matched pool as possible (up to holdout_frac), capped + so every π₀ in pi0_grid can draw its null component from S_e without + replacement. The null component is drawn without replacement, so mixture + peptide keys are unique and every method realises the same π₀ on the same + rows. + """ + rng = np.random.default_rng(seed) + correct, correct_keys, effective_holdout_frac = _sample_correct_component( + dataset=dataset, + matched=matched, + external=external, + pi0_grid=pi0_grid, + holdout_frac=holdout_frac, + rng=rng, + ) + decoy_by_pair = prepare_novoboard_decoy_by_pair(nb_decoy, already_filtered=True) + + rows: list[dict[str, object]] = [] + for pi0 in pi0_grid: + n_external = _external_draw_size( + dataset=dataset, + pi0=pi0, + n_correct=len(correct_keys), + n_external_pool=len(external), + ) + if n_external is None: + continue + + for iteration in range(n_iterations): + iter_rng = np.random.default_rng( + seed + 10_000 * iteration + int(1000 * pi0) + ) + ext_sample = external.sample( + n=n_external, + replace=False, + random_state=int(iter_rng.integers(0, 2**32 - 1)), + ).copy() + ext_sample["source"] = "external" + mixed = pd.concat([ext_sample, correct], ignore_index=True, sort=False) + if mixed["peptide_key"].duplicated().any(): + raise AssertionError( + f"{dataset} mixture has duplicate peptide keys at " + f"pi0={pi0:.3g} iter={iteration}" + ) + true_pi0 = len(ext_sample) / (len(ext_sample) + len(correct_keys)) + rows.extend( + _evaluate_mixture_methods( + mixed=mixed, + estimator_reference=glissade_reference, + decoy_by_pair=decoy_by_pair, + glissade_repo=glissade_repo, + n_bootstraps=n_bootstraps, + iter_rng=iter_rng, + dataset=dataset, + pi0=pi0, + true_pi0=true_pi0, + holdout_frac=effective_holdout_frac, + seed=seed, + iteration=iteration, + correct_keys=correct_keys, + n_external=len(ext_sample), + thresholds=thresholds, + ) + ) + return rows + + +def _plot_metric_by_pi0( + dataset_results: pd.DataFrame, + *, + dataset_label: str, + dataset_slug: str, + metric: str, + ylabel: str, + base_name: str, + percent_axis: bool, + output_dir: Path, +) -> None: + method_order = list(METHODS) + colors = {"Winnow": _PALETTE[0], "NovoBoard": _PALETTE[2], "Glissade": _PALETTE[4]} + pi0_values = sorted(dataset_results["pi0_target"].dropna().unique()) + n_pi0 = max(1, len(pi0_values)) + fig, axes = plt.subplots( + 1, n_pi0, figsize=(4.2 * n_pi0, 5.5), sharey=True, squeeze=False + ) + for ax, pi0 in zip(axes[0], pi0_values): + sub = dataset_results[dataset_results["pi0_target"] == pi0] + for method in method_order: + msub = sub[sub["method"] == method] + if msub.empty: + continue + summary = ( + msub.groupby("q_value_threshold", as_index=False)[metric] + .mean(numeric_only=True) + .sort_values("q_value_threshold") + ) + ax.plot( + summary["q_value_threshold"], + summary[metric], + lw=1.5, + label=method, + color=colors[method], + ) + if metric == "observed_fdp": + max_threshold = float(sub["q_value_threshold"].max()) + ax.plot( + [0.0, max_threshold], + [0.0, max_threshold], + color="#666666", + lw=1, + ls="--", + label="Nominal FDR", + ) + ax.set_ylim(bottom=0) + if percent_axis: + ax.set_ylim(0, 100) + ax.set_xlim(0, float(sub["q_value_threshold"].max())) + ax.set_xlabel("Estimated q-value threshold") + ax.set_title(f"π₀={pi0:g}") + _style_ax(ax) + axes[0][0].set_ylabel(ylabel) + axes[0][0].legend(loc="best", fontsize=9) + fig.suptitle(f"{dataset_label} external peptide score-mixture benchmark", y=1.02) + fig.tight_layout() + _save_fig(fig, output_dir / f"{base_name}_{dataset_slug}") + + +def plot_benchmark_results(results: pd.DataFrame, output_dir: Path) -> None: + """Save FDP and recovery plots faceted by π₀.""" + if results.empty: + return + output_dir.mkdir(parents=True, exist_ok=True) + specs = [ + ("observed_fdp", "Observed FDP", "external_peptide_score_mixture_fdp", False), + ( + "correct_discovery_pct", + "Correct peptide recovery\n(% of held-out correct peptides)", + "external_peptide_score_mixture_correct_discovery_pct", + True, + ), + ] + for dataset, dataset_results in results.groupby("dataset", sort=False): + for metric, ylabel, base_name, percent_axis in specs: + _plot_metric_by_pi0( + dataset_results, + dataset_label=_display_name(str(dataset)), + dataset_slug=str(dataset).replace("/", "_"), + metric=metric, + ylabel=ylabel, + base_name=base_name, + percent_axis=percent_axis, + output_dir=output_dir, + ) + + +def write_holdout_summary_tables( + results: pd.DataFrame, + output_dir: Path, + *, + thresholds: list[float] | None = None, +) -> tuple[Path, Path]: + """Aggregate raw mixture rows into acceptance and error/gain CSVs.""" + acceptance, error_gain = summarise_holdout_results( + results, + thresholds=thresholds if thresholds is not None else SUMMARY_THRESHOLDS, + group_extra=("pi0_target",), + ) + return write_summary_tables( + acceptance, error_gain, output_dir, "external_peptide_holdout" + ) + + +@app.command() +def main( + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for benchmark outputs."), + ] = DEFAULT_OUTPUT_DIR, + datasets: Annotated[ + Optional[list[str]], + typer.Option("--datasets", help="Dataset keys to benchmark."), + ] = None, + pi0_grid: Annotated[ + Optional[list[float]], + typer.Option("--pi0-grid", help="Target mixture null fractions."), + ] = None, + holdout_frac: Annotated[ + float, + typer.Option( + "--holdout-frac", + help=( + "Maximum fraction of shared matched peptides used as S_c " + "(default 1 = prefer the full pool). |S_c| is further capped so " + "every π₀ in --pi0-grid fits in S_e without replacement." + ), + ), + ] = DEFAULT_HOLDOUT_FRAC, + q_thresholds: Annotated[ + Optional[list[float]], + typer.Option( + "--q-thresholds", help="Estimated q-value thresholds to evaluate." + ), + ] = None, + seed: Annotated[int, typer.Option("--seed", help="Random seed.")] = DEFAULT_SEED, + n_iterations: Annotated[ + int, + typer.Option( + "--n-iterations", help="Number of external-score resampling iterations." + ), + ] = DEFAULT_N_ITERATIONS, + winnow_results: Annotated[ + Path, + typer.Option("--winnow-results", help="Winnow results directory."), + ] = DEFAULT_WINNOW_RESULTS, + novoboard_root: Annotated[ + Path, + typer.Option("--novoboard-root", help="NovoBoard datasets root."), + ] = DEFAULT_NOVOBOARD_ROOT, + model_root: Annotated[ + Path, + typer.Option( + "--model-root", + help="Per-dataset calibrator directories, used for Glissade's anchor.", + ), + ] = DEFAULT_MODEL_ROOT, + glissade_repo: Annotated[ + Path, + typer.Option("--glissade-repo", help="Glissade repository root."), + ] = DEFAULT_GLISSADE_REPO, + n_bootstraps: Annotated[ + int, + typer.Option( + "--n-bootstraps", help="Glissade bootstraps per mixture iteration." + ), + ] = DEFAULT_N_BOOTSTRAPS, + min_peptide_length: Annotated[ + int, + typer.Option( + "--min-peptide-length", + help=( + "Minimum normalised peptide length for unlabelled / external " + "pools. Labelled matched pools and Glissade's training reference " + "use the labelled floor (non-empty key only)." + ), + ), + ] = MIN_PEPTIDE_LENGTH, + plot: Annotated[bool, typer.Option(help="Create summary plots.")] = True, + summarise_only: Annotated[ + Optional[Path], + typer.Option( + "--summarise-only", + help="Only write summary CSVs/plots from an existing results CSV.", + ), + ] = None, +) -> None: + """Run the controlled-π₀ external peptide score-mixture benchmark.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + output_dir.mkdir(parents=True, exist_ok=True) + + if summarise_only is not None: + results = pd.read_csv(summarise_only) + write_holdout_summary_tables(results, output_dir) + if plot: + plot_benchmark_results(results, output_dir / "plots") + return + + dataset_keys = datasets if datasets is not None else DEFAULT_DATASETS + pi0s = pi0_grid if pi0_grid is not None else list(DEFAULT_PI0_GRID) + thresholds = ( + q_thresholds if q_thresholds is not None else list(DEFAULT_Q_THRESHOLDS) + ) + + rows: list[dict[str, object]] = [] + for dataset in dataset_keys: + logger.info("Building shared score tables for %s", _display_name(dataset)) + matched, external, nb_decoy, glissade_reference = build_shared_score_tables( + dataset=dataset, + winnow_results=winnow_results, + novoboard_root=novoboard_root, + model_root=model_root, + unlabelled_min_length=min_peptide_length, + labelled_min_length=LABELLED_MIN_PEPTIDE_LENGTH, + ) + rows.extend( + evaluate_controlled_mixtures( + dataset=dataset, + matched=matched, + external=external, + nb_decoy=nb_decoy, + glissade_reference=glissade_reference, + pi0_grid=pi0s, + holdout_frac=holdout_frac, + seed=seed, + n_iterations=n_iterations, + thresholds=thresholds, + glissade_repo=glissade_repo, + n_bootstraps=n_bootstraps, + ) + ) + + results = pd.DataFrame(rows) + results_path = output_dir / "external_peptide_holdout_results.csv" + results.to_csv(results_path, index=False) + logger.info("Wrote %s (%d rows)", results_path, len(results)) + write_holdout_summary_tables(results, output_dir) + if plot: + plot_benchmark_results(results, output_dir / "plots") + + +if __name__ == "__main__": + app() diff --git a/scripts/run_feature_ablations.py b/scripts/run_feature_ablations.py new file mode 100644 index 00000000..9da59237 --- /dev/null +++ b/scripts/run_feature_ablations.py @@ -0,0 +1,1532 @@ +"""Feature ablation study for Winnow calibrator. + +Trains MLP calibrators on subsets of pre-computed training feature matrices, +computes features from raw spectra for evaluation datasets, and produces +publication-quality plots of calibration, discrimination, and FDR behavior. +""" + +from __future__ import annotations + +import json +import logging +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Iterable, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import torch +import typer +from rich.logging import RichHandler + +_REPO_ROOT = Path(__file__).resolve().parent.parent +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from scripts.feature_subsets import FEATURE_SUBSETS # noqa: E402 +from scripts.plot_ablation_summary import ( # noqa: E402 + FDR_BIAS_COLUMN_BY_THRESHOLD, + Q_DEV_COLUMN_BY_THRESHOLD, + TAIL_ECE_COLUMN_BY_THRESHOLD, + assign_ablation_colors, + compute_ece, + compute_fdr_bias_at_fdr_thresholds, + compute_pr_auc, + compute_q_value_deviations, + compute_tail_ece_at_fdr, + ordered_ablation_configs, +) + +from winnow.calibration.calibrator import ProbabilityCalibrator # noqa: E402 +from winnow.datasets.feature_dataset import FeatureDataset # noqa: E402 +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# --------------------------------------------------------------------------- +# Plot theme — Paul Tol qualitative (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "HCT116": "Human colon", + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD004732": "ProteomeTools-1", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "Arabidopsis": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", +} + +# --------------------------------------------------------------------------- +# Feature group definitions (reduced set: no xcorr, spectral_angle, gap/similarity, edit_distance) +# --------------------------------------------------------------------------- +_EXCLUDED_REDUCED = frozenset( + { + "xcorr", + "spectral_angle", + "complementary_ion_count", + "max_ion_gap", + "edit_distance", + } +) + +BEAM_COLUMNS = ["margin", "median_margin", "entropy", "z-score", "edit_distance"] +TOKEN_COLUMNS = ["min_token_probability", "std_token_probability"] +FRAGMENT_MATCH_COLUMNS = [ + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", +] +RETENTION_TIME_COLUMNS = ["irt_error"] +MASS_ERROR_PPM = "mass_error_ppm" +MASS_ERROR_DA = "mass_error_da" + +REDUCED_BEAM_COLUMNS = [c for c in BEAM_COLUMNS if c not in _EXCLUDED_REDUCED] +REDUCED_FRAGMENT_COLUMNS = [ + c for c in FRAGMENT_MATCH_COLUMNS if c not in _EXCLUDED_REDUCED +] + +# Default training matrix columns (train_extra_small_matrix.parquet). +REDUCED_TRAIN_COLUMNS: list[str] = FEATURE_SUBSETS["no_fragment_similarity"]["columns"] + +# Hydra overrides aligned with Makefile ANALYSIS_REDUCED_FEATURE_OVERRIDES (mass_error_da model). +REDUCED_FEATURE_COMPUTE_OVERRIDES: list[str] = [ + "~calibrator.features.mass_error", + "+calibrator.features.mass_error_da._target_=winnow.calibration.calibration_features.MassErrorDaFeature", + "+calibrator.features.mass_error_da.residue_masses=${residue_masses}", + "+calibrator.features.fragment_match_features.excluded_columns=[spectral_angle,xcorr,complementary_ion_count,max_ion_gap]", + "+calibrator.features.beam_features.excluded_columns=[edit_distance]", +] + + +def _reference_model_columns(model_dir: Path | None) -> list[str] | None: + """Return ``feature_columns`` from a saved calibrator, if present.""" + if model_dir is None: + return None + config_path = model_dir / "config.json" + if not config_path.is_file(): + return None + with open(config_path) as f: + config = json.load(f) + cols = config.get("feature_columns") + return list(cols) if cols else None + + +def _resolve_mass_error_column( + df: pl.DataFrame, + reference_model_dir: Path | None, +) -> str: + """Pick mass-error column present in *df*, preferring the reference model.""" + ref_cols = _reference_model_columns(reference_model_dir) + if MASS_ERROR_DA in df.columns: + return MASS_ERROR_DA + if MASS_ERROR_PPM in df.columns: + if ref_cols and MASS_ERROR_DA in ref_cols: + logger.warning( + "Reference model uses %s but data has %s; using %s for ablations.", + MASS_ERROR_DA, + MASS_ERROR_PPM, + MASS_ERROR_PPM, + ) + return MASS_ERROR_PPM + raise ValueError( + f"No mass error column in data (tried {MASS_ERROR_DA}, {MASS_ERROR_PPM})" + ) + + +def _columns_available(df: pl.DataFrame, columns: list[str]) -> list[str]: + missing = [c for c in columns if c not in df.columns] + if missing: + raise ValueError(f"Missing columns: {missing}. Available: {df.columns}") + return columns + + +def resolve_all_feature_columns( + df: pl.DataFrame, + reference_model_dir: Path | None, +) -> list[str]: + """Full reduced feature set for the 'All features' ablation config.""" + ref_cols = _reference_model_columns(reference_model_dir) + if ref_cols: + cols = ["confidence"] + for col in ref_cols: + if ( + col == MASS_ERROR_DA + and col not in df.columns + and MASS_ERROR_PPM in df.columns + ): + cols.append(MASS_ERROR_PPM) + elif col in df.columns: + cols.append(col) + else: + cols = [c for c in REDUCED_TRAIN_COLUMNS if c in df.columns] + return _columns_available(df, cols) + + +def build_ablation_configs( + df: pl.DataFrame, + reference_model_dir: Path | None, +) -> dict[str, list[str]]: + """Build ablation configs using columns available in *df*.""" + mass_col = _resolve_mass_error_column(df, reference_model_dir) + all_features = resolve_all_feature_columns(df, reference_model_dir) + return { + "Confidence only": ["confidence"], + "Confidence + mass error": ["confidence", mass_col], + "Confidence + iRT error": ["confidence", *RETENTION_TIME_COLUMNS], + "Confidence + token-level": ["confidence", *TOKEN_COLUMNS], + "Confidence + beam search": ["confidence", *REDUCED_BEAM_COLUMNS], + "Confidence + fragment matching": ["confidence", *REDUCED_FRAGMENT_COLUMNS], + "All features": all_features, + } + + +ABLATION_CONFIGS: dict[str, list[str]] = {} + +ABLATION_COLORS: dict[str, str] = {} + + +def _dataset_display_name(key: str) -> str: + """Publication-ready dataset label for plot titles.""" + if key in EVAL_DATASETS: + return str(EVAL_DATASETS[key]["label"]) + return DATASET_DISPLAY_NAMES.get(key, key) + + +def _configure_ablation_colors(config_names: Iterable[str]) -> None: + global ABLATION_COLORS + ABLATION_COLORS = assign_ablation_colors( + ordered_ablation_configs(set(config_names)) + ) + + +# Default training hyperparameters (overridden by --hyperparams-from-model). +TRAIN_HYPERPARAMS = { + "hidden_dims": [128, 64], + "learning_rate": 0.0001, + "weight_decay": 0.0001, + "batch_size": 4096, + "max_epochs": 200, + "n_iter_no_change": 10, + "tol": 1e-4, +} + + +def train_hyperparams_from_model(model_dir: Path) -> dict[str, object]: + """Load MLP training hyperparameters from a saved calibrator ``config.json``.""" + config_path = model_dir / "config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"No config.json at {model_dir}") + with open(config_path) as f: + config = json.load(f) + return { + "hidden_dims": tuple(config["hidden_dims"]), + "dropout": config["dropout"], + "learning_rate": config["learning_rate"], + "weight_decay": config["weight_decay"], + "batch_size": config["batch_size"], + "max_epochs": config["max_epochs"], + "n_iter_no_change": config["n_iter_no_change"], + "tol": config["tol"], + } + + +EVAL_DATASETS = { + "HCT116": { + "label": "Human colon", + "spectra": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.parquet", + "predictions": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.csv", + "koina_mode": "columns", + }, + "Arabidopsis": { + "label": "Arabidopsis", + "spectra": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.parquet", + "predictions": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.csv", + "koina_mode": "columns", + }, + "PXD023064": { + "label": "Immunopeptidomics-2", + "spectra": "held_out_projects/lcfm/PXD023064/", + "predictions": "held_out_projects/lcfm/PXD023064_predictions/PXD023064.csv", + "koina_mode": "columns", + }, +} + +# Residue masses for DatabaseGroundedFDRControl (loaded from config at runtime) +_RESIDUE_MASSES: dict[str, float] | None = None + + +def _get_residue_masses() -> dict[str, float]: + """Load residue masses from the winnow residues config.""" + global _RESIDUE_MASSES + if _RESIDUE_MASSES is None: + import yaml + + config_path = ( + Path(__file__).resolve().parent.parent + / "winnow" + / "configs" + / "residues.yaml" + ) + with open(config_path) as f: + cfg = yaml.safe_load(f) + _RESIDUE_MASSES = cfg["residue_masses"] + return _RESIDUE_MASSES + + +# --------------------------------------------------------------------------- +# Eval feature computation +# --------------------------------------------------------------------------- +def _feature_compute_overrides(reference_model_dir: Path | None) -> list[str]: + """Hydra overrides so eval features match a mass_error_da / reduced-feature model.""" + ref_cols = _reference_model_columns(reference_model_dir) + if ref_cols and MASS_ERROR_DA in ref_cols: + return list(REDUCED_FEATURE_COMPUTE_OVERRIDES) + return [ + "+calibrator.features.fragment_match_features.excluded_columns=[spectral_angle,xcorr,complementary_ion_count,max_ion_gap]", + "+calibrator.features.beam_features.excluded_columns=[edit_distance]", + ] + + +def _compute_eval_features_for_dataset( + name: str, + spectra_path: str, + predictions_path: str, + cache_dir: Path, + koina_url: str, + koina_ssl: bool, + koina_mode: str = "columns", + feature_overrides: list[str] | None = None, +) -> Path: + """Compute the full feature matrix for an eval dataset and cache as Parquet. + + Args: + koina_mode: ``"columns"`` to read collision_energy / frag_type from + per-row metadata columns, or ``"constants"`` to use fixed values + (CE=27, HCD). + """ + cache_path = cache_dir / f"{name}.parquet" + if cache_path.exists(): + logger.info("Using cached eval features for %s at %s", name, cache_path) + return cache_path + + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + + logger.info("Computing features for eval dataset %s ...", name) + + if koina_mode == "columns": + koina_overrides = [ + "+koina.input_columns.collision_energies=collision_energy", + "+koina.input_columns.fragmentation_types=frag_type", + "+calibrator.features.fragment_match_features.model_input_columns.collision_energies=collision_energy", + "+calibrator.features.fragment_match_features.model_input_columns.fragmentation_types=frag_type", + ] + else: + koina_overrides = [ + "+koina.input_constants.collision_energies=27", + "+koina.input_constants.fragmentation_types=HCD", + "+calibrator.features.fragment_match_features.model_input_constants.collision_energies=27", + "+calibrator.features.fragment_match_features.model_input_constants.fragmentation_types=HCD", + ] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name=f"winnow_ablation_features_{name}", + ): + cfg = compose( + config_name="compute_features", + overrides=[ + f"dataset.spectrum_path_or_directory={spectra_path}", + f"dataset.predictions_path={predictions_path}", + f"koina.server_url={koina_url}", + f"koina.ssl={koina_ssl}", + *koina_overrides, + *(feature_overrides or []), + "labelled=true", + "filter_empty_predictions=true", + ], + ) + + data_loader = instantiate(cfg.data_loader) + calibrator = instantiate(cfg.calibrator) + + from winnow.scripts.main import ( + _compute_features_batched_metadata, + ) + + spectrum_path = Path(spectra_path) + preds_path = cfg.dataset.get("predictions_path", predictions_path) + + all_metadata = _compute_features_batched_metadata( + spectrum_path, + preds_path, + data_loader, + calibrator, + labelled=True, + filter_empty=True, + ) + + combined_metadata = pd.concat(all_metadata, ignore_index=True) + logger.info( + " %s: %d spectra after feature computation", name, len(combined_metadata) + ) + + # Write the training matrix parquet with all feature columns + correct + extra cols for FDR + feature_columns = ["confidence"] + calibrator.columns + keep_cols = list(feature_columns) + if "correct" in combined_metadata.columns: + keep_cols.append("correct") + if "sequence" in combined_metadata.columns: + keep_cols.append("sequence") + if "prediction" in combined_metadata.columns: + keep_cols.append("prediction") + if "precursor_mz" in combined_metadata.columns: + keep_cols.append("precursor_mz") + if "precursor_charge" in combined_metadata.columns: + keep_cols.append("precursor_charge") + + # Deduplicate while preserving order + seen = set() + unique_cols = [] + for c in keep_cols: + if c not in seen and c in combined_metadata.columns: + seen.add(c) + unique_cols.append(c) + + training_df = pl.from_pandas(combined_metadata[unique_cols]) + cache_dir.mkdir(parents=True, exist_ok=True) + training_df.write_parquet(cache_path) + logger.info( + " Cached eval features to %s (%d rows, %d cols)", + cache_path, + len(training_df), + len(training_df.columns), + ) + return cache_path + + +def compute_all_eval_features( + output_dir: Path, + koina_url: str, + koina_ssl: bool, + astral_spectra: str | None, + astral_predictions: str | None, + skip_feature_compute: bool, + reference_model_dir: Path | None = None, +) -> dict[str, Path]: + """Compute (or locate cached) eval feature Parquets for all datasets.""" + cache_dir = output_dir / "eval_feature_cache" + result: dict[str, Path] = {} + feature_overrides = _feature_compute_overrides(reference_model_dir) + + for name, info in EVAL_DATASETS.items(): + if skip_feature_compute: + cache_path = cache_dir / f"{name}.parquet" + if not cache_path.exists(): + raise FileNotFoundError( + f"--skip-feature-compute set but cache not found: {cache_path}" + ) + result[name] = cache_path + else: + result[name] = _compute_eval_features_for_dataset( + name, + info["spectra"], + info["predictions"], + cache_dir, + koina_url, + koina_ssl, + koina_mode=info.get("koina_mode", "columns"), + feature_overrides=feature_overrides, + ) + + if astral_spectra and astral_predictions: + name = "Astral" + if skip_feature_compute: + cache_path = cache_dir / f"{name}.parquet" + if not cache_path.exists(): + raise FileNotFoundError( + f"--skip-feature-compute set but cache not found: {cache_path}" + ) + result[name] = cache_path + else: + result[name] = _compute_eval_features_for_dataset( + name, + astral_spectra, + astral_predictions, + cache_dir, + koina_url, + koina_ssl, + feature_overrides=feature_overrides, + ) + + return result + + +# --------------------------------------------------------------------------- +# Training +# --------------------------------------------------------------------------- +def _load_parquet_as_polars(path: str | Path) -> pl.DataFrame: + """Load a Parquet file or directory of Parquets into a single Polars DataFrame.""" + path = Path(path) + if path.is_dir(): + parquet_files = sorted(path.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError(f"No .parquet files in {path}") + return pl.concat([pl.read_parquet(f) for f in parquet_files]) + return pl.read_parquet(path) + + +def split_train_val_frames( + df: pl.DataFrame, + validation_fraction: float, + seed: int, +) -> tuple[pl.DataFrame, pl.DataFrame]: + """Random train/validation split with a fixed index permutation. + + Uses the same scheme as ``winnow.scripts.main._maybe_split_calibration_dataset``: + shuffle all row indices with *seed*, then assign the last ``validation_fraction`` + fraction to validation. The same split is reused for every ablation config. + """ + if "correct" not in df.columns: + raise ValueError("Training Parquet must contain a 'correct' column") + if not 0 < validation_fraction < 1: + raise ValueError( + f"validation_fraction must be in (0, 1), got {validation_fraction}" + ) + + n = len(df) + n_val = max(1, int(n * validation_fraction)) + rng = np.random.default_rng(seed) + perm = rng.permutation(n) + train_df = df[perm[: n - n_val].tolist()] + val_df = df[perm[n - n_val :].tolist()] + logger.info( + "Train/val split: %d train, %d val (fraction=%.2f, seed=%d)", + len(train_df), + len(val_df), + validation_fraction, + seed, + ) + return train_df, val_df + + +def _column_slice_to_feature_dataset( + df: pl.DataFrame, columns: list[str] +) -> FeatureDataset: + """Select columns from a Polars DataFrame and build a FeatureDataset.""" + if "correct" not in df.columns: + raise ValueError("Parquet must contain a 'correct' column") + missing = [c for c in columns if c not in df.columns] + if missing: + raise ValueError( + f"Missing columns in Parquet: {missing}. Available: {df.columns}" + ) + features = df.select(columns).to_numpy().astype(np.float32) + labels = df["correct"].to_numpy().astype(np.float32) + return FeatureDataset(features=features, labels=labels) + + +def _config_dir_name(config_name: str) -> str: + """Derive the on-disk directory name for an ablation config.""" + return config_name.lower().replace(" ", "_").replace("+", "and") + + +_LEGACY_DIR_NAMES: dict[str, list[str]] = { + "Confidence only": ["confidence_only"], + "Confidence + mass error": [ + "confidence_and_mass_error", + "confidence_and_mass_error_and_rt", + ], + "Confidence + iRT error": [ + "confidence_and_irt_error", + "confidence_and_mass_error_and_rt", + "confidence_and_fragment_matching", + "prosit", + ], + "Confidence + token-level": [ + "confidence_and_token_level", + "confidence_and_beam_search", + "beam_and_token", + ], + "Confidence + beam search": ["confidence_and_beam_search", "beam_and_token"], + "Confidence + fragment matching": [ + "confidence_and_fragment_matching", + "prosit", + ], + "All features": ["all_features", "full_model"], +} + + +def _resolve_model_dir(output_dir: Path, config_name: str) -> Path: + """Find the model directory for a config, falling back to legacy names.""" + candidates = _LEGACY_DIR_NAMES.get(config_name, [_config_dir_name(config_name)]) + for candidate in candidates: + model_dir = output_dir / "models" / candidate + if model_dir.exists(): + return model_dir + raise FileNotFoundError( + f"No saved model found for '{config_name}'. " + f"Searched: {[str(output_dir / 'models' / c) for c in candidates]}" + ) + + +def train_ablation_models( + train_df: pl.DataFrame, + val_df: pl.DataFrame, + output_dir: Path, + seed: int, + train_hyperparams: dict[str, object] | None = None, +) -> dict[str, ProbabilityCalibrator]: + """Train one calibrator per ablation config, return dict of fitted calibrators.""" + models: dict[str, ProbabilityCalibrator] = {} + hp = {**TRAIN_HYPERPARAMS, **(train_hyperparams or {})} + + for config_name, columns in ABLATION_CONFIGS.items(): + logger.info( + "Training ablation config: %s (%d features)", config_name, len(columns) + ) + + train_ds = _column_slice_to_feature_dataset(train_df, columns) + val_ds = _column_slice_to_feature_dataset(val_df, columns) + + calibrator = ProbabilityCalibrator( + seed=seed, + **hp, # type: ignore[arg-type] + ) + history = calibrator.fit_from_features(train_ds, val_ds) + + model_dir = output_dir / "models" / _config_dir_name(config_name) + ProbabilityCalibrator.save(calibrator, model_dir) + logger.info( + " Trained %s: %d epochs, best_epoch=%d", + config_name, + history.epochs_trained, + history.best_epoch, + ) + + models[config_name] = calibrator + + return models + + +def load_ablation_models( + output_dir: Path, +) -> dict[str, ProbabilityCalibrator]: + """Load pre-trained ablation calibrators from ``{output_dir}/models/``.""" + models: dict[str, ProbabilityCalibrator] = {} + + for config_name in ABLATION_CONFIGS: + model_dir = _resolve_model_dir(output_dir, config_name) + calibrator = ProbabilityCalibrator.load(model_dir) + logger.info(" Loaded %s from %s", config_name, model_dir) + models[config_name] = calibrator + + return models + + +# --------------------------------------------------------------------------- +# Evaluation helpers +# --------------------------------------------------------------------------- +def _predict_calibrated_scores( + calibrator: ProbabilityCalibrator, + features: np.ndarray, +) -> np.ndarray: + """Run forward pass through a fitted calibrator and return calibrated probabilities.""" + assert calibrator.network is not None + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + + device = next(calibrator.network.parameters()).device + x = torch.as_tensor(features, dtype=torch.float32, device=device) + x = (x - calibrator.feature_mean) / calibrator.feature_std + + calibrator.network.eval() + with torch.no_grad(): + logits = calibrator.network(x) + probs = torch.sigmoid(logits).cpu().numpy().flatten() + + return probs + + +def compute_precision_recall_curve( + dataset: pd.DataFrame, + confidence_column: str, + label_column: str, + name: str, +) -> pd.DataFrame: + """Non-standard cumulative PR curve matching the casanovo notebook.""" + original = dataset[[confidence_column, label_column]] + original = original.sort_values(by=confidence_column, ascending=False) + cum_correct = np.cumsum(original[label_column].values) + precision = cum_correct / np.arange(1, len(original) + 1) + recall = cum_correct / len(original) + metrics = pd.DataFrame({"precision": precision, "recall": recall}).reset_index( + drop=True + ) + metrics["name"] = name + return metrics + + +def compute_calibration_curve( + df: pd.DataFrame, + pred_col: str, + label_col: str, + name: str, + n_bins: int = 10, +) -> pd.DataFrame: + """Fixed-width bin calibration curve matching the casanovo notebook.""" + data = df[[pred_col, label_col]].dropna().copy(deep=True) + data[pred_col] = data[pred_col].clip(0.0, 1.0) + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_cats = pd.cut(data[pred_col], bins=bins, include_lowest=True) + bin_cats.name = "bin" + grouped = ( + data.groupby(bin_cats, observed=True) + .agg( + pred_mean=(pred_col, "mean"), + empirical=(label_col, "mean"), + count=(label_col, "size"), + ) + .reset_index() + ) + grouped = grouped[grouped["count"] > 0] + grouped["bin_center"] = grouped["bin"].apply(lambda iv: (iv.left + iv.right) / 2) + grouped["name"] = name + return grouped[["pred_mean", "empirical", "count", "bin_center", "name"]] + + +def compute_brier_score(pred: np.ndarray, labels: np.ndarray) -> float: + """Brier score.""" + return float(np.mean((pred - labels) ** 2)) + + +def compute_ids_at_fdr( + calibrated_scores: np.ndarray, + labels: np.ndarray, + fdr_threshold: float, +) -> int: + """Count PSMs accepted at a given FDR threshold using NonParametricFDRControl.""" + fdr_ctrl = NonParametricFDRControl() + scores_series = pd.Series(calibrated_scores, name="score") + fdr_ctrl.fit(dataset=scores_series) + cutoff = fdr_ctrl.get_confidence_cutoff(threshold=fdr_threshold) + if np.isnan(cutoff): + return 0 + return int((calibrated_scores >= cutoff).sum()) + + +@dataclass +class EvalResult: + """Metrics and curves for a single ablation config evaluated on one dataset.""" + + config_name: str + dataset_name: str + ece: float + tail_ece_at_5pct: float + tail_ece_at_10pct: float + brier: float + ids_at_1pct: int + ids_at_5pct: int + ids_at_10pct: int + pr_auc: float + fdr_bias_at_5pct: float + fdr_bias_at_10pct: float + q_dev_at_5pct: float + q_dev_at_10pct: float + pr_curve: pd.DataFrame = field(repr=False) + calibration_curve: pd.DataFrame = field(repr=False) + calibrated_scores: np.ndarray = field(repr=False) + labels: np.ndarray = field(repr=False) + raw_confidence: np.ndarray = field(repr=False) + eval_df: pd.DataFrame = field(repr=False) + + +def evaluate_single( + config_name: str, + calibrator: ProbabilityCalibrator, + columns: list[str], + eval_df: pl.DataFrame, + dataset_name: str, +) -> EvalResult: + """Evaluate a single ablation config on a single eval dataset.""" + features = eval_df.select(columns).to_numpy().astype(np.float32) + labels = eval_df["correct"].to_numpy().astype(np.float32) + raw_confidence = eval_df["confidence"].to_numpy().astype(np.float64) + + calibrated = _predict_calibrated_scores(calibrator, features) + + # Build a pandas DataFrame for PR / calibration / FDR computations + meta = pd.DataFrame( + { + "confidence": raw_confidence, + "calibrated_confidence": calibrated, + "correct": labels, + } + ) + + # Carry over sequence and prediction for database-grounded FDR if available + if "sequence" in eval_df.columns: + meta["sequence"] = eval_df["sequence"].to_pandas() + if "prediction" in eval_df.columns: + meta["prediction"] = eval_df["prediction"].to_pandas() + + pr = compute_precision_recall_curve( + meta, "calibrated_confidence", "correct", config_name + ) + + cal = compute_calibration_curve( + meta, "calibrated_confidence", "correct", config_name + ) + + ece = compute_ece(calibrated, labels) + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=pd.Series(calibrated, name="score")) + tail_ece_5 = compute_tail_ece_at_fdr(calibrated, labels, 0.05, fdr_ctrl=fdr_ctrl) + tail_ece_10 = compute_tail_ece_at_fdr(calibrated, labels, 0.10, fdr_ctrl=fdr_ctrl) + brier = compute_brier_score(calibrated, labels) + + ids_1 = compute_ids_at_fdr(calibrated, labels, 0.01) + ids_5 = compute_ids_at_fdr(calibrated, labels, 0.05) + ids_10 = compute_ids_at_fdr(calibrated, labels, 0.10) + + pr_auc = compute_pr_auc(meta) + fdr_bias = compute_fdr_bias_at_fdr_thresholds(meta) + q_dev = compute_q_value_deviations(meta) + + return EvalResult( + config_name=config_name, + dataset_name=dataset_name, + ece=ece, + tail_ece_at_5pct=tail_ece_5, + tail_ece_at_10pct=tail_ece_10, + brier=brier, + ids_at_1pct=ids_1, + ids_at_5pct=ids_5, + ids_at_10pct=ids_10, + pr_auc=pr_auc, + fdr_bias_at_5pct=fdr_bias[0.05], + fdr_bias_at_10pct=fdr_bias[0.10], + q_dev_at_5pct=q_dev[0.05], + q_dev_at_10pct=q_dev[0.10], + pr_curve=pr, + calibration_curve=cal, + calibrated_scores=calibrated, + labels=labels, + raw_confidence=raw_confidence, + eval_df=meta, + ) + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- +def _style_axes(ax: plt.Axes) -> None: + """Apply standard axes formatting: no grid, black spines.""" + ax.set_axisbelow(True) + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _save_fig(fig: plt.Figure, base_path: Path, plot_format: str) -> None: + """Save figure in the requested format(s).""" + if plot_format in ("pdf", "both"): + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + if plot_format in ("png", "both"): + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _lineplot( + ax: plt.Axes, + data: pd.DataFrame, + *, + x: str, + y: str, + label: str, + color: str, + linestyle: str = "-", + linewidth: float = 0.5, + marker: str | None = None, +) -> None: + """Line plot with consistent linewidth (seaborn, no auto legend).""" + kwargs: dict = { + "data": data, + "x": x, + "y": y, + "label": label, + "color": color, + "linestyle": linestyle, + "linewidth": linewidth, + "ax": ax, + "legend": False, + } + if marker is not None: + kwargs["marker"] = marker + sns.lineplot(**kwargs) + + +def _generate_plots_for_dataset( + ds_results: list[EvalResult], + ds_name: str, + plots_dir: Path, + plot_format: str, +) -> None: + """Generate all ablation figures for one dataset.""" + plot_precision_recall(ds_results, ds_name, plots_dir, plot_format) + plot_calibration(ds_results, ds_name, plots_dir, plot_format) + plot_fdr_vs_confidence(ds_results, ds_name, plots_dir, plot_format) + plot_fdr_accepted_psms(ds_results, ds_name, plots_dir, plot_format) + + +def plot_precision_recall( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """PR curve: one line per ablation config.""" + fig, ax = plt.subplots(figsize=(6, 4)) + + for r in results: + _lineplot( + ax, + r.pr_curve, + x="recall", + y="precision", + label=r.config_name, + color=ABLATION_COLORS[r.config_name], + ) + + display = _dataset_display_name(dataset_name) + ax.set( + xlabel="Recall", + ylabel="Precision", + title=f"{display} precision-recall by feature set", + ) + ax.legend(loc="lower left", fontsize=7) + _style_axes(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"pr_curve_{dataset_name}", plot_format) + + +def plot_calibration( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Calibration diagram: reliability curves + diagonal.""" + fig, ax = plt.subplots(figsize=(6, 4)) + + for r in results: + _lineplot( + ax, + r.calibration_curve, + x="pred_mean", + y="empirical", + label=r.config_name, + color=ABLATION_COLORS[r.config_name], + marker="o", + ) + + display = _dataset_display_name(dataset_name) + ax.plot([0, 1], [0, 1], ls="--", color="gray", lw=0.5) + ax.set( + xlabel="Mean predicted probability", + ylabel="Empirical accuracy\n(database label)", + title=f"{display} probability calibration by feature set", + ) + ax.legend(loc="lower right", fontsize=7) + _style_axes(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"calibration_{dataset_name}", plot_format) + + +def plot_fdr_vs_confidence( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """PSM FDR vs calibrated confidence: non-parametric vs database-grounded per config.""" + n_configs = len(results) + fig, axes = plt.subplots(1, n_configs, figsize=(5 * n_configs, 4), squeeze=False) + + residue_masses = _get_residue_masses() + + for i, r in enumerate(results): + ax = axes[0, i] + + np_fdr = NonParametricFDRControl() + np_fdr.fit(dataset=r.eval_df["calibrated_confidence"]) + winnow_metrics = np_fdr.add_psm_fdr( + r.eval_df.copy(), confidence_col="calibrated_confidence" + ) + + has_sequence = ( + "sequence" in r.eval_df.columns and "prediction" in r.eval_df.columns + ) + + if has_sequence: + dbg_fdr = DatabaseGroundedFDRControl( + confidence_feature="calibrated_confidence", + residue_masses=residue_masses, + ) + try: + dbg_fdr.fit(dataset=r.eval_df.copy()) + dbg_metrics = dbg_fdr.add_psm_fdr( + r.eval_df.copy(), confidence_col="calibrated_confidence" + ) + + sns.lineplot( + x=np.asarray(dbg_metrics["calibrated_confidence"], dtype=float), + y=np.asarray(dbg_metrics["psm_fdr"], dtype=float), + label="Database-grounded", + ax=ax, + color=_PALETTE[3], + linewidth=0.5, + legend=False, + ) + except Exception as e: + logger.warning( + "Database-grounded FDR failed for %s/%s: %s", + r.config_name, + dataset_name, + e, + ) + + sns.lineplot( + x=np.asarray(winnow_metrics["calibrated_confidence"], dtype=float), + y=np.asarray(winnow_metrics["psm_fdr"], dtype=float), + label="Winnow (non-parametric)", + ax=ax, + color=_PALETTE[0], + linewidth=0.5, + legend=False, + ) + + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM FDR") + ax.set_title(r.config_name) + ax.legend(fontsize=7) + _style_axes(ax) + + display = _dataset_display_name(dataset_name) + fig.suptitle( + f"{display} PSM FDR vs calibrated confidence by feature set", fontsize=12 + ) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_vs_confidence_{dataset_name}", plot_format) + + +def plot_fdr_accepted_psms( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Number of accepted PSMs vs q-value threshold.""" + fig, ax = plt.subplots(figsize=(6, 4)) + + thresholds = np.linspace(0.001, 0.10, 200) + + for r in results: + np_fdr = NonParametricFDRControl() + scores_series = pd.Series(r.calibrated_scores, name="score") + np_fdr.fit(dataset=scores_series) + + meta_with_q = np_fdr.add_psm_q_value( + pd.DataFrame({"calibrated_confidence": r.calibrated_scores}), + confidence_col="calibrated_confidence", + ) + + q_values = meta_with_q["psm_q_value"].values + counts = [] + for t in thresholds: + counts.append(int((q_values <= t).sum())) + + ax.plot( + thresholds, + counts, + label=r.config_name, + color=ABLATION_COLORS[r.config_name], + linewidth=0.5, + ) + + for fdr_line in [0.01, 0.05, 0.10]: + ax.axvline(fdr_line, ls="--", color="gray", lw=0.5, alpha=0.7) + + ax.relim() + ax.autoscale_view() + y_text = ax.get_ylim()[1] * 0.02 + for fdr_line in [0.01, 0.05, 0.10]: + ax.text( + fdr_line - 0.002, + y_text, + f"{fdr_line:.0%}", + ha="right", + va="bottom", + fontsize=7, + color="gray", + ) + + display = _dataset_display_name(dataset_name) + ax.set_xlabel("Non-parametric q-value threshold") + ax.set_ylabel("Accepted PSMs") + ax.set_title(f"{display} accepted PSMs at non-parametric q-value threshold") + ax.legend(loc="upper left", fontsize=7) + _style_axes(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_accepted_psms_{dataset_name}", plot_format) + + +# --------------------------------------------------------------------------- +# Saving eval results +# --------------------------------------------------------------------------- +def save_eval_results(all_results: list[EvalResult], output_dir: Path) -> None: + """Persist per-PSM eval DataFrames so plots can be reproduced without re-inference.""" + results_dir = output_dir / "eval_results" + results_dir.mkdir(parents=True, exist_ok=True) + + for r in all_results: + safe_config = r.config_name.lower().replace(" ", "_").replace("+", "and") + path = results_dir / f"{r.dataset_name}_{safe_config}.parquet" + df = r.eval_df.copy() + df["config_name"] = r.config_name + df["dataset_name"] = r.dataset_name + df.to_parquet(path, index=False) + + logger.info("Saved %d eval result Parquets to %s", len(all_results), results_dir) + + +def load_eval_results_for_plotting( + output_dir: Path, +) -> dict[str, list[EvalResult]]: + """Load saved eval Parquets and rebuild curve data for plotting.""" + results_dir = output_dir / "eval_results" + if not results_dir.is_dir(): + raise FileNotFoundError(f"No eval_results directory at {results_dir}") + + paths = sorted(results_dir.glob("*.parquet")) + if not paths: + raise FileNotFoundError(f"No eval result Parquets in {results_dir}") + + grouped: dict[str, list[EvalResult]] = defaultdict(list) + for path in paths: + df = pd.read_parquet(path) + config_name = str(df["config_name"].iloc[0]) + dataset_name = str(df["dataset_name"].iloc[0]) + meta = df.drop(columns=["config_name", "dataset_name"], errors="ignore") + calibrated = meta["calibrated_confidence"].to_numpy(dtype=np.float64) + labels = meta["correct"].to_numpy(dtype=np.float32) + pr = compute_precision_recall_curve( + meta, "calibrated_confidence", "correct", config_name + ) + cal = compute_calibration_curve( + meta, "calibrated_confidence", "correct", config_name + ) + grouped[dataset_name].append( + EvalResult( + config_name=config_name, + dataset_name=dataset_name, + ece=0.0, + tail_ece_at_5pct=float("nan"), + tail_ece_at_10pct=float("nan"), + brier=0.0, + ids_at_1pct=0, + ids_at_5pct=0, + ids_at_10pct=0, + pr_auc=0.0, + fdr_bias_at_5pct=float("nan"), + fdr_bias_at_10pct=float("nan"), + q_dev_at_5pct=float("nan"), + q_dev_at_10pct=float("nan"), + pr_curve=pr, + calibration_curve=cal, + calibrated_scores=calibrated, + labels=labels, + raw_confidence=meta["confidence"].to_numpy(dtype=np.float64), + eval_df=meta, + ) + ) + + for ds_name in grouped: + grouped[ds_name].sort(key=lambda r: r.config_name) + + return dict(grouped) + + +def _run_plots_only(output_dir: Path, plot_format: str) -> None: + """Regenerate plots from ``{output_dir}/eval_results`` without inference.""" + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + grouped = load_eval_results_for_plotting(output_dir) + config_names = [r.config_name for results in grouped.values() for r in results] + _configure_ablation_colors(config_names) + + for ds_name in sorted(grouped): + ds_results = grouped[ds_name] + logger.info("Generating plots for %s (%d configs)...", ds_name, len(ds_results)) + _generate_plots_for_dataset(ds_results, ds_name, plots_dir, plot_format) + + logger.info("Plots saved to %s", plots_dir) + + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- +def build_summary_table(all_results: list[EvalResult]) -> pd.DataFrame: + """Aggregate all EvalResults into a single summary DataFrame.""" + rows = [] + for r in all_results: + rows.append( + { + "config": r.config_name, + "dataset": r.dataset_name, + "ECE": round(r.ece, 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.05]: round(r.tail_ece_at_5pct, 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.10]: round(r.tail_ece_at_10pct, 5), + "Brier": round(r.brier, 5), + "PR_AUC": round(r.pr_auc, 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.05]: round(r.fdr_bias_at_5pct, 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.10]: round(r.fdr_bias_at_10pct, 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.05]: round(r.q_dev_at_5pct, 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.10]: round(r.q_dev_at_10pct, 5), + "IDs@1%FDR": r.ids_at_1pct, + "IDs@5%FDR": r.ids_at_5pct, + "IDs@10%FDR": r.ids_at_10pct, + } + ) + return pd.DataFrame(rows) + + +_DEFAULT_OUTPUT_DIR = Path("analysis/hpo_ablation") + + +def _validate_training_inputs( + *, + skip_training: bool, + train_features: Path | None, + val_features: Path | None, + validation_fraction: float | None, +) -> None: + if skip_training: + return + if train_features is None: + raise typer.BadParameter( + "--train-features is required unless --skip-training is set." + ) + if val_features is None and validation_fraction is None: + raise typer.BadParameter( + "Provide --val-features or --validation-fraction when training." + ) + if val_features is not None and validation_fraction is not None: + logger.warning( + "Both --val-features and --validation-fraction set; using --val-features." + ) + + +def _configure_ablation_configs( + *, + skip_training: bool, + train_features: Path | None, + eval_dfs: dict[str, pl.DataFrame], + hyperparams_from_model: Path | None, +) -> None: + global ABLATION_CONFIGS + + if not skip_training: + assert train_features is not None + train_schema_df = _load_parquet_as_polars(train_features) + ABLATION_CONFIGS = build_ablation_configs( + train_schema_df, hyperparams_from_model + ) + else: + first_eval = next(iter(eval_dfs.values())) + ABLATION_CONFIGS = build_ablation_configs(first_eval, hyperparams_from_model) + _configure_ablation_colors(ABLATION_CONFIGS.keys()) + logger.info("Ablation configs: %s", list(ABLATION_CONFIGS.keys())) + + +def _train_or_load_ablation_models( + *, + skip_training: bool, + train_features: Path | None, + val_features: Path | None, + validation_fraction: float | None, + output_dir: Path, + seed: int, + hyperparams_from_model: Path | None, +) -> dict[str, ProbabilityCalibrator]: + if skip_training: + logger.info("Step 3: Loading pre-trained ablation models...") + return load_ablation_models(output_dir) + + logger.info("Step 3: Training ablation models...") + assert train_features is not None + full_train_df = _load_parquet_as_polars(train_features) + if val_features is not None: + train_df = full_train_df + val_df = _load_parquet_as_polars(val_features) + else: + assert validation_fraction is not None + train_df, val_df = split_train_val_frames( + full_train_df, validation_fraction, seed + ) + train_hp = None + if hyperparams_from_model is not None: + train_hp = train_hyperparams_from_model(hyperparams_from_model) + logger.info( + "Using training hyperparameters from %s: %s", + hyperparams_from_model, + train_hp, + ) + return train_ablation_models( + train_df, val_df, output_dir, seed, train_hyperparams=train_hp + ) + + +def _evaluate_ablations( + *, + eval_dfs: dict[str, pl.DataFrame], + models: dict[str, ProbabilityCalibrator], + plots_dir: Path, + plot_format: str, +) -> list[EvalResult]: + logger.info("Step 4: Evaluating ablation models...") + all_results: list[EvalResult] = [] + + for ds_name, ds_df in eval_dfs.items(): + ds_results: list[EvalResult] = [] + for config_name, columns in ABLATION_CONFIGS.items(): + result = evaluate_single( + config_name, models[config_name], columns, ds_df, ds_name + ) + ds_results.append(result) + all_results.append(result) + logger.info( + " %s / %s: ECE=%.4f, Brier=%.4f, IDs@1%%=%d, IDs@5%%=%d, IDs@10%%=%d", + ds_name, + config_name, + result.ece, + result.brier, + result.ids_at_1pct, + result.ids_at_5pct, + result.ids_at_10pct, + ) + + logger.info("Step 5: Generating plots for %s...", ds_name) + _generate_plots_for_dataset(ds_results, ds_name, plots_dir, plot_format) + + return all_results + + +def _write_ablation_summary(output_dir: Path, all_results: list[EvalResult]) -> None: + logger.info("Step 7: Writing summary...") + summary = build_summary_table(all_results) + summary.to_csv(output_dir / "ablation_summary.csv", index=False) + + summary_json = summary.to_dict(orient="records") + with open(output_dir / "ablation_summary.json", "w") as f: + json.dump(summary_json, f, indent=2) + + logger.info("Summary table:\n%s", summary.to_string(index=False)) + + +# --------------------------------------------------------------------------- +# Main CLI +# --------------------------------------------------------------------------- +@app.command() +def main( + train_features: Annotated[ + Optional[Path], + typer.Option( + help="Path to pre-computed training Parquet file or directory. " + "Required unless --skip-training is set.", + ), + ] = None, + val_features: Annotated[ + Optional[Path], + typer.Option( + help="Pre-computed validation Parquet. Omit if using --validation-fraction.", + ), + ] = None, + validation_fraction: Annotated[ + Optional[float], + typer.Option( + "--validation-fraction", + min=0.0, + max=1.0, + help=( + "Hold out this fraction of --train-features for validation " + "(same row split for every ablation model). Alternative to --val-features." + ), + ), + ] = None, + output_dir: Annotated[ + Path, + typer.Option(help="Directory for cached features, models, metrics, and plots."), + ] = _DEFAULT_OUTPUT_DIR, + astral_spectra: Annotated[ + Optional[str], + typer.Option(help="Optional: path to Astral spectra directory."), + ] = None, + astral_predictions: Annotated[ + Optional[str], + typer.Option(help="Optional: path to Astral predictions CSV."), + ] = None, + plot_format: Annotated[ + str, + typer.Option(help="Plot format: 'pdf', 'png', or 'both'."), + ] = "both", + seed: Annotated[ + int, + typer.Option(help="Random seed."), + ] = 42, + koina_url: Annotated[ + str, + typer.Option(help="Koina server URL for eval feature computation."), + ] = "koina.wilhelmlab.org:443", + koina_ssl: Annotated[ + bool, + typer.Option(help="Use SSL for Koina server."), + ] = True, + skip_feature_compute: Annotated[ + bool, + typer.Option( + "--skip-feature-compute", + help="Skip eval feature computation; assume cache exists.", + ), + ] = False, + skip_training: Annotated[ + bool, + typer.Option( + "--skip-training", + help="Load pre-trained ablation models from {output-dir}/models/ " + "instead of training from scratch.", + ), + ] = False, + hyperparams_from_model: Annotated[ + Optional[Path], + typer.Option( + help="Use training hyperparameters from this saved calibrator directory " + "(e.g. HPO best model). Reads config.json.", + ), + ] = None, + plots_only: Annotated[ + bool, + typer.Option( + "--plots-only", + help="Regenerate plots from {output-dir}/eval_results only " + "(no feature compute, training, or evaluation).", + ), + ] = False, +) -> None: + """Run feature ablation study for the Winnow calibrator.""" + output_dir.mkdir(parents=True, exist_ok=True) + + if plots_only: + _run_plots_only(output_dir, plot_format) + logger.info("Feature ablation plots complete.") + return + + _validate_training_inputs( + skip_training=skip_training, + train_features=train_features, + val_features=val_features, + validation_fraction=validation_fraction, + ) + + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + logger.info("Step 1: Computing eval features...") + eval_parquets = compute_all_eval_features( + output_dir, + koina_url, + koina_ssl, + astral_spectra, + astral_predictions, + skip_feature_compute, + reference_model_dir=hyperparams_from_model, + ) + + logger.info("Step 2: Loading Parquets...") + eval_dfs: dict[str, pl.DataFrame] = {} + for name, path in eval_parquets.items(): + eval_dfs[name] = _load_parquet_as_polars(path) + logger.info(" Loaded eval %s: %d rows", name, len(eval_dfs[name])) + + _configure_ablation_configs( + skip_training=skip_training, + train_features=train_features, + eval_dfs=eval_dfs, + hyperparams_from_model=hyperparams_from_model, + ) + models = _train_or_load_ablation_models( + skip_training=skip_training, + train_features=train_features, + val_features=val_features, + validation_fraction=validation_fraction, + output_dir=output_dir, + seed=seed, + hyperparams_from_model=hyperparams_from_model, + ) + all_results = _evaluate_ablations( + eval_dfs=eval_dfs, + models=models, + plots_dir=plots_dir, + plot_format=plot_format, + ) + + logger.info("Step 6: Saving eval results...") + save_eval_results(all_results, output_dir) + _write_ablation_summary(output_dir, all_results) + logger.info("Results saved to %s", output_dir) + logger.info("Feature ablation study complete.") + + +if __name__ == "__main__": + app() From 52c00289b01c2641967b3b9427e4e8c5d2153f2a Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:02:05 +0100 Subject: [PATCH 02/26] chore: pass pre-commit and fix code drift --- scripts/analyze_fdr_overlap.py | 384 ++++++++++-------- scripts/analyze_features.py | 8 +- scripts/analyze_novelty.py | 2 - scripts/analyze_upscored_fps.py | 30 +- scripts/benchmark_runtime.py | 8 +- scripts/benchmark_scaling.py | 4 +- scripts/evaluate_calibrator_generalisation.py | 27 +- scripts/plot_acfm_minus_lcfm_fdr.py | 113 ++++-- scripts/plot_analysis.py | 39 +- scripts/plot_fdr_method_comparison.py | 25 +- scripts/run_feature_ablations.py | 24 +- 11 files changed, 356 insertions(+), 308 deletions(-) diff --git a/scripts/analyze_fdr_overlap.py b/scripts/analyze_fdr_overlap.py index cbe44ef2..6e51fc2d 100644 --- a/scripts/analyze_fdr_overlap.py +++ b/scripts/analyze_fdr_overlap.py @@ -30,7 +30,7 @@ import re from collections import defaultdict from pathlib import Path -from typing import Annotated, cast +from typing import Annotated, Callable, cast import matplotlib.pyplot as plt import numpy as np @@ -448,74 +448,61 @@ def _load_from_folder(folder: Path) -> pd.DataFrame: return preds_df -def _discover_unlabelled_folders(root: Path) -> dict[str, Path]: - """Map project key -> full-search folder under ``root``. +def _register_project_folder(projects: dict[str, Path], key: str, folder: Path) -> None: + """Register *folder* under *key*, warning on duplicate keys.""" + if key in projects: + logger.warning( + "Duplicate project key %r: %s and %s", + key, + projects[key], + folder, + ) + return + projects[key] = folder - Supports flat project folders (``{root}/PXD004732/``) and nested per-run - layouts (``{root}/PXD004452//``) used by new eval sets. - """ + +def _discover_project_folders( + root: Path, + *, + is_preds_folder: Callable[[Path], bool], +) -> dict[str, Path]: + """Map project key -> preds folder under *root* (flat or ``PXD*//``).""" projects: dict[str, Path] = {} if not root.is_dir(): return projects - def _register(key: str, folder: Path) -> None: - if key in projects: - logger.warning( - "Duplicate unlabelled project key %r: %s and %s", - key, - projects[key], - folder, - ) - return - projects[key] = folder - for child in sorted(root.iterdir()): if not child.is_dir(): continue - if _is_unlabelled_preds_folder(child): - _register(_project_key_from_folder(child.name), child) + if is_preds_folder(child): + _register_project_folder( + projects, _project_key_from_folder(child.name), child + ) continue if not child.name.startswith(_PXD_ACCESSION_PREFIX): continue for run_dir in sorted(child.iterdir()): - if run_dir.is_dir() and _is_unlabelled_preds_folder(run_dir): - _register(run_dir.name, run_dir) + if run_dir.is_dir() and is_preds_folder(run_dir): + _register_project_folder(projects, run_dir.name, run_dir) return projects -def _discover_labelled_folders(root: Path) -> dict[str, Path]: - """Map project key -> database-reference folder under ``root``. +def _discover_unlabelled_folders(root: Path) -> dict[str, Path]: + """Map project key -> full-search folder under ``root``. Supports flat project folders (``{root}/PXD004732/``) and nested per-run layouts (``{root}/PXD004452//``) used by new eval sets. """ - projects: dict[str, Path] = {} - if not root.is_dir(): - return projects + return _discover_project_folders(root, is_preds_folder=_is_unlabelled_preds_folder) - def _register(key: str, folder: Path) -> None: - if key in projects: - logger.warning( - "Duplicate labelled project key %r: %s and %s", - key, - projects[key], - folder, - ) - return - projects[key] = folder - for child in sorted(root.iterdir()): - if not child.is_dir(): - continue - if _is_labelled_preds_folder(child): - _register(_project_key_from_folder(child.name), child) - continue - if not child.name.startswith(_PXD_ACCESSION_PREFIX): - continue - for run_dir in sorted(child.iterdir()): - if run_dir.is_dir() and _is_labelled_preds_folder(run_dir): - _register(run_dir.name, run_dir) - return projects +def _discover_labelled_folders(root: Path) -> dict[str, Path]: + """Map project key -> database-reference folder under ``root``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts (``{root}/PXD004452//``) used by new eval sets. + """ + return _discover_project_folders(root, is_preds_folder=_is_labelled_preds_folder) def discover_project_pairs( @@ -584,7 +571,6 @@ def _add_q_values( def _add_database_grounded_q_values( df: pd.DataFrame, - residue_masses: dict[str, float], confidence_col: str = RAW_CONFIDENCE_COL, *, q_col: str = DB_Q_VALUE_COL, @@ -606,12 +592,18 @@ def _add_database_grounded_q_values( columns=[q_col, "psm_q_value", "psm_fdr", "fdr"], errors="ignore", ).copy() + drop = _effective_db_grounded_drop(len(work)) ctrl = DatabaseGroundedFDRControl( confidence_feature=confidence_col, - residue_masses=residue_masses, - drop=_effective_db_grounded_drop(len(work)), + drop=drop, ) - ctrl.fit(dataset=work.copy(), correct_column=correct_col) + # DataFrame path: mirror DatabaseGroundedFDRControl.fit without CalibrationDataset. + sorted_df = work.sort_values(confidence_col, ascending=False) + labels = sorted_df[correct_col].astype(float).to_numpy() + conf = sorted_df[confidence_col].to_numpy() + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + ctrl._fdr_values = np.array(1.0 - precision)[drop:] + ctrl._confidence_scores = conf[drop:] q_df = ctrl.add_psm_q_value( work[[confidence_col]].copy(), confidence_col=confidence_col ) @@ -711,7 +703,6 @@ def compute_overlap_table( df = _add_q_values(df.copy()) db_scored = _add_database_grounded_q_values( db_df.copy(), - residue_masses, confidence_col=RAW_CONFIDENCE_COL, q_col=DB_Q_VALUE_COL, ) @@ -739,7 +730,7 @@ def compute_overlap_table( is_match = retained["pred_key"] == retained["seq_key"] else: is_match = retained["prediction"].map( - lambda p: _is_db_match(str(p), db_keys_at_fdr) + lambda p, keys=db_keys_at_fdr: _is_db_match(str(p), keys) ) n_matching = int(is_match.sum()) @@ -789,6 +780,46 @@ def compute_overlap_table( # --------------------------------------------------------------------------- # Plots # --------------------------------------------------------------------------- +def _draw_venn_panel( + ax, + *, + winnow_peptides: set[str], + db_peptides: set[str], + winnow_label: str, + fdr_t: float, +) -> None: + """Draw one FDR-threshold Venn panel onto *ax*.""" + pct = int(fdr_t * 100) + if not winnow_peptides and not db_peptides: + ax.set_title(f"No peptides retained at {pct}% FDR") + ax.axis("off") + return + + if not winnow_peptides or not db_peptides: + missing = "Winnow" if not winnow_peptides else "Database search" + ax.text( + 0.5, + 0.5, + f"No {missing} peptides at {pct}% FDR", + ha="center", + va="center", + transform=ax.transAxes, + ) + ax.set_title(f"{pct}% FDR") + ax.axis("off") + return + + venn2( + [db_peptides, winnow_peptides], + set_labels=("Database search", winnow_label), + set_colors=(_CORRECT_COLOUR, _INCORRECT_COLOUR), + alpha=0.6, + ax=ax, + ) + ax.set_title(f"Unique peptides at {pct}% FDR") + _style_ax(ax) + + def _plot_venn_panels( winnow_df: pd.DataFrame, project: str, @@ -800,17 +831,16 @@ def _plot_venn_panels( db_df: pd.DataFrame | None = None, db_peptides_static: set[str] | None = None, ) -> None: + del residue_masses # kept for call-site compatibility + del project if db_df is None and db_peptides_static is None: raise ValueError("Provide db_df or db_peptides_static for Venn panels") winnow_scored = _add_q_values(winnow_df.copy()) db_scored: pd.DataFrame | None = None if db_df is not None: - if residue_masses is None: - raise ValueError("residue_masses required when db_df is provided") db_scored = _add_database_grounded_q_values( db_df.copy(), - residue_masses, confidence_col=RAW_CONFIDENCE_COL, q_col=DB_Q_VALUE_COL, ) @@ -832,35 +862,13 @@ def _plot_venn_panels( assert db_peptides_static is not None db_peptides = db_peptides_static - pct = int(fdr_t * 100) - if not winnow_peptides and not db_peptides: - ax.set_title(f"No peptides retained at {pct}% FDR") - ax.axis("off") - continue - - if not winnow_peptides or not db_peptides: - missing = "Winnow" if not winnow_peptides else "Database search" - ax.text( - 0.5, - 0.5, - f"No {missing} peptides at {pct}% FDR", - ha="center", - va="center", - transform=ax.transAxes, - ) - ax.set_title(f"{pct}% FDR") - ax.axis("off") - continue - - venn2( - [db_peptides, winnow_peptides], - set_labels=("Database search", winnow_label), - set_colors=(_CORRECT_COLOUR, _INCORRECT_COLOUR), - alpha=0.6, - ax=ax, + _draw_venn_panel( + ax, + winnow_peptides=winnow_peptides, + db_peptides=db_peptides, + winnow_label=winnow_label, + fdr_t=fdr_t, ) - ax.set_title(f"Unique peptides at {pct}% FDR") - _style_ax(ax) fig.suptitle(suptitle, fontsize=12) fig.tight_layout() @@ -962,6 +970,109 @@ def _subsample_violin_groups(df: pd.DataFrame, category_col: str) -> pd.DataFram return pd.concat(parts, ignore_index=True) if parts else df +def _plot_novel_violins_at_fdr( + df: pd.DataFrame, + db_scored: pd.DataFrame | None, + discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, + available: list[tuple[str, str]], + *, + project: str, + display: str, + eval_type: str, + plots_dir: Path, + fdr_t: float, + labelled_subset: bool, +) -> None: + """Render one FDR-threshold novel-feature violin figure, or skip if too sparse.""" + retained = df[df["psm_q_value"] <= fdr_t].copy() + if len(retained) < _MIN_VIOLIN_GROUP_SIZE: + logger.info( + "%s: skip violins at %d%% FDR (n=%d retained)", + project, + int(fdr_t * 100), + len(retained), + ) + return + + match_keys: set[str] | None = None + if not labelled_subset: + assert db_scored is not None + match_keys = _unique_peptides_at_fdr( + db_scored, "sequence", DB_Q_VALUE_COL, fdr_t + ) + + groups = _assign_retained_groups( + retained, + match_keys, + discordance_cache, + labelled_subset=labelled_subset, + ) + retained = retained.assign(_overlap_group=groups) + plot_df = retained[ + retained["_overlap_group"].isin(["Database match", "fully_discordant"]) + ].copy() + plot_df["Category"] = plot_df["_overlap_group"].map( + { + "Database match": "Database match", + "fully_discordant": "Novel", + } + ) + plot_df = _subsample_violin_groups(plot_df, "Category") + + n_match = (plot_df["Category"] == "Database match").sum() + n_novel = (plot_df["Category"] == "Novel").sum() + if n_match < _MIN_VIOLIN_GROUP_SIZE or n_novel < _MIN_VIOLIN_GROUP_SIZE: + logger.info( + "%s: skip violins at %d%% FDR (match=%d, novel=%d)", + project, + int(fdr_t * 100), + n_match, + n_novel, + ) + return + + n_feats = len(available) + n_cols = 4 + n_rows = int(np.ceil(n_feats / n_cols)) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows)) + axes_flat = np.atleast_1d(axes).flatten() + palette = {"Database match": _CORRECT_COLOUR, "Novel": _NOVEL_COLOUR} + cat_order = ["Database match", "Novel"] + + for ax, (col, label) in zip(axes_flat, available): + sub = plot_df[[col, "Category"]].dropna() + if sub["Category"].nunique() < 2: + ax.set_visible(False) + continue + sns.violinplot( + data=sub, + x="Category", + y=col, + order=cat_order, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + ax.set_ylabel(label) + ax.tick_params(axis="x", rotation=15) + _style_ax(ax) + + for ax in axes_flat[len(available) :]: + ax.set_visible(False) + + pct = int(fdr_t * 100) + fig.suptitle( + f"{display} ({_eval_type_display(eval_type)}): " + f"database-matched vs novel features at {pct}% FDR", + fontsize=12, + ) + fig.tight_layout() + _save_fig(fig, plots_dir / f"novel_feature_violins_{project}_fdr{pct}") + + def plot_novel_feature_violins( df: pd.DataFrame, db_df: pd.DataFrame | None, @@ -979,6 +1090,7 @@ def plot_novel_feature_violins( full-search comparisons, database match uses peptides retained at the same nominal FDR via database-grounded FDR on raw confidence (see overlap table). """ + del residue_masses # kept for call-site compatibility available = [ (col, label) for col, label in NOVEL_FEATURE_COLUMNS if col in df.columns ] @@ -994,107 +1106,27 @@ def plot_novel_feature_violins( db_scored: pd.DataFrame | None = None if not labelled_subset: - if db_df is None or residue_masses is None: - raise ValueError( - "Full-search violin plots require db_df and residue_masses" - ) + if db_df is None: + raise ValueError("Full-search violin plots require db_df") db_scored = _add_database_grounded_q_values( db_df.copy(), - residue_masses, confidence_col=RAW_CONFIDENCE_COL, q_col=DB_Q_VALUE_COL, ) for fdr_t in FDR_THRESHOLDS: - retained = df[df["psm_q_value"] <= fdr_t].copy() - if len(retained) < _MIN_VIOLIN_GROUP_SIZE: - logger.info( - "%s: skip violins at %d%% FDR (n=%d retained)", - project, - int(fdr_t * 100), - len(retained), - ) - continue - - if labelled_subset: - match_keys: set[str] | None = None - else: - assert db_scored is not None - match_keys = _unique_peptides_at_fdr( - db_scored, "sequence", DB_Q_VALUE_COL, fdr_t - ) - - groups = _assign_retained_groups( - retained, - match_keys, + _plot_novel_violins_at_fdr( + df, + db_scored, discordance_cache, + available, + project=project, + display=display, + eval_type=eval_type, + plots_dir=plots_dir, + fdr_t=fdr_t, labelled_subset=labelled_subset, ) - retained = retained.assign(_overlap_group=groups) - plot_df = retained[ - retained["_overlap_group"].isin(["Database match", "fully_discordant"]) - ].copy() - plot_df["Category"] = plot_df["_overlap_group"].map( - { - "Database match": "Database match", - "fully_discordant": "Novel", - } - ) - plot_df = _subsample_violin_groups(plot_df, "Category") - - n_match = (plot_df["Category"] == "Database match").sum() - n_novel = (plot_df["Category"] == "Novel").sum() - if n_match < _MIN_VIOLIN_GROUP_SIZE or n_novel < _MIN_VIOLIN_GROUP_SIZE: - logger.info( - "%s: skip violins at %d%% FDR (match=%d, novel=%d)", - project, - int(fdr_t * 100), - n_match, - n_novel, - ) - continue - - n_feats = len(available) - n_cols = 4 - n_rows = int(np.ceil(n_feats / n_cols)) - fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows)) - axes_flat = np.atleast_1d(axes).flatten() - - palette = {"Database match": _CORRECT_COLOUR, "Novel": _NOVEL_COLOUR} - cat_order = ["Database match", "Novel"] - - for ax, (col, label) in zip(axes_flat, available): - sub = plot_df[[col, "Category"]].dropna() - if sub["Category"].nunique() < 2: - ax.set_visible(False) - continue - sns.violinplot( - data=sub, - x="Category", - y=col, - order=cat_order, - palette=palette, - ax=ax, - inner="quartile", - cut=0, - linewidth=0.8, - ) - ax.set_xlabel("") - ax.set_ylabel(label) - ax.tick_params(axis="x", rotation=15) - _style_ax(ax) - - for ax in axes_flat[len(available) :]: - ax.set_visible(False) - - pct = int(fdr_t * 100) - fig.suptitle( - f"{display} ({_eval_type_display(eval_type)}): " - f"database-matched vs novel features at {pct}% FDR", - fontsize=12, - ) - fig.tight_layout() - _save_fig(fig, plots_dir / f"novel_feature_violins_{project}_fdr{pct}") # --------------------------------------------------------------------------- diff --git a/scripts/analyze_features.py b/scripts/analyze_features.py index e8a6ec51..6fc4fec7 100644 --- a/scripts/analyze_features.py +++ b/scripts/analyze_features.py @@ -619,9 +619,11 @@ def main( calibrator = ProbabilityCalibrator.load(model_path) if koina_url is not None or not koina_ssl: - calibrator.apply_koina_server_overrides( - server_url=koina_url, - ssl=koina_ssl, + logger.warning( + "Ignoring --koina-url/--koina-ssl; Koina server overrides are no longer " + "supported on ProbabilityCalibrator (url=%s, ssl=%s).", + koina_url, + koina_ssl, ) koina_constants = _parse_koina_constants(koina_input_constant) diff --git a/scripts/analyze_novelty.py b/scripts/analyze_novelty.py index 8b87cdc4..46efbd68 100644 --- a/scripts/analyze_novelty.py +++ b/scripts/analyze_novelty.py @@ -523,7 +523,6 @@ def _plot_nontryptic_score_by_terminus( linewidth=0.8, ) ax.set_xlabel("") - pct = int(fdr_t * 100) ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) ax.grid(False) _spine_fmt(ax) @@ -925,7 +924,6 @@ def _plot_nontryptic_delta_by_terminus( ) ax.axhline(0.0, ls="--", color=_IDEAL_LINE_COLOUR, lw=1) ax.set_xlabel("") - pct = int(fdr_t * 100) ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) ax.grid(False) _spine_fmt(ax) diff --git a/scripts/analyze_upscored_fps.py b/scripts/analyze_upscored_fps.py index eaeef641..184ae2ae 100644 --- a/scripts/analyze_upscored_fps.py +++ b/scripts/analyze_upscored_fps.py @@ -229,6 +229,19 @@ def _is_labelled_preds_folder(folder: Path) -> bool: return required.issubset(header) +def _register_labelled_folder(results: dict[str, Path], key: str, folder: Path) -> None: + """Register *folder* under *key*, warning on duplicate keys.""" + if key in results: + logger.warning( + "Duplicate labelled project key %r: %s and %s", + key, + results[key], + folder, + ) + return + results[key] = folder + + def _discover_labelled_folders(root: Path) -> dict[str, Path]: """Find folders with labelled ``preds_and_fdr_metrics.csv``. @@ -239,28 +252,19 @@ def _discover_labelled_folders(root: Path) -> dict[str, Path]: if not root.is_dir(): return results - def _register(key: str, folder: Path) -> None: - if key in results: - logger.warning( - "Duplicate labelled project key %r: %s and %s", - key, - results[key], - folder, - ) - return - results[key] = folder - for child in sorted(root.iterdir()): if not child.is_dir(): continue if _is_labelled_preds_folder(child): - _register(_project_key_from_folder(child.name), child) + _register_labelled_folder( + results, _project_key_from_folder(child.name), child + ) continue if not child.name.startswith(_PXD_ACCESSION_PREFIX): continue for run_dir in sorted(child.iterdir()): if run_dir.is_dir() and _is_labelled_preds_folder(run_dir): - _register(run_dir.name, run_dir) + _register_labelled_folder(results, run_dir.name, run_dir) return results diff --git a/scripts/benchmark_runtime.py b/scripts/benchmark_runtime.py index b70998cf..a287959f 100644 --- a/scripts/benchmark_runtime.py +++ b/scripts/benchmark_runtime.py @@ -168,9 +168,9 @@ def load_dataset( predictions_path=predictions_path, ) - from winnow.scripts.main import filter_dataset + from winnow.scripts.main import _filter_dataset - dataset = filter_dataset(dataset) + dataset = _filter_dataset(dataset) return dataset @@ -316,10 +316,6 @@ def run_benchmark( # Load calibrator calibrator = ProbabilityCalibrator.load(pretrained_model_name_or_path=model_path) - # Apply Koina server overrides if provided - if koina_url is not None or koina_ssl is not None: - calibrator.apply_koina_server_overrides(server_url=koina_url, ssl=koina_ssl) - # Remove Prosit features if requested if not include_prosit: to_remove = [ diff --git a/scripts/benchmark_scaling.py b/scripts/benchmark_scaling.py index 4fbf4596..63d07852 100644 --- a/scripts/benchmark_scaling.py +++ b/scripts/benchmark_scaling.py @@ -131,9 +131,9 @@ def load_dataset_timed( data_path=spectrum_path, predictions_path=predictions_path, ) - from winnow.scripts.main import filter_dataset + from winnow.scripts.main import _filter_dataset - dataset = filter_dataset(dataset) + dataset = _filter_dataset(dataset) return dataset, m["wall_time_s"] diff --git a/scripts/evaluate_calibrator_generalisation.py b/scripts/evaluate_calibrator_generalisation.py index 8947b51e..e443c530 100644 --- a/scripts/evaluate_calibrator_generalisation.py +++ b/scripts/evaluate_calibrator_generalisation.py @@ -100,10 +100,13 @@ def initialise_calibrator( train_project: Optional[str] = None, ) -> ProbabilityCalibrator: """Create a fresh calibrator matching train-extra-small-mass-error-da.""" - koina_kwargs: Dict = {} - if koina_server_url is not None: - koina_kwargs["koina_server_url"] = koina_server_url - koina_kwargs["koina_ssl"] = koina_ssl + if koina_server_url is not None or not koina_ssl: + logger.warning( + "Ignoring koina_server_url/koina_ssl; Koina server overrides are no " + "longer supported (url=%s, ssl=%s).", + koina_server_url, + koina_ssl, + ) irt_train_fraction = _IRT_TRAIN_FRACTION_OVERRIDES.get(train_project or "", 0.1) @@ -123,15 +126,14 @@ def initialise_calibrator( calibrator.add_feature(MassErrorDaFeature(residue_masses=RESIDUE_MASSES)) calibrator.add_feature( FragmentMatchFeatures( - mz_tolerance_ppm=20, + mz_tolerance=20, + mz_tolerance_unit="ppm", learn_from_missing=False, intensity_model_name=_INTENSITY_MODEL, max_precursor_charge=_MAX_PRECURSOR_CHARGE, max_peptide_length=_MAX_PEPTIDE_LENGTH, unsupported_residues=_UNSUPPORTED_RESIDUES, model_input_constants=_KOINA_INPUT_CONSTANTS, - excluded_columns=_EXTRA_SMALL_FRAGMENT_EXCLUDE, - **koina_kwargs, ) ) calibrator.add_feature( @@ -142,11 +144,18 @@ def initialise_calibrator( irt_model_name=_IRT_MODEL, max_peptide_length=_MAX_PEPTIDE_LENGTH, unsupported_residues=_UNSUPPORTED_RESIDUES, - **koina_kwargs, ) ) - calibrator.add_feature(BeamFeatures(excluded_columns=_EXTRA_SMALL_BEAM_EXCLUDE)) + calibrator.add_feature(BeamFeatures()) calibrator.add_feature(TokenScoreFeatures()) + # Former excluded_columns behaviour: train on a reduced feature subset. + training_columns = [ + col + for col in calibrator.columns + if col not in _EXTRA_SMALL_FRAGMENT_EXCLUDE + and col not in _EXTRA_SMALL_BEAM_EXCLUDE + ] + calibrator.set_training_feature_columns(training_columns) return calibrator diff --git a/scripts/plot_acfm_minus_lcfm_fdr.py b/scripts/plot_acfm_minus_lcfm_fdr.py index 8381eacb..9dad0ebd 100644 --- a/scripts/plot_acfm_minus_lcfm_fdr.py +++ b/scripts/plot_acfm_minus_lcfm_fdr.py @@ -47,6 +47,19 @@ def _safe_basename(project: str) -> str: _PXD_RUN_PARENTS: tuple[str, ...] = ("PXD004452", "PXD006939", "PXD013868") +def _add_preds_candidate( + candidates: list[Path], + seen: set[Path], + root: Path, + *relative: str, + fname: str, +) -> None: + path = root.joinpath(*relative, fname) + if path not in seen: + seen.add(path) + candidates.append(path) + + def _preds_csv_candidates(root: Path, project: str, *, role: str) -> list[Path]: """Paths to try for ``preds_and_fdr_metrics.csv`` under *root*. @@ -62,34 +75,42 @@ def _preds_csv_candidates(root: Path, project: str, *, role: str) -> list[Path]: candidates: list[Path] = [] def add(*relative: str) -> None: - path = root.joinpath(*relative, fname) - if path not in seen: - seen.add(path) - candidates.append(path) + _add_preds_candidate(candidates, seen, root, *relative, fname=fname) add(project) base = project.split("/")[-1] add(f"{base}{role_suffix}") if "/" in project: add(*project.split("/")) + return candidates # S3 and download-new-eval-results: {root}/PXD*/{run}/; legacy flat {root}/{run}/. - if "/" not in project: - pxd_seen: set[str] = set() - for pxd in _PXD_RUN_PARENTS: - add(pxd, project) - add(pxd, f"{project}{role_suffix}") - pxd_seen.add(pxd) - if root.is_dir(): - for child in sorted(root.iterdir()): - if child.is_dir() and child.name.startswith("PXD"): - if child.name not in pxd_seen: - add(child.name, project) - add(child.name, f"{project}{role_suffix}") + pxd_seen: set[str] = set() + for pxd in _PXD_RUN_PARENTS: + add(pxd, project) + add(pxd, f"{project}{role_suffix}") + pxd_seen.add(pxd) + if root.is_dir(): + for child in sorted(root.iterdir()): + if child.is_dir() and child.name.startswith("PXD"): + if child.name not in pxd_seen: + add(child.name, project) + add(child.name, f"{project}{role_suffix}") return candidates +def _collect_alt_roots(root: Path, alt_roots: list[Path] | None) -> list[Path]: + """Build ordered list of roots to search, deduplicating by resolve().""" + roots_to_try: list[Path] = [root] + if not alt_roots: + return roots_to_try + for alt in alt_roots: + if alt.resolve() != root.resolve() and alt not in roots_to_try: + roots_to_try.append(alt) + return roots_to_try + + def _resolve_preds_csv( root: Path, project: str, @@ -98,25 +119,20 @@ def _resolve_preds_csv( alt_roots: list[Path] | None = None, ) -> Path: """Resolve ``preds_and_fdr_metrics.csv`` for lcfm (labelled) or acfm (unlabelled).""" - roots_to_try: list[Path] = [root] - if alt_roots: - for alt in alt_roots: - if alt.resolve() != root.resolve() and alt not in roots_to_try: - roots_to_try.append(alt) - tried: list[Path] = [] - for base in roots_to_try: + for base in _collect_alt_roots(root, alt_roots): for path in _preds_csv_candidates(base, project, role=role): tried.append(path) - if path.is_file(): - if base.resolve() != root.resolve(): - logger.info( - "Using %s predictions at %s (not under %s)", - role, - path, - root, - ) - return path + if not path.is_file(): + continue + if base.resolve() != root.resolve(): + logger.info( + "Using %s predictions at %s (not under %s)", + role, + path, + root, + ) + return path hint = "" if root.is_dir(): @@ -145,6 +161,27 @@ def _infer_predictions_root( ) +def _resolve_explicit_tree_root( + predictions_root: Path, + explicit: Path, + *, + sub: str, +) -> tuple[Path, list[Path]]: + """Resolve an explicit tree path that may be relative to *predictions_root*.""" + alt_roots: list[Path] = [] + if len(explicit.parts) == 1: + under_predictions = predictions_root / explicit + if under_predictions.is_dir(): + root = under_predictions + if explicit.is_dir() and explicit.resolve() != root.resolve(): + alt_roots.append(explicit) + nested = predictions_root / sub + if nested.is_dir() and nested.resolve() != root.resolve(): + alt_roots.append(nested) + return root, alt_roots + return explicit, alt_roots + + def _resolve_tree_root( predictions_root: Path | None, explicit: Path | None, @@ -157,16 +194,8 @@ def _resolve_tree_root( alt_roots: list[Path] = [] if explicit is not None: - if predictions_root is not None and len(explicit.parts) == 1: - under_predictions = predictions_root / explicit - if under_predictions.is_dir(): - root = under_predictions - if explicit.is_dir() and explicit.resolve() != root.resolve(): - alt_roots.append(explicit) - nested = predictions_root / sub - if nested.is_dir() and nested.resolve() != root.resolve(): - alt_roots.append(nested) - return root, alt_roots + if predictions_root is not None: + return _resolve_explicit_tree_root(predictions_root, explicit, sub=sub) return explicit, alt_roots if predictions_root is None: diff --git a/scripts/plot_analysis.py b/scripts/plot_analysis.py index 40748198..ee4f0e78 100644 --- a/scripts/plot_analysis.py +++ b/scripts/plot_analysis.py @@ -252,26 +252,20 @@ def plot_confidence_histogram( def _fit_db_fdr( df: pl.DataFrame, correct_col: str, - residue_masses: dict, confidence_feature: str = "calibrated_confidence", drop: int = 10, - use_proteome_shortcut: bool = False, ) -> DatabaseGroundedFDRControl: - """Fit a DatabaseGroundedFDRControl, using a proteome shortcut if labels lack sequences.""" + """Fit a DatabaseGroundedFDRControl from per-row correctness labels.""" ctrl = DatabaseGroundedFDRControl( confidence_feature=confidence_feature, - residue_masses=residue_masses, drop=drop, ) - if use_proteome_shortcut: - sorted_df = df.sort(confidence_feature, descending=True) - correct_vals = sorted_df[correct_col].to_numpy().astype(float) - confidence_vals = sorted_df[confidence_feature].to_numpy() - precision = np.cumsum(correct_vals) / np.arange(1, len(sorted_df) + 1) - ctrl._fdr_values = np.array(1 - precision[drop:]) - ctrl._confidence_scores = confidence_vals[drop:] - else: - ctrl.fit(dataset=df.to_pandas(), correct_column=correct_col) + sorted_df = df.sort(confidence_feature, descending=True) + correct_vals = sorted_df[correct_col].to_numpy().astype(float) + confidence_vals = sorted_df[confidence_feature].to_numpy() + precision = np.cumsum(correct_vals) / np.arange(1, len(sorted_df) + 1) + ctrl._fdr_values = np.array(1 - precision[drop:]) + ctrl._confidence_scores = confidence_vals[drop:] return ctrl @@ -284,13 +278,12 @@ def plot_fdr_accuracy( use_proteome_shortcut: bool = False, ) -> plt.Figure: """Compare non-parametric vs database-grounded FDR or q-value vs confidence.""" + del residue_masses, use_proteome_shortcut # retained for call-site compatibility fig, ax = plt.subplots(figsize=(8, 6)) col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" winnow_col = col_name - ctrl = _fit_db_fdr( - df, correct_col, residue_masses, use_proteome_shortcut=use_proteome_shortcut - ) + ctrl = _fit_db_fdr(df, correct_col) if metric == "fdr": db_pd = ctrl.add_psm_fdr(df.to_pandas(), "calibrated_confidence") @@ -340,7 +333,7 @@ def plot_ranked_qvalue( title: str, ) -> plt.Figure: """Ranked predictions vs q-value (non-parametric & database-grounded).""" - ctrl = _fit_db_fdr(df, correct_col, residue_masses) + ctrl = _fit_db_fdr(df, correct_col) test_pd = df.to_pandas() test_pd_no_q = test_pd.drop(columns=["psm_q_value"], errors="ignore") db_q = ctrl.add_psm_q_value(test_pd_no_q, "calibrated_confidence") @@ -421,12 +414,10 @@ def plot_ranked_fdr_raw_vs_cal( ) db_cal_ctrl = _fit_db_fdr( - df, correct_col, residue_masses, confidence_feature="calibrated_confidence" + df, correct_col, confidence_feature="calibrated_confidence" ) raw_df = df_raw if df_raw is not None else df - db_raw_ctrl = _fit_db_fdr( - raw_df, correct_col, residue_masses, confidence_feature="confidence" - ) + db_raw_ctrl = _fit_db_fdr(raw_df, correct_col, confidence_feature="confidence") col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" add_fn = "add_psm_fdr" if metric == "fdr" else "add_psm_q_value" @@ -489,12 +480,10 @@ def plot_bar_psms_fdr( test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") db_cal_ctrl = _fit_db_fdr( - df, correct_col, residue_masses, confidence_feature="calibrated_confidence" + df, correct_col, confidence_feature="calibrated_confidence" ) raw_df = df_raw if df_raw is not None else df - db_raw_ctrl = _fit_db_fdr( - raw_df, correct_col, residue_masses, confidence_feature="confidence" - ) + db_raw_ctrl = _fit_db_fdr(raw_df, correct_col, confidence_feature="confidence") db_cal = db_cal_ctrl.add_psm_q_value(test_pd_no_q.copy(), "calibrated_confidence") raw_pd_no_q = raw_df.to_pandas().drop( diff --git a/scripts/plot_fdr_method_comparison.py b/scripts/plot_fdr_method_comparison.py index bd83385b..d4444483 100644 --- a/scripts/plot_fdr_method_comparison.py +++ b/scripts/plot_fdr_method_comparison.py @@ -197,30 +197,20 @@ def _fit_database_grounded_fdr( df: pd.DataFrame, correct_col: str, confidence_col: str, - residue_masses: dict[str, float], *, drop: int = _DB_GROUNDED_DROP, ) -> DatabaseGroundedFDRControl: - """Fit ``DatabaseGroundedFDRControl`` (proteome shortcut or labelled sequence fit).""" + """Fit ``DatabaseGroundedFDRControl`` from per-row correctness labels.""" ctrl = DatabaseGroundedFDRControl( confidence_feature=confidence_col, - residue_masses=residue_masses, drop=drop, ) - if correct_col == "proteome_hit": - sorted_df = df.sort_values(confidence_col, ascending=False) - labels = sorted_df[correct_col].astype(float).to_numpy() - conf = sorted_df[confidence_col].to_numpy() - precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) - ctrl._fdr_values = np.array(1.0 - precision)[drop:] - ctrl._confidence_scores = conf[drop:] - else: - fit_df = df.copy() - if "sequence" not in fit_df.columns or "prediction" not in fit_df.columns: - raise ValueError( - "Labelled database-grounded FDR requires 'sequence' and 'prediction' columns" - ) - ctrl.fit(dataset=fit_df, correct_column=correct_col) + sorted_df = df.sort_values(confidence_col, ascending=False) + labels = sorted_df[correct_col].astype(float).to_numpy() + conf = sorted_df[confidence_col].to_numpy() + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + ctrl._fdr_values = np.array(1.0 - precision)[drop:] + ctrl._confidence_scores = conf[drop:] return ctrl @@ -332,7 +322,6 @@ def _add_database_grounded_qvalues( reference, correct_col, confidence_col, - residue_masses, drop=_effective_db_grounded_drop(len(reference), drop), ) return _assign_q_values_fast(work, confidence_col, ctrl, out_col) diff --git a/scripts/run_feature_ablations.py b/scripts/run_feature_ablations.py index 9da59237..17211d9b 100644 --- a/scripts/run_feature_ablations.py +++ b/scripts/run_feature_ablations.py @@ -127,8 +127,6 @@ "~calibrator.features.mass_error", "+calibrator.features.mass_error_da._target_=winnow.calibration.calibration_features.MassErrorDaFeature", "+calibrator.features.mass_error_da.residue_masses=${residue_masses}", - "+calibrator.features.fragment_match_features.excluded_columns=[spectral_angle,xcorr,complementary_ion_count,max_ion_gap]", - "+calibrator.features.beam_features.excluded_columns=[edit_distance]", ] @@ -315,10 +313,7 @@ def _feature_compute_overrides(reference_model_dir: Path | None) -> list[str]: ref_cols = _reference_model_columns(reference_model_dir) if ref_cols and MASS_ERROR_DA in ref_cols: return list(REDUCED_FEATURE_COMPUTE_OVERRIDES) - return [ - "+calibrator.features.fragment_match_features.excluded_columns=[spectral_angle,xcorr,complementary_ion_count,max_ion_gap]", - "+calibrator.features.beam_features.excluded_columns=[edit_distance]", - ] + return [] def _compute_eval_features_for_dataset( @@ -402,7 +397,6 @@ def _compute_eval_features_for_dataset( data_loader, calibrator, labelled=True, - filter_empty=True, ) combined_metadata = pd.concat(all_metadata, ignore_index=True) @@ -562,7 +556,8 @@ def _column_slice_to_feature_dataset( ) features = df.select(columns).to_numpy().astype(np.float32) labels = df["correct"].to_numpy().astype(np.float32) - return FeatureDataset(features=features, labels=labels) + non_confidence = [c for c in columns if c != "confidence"] + return FeatureDataset(features=features, labels=labels, columns=non_confidence) def _config_dir_name(config_name: str) -> str: @@ -993,8 +988,6 @@ def plot_fdr_vs_confidence( n_configs = len(results) fig, axes = plt.subplots(1, n_configs, figsize=(5 * n_configs, 4), squeeze=False) - residue_masses = _get_residue_masses() - for i, r in enumerate(results): ax = axes[0, i] @@ -1011,10 +1004,17 @@ def plot_fdr_vs_confidence( if has_sequence: dbg_fdr = DatabaseGroundedFDRControl( confidence_feature="calibrated_confidence", - residue_masses=residue_masses, ) try: - dbg_fdr.fit(dataset=r.eval_df.copy()) + sorted_df = r.eval_df.sort_values( + "calibrated_confidence", ascending=False + ) + labels = sorted_df["correct"].astype(float).to_numpy() + conf = sorted_df["calibrated_confidence"].to_numpy() + drop = 10 + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + dbg_fdr._fdr_values = np.array(1.0 - precision)[drop:] + dbg_fdr._confidence_scores = conf[drop:] dbg_metrics = dbg_fdr.add_psm_fdr( r.eval_df.copy(), confidence_col="calibrated_confidence" ) From 414f9699564fca88c2b386646abf1a21aa7fbf58 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:17:58 +0100 Subject: [PATCH 03/26] chore: remove local path hardcoding and detail external repos for FDR tool benchmarking --- scripts/plot_fdr_method_comparison.py | 34 ++++++++++--- .../run_external_peptide_holdout_benchmark.py | 50 +++++++++++++++---- 2 files changed, 66 insertions(+), 18 deletions(-) diff --git a/scripts/plot_fdr_method_comparison.py b/scripts/plot_fdr_method_comparison.py index d4444483..02b76baa 100644 --- a/scripts/plot_fdr_method_comparison.py +++ b/scripts/plot_fdr_method_comparison.py @@ -13,6 +13,16 @@ A long-form ``fdr_method_comparison_curves.csv`` (per spectrum x method) is written so plots and summary tables can be regenerated with ``--summarise-only``. + +External NovoBoard inputs (``--novoboard-root``) must follow +``{root}/{dataset}/novoboard/`` with target/decoy CSVs such as +``annotated_test.csv``, ``annotated_test_decoy_{rate}.csv``, +``raw_unlabelled.csv`` and ``raw_unlabelled_decoy_{rate}.csv``. Point +``--novoboard-root`` at the ``datasets`` directory of a NovoBoard checkout. +Local results were produced from the fork +``git@github.com:JemmaLDaniel/NovoBoard.git``, branch +``feat/adapt-to-instanovo`` at commit +``a9faab3ef1af06987599c2f01e6ba96072c80172``. """ from __future__ import annotations @@ -86,7 +96,6 @@ FDR_THRESHOLDS = [0.01, 0.05, 0.10] _DB_GROUNDED_DROP = 10 -DEFAULT_NOVOBOARD_ROOT = Path("/home/j-daniel/repos/NovoBoard/datasets") DEFAULT_WINNOW_RESULTS = _REPO_ROOT / "results" DEFAULT_MODEL_ROOT = _REPO_ROOT / "models" DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/fdr_method_comparison_psm" @@ -139,7 +148,8 @@ class DatasetConfig: def build_dataset_configs( winnow_results: Path = DEFAULT_WINNOW_RESULTS, - novoboard_root: Path = DEFAULT_NOVOBOARD_ROOT, + *, + novoboard_root: Path, model_root: Path = DEFAULT_MODEL_ROOT, ) -> dict[str, DatasetConfig]: """Build per-dataset path bundles from repo roots.""" @@ -1011,6 +1021,20 @@ def process_dataset(cfg: DatasetConfig, output_dir: Path) -> pd.DataFrame: @app.command() def main( + novoboard_root: Annotated[ + Path, + typer.Option( + "--novoboard-root", + help=( + "Root of NovoBoard per-dataset tables: " + "{root}/{dataset}/novoboard/ with annotated_test*.csv and " + "raw_unlabelled*.csv target/decoy pairs (the datasets/ dir of " + "a NovoBoard checkout). Local runs used fork " + "JemmaLDaniel/NovoBoard, branch feat/adapt-to-instanovo " + "(commit a9faab3ef1af06987599c2f01e6ba96072c80172)." + ), + ), + ], output_dir: Annotated[ Path, typer.Option("--output-dir", help="Directory for PNG/PDF outputs."), @@ -1019,10 +1043,6 @@ def main( Optional[list[str]], typer.Option("--datasets", help="Dataset keys to plot."), ] = None, - novoboard_root: Annotated[ - Path, - typer.Option("--novoboard-root", help="NovoBoard datasets root."), - ] = DEFAULT_NOVOBOARD_ROOT, winnow_results: Annotated[ Path, typer.Option("--winnow-results", help="Winnow results directory."), @@ -1064,7 +1084,7 @@ def main( return dataset_keys = datasets if datasets is not None else list(DEFAULT_DATASETS) - configs = build_dataset_configs(winnow_results, novoboard_root) + configs = build_dataset_configs(winnow_results, novoboard_root=novoboard_root) curve_parts: list[pd.DataFrame] = [] for key in dataset_keys: diff --git a/scripts/run_external_peptide_holdout_benchmark.py b/scripts/run_external_peptide_holdout_benchmark.py index d535cac6..fca74480 100644 --- a/scripts/run_external_peptide_holdout_benchmark.py +++ b/scripts/run_external_peptide_holdout_benchmark.py @@ -26,6 +26,17 @@ NovoBoard peptide FDR uses max-target → twin-decoy TDC. Winnow uses max calibrated confidence then nonparametric FDR (PSM-calibrator proxy). Glissade uses native bootstrap FDR with NumPy seeded from the benchmark RNG. + +External tool checkouts for local results: + +- ``--novoboard-root``: ``{root}/{dataset}/novoboard/`` target/decoy CSVs (the + ``datasets`` dir of a NovoBoard checkout). Local runs used fork + ``git@github.com:JemmaLDaniel/NovoBoard.git``, branch + ``feat/adapt-to-instanovo`` at + ``a9faab3ef1af06987599c2f01e6ba96072c80172``. +- ``--glissade-repo``: clone root with importable ``glissade.glissade``. Local + runs used fork ``git@github.com:JemmaLDaniel/glissade.git``, branch + ``winnow-benchmark`` at ``6ee11b51b5f21ba8fdc1eb5821608352b082a533``. """ from __future__ import annotations @@ -70,7 +81,6 @@ from scripts.plot_eval_results import _PALETTE, _display_name, _save_fig, _style_ax # noqa: E402 from scripts.plot_fdr_method_comparison import ( # noqa: E402 DEFAULT_MODEL_ROOT, - DEFAULT_NOVOBOARD_ROOT, DEFAULT_WINNOW_RESULTS, build_dataset_configs, load_novoboard_target_decoy, @@ -82,7 +92,6 @@ app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark_v2" -DEFAULT_GLISSADE_REPO = Path("/home/j-daniel/repos/glissade") DEFAULT_DATASETS = ["helaqc", "celegans"] DEFAULT_Q_THRESHOLDS = [round(float(x), 2) for x in np.linspace(0.0, 0.25, 26)] DEFAULT_PI0_GRID = [0.5, 0.6, 0.7, 0.8, 0.9] @@ -224,7 +233,9 @@ def build_shared_score_tables( nb_decoy_combined: namespaced twin-valid decoys for twin TDC. glissade_reference: training-split matched scores for Glissade's anchor. """ - cfg = build_dataset_configs(winnow_results, novoboard_root, model_root)[dataset] + cfg = build_dataset_configs( + winnow_results, novoboard_root=novoboard_root, model_root=model_root + )[dataset] winnow_test = _load_winnow_with_raw_confidence( cfg.winnow_test, cfg.fasta, "labelled" @@ -976,6 +987,31 @@ def write_holdout_summary_tables( @app.command() def main( + novoboard_root: Annotated[ + Path, + typer.Option( + "--novoboard-root", + help=( + "Root of NovoBoard per-dataset tables: " + "{root}/{dataset}/novoboard/ with annotated_test*.csv and " + "raw_unlabelled*.csv target/decoy pairs (the datasets/ dir of " + "a NovoBoard checkout). Local runs used fork " + "JemmaLDaniel/NovoBoard, branch feat/adapt-to-instanovo " + "(commit a9faab3ef1af06987599c2f01e6ba96072c80172)." + ), + ), + ], + glissade_repo: Annotated[ + Path, + typer.Option( + "--glissade-repo", + help=( + "Glissade clone root (must import as glissade.glissade). Local " + "runs used fork JemmaLDaniel/glissade, branch winnow-benchmark " + "(commit 6ee11b51b5f21ba8fdc1eb5821608352b082a533)." + ), + ), + ], output_dir: Annotated[ Path, typer.Option("--output-dir", help="Directory for benchmark outputs."), @@ -1016,10 +1052,6 @@ def main( Path, typer.Option("--winnow-results", help="Winnow results directory."), ] = DEFAULT_WINNOW_RESULTS, - novoboard_root: Annotated[ - Path, - typer.Option("--novoboard-root", help="NovoBoard datasets root."), - ] = DEFAULT_NOVOBOARD_ROOT, model_root: Annotated[ Path, typer.Option( @@ -1027,10 +1059,6 @@ def main( help="Per-dataset calibrator directories, used for Glissade's anchor.", ), ] = DEFAULT_MODEL_ROOT, - glissade_repo: Annotated[ - Path, - typer.Option("--glissade-repo", help="Glissade repository root."), - ] = DEFAULT_GLISSADE_REPO, n_bootstraps: Annotated[ int, typer.Option( From 700962142b0ee130d574731dc11579279bd0dd68 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:51:43 +0100 Subject: [PATCH 04/26] chore: update proteome mapping functions --- scripts/analyze_novelty.py | 2 +- scripts/fdr_tool_comparison_preprocess.py | 37 ++++++++++++++++++- scripts/plot_analysis.py | 8 ++-- scripts/plot_fdr_method_comparison.py | 4 +- .../run_external_peptide_holdout_benchmark.py | 2 +- 5 files changed, 44 insertions(+), 9 deletions(-) diff --git a/scripts/analyze_novelty.py b/scripts/analyze_novelty.py index 46efbd68..81a619b5 100644 --- a/scripts/analyze_novelty.py +++ b/scripts/analyze_novelty.py @@ -261,7 +261,7 @@ def _nontryptic_annotate( """Annotate full-search predictions with proteome and tryptic-terminus flags. ``proteome_hit`` is True when the mod-stripped prediction is a substring of - the reference proteome FASTA (same rule as ``annotate_preds_proteome_hits``). + the reference proteome FASTA (same rule as ``winnow.utils.proteome``). """ haystack = _load_proteome_haystack(fasta_path) diff --git a/scripts/fdr_tool_comparison_preprocess.py b/scripts/fdr_tool_comparison_preprocess.py index 2eed87de..2f936042 100644 --- a/scripts/fdr_tool_comparison_preprocess.py +++ b/scripts/fdr_tool_comparison_preprocess.py @@ -23,12 +23,15 @@ import numpy as np import pandas as pd +import polars as pl import yaml from instanovo.utils.metrics import Metrics from instanovo.utils.residues import ResidueSet -from scripts.annotate_preds_proteome_hits import ( +from winnow.utils.proteome import ( _batch_peptide_substring_hits, + processed_peptide_for_match, + residue_token_count, ) logger = logging.getLogger(__name__) @@ -489,6 +492,38 @@ def proteome_hit_mask( ) +def filter_and_annotate_preds( + preds: pl.DataFrame, + haystack: str, + metrics: Metrics, + min_residue_length: int, +) -> pl.DataFrame: + """Filter short peptides and annotate ``proteome_hit`` via ``winnow.utils.proteome``. + + Args: + preds: Polars frame with a ``prediction`` column. + haystack: I/L-normalised FASTA haystack from ``load_proteome_haystack``. + metrics: InstaNovo ``Metrics`` (uses ``metrics.residue_set`` for length). + min_residue_length: Drop PSMs with fewer than this many residue tokens. + """ + residue_set = metrics.residue_set + n_tok = preds["prediction"].map_elements( + lambda x: residue_token_count(x, residue_set), + return_dtype=pl.Int32, + ) + filtered = preds.with_columns(n_tok.alias("_n_residue_tokens")).filter( + pl.col("_n_residue_tokens") >= min_residue_length + ) + processed = filtered["prediction"].map_elements( + lambda x: processed_peptide_for_match(x) if isinstance(x, str) else "", + return_dtype=pl.Utf8, + ) + hits = _batch_peptide_substring_hits(processed.to_list(), haystack) + return filtered.drop("_n_residue_tokens").with_columns( + pl.Series("proteome_hit", hits, dtype=pl.Boolean) + ) + + def max_score_per_peptide( df: pd.DataFrame, key_col: str, diff --git a/scripts/plot_analysis.py b/scripts/plot_analysis.py index ee4f0e78..74d49700 100644 --- a/scripts/plot_analysis.py +++ b/scripts/plot_analysis.py @@ -31,12 +31,12 @@ REPO_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(REPO_ROOT)) -from winnow.calibration.calibrator import TrainingHistory # noqa: E402 -from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 -from scripts.annotate_preds_proteome_hits import ( # noqa: E402 +from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 filter_and_annotate_preds, - load_proteome_haystack, ) +from winnow.utils.proteome import load_proteome_haystack # noqa: E402 +from winnow.calibration.calibrator import TrainingHistory # noqa: E402 +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 # ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── _PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] diff --git a/scripts/plot_fdr_method_comparison.py b/scripts/plot_fdr_method_comparison.py index 02b76baa..477738e7 100644 --- a/scripts/plot_fdr_method_comparison.py +++ b/scripts/plot_fdr_method_comparison.py @@ -43,7 +43,7 @@ _REPO_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(_REPO_ROOT)) -from scripts.annotate_preds_proteome_hits import load_proteome_haystack # noqa: E402 +from winnow.utils.proteome import load_proteome_haystack # noqa: E402 from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 LABELLED_MIN_PEPTIDE_LENGTH, MIN_PEPTIDE_LENGTH, @@ -116,7 +116,7 @@ "winnow_suffix": "helaqc", }, "celegans": { - "fasta": "fasta/Celegans.fasta", + "fasta": "fasta/celegans.fasta", "novoboard_decoy": "0.70", "winnow_suffix": "celegans", }, diff --git a/scripts/run_external_peptide_holdout_benchmark.py b/scripts/run_external_peptide_holdout_benchmark.py index fca74480..c84f91f8 100644 --- a/scripts/run_external_peptide_holdout_benchmark.py +++ b/scripts/run_external_peptide_holdout_benchmark.py @@ -91,7 +91,7 @@ logger = logging.getLogger(__name__) app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) -DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark_v2" +DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark" DEFAULT_DATASETS = ["helaqc", "celegans"] DEFAULT_Q_THRESHOLDS = [round(float(x), 2) for x in np.linspace(0.0, 0.25, 26)] DEFAULT_PI0_GRID = [0.5, 0.6, 0.7, 0.8, 0.9] From d48c7a29d8feb2a96468586d7e9c9d49c6e1a3c6 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:10:17 +0100 Subject: [PATCH 05/26] chore: pin glissade fork as dependency --- pyproject.toml | 7 +++++++ uv.lock | 28 ++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 99ac5e59..55e5ec5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,13 @@ notebook = [ "seaborn>=0.13.2", "tqdm>=4.67.1", ] +# Paper reproduction extras (not published as package extras). +paper = [ + "glissade", +] + +[tool.uv.sources] +glissade = { git = "https://github.com/JemmaLDaniel/glissade.git", rev = "7c723a2af4a88fda84a6bd4f223b351179bd36da" } [project.scripts] winnow='winnow.scripts.main:app' diff --git a/uv.lock b/uv.lock index 4c076925..36be024e 100644 --- a/uv.lock +++ b/uv.lock @@ -1238,6 +1238,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/6e/2139de986d9c7c3ac86f1f8be43858ce90bdfe2f7175e6c80c650ba15242/gitpython-3.1.57-py3-none-any.whl", hash = "sha256:4ccf7d73c10f5c9e76043fbb2675ac5a1b3ff5b41e648f56bcbed5f63792ecaf", size = 217151, upload-time = "2026-07-26T07:33:24.838Z" }, ] +[[package]] +name = "glissade" +version = "0.0.1" +source = { git = "https://github.com/JemmaLDaniel/glissade.git?rev=7c723a2af4a88fda84a6bd4f223b351179bd36da#7c723a2af4a88fda84a6bd4f223b351179bd36da" } +dependencies = [ + { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pyarrow" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] + [[package]] name = "greenlet" version = "3.5.4" @@ -1247,14 +1263,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/9d/58f80897f4121f5c218bb931cf6d3b6514873f02ad0b729f744352926b9f/greenlet-3.5.4-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:ac5bf81d79d2c8eeb2ef6359b2e1687a1e9ebf46c2b1f970da9a9255df51d190", size = 293072, upload-time = "2026-07-22T11:38:14.299Z" }, { url = "https://files.pythonhosted.org/packages/dd/9f/b4bc9bbd6a7855cbd8ad8a83c874eeeca56c24de9132b3323f81c03a30ba/greenlet-3.5.4-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:89f3738167bab8c1084b94e23023d41d247117ac149fa0fbcb5bd4cf6262b353", size = 609393, upload-time = "2026-07-22T12:26:37.964Z" }, { url = "https://files.pythonhosted.org/packages/05/0e/744b5e063af127d2e3c74fe0f1aef15573064c83b6066883524f5b258b17/greenlet-3.5.4-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e9a5e3406e3ed8125ae1a3b37c12f3434e2b1f0fa053197c5557895b4fb09606", size = 622750, upload-time = "2026-07-22T12:28:59.546Z" }, + { url = "https://files.pythonhosted.org/packages/5d/5c/53d6b94742a6f1ee1877c7ff76262c909e137f3f7383ce96a8ab78e1ae31/greenlet-3.5.4-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a2d614cb2372c7101a12ea8b96dd56f81c986d247c5a73db67063f3ed1ca4a52", size = 629659, upload-time = "2026-07-22T12:43:40.451Z" }, { url = "https://files.pythonhosted.org/packages/d6/6b/d78ea2908e8e08985348f28ac396c2950be7ab66321dfe0054c73bd1f456/greenlet-3.5.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4ab9f0704bccf6d3b38e0d2130b7b33271cff11453690da074fa280c3aa8e8e7", size = 622920, upload-time = "2026-07-22T11:51:06.83Z" }, + { url = "https://files.pythonhosted.org/packages/f2/34/957fc5577180ef2f57be82580ee1f59fdefad4f6c623c7d5e1b6980a76fb/greenlet-3.5.4-cp310-cp310-manylinux_2_39_riscv64.whl", hash = "sha256:188e4d142f243051d92a1f5c244a741da02dddc070a0620c842804d7b56d008c", size = 425580, upload-time = "2026-07-22T12:39:48.326Z" }, { url = "https://files.pythonhosted.org/packages/eb/e4/3ce7009c948920b01527f8d9da29f501a31ac3d98318829e981fd879b850/greenlet-3.5.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:2cdaadc3d31445a8f782bde3cd37e49a2c2a9c6da6daf76a3e34c683b271a3c7", size = 1582262, upload-time = "2026-07-22T12:25:01.131Z" }, { url = "https://files.pythonhosted.org/packages/f2/4c/0408366102a33829f7bdd6a992dad75abbf75e86cc1e76caf19e57311d29/greenlet-3.5.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:70bdfacdc183dac838b2a0aaff2dd6134a457c52fe68a9c6bbab435483d2b9df", size = 1648906, upload-time = "2026-07-22T11:51:08.627Z" }, { url = "https://files.pythonhosted.org/packages/13/52/ebfe8f6a1aeb8e430540b406c844ecc4e3367072b0192f69dcb85eeeec2b/greenlet-3.5.4-cp310-cp310-win_amd64.whl", hash = "sha256:69173331fbc5d64bfac0065d7e22c39cfcd089e9b18d125bdcd5079363b09616", size = 246036, upload-time = "2026-07-22T11:38:30.073Z" }, { url = "https://files.pythonhosted.org/packages/61/16/71eefcf68267bbf06a9b6bff57d0b222e49432326e85d74348b67694b8d4/greenlet-3.5.4-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:e883de250e299654b1f1680f72a1a9f9ba62c9bd1bce84099c90657349a8dfbb", size = 294266, upload-time = "2026-07-22T11:37:56.142Z" }, { url = "https://files.pythonhosted.org/packages/36/ea/a0b19adfc35d07e10acb626e9d22a3893b95f1309c42c4a20161dec16800/greenlet-3.5.4-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:32802705c2c1ff25e8237b3bdacf2594fa02be80af8a66703eb7853ea7e68686", size = 613712, upload-time = "2026-07-22T12:26:39.375Z" }, { url = "https://files.pythonhosted.org/packages/54/76/a121978b3337407d05a1ce5f79b4aa5998a43a9d8422f9726029b90b4471/greenlet-3.5.4-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:57aa201b351f7c7c75627c60d29e4d5b97a07d37efeb62b903466fca42c097d7", size = 625582, upload-time = "2026-07-22T12:29:00.814Z" }, + { url = "https://files.pythonhosted.org/packages/d1/4a/f301f1d85c69a86b90b5d581a73e8927bba4e79450037e6e2cbca05eb4fd/greenlet-3.5.4-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9667862a2e38ad379f11b845daeda22c8989186def44f06962c9c4c05e556da7", size = 633429, upload-time = "2026-07-22T12:43:42.073Z" }, { url = "https://files.pythonhosted.org/packages/34/c2/080f16cf870e929e592f55767f01d6c98d2ee83bfdc36c3b892f2d0459ab/greenlet-3.5.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c3fe76c2cac86b4f7a1e92865ac0a54384deb05c92986287c1a7110d9bd53071", size = 624663, upload-time = "2026-07-22T11:51:08.016Z" }, + { url = "https://files.pythonhosted.org/packages/6f/2e/26884072b0eb343a4d5fee903341bfe5171b32b7f14553886e2b6349135a/greenlet-3.5.4-cp311-cp311-manylinux_2_39_riscv64.whl", hash = "sha256:ae53534b5dec0f4c2ec26f898f538dc8ea1ca3ef2927d597a9439e40a09da937", size = 428238, upload-time = "2026-07-22T12:39:49.973Z" }, { url = "https://files.pythonhosted.org/packages/9e/bb/8f3ca88370b817369008faeceeee85970adc16c92a70a3e5fe5fea495a57/greenlet-3.5.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1e1a4a684b16c45ba324e60b32a4386a87722bcb815d2a149d2182f9b401ca72", size = 1585010, upload-time = "2026-07-22T12:25:02.539Z" }, { url = "https://files.pythonhosted.org/packages/51/c2/45877154689709ebce9a0b83c2235e6ca0f31577889b02af308c8cc5f8fb/greenlet-3.5.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e849e6e139b9671adeac505f72fc05f4af7fd1921faef40295e214fc3b361b59", size = 1651283, upload-time = "2026-07-22T11:51:10.408Z" }, { url = "https://files.pythonhosted.org/packages/cd/7d/8711a75cb61d85246277c07ff6e1a6504621ba473d808c11ad225ffca43f/greenlet-3.5.4-cp311-cp311-win_amd64.whl", hash = "sha256:dc418cf4c873357964d6624445ed09472e50def990c65dd4e76fc3ba8cd9cef6", size = 246434, upload-time = "2026-07-22T11:43:15.557Z" }, @@ -1262,7 +1282,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f3/04/81bd731d6d1e3a469d9a4c36f5eb069bcf0cbb2d5d342c9fec22245b91fc/greenlet-3.5.4-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:3d66250e8b09f182ede05490998c818b5961f7a3640332d44c4927caec7bbfe4", size = 295909, upload-time = "2026-07-22T11:38:09.261Z" }, { url = "https://files.pythonhosted.org/packages/cc/dd/f5f22903a6ae70f5ea328ed0beaec92ad903f0e3b7d2845133b354abc4b8/greenlet-3.5.4-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c90e930c9c192e5b3ee9fb8bcd920ea3926155e2e3ded39fc697323addecee17", size = 612011, upload-time = "2026-07-22T12:26:40.69Z" }, { url = "https://files.pythonhosted.org/packages/8e/10/92a4a88d12b915d74ea5b6d288e4afefda4771647caa34442c156f7a454f/greenlet-3.5.4-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:791fdfeeb9c6e0c7b10fa151bf110d2a6974866f13dcb5b1c7efae698245893a", size = 624299, upload-time = "2026-07-22T12:29:02.089Z" }, + { url = "https://files.pythonhosted.org/packages/6c/f9/03e26be3487c5238e81f2b84714959a86ea8515a869828cf41f4fc54b34e/greenlet-3.5.4-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b7c895310363f310361e0fe2072af85269d2a2a285cd04c0c59e79a5e3670dcf", size = 629603, upload-time = "2026-07-22T12:43:43.456Z" }, { url = "https://files.pythonhosted.org/packages/50/6d/0b14bb9db2989f32cd9fe7f76afedea01ee8bee3f87c07e69f24adfe7e63/greenlet-3.5.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f88193799d43dbf8c8a806d6405c9c52fe2af40bf75072a606357b33cc336c7f", size = 621541, upload-time = "2026-07-22T11:51:09.464Z" }, + { url = "https://files.pythonhosted.org/packages/57/6b/7c55ca72ef80d57c16c4a55210f82582622462dc4485799a30f4ec6f3372/greenlet-3.5.4-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:13b980043cb1b3134e81ea469da1250ddcc6bfe6d245bbaa59168d9cdc8f228f", size = 432554, upload-time = "2026-07-22T12:39:51.379Z" }, { url = "https://files.pythonhosted.org/packages/48/3d/25e9a2d9eb6b2e8b7ca4e80a3a26cb887cce6c8e0a87c921164f11bc5574/greenlet-3.5.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b7a5f095767c4493afcd06067f2bb3b8716e3f3f9e92b99c88e7e99f885b3d4d", size = 1581444, upload-time = "2026-07-22T12:25:03.818Z" }, { url = "https://files.pythonhosted.org/packages/b9/96/4c9bf2e2c408dcc0556edce69efa9f802e82223573c53240136a086821f1/greenlet-3.5.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:42afdc1ab5f66da8c586c32af9224a74a706b4f0ea0dc3a4188a0860a09c65c9", size = 1645842, upload-time = "2026-07-22T11:51:12.295Z" }, { url = "https://files.pythonhosted.org/packages/b5/41/303ecb26a3a56122c0f4d4073ee078881847bd6b6f463ae0ec57ec20223b/greenlet-3.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:60149df8f462d1b230038e6590c23c3b4768bb5d6c022b3b6e82532b34b0b8a3", size = 247169, upload-time = "2026-07-22T11:38:19.893Z" }, @@ -1270,7 +1292,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/9a/e51225dcd58713f16ccbdcc501a8da21098ea14515b7870f1f94459e5ff5/greenlet-3.5.4-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:24e61b88cb7e1b1d794b32a10cc346ac779681d6d74ff137a3e0a444d2bf1f02", size = 294831, upload-time = "2026-07-22T11:38:53.389Z" }, { url = "https://files.pythonhosted.org/packages/9f/ea/de50a50fadf979713ab18b46f22ad5ff5f2dcfc637a3ebdecf669801e1a5/greenlet-3.5.4-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:870d730fec833f5a06906a32596cc099b9161594642a92a520b7a88911c95356", size = 614619, upload-time = "2026-07-22T12:26:42.282Z" }, { url = "https://files.pythonhosted.org/packages/db/c7/2aae27fea41205b8650294c301f042a2a4bb6155eea48c995b890a92f2c1/greenlet-3.5.4-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ec5ff0d1878df6af3bf9b638a5a92a7d5693291de77c91bff10fa48519c604ef", size = 627021, upload-time = "2026-07-22T12:29:03.445Z" }, + { url = "https://files.pythonhosted.org/packages/1b/80/fb4d4788bbc8e54761f1fc88533af9523a6e86299fa113d6e8a8503ed9fc/greenlet-3.5.4-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:07bd44616608d873d06735b63ef1a88191d6ca57c8d291d6559c71bc14c0893c", size = 632845, upload-time = "2026-07-22T12:43:45.19Z" }, { url = "https://files.pythonhosted.org/packages/eb/56/79fd826f9ccaae0b84e1b4ef68dabba5e105bb044ffcd448a0b782fcba9a/greenlet-3.5.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d84d993f6e575c950d91a23c1345d18fe1a4310d447bf630849d7809196b52f0", size = 624002, upload-time = "2026-07-22T11:51:11.391Z" }, + { url = "https://files.pythonhosted.org/packages/42/e3/6086fa578ebb72772722cdc4bcd628459814b42e0c2db1e3cbd6552b3271/greenlet-3.5.4-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:3529a8a933582ad19e224792cac7372489526576b75b4c124e8e4f29948f4861", size = 435053, upload-time = "2026-07-22T12:39:52.715Z" }, { url = "https://files.pythonhosted.org/packages/0a/1a/27319f97e731298513dcba1a2e91b63e9d8811d9de22130f960b129b1bf1/greenlet-3.5.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:58023945f421093de5e6fa108c0985a8659d43f49e0216da25099369a121bcbd", size = 1581533, upload-time = "2026-07-22T12:25:05.322Z" }, { url = "https://files.pythonhosted.org/packages/b1/6d/24240bf562e9786dd2799ee0a4a4dadb4ded22510f41b20245099159ac8c/greenlet-3.5.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:bae2728e1897aa8df8cb1af38cd48b3a743aefe29372de7b8b7a9f532501e69f", size = 1645781, upload-time = "2026-07-22T11:51:14.805Z" }, { url = "https://files.pythonhosted.org/packages/c1/5a/442ab1a9ef7ca6bf7210e5397a95972206a91a31033a03c8900866a10039/greenlet-3.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:ca5726c0b08ca35ae873557266a78b2c3f3b2b7d7401aa5ff886c2045dd0111c", size = 247133, upload-time = "2026-07-22T11:39:20.661Z" }, @@ -5555,6 +5579,9 @@ notebook = [ { name = "seaborn" }, { name = "tqdm" }, ] +paper = [ + { name = "glissade" }, +] [package.metadata] requires-dist = [ @@ -5595,6 +5622,7 @@ notebook = [ { name = "seaborn", specifier = ">=0.13.2" }, { name = "tqdm", specifier = ">=4.67.1" }, ] +paper = [{ name = "glissade", git = "https://github.com/JemmaLDaniel/glissade.git?rev=7c723a2af4a88fda84a6bd4f223b351179bd36da" }] [[package]] name = "wrapt" From 1501e570176d3d109939f3a3fcc7079038d8f1d0 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:41:56 +0100 Subject: [PATCH 06/26] chore: update paper reproduction dependency group --- pyproject.toml | 3 + uv.lock | 562 +++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 478 insertions(+), 87 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 55e5ec5f..9cbf26b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,9 @@ notebook = [ # Paper reproduction extras (not published as package extras). paper = [ "glissade", + "matplotlib-venn>=1.1.1", + "shap>=0.44.0", + "xgboost>=3.2.0", ] [tool.uv.sources] diff --git a/uv.lock b/uv.lock index 36be024e..b8437afc 100644 --- a/uv.lock +++ b/uv.lock @@ -7,10 +7,14 @@ resolution-markers = [ "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version < '3.11' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -34,8 +38,8 @@ version = "1.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "packaging" }, { name = "psutil" }, { name = "pyyaml" }, @@ -361,8 +365,8 @@ name = "biopython" version = "1.87" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/df/3e/3c6aa8b2a7e6b791a34407736db32f59657001f0446ada31db73a3e0b7d5/biopython-1.87.tar.gz", hash = "sha256:8456c803459b679a9712422e5a7fd9809f2f089bf69bb085f3b077946ac9bdbf", size = 19855264, upload-time = "2026-03-30T11:28:29.823Z" } wheels = [ @@ -627,6 +631,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -726,17 +739,22 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } wheels = [ @@ -884,8 +902,8 @@ dependencies = [ { name = "httpx" }, { name = "huggingface-hub" }, { name = "multiprocess" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "packaging" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -1245,8 +1263,8 @@ source = { git = "https://github.com/JemmaLDaniel/glissade.git?rev=7c723a2af4a88 dependencies = [ { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pyarrow" }, @@ -1482,8 +1500,8 @@ dependencies = [ { name = "jiwer" }, { name = "matchms" }, { name = "neptune" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "omegaconf" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -1569,10 +1587,14 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -1664,10 +1686,14 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -2139,10 +2165,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/82/3d/14ce75ef66813643812f3093ab17e46d3a206942ce7376d31ec2d36229e7/lark-1.3.1-py3-none-any.whl", hash = "sha256:c629b661023a014c37da873b4ff58a817398d12635d3bbb2c5a03be7fe5d1e12", size = 113151, upload-time = "2025-10-27T18:25:54.882Z" }, ] +[[package]] +name = "llvmlite" +version = "0.45.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +sdist = { url = "https://files.pythonhosted.org/packages/99/8d/5baf1cef7f9c084fb35a8afbde88074f0d6a727bc63ef764fe0e7543ba40/llvmlite-0.45.1.tar.gz", hash = "sha256:09430bb9d0bb58fc45a45a57c7eae912850bedc095cd0810a57de109c69e1c32", size = 185600, upload-time = "2025-10-01T17:59:52.046Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cf/6d/585c84ddd9d2a539a3c3487792b3cf3f988e28ec4fa281bf8b0e055e1166/llvmlite-0.45.1-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:1b1af0c910af0978aa55fa4f60bbb3e9f39b41e97c2a6d94d199897be62ba07a", size = 43043523, upload-time = "2025-10-01T18:02:58.621Z" }, + { url = "https://files.pythonhosted.org/packages/04/ad/9bdc87b2eb34642c1cfe6bcb4f5db64c21f91f26b010f263e7467e7536a3/llvmlite-0.45.1-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:60f92868d5d3af30b4239b50e1717cb4e4e54f6ac1c361a27903b318d0f07f42", size = 43043526, upload-time = "2025-10-01T18:03:15.051Z" }, + { url = "https://files.pythonhosted.org/packages/e2/7c/82cbd5c656e8991bcc110c69d05913be2229302a92acb96109e166ae31fb/llvmlite-0.45.1-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:28e763aba92fe9c72296911e040231d486447c01d4f90027c8e893d89d49b20e", size = 43043524, upload-time = "2025-10-01T18:03:30.666Z" }, + { url = "https://files.pythonhosted.org/packages/1d/e2/c185bb7e88514d5025f93c6c4092f6120c6cea8fe938974ec9860fb03bbb/llvmlite-0.45.1-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:d9ea9e6f17569a4253515cc01dade70aba536476e3d750b2e18d81d7e670eb15", size = 43043524, upload-time = "2025-10-01T18:03:43.249Z" }, +] + [[package]] name = "llvmlite" version = "0.48.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version >= '3.13' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'win32')", + "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] sdist = { url = "https://files.pythonhosted.org/packages/dc/a0/acc8ffcd5bdc63df0097e22c719bfcd61b604358343089313a8aebbb24ab/llvmlite-0.48.0.tar.gz", hash = "sha256:543b19f9ef8f3c7c60d1468191e4ee1b1537bf9f8a3d56f64c0ddd98de92edd2", size = 184016, upload-time = "2026-07-02T20:20:05.308Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/a2/4e/32543c42568fb321b3bdfcf9106e4116ab8f5a7bbcfd9ecf5569b0c07d83/llvmlite-0.48.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:614aad57df707e3172efd5165f2aa7da6a0c6897e40dce590bf756396815ba76", size = 40480650, upload-time = "2026-07-01T18:41:01.945Z" }, @@ -2351,9 +2412,10 @@ dependencies = [ { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "networkx" }, - { name = "numba" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or platform_machine != 'x86_64' or sys_platform != 'darwin'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pickydict" }, @@ -2447,10 +2509,14 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -2461,7 +2527,8 @@ dependencies = [ { name = "cycler", marker = "python_full_version >= '3.11'" }, { name = "fonttools", marker = "python_full_version >= '3.11'" }, { name = "kiwisolver", marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "packaging", marker = "python_full_version >= '3.11'" }, { name = "pillow", marker = "python_full_version >= '3.11'" }, { name = "pyparsing", marker = "python_full_version >= '3.11'" }, @@ -2514,6 +2581,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" }, ] +[[package]] +name = "matplotlib-venn" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e8/f7/47bddf95492f4d1370ed7164d2b16407805e8eeb38231361de65d387a562/matplotlib-venn-1.1.2.tar.gz", hash = "sha256:6f2b07a03e9bb5a62de2f32f965216739e175176f9d654dd19e7af2c22ec36e3", size = 40821, upload-time = "2025-02-25T10:44:24.294Z" } + [[package]] name = "mdurl" version = "0.1.2" @@ -2618,8 +2699,8 @@ name = "ml-dtypes" version = "0.5.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } wheels = [ @@ -2985,14 +3066,53 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f9/33/bd5b9137445ea4b680023eb0469b2bb969d61303dedb2aac6560ff3d14a1/notebook_shim-0.2.4-py3-none-any.whl", hash = "sha256:411a5be4e9dc882a074ccbcae671eda64cceb068767e9a3419096986560e1cef", size = 13307, upload-time = "2024-02-14T23:35:16.286Z" }, ] +[[package]] +name = "numba" +version = "0.62.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "llvmlite", version = "0.45.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/20/33dbdbfe60e5fd8e3dbfde299d106279a33d9f8308346022316781368591/numba-0.62.1.tar.gz", hash = "sha256:7b774242aa890e34c21200a1fc62e5b5757d5286267e71103257f4e2af0d5161", size = 2749817, upload-time = "2025-09-29T10:46:31.551Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f5/27/a5a9a58f267ec3b72f609789b2a8eefd6156bd7117e41cc9b7cf5de30490/numba-0.62.1-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:a323df9d36a0da1ca9c592a6baaddd0176d9f417ef49a65bb81951dce69d941a", size = 2684281, upload-time = "2025-09-29T10:43:31.863Z" }, + { url = "https://files.pythonhosted.org/packages/dd/5f/8b3491dd849474f55e33c16ef55678ace1455c490555337899c35826836c/numba-0.62.1-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:f43e24b057714e480fe44bc6031de499e7cf8150c63eb461192caa6cc8530bc8", size = 2684279, upload-time = "2025-09-29T10:43:37.213Z" }, + { url = "https://files.pythonhosted.org/packages/5e/fa/30fa6873e9f821c0ae755915a3ca444e6ff8d6a7b6860b669a3d33377ac7/numba-0.62.1-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:1b743b32f8fa5fff22e19c2e906db2f0a340782caf024477b97801b918cf0494", size = 2685346, upload-time = "2025-09-29T10:43:43.677Z" }, + { url = "https://files.pythonhosted.org/packages/22/76/501ea2c07c089ef1386868f33dff2978f43f51b854e34397b20fc55e0a58/numba-0.62.1-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:b72489ba8411cc9fdcaa2458d8f7677751e94f0109eeb53e5becfdc818c64afb", size = 2685766, upload-time = "2025-09-29T10:43:49.161Z" }, +] + [[package]] name = "numba" version = "0.66.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version >= '3.13' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'win32')", + "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] dependencies = [ - { name = "llvmlite" }, + { name = "llvmlite", version = "0.48.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or platform_machine != 'x86_64' or sys_platform != 'darwin'" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ae/a0/570e3dc53e5602b49108f62a13e529f1eec8bfc7ef37d49c825924dcf546/numba-0.66.0.tar.gz", hash = "sha256:b900e63a0e26c05ea9a6d5a3a5a0a177cb64c5011887bf43edb8c3ed2c38d363", size = 2806181, upload-time = "2026-07-01T23:12:46.36Z" } wheels = [ @@ -3020,6 +3140,10 @@ version = "2.2.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version < '3.11' and sys_platform == 'darwin'", "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] @@ -3090,10 +3214,10 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -3260,9 +3384,19 @@ name = "nvidia-nccl-cu12" version = "2.27.3" source = { registry = "https://pypi.org/simple" } wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/7b/8354b784cf73b0ba51e566b4baba3ddd44fe8288a3d39ef1e06cd5417226/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9ddf1a245abc36c550870f26d537a9b6087fb2e2e3d6e0ef03374c6fd19d984f", size = 322397768, upload-time = "2025-06-03T21:57:30.234Z" }, { url = "https://files.pythonhosted.org/packages/5c/5b/4e4fff7bad39adf89f735f2bc87248c81db71205b62bcc0d5ca5b606b3c3/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039", size = 322364134, upload-time = "2025-06-03T21:58:04.013Z" }, ] +[[package]] +name = "nvidia-nccl-cu13" +version = "2.31.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/a0/530efd7db8857c0436868bb7df9764f09fde2bd4d1f0bae546eec9fc40d0/nvidia_nccl_cu13-2.31.2-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:b5563f8e2534f363d93ace022670ba016d3717e190ac4eba564d05fbbe8495b1", size = 252479893, upload-time = "2026-08-11T23:22:01.53Z" }, + { url = "https://files.pythonhosted.org/packages/14/fb/94933e00bb3dcfdf66ea3456739c6a51d322353f7cc64fa1f5f660e695ac/nvidia_nccl_cu13-2.31.2-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:0bcaf0308854cb55fcc35af72e2c83143f3b71e65a4e865e2c586b1cdcdb5ae0", size = 252442223, upload-time = "2026-08-11T23:22:40.341Z" }, +] + [[package]] name = "nvidia-nvjitlink-cu12" version = "12.8.93" @@ -3390,17 +3524,22 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, { name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, ] @@ -3731,8 +3870,8 @@ version = "1.3.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "lxml" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "six" }, { name = "sqlalchemy" }, ] @@ -3911,8 +4050,10 @@ version = "0.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "joblib" }, - { name = "llvmlite" }, - { name = "numba" }, + { name = "llvmlite", version = "0.45.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "llvmlite", version = "0.48.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or platform_machine != 'x86_64' or sys_platform != 'darwin'" }, + { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or platform_machine != 'x86_64' or sys_platform != 'darwin'" }, { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scikit-learn", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -4285,8 +4426,8 @@ name = "rdkit" version = "2026.3.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "pillow" }, ] wheels = [ @@ -4504,10 +4645,14 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -4677,10 +4822,14 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", @@ -4689,7 +4838,8 @@ resolution-markers = [ dependencies = [ { name = "joblib", marker = "python_full_version >= '3.11'" }, { name = "narwhals", marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "threadpoolctl", marker = "python_full_version >= '3.11'" }, ] @@ -4785,17 +4935,22 @@ resolution-markers = [ "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0a/ca/d8ace4f98322d01abcd52d381134344bf7b431eba7ed8b42bdea5a3c2ac9/scipy-1.16.3.tar.gz", hash = "sha256:01e87659402762f43bd2fee13370553a17ada367d42e7487800bf2916535aecb", size = 30597883, upload-time = "2025-10-28T17:38:54.068Z" } wheels = [ @@ -4848,8 +5003,8 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] @@ -4876,6 +5031,150 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, ] +[[package]] +name = "shap" +version = "0.49.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "cloudpickle", marker = "python_full_version < '3.11'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "packaging", marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "slicer", marker = "python_full_version < '3.11'" }, + { name = "tqdm", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/c6/9823a7f483aa9f3179fc359c10d22da9e418b1a7a3fc99a42b705d05e82a/shap-0.49.1.tar.gz", hash = "sha256:1114ecd804fff29f50d522ce6031082fcf42fe4a32fb1b5da233b2415d784c8c", size = 4084725, upload-time = "2025-10-14T10:04:49.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/a1/66b4f04995ee23ff8638c21294f1a3a6dc87397af54c87aeeb037500f71f/shap-0.49.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:40140ec5d306719f89daee1df27805a71bcc1ac39630832455d316d0306d1283", size = 558950, upload-time = "2025-10-14T10:04:08.441Z" }, + { url = "https://files.pythonhosted.org/packages/06/76/2142615fa5cc745fd66beb066d00db123cc86d614a31ca8029b29537a959/shap-0.49.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6e9977f1e0b6bba57967de600e8e6047b3e4643d06a4671f2dba1a97c1b5ab3e", size = 556605, upload-time = "2025-10-14T10:04:10.049Z" }, + { url = "https://files.pythonhosted.org/packages/a8/3a/e28014ffc23f386da3d69abd978838e653fff5641831e5a34aade3f4dfe7/shap-0.49.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:54ad4c38e6af56eaa1c892bb3af550a35df15ca0d27d2d41c1d1619ca6a2ba75", size = 1000329, upload-time = "2025-10-14T10:04:11.359Z" }, + { url = "https://files.pythonhosted.org/packages/bd/09/734325f0a9ab9d3dfa5c0908a927027b3d95b3f6929bb62d88e840b85abf/shap-0.49.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fcd832e97038648ba89f659863322d5cd3ea0815e18c36dd48cd7ae1ca9f264b", size = 1000713, upload-time = "2025-10-14T10:04:12.573Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a2/0518acabb104e21fecda65b0202e41edd06637c44dac15e2197e7d13a002/shap-0.49.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:7fc2e864908277dca2b1d9c59a18b3f31576b985bd024f39b0c3cb7e2c7441db", size = 2065477, upload-time = "2025-10-14T10:04:13.874Z" }, + { url = "https://files.pythonhosted.org/packages/af/e1/3d52717b617b9ad1e4d0c9634d3b7c52a913540fde27c4b4663a7ee76b87/shap-0.49.1-cp310-cp310-win_amd64.whl", hash = "sha256:4f5bec3d061b4f4889e1ac4e9b676aede2875778ff44b9d5f5a844cbe6788fd2", size = 547034, upload-time = "2025-10-14T10:04:16.17Z" }, + { url = "https://files.pythonhosted.org/packages/1d/08/d433b7d18a8b51a7d10477120f78877d806d2eb86283cb1661318d865f3d/shap-0.49.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1e208a0129c721bd0eba6268a9ffac4610dbc8a833d07d2ad9f39541bb737f06", size = 558742, upload-time = "2025-10-14T10:04:17.45Z" }, + { url = "https://files.pythonhosted.org/packages/c2/35/72929fdad25e055aff9dfbeb48c044682fc3b815d90cee4036b90bd65f4c/shap-0.49.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0b878470bdf6800069c25d2a8598eb0548aa1e6826becd39cca253521cc14866", size = 556486, upload-time = "2025-10-14T10:04:18.934Z" }, + { url = "https://files.pythonhosted.org/packages/02/be/d92623be2c584784e99a8eb9a6cd02263b4eb363c9e49fa14c20f824bcbb/shap-0.49.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:118577d40c53f005268024e59f6a10cbcafbb6d03b3d97dce7c0c7510190ebaa", size = 1025978, upload-time = "2025-10-14T10:04:20.096Z" }, + { url = "https://files.pythonhosted.org/packages/14/e9/e4079b5de26a8269121ce38125e130c147dac7b59611e0bd94be10f9444e/shap-0.49.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f424465699aa2dda8057656c6b6d3cb927cf29b054c5bb01cfffcb9efa5dbf98", size = 1027831, upload-time = "2025-10-14T10:04:21.666Z" }, + { url = "https://files.pythonhosted.org/packages/49/ff/e22e1d899ed56384a2395d6121d6e21833c518c01c5b6c52fce3c0b0cbab/shap-0.49.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d505834fdf2a159e88b1dcdeddfd79f101fd789ba89d589faf0aaec060c0bad9", size = 2092627, upload-time = "2025-10-14T10:04:22.894Z" }, + { url = "https://files.pythonhosted.org/packages/17/48/bbcd638a391ac0fb30033398a3cca60ba5c36941d962dd74958e67069108/shap-0.49.1-cp311-cp311-win_amd64.whl", hash = "sha256:897c7e6fa98d66482282c8f898c97ade181d714ecaf581da0dab5c49adb9f62c", size = 546845, upload-time = "2025-10-14T10:04:24.238Z" }, + { url = "https://files.pythonhosted.org/packages/92/7a/ccecf7a9158baa10bdc5146907c72dd5f85c762cb5f16cdc74d15cebb8a1/shap-0.49.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c652dc77f1fffe73f5a3def3356c5090e2e6401c261e4fe5329d83cb6251e772", size = 559663, upload-time = "2025-10-14T10:04:25.412Z" }, + { url = "https://files.pythonhosted.org/packages/ee/c6/c43382d6c891fcf067d0a9f6d954351e3c7d330f4328c5816769b796aa27/shap-0.49.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c23f1493205e648634680c8974e82e7f4b2e96ae3a7eca2251680172bd197ae9", size = 556265, upload-time = "2025-10-14T10:04:27.098Z" }, + { url = "https://files.pythonhosted.org/packages/c0/71/f7db7a5a2cedaa3ac52f58f453172d613be041bedd9509ce5b5cba2096a6/shap-0.49.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:41147740c42821023e1b60185ce8be989656ccac266cc9490d7a8e3ad53c556a", size = 1022419, upload-time = "2025-10-14T10:04:28.793Z" }, + { url = "https://files.pythonhosted.org/packages/c2/a4/96ca9a69dd669ff835ddef875c5dd8e07599103769417d3e9051fd97d470/shap-0.49.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ef9952929d4a7e6763d2716938067bdad762217e3afb46cabfc15a62c012b364", size = 1027074, upload-time = "2025-10-14T10:04:30.2Z" }, + { url = "https://files.pythonhosted.org/packages/fc/9a/89ed1ac8beffe8ff8e09c12cb351bc3c79ddaadcc47ca6ee434d76e464d7/shap-0.49.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e823417eb0a01947cd9bd763bef2e534c5aef7a7c2952b1badfa969c7d59d3b3", size = 2088172, upload-time = "2025-10-14T10:04:31.725Z" }, + { url = "https://files.pythonhosted.org/packages/4a/28/11422c1c3aa022a06e76cbfa3267e1750cedc00c1e02ef1ccae9c88cd6f4/shap-0.49.1-cp312-cp312-win_amd64.whl", hash = "sha256:cb28043decfec3f35f795421eb5a81545f629b7f60bbf7449cd2843a7f1c8cc6", size = 548036, upload-time = "2025-10-14T10:04:33.087Z" }, + { url = "https://files.pythonhosted.org/packages/e9/5c/030bbfa19605ca4ad66a753d55e76aee5093be6748a6d33eda89e5613995/shap-0.49.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:333cd8e8c427badda92d5ada9e7aad1e3e1e8e7e0398da51a18b7ffb03514e45", size = 558604, upload-time = "2025-10-14T10:04:34.298Z" }, + { url = "https://files.pythonhosted.org/packages/2c/7f/7e7b78e9fac6f891096fb6a59a6d4db23243b0af2369ae54e161f513c485/shap-0.49.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f4faf61560f73a66f4f26bc027c91f8939201979c4db24949dca305ba0a2ad36", size = 555311, upload-time = "2025-10-14T10:04:35.582Z" }, + { url = "https://files.pythonhosted.org/packages/f2/be/25283a0f8c30deaf897b89a0dbfd490d330f6fc68caa6f19db6e130832e9/shap-0.49.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b440da658d9aee7711bf642c9b4826d81f588fb478cd9e90c068646e90f56669", size = 1016897, upload-time = "2025-10-14T10:04:36.856Z" }, + { url = "https://files.pythonhosted.org/packages/5c/91/a63e563f3dc8e134db12dd155a1a6ed5e0649f79fc8ac651aac1088e8652/shap-0.49.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d8dfa5654eccf4d13dcb262a10314a4e0eb1060db842b2ef31e9fb0038168bc1", size = 1022476, upload-time = "2025-10-14T10:04:38.171Z" }, + { url = "https://files.pythonhosted.org/packages/15/a2/89303c1f7eb206658bf9ec974dc6e69b0a6bd309cf5de0cfa8f92f5a8eb3/shap-0.49.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ed3080030a6000d3737841c5770ed555b8a922b794fa0ba5aae1e45655eda1fa", size = 2087940, upload-time = "2025-10-14T10:04:39.497Z" }, + { url = "https://files.pythonhosted.org/packages/84/bd/0b9b3e19b9b8cda51463f8a749dc354eb9c87f42eddcbfdf742dceb3746b/shap-0.49.1-cp313-cp313-win_amd64.whl", hash = "sha256:6af779344c23b12a47063aab7fc135fefbdb5849233c1813f11dd8cf2fc73bea", size = 547806, upload-time = "2025-10-14T10:04:40.712Z" }, +] + +[[package]] +name = "shap" +version = "0.51.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "cloudpickle", marker = "python_full_version == '3.11.*'" }, + { name = "llvmlite", version = "0.45.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "llvmlite", version = "0.48.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version == '3.11.*' and platform_machine != 'x86_64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')" }, + { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version == '3.11.*' and platform_machine != 'x86_64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version == '3.11.*' and platform_machine != 'x86_64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')" }, + { name = "packaging", marker = "python_full_version == '3.11.*'" }, + { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "scikit-learn", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "slicer", marker = "python_full_version == '3.11.*'" }, + { name = "tqdm", marker = "python_full_version == '3.11.*'" }, + { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a4/0a/4a3ee4b1a3654f2a9ae038a64bb3e91a42af3da07577d69b65241f010970/shap-0.51.0.tar.gz", hash = "sha256:cfa17ff213657c9d50285aa923d79b0037a62e2ee1a31bc3eec7e196b00bdb59", size = 4108336, upload-time = "2026-03-04T09:18:19.985Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1b/4f/513ffc1c27242488be2269d5704c7bd8e82c6b28a04c297533d616003948/shap-0.51.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4b4b81d7401bc821490148a694a9f149f8ef488fa973b49fdc8a7916a572750f", size = 565103, upload-time = "2026-03-04T09:17:24.11Z" }, + { url = "https://files.pythonhosted.org/packages/14/bf/bafa92ae6606f1ee38f14e866045fbcae21412a3a5766fb493b951a6e6e1/shap-0.51.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9b92f7371644c4a19660f734e0b0fe299d0eff273102230c9cf191e310576619", size = 562798, upload-time = "2026-03-04T09:17:25.909Z" }, + { url = "https://files.pythonhosted.org/packages/fc/09/a952ef8a1fe64b8a7bb909b2d3ab614d33cff7fdf646d62c3bd35d84de02/shap-0.51.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ddd314c9dc766f72027b3b6d9a8b975c7df67f1a27af7f857267c716b7d3dd8", size = 1052516, upload-time = "2026-03-04T09:17:27.816Z" }, + { url = "https://files.pythonhosted.org/packages/42/48/541bd5b7f068804f33fac459274c06a46deb8f0e1edf3a9b176c00137629/shap-0.51.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5a909d6e306d4083e4862b6d097d53f3b4aa4d5cdf2ef4dbac5e8245160d9abd", size = 1060070, upload-time = "2026-03-04T09:17:29.751Z" }, + { url = "https://files.pythonhosted.org/packages/5f/f9/02269c83b3056c5fbaa13911e59b844146ed80c97a6f78aa0d374a9e8368/shap-0.51.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ff5534e92f5876c74975b53cc6b251126bbe092e5032f81e1fe8583c4a74d58c", size = 2023431, upload-time = "2026-03-04T09:17:31.97Z" }, + { url = "https://files.pythonhosted.org/packages/72/ac/751f3070be48c4b365f565ca5b4ba4c93b0a371d8ee595859c62eab9885b/shap-0.51.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:6f3a83f28d2304f81645392a1d346154d9d992705e762fb949fa5371925399b2", size = 2092764, upload-time = "2026-03-04T09:17:33.991Z" }, + { url = "https://files.pythonhosted.org/packages/a5/8e/cee1ee136a4e54fe2fbb63a60d72d7c25e21a4ffe6aa05779cab7669cb31/shap-0.51.0-cp311-cp311-win_amd64.whl", hash = "sha256:ca2a9171e6c5b9a700d585982d7fd8336856ad818e24264420c56f9d8f02a961", size = 554870, upload-time = "2026-03-04T09:17:35.856Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ba/8c8fac8506327febada7dc58f90dc459287995bb7b8aadbc44506e61be55/shap-0.51.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:07b0367408b1b9fc51556f2ddac5ee4209cc51be592099e6d51d0834c9b037d8", size = 565741, upload-time = "2026-03-04T09:17:37.286Z" }, + { url = "https://files.pythonhosted.org/packages/49/fd/07f7c454ff5dff455576e5bb08cdb2cab05a4c1eb5e1b9959ef2ac28366d/shap-0.51.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e412bb475c9074ffd6684abb88d86d93275729b344cbfb37b4e4db37db759fbf", size = 562281, upload-time = "2026-03-04T09:17:39.006Z" }, + { url = "https://files.pythonhosted.org/packages/93/a1/37e7229be000cf608ece024dcd76edae4cc618b22b402ea78270849cac3f/shap-0.51.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1aa9f659d2028e26ac7ec34cafbc14585fdc14d0c8973e9442c65af1af1ff781", size = 1051009, upload-time = "2026-03-04T09:17:40.989Z" }, + { url = "https://files.pythonhosted.org/packages/af/c1/a9152876b04f9a05ca18bd3e8bc4bc72468ae32429bfbb30a9cbd4ad35b9/shap-0.51.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0b1f9e62d6a3fa28765d7b61abda7caf76aba21e769423fbf3ce8a7a5e498243", size = 1062849, upload-time = "2026-03-04T09:17:42.587Z" }, + { url = "https://files.pythonhosted.org/packages/80/89/38c903c438b33063b006f41d00684af8b424bb95f0fcfd8963d1501bf427/shap-0.51.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:dee16e81082dec5ce2a37c41c2b9cbebcb4bf7de79133a72d84a4093b7d4158c", size = 2014842, upload-time = "2026-03-04T09:17:44.212Z" }, + { url = "https://files.pythonhosted.org/packages/c9/fd/9b295ad15420566dca713b792d9beb65692804b96c69cc99ffec5e31db58/shap-0.51.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c3b878f6414213a12247faa00d609957fdfbcc33cfd48a6751500c4708b5666d", size = 2090611, upload-time = "2026-03-04T09:17:46.728Z" }, + { url = "https://files.pythonhosted.org/packages/20/0e/6f581645b66efff6bf091953f474eb16e64da499cfac0c552dd77559f205/shap-0.51.0-cp312-cp312-win_amd64.whl", hash = "sha256:ee76aa705927ac64acd4f506722f52596e77d3ced87078bc86bfcb4571c7b976", size = 556117, upload-time = "2026-03-04T09:17:48.68Z" }, + { url = "https://files.pythonhosted.org/packages/36/4a/77f8dc2c6874d8a27123afffcb79f540a80ed3ccfd640604e4d8beb9cc5e/shap-0.51.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f3efeef8a76e2ee735fad50f82e5a8e56ba8639f42e2fa50dc1997caed77c488", size = 564931, upload-time = "2026-03-04T09:17:51.022Z" }, + { url = "https://files.pythonhosted.org/packages/e6/5e/5c6d37992e93b3fa44509d8544281cd5ae357c8946bc0e756e78139b4baf/shap-0.51.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:724cdd8298450ae22b08ca40e07136c73e4e75dbdf8e3f07d741a291bf636dad", size = 561686, upload-time = "2026-03-04T09:17:52.462Z" }, + { url = "https://files.pythonhosted.org/packages/ec/b7/76dbca9c4b83602841c016fec7201e4146c5e6347a8b0428e7c0617ba424/shap-0.51.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:13eb21626ea671604769c847ac0604871a03df0842522087ffc00181683780a4", size = 1050719, upload-time = "2026-03-04T09:17:54.252Z" }, + { url = "https://files.pythonhosted.org/packages/64/b1/472ca0adf25215dfdbca9f398d853536413091fe47dd69bb3f67dbd445f4/shap-0.51.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f32d006680708513efff07f67dcbe44a531b242ae042dee99b7024e210391ac2", size = 1063794, upload-time = "2026-03-04T09:17:56.154Z" }, + { url = "https://files.pythonhosted.org/packages/c6/8e/9072acfcbe6abc79fbfe87360c7dcfe16d7498cdb13dc560820912eb5dd5/shap-0.51.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:33540a7e3c70bd1a742a4d0576c19b5e000165038de953d3ebf31a7bb53d01f4", size = 2012838, upload-time = "2026-03-04T09:17:57.835Z" }, + { url = "https://files.pythonhosted.org/packages/f3/65/b95588a1f48eb9e98aa61e6db31cf63a388970e4c11341d40ddece3b54f8/shap-0.51.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f461e5a1d6b0cae3fd9c6bd00c95111ed95b9e0020ec71f14291429ed17d49f4", size = 2090245, upload-time = "2026-03-04T09:17:59.501Z" }, + { url = "https://files.pythonhosted.org/packages/79/d5/1ec3120f461f31a03d1d2f1d339f5058f12c7a542d22bfcc350511eccc8a/shap-0.51.0-cp313-cp313-win_amd64.whl", hash = "sha256:5f51ca55bda10b3fa2125f2b8e08e9d6a6edcbb0e752c67050e87a6d3ca7d53b", size = 555927, upload-time = "2026-03-04T09:18:01.274Z" }, +] + +[[package]] +name = "shap" +version = "0.52.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version >= '3.13' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'win32')", + "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "cloudpickle", marker = "python_full_version >= '3.12'" }, + { name = "llvmlite", version = "0.45.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "llvmlite", version = "0.48.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and platform_machine != 'x86_64') or (python_full_version >= '3.12' and sys_platform != 'darwin')" }, + { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and platform_machine != 'x86_64') or (python_full_version >= '3.12' and sys_platform != 'darwin')" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and platform_machine != 'x86_64') or (python_full_version >= '3.12' and sys_platform != 'darwin')" }, + { name = "packaging", marker = "python_full_version >= '3.12'" }, + { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "scikit-learn", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "slicer", marker = "python_full_version >= '3.12'" }, + { name = "tqdm", marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fc/0a/aa278f42c08cb47f2bb503085be0c521da2886929c6605b6105748a7590f/shap-0.52.0.tar.gz", hash = "sha256:81d4ae478f67f8122de1bb411dc4e3ddff0604cbc27dc9cb8ea66d5c73462fd2", size = 4192842, upload-time = "2026-05-28T14:17:49.011Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/09/61/ddbe6fb40120fc7d3dbd702f5f4e0ef1c0795e39f208afbc928ce46a0246/shap-0.52.0-cp312-abi3-macosx_10_13_x86_64.whl", hash = "sha256:334cdc36a925db2242875f69267b88eb6108ec47a6c259f4f87a6b022b6188dc", size = 496146, upload-time = "2026-05-28T14:17:33.352Z" }, + { url = "https://files.pythonhosted.org/packages/19/d1/b020cb524513496d046a9711ec466c0fdd479b722c09ceb1162c138d0db7/shap-0.52.0-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:9a1116361c01fc5a045cf34681673f6c79100e2e648a6fbde7da084d18193d2e", size = 490868, upload-time = "2026-05-28T14:17:35.075Z" }, + { url = "https://files.pythonhosted.org/packages/88/58/6e7f8d13b6078485a4bc3c5e6ae97ef52c9208315206982fd0fedabe2db4/shap-0.52.0-cp312-abi3-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61e431aed9de5f2deae1b7847b00edd739b8d85df9c4d04b137230d20dbbd4d3", size = 495268, upload-time = "2026-05-28T14:17:36.312Z" }, + { url = "https://files.pythonhosted.org/packages/fe/1c/8aad9c7cc4c09a3841496def5434f56d1b172ad55dbb42fc839c25798f1c/shap-0.52.0-cp312-abi3-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:20b3d09fdff3e94e418abec1f6033522189e4ca554f1014a2b320f80b5454ad2", size = 498042, upload-time = "2026-05-28T14:17:37.664Z" }, + { url = "https://files.pythonhosted.org/packages/8e/78/f8f86c768a2fae213d99721666eb01653347b1398cf6e24afd84743899a8/shap-0.52.0-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:6b843d61e18ad4659584e004a4e681fff0648d4cf371830e3019709722190467", size = 1570560, upload-time = "2026-05-28T14:17:39.058Z" }, + { url = "https://files.pythonhosted.org/packages/5d/20/f5824640d8e7bf6bffb2ed8f6221c8c6fb2d39b638ba72f01b60e934e40f/shap-0.52.0-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:2828d18366db812599a5d8340ce93fed9c95d97337d90568597b2e96d01ce516", size = 1627401, upload-time = "2026-05-28T14:17:40.435Z" }, + { url = "https://files.pythonhosted.org/packages/58/bf/6be16d28ef1b6ff69078a1d7ea58892e9d40a4680c1077563f74ebd31c9e/shap-0.52.0-cp312-abi3-win_amd64.whl", hash = "sha256:07d44ace491ca6204dac6ce4fda128bcaeff27553a279bca67c55a14987cc957", size = 499853, upload-time = "2026-05-28T14:17:41.786Z" }, +] + [[package]] name = "shellingham" version = "1.5.4" @@ -4959,6 +5258,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "slicer" +version = "0.0.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/f9/b4bce2825b39b57760b361e6131a3dacee3d8951c58cb97ad120abb90317/slicer-0.0.8.tar.gz", hash = "sha256:2e7553af73f0c0c2d355f4afcc3ecf97c6f2156fcf4593955c3f56cf6c4d6eb7", size = 14894, upload-time = "2024-03-09T23:35:26.826Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/81/9ef641ff4e12cbcca30e54e72fb0951a2ba195d0cda0ba4100e532d929db/slicer-0.0.8-py3-none-any.whl", hash = "sha256:6c206258543aecd010d497dc2eca9d2805860a0b3758673903456b7df7934dc3", size = 15251, upload-time = "2024-03-09T07:03:07.708Z" }, +] + [[package]] name = "smmap" version = "5.0.3" @@ -4982,9 +5290,10 @@ name = "sparsestack" version = "0.7.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numba" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or platform_machine != 'x86_64' or sys_platform != 'darwin'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] @@ -5002,9 +5311,10 @@ dependencies = [ { name = "lark" }, { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "numba" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numba", version = "0.66.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or platform_machine != 'x86_64' or sys_platform != 'darwin'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "platformdirs" }, { name = "pyteomics" }, ] @@ -5103,8 +5413,8 @@ dependencies = [ { name = "absl-py" }, { name = "grpcio" }, { name = "markdown" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "packaging" }, { name = "pillow" }, { name = "protobuf" }, @@ -5202,10 +5512,14 @@ name = "torch" version = "2.7.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version < '3.11' and sys_platform == 'darwin'", ] dependencies = [ @@ -5351,8 +5665,8 @@ version = "2.70.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ml-dtypes" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "python-rapidjson" }, { name = "urllib3" }, ] @@ -5363,8 +5677,8 @@ wheels = [ [package.optional-dependencies] grpc = [ { name = "grpcio" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.11' and sys_platform != 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'x86_64') or (python_full_version >= '3.11' and sys_platform != 'darwin')" }, { name = "packaging" }, { name = "protobuf" }, { name = "python-rapidjson" }, @@ -5581,6 +5895,12 @@ notebook = [ ] paper = [ { name = "glissade" }, + { name = "matplotlib-venn" }, + { name = "shap", version = "0.49.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "shap", version = "0.51.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "shap", version = "0.52.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "xgboost", version = "3.2.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "xgboost", version = "3.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] [package.metadata] @@ -5622,7 +5942,12 @@ notebook = [ { name = "seaborn", specifier = ">=0.13.2" }, { name = "tqdm", specifier = ">=4.67.1" }, ] -paper = [{ name = "glissade", git = "https://github.com/JemmaLDaniel/glissade.git?rev=7c723a2af4a88fda84a6bd4f223b351179bd36da" }] +paper = [ + { name = "glissade", git = "https://github.com/JemmaLDaniel/glissade.git?rev=7c723a2af4a88fda84a6bd4f223b351179bd36da" }, + { name = "matplotlib-venn", specifier = ">=1.1.1" }, + { name = "shap", specifier = ">=0.44.0" }, + { name = "xgboost", specifier = ">=3.2.0" }, +] [[package]] name = "wrapt" @@ -5688,6 +6013,69 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6e/d2/6317eb6d4554855bbf12d61857774af34747bf88a42c19bf306de67e2fa3/wrapt-2.2.2-py3-none-any.whl", hash = "sha256:5bad217350f19ce99ca5b5e71d406765ea86fe541628426772b657375ee1c048", size = 61460, upload-time = "2026-06-20T23:49:42.966Z" }, ] +[[package]] +name = "xgboost" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'linux') or (python_full_version > '3.11' and python_full_version < '3.12' and sys_platform == 'win32')", + "(python_full_version == '3.11' and sys_platform == 'linux') or (python_full_version == '3.11' and sys_platform == 'win32')", + "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version > '3.11' and python_full_version < '3.12' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'x86_64') or (python_full_version < '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform != 'darwin')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version == '3.11.*' and platform_machine != 'x86_64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')" }, + { name = "nvidia-nccl-cu12", marker = "python_full_version < '3.12' and sys_platform == 'linux'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/bb/1eb0242409d22db725d7a88088e6cfd6556829fb0736f9ff69aa9f1e9455/xgboost-3.2.0.tar.gz", hash = "sha256:99b0e9a2a64896cdaf509c5e46372d336c692406646d20f2af505003c0c5d70d", size = 1263936, upload-time = "2026-02-10T11:03:05.542Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2d/49/6e4cdd877c24adf56cb3586bc96d93d4dcd780b5ea1efb32e1ee0de08bae/xgboost-3.2.0-py3-none-macosx_10_15_x86_64.whl", hash = "sha256:2f661966d3e322536d9c448090a870fcba1e32ee5760c10b7c46bac7a342079a", size = 2507014, upload-time = "2026-02-10T10:50:57.44Z" }, + { url = "https://files.pythonhosted.org/packages/93/f1/c09ef1add609453aa3ba5bafcd0d1c1a805c1263c0b60138ec968f8ec296/xgboost-3.2.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:eabbd40d474b8dbf6cb3536325f9150b9e6f0db32d18de9914fb3227d0bef5b7", size = 2328527, upload-time = "2026-02-10T10:51:17.502Z" }, + { url = "https://files.pythonhosted.org/packages/96/9f/d9914a7b8df842832850b1a18e5f47aaa071c217cdd1da2ae9deb291018b/xgboost-3.2.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:852eabc6d3b3702a59bf78dbfdcd1cb9c4d3a3b6e5ed1f8781d8b9512354fdd2", size = 131100954, upload-time = "2026-02-10T11:02:42.704Z" }, + { url = "https://files.pythonhosted.org/packages/79/98/679de17c2caa4fd3b0b4386ecf7377301702cb0afb22930a07c142fcb1d8/xgboost-3.2.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:99b4a6bbcb47212fec5cf5fbe12347215f073c08967431b0122cfbd1ee70312c", size = 131748579, upload-time = "2026-02-10T10:54:40.424Z" }, + { url = "https://files.pythonhosted.org/packages/1f/3d/1661dd114a914a67e3f7ab66fa1382e7599c2a8c340f314ad30a3e2b4d08/xgboost-3.2.0-py3-none-win_amd64.whl", hash = "sha256:0d169736fd836fc13646c7ab787167b3a8110351c2c6bc770c755ee1618f0442", size = 101681668, upload-time = "2026-02-10T10:59:31.202Z" }, +] + +[[package]] +name = "xgboost" +version = "3.4.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version >= '3.13' and sys_platform == 'linux') or (python_full_version >= '3.13' and sys_platform == 'win32')", + "(python_full_version == '3.12.*' and sys_platform == 'linux') or (python_full_version == '3.12.*' and sys_platform == 'win32')", + "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'darwin'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and platform_machine != 'x86_64') or (python_full_version >= '3.12' and sys_platform != 'darwin')" }, + { name = "nvidia-nccl-cu13", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/a9/295320f741c5be4be996c73ee65a2a11852028c50daa7229adb0d61c330b/xgboost-3.4.1.tar.gz", hash = "sha256:6968a4c71efdfa859df0dfcad0d99211c95c28c4ffd6aecff46efff77d18026a", size = 1231819, upload-time = "2026-08-15T08:39:21.197Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/ea/0bdcd374241a86f1986e87e272516f0a70d841c3aa86aa9ca167fb651573/xgboost-3.4.1-py3-none-macosx_10_15_x86_64.whl", hash = "sha256:1ea15f15f661825b6a67d87674fb9604a1abb38dd0d4c5cf0486fc85f5203e83", size = 2541584, upload-time = "2026-08-15T08:38:48.484Z" }, + { url = "https://files.pythonhosted.org/packages/f7/94/e5c37a8972ad780edc1d8459d1931356344ca133f7f99ba9cfda516b5bba/xgboost-3.4.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:a7afd7dbace0951c93aa85ffe046e54bc40893f5b51cd3e7991eb157bf9c7c7c", size = 2365501, upload-time = "2026-08-15T08:38:52.366Z" }, + { url = "https://files.pythonhosted.org/packages/a7/11/4ff1f36ca5c32c642c71c88bec1508ee98b2c3b1e9eb169e8c82de303522/xgboost-3.4.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7faaf99de26719c22bfae883a02bd56b5a3c2203122616e563cc72b7191b5c96", size = 57196172, upload-time = "2026-08-15T08:39:03.288Z" }, + { url = "https://files.pythonhosted.org/packages/99/c7/bd05c5c430feb347aa040fcc8870135d70b256718deee9bc7d2ca74a77ff/xgboost-3.4.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:6adf2afa396da2ae8ed30295b50b99d4712eed9a6e0ce6cfe069290e4335e51f", size = 57615456, upload-time = "2026-08-15T08:39:09.983Z" }, + { url = "https://files.pythonhosted.org/packages/2f/3c/925394671f6a1668e2a71886de66e80be694eaf37f615cec74eefaf43107/xgboost-3.4.1-py3-none-win_amd64.whl", hash = "sha256:2d30fa513673101f542fdcbd18f30c8f96c064046f798635ac08663e9969f81b", size = 48942686, upload-time = "2026-08-15T08:39:16.182Z" }, + { url = "https://files.pythonhosted.org/packages/90/2f/f2fbe984ca095709fd246546125e78834740f347e3aa7561a22a1e928510/xgboost-3.4.1-py3-none-win_arm64.whl", hash = "sha256:e9312b30e5679d27c1d8b9ee97e092b964d960a672d5d406d9fb3cd0845c9797", size = 2094178, upload-time = "2026-08-15T08:39:19.308Z" }, +] + [[package]] name = "xxhash" version = "3.8.1" From 82f376692075f2f0508325c6d0b937e758a51305 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:51:14 +0100 Subject: [PATCH 07/26] feat: add paper reproduction scripts --- paper_scripts/analyze_fdr_overlap.py | 1056 ++++++++++ paper_scripts/analyze_features.py | 871 +++++++++ paper_scripts/analyze_novelty.py | 1709 +++++++++++++++++ paper_scripts/analyze_upscored_fps.py | 713 +++++++ paper_scripts/annotate_preds_proteome_hits.py | 337 ++++ paper_scripts/benchmark_runtime.py | 908 +++++++++ paper_scripts/benchmark_scaling.py | 480 +++++ .../calibrator_generalisation_utils.py | 115 ++ paper_scripts/download_figshare_article.py | 365 ++++ .../evaluate_calibrator_generalisation.py | 384 ++++ .../fdr_tool_comparison_preprocess.py | 809 ++++++++ .../fdr_tool_comparison_summaries.py | 538 ++++++ paper_scripts/feature_subsets.py | 74 + paper_scripts/no_prosit_dummy.py | 125 ++ paper_scripts/plot_ablation_summary.py | 674 +++++++ paper_scripts/plot_analysis.py | 1201 ++++++++++++ .../plot_calibrator_generalisation_heatmap.py | 270 +++ paper_scripts/plot_eval_results.py | 995 ++++++++++ paper_scripts/plot_fdr_method_comparison.py | 1121 +++++++++++ paper_scripts/plot_feature_investigation.py | 1349 +++++++++++++ .../run_external_peptide_holdout_benchmark.py | 1160 +++++++++++ paper_scripts/run_feature_ablations.py | 1542 +++++++++++++++ paper_scripts/subset_eval_by_experiment.py | 107 ++ 23 files changed, 16903 insertions(+) create mode 100644 paper_scripts/analyze_fdr_overlap.py create mode 100644 paper_scripts/analyze_features.py create mode 100644 paper_scripts/analyze_novelty.py create mode 100644 paper_scripts/analyze_upscored_fps.py create mode 100644 paper_scripts/annotate_preds_proteome_hits.py create mode 100644 paper_scripts/benchmark_runtime.py create mode 100644 paper_scripts/benchmark_scaling.py create mode 100644 paper_scripts/calibrator_generalisation_utils.py create mode 100644 paper_scripts/download_figshare_article.py create mode 100644 paper_scripts/evaluate_calibrator_generalisation.py create mode 100644 paper_scripts/fdr_tool_comparison_preprocess.py create mode 100644 paper_scripts/fdr_tool_comparison_summaries.py create mode 100644 paper_scripts/feature_subsets.py create mode 100644 paper_scripts/no_prosit_dummy.py create mode 100644 paper_scripts/plot_ablation_summary.py create mode 100644 paper_scripts/plot_analysis.py create mode 100644 paper_scripts/plot_calibrator_generalisation_heatmap.py create mode 100644 paper_scripts/plot_eval_results.py create mode 100644 paper_scripts/plot_fdr_method_comparison.py create mode 100644 paper_scripts/plot_feature_investigation.py create mode 100644 paper_scripts/run_external_peptide_holdout_benchmark.py create mode 100644 paper_scripts/run_feature_ablations.py create mode 100644 paper_scripts/subset_eval_by_experiment.py diff --git a/paper_scripts/analyze_fdr_overlap.py b/paper_scripts/analyze_fdr_overlap.py new file mode 100644 index 00000000..fd433c7b --- /dev/null +++ b/paper_scripts/analyze_fdr_overlap.py @@ -0,0 +1,1056 @@ +#!/usr/bin/env python3 +"""Post-FDR overlap analysis: Winnow-filtered identifications vs database search. + +For each project (see ``plot_eval_results.py`` CLI pattern), at 1 %, 5 %, and 10 % +nominal FDR: + + * Count retained PSMs / unique peptides vs database-search reference peptides at + the same nominal FDR (Winnow: non-parametric on calibrated confidence; + database search: database-grounded on raw confidence). + * Match rule: exact ProForma sequence after I/L equivalence; PTM differences + are not a match. + * Categorise discordant calls vs the database peptide set (PTM candidate, + single-AA variant, near-miss edit distance 2-3, fully discordant). + * Full-search Venns: Winnow (non-parametric calibrated confidence) vs database + search unique peptides at the same nominal FDR (database-grounded raw confidence). + * Violin plots comparing database-matched vs fully novel retained PSMs at + matched FDR (same retention rules as overlap summaries). + +Inputs are ``winnow predict`` output folders arranged as subdirectories under two +roots: an **unlabelled** tree (full-search Winnow predictions) and a **labelled** +tree (database-search reference with ``sequence``). Each project folder (flat or +``PXD*//`` nested) must contain ``preds_and_fdr_metrics.csv``; +``metadata.csv`` is merged when present for violin plots. +""" + +from __future__ import annotations + +import json +import logging +import re +from collections import defaultdict +from pathlib import Path +from typing import Annotated, Callable + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer +from matplotlib_venn import venn2 +from rich.logging import RichHandler + +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl +from winnow.fdr.nonparametric import NonParametricFDRControl + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_NOVEL_COLOUR = _PALETTE[2] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)-?") +_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") + +_PXD_ACCESSION_PREFIX = "PXD" + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] +_DB_GROUNDED_DROP = 10 +RAW_CONFIDENCE_COL = "confidence" +DB_Q_VALUE_COL = "db_psm_q_value" + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", + "PXD004732": "ProteomeTools-1", +} + +_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") +_UNLABELLED_FOLDER_SUFFIXES = ("_raw", "_unlabelled") +_LABELLED_FOLDER_SUFFIXES = ("_annotated", "_labelled") + +NOVEL_FEATURE_COLUMNS: list[tuple[str, str]] = [ + ("spectral_angle", "Spectral angle"), + ("ion_matches", "Ion match rate"), + ("ion_match_intensity", "Ion match intensity"), + ("precursor_charge", "Precursor charge"), + ("mass_error_da", "Precursor mass error (Da)"), + ("irt_error", "iRT error"), + ("confidence", "Raw confidence"), + ("margin", "Beam margin"), +] + +_DISCORDANCE_COUNT_COLS = [ + "n_ptm_candidate", + "n_single_aa_variant", + "n_near_miss_edit_dist", + "n_fully_discordant", +] + +_MIN_VIOLIN_GROUP_SIZE = 5 +_MAX_VIOLIN_PSMs_PER_GROUP = 5000 +_VIOLIN_SUBSAMPLE_SEED = 42 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _display_name(key: str) -> str: + return DATASET_DISPLAY_NAMES.get(key, key) + + +def _project_key_from_folder(folder_name: str) -> str: + """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" + for suffix in _FOLDER_SUFFIXES: + if folder_name.endswith(suffix): + return folder_name[: -len(suffix)] + return folder_name + + +def _search_space_tag_from_folder(folder_name: str) -> str: + """Infer eval-type label for tables/plots from the unlabelled subfolder name.""" + for suffix in _UNLABELLED_FOLDER_SUFFIXES: + if folder_name.endswith(suffix): + return suffix[1:] # raw | unlabelled + return "full_search" + + +def _eval_type_display(eval_type: str) -> str: + """Human-readable eval-type label for plot titles.""" + return { + "full_search": "full search space", + "raw": "raw", + "unlabelled": "unlabelled", + }.get(eval_type, eval_type) + + +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _sequence_match_key(seq: str) -> str: + """Exact match key: ProForma with mods preserved, I/L equivalent.""" + if not seq or not isinstance(seq, str): + return "" + return seq.replace("I", "L") + + +def _strip_mods(seq: str) -> str: + """Strip PTM annotations and normalise I -> L (discordance subtyping only).""" + if not seq or not isinstance(seq, str): + return "" + s = _MOD_PLUS.sub("", seq) + s = _MOD_UNIMOD.sub("", s) + return s.replace("I", "L") + + +def _levenshtein(s: str, t: str) -> int: + n, m = len(s), len(t) + if n == 0: + return m + if m == 0: + return n + prev = list(range(m + 1)) + for i in range(1, n + 1): + curr = [i] + [0] * m + for j in range(1, m + 1): + cost = 0 if s[i - 1] == t[j - 1] else 1 + curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost) + prev = curr + return prev[m] + + +def _db_stripped_by_length(db_stripped_list: list[str]) -> dict[int, list[str]]: + by_len: dict[int, list[str]] = defaultdict(list) + for s in db_stripped_list: + by_len[len(s)].append(s) + return by_len + + +def _min_edit_distance_to_db( + pred_stripped: str, db_stripped_by_len: dict[int, list[str]] +) -> int: + if not pred_stripped: + return 999 + plen = len(pred_stripped) + best = 999 + for length in range(max(0, plen - 3), plen + 4): + for db in db_stripped_by_len.get(length, []): + d = _levenshtein(pred_stripped, db) + if d < best: + best = d + if best == 0: + return 0 + return best + + +def _build_db_reference_sets( + db_df: pd.DataFrame, +) -> tuple[set[str], set[str], list[str], dict[int, list[str]]]: + sequences = db_df["sequence"].dropna().astype(str) + db_keys = {_sequence_match_key(s) for s in sequences if _sequence_match_key(s)} + db_stripped = [_strip_mods(s) for s in sequences if _strip_mods(s)] + db_stripped_set = set(db_stripped) + db_stripped_unique = list(dict.fromkeys(db_stripped)) + return ( + db_keys, + db_stripped_set, + db_stripped_unique, + _db_stripped_by_length(db_stripped_unique), + ) + + +def _is_db_match(pred: str, db_keys: set[str]) -> bool: + return _sequence_match_key(pred) in db_keys + + +# --------------------------------------------------------------------------- +# Discordance classification +# --------------------------------------------------------------------------- +def classify_discordance_vs_db( + pred_str: str, + db_keys: set[str], + db_stripped_set: set[str], + db_stripped_by_len: dict[int, list[str]], +) -> str: + """Classify a discordant full-search PSM vs the database peptide reference set.""" + pred_key = _sequence_match_key(pred_str) + if pred_key in db_keys: + return "db_match" + + pred_stripped = _strip_mods(pred_str) + if pred_stripped in db_stripped_set: + return "ptm_candidate" + + ed = _min_edit_distance_to_db(pred_stripped, db_stripped_by_len) + if ed == 1: + return "single_aa_variant" + if ed in (2, 3): + return "near_miss_edit_dist" + return "fully_discordant" + + +def _build_discordance_cache( + predictions: pd.Series, + db_keys: set[str], + db_stripped_set: set[str], + db_stripped_by_len: dict[int, list[str]], +) -> dict[str, str]: + """Classify each unique discordant prediction string once.""" + cache: dict[str, str] = {} + for pred in predictions.dropna().unique(): + key = str(pred) + if _is_db_match(key, db_keys): + cache[key] = "db_match" + elif key not in cache: + cache[key] = classify_discordance_vs_db( + key, db_keys, db_stripped_set, db_stripped_by_len + ) + return cache + + +def _classify_predictions_vs_db( + predictions: pd.Series, + cache: dict[str, str], +) -> pd.Series: + return predictions.map(lambda p: cache.get(str(p), "fully_discordant")) + + +# --------------------------------------------------------------------------- +# Data loading +# --------------------------------------------------------------------------- +def _preds_header(folder: Path) -> set[str] | None: + preds_path = folder / "preds_and_fdr_metrics.csv" + if not preds_path.is_file(): + return None + return set(pd.read_csv(preds_path, nrows=0).columns.tolist()) + + +def _is_unlabelled_preds_folder(folder: Path) -> bool: + """Return True for full-search / proteome-hit prediction folders. + + Figshare ``general_results/full/`` trees may still carry a ``sequence`` + column. Prefer ``proteome_hit`` without labelled ``correct`` when present; + otherwise require ``sequence`` to be absent. + """ + header = _preds_header(folder) + if header is None: + return False + if "proteome_hit" in header and "correct" not in header: + return True + return "sequence" not in header + + +def _is_labelled_preds_folder(folder: Path) -> bool: + """Return True for database-labelled prediction folders with ``sequence``.""" + header = _preds_header(folder) + if header is None: + return False + if "correct" in header and "sequence" in header: + return True + return "sequence" in header and "proteome_hit" not in header + + +def _load_from_folder(folder: Path) -> pd.DataFrame: + """Load preds_and_fdr_metrics.csv merged with metadata.csv from a project folder.""" + preds_path = folder / "preds_and_fdr_metrics.csv" + if not preds_path.is_file(): + raise FileNotFoundError(f"Missing predictions file: {preds_path}") + + preds_df = pd.read_csv(preds_path) + meta_path = folder / "metadata.csv" + if meta_path.is_file(): + meta_df = pd.read_csv(meta_path) + overlap_cols = [ + c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" + ] + if overlap_cols: + meta_df = meta_df.drop(columns=overlap_cols) + return preds_df.merge(meta_df, on="spectrum_id", how="left") + return preds_df + + +def _register_project_folder(projects: dict[str, Path], key: str, folder: Path) -> None: + """Register *folder* under *key*, warning on duplicate keys.""" + if key in projects: + logger.warning( + "Duplicate project key %r: %s and %s", + key, + projects[key], + folder, + ) + return + projects[key] = folder + + +def _discover_project_folders( + root: Path, + *, + is_preds_folder: Callable[[Path], bool], +) -> dict[str, Path]: + """Map project key -> preds folder under *root* (flat or ``PXD*//``).""" + projects: dict[str, Path] = {} + if not root.is_dir(): + return projects + + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if is_preds_folder(child): + _register_project_folder( + projects, _project_key_from_folder(child.name), child + ) + continue + if not child.name.startswith(_PXD_ACCESSION_PREFIX): + continue + for run_dir in sorted(child.iterdir()): + if run_dir.is_dir() and is_preds_folder(run_dir): + _register_project_folder(projects, run_dir.name, run_dir) + return projects + + +def _discover_unlabelled_folders(root: Path) -> dict[str, Path]: + """Map project key -> full-search folder under ``root``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts (``{root}/PXD004452//``) used by new eval sets. + """ + return _discover_project_folders(root, is_preds_folder=_is_unlabelled_preds_folder) + + +def _discover_labelled_folders(root: Path) -> dict[str, Path]: + """Map project key -> database-reference folder under ``root``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts (``{root}/PXD004452//``) used by new eval sets. + """ + return _discover_project_folders(root, is_preds_folder=_is_labelled_preds_folder) + + +def discover_project_pairs( + unlabelled_dir: Path, + labelled_dir: Path, + *, + projects_filter: set[str] | None = None, +) -> list[tuple[str, Path, Path, str]]: + """Return ``(project, unlabelled_folder, labelled_folder, search_space_tag)`` pairs.""" + unlabelled = _discover_unlabelled_folders(unlabelled_dir) + labelled = _discover_labelled_folders(labelled_dir) + + keys = sorted(unlabelled.keys() & labelled.keys()) + if projects_filter is not None: + keys = [k for k in keys if k in projects_filter] + + pairs: list[tuple[str, Path, Path, str]] = [] + for key in keys: + pairs.append( + ( + key, + unlabelled[key], + labelled[key], + _search_space_tag_from_folder(unlabelled[key].name), + ) + ) + + for key in sorted(unlabelled.keys() - labelled.keys()): + if projects_filter is None or key in projects_filter: + logger.warning("No labelled folder for unlabelled project %r", key) + for key in sorted(labelled.keys() - unlabelled.keys()): + if projects_filter is None or key in projects_filter: + logger.warning("No unlabelled folder for labelled project %r", key) + + return pairs + + +def _effective_db_grounded_drop(n_rows: int, drop: int = _DB_GROUNDED_DROP) -> int: + return min(drop, max(0, n_rows - 1)) + + +def _add_q_values( + df: pd.DataFrame, + conf_col: str = "calibrated_confidence", + *, + q_col: str = "psm_q_value", +) -> pd.DataFrame: + """Attach PSM q-values from a non-parametric FDR fit on *conf_col*.""" + if q_col in df.columns: + return df + if conf_col not in df.columns: + raise ValueError(f"Missing confidence column {conf_col!r}") + + existing_q = df["psm_q_value"] if "psm_q_value" in df.columns else None + work = df.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") + + fdr = NonParametricFDRControl() + fdr.fit(dataset=work[conf_col]) + out = fdr.add_psm_q_value(work, confidence_col=conf_col) + if q_col != "psm_q_value": + out = out.rename(columns={"psm_q_value": q_col}) + if existing_q is not None and q_col != "psm_q_value": + out["psm_q_value"] = existing_q + return out + + +def _add_database_grounded_q_values( + df: pd.DataFrame, + confidence_col: str = RAW_CONFIDENCE_COL, + *, + q_col: str = DB_Q_VALUE_COL, + correct_col: str = "correct", +) -> pd.DataFrame: + """Attach PSM q-values from database-grounded FDR on *confidence_col*.""" + if q_col in df.columns: + return df + if confidence_col not in df.columns: + raise ValueError(f"Missing confidence column {confidence_col!r}") + if "sequence" not in df.columns or "prediction" not in df.columns: + raise ValueError( + "Database-grounded FDR requires 'sequence' and 'prediction' columns" + ) + + # Drop Winnow NP q-values before fitting; merge-based add_psm_q_value can + # leave duplicate psm_q_value_* columns and skip the rename to db_psm_q_value. + work = df.drop( + columns=[q_col, "psm_q_value", "psm_fdr", "fdr"], + errors="ignore", + ).copy() + drop = _effective_db_grounded_drop(len(work)) + ctrl = DatabaseGroundedFDRControl( + confidence_feature=confidence_col, + drop=drop, + ) + # DataFrame path: mirror DatabaseGroundedFDRControl.fit without CalibrationDataset. + sorted_df = work.sort_values(confidence_col, ascending=False) + labels = sorted_df[correct_col].astype(float).to_numpy() + conf = sorted_df[confidence_col].to_numpy() + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + ctrl._fdr_values = np.array(1.0 - precision)[drop:] + ctrl._confidence_scores = conf[drop:] + q_df = ctrl.add_psm_q_value( + work[[confidence_col]].copy(), confidence_col=confidence_col + ) + work[q_col] = q_df["psm_q_value"].values + return work + + +def _unique_peptides_at_fdr( + df: pd.DataFrame, + sequence_col: str, + q_col: str, + fdr_t: float, +) -> set[str]: + retained = df[df[q_col] <= fdr_t] + return set(retained[sequence_col].dropna().map(_sequence_match_key)) - {""} + + +def _empty_overlap_row( + project: str, + eval_type: str, + fdr_t: float, + n_db_peptides: int, +) -> dict: + row: dict = { + "project": project, + "eval_type": eval_type, + "fdr_threshold": fdr_t, + "n_psms_retained": 0, + "n_unique_peptides_retained": 0, + "n_db_search_peptides": n_db_peptides, + "n_matching": 0, + "pct_matching": 0.0, + "n_discordant": 0, + "pct_discordant": 0.0, + } + for col in _DISCORDANCE_COUNT_COLS: + row[col] = 0 + return row + + +def _discordance_cache_for_fdr_retained( + df: pd.DataFrame, + db_keys: set[str], + db_stripped_set: set[str], + db_stripped_by_len: dict[int, list[str]], +) -> dict[str, str]: + """Build discordance lookup for all predictions retained at any FDR threshold.""" + df = _add_q_values(df) + retained = df[df["psm_q_value"] <= max(FDR_THRESHOLDS)] + return _build_discordance_cache( + retained["prediction"], db_keys, db_stripped_set, db_stripped_by_len + ) + + +def compute_overlap_table( + df: pd.DataFrame, + project: str, + eval_type: str, + db_df: pd.DataFrame, + discordance_cache: dict[str, str], +) -> pd.DataFrame: + """Overlap summary at each FDR threshold. + + Winnow uses non-parametric FDR on calibrated confidence (full-search run). + Database reference peptides use database-grounded FDR on raw confidence in the + database-labelled run (see ``plot_full_search_venn``). + """ + df = _add_q_values(df.copy()) + db_scored = _add_database_grounded_q_values( + db_df.copy(), + confidence_col=RAW_CONFIDENCE_COL, + q_col=DB_Q_VALUE_COL, + ) + + rows: list[dict] = [] + for fdr_t in FDR_THRESHOLDS: + db_keys_at_fdr = _unique_peptides_at_fdr( + db_scored, "sequence", DB_Q_VALUE_COL, fdr_t + ) + n_db_peptides = len(db_keys_at_fdr) + + retained = df[df["psm_q_value"] <= fdr_t].copy() + n_retained = len(retained) + if n_retained == 0: + rows.append(_empty_overlap_row(project, eval_type, fdr_t, n_db_peptides)) + continue + + is_match = retained["prediction"].map( + lambda p, keys=db_keys_at_fdr: _is_db_match(str(p), keys) + ) + + n_matching = int(is_match.sum()) + n_discordant = n_retained - n_matching + n_unique_peptides = int( + retained["prediction"].map(_sequence_match_key).replace("", pd.NA).nunique() + ) + + disc = retained[~is_match] + cat_counts: dict[str, int] = {} + if len(disc) > 0: + cats = _classify_predictions_vs_db(disc["prediction"], discordance_cache) + cat_counts = cats.value_counts().to_dict() + + rows.append( + { + "project": project, + "eval_type": eval_type, + "fdr_threshold": fdr_t, + "n_psms_retained": n_retained, + "n_unique_peptides_retained": n_unique_peptides, + "n_db_search_peptides": n_db_peptides, + "n_matching": n_matching, + "pct_matching": round(n_matching / n_retained * 100, 2), + "n_discordant": n_discordant, + "pct_discordant": round(n_discordant / n_retained * 100, 2), + "n_ptm_candidate": cat_counts.get("ptm_candidate", 0), + "n_single_aa_variant": cat_counts.get("single_aa_variant", 0), + "n_near_miss_edit_dist": cat_counts.get("near_miss_edit_dist", 0), + "n_fully_discordant": cat_counts.get("fully_discordant", 0), + } + ) + + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Plots +# --------------------------------------------------------------------------- +def _draw_venn_panel( + ax, + *, + winnow_peptides: set[str], + db_peptides: set[str], + winnow_label: str, + fdr_t: float, +) -> None: + """Draw one FDR-threshold Venn panel onto *ax*.""" + pct = int(fdr_t * 100) + if not winnow_peptides and not db_peptides: + ax.set_title(f"No peptides retained at {pct}% FDR") + ax.axis("off") + return + + if not winnow_peptides or not db_peptides: + missing = "Winnow" if not winnow_peptides else "Database search" + ax.text( + 0.5, + 0.5, + f"No {missing} peptides at {pct}% FDR", + ha="center", + va="center", + transform=ax.transAxes, + ) + ax.set_title(f"{pct}% FDR") + ax.axis("off") + return + + venn2( + [db_peptides, winnow_peptides], + set_labels=("Database search", winnow_label), + set_colors=(_CORRECT_COLOUR, _INCORRECT_COLOUR), + alpha=0.6, + ax=ax, + ) + ax.set_title(f"Unique peptides at {pct}% FDR") + _style_ax(ax) + + +def _plot_venn_panels( + winnow_df: pd.DataFrame, + output_path: Path, + *, + winnow_label: str, + suptitle: str, + db_df: pd.DataFrame, +) -> None: + winnow_scored = _add_q_values(winnow_df.copy()) + db_scored = _add_database_grounded_q_values( + db_df.copy(), + confidence_col=RAW_CONFIDENCE_COL, + q_col=DB_Q_VALUE_COL, + ) + + n_thresholds = len(FDR_THRESHOLDS) + fig, axes = plt.subplots(1, n_thresholds, figsize=(5 * n_thresholds, 5)) + if n_thresholds == 1: + axes = [axes] + + for ax, fdr_t in zip(axes, FDR_THRESHOLDS): + winnow_peptides = _unique_peptides_at_fdr( + winnow_scored, "prediction", "psm_q_value", fdr_t + ) + db_peptides = _unique_peptides_at_fdr( + db_scored, "sequence", DB_Q_VALUE_COL, fdr_t + ) + _draw_venn_panel( + ax, + winnow_peptides=winnow_peptides, + db_peptides=db_peptides, + winnow_label=winnow_label, + fdr_t=fdr_t, + ) + + fig.suptitle(suptitle, fontsize=12) + fig.tight_layout() + _save_fig(fig, output_path) + + +def plot_full_search_venn( + df: pd.DataFrame, + db_df: pd.DataFrame, + project: str, + plots_dir: Path, +) -> None: + """Venn diagrams of FDR-filtered DB vs Winnow full-search unique peptides. + + Database peptides use database-grounded q-values on raw ``confidence`` in the + database-labelled run. Winnow peptides use non-parametric q-values on + ``calibrated_confidence`` in the full-search run. + """ + display = _display_name(project) + _plot_venn_panels( + df, + plots_dir / f"venn_{project}_full_search", + winnow_label="Winnow", + suptitle=f"Database search vs Winnow full search at matched FDR for {display}", + db_df=db_df, + ) + + +def _assign_retained_groups( + retained: pd.DataFrame, + db_keys: set[str], + discordance_cache: dict[str, str], +) -> pd.Series: + is_match = retained["prediction"].map(lambda p: _is_db_match(str(p), db_keys)) + groups = pd.Series("Database match", index=retained.index) + disc_mask = ~is_match + if disc_mask.any(): + groups.loc[disc_mask] = _classify_predictions_vs_db( + retained.loc[disc_mask, "prediction"], + discordance_cache, + ).values + return groups + + +def _subsample_violin_groups(df: pd.DataFrame, category_col: str) -> pd.DataFrame: + """Limit points per category so violin plots stay responsive.""" + parts: list[pd.DataFrame] = [] + for _cat, group in df.groupby(category_col, observed=True): + if len(group) > _MAX_VIOLIN_PSMs_PER_GROUP: + group = group.sample( + n=_MAX_VIOLIN_PSMs_PER_GROUP, + random_state=_VIOLIN_SUBSAMPLE_SEED, + ) + parts.append(group) + return pd.concat(parts, ignore_index=True) if parts else df + + +def _plot_novel_violins_at_fdr( + df: pd.DataFrame, + db_scored: pd.DataFrame, + discordance_cache: dict[str, str], + available: list[tuple[str, str]], + *, + project: str, + display: str, + eval_type: str, + plots_dir: Path, + fdr_t: float, +) -> None: + """Render one FDR-threshold novel-feature violin figure, or skip if too sparse.""" + retained = df[df["psm_q_value"] <= fdr_t].copy() + if len(retained) < _MIN_VIOLIN_GROUP_SIZE: + logger.info( + "%s: skip violins at %d%% FDR (n=%d retained)", + project, + int(fdr_t * 100), + len(retained), + ) + return + + match_keys = _unique_peptides_at_fdr(db_scored, "sequence", DB_Q_VALUE_COL, fdr_t) + + groups = _assign_retained_groups(retained, match_keys, discordance_cache) + retained = retained.assign(_overlap_group=groups) + plot_df = retained[ + retained["_overlap_group"].isin(["Database match", "fully_discordant"]) + ].copy() + plot_df["Category"] = plot_df["_overlap_group"].map( + { + "Database match": "Database match", + "fully_discordant": "Novel", + } + ) + plot_df = _subsample_violin_groups(plot_df, "Category") + + n_match = (plot_df["Category"] == "Database match").sum() + n_novel = (plot_df["Category"] == "Novel").sum() + if n_match < _MIN_VIOLIN_GROUP_SIZE or n_novel < _MIN_VIOLIN_GROUP_SIZE: + logger.info( + "%s: skip violins at %d%% FDR (match=%d, novel=%d)", + project, + int(fdr_t * 100), + n_match, + n_novel, + ) + return + + n_feats = len(available) + n_cols = 4 + n_rows = int(np.ceil(n_feats / n_cols)) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows)) + axes_flat = np.atleast_1d(axes).flatten() + palette = {"Database match": _CORRECT_COLOUR, "Novel": _NOVEL_COLOUR} + cat_order = ["Database match", "Novel"] + + for ax, (col, label) in zip(axes_flat, available): + sub = plot_df[[col, "Category"]].dropna() + if sub["Category"].nunique() < 2: + ax.set_visible(False) + continue + sns.violinplot( + data=sub, + x="Category", + y=col, + order=cat_order, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + ax.set_ylabel(label) + ax.tick_params(axis="x", rotation=15) + _style_ax(ax) + + for ax in axes_flat[len(available) :]: + ax.set_visible(False) + + pct = int(fdr_t * 100) + fig.suptitle( + f"{display} ({_eval_type_display(eval_type)}): " + f"database-matched vs novel features at {pct}% FDR", + fontsize=12, + ) + fig.tight_layout() + _save_fig(fig, plots_dir / f"novel_feature_violins_{project}_fdr{pct}") + + +def plot_novel_feature_violins( + df: pd.DataFrame, + db_df: pd.DataFrame, + discordance_cache: dict[str, str], + project: str, + eval_type: str, + plots_dir: Path, +) -> None: + """Violin plots: database-matched vs fully discordant retained PSMs. + + Winnow retention uses non-parametric FDR on calibrated confidence. For + full-search comparisons, database match uses peptides retained at the same + nominal FDR via database-grounded FDR on raw confidence (see overlap table). + """ + available = [ + (col, label) for col, label in NOVEL_FEATURE_COLUMNS if col in df.columns + ] + if not available: + logger.warning( + "%s: no feature columns for violin plots (metadata merge missing?)", + project, + ) + return + + df = _add_q_values(df.copy()) + display = _display_name(project) + db_scored = _add_database_grounded_q_values( + db_df.copy(), + confidence_col=RAW_CONFIDENCE_COL, + q_col=DB_Q_VALUE_COL, + ) + + for fdr_t in FDR_THRESHOLDS: + _plot_novel_violins_at_fdr( + df, + db_scored, + discordance_cache, + available, + project=project, + display=display, + eval_type=eval_type, + plots_dir=plots_dir, + fdr_t=fdr_t, + ) + + +# --------------------------------------------------------------------------- +# Per-project orchestration +# --------------------------------------------------------------------------- +def generate_all_analyses( + df: pd.DataFrame, + project: str, + eval_type: str, + results_dir: Path, + plots_dir: Path, + db_df: pd.DataFrame, +) -> pd.DataFrame: + """Tables and plots for one project.""" + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + db_keys, db_stripped_set, _, db_by_len = _build_db_reference_sets(db_df) + disc_cache = _discordance_cache_for_fdr_retained( + df, db_keys, db_stripped_set, db_by_len + ) + overlap = compute_overlap_table(df, project, eval_type, db_df, disc_cache) + plot_full_search_venn(df, db_df, project, plots_dir) + plot_novel_feature_violins(df, db_df, disc_cache, project, eval_type, plots_dir) + + overlap.to_csv(results_dir / f"{project}_overlap_summary.csv", index=False) + with open(results_dir / f"{project}_overlap_summary.json", "w") as f: + json.dump(overlap.to_dict(orient="records"), f, indent=2) + + logger.info("\n%s", overlap.to_string(index=False)) + return overlap + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +@app.command() +def main( + unlabelled_dir: Annotated[ + Path, + typer.Option( + "--unlabelled-dir", + help="Root with per-project full-search folders (e.g. gluc_raw/, PXD014877_unlabelled/).", + ), + ], + labelled_dir: Annotated[ + Path, + typer.Option( + "--labelled-dir", + help="Root with per-project database-search folders (e.g. gluc_annotated/, PXD014877_labelled/).", + ), + ], + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for CSV/JSON tables."), + ], + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for png/pdf figures."), + ], + projects: Annotated[ + str | None, + typer.Option( + "--projects", + help="Optional space- or comma-separated project keys to restrict analysis.", + ), + ] = None, +) -> None: + """Post-FDR overlap: full-search Winnow identifications vs database search.""" + logging.basicConfig(level=logging.INFO, format="%(message)s", datefmt="%H:%M:%S") + + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + projects_filter: set[str] | None = None + if projects is not None: + project_list = [ + p.strip() for p in projects.replace(",", " ").split() if p.strip() + ] + if not project_list: + raise typer.BadParameter("No projects specified in --projects.") + projects_filter = set(project_list) + + pairs = discover_project_pairs( + unlabelled_dir, labelled_dir, projects_filter=projects_filter + ) + if not pairs: + logger.error( + "No paired projects found under unlabelled-dir=%s and labelled-dir=%s", + unlabelled_dir, + labelled_dir, + ) + raise typer.Exit(code=1) + + all_tables: list[pd.DataFrame] = [] + for project, unlabelled_folder, labelled_folder, search_tag in pairs: + display = _display_name(project) + logger.info( + "Processing %s (%s): %s vs %s", + project, + display, + unlabelled_folder.name, + labelled_folder.name, + ) + + try: + df = _load_from_folder(unlabelled_folder) + db_df = _load_from_folder(labelled_folder) + except FileNotFoundError as exc: + logger.warning("Skipping %s: %s", project, exc) + continue + + logger.info( + " Full search: %d rows; DB reference: %d rows", len(df), len(db_df) + ) + + try: + table = generate_all_analyses( + df, + project, + search_tag, + results_dir, + plots_dir, + db_df, + ) + all_tables.append(table) + except ValueError as exc: + logger.warning("Skipping %s: %s", project, exc) + + if not all_tables: + logger.error("No projects produced overlap output under %s", results_dir) + raise typer.Exit(code=1) + + combined = pd.concat(all_tables, ignore_index=True) + combined.to_csv(results_dir / "all_projects_overlap_summary.csv", index=False) + with open(results_dir / "all_projects_overlap_summary.json", "w") as f: + json.dump(combined.to_dict(orient="records"), f, indent=2) + + logger.info( + "FDR overlap analysis complete. Tables in %s; plots in %s", + results_dir, + plots_dir, + ) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/analyze_features.py b/paper_scripts/analyze_features.py new file mode 100644 index 00000000..0e34ed56 --- /dev/null +++ b/paper_scripts/analyze_features.py @@ -0,0 +1,871 @@ +"""Analyze feature importance and correlations for a pretrained calibrator. + +This script provides comprehensive analysis of feature importance: + - Permutation importance on test set + - SHAP values with training background on test set + - Feature correlation analysis on training data + - Optional visualization of results +""" + +import logging +import pickle +from pathlib import Path +from typing import Annotated, Any, Dict, List, Optional + +import matplotlib.pyplot as plt +from matplotlib.colors import LinearSegmentedColormap +import numpy as np +import pandas as pd +import seaborn as sns +import shap +import torch +import typer +import yaml +from rich.console import Console +from rich.theme import Theme +from sklearn.inspection import permutation_importance + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.datasets.data_loaders import InstaNovoDatasetLoader + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette + "sunset" diverging colourmap +# --------------------------------------------------------------------------- +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +_SUNSET_COLORS = [ + "#364B9A", + "#4A7BB7", + "#6EA6CD", + "#98CAE1", + "#C2E4EF", + "#EAECCC", + "#FEDA8B", + "#FDB366", + "#F67E4B", + "#DD3D2D", + "#A50026", +] +_BAD_COLOUR = "#FFFFFF" + + +def _sunset_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) + cmap.set_bad(color=_BAD_COLOUR) + return cmap + + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + + +def _style_axes(ax: plt.Axes) -> None: + """Apply standard axes formatting: no grid, black spines.""" + ax.set_axisbelow(True) + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +logger = logging.getLogger("winnow") +logger.setLevel(logging.INFO) + +logging.getLogger("shap").setLevel(logging.WARNING) + +# --------------------------------------------------------------------------- +# Constants — loaded from the canonical Winnow YAML configs +# --------------------------------------------------------------------------- +SEED = 42 + +_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "winnow" / "configs" + +with open(_CONFIGS_DIR / "residues.yaml") as _f: + RESIDUE_MASSES: dict[str, float] = yaml.safe_load(_f)["residue_masses"] + +with open(_CONFIGS_DIR / "data_loader" / "instanovo.yaml") as _f: + _instanovo_cfg = yaml.safe_load(_f) + RESIDUE_REMAPPING: dict[str, str] = _instanovo_cfg.get("residue_remapping", {}) + BEAM_COLUMNS: dict[str, str] | None = _instanovo_cfg.get("beam_columns") + +COLUMN_DISPLAY_NAMES = { + "confidence": "Raw confidence", + "mass_error": "Mass error", + "mass_error_ppm": "Mass error (ppm)", + "mass_error_da": "Mass error (Da)", + "spectral_angle": "Spectral angle", + "ion_matches": "Ion matches", + "ion_match_intensity": "Ion match intensity", + "chimeric_ion_matches": "Chimeric ion matches", + "chimeric_ion_match_intensity": "Chimeric ion match intensity", + "irt_error": "iRT error", + "margin": "Margin", + "median_margin": "Median margin", + "entropy": "Entropy", + "z-score": "Z-score", + "edit_distance": "Edit distance", + "xcorr": "XCorr", + "chimeric_xcorr": "Chimeric XCorr", + "longest_ion_series": "Longest ion series", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Max ion gap", + "chimeric_longest_ion_series": "Chimeric longest ion series", + "chimeric_complementary_ion_count": "Chimeric complementary ion count", + "chimeric_max_ion_gap": "Chimeric max ion gap", + "is_missing_fragment_match_features": "Missing fragment match", + "is_missing_chimeric_features": "Missing chimeric", + "is_missing_irt_error": "Missing iRT", + "sequence_length": "Sequence length", + "precursor_charge": "Precursor charge", + "min_token_probability": "Min token probability", + "std_token_probability": "Std token probability", +} + +error_theme = Theme({"error": "red bold", "error_highlight": "red bold underline"}) +console = Console(theme=error_theme) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +def feature_display_name(column: str) -> str: + """Human-readable label for a feature column.""" + if column in COLUMN_DISPLAY_NAMES: + return COLUMN_DISPLAY_NAMES[column] + return column.replace("_", " ").replace("-", " ").strip().title() + + +def to_sentence_case(name: str) -> str: + """Convert a feature display name to sentence case.""" + return name.lower() + + +_SUPPORTED_EXTENSIONS = {".parquet", ".ipc", ".mgf"} + + +def _load_and_compute_features( + spectrum_path: Path, + predictions_path: Path, + loader: InstaNovoDatasetLoader, + calibrator: ProbabilityCalibrator, +) -> CalibrationDataset: + """Load spectra (single file or directory) and compute calibration features. + + When *spectrum_path* is a directory the contained spectrum files are + processed one at a time and the resulting metadata frames are + concatenated, mirroring the batched logic in + ``winnow.scripts.main._compute_features_directory``. + """ + if spectrum_path.is_dir(): + files = sorted( + f for f in spectrum_path.iterdir() if f.suffix in _SUPPORTED_EXTENSIONS + ) + if not files: + raise FileNotFoundError( + f"No spectrum files found in {spectrum_path}. " + f"Supported extensions: {', '.join(sorted(_SUPPORTED_EXTENSIONS))}" + ) + all_metadata: list[pd.DataFrame] = [] + for file_path in files: + logger.info(" Processing experiment file: %s", file_path.name) + ds = loader.load(data_path=file_path, predictions_path=predictions_path) + calibrator.compute_features(ds) + all_metadata.append(ds.metadata) + combined = pd.concat(all_metadata, ignore_index=True) + return CalibrationDataset(metadata=combined, predictions=None) + + dataset = loader.load(data_path=spectrum_path, predictions_path=predictions_path) + calibrator.compute_features(dataset) + return dataset + + +# --------------------------------------------------------------------------- +# Model wrapper for sklearn-compatible predict_proba +# --------------------------------------------------------------------------- +class _CalibratorPredictor: + """Wraps a fitted ProbabilityCalibrator as an sklearn-style estimator. + + Provides ``predict_proba`` and ``predict`` on *pre-normalised* feature + arrays so that permutation importance and SHAP can treat it like a + classifier. The ``classes_`` attribute is set to ``[0, 1]``. + """ + + def __init__(self, calibrator: ProbabilityCalibrator) -> None: + assert calibrator.network is not None + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + + self.network = calibrator.network + self.feature_mean = calibrator.feature_mean.cpu() + self.feature_std = calibrator.feature_std.cpu() + self.classes_ = np.array([0, 1]) + + def fit(self, x_input: np.ndarray, y: np.ndarray) -> "_CalibratorPredictor": + """No-op fit to satisfy sklearn estimator interface.""" + return self + + def score(self, x_input: np.ndarray, y: np.ndarray) -> float: + """Return accuracy to satisfy sklearn estimator interface.""" + return float(np.mean(self.predict(x_input) == y)) + + def predict_proba(self, x_input: np.ndarray) -> np.ndarray: # noqa: N803 + """Return class probabilities for each sample.""" + x = torch.as_tensor(x_input, dtype=torch.float32) + self.network.eval() + with torch.no_grad(): + logits = self.network(x) + probs = torch.sigmoid(logits).numpy().ravel() + return np.column_stack([1 - probs, probs]) + + def predict(self, x_input: np.ndarray) -> np.ndarray: # noqa: N803 + """Return binary predictions for each sample.""" + return (self.predict_proba(x_input)[:, 1] >= 0.5).astype(int) + + +# --------------------------------------------------------------------------- +# Plotting functions +# --------------------------------------------------------------------------- +def plot_feature_importance( + importance_scores: Dict[str, float], + title: str, + output_path_base: Path, +) -> None: + """Plot horizontal bar chart of permutation feature importance scores.""" + plt.figure(figsize=(8, 6)) + features = list(importance_scores.keys()) + scores = list(importance_scores.values()) + + sorted_idx = np.argsort(scores) + features = [features[i] for i in sorted_idx] + scores = [scores[i] for i in sorted_idx] + + plt.barh(range(len(features)), scores, color=_PALETTE[0]) + plt.yticks(range(len(features)), features) + plt.xlabel("Importance score") + plt.title(title) + _style_axes(plt.gca()) + + plt.savefig(f"{output_path_base}.pdf", bbox_inches="tight", dpi=300) + plt.savefig(f"{output_path_base}.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_feature_correlations(features: pd.DataFrame, output_path_base: Path) -> None: + """Plot lower-triangle feature correlation heatmap.""" + plt.figure(figsize=(12, 10)) + corr_matrix = features.corr() + mask = np.triu(np.ones_like(corr_matrix, dtype=bool)) + + sns.heatmap( + corr_matrix, + mask=mask, + cmap=_sunset_cmap(), + vmin=-1, + vmax=1, + center=0, + square=True, + annot=True, + fmt=".2f", + cbar_kws={"shrink": 0.5}, + ) + plt.title("Feature correlation matrix") + _style_axes(plt.gca()) + + plt.savefig(f"{output_path_base}.pdf", bbox_inches="tight", dpi=300) + plt.savefig(f"{output_path_base}.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_shap_summary(shap_values, correct_class_idx: int, output_dir: Path) -> None: + """Plot SHAP beeswarm summary for the correct class.""" + plt.figure(figsize=(8, 6)) + shap.plots.beeswarm( + shap_values[:, :, correct_class_idx], + show=False, + max_display=12, + color=_sunset_cmap(), + ) + plt.title(r"SHAP feature impact on $P(\text{correct})$") + _style_axes(plt.gca()) + + plt.savefig(output_dir / "shap_summary.pdf", bbox_inches="tight", dpi=300) + plt.savefig(output_dir / "shap_summary.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_shap_bar( + shap_values, + test_features_scaled, + test_labels, + correct_class_idx: int, + output_dir: Path, +) -> None: + """Plot SHAP bar chart with hierarchical clustering.""" + plt.figure(figsize=(8, 6)) + clustering = shap.utils.hclust(test_features_scaled, test_labels) + shap.plots.bar( + shap_values[:, :, correct_class_idx], + clustering=clustering, + show=False, + clustering_cutoff=0.5, + max_display=12, + ) + ax = plt.gca() + for patch in ax.patches: + patch.set_facecolor(_PALETTE[1]) + plt.title(r"SHAP feature importance for $P(\text{correct})$") + _style_axes(plt.gca()) + + plt.savefig(output_dir / "shap_importance.pdf", bbox_inches="tight", dpi=300) + plt.savefig(output_dir / "shap_importance.png", bbox_inches="tight", dpi=300) + plt.close() + + +def plot_shap_dependence( + shap_values, + feature_names: list, + display_feature_names: list, + correct_class_idx: int, + output_dir: Path, + top_n: int = 3, +) -> None: + """Plot SHAP dependence scatter for the top-N most important features.""" + mean_abs_shap = np.abs(shap_values.values[:, :, correct_class_idx]).mean(axis=0) + top_features_idx = np.argsort(mean_abs_shap)[-top_n:][::-1] + + for idx in top_features_idx: + feature_name = display_feature_names[idx] + original_feature_name = feature_names[idx] + plt.figure(figsize=(8, 6)) + shap.plots.scatter( + shap_values[:, idx, correct_class_idx], + show=False, + color=_PALETTE[2], + ) + plt.title( + "SHAP dependence plot for " + + to_sentence_case(feature_name) + + "\n" + + r"(impact on $P(\text{correct})$)" + ) + + ax = plt.gca() + ylabel = ax.get_ylabel() + if "SHAP value for" in ylabel: + ax.set_ylabel(ylabel.replace(feature_name, to_sentence_case(feature_name))) + + _style_axes(ax) + plt.savefig( + output_dir / f"shap_dependence_{original_feature_name}.pdf", + bbox_inches="tight", + dpi=300, + ) + plt.savefig( + output_dir / f"shap_dependence_{original_feature_name}.png", + bbox_inches="tight", + dpi=300, + ) + plt.close() + + +def plot_shap_interactions( + shap_values, + feature_names: list, + display_feature_names: list, + correct_class_idx: int, + output_dir: Path, + top_n: int = 3, +) -> None: + """Plot pairwise SHAP interaction scatter plots for top-N features.""" + mean_abs_shap = np.abs(shap_values.values[:, :, correct_class_idx]).mean(axis=0) + top_features_idx = np.argsort(mean_abs_shap)[-top_n:][::-1] + + for i, idx1 in enumerate(top_features_idx): + f1_display = display_feature_names[idx1] + f1_orig = feature_names[idx1] + for j, idx2 in enumerate(top_features_idx): + if i == j: + continue + f2_display = display_feature_names[idx2] + f2_orig = feature_names[idx2] + + plt.figure(figsize=(8, 6)) + shap.plots.scatter( + shap_values[:, f1_display, correct_class_idx], + color=shap_values[:, f2_display, correct_class_idx], + show=False, + cmap=_sunset_cmap(), + ) + plt.title( + "SHAP interaction plot for " + + to_sentence_case(f1_display) + + " vs " + + to_sentence_case(f2_display) + + "\n" + + r" (impact on $P(\text{correct})$)" + ) + + ax = plt.gca() + ylabel = ax.get_ylabel() + if "SHAP value for" in ylabel: + ax.set_ylabel(ylabel.replace(f1_display, to_sentence_case(f1_display))) + + _style_axes(ax) + plt.savefig( + output_dir / f"shap_interaction_{f1_orig}_vs_{f2_orig}.pdf", + bbox_inches="tight", + dpi=300, + ) + plt.savefig( + output_dir / f"shap_interaction_{f1_orig}_vs_{f2_orig}.png", + bbox_inches="tight", + dpi=300, + ) + plt.close() + + +def plot_shap_heatmap(shap_values, correct_class_idx: int, output_dir: Path) -> None: + """Plot SHAP heatmap showing per-sample feature contributions.""" + plt.figure(figsize=(8, 6)) + shap.plots.heatmap( + shap_values[:, :, correct_class_idx], + max_display=12, + show=False, + cmap=_sunset_cmap(), + ) + plt.title("SHAP feature impact heatmap\n" + r"(impact on $P(\text{correct})$)") + _style_axes(plt.gca()) + + plt.savefig(output_dir / "shap_heatmap.pdf", bbox_inches="tight", dpi=300) + plt.savefig(output_dir / "shap_heatmap.png", bbox_inches="tight", dpi=300) + plt.close() + + +def create_all_plots( + perm_importance_dict: Dict[str, float], + shap_values, + train_features_scaled_df: pd.DataFrame, + test_features_scaled: np.ndarray, + test_labels: np.ndarray, + feature_names: list, + display_feature_names: list, + correct_class_idx: int, + output_dir: Path, +) -> None: + """Generate all analysis plots (importance, correlations, SHAP).""" + logger.info("Creating plots...") + + plot_feature_importance( + perm_importance_dict, + "Permutation feature importance", + output_dir / "permutation_importance", + ) + plot_shap_summary(shap_values, correct_class_idx, output_dir) + plot_shap_bar( + shap_values, test_features_scaled, test_labels, correct_class_idx, output_dir + ) + plot_shap_dependence( + shap_values, feature_names, display_feature_names, correct_class_idx, output_dir + ) + plot_shap_interactions( + shap_values, feature_names, display_feature_names, correct_class_idx, output_dir + ) + plot_shap_heatmap(shap_values, correct_class_idx, output_dir) + plot_feature_correlations( + train_features_scaled_df, output_dir / "feature_correlations" + ) + + +def replot_from_pickles( + results_dir: Path, + output_dir: Path, + *, + correct_class_idx: int = 1, +) -> None: + """Regenerate plots that only need deposited ``*.pkl`` artefacts. + + Loads ``perm_importance.pkl`` and ``shap_values.pkl`` from ``results_dir`` + (e.g. Figshare ``feature_importance/PXD014877/``). Writes permutation + importance, SHAP summary / dependence / interactions / heatmap under + ``output_dir``. + + Skips the SHAP bar plot (needs scaled test features for clustering) and + the correlation matrix (needs the training feature matrix); neither is in + the Figshare deposit. + """ + perm_path = results_dir / "perm_importance.pkl" + shap_path = results_dir / "shap_values.pkl" + missing = [str(p) for p in (perm_path, shap_path) if not p.is_file()] + if missing: + raise FileNotFoundError( + "Replot requires both pickles under " + f"{results_dir}: missing {', '.join(missing)}" + ) + + output_dir.mkdir(parents=True, exist_ok=True) + + logger.info("Loading permutation importance from %s", perm_path) + with open(perm_path, "rb") as handle: + perm_importance = pickle.load(handle) + + logger.info("Loading SHAP values from %s", shap_path) + with open(shap_path, "rb") as handle: + shap_values = pickle.load(handle) + + display_names = list(shap_values.feature_names) + if len(display_names) != len(perm_importance.importances_mean): + raise ValueError( + "Feature count mismatch between shap_values.feature_names " + f"({len(display_names)}) and perm_importance.importances_mean " + f"({len(perm_importance.importances_mean)})" + ) + + perm_importance_dict = dict(zip(display_names, perm_importance.importances_mean)) + + logger.info("Replotting from pickles (skipping SHAP bar and feature correlations)") + plot_feature_importance( + perm_importance_dict, + "Permutation feature importance", + output_dir / "permutation_importance", + ) + plot_shap_summary(shap_values, correct_class_idx, output_dir) + plot_shap_dependence( + shap_values, display_names, display_names, correct_class_idx, output_dir + ) + plot_shap_interactions( + shap_values, display_names, display_names, correct_class_idx, output_dir + ) + plot_shap_heatmap(shap_values, correct_class_idx, output_dir) + logger.info("Replot complete; plots written to %s", output_dir) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _load_features_from_parquet( + path: Path, + feature_columns: list[str], +) -> tuple[np.ndarray, np.ndarray]: + """Load feature matrix and labels from a Parquet file or directory.""" + import polars as pl + + p = Path(path) + if p.is_dir(): + parquet_files = sorted(p.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError(f"No .parquet files found in {p}") + df = pl.concat([pl.read_parquet(f) for f in parquet_files]) + else: + df = pl.read_parquet(p) + + if "correct" not in df.columns: + raise ValueError(f"Parquet at {path} must contain a 'correct' column") + missing = [c for c in feature_columns if c not in df.columns] + if missing: + raise ValueError(f"Missing feature columns in Parquet: {missing}") + + features = df.select(feature_columns).to_numpy().astype(np.float32) + labels = df["correct"].to_numpy().astype(np.float32) + return features, labels + + +def _parse_koina_constants(raw: Optional[List[str]]) -> Optional[Dict[str, Any]]: + """Parse ``KEY=VALUE`` pairs into a dict, casting numeric strings.""" + if not raw: + return None + out: Dict[str, Any] = {} + for item in raw: + if "=" not in item: + raise typer.BadParameter( + f"Invalid --koina-input-constant format: '{item}'. Expected KEY=VALUE." + ) + key, value = item.split("=", 1) + try: + out[key] = int(value) + except ValueError: + try: + out[key] = float(value) + except ValueError: + out[key] = value + return out + + +@app.command() +def main( + output_dir: Annotated[ + Path, typer.Option(help="Directory to save analysis results and plots.") + ], + model_path: Annotated[ + Optional[Path], + typer.Option( + help=( + "Path to pretrained calibrator model directory " + "(required unless --replot-dir is set)." + ), + ), + ] = None, + replot_dir: Annotated[ + Optional[Path], + typer.Option( + "--replot-dir", + help=( + "Directory with deposited perm_importance.pkl and shap_values.pkl " + "(e.g. paper_data/feature_importance/PXD014877). Skips recomputation " + "and regenerates plots that do not need feature matrices." + ), + ), + ] = None, + data_dir: Annotated[ + Optional[Path], + typer.Option( + help="Directory containing train and test data files (raw spectra path)." + ), + ] = None, + train_features_path: Annotated[ + Optional[Path], + typer.Option( + help="Path to pre-computed training feature Parquet (alternative to --data-dir)." + ), + ] = None, + test_features_path: Annotated[ + Optional[Path], + typer.Option( + help="Path to pre-computed test feature Parquet (alternative to --data-dir)." + ), + ] = None, + train_spectra: Annotated[ + str, typer.Option(help="Filename of training spectra parquet inside data-dir.") + ] = "general_train.parquet", + train_preds: Annotated[ + str, typer.Option(help="Filename of training predictions CSV inside data-dir.") + ] = "general_train_beams.csv", + test_spectra: Annotated[ + str, typer.Option(help="Filename of test spectra parquet inside data-dir.") + ] = "general_test.parquet", + test_preds: Annotated[ + str, typer.Option(help="Filename of test predictions CSV inside data-dir.") + ] = "general_test_beams.csv", + koina_input_constant: Annotated[ + Optional[List[str]], + typer.Option( + help="Koina model input constant as KEY=VALUE (repeatable). " + "E.g. --koina-input-constant collision_energies=27 " + "--koina-input-constant fragmentation_types=HCD", + ), + ] = None, + n_background_samples: Annotated[ + int, typer.Option(help="Background samples for SHAP.", min=1, max=10000) + ] = 500, + n_test_samples: Annotated[ + int, typer.Option(help="Test samples for SHAP.", min=1, max=10000) + ] = 1000, + create_plots: Annotated[ + bool, typer.Option("--create-plots/--no-plots", help="Whether to create plots.") + ] = True, +) -> None: + """Analyze feature importance and correlations for a pretrained calibrator.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + if replot_dir is not None: + replot_from_pickles(replot_dir, output_dir) + return + + if model_path is None: + raise typer.BadParameter("--model-path is required unless --replot-dir is set") + + use_parquet = train_features_path is not None or test_features_path is not None + if use_parquet and (train_features_path is None or test_features_path is None): + raise typer.BadParameter( + "--train-features-path and --test-features-path must both be provided." + ) + if not use_parquet and data_dir is None: + raise typer.BadParameter( + "Either --data-dir or --train-features-path/--test-features-path must be provided." + ) + + output_dir.mkdir(parents=True, exist_ok=True) + + # Load calibrator + logger.info("Loading pretrained calibrator from %s", model_path) + calibrator = ProbabilityCalibrator.load(model_path) + + koina_constants = _parse_koina_constants(koina_input_constant) + if koina_constants: + logger.info("Applying Koina input constant overrides: %s", koina_constants) + calibrator.apply_koina_model_input_overrides( + model_input_constants=koina_constants, + ) + + # Build sklearn-compatible predictor wrapper + predictor = _CalibratorPredictor(calibrator) + + if use_parquet: + assert train_features_path is not None + assert test_features_path is not None + + feature_columns = ["confidence"] + calibrator.columns + + logger.info("Loading training features from Parquet: %s", train_features_path) + train_features, train_labels = _load_features_from_parquet( + train_features_path, + feature_columns, + ) + logger.info( + " %d training samples, %d features", + len(train_labels), + train_features.shape[1], + ) + + logger.info("Loading test features from Parquet: %s", test_features_path) + test_features, test_labels = _load_features_from_parquet( + test_features_path, + feature_columns, + ) + logger.info( + " %d test samples, %d features", len(test_labels), test_features.shape[1] + ) + + feature_names = feature_columns + else: + assert data_dir is not None + loader = InstaNovoDatasetLoader( + residue_masses=RESIDUE_MASSES, + residue_remapping=RESIDUE_REMAPPING, + beam_columns=BEAM_COLUMNS, + add_index_cols=False, + ) + + logger.info("Loading and computing features for training set...") + train_dataset = _load_and_compute_features( + data_dir / train_spectra, + data_dir / train_preds, + loader, + calibrator, + ) + train_features, train_labels = calibrator._extract_feature_matrix( + train_dataset, labelled=True + ) + + logger.info("Loading and computing features for test set...") + test_dataset = _load_and_compute_features( + data_dir / test_spectra, + data_dir / test_preds, + loader, + calibrator, + ) + test_features, test_labels = calibrator._extract_feature_matrix( + test_dataset, labelled=True + ) + + feature_names = [train_dataset.confidence_column] + calibrator.columns + + display_feature_names = [feature_display_name(name) for name in feature_names] + + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + feature_mean = calibrator.feature_mean.cpu().numpy() + feature_std = calibrator.feature_std.cpu().numpy() + train_features_scaled = (train_features - feature_mean) / feature_std + test_features_scaled = (test_features - feature_mean) / feature_std + + correct_class_idx = 1 # class 1 = correct + + # 1. Permutation importance on test set + logger.info("Computing permutation importance on test set...") + perm_importance = permutation_importance( + predictor, + test_features_scaled, + test_labels, + n_repeats=10, + random_state=SEED, + n_jobs=-1, + ) + perm_importance_dict = dict( + zip(display_feature_names, perm_importance.importances_mean) + ) + + # 2. SHAP values + logger.info("Computing SHAP values...") + background = shap.sample( + train_features_scaled, + min(n_background_samples, len(train_features_scaled)), + random_state=SEED, + ) + + explainer = shap.KernelExplainer( + model=predictor.predict_proba, + data=background, + seed=SEED, + link="identity", + ) + + np.random.seed(SEED) + n_samples = min(n_test_samples, test_features_scaled.shape[0]) + indices = np.random.choice( + test_features_scaled.shape[0], size=n_samples, replace=False + ) + + shap_values = explainer(test_features_scaled[indices]) + + # Switch to original feature space for visualisation + shap_values.data = test_features[indices] + shap_values.feature_names = display_feature_names + + # 3. Feature correlations on training data + logger.info("Computing feature correlations on training data...") + train_features_scaled_df = pd.DataFrame( + train_features_scaled, columns=display_feature_names + ) + + if create_plots: + create_all_plots( + perm_importance_dict=perm_importance_dict, + shap_values=shap_values, + train_features_scaled_df=train_features_scaled_df, + test_features_scaled=test_features_scaled, + test_labels=test_labels, + feature_names=feature_names, + display_feature_names=display_feature_names, + correct_class_idx=correct_class_idx, + output_dir=output_dir, + ) + + # Save raw objects + logger.info("Saving raw analysis objects...") + + with open(output_dir / "perm_importance.pkl", "wb") as f: + pickle.dump(perm_importance, f) + + with open(output_dir / "shap_values.pkl", "wb") as f: + pickle.dump(shap_values, f) + + logger.info("Analysis complete!") + logger.info("Results saved to %s", output_dir) + logger.info( + "Permutation Feature Importance: computed on %d test samples", len(test_labels) + ) + logger.info( + "SHAP values: computed on %d test samples with %d training samples as background", + n_samples, + len(background), + ) + logger.info( + "Correlation matrix: computed on %d training samples", len(train_labels) + ) + + saved_files = ["perm_importance.pkl", "shap_values.pkl"] + if create_plots: + saved_files.append("All plots in PDF and PNG formats") + else: + logger.info("Plots were skipped (--no-plots flag used)") + + logger.info("Saved files: %s", ", ".join(saved_files)) + print(f"\nResults saved to {output_dir}") + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/analyze_novelty.py b/paper_scripts/analyze_novelty.py new file mode 100644 index 00000000..b366d65d --- /dev/null +++ b/paper_scripts/analyze_novelty.py @@ -0,0 +1,1709 @@ +#!/usr/bin/env python3 +"""Analyse Winnow calibrator behaviour on out-of-distribution / novel peptides. + +Two analyses demonstrate that the calibrator does not penalise peptides absent +from the standard tryptic database-search training distribution: + +1. **Non-tryptic enzyme digest (``nontryptic_digest`` subcommand)** -- The model + was trained on tryptic data. Enzymes such as GluC, AspN, LysC or chymotrypsin cleave at + non-K/R sites, so retained peptides often have a C-terminus that is *not* + K or R. We classify predictions by whether their C-terminal residue + is tryptic (K/R) or non-tryptic, report terminus proportions before and + after FDR, compare raw InstaNovo versus Winnow calibrated scores, and + quantify calibration shifts (``calibrated_confidence - confidence``). + + Inputs are **full search space** Winnow predictions (acfm / unlabelled eval: + all candidate spectra, not the labelled database-search subset and not + acfm-minus-lcfm). The ``proteome_hit`` column flags predictions whose + stripped sequence occurs in the reference proteome FASTA; plots and tables + label this cohort **full search space**. + + *Non-tryptic* is defined solely by the C-terminal residue of the + mod-stripped, I/L-normalised prediction. N-terminal context is not + checked because positional information is lost in the substring proteome + match. + +2. **ProteomeTools-1 PXD004732 (``proteometools`` subcommand)** -- Synthetic + peptide library. The *lcfm* set contains database-search-confirmed + peptides; the *acfm* set contains all candidates. For each acfm + prediction we check whether it exactly matches, is a subsequence of, or + shares no overlap with any lcfm peptide. Subsequence matches are + validated novel identifications the search engine missed. +""" + +from __future__ import annotations + +import re +import warnings +from pathlib import Path +from typing import Annotated + +import ahocorasick +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import typer +from Bio import SeqIO +from scipy.stats import gaussian_kde + +from winnow.fdr.nonparametric import NonParametricFDRControl + +warnings.filterwarnings("ignore", module="winnow") + +# ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_NOVEL_COLOUR = _PALETTE[2] +_MAIN_LINE_COLOUR = _PALETTE[3] +_RAW_LINE_COLOUR = _PALETTE[5] +_IDEAL_LINE_COLOUR = _PALETTE[6] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)") +_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") +_PROTEOME_JOIN_SEP = "\x1f" + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] + +# Display label for the ``proteome_hit`` cohort (full-search / acfm evaluation). +FULL_SEARCH_SPACE_LABEL = "full search space" +COHORT_FULL_SEARCH = "full_search_space" +COHORT_FULL_SEARCH_AT_FDR = "full_search_space_at_fdr" + +_NONTRYPTIC_CALIBRATION_SCATTER_MAX_POINTS = 10_000 +_NONTRYPTIC_CALIBRATION_SCATTER_RANDOM_STATE = 42 +_C_TERMINUS_ORDER = ("Tryptic (K/R)", "Non-tryptic") +_NONTRYPTIC_GROUP_ALPHA = 0.6 +_TRYPTIC_GROUP_ALPHA = 0.7 + +FEATURE_COLUMNS = [ + "spectral_angle", + "xcorr", + "ion_matches", + "ion_match_intensity", + "irt_error", + "mass_error_ppm", +] + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "PXD004732": "ProteomeTools-1", +} + +app = typer.Typer( + add_completion=False, + no_args_is_help=True, + pretty_exceptions_show_locals=False, +) + + +# ── Shared helpers ──────────────────────────────────────────────────── + + +def _spine_fmt(ax: plt.Axes) -> None: + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _save(fig: plt.Figure, out_dir: Path, name: str) -> None: + base = out_dir / name + fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + print(f" saved {name}") + + +def _subsample_psms( + df: pd.DataFrame, + max_points: int, + random_state: int = _NONTRYPTIC_CALIBRATION_SCATTER_RANDOM_STATE, +) -> pd.DataFrame: + """Return up to ``max_points`` rows without replacement.""" + if len(df) <= max_points: + return df + return df.sample(n=max_points, random_state=random_state) + + +def _strip_mods(seq: str) -> str: + """Strip PTM annotations and normalise I -> L.""" + if not seq or not isinstance(seq, str): + return "" + s = _MOD_PLUS.sub("", seq) + s = _MOD_UNIMOD.sub("", s) + return s.replace("I", "L") + + +def _load_data(predictions_dir: Path) -> pl.DataFrame: + """Load and join ``preds_and_fdr_metrics.csv`` + ``metadata.csv``.""" + preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") + meta_path = predictions_dir / "metadata.csv" + if meta_path.exists(): + meta = pl.read_csv(meta_path) + join_cols = ["spectrum_id"] + [ + c for c in meta.columns if c != "spectrum_id" and c not in preds.columns + ] + if len(join_cols) > 1: + preds = preds.join(meta.select(join_cols), on="spectrum_id", how="inner") + return preds + + +def _add_q_values( + df: pd.DataFrame, + conf_col: str = "calibrated_confidence", +) -> pd.DataFrame: + """Fit non-parametric FDR and append ``psm_q_value`` if missing.""" + if "psm_q_value" in df.columns: + return df + fdr = NonParametricFDRControl() + fdr.fit(dataset=df[conf_col]) + return fdr.add_psm_q_value(df, confidence_col=conf_col) + + +def _load_proteome_haystack(fasta_file: Path) -> str: + """Load a FASTA proteome into a single string for substring matching.""" + parts: list[str] = [] + for record in SeqIO.parse(fasta_file, "fasta"): + s = str(record.seq).replace("I", "L") + if s: + parts.append(s) + return _PROTEOME_JOIN_SEP.join(parts) + + +def _batch_substring_hits( + needles: list[str], + haystack: str, +) -> list[bool]: + """Aho-Corasick batch substring matching.""" + n = len(needles) + out = [False] * n + if not haystack: + return out + + by_needle: dict[str, list[int]] = {} + for i, p in enumerate(needles): + if not p: + continue + by_needle.setdefault(p, []).append(i) + if not by_needle: + return out + + auto = ahocorasick.Automaton() + needle_for_pid: list[str] = [] + for pid, needle in enumerate(by_needle): + auto.add_word(needle, pid) + needle_for_pid.append(needle) + auto.make_automaton() + + matched_pids: set[int] = set() + for _end_idx, pid in auto.iter(haystack): + matched_pids.add(pid) + + for pid in matched_pids: + needle = needle_for_pid[pid] + for row_i in by_needle[needle]: + out[row_i] = True + return out + + +def _nice_feature_label(col: str) -> str: + labels = { + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation", + "ion_matches": "Ion match rate", + "ion_match_intensity": "Ion match intensity", + "irt_error": "iRT prediction error", + "mass_error_ppm": "Mass error (ppm)", + } + return labels.get(col, col.replace("_", " ").capitalize()) + + +# ── Non-tryptic digest analysis ───────────────────────────────────────── + + +def _is_tryptic_cterm(seq: str) -> bool: + """Return True if the mod-stripped C-terminal residue is K or R.""" + stripped = _strip_mods(seq) + if not stripped: + return False + return stripped[-1] in ("K", "R") + + +def _nontryptic_full_search_panel_title(fdr_t: float, n_psms: int) -> str: + """Subplot title for FDR-filtered full-search-space PSMs.""" + pct = int(fdr_t * 100) + return ( + f"{FULL_SEARCH_SPACE_LABEL.title()} identifications at {pct}% FDR\n" + f"(n={n_psms:,})" + ) + + +def _nontryptic_annotate( + df: pl.DataFrame, + fasta_path: Path, +) -> pl.DataFrame: + """Annotate full-search predictions with proteome and tryptic-terminus flags. + + ``proteome_hit`` is True when the mod-stripped prediction is a substring of + the reference proteome FASTA (same rule as ``winnow.utils.proteome``). + """ + haystack = _load_proteome_haystack(fasta_path) + + processed = df["prediction"].map_elements( + lambda x: _strip_mods(x) if isinstance(x, str) else "", + return_dtype=pl.Utf8, + ) + hits = _batch_substring_hits(processed.to_list(), haystack) + tryptic = df["prediction"].map_elements( + lambda x: _is_tryptic_cterm(x) if isinstance(x, str) else False, + return_dtype=pl.Boolean, + ) + return df.with_columns( + pl.Series("proteome_hit", hits, dtype=pl.Boolean), + tryptic.alias("tryptic_cterm"), + ) + + +def _terminus_count_row( + sub: pd.DataFrame, + *, + cohort: str, + fdr_threshold: float | None, +) -> dict: + """Return one row of tryptic / non-tryptic counts for *sub*.""" + n = len(sub) + n_tryp = int(sub["tryptic_cterm"].sum()) if n > 0 else 0 + n_non = n - n_tryp + return { + "cohort": cohort, + "fdr_threshold": fdr_threshold, + "n": n, + "n_tryptic": n_tryp, + "n_non_tryptic": n_non, + "pct_tryptic": round(n_tryp / n * 100, 2) if n > 0 else 0.0, + "pct_non_tryptic": round(n_non / n * 100, 2) if n > 0 else 0.0, + } + + +def _nontryptic_terminus_proportions_table(df: pd.DataFrame) -> pd.DataFrame: + """Tryptic versus non-tryptic counts across cohorts and FDR cutoffs.""" + df = _add_q_values(df) + rows: list[dict] = [ + _terminus_count_row(df, cohort="all_predictions", fdr_threshold=None), + _terminus_count_row( + df[df["proteome_hit"]], + cohort=COHORT_FULL_SEARCH, + fdr_threshold=None, + ), + ] + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t] + rows.append( + _terminus_count_row( + retained, + cohort="retained_at_fdr", + fdr_threshold=fdr_t, + ) + ) + rows.append( + _terminus_count_row( + retained[retained["proteome_hit"]], + cohort=COHORT_FULL_SEARCH_AT_FDR, + fdr_threshold=fdr_t, + ) + ) + return pd.DataFrame(rows) + + +def _nontryptic_calibration_delta_rows( + sub: pd.DataFrame, + *, + cohort: str, + fdr_threshold: float | None, +) -> list[dict]: + """Build calibration-shift summary rows for tryptic and non-tryptic groups.""" + out: list[dict] = [] + for tryptic, label in ((True, "tryptic"), (False, "non_tryptic")): + grp = sub[sub["tryptic_cterm"] == tryptic] + n = len(grp) + out.append( + { + "cohort": cohort, + "fdr_threshold": fdr_threshold, + "terminus_group": label, + "n": n, + "mean_confidence": ( + round(float(grp["confidence"].mean()), 4) if n > 0 else float("nan") + ), + "mean_calibrated_confidence": ( + round(float(grp["calibrated_confidence"].mean()), 4) + if n > 0 + else float("nan") + ), + "mean_delta_confidence": ( + round(float(grp["delta_confidence"].mean()), 4) + if n > 0 + else float("nan") + ), + "median_delta_confidence": ( + round(float(grp["delta_confidence"].median()), 4) + if n > 0 + else float("nan") + ), + } + ) + return out + + +def _nontryptic_calibration_delta_table(df: pd.DataFrame) -> pd.DataFrame: + """Mean calibration shift by terminus and cohort.""" + if "confidence" not in df.columns: + return pd.DataFrame() + + work = _add_q_values(df.copy()) + work["delta_confidence"] = work["calibrated_confidence"] - work["confidence"] + work = work.dropna( + subset=["confidence", "calibrated_confidence", "delta_confidence"] + ) + + rows: list[dict] = [] + rows.extend( + _nontryptic_calibration_delta_rows( + work, cohort="all_predictions", fdr_threshold=None + ) + ) + rows.extend( + _nontryptic_calibration_delta_rows( + work[work["proteome_hit"]], + cohort=COHORT_FULL_SEARCH, + fdr_threshold=None, + ) + ) + for fdr_t in FDR_THRESHOLDS: + retained = work[(work["psm_q_value"] <= fdr_t) & work["proteome_hit"]] + rows.extend( + _nontryptic_calibration_delta_rows( + retained, + cohort=COHORT_FULL_SEARCH_AT_FDR, + fdr_threshold=fdr_t, + ) + ) + return pd.DataFrame(rows) + + +def _nontryptic_summary_table(df: pd.DataFrame) -> pd.DataFrame: + """Build the tryptic-summary table at each FDR threshold.""" + df = _add_q_values(df) + rows: list[dict] = [] + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t] + n_retained = len(retained) + hits = retained[retained["proteome_hit"]] + n_hit = len(hits) + tryptic_hits = hits[hits["tryptic_cterm"]] + non_tryptic_hits = hits[~hits["tryptic_cterm"]] + n_tryp = len(tryptic_hits) + n_non = len(non_tryptic_hits) + rows.append( + { + "fdr_threshold": fdr_t, + "n_retained": n_retained, + "n_full_search_space": n_hit, + "n_tryptic_hit": n_tryp, + "n_non_tryptic_hit": n_non, + "pct_non_tryptic_among_hits": ( + round(n_non / n_hit * 100, 2) if n_hit > 0 else 0.0 + ), + "mean_cal_conf_tryptic": ( + round(float(tryptic_hits["calibrated_confidence"].mean()), 4) + if n_tryp > 0 + else float("nan") + ), + "mean_cal_conf_non_tryptic": ( + round(float(non_tryptic_hits["calibrated_confidence"].mean()), 4) + if n_non > 0 + else float("nan") + ), + } + ) + return pd.DataFrame(rows) + + +def _nontryptic_score_label(score_col: str) -> str: + if score_col == "confidence": + return "Raw InstaNovo confidence" + return "Calibrated confidence" + + +def _finalize_nontryptic_violin_figure( + fig: plt.Figure, + axes: list[plt.Axes], + *, + y_label: str, + suptitle: str, +) -> None: + """Apply shared y-axis label and suptitle after multi-panel violin plots.""" + if len(axes) > 0: + axes[0].set_ylabel(y_label) + fig.tight_layout(rect=[0, 0, 1, 0.92]) + fig.suptitle(suptitle, fontsize=13, y=0.98) + + +def _plot_nontryptic_score_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + score_col: str, + save_name: str, + suptitle: str, +) -> None: + """Violin plot of a score column split by C-terminal residue at each FDR.""" + if score_col not in df.columns: + print(f" skipping {save_name} (missing {score_col})") + return + + df = _add_q_values(df) + y_label = _nontryptic_score_label(score_col) + palette = {"Tryptic (K/R)": _MAIN_LINE_COLOUR, "Non-tryptic": _NOVEL_COLOUR} + + panels: list[tuple[float, pd.DataFrame]] = [] + for fdr_t in FDR_THRESHOLDS: + retained = df[(df["psm_q_value"] <= fdr_t) & df["proteome_hit"]] + retained = retained.dropna(subset=[score_col]) + if len(retained) < 5: + print( + f" skipping {int(fdr_t * 100)}% FDR panel in {save_name} " + f"(n={len(retained):,})" + ) + continue + retained = retained.copy() + retained["C-terminus"] = retained["tryptic_cterm"].map( + {True: "Tryptic (K/R)", False: "Non-tryptic"}, + ) + panels.append((fdr_t, retained)) + + if not panels: + print( + f" skipping {save_name} (no FDR panels with enough " + f"{FULL_SEARCH_SPACE_LABEL} PSMs)" + ) + return + + n_cols = len(panels) + fig, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 5), sharey=True) + if n_cols == 1: + axes = [axes] + + for ax, (fdr_t, retained) in zip(axes, panels): + sns.violinplot( + data=retained, + x="C-terminus", + y=score_col, + order=list(_C_TERMINUS_ORDER), + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) + ax.grid(False) + _spine_fmt(ax) + + _finalize_nontryptic_violin_figure(fig, axes, y_label=y_label, suptitle=suptitle) + _save(fig, out_dir, save_name) + + +def _plot_nontryptic_conf_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Violin plot of calibrated confidence split by C-terminal residue.""" + _plot_nontryptic_score_by_terminus( + df, + out_dir, + score_col="calibrated_confidence", + save_name=f"{prefix}_conf_by_terminus", + suptitle=( + f"Calibrated confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" + "by C-terminal residue" + ), + ) + + +def _plot_nontryptic_raw_conf_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Violin plot of raw InstaNovo confidence split by C-terminal residue.""" + _plot_nontryptic_score_by_terminus( + df, + out_dir, + score_col="confidence", + save_name=f"{prefix}_raw_conf_by_terminus", + suptitle=( + f"Raw InstaNovo confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" + "by C-terminal residue" + ), + ) + + +def _plot_nontryptic_overlapping_score_histogram( + tryp: np.ndarray, + non_tryp: np.ndarray, + *, + score_col: str, + title: str, + out_dir: Path, + save_name: str, +) -> None: + """Overlapping tryptic / non-tryptic histogram with KDE overlays.""" + if len(tryp) + len(non_tryp) < 2: + print(f" skipping {save_name} (too few PSMs)") + return + + fig, ax = plt.subplots(figsize=(7, 5)) + bins = 50 + + ax.hist( + non_tryp, + bins=bins, + alpha=_NONTRYPTIC_GROUP_ALPHA, + label=f"Non-tryptic (n={len(non_tryp):,})", + density=False, + edgecolor="black", + color=_NOVEL_COLOUR, + ) + ax.hist( + tryp, + bins=bins, + alpha=_TRYPTIC_GROUP_ALPHA, + label=f"Tryptic (K/R, n={len(tryp):,})", + density=False, + edgecolor="black", + color=_MAIN_LINE_COLOUR, + ) + + all_vals = np.concatenate([tryp, non_tryp]) if len(non_tryp) else tryp + x_min, x_max = float(all_vals.min()), float(all_vals.max()) + if x_max <= x_min: + x_max = x_min + 1e-6 + x_grid = np.linspace(x_min, x_max, 300) + bin_width = (x_max - x_min) / bins if bins > 1 else 1.0 + + if len(non_tryp) > 1: + y_non = gaussian_kde(non_tryp)(x_grid) * len(non_tryp) * bin_width + ax.plot(x_grid, y_non, color=_NOVEL_COLOUR, lw=1.5) + if len(tryp) > 1: + y_tryp = gaussian_kde(tryp)(x_grid) * len(tryp) * bin_width + ax.plot(x_grid, y_tryp, color=_MAIN_LINE_COLOUR, lw=1.5) + + ax.set_xlabel(_nontryptic_score_label(score_col)) + ax.set_ylabel("Frequency") + ax.set_title(title) + ax.legend(loc="upper center") + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, save_name) + + +def _plot_nontryptic_score_histograms( + df: pd.DataFrame, + out_dir: Path, + *, + score_col: str, + save_name: str, + title: str, + subset: pd.DataFrame | None = None, + min_psms: int = 10, +) -> None: + """Overlapping tryptic / non-tryptic histograms for a score column.""" + if score_col not in df.columns: + print(f" skipping {save_name} (missing {score_col})") + return + + retained = df if subset is None else subset + retained = retained.dropna(subset=[score_col]) + if len(retained) < min_psms: + print(f" skipping {save_name} (too few PSMs)") + return + + tryp = retained.loc[retained["tryptic_cterm"], score_col].to_numpy() + non_tryp = retained.loc[~retained["tryptic_cterm"], score_col].to_numpy() + _plot_nontryptic_overlapping_score_histogram( + tryp, + non_tryp, + score_col=score_col, + title=title, + out_dir=out_dir, + save_name=save_name, + ) + + +def _plot_nontryptic_score_histograms_at_fdr( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Calibrated confidence histogram at 5% FDR for full search space PSMs.""" + df = _add_q_values(df) + retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="calibrated_confidence", + save_name=f"{prefix}_score_histograms", + title=( + f"Calibrated confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} " + "at 5% FDR" + ), + subset=retained, + ) + + +def _plot_nontryptic_raw_score_histograms_at_fdr( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Raw InstaNovo confidence histogram at 5% FDR for full search space PSMs.""" + df = _add_q_values(df) + retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="confidence", + save_name=f"{prefix}_raw_score_histograms", + title=( + f"Raw InstaNovo confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} " + "at 5% FDR" + ), + subset=retained, + ) + + +def _plot_nontryptic_full_score_histograms( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Full-dataset raw and calibrated histograms by C-terminal residue.""" + all_preds = df.dropna(subset=["calibrated_confidence"]) + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="calibrated_confidence", + save_name=f"{prefix}_calibrated_score_histogram_full", + title=( + f"Calibrated confidence for all {dataset_label} predictions\n" + "by C-terminal residue" + ), + subset=all_preds, + min_psms=2, + ) + if "confidence" not in df.columns: + print(f" skipping {prefix}_raw_score_histogram_full (missing confidence)") + return + raw_preds = df.dropna(subset=["confidence"]) + _plot_nontryptic_score_histograms( + df, + out_dir, + score_col="confidence", + save_name=f"{prefix}_raw_score_histogram_full", + title=( + f"Raw InstaNovo confidence for all {dataset_label} predictions\n" + "by C-terminal residue" + ), + subset=raw_preds, + min_psms=2, + ) + + has_raw = "confidence" in df.columns + panels: list[tuple[str, str, pd.DataFrame]] = [ + ( + "calibrated_confidence", + "Calibrated confidence", + all_preds, + ), + ] + if has_raw: + panels.append(("confidence", "Raw InstaNovo confidence", raw_preds)) + + n_cols = len(panels) + fig, axes = plt.subplots(1, n_cols, figsize=(7 * n_cols, 5), sharey=True) + if n_cols == 1: + axes = [axes] + + bins = 50 + for ax, (score_col, y_label, subset) in zip(axes, panels): + tryp = subset.loc[subset["tryptic_cterm"], score_col].to_numpy() + non_tryp = subset.loc[~subset["tryptic_cterm"], score_col].to_numpy() + ax.hist( + non_tryp, + bins=bins, + alpha=_NONTRYPTIC_GROUP_ALPHA, + label=f"Non-tryptic (n={len(non_tryp):,})", + density=False, + edgecolor="black", + color=_NOVEL_COLOUR, + ) + ax.hist( + tryp, + bins=bins, + alpha=_TRYPTIC_GROUP_ALPHA, + label=f"Tryptic (K/R, n={len(tryp):,})", + density=False, + edgecolor="black", + color=_MAIN_LINE_COLOUR, + ) + ax.set_xlabel(y_label) + ax.set_ylabel("Frequency") + ax.set_title(f"All predictions (n={len(subset):,})") + ax.legend(loc="upper center", fontsize=9) + ax.grid(False) + _spine_fmt(ax) + + fig.tight_layout(rect=[0, 0, 1, 0.92]) + fig.suptitle( + f"Score distributions for all {dataset_label} predictions by C-terminal residue", + fontsize=13, + y=0.98, + ) + _save(fig, out_dir, f"{prefix}_score_histogram_full_panel") + + +def _plot_nontryptic_calibration_scatter( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Subsampled scatter of raw versus calibrated confidence by C-terminus.""" + if "confidence" not in df.columns: + print(f" skipping {prefix}_calibration_scatter (missing confidence)") + return + + work = df.dropna(subset=["confidence", "calibrated_confidence"]) + n_total = len(work) + if n_total < 10: + print(f" skipping {prefix}_calibration_scatter (too few PSMs)") + return + + plot_df = _subsample_psms(work, _NONTRYPTIC_CALIBRATION_SCATTER_MAX_POINTS) + n_show = len(plot_df) + fig, ax = plt.subplots(figsize=(7.5, 7)) + panels = [ + ("Non-tryptic", _NOVEL_COLOUR, False, _NONTRYPTIC_GROUP_ALPHA), + ("Tryptic (K/R)", _MAIN_LINE_COLOUR, True, _TRYPTIC_GROUP_ALPHA), + ] + for label, colour, tryptic, alpha in panels: + sub = plot_df.loc[plot_df["tryptic_cterm"] == tryptic] + if len(sub) == 0: + continue + ax.scatter( + sub["confidence"], + sub["calibrated_confidence"], + c=colour, + s=12, + alpha=alpha, + rasterized=True, + label=f"{label}", + ) + + ax.plot( + [-0.01, 1.01], + [-0.01, 1.01], + ls="--", + color="black", + lw=1, + label="No recalibration", + zorder=5, + ) + ax.set_xlim(-0.01, 1.01) + ax.set_ylim(-0.01, 1.01) + ax.set_xlabel("Raw InstaNovo confidence") + ax.set_ylabel("Calibrated confidence") + if n_show < n_total: + ax.set_title( + f"Raw vs calibrated confidence for all {dataset_label} predictions" + ) + else: + ax.set_title(f"All {dataset_label} predictions") + ax.legend(loc="lower right", fontsize=9) + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, f"{prefix}_calibration_scatter") + + +def _plot_nontryptic_delta_by_terminus( + df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Calibration shift by C-terminal residue.""" + if "confidence" not in df.columns: + print(f" skipping {prefix}_delta_by_terminus (missing confidence)") + return + + df = _add_q_values(df) + y_label = "Calibration shift" + palette = {"Tryptic (K/R)": _MAIN_LINE_COLOUR, "Non-tryptic": _NOVEL_COLOUR} + work = df.copy() + work["delta_confidence"] = work["calibrated_confidence"] - work["confidence"] + + panels: list[tuple[float, pd.DataFrame]] = [] + for fdr_t in FDR_THRESHOLDS: + retained = work[(work["psm_q_value"] <= fdr_t) & work["proteome_hit"]] + retained = retained.dropna(subset=["delta_confidence"]) + if len(retained) < 5: + print( + f" skipping {int(fdr_t * 100)}% FDR panel in " + f"{prefix}_delta_by_terminus (n={len(retained):,})" + ) + continue + retained = retained.copy() + retained["C-terminus"] = retained["tryptic_cterm"].map( + {True: "Tryptic (K/R)", False: "Non-tryptic"}, + ) + panels.append((fdr_t, retained)) + + if not panels: + print(f" skipping {prefix}_delta_by_terminus (no FDR panels with enough PSMs)") + return + + n_cols = len(panels) + fig, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 5), sharey=True) + if n_cols == 1: + axes = [axes] + + for ax, (fdr_t, retained) in zip(axes, panels): + sns.violinplot( + data=retained, + x="C-terminus", + y="delta_confidence", + order=list(_C_TERMINUS_ORDER), + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.axhline(0.0, ls="--", color=_IDEAL_LINE_COLOUR, lw=1) + ax.set_xlabel("") + ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) + ax.grid(False) + _spine_fmt(ax) + + _finalize_nontryptic_violin_figure( + fig, + axes, + y_label=y_label, + suptitle=( + f"Winnow calibration shift for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" + "by C-terminal residue" + ), + ) + _save(fig, out_dir, f"{prefix}_delta_by_terminus") + + +def _pooled_feature_mean_std(retained: pd.DataFrame, col: str) -> tuple[float, float]: + vals = retained[col].dropna() + if len(vals) == 0: + return float("nan"), float("nan") + if len(vals) == 1: + return float(vals.iloc[0]), float("nan") + return float(vals.mean()), float(vals.std()) + + +def _feature_group_median_z_row( + sub: pd.DataFrame, + available: list[str], + pooled: dict[str, tuple[float, float]], +) -> dict[str, float]: + row: dict[str, float] = {} + for col in available: + vals = sub[col].dropna() + if len(vals) == 0: + row[f"median_{col}"] = float("nan") + row[f"z_median_{col}"] = float("nan") + continue + med = float(vals.median()) + row[f"median_{col}"] = round(med, 4) + mu, std = pooled[col] + if np.isnan(std) or std == 0: + row[f"z_median_{col}"] = float("nan") + else: + row[f"z_median_{col}"] = round((med - mu) / std, 4) + return row + + +def _feature_median_z_score_table( + retained: pd.DataFrame, + available: list[str], + groups: list[tuple[str, str, pd.Series]], + *, + group_col: str, + reference: pd.DataFrame | None = None, +) -> pd.DataFrame: + """Per-group feature medians and z-scores relative to *reference* or *retained* PSMs.""" + pool_from = reference if reference is not None else retained + pooled = {col: _pooled_feature_mean_std(pool_from, col) for col in available} + + rows: list[dict] = [] + for group_key, _label, mask in groups: + sub = retained[mask] + row: dict = {group_col: group_key, "n": len(sub)} + row.update(_feature_group_median_z_row(sub, available, pooled)) + rows.append(row) + return pd.DataFrame(rows) + + +def _nontryptic_feature_table(df: pd.DataFrame) -> pd.DataFrame: + """Median feature values for tryptic vs non-tryptic full search space PSMs at 5% FDR.""" + df = _add_q_values(df) + retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] + available = [c for c in FEATURE_COLUMNS if c in retained.columns] + if not available: + return pd.DataFrame() + + return _feature_median_z_score_table( + retained, + available, + [ + ("tryptic", "Tryptic (K/R)", retained["tryptic_cterm"]), + ("non_tryptic", "Non-tryptic", ~retained["tryptic_cterm"]), + ], + group_col="group", + ) + + +def _plot_grouped_feature_z_scores( + feat_df: pd.DataFrame, + *, + group_col: str, + out_dir: Path, + save_name: str, + title: str, + group_style: list[tuple[str, str, str]], + z_score_ylabel: str = "Median z-score", +) -> None: + """Grouped bar chart of pooled z-scored feature medians.""" + if feat_df.empty: + print(f" skipping {save_name} (no feature data)") + return + + z_cols = [c for c in feat_df.columns if c.startswith("z_median_")] + if not z_cols: + print(f" skipping {save_name} (no z-scored feature columns)") + return + + plot_df = feat_df.set_index(group_col) + feature_labels = [_nice_feature_label(c.replace("z_median_", "")) for c in z_cols] + x = np.arange(len(z_cols)) + + present = [ + (key, label, colour) + for key, label, colour in group_style + if key in plot_df.index + ] + n_groups = len(present) + total_width = 0.7 + bar_w = total_width / max(n_groups, 1) + + fig, ax = plt.subplots(figsize=(9, 6.5)) + ax.axhline(0.0, color=_IDEAL_LINE_COLOUR, lw=0.8, zorder=0) + bar_groups = [] + for plot_i, (group_key, label, colour) in enumerate(present): + offset = (plot_i - (n_groups - 1) / 2) * bar_w + vals = plot_df.loc[group_key, z_cols].to_numpy(dtype=float) + bars = ax.bar( + x + offset, + vals, + bar_w, + label=label, + color=colour, + edgecolor="black", + linewidth=1, + ) + bar_groups.append(bars) + + if not bar_groups: + print(f" skipping {save_name} (no groups to plot)") + plt.close(fig) + return + + ax.set_xticks(x) + ax.set_xticklabels(feature_labels, rotation=30, ha="right") + ax.set_ylabel(z_score_ylabel) + ax.set_title(title) + ax.legend(loc="upper right", fontsize=9) + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, save_name) + + +def _plot_nontryptic_feature_comparison( + feat_df: pd.DataFrame, + out_dir: Path, + *, + prefix: str = "nontryptic_digest", + dataset_label: str = "Non-tryptic digest", +) -> None: + """Grouped bar chart of median features, tryptic vs non-tryptic.""" + _plot_grouped_feature_z_scores( + feat_df, + group_col="group", + out_dir=out_dir, + save_name=f"{prefix}_feature_comparison", + title=( + f"Median feature values for {dataset_label} tryptic versus " + f"non-tryptic {FULL_SEARCH_SPACE_LABEL} PSMs at 5% FDR" + ), + group_style=[ + ("tryptic", "Tryptic (K/R)", _MAIN_LINE_COLOUR), + ("non_tryptic", "Non-tryptic", _NOVEL_COLOUR), + ], + ) + + +def _nontryptic_digest_analysis( + predictions_dir: Path, + fasta: Path, + results_dir: Path, + plots_dir: Path, + *, + file_prefix: str, + dataset_label: str, +) -> None: + """Shared tryptic vs non-tryptic digest analysis (any non-tryptic enzyme dataset).""" + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + print(f"Loading predictions from {predictions_dir}") + df_pl = _load_data(predictions_dir) + print(f" {df_pl.height:,} rows loaded") + + print(f"Annotating {FULL_SEARCH_SPACE_LABEL} predictions against {fasta}") + df_pl = _nontryptic_annotate(df_pl, fasta) + n_hits = df_pl.filter(pl.col("proteome_hit")).height + print(f" {n_hits:,} PSMs in {FULL_SEARCH_SPACE_LABEL} (proteome substring match)") + + df = df_pl.to_pandas() + + print("Building tryptic summary table") + summary = _nontryptic_summary_table(df) + summary.to_csv(results_dir / f"{file_prefix}_tryptic_summary.csv", index=False) + print(summary.to_string(index=False)) + + print("Building terminus proportion table") + prop_df = _nontryptic_terminus_proportions_table(df) + prop_df.to_csv(results_dir / f"{file_prefix}_terminus_proportions.csv", index=False) + print(prop_df.to_string(index=False)) + + if "confidence" in df.columns: + print("Building calibration shift table") + delta_df = _nontryptic_calibration_delta_table(df) + delta_df.to_csv( + results_dir / f"{file_prefix}_calibration_delta_summary.csv", index=False + ) + print(delta_df.to_string(index=False)) + else: + print(" skipping calibration shift table (missing raw confidence)") + + print("Building feature comparison table") + feat_df = _nontryptic_feature_table(df) + if not feat_df.empty: + feat_df.to_csv( + results_dir / f"{file_prefix}_feature_comparison.csv", index=False + ) + + plot_kw = {"prefix": file_prefix, "dataset_label": dataset_label} + print("Plotting") + _plot_nontryptic_conf_by_terminus(df, plots_dir, **plot_kw) + _plot_nontryptic_raw_conf_by_terminus(df, plots_dir, **plot_kw) + _plot_nontryptic_score_histograms_at_fdr(df, plots_dir, **plot_kw) + _plot_nontryptic_raw_score_histograms_at_fdr(df, plots_dir, **plot_kw) + _plot_nontryptic_full_score_histograms(df, plots_dir, **plot_kw) + _plot_nontryptic_calibration_scatter(df, plots_dir, **plot_kw) + _plot_nontryptic_delta_by_terminus(df, plots_dir, **plot_kw) + _plot_nontryptic_feature_comparison(feat_df, plots_dir, **plot_kw) + + print( + f"\n{dataset_label} analysis complete. " + f"Tables in {results_dir}; plots in {plots_dir}" + ) + + +@app.command() +def nontryptic_digest( + predictions_dir: Annotated[ + Path, + typer.Option( + "--predictions-dir", + help=( + "winnow predict output folder for the non-tryptic enzyme digest " + "(full search space / acfm)." + ), + ), + ], + fasta: Annotated[ + Path, + typer.Option("--fasta", help="Proteome FASTA for substring matching."), + ], + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for CSV tables."), + ], + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for png/pdf figures."), + ], + file_prefix: Annotated[ + str, + typer.Option( + "--file-prefix", + help="Prefix for output CSV and plot filenames (e.g. chymotrypsin).", + ), + ] = "nontryptic_digest", + dataset_label: Annotated[ + str, + typer.Option( + "--dataset-label", + help="Human-readable dataset name used in plot titles.", + ), + ] = "Non-tryptic digest", +) -> None: + """Analyse calibrator behaviour on non-tryptic enzyme digest peptides.""" + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + _nontryptic_digest_analysis( + predictions_dir, + fasta, + results_dir, + plots_dir, + file_prefix=file_prefix, + dataset_label=dataset_label, + ) + + +# ── ProteomeTools-1 analysis ──────────────────────────────────────────── + + +def _exact_match_category(fits_precursor: bool) -> str: + if fits_precursor: + return "exact_match_and_fits_precursor" + return "exact_match_and_no_precursor_fit" + + +def _subsequence_category(fits_precursor: bool) -> str: + if fits_precursor: + return "subsequence_and_fits_precursor" + return "subsequence_and_no_precursor_fit" + + +def _neither_category(fits_precursor: bool) -> str: + if fits_precursor: + return "neither_and_fits_precursor" + return "neither_and_no_precursor_fit" + + +def _classify_exact_matches( + categories: list[str], + predictions: list[str], + fits_precursor: list[bool], + lcfm_peptide_set: set[str], +) -> None: + for i, (peptide, fit) in enumerate(zip(predictions, fits_precursor)): + if peptide in lcfm_peptide_set: + categories[i] = _exact_match_category(fit) + + +def _classify_subsequence_matches( + categories: list[str], + predictions: list[str], + fits_precursor: list[bool], + lcfm_haystack: str, +) -> None: + remaining_indices = [ + i for i, category in enumerate(categories) if category == "neither" + ] + remaining_peps = [predictions[i] for i in remaining_indices] + remaining_fits = [fits_precursor[i] for i in remaining_indices] + if not remaining_peps: + return + + hits = _batch_substring_hits(remaining_peps, lcfm_haystack) + for idx, fit, hit in zip(remaining_indices, remaining_fits, hits): + if hit: + categories[idx] = _subsequence_category(fit) + + +def _classify_remaining_neither( + categories: list[str], + fits_precursor: list[bool], +) -> None: + for i, fit in enumerate(fits_precursor): + if categories[i] == "neither": + categories[i] = _neither_category(fit) + + +def _classify_predictions( + predictions: list[str], + fits_precursor: list[bool], + lcfm_peptide_set: set[str], + lcfm_haystack: str, +) -> list[str]: + """Classify each prediction by lcfm overlap and precursor mass fit (<20 ppm). + + Uses Aho-Corasick to find which predictions are substrings of at least + one lcfm peptide (the haystack is built by joining all lcfm peptides + with a separator). Unmatched predictions are split by precursor fit. + """ + categories = ["neither"] * len(predictions) + _classify_exact_matches(categories, predictions, fits_precursor, lcfm_peptide_set) + _classify_subsequence_matches( + categories, predictions, fits_precursor, lcfm_haystack + ) + _classify_remaining_neither(categories, fits_precursor) + return categories + + +def _proteometools_summary_table(df: pd.DataFrame) -> pd.DataFrame: + """Build novelty summary table at each FDR threshold.""" + df = _add_q_values(df) + rows: list[dict] = [] + for fdr_t in FDR_THRESHOLDS: + retained = df[df["psm_q_value"] <= fdr_t] + n = len(retained) + if n == 0: + rows.append( + { + "fdr_threshold": fdr_t, + "n_retained": 0, + "n_exact_match_and_no_precursor_fit": 0, + "n_exact_match_and_fits_precursor": 0, + "n_subsequence_and_no_precursor_fit": 0, + "n_subsequence_and_fits_precursor": 0, + "n_neither_and_no_precursor_fit": 0, + "n_neither_and_fits_precursor": 0, + "pct_exact_or_sub_and_fit_among_retained": 0.0, + } + ) + continue + + cats = retained["novelty_category"] + n_exact_and_no_fit = int((cats == "exact_match_and_no_precursor_fit").sum()) + n_exact_and_fits_precursor = int( + (cats == "exact_match_and_fits_precursor").sum() + ) + n_sub_and_no_fit = int((cats == "subsequence_and_no_precursor_fit").sum()) + n_sub_and_fits_precursor = int((cats == "subsequence_and_fits_precursor").sum()) + n_neither_and_no_fit = int((cats == "neither_and_no_precursor_fit").sum()) + n_neither_and_fits_precursor = int((cats == "neither_and_fits_precursor").sum()) + + rows.append( + { + "fdr_threshold": fdr_t, + "n_retained": n, + "n_exact_match_and_no_precursor_fit": n_exact_and_no_fit, + "n_exact_match_and_fits_precursor": n_exact_and_fits_precursor, + "n_subsequence_and_no_precursor_fit": n_sub_and_no_fit, + "n_subsequence_and_fits_precursor": n_sub_and_fits_precursor, + "n_neither_and_no_precursor_fit": n_neither_and_no_fit, + "n_neither_and_fits_precursor": n_neither_and_fits_precursor, + "pct_exact_or_sub_and_fit_among_retained": round( + (n_exact_and_fits_precursor + n_sub_and_fits_precursor) / n * 100, 2 + ) + if n > 0 + else 0.0, + } + ) + return pd.DataFrame(rows) + + +_PROTEOMETOOLS_CONF_PLOT_LABELS: dict[str, str] = { + "exact_match_and_no_precursor_fit": "ID-", + "exact_match_and_fits_precursor": "ID+", + "subsequence_and_no_precursor_fit": "Sub-", + "subsequence_and_fits_precursor": "Sub+", + "neither_and_no_precursor_fit": "Novel-", + "neither_and_fits_precursor": "Novel+", +} + +_PROTEOMETOOLS_CONF_CATEGORY_ORDER = list(_PROTEOMETOOLS_CONF_PLOT_LABELS.keys()) + + +def _proteometools_conf_category_legend(ax: plt.Axes) -> None: + handles = [ + Line2D([], [], color="none", label="ID: Exact sequence match to labelled set."), + Line2D([], [], color="none", label="Sub: Subsequence of labelled set peptide."), + Line2D([], [], color="none", label="Novel: No sequence match to labelled set."), + Line2D( + [], + [], + color="none", + label="+: Matches precursor mass within 20 ppm.", + ), + ] + ( + Line2D( + [], + [], + color="none", + label="-: Does not match precursor mass within 20 ppm.", + ), + ) + ax.legend( + handles=handles, + loc="upper left", + bbox_to_anchor=(1.02, 1.0), + borderaxespad=0, + frameon=True, + fontsize=9, + ) + + +def _plot_proteometools_conf_by_category( + df: pd.DataFrame, + out_dir: Path, +) -> None: + """Violin plot of calibrated confidence by novelty category (all unlabelled PSMs).""" + plot_df = df.dropna(subset=["calibrated_confidence"]).copy() + if len(plot_df) < 5: + print(" skipping proteometools_conf_by_category (too few PSMs)") + return + + palette = { + _PROTEOMETOOLS_CONF_PLOT_LABELS[k]: _PALETTE[i] + for i, k in enumerate(_PROTEOMETOOLS_CONF_CATEGORY_ORDER) + } + + plot_df["Category"] = plot_df["novelty_category"].map( + _PROTEOMETOOLS_CONF_PLOT_LABELS + ) + present_cats = [ + _PROTEOMETOOLS_CONF_PLOT_LABELS[c] + for c in _PROTEOMETOOLS_CONF_CATEGORY_ORDER + if _PROTEOMETOOLS_CONF_PLOT_LABELS[c] in plot_df["Category"].values + ] + if not present_cats: + print(" skipping proteometools_conf_by_category (no categories present)") + return + + conf = plot_df["calibrated_confidence"] + y_min, y_max = float(conf.min()), float(conf.max()) + + fig, ax = plt.subplots(figsize=(10, 5)) + sns.violinplot( + data=plot_df, + x="Category", + y="calibrated_confidence", + order=present_cats, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_ylim(y_min, y_max) + ax.set_xlabel("") + ax.set_ylabel("Calibrated confidence") + ax.set_title( + "Calibrated confidence for ProteomeTools-1 predictions\nby novelty category" + ) + ax.grid(False) + _spine_fmt(ax) + _proteometools_conf_category_legend(ax) + + fig.tight_layout() + _save(fig, out_dir, "proteometools_conf_by_category") + + +def _plot_proteometools_hit_rate( + df: pd.DataFrame, + out_dir: Path, +) -> None: + """Line plot: validated hit rate (exact or subsequence match and fits precursor mass within 20ppm) by calibrated confidence decile.""" + if len(df) < 20: + print(" skipping proteometools_hit_rate_vs_conf (too few PSMs)") + return + + df = df.copy() + df["is_validated"] = df["novelty_category"].isin( + ["exact_match_and_fits_precursor", "subsequence_and_fits_precursor"] + ) + df["conf_decile"] = pd.qcut( + df["calibrated_confidence"], + q=10, + duplicates="drop", + ) + grouped = ( + df.groupby("conf_decile", observed=True) + .agg( + hit_rate=("is_validated", "mean"), + mid=("calibrated_confidence", "mean"), + ) + .sort_values("mid") + ) + + fig, ax = plt.subplots(figsize=(8, 6)) + ax.plot( + grouped["mid"], + grouped["hit_rate"], + color=_MAIN_LINE_COLOUR, + linewidth=1.5, + marker="o", + markersize=6, + label="Validated hit rate", + ) + overall = float(df["is_validated"].mean()) + ax.axhline( + overall, + color=_IDEAL_LINE_COLOUR, + lw=1, + linestyle="--", + label=f"Overall mean ({overall:.2%})", + ) + ax.set_xlabel("Mean calibrated confidence per decile") + ax.set_ylabel( + "Fraction validated\n(exact match or subsequence fitting precursor mass)" + ) + ax.set_title( + "Validated hit rate by calibrated confidence decile for ProteomeTools-1" + ) + ax.legend(loc="lower right") + ax.grid(False) + _spine_fmt(ax) + fig.tight_layout() + _save(fig, out_dir, "proteometools_hit_rate_vs_conf") + + +def _proteometools_feature_table( + df: pd.DataFrame, + labelled_df: pd.DataFrame, +) -> pd.DataFrame: + """Median features for exact / novel / neither at 5% FDR.""" + df = _add_q_values(df) + labelled_df = _add_q_values(labelled_df) + retained = df[df["psm_q_value"] <= 0.05] + labelled_ref = labelled_df[labelled_df["psm_q_value"] <= 0.05] + available = [c for c in FEATURE_COLUMNS if c in retained.columns] + if not available: + return pd.DataFrame() + + return _feature_median_z_score_table( + retained, + available, + [ + ( + "exact_match_and_fits_precursor", + "Exact match, fits precursor mass", + retained["novelty_category"] == "exact_match_and_fits_precursor", + ), + ( + "exact_match_and_no_precursor_fit", + "Exact match, no precursor mass fit", + retained["novelty_category"] == "exact_match_and_no_precursor_fit", + ), + ( + "subsequence_and_fits_precursor", + "Subsequence, precursor mass fit", + retained["novelty_category"] == "subsequence_and_fits_precursor", + ), + ( + "subsequence_and_no_precursor_fit", + "Subsequence, no precursor mass fit", + retained["novelty_category"] == "subsequence_and_no_precursor_fit", + ), + ( + "neither_and_fits_precursor", + "Novel, precursor mass fit", + retained["novelty_category"] == "neither_and_fits_precursor", + ), + ( + "neither_and_no_precursor_fit", + "Novel, no precursor mass fit", + retained["novelty_category"] == "neither_and_no_precursor_fit", + ), + ], + group_col="category", + reference=labelled_ref, + ) + + +def _plot_proteometools_feature_comparison( + feat_df: pd.DataFrame, + out_dir: Path, +) -> None: + """Grouped bar chart of median features by category.""" + _plot_grouped_feature_z_scores( + feat_df, + group_col="category", + out_dir=out_dir, + save_name="proteometools_feature_comparison", + title=( + "Median feature values for ProteomeTools-1 predictions by novelty " + "category at 5% FDR" + ), + group_style=[ + ( + "exact_match_and_fits_precursor", + "Exact match, precursor mass fit", + _PALETTE[1], + ), + ( + "exact_match_and_no_precursor_fit", + "Exact match, no precursor mass fit", + _PALETTE[0], + ), + ( + "subsequence_and_fits_precursor", + "Subsequence, precursor mass fit", + _PALETTE[3], + ), + ( + "subsequence_and_no_precursor_fit", + "Subsequence, no precursor mass fit", + _PALETTE[2], + ), + ("neither_and_fits_precursor", "Novel, precursor mass fit", _PALETTE[5]), + ( + "neither_and_no_precursor_fit", + "Novel, no precursor mass fit", + _PALETTE[4], + ), + ], + z_score_ylabel="Median z-score (vs labelled PSMs at 5% FDR)", + ) + + +@app.command() +def proteometools( + lcfm_predictions_dir: Annotated[ + Path, + typer.Option( + "--lcfm-predictions-dir", + help="winnow predict output for PXD004732 lcfm (labelled).", + ), + ], + acfm_predictions_dir: Annotated[ + Path, + typer.Option( + "--acfm-predictions-dir", + help="winnow predict output for PXD004732 acfm (unlabelled).", + ), + ], + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for CSV tables."), + ], + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for png/pdf figures."), + ], +) -> None: + """Analyse calibrator behaviour on ProteomeTools-1 novel identifications.""" + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + print(f"Loading lcfm predictions from {lcfm_predictions_dir}") + lcfm_pl = _load_data(lcfm_predictions_dir) + print(f" {lcfm_pl.height:,} lcfm rows") + + print(f"Loading acfm predictions from {acfm_predictions_dir}") + acfm_pl = _load_data(acfm_predictions_dir) + print(f" {acfm_pl.height:,} acfm rows") + + print("Building lcfm peptide set for subsequence matching") + lcfm_sequences = lcfm_pl["sequence"].drop_nulls().to_list() + lcfm_peptide_set: set[str] = set() + for seq in lcfm_sequences: + stripped = _strip_mods(seq) + if stripped: + lcfm_peptide_set.add(stripped) + print(f" {len(lcfm_peptide_set):,} unique lcfm peptides") + + lcfm_haystack = _PROTEOME_JOIN_SEP.join(sorted(lcfm_peptide_set)) + + print("Classifying unlabelled predictions") + unlabelled_pl = acfm_pl.join( + lcfm_pl.select("spectrum_id"), on="spectrum_id", how="anti" + ) + unlabelled_pl = unlabelled_pl.with_columns( + (pl.col("delta_mass_ppm").abs() < 20).alias("fits_precursor") + ) + unlabelled_preds_raw = unlabelled_pl["prediction"].to_list() + unlabelled_preds_stripped = [ + _strip_mods(p) if isinstance(p, str) else "" for p in unlabelled_preds_raw + ] + fits_precursor = unlabelled_pl["fits_precursor"].to_list() + + categories = _classify_predictions( + unlabelled_preds_stripped, + fits_precursor, + lcfm_peptide_set, + lcfm_haystack, + ) + unlabelled_pl = unlabelled_pl.with_columns( + pl.Series("novelty_category", categories, dtype=pl.Utf8), + ) + df = unlabelled_pl.to_pandas() + labelled_df = lcfm_pl.to_pandas() + + print("Building novelty summary table") + summary = _proteometools_summary_table(df) + summary.to_csv(results_dir / "proteometools_novelty_summary.csv", index=False) + print(summary.to_string(index=False)) + + print("Building feature comparison table") + feat_df = _proteometools_feature_table(df, labelled_df) + if not feat_df.empty: + feat_df.to_csv( + results_dir / "proteometools_feature_comparison.csv", index=False + ) + + print("Plotting") + _plot_proteometools_conf_by_category(df, plots_dir) + _plot_proteometools_hit_rate(df, plots_dir) + _plot_proteometools_feature_comparison(feat_df, plots_dir) + + print( + f"\nProteomeTools-1 analysis complete. " + f"Tables in {results_dir}; plots in {plots_dir}" + ) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/analyze_upscored_fps.py b/paper_scripts/analyze_upscored_fps.py new file mode 100644 index 00000000..3ea56dbf --- /dev/null +++ b/paper_scripts/analyze_upscored_fps.py @@ -0,0 +1,713 @@ +#!/usr/bin/env python3 +"""Characterise up-scored false positives from Winnow calibration. + +For each labelled evaluation dataset (where both ``sequence`` and ``prediction`` +are available), this script quantifies the false positives that calibration +"rescues" into high-confidence regions and compares their feature profiles to +true positives. + +Inputs are ``winnow predict`` output folders, each containing +``preds_and_fdr_metrics.csv`` and ``metadata.csv``. +""" + +from __future__ import annotations + +import json +import logging +import re +from pathlib import Path +from typing import Annotated + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer +import yaml +from instanovo.utils.metrics import Metrics +from instanovo.utils.residues import ResidueSet +from rich.logging import RichHandler + +from winnow.fdr.nonparametric import NonParametricFDRControl + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] + +TP_COLOR = _CORRECT_COLOUR +FP_COLOR = _INCORRECT_COLOUR + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)-?") +_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] +_FEATURE_VIOLIN_FDR_THRESHOLDS = [0.05, 0.10] +_MIN_FEATURE_VIOLIN_PSMs = 20 + +# Max PSMs per correctness panel in the raw-vs-calibrated confidence scatter. +_CONFIDENCE_SCATTER_MAX_POINTS = 10_000 +_CONFIDENCE_SCATTER_RANDOM_STATE = 42 + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD004732": "ProteomeTools-1", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", +} + +_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") + +# new_eval_sets_results layout: lcfm/PXD004452//preds_and_fdr_metrics.csv +_PXD_ACCESSION_PREFIX = "PXD" + +FEATURE_COLUMNS_OF_INTEREST = [ + "spectral_angle", + "xcorr", + "ion_matches", + "ion_match_intensity", + "irt_error", + "mass_error_ppm", + "margin", + "entropy", + "confidence", +] + +_MASS_ERROR_COLUMNS = ("mass_error_da", "mass_error_ppm") + +_NICE_LABELS: dict[str, str] = { + "ion_matches": "Ion match rate", + "ion_match_intensity": "Ion match intensity", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Max ion gap", + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation (XCorr)", + "mass_error_ppm": "Precursor mass error (ppm)", + "mass_error_da": "Precursor mass error (Da)", + "irt_error": "iRT prediction error", + "confidence": "Model confidence", + "margin": "Beam margin", + "median_margin": "Beam median margin", + "entropy": "Beam entropy", + "z-score": "Beam z-score", + "edit_distance": "Runner-up edit distance", + "min_token_probability": "Min. token probability", + "std_token_probability": "Std. token probability", +} + + +def _nice_label(col: str) -> str: + return _NICE_LABELS.get(col, col.replace("_", " ").capitalize()) + + +def _mass_error_column(df: pd.DataFrame, *, min_count: int = 10) -> str | None: + """Return ``mass_error_da`` or ``mass_error_ppm`` when present with enough data.""" + for col in _MASS_ERROR_COLUMNS: + if col in df.columns and df[col].notna().sum() > min_count: + return col + return None + + +def _violin_feature_columns(df: pd.DataFrame) -> list[str]: + """Feature list for violin plots, resolving Da vs ppm mass error.""" + cols: list[str] = [] + for col in FEATURE_COLUMNS_OF_INTEREST: + if col == "mass_error_ppm": + mass_col = _mass_error_column(df) + if mass_col is not None: + cols.append(mass_col) + else: + cols.append(col) + return cols + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _get_residue_masses() -> dict[str, float]: + config_path = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" + with open(config_path) as f: + cfg = yaml.safe_load(f) + return cfg["residue_masses"] + + +def _save_fig(fig: plt.Figure, base_path: Path, fmt: str = "both") -> None: + if fmt in ("pdf", "both"): + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + if fmt in ("png", "both"): + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _folder_display_name(folder_name: str) -> str: + """Map an evaluation folder name to a publication-ready dataset label.""" + key = folder_name + for suffix in _FOLDER_SUFFIXES: + if key.endswith(suffix): + key = key[: -len(suffix)] + break + return DATASET_DISPLAY_NAMES.get(key, key) + + +def _subsample_psms( + df: pd.DataFrame, + max_points: int, + random_state: int = _CONFIDENCE_SCATTER_RANDOM_STATE, +) -> pd.DataFrame: + """Return up to ``max_points`` rows without replacement.""" + if len(df) <= max_points: + return df + return df.sample(n=max_points, random_state=random_state) + + +def _strip_mods(seq: str) -> str: + if not seq or not isinstance(seq, str): + return "" + s = _MOD_PLUS.sub("", seq) + s = _MOD_UNIMOD.sub("", s) + return s.replace("I", "L") + + +def _project_key_from_folder(folder_name: str) -> str: + """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" + for suffix in _FOLDER_SUFFIXES: + if folder_name.endswith(suffix): + return folder_name[: -len(suffix)] + return folder_name + + +def _is_labelled_preds_folder(folder: Path) -> bool: + preds_csv = folder / "preds_and_fdr_metrics.csv" + if not preds_csv.is_file(): + return False + header = pd.read_csv(preds_csv, nrows=0).columns.tolist() + required = {"sequence", "prediction", "calibrated_confidence"} + return required.issubset(header) + + +def _register_labelled_folder(results: dict[str, Path], key: str, folder: Path) -> None: + """Register *folder* under *key*, warning on duplicate keys.""" + if key in results: + logger.warning( + "Duplicate labelled project key %r: %s and %s", + key, + results[key], + folder, + ) + return + results[key] = folder + + +def _discover_labelled_folders(root: Path) -> dict[str, Path]: + """Find folders with labelled ``preds_and_fdr_metrics.csv``. + + Supports flat project folders (``{root}/PXD004732/``) and nested per-run + layouts used by new eval sets (``{root}/PXD004452//``). + """ + results: dict[str, Path] = {} + if not root.is_dir(): + return results + + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if _is_labelled_preds_folder(child): + _register_labelled_folder( + results, _project_key_from_folder(child.name), child + ) + continue + if not child.name.startswith(_PXD_ACCESSION_PREFIX): + continue + for run_dir in sorted(child.iterdir()): + if run_dir.is_dir() and _is_labelled_preds_folder(run_dir): + _register_labelled_folder(results, run_dir.name, run_dir) + return results + + +def _load_dataset(folder: Path) -> pd.DataFrame: + """Load and merge preds + metadata CSVs for a single evaluation folder.""" + preds = pd.read_csv(folder / "preds_and_fdr_metrics.csv") + meta_path = folder / "metadata.csv" + if meta_path.is_file(): + meta = pd.read_csv(meta_path) + join_cols = ["spectrum_id"] + [ + c for c in meta.columns if c != "spectrum_id" and c not in preds.columns + ] + if len(join_cols) > 1: + preds = preds.merge( + meta[join_cols].drop_duplicates(subset=["spectrum_id"]), + on="spectrum_id", + how="left", + ) + if "correct" not in preds.columns and {"sequence", "prediction"}.issubset( + preds.columns + ): + preds = preds.copy() + preds["correct"] = preds["sequence"] == preds["prediction"] + return preds + + +def _add_q_values( + df: pd.DataFrame, conf_col: str = "calibrated_confidence" +) -> pd.DataFrame: + """Fit non-parametric FDR and append ``psm_q_value`` if missing.""" + if "psm_q_value" in df.columns: + return df + fdr = NonParametricFDRControl() + fdr.fit(dataset=df[conf_col]) + df = fdr.add_psm_q_value(df, confidence_col=conf_col) + return df + + +# --------------------------------------------------------------------------- +# Analysis +# --------------------------------------------------------------------------- +def _upscored_summary_table( + df: pd.DataFrame, + delta_threshold: float, + dataset_name: str, +) -> pd.DataFrame: + """Build per-FDR-threshold summary of up-scored TP / FP counts.""" + df = _add_q_values(df) + upscored = df["delta_confidence"] > delta_threshold + + rows = [] + for fdr_t in FDR_THRESHOLDS: + passing = df["psm_q_value"] <= fdr_t + for label, mask in [ + ("all", pd.Series(True, index=df.index)), + ("up-scored", upscored), + ("not up-scored", ~upscored), + ]: + sub = df[mask & passing] + n = len(sub) + n_correct = int(sub["correct"].sum()) if "correct" in sub.columns else 0 + n_incorrect = n - n_correct + rows.append( + { + "dataset": dataset_name, + "fdr_threshold": fdr_t, + "subset": label, + "n_passing": n, + "n_correct": n_correct, + "n_incorrect": n_incorrect, + "pct_correct": round(n_correct / n * 100, 2) if n > 0 else 0.0, + } + ) + return pd.DataFrame(rows) + + +def _plot_confidence_scatter( + df: pd.DataFrame, + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Subsampled scatter of raw vs calibrated confidence, colored by correctness.""" + display = _folder_display_name(dataset_name) + fig, ax = plt.subplots(figsize=(8, 6)) + + # Define panels as before + panels = [ + ("Correct", TP_COLOR, df["correct"].astype(bool)), + ("Incorrect", FP_COLOR, ~df["correct"].astype(bool)), + ] + + handles = [] + for label, colour, mask in panels: + sub = df.loc[mask, ["confidence", "calibrated_confidence"]].dropna() + n_total = len(sub) + if n_total < 2: + # Only skip plotting, no data for this class + continue + + plot_df = _subsample_psms(sub, _CONFIDENCE_SCATTER_MAX_POINTS) + handle = ax.scatter( + plot_df["confidence"], + plot_df["calibrated_confidence"], + c=colour, + s=10, + alpha=0.3, + rasterized=True, + label=f"{label}", + ) + handles.append(handle) + + ax.plot( + [-0.01, 1.01], + [-0.01, 1.01], + ls="--", + color="black", + lw=1, + label="No recalibration", + zorder=5, + ) + ax.set_xlim(-0.01, 1.01) + ax.set_ylim(-0.01, 1.01) + ax.set_xlabel("Raw confidence") + ax.set_ylabel("Calibrated confidence") + ax.legend(loc="lower right", fontsize=9) + _style_ax(ax) + + fig.suptitle( + f"Raw versus calibrated confidence for {display}", + fontsize=13, + ) + fig.tight_layout() + _save_fig(fig, output_dir / f"confidence_scatter_{dataset_name}", plot_format) + + +def _plot_feature_distributions_at_fdr( + upscored: pd.DataFrame, + *, + fdr_t: float, + delta_threshold: float, + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Violin plots of features for up-scored TPs vs FPs retained at one FDR cutoff.""" + n = len(upscored) + pct = int(fdr_t * 100) + if n < _MIN_FEATURE_VIOLIN_PSMs: + logger.info( + "Skipping feature violins for %s at %d%% FDR (n=%d up-scored retained)", + dataset_name, + pct, + n, + ) + return + + available = [ + c + for c in _violin_feature_columns(upscored) + if c in upscored.columns + and upscored[c].notna().sum() >= _MIN_FEATURE_VIOLIN_PSMs + ] + if not available: + logger.warning( + "No feature columns with >=%d values for %s at %d%% FDR", + _MIN_FEATURE_VIOLIN_PSMs, + dataset_name, + pct, + ) + return + + plot_df = upscored.copy() + plot_df["label"] = plot_df["correct"].map({True: "Correct", False: "Incorrect"}) + n_features = len(available) + n_cols = min(3, n_features) + n_rows = (n_features + n_cols - 1) // n_cols + fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows)) + axes = np.atleast_1d(axes).flatten() + + palette = {"Correct": TP_COLOR, "Incorrect": FP_COLOR} + n_plotted = 0 + + for i, col in enumerate(available): + ax = axes[i] + sub = plot_df[[col, "label"]].dropna(subset=[col]) + if len(sub) < _MIN_FEATURE_VIOLIN_PSMs: + ax.set_visible(False) + continue + sns.violinplot( + data=sub, + x="label", + y=col, + palette=palette, + ax=ax, + inner="quartile", + cut=0, + linewidth=0.8, + ) + ax.set_xlabel("") + ax.set_ylabel(_nice_label(col)) + ax.set_title(_nice_label(col)) + _style_ax(ax) + n_plotted += 1 + + for i in range(len(available), len(axes)): + axes[i].set_visible(False) + + if n_plotted == 0: + plt.close(fig) + logger.info( + "Skipping feature violins for %s at %d%% FDR (no feature panels with n>=%d)", + dataset_name, + pct, + _MIN_FEATURE_VIOLIN_PSMs, + ) + return + + display = _folder_display_name(dataset_name) + fig.suptitle( + f"Feature distributions for up-scored PSMs retained at {pct}% FDR " + f"(calibration increase > {delta_threshold:.2f}) on {display}\n" + f"(n={n:,})", + fontsize=13, + ) + fig.tight_layout() + _save_fig( + fig, + output_dir / f"upscored_features_{dataset_name}_fdr{pct}", + plot_format, + ) + + +def _plot_feature_distributions( + df: pd.DataFrame, + delta_threshold: float, + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Violin plots at 5% and 10% FDR for up-scored correct vs incorrect PSMs.""" + df = _add_q_values(df) + upscored_mask = df["delta_confidence"] > delta_threshold + for fdr_t in _FEATURE_VIOLIN_FDR_THRESHOLDS: + retained = df[upscored_mask & (df["psm_q_value"] <= fdr_t)].copy() + _plot_feature_distributions_at_fdr( + retained, + fdr_t=fdr_t, + delta_threshold=delta_threshold, + dataset_name=dataset_name, + output_dir=output_dir, + plot_format=plot_format, + ) + + +def _upscored_fp_detail( + df: pd.DataFrame, + delta_threshold: float, + dataset_name: str, + metrics: Metrics, +) -> pd.DataFrame: + """Detailed characterisation of up-scored FPs that pass FDR thresholds.""" + df = _add_q_values(df) + upscored_fps = df[ + (df["delta_confidence"] > delta_threshold) & (~df["correct"].astype(bool)) + ].copy() + + if len(upscored_fps) == 0: + return pd.DataFrame() + + def _match_fraction(row: pd.Series) -> float: + nm = row.get("num_matches", 0) + seq = row.get("sequence", "") + if isinstance(seq, str): + tokens = metrics._split_peptide(seq) + else: + tokens = seq if seq else [] + return nm / len(tokens) if tokens else 0.0 + + upscored_fps["match_fraction"] = upscored_fps.apply(_match_fraction, axis=1) + + def _edit_dist(row: pd.Series) -> int: + s = _strip_mods(str(row.get("sequence", ""))) + p = _strip_mods(str(row.get("prediction", ""))) + if not s or not p: + return -1 + return _levenshtein(s, p) + + upscored_fps["edit_distance_norm"] = upscored_fps.apply(_edit_dist, axis=1) + + rows = [] + for fdr_t in FDR_THRESHOLDS: + sub = upscored_fps[upscored_fps["psm_q_value"] <= fdr_t] + if len(sub) == 0: + rows.append( + {"dataset": dataset_name, "fdr_threshold": fdr_t, "n_upscored_fps": 0} + ) + continue + rows.append( + { + "dataset": dataset_name, + "fdr_threshold": fdr_t, + "n_upscored_fps": len(sub), + "mean_match_fraction": round(float(sub["match_fraction"].mean()), 4), + "median_edit_distance": int(sub["edit_distance_norm"].median()), + "n_edit_dist_le2": int((sub["edit_distance_norm"] <= 2).sum()), + "n_partial_match": int((sub["match_fraction"] > 0).sum()), + } + ) + return pd.DataFrame(rows) + + +def _levenshtein(s: str, t: str) -> int: + """Simple Levenshtein distance for short peptide strings.""" + n, m = len(s), len(t) + if n == 0: + return m + if m == 0: + return n + prev = list(range(m + 1)) + for i in range(1, n + 1): + curr = [i] + [0] * m + for j in range(1, m + 1): + cost = 0 if s[i - 1] == t[j - 1] else 1 + curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost) + prev = curr + return prev[m] + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +_DEFAULT_PREDICTIONS_ROOT = Path("predictions/general_model") +_DEFAULT_RESULTS_DIR = Path("analysis/upscored_fps") +_DEFAULT_PLOTS_DIR = Path("analysis/upscored_fps/plots") + + +@app.command() +def main( + predictions_root: Annotated[ + Path, + typer.Option( + help="Root directory containing winnow predict output folders.", + ), + ] = _DEFAULT_PREDICTIONS_ROOT, + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for CSV/JSON tables."), + ] = _DEFAULT_RESULTS_DIR, + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for png/pdf figures."), + ] = _DEFAULT_PLOTS_DIR, + delta_threshold: Annotated[ + float, + typer.Option( + help="Minimum delta (calibrated - raw) to classify a PSM as up-scored. " + "Default 0.2 (20 percentage-point increase).", + ), + ] = 0.2, + plot_format: Annotated[ + str, + typer.Option(help="Plot format: 'pdf', 'png', or 'both'."), + ] = "both", +) -> None: + """Characterise false positives that calibration up-scores into high-confidence regions.""" + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + residue_masses = _get_residue_masses() + metrics = Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + folders = _discover_labelled_folders(predictions_root) + if not folders: + logger.error("No labelled output folders found under %s", predictions_root) + raise typer.Exit(code=1) + + logger.info("Found %d labelled folder(s): %s", len(folders), list(folders.keys())) + + all_summary: list[pd.DataFrame] = [] + all_detail: list[pd.DataFrame] = [] + + for name, folder in folders.items(): + logger.info("Processing %s ...", name) + df = _load_dataset(folder) + + missing_conf = [ + c for c in ("confidence", "calibrated_confidence") if c not in df.columns + ] + if missing_conf: + logger.warning( + "Skipping %s: missing confidence columns %s", + name, + missing_conf, + ) + continue + if "correct" not in df.columns: + logger.warning("Skipping %s: missing 'correct' column", name) + continue + + df["delta_confidence"] = df["calibrated_confidence"] - df["confidence"] + + logger.info( + " %s: %d PSMs, %d correct, delta stats: mean=%.3f, q75=%.3f", + name, + len(df), + int(df["correct"].sum()), + df["delta_confidence"].mean(), + df["delta_confidence"].quantile(0.75), + ) + + summary = _upscored_summary_table(df, delta_threshold, name) + all_summary.append(summary) + + _plot_confidence_scatter(df, name, plots_dir, plot_format) + _plot_feature_distributions(df, delta_threshold, name, plots_dir, plot_format) + + detail = _upscored_fp_detail(df, delta_threshold, name, metrics) + if len(detail) > 0: + all_detail.append(detail) + + if all_summary: + combined = pd.concat(all_summary, ignore_index=True) + combined.to_csv(results_dir / "upscored_summary.csv", index=False) + logger.info("Summary table:\n%s", combined.to_string(index=False)) + + with open(results_dir / "upscored_summary.json", "w") as f: + json.dump(combined.to_dict(orient="records"), f, indent=2) + + if all_detail: + detail_df = pd.concat(all_detail, ignore_index=True) + detail_df.to_csv(results_dir / "upscored_fp_detail.csv", index=False) + logger.info("FP detail table:\n%s", detail_df.to_string(index=False)) + + logger.info( + "Up-scored FP analysis complete. Tables in %s; plots in %s", + results_dir, + plots_dir, + ) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/annotate_preds_proteome_hits.py b/paper_scripts/annotate_preds_proteome_hits.py new file mode 100644 index 00000000..a4e27e40 --- /dev/null +++ b/paper_scripts/annotate_preds_proteome_hits.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +"""Post-process ``winnow predict`` outputs: min residue length + proteome substring hit. + +Paper full-search recompute uses this **after** ``winnow predict`` so deposited +``general_results/full/`` trees stay aligned (FDR on all PSMs, then drop short +peptides and add ``proteome_hit``). That differs from the package CLI +``winnow annotate-proteome-hits``, which annotates spectra + de novo preds into a +Winnow dataset directory *before* predict. See ``paper_scripts/README.md``. + +Each project is a subfolder of ``--predictions-root`` containing +``preds_and_fdr_metrics.csv`` and ``metadata.csv``. +""" + +from __future__ import annotations + +import logging +import math +import re +from pathlib import Path +from typing import Annotated, Any + +import ahocorasick +import polars as pl +import typer +import yaml +from Bio import SeqIO +from instanovo.utils.metrics import Metrics +from instanovo.utils.residues import ResidueSet + +logger = logging.getLogger(__name__) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_DEFAULT_RESIDUES = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" + +_PROTEOME_JOIN_SEP = "\x1f" +_MOD_ROUND = re.compile(r"\([^)]*\)-?") +_MOD_SQUARE = re.compile(r"\[[^\]]*\]-?") + +app = typer.Typer(add_completion=False, no_args_is_help=True) + + +def normalize_sequence(seq: str) -> str: + """Normalise a peptide sequence by replacing I with L. + + Args: + seq: Peptide sequence to normalise. + """ + if seq: + return seq.replace("I", "L") + return seq + + +def load_proteome_haystack(fasta_file: Path | str) -> str: + """Load a FASTA file into a string for substring matching. + + Args: + fasta_file: Path to the FASTA file containing the proteome. + """ + path = Path(fasta_file) + if not path.is_file(): + raise FileNotFoundError(f"FASTA file not found: {path}") + + parts: list[str] = [] + for record in SeqIO.parse(path, "fasta"): + s = normalize_sequence(str(record.seq)) + if s: + parts.append(s) + return _PROTEOME_JOIN_SEP.join(parts) + + +def processed_peptide_for_match(prediction: str) -> str: + """Process a peptide string for substring matching. + + Args: + prediction: Peptide string to process. + """ + if not prediction or not isinstance(prediction, str): + return "" + s = _MOD_ROUND.sub("", prediction) + s = _MOD_SQUARE.sub("", s) + return s.replace("I", "L") + + +def _batch_peptide_substring_hits(peptides: list[str], haystack: str) -> list[bool]: + n = len(peptides) + out = [False] * n + if not haystack: + return out + + by_peptide: dict[str, list[int]] = {} + for i, p in enumerate(peptides): + if not p: + continue + by_peptide.setdefault(p, []).append(i) + + if not by_peptide: + return out + + auto = ahocorasick.Automaton() + peptide_for_pid: list[str] = [] + for pid, pep in enumerate(by_peptide): + auto.add_word(pep, pid) + peptide_for_pid.append(pep) + + auto.make_automaton() + matched_pids: set[int] = set() + for _end_idx, pid in auto.iter(haystack): + matched_pids.add(pid) + + for pid in matched_pids: + pep = peptide_for_pid[pid] + for row_i in by_peptide[pep]: + out[row_i] = True + return out + + +def residue_token_count(prediction: Any, metrics: Metrics) -> int: + """Tokenizer residue count (``_split_peptide``), not raw string length.""" + if prediction is None: + return 0 + if isinstance(prediction, float) and math.isnan(prediction): + return 0 + if isinstance(prediction, list): + return len(prediction) + if not isinstance(prediction, str): + return 0 + text = prediction.strip() + if not text: + return 0 + try: + return len(metrics._split_peptide(text)) + except Exception: + return 0 + + +def filter_and_annotate_preds( + preds: pl.DataFrame, + haystack: str, + metrics: Metrics, + min_residue_length: int, +) -> pl.DataFrame: + """Filter and annotate Winnow predictions with proteome substring hits. + + Args: + preds: Polars DataFrame containing Winnow predictions. + haystack: String containing the proteome. + metrics: Metrics object for InstaNovo ``Metrics`` / ``_split_peptide``. + min_residue_length: Drop PSMs with fewer than this many tokenizer residues. + """ + n_tok = preds["prediction"].map_elements( + lambda x: residue_token_count(x, metrics), + return_dtype=pl.Int32, + ) + filtered = preds.with_columns(n_tok.alias("_n_residue_tokens")).filter( + pl.col("_n_residue_tokens") >= min_residue_length + ) + + processed = filtered["prediction"].map_elements( + lambda x: processed_peptide_for_match(x) if isinstance(x, str) else "", + return_dtype=pl.Utf8, + ) + hits = _batch_peptide_substring_hits(processed.to_list(), haystack) + return filtered.drop("_n_residue_tokens").with_columns( + pl.Series("proteome_hit", hits, dtype=pl.Boolean) + ) + + +def annotate_prediction_folder( + output_folder: Path | str, + fasta_path: Path | str, + metrics: Metrics, + *, + min_residue_length: int = 7, +) -> tuple[int, int, int]: + """Annotate Winnow predictions with proteome substring hits. + + Args: + output_folder: Path to the output folder containing predictions and metadata. + fasta_path: Path to the FASTA file containing the proteome. + metrics: Metrics object for InstaNovo ``Metrics`` / ``_split_peptide``. + min_residue_length: Drop PSMs with fewer than this many tokenizer residues. + """ + folder = Path(output_folder) + preds_path = folder / "preds_and_fdr_metrics.csv" + meta_path = folder / "metadata.csv" + if not preds_path.is_file(): + raise FileNotFoundError(f"Missing predictions file: {preds_path}") + if not meta_path.is_file(): + raise FileNotFoundError(f"Missing metadata file: {meta_path}") + + preds = pl.read_csv(preds_path) + if "prediction" not in preds.columns: + raise ValueError(f"'prediction' column missing in {preds_path}") + if "spectrum_id" not in preds.columns: + raise ValueError(f"'spectrum_id' column missing in {preds_path}") + + n_in = preds.height + haystack = load_proteome_haystack(fasta_path) + annotated = filter_and_annotate_preds( + preds, haystack, metrics, min_residue_length=min_residue_length + ) + n_kept = annotated.height + n_short = n_in - n_kept + + keep_ids = annotated.select("spectrum_id").unique() + meta = pl.read_csv(meta_path) + if "spectrum_id" not in meta.columns: + raise ValueError(f"'spectrum_id' column missing in {meta_path}") + meta_kept = meta.join(keep_ids, on="spectrum_id", how="inner") + + annotated.write_csv(preds_path) + meta_kept.write_csv(meta_path) + return n_in, n_short, n_kept + + +def _metrics_from_residues_yaml(residues_path: Path) -> Metrics: + with residues_path.open() as f: + data = yaml.safe_load(f) + residue_masses = data["residue_masses"] + return Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + +def _resolve_fasta_path(raw: str) -> Path: + p = Path(raw).expanduser() + if p.is_absolute(): + return p + return (_REPO_ROOT / p).resolve() + + +_PXD_RUN_PARENTS = ("PXD004452", "PXD006939", "PXD013868") + + +def _resolve_project_dir(predictions_root: Path, project: str) -> Path: + """Resolve per-run folder under *predictions_root* (flat or PXD*/run).""" + if "/" in project: + return predictions_root / project + direct = predictions_root / project + if direct.is_dir(): + return direct + for pxd in _PXD_RUN_PARENTS: + nested = predictions_root / pxd / project + if nested.is_dir(): + return nested + return direct + + +@app.command() +def main( + projects: Annotated[ + list[str], + typer.Argument(help="Project folder names under --predictions-root."), + ], + predictions_root: Annotated[ + Path, + typer.Option( + "--predictions-root", + "-p", + help="Root directory containing per-project prediction folders.", + ), + ], + fasta: Annotated[ + Path, + typer.Option( + "--fasta", + "-f", + help="FASTA file for proteome substring matching.", + ), + ], + residues_config: Annotated[ + Path, + typer.Option( + "--residues-config", + help="``residues.yaml`` for InstaNovo ``Metrics`` / ``_split_peptide``.", + ), + ] = _DEFAULT_RESIDUES, + min_residue_length: Annotated[ + int, + typer.Option( + "--min-residue-length", + "-m", + help="Drop PSMs with fewer than this many tokenizer residues.", + ), + ] = 7, + dry_run: Annotated[ + bool, + typer.Option( + "--dry-run", + help="Log only; do not write CSVs.", + ), + ] = False, +) -> None: + """Annotate Winnow predictions with proteome substring hits.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + project_list = [p.strip() for p in projects if p.strip()] + if not project_list: + raise typer.BadParameter("No projects specified.") + + fasta_path = _resolve_fasta_path(str(fasta)) + metrics = _metrics_from_residues_yaml(residues_config) + + for project in project_list: + out_dir = _resolve_project_dir(predictions_root, project) + logger.info( + "project=%s folder=%s fasta=%s min_residues=%s dry_run=%s", + project, + out_dir, + fasta_path, + min_residue_length, + dry_run, + ) + if dry_run: + continue + n_in, n_short, n_kept = annotate_prediction_folder( + out_dir, + fasta_path, + metrics, + min_residue_length=min_residue_length, + ) + logger.info( + "done project=%s rows_in=%d removed_short=%d kept=%d", + project, + n_in, + n_short, + n_kept, + ) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/benchmark_runtime.py b/paper_scripts/benchmark_runtime.py new file mode 100644 index 00000000..1997954d --- /dev/null +++ b/paper_scripts/benchmark_runtime.py @@ -0,0 +1,908 @@ +#!/usr/bin/env python3 +"""Benchmark wall-clock time and memory for the winnow prediction pipeline. + +Measures end-to-end processing time and peak memory, broken down by data loading, +per-feature computation, MLP calibration inference, and FDR / q-value computation. +Benchmarks full and/or no-Prosit feature configurations. +""" + +from __future__ import annotations + +import json +import logging +import os +import platform +import resource +import shutil +import sys +import tempfile +import time +import tracemalloc +from contextlib import contextmanager +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Annotated, Any, Dict, List, Optional, Sequence, Tuple + +import polars as pl +import torch +import typer + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.calibration.features.fragment_match import FragmentMatchFeatures +from winnow.calibration.features.retention_time import RetentionTimeFeature +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.fdr.nonparametric import NonParametricFDRControl + +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + +from no_prosit_dummy import train_or_load_dummy_calibrator # noqa: E402 + +logger = logging.getLogger(__name__) +app = typer.Typer( + add_completion=False, pretty_exceptions_show_locals=False, no_args_is_help=True +) + +PROSIT_FEATURE_CLASSES = (FragmentMatchFeatures, RetentionTimeFeature) + + +# --------------------------------------------------------------------------- +# Measurement helpers +# --------------------------------------------------------------------------- + + +@dataclass +class StageResult: + """Timing and memory result for a single pipeline stage.""" + + name: str + device: str + wall_time_s: float + peak_mem_mb: float + is_feature: bool = False + is_prosit: bool = False + columns: List[str] = field(default_factory=list) + + +@contextmanager +def measure(): + """Context manager that yields a dict populated with wall_time_s and peak_mem_mb on exit.""" + result: Dict[str, float] = {} + tracemalloc.start() + # Reset the peak so we measure only this block + tracemalloc.reset_peak() + t0 = time.perf_counter() + try: + yield result + finally: + result["wall_time_s"] = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + result["peak_mem_mb"] = peak / (1024 * 1024) + + +# --------------------------------------------------------------------------- +# Hardware info +# --------------------------------------------------------------------------- + + +def get_hardware_info() -> Dict[str, str]: + """Collect CPU, RAM, and GPU identifiers for the benchmark report.""" + info: Dict[str, str] = {} + + # CPU + cpu_name = None + try: + with open("/proc/cpuinfo") as f: + for line in f: + if line.startswith("model name"): + cpu_name = line.split(":", 1)[1].strip() + break + except OSError: + pass + if not cpu_name: + cpu_name = platform.processor() or "unknown" + info["cpu"] = cpu_name + info["cpu_cores"] = str(os.cpu_count() or "unknown") + + # RAM + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal"): + kb = int(line.split()[1]) + info["ram_gb"] = f"{kb / (1024**2):.0f}" + break + except OSError: + info["ram_gb"] = "unknown" + + # GPU + if torch.cuda.is_available(): + info["gpu"] = torch.cuda.get_device_name(0) + else: + info["gpu"] = "none" + + return info + + +# --------------------------------------------------------------------------- +# Pipeline stages +# --------------------------------------------------------------------------- + + +def write_combined_inputs( + spectrum_paths: Sequence[Path], + predictions_paths: Sequence[Path], + output_dir: Path, +) -> Tuple[Path, Path, str]: + """Concatenate spectrum/prediction inputs; return paths and a display label. + + A single pair is used as-is. Multiple pairs are joined with + ``diagonal_relaxed`` so labelled (with ``sequence``) and unlabelled schemas + can be combined. + """ + if len(spectrum_paths) != len(predictions_paths): + raise ValueError( + "The number of --spectrum-path and --predictions-path values must match" + ) + if not spectrum_paths: + raise ValueError( + "At least one --spectrum-path / --predictions-path pair is required" + ) + + if len(spectrum_paths) == 1: + return spectrum_paths[0], predictions_paths[0], str(spectrum_paths[0]) + + spectra = pl.concat( + [pl.read_parquet(path) for path in spectrum_paths], + how="diagonal_relaxed", + ) + preds = pl.concat( + [pl.read_csv(path) for path in predictions_paths], + how="diagonal_relaxed", + ) + output_dir.mkdir(parents=True, exist_ok=True) + spec_path = output_dir / "combined_spectra.parquet" + pred_path = output_dir / "combined_preds.csv" + spectra.write_parquet(spec_path) + preds.write_csv(pred_path) + label = " + ".join(str(path) for path in spectrum_paths) + logger.info( + "Combined %s spectra from %d inputs -> %s", + f"{len(spectra):,}", + len(spectrum_paths), + spec_path, + ) + return spec_path, pred_path, label + + +def load_dataset( + spectrum_path: str, + predictions_path: Optional[str], + data_loader_name: str, +) -> CalibrationDataset: + """Load and filter the dataset, returning a CalibrationDataset.""" + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + overrides = [f"data_loader={data_loader_name}"] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name="benchmark", + ): + cfg = compose(config_name="predict", overrides=overrides) + + data_loader = instantiate(cfg.data_loader) + dataset = data_loader.load( + data_path=spectrum_path, + predictions_path=predictions_path, + ) + + from winnow.scripts.main import _filter_dataset + + dataset = _filter_dataset(dataset) + return dataset + + +def compute_features_individually( + calibrator: ProbabilityCalibrator, + dataset: CalibrationDataset, +) -> List[StageResult]: + """Run each feature's prepare+compute with individual timing.""" + results: List[StageResult] = [] + + # Dependencies (currently all features return [], but measure for completeness) + for dep in calibrator.dependencies.values(): + with measure() as m: + dep.compute(dataset=dataset) + results.append( + StageResult( + name=f"Dependency: {dep.name}", + device="CPU", + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + ) + + for name, feat in calibrator.feature_dict.items(): + is_prosit = isinstance(feat, PROSIT_FEATURE_CLASSES) + device = "CPU + network" if is_prosit else "CPU" + + with measure() as m: + feat.prepare(dataset=dataset) + feat.compute(dataset=dataset) + + results.append( + StageResult( + name=f"Feature: {name}", + device=device, + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + is_feature=True, + is_prosit=is_prosit, + columns=list(feat.columns), + ) + ) + + return results + + +def run_mlp_inference( + calibrator: ProbabilityCalibrator, + dataset: CalibrationDataset, +) -> StageResult: + """Run MLP calibration inference.""" + if calibrator.network is None: + raise RuntimeError("Calibrator network is not loaded") + device = str(next(calibrator.network.parameters()).device) + with measure() as m: + calibrator.predict(dataset) + return StageResult( + name="MLP calibration inference", + device=device.upper() if device == "cpu" else device, + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + + +def run_fdr( + dataset: CalibrationDataset, + confidence_column: str = "calibrated_confidence", + fdr_threshold: float = 0.05, +) -> StageResult: + """Run FDR / q-value computation.""" + fdr_control = NonParametricFDRControl() + + with measure() as m: + fdr_control.fit(dataset=dataset.metadata[confidence_column]) + dataset.metadata = fdr_control.add_psm_pep(dataset.metadata, confidence_column) + dataset.metadata = fdr_control.add_psm_fdr(dataset.metadata, confidence_column) + dataset.metadata = fdr_control.add_psm_q_value( + dataset.metadata, confidence_column + ) + confidence_cutoff = fdr_control.get_confidence_cutoff(threshold=fdr_threshold) + _ = dataset.metadata[dataset.metadata[confidence_column] >= confidence_cutoff] + + return StageResult( + name="FDR / q-value computation", + device="CPU", + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + + +# --------------------------------------------------------------------------- +# Single benchmark run +# --------------------------------------------------------------------------- + + +@dataclass +class BenchmarkRun: + """Results from a single pipeline configuration.""" + + config_label: str + n_spectra: int + n_features: int + n_columns: int + stages: List[StageResult] + + @property + def total_wall_time_s(self) -> float: + """Sum of wall times across all recorded stages.""" + return sum(s.wall_time_s for s in self.stages) + + @property + def feature_wall_time_s(self) -> float: + """Sum of wall times for feature computation stages only.""" + return sum(s.wall_time_s for s in self.stages if s.is_feature) + + @property + def peak_mem_mb(self) -> float: + """Maximum peak memory across stages, in megabytes.""" + return max(s.peak_mem_mb for s in self.stages) if self.stages else 0.0 + + +def _model_matches_features(calibrator: ProbabilityCalibrator) -> bool: + """Check whether the loaded MLP input dim matches the current feature set.""" + if calibrator.network is None or calibrator.feature_mean is None: + return False + expected_dim = calibrator.feature_mean.shape[0] + actual_dim = 1 + len(calibrator.columns) # confidence + feature columns + return expected_dim == actual_dim + + +def _parse_koina_constants(raw: Optional[List[str]]) -> Optional[Dict[str, Any]]: + """Parse ``KEY=VALUE`` pairs into a dict, casting numeric strings.""" + if not raw: + return None + out: Dict[str, Any] = {} + for item in raw: + if "=" not in item: + raise typer.BadParameter( + f"Invalid --koina-input-constant format: '{item}'. Expected KEY=VALUE." + ) + key, value = item.split("=", 1) + try: + out[key] = int(value) + except ValueError: + try: + out[key] = float(value) + except ValueError: + out[key] = value + return out + + +def run_benchmark( + spectrum_path: str, + predictions_path: Optional[str], + model_path: str, + data_loader_name: str, + include_prosit: bool, + koina_input_constants: Optional[Dict[str, Any]] = None, +) -> BenchmarkRun: + """Execute the full prediction pipeline with per-stage timing.""" + config_label = "Full feature set" if include_prosit else "Without Prosit features" + + # Load calibrator + calibrator = ProbabilityCalibrator.load(pretrained_model_name_or_path=model_path) + + if koina_input_constants: + calibrator.apply_koina_model_input_overrides( + model_input_constants=koina_input_constants, + ) + + # Remove Prosit features if requested + if not include_prosit: + to_remove = [ + name + for name, feat in calibrator.feature_dict.items() + if isinstance(feat, PROSIT_FEATURE_CLASSES) + ] + for name in to_remove: + calibrator.remove_feature(name) + + stages: List[StageResult] = [] + + # Stage 1: Data loading + with measure() as m: + dataset = load_dataset(spectrum_path, predictions_path, data_loader_name) + n_spectra = len(dataset.metadata) + stages.append( + StageResult( + name="Data loading", + device="CPU", + wall_time_s=m["wall_time_s"], + peak_mem_mb=m["peak_mem_mb"], + ) + ) + + # Stage 2: Per-feature computation + feature_results = compute_features_individually(calibrator, dataset) + stages.extend(feature_results) + + n_features = len(calibrator.feature_dict) + n_columns = len(calibrator.columns) + + # Stage 3: MLP calibration inference + # The MLP input dimension must match the feature set. If features were + # removed (e.g. Prosit features dropped) but the model was trained with + # the full set, the dimensions won't match. In that case we skip MLP + + # FDR and note the mismatch -- these stages are sub-millisecond anyway + # and their cost is independent of the feature set used. + can_infer = _model_matches_features(calibrator) + if can_infer: + stages.append(run_mlp_inference(calibrator, dataset)) + # Stage 4: FDR / q-value + stages.append(run_fdr(dataset)) + else: + mean_dim = ( + calibrator.feature_mean.shape[0] + if calibrator.feature_mean is not None + else "unknown" + ) + print( + f" [note] MLP input dim ({mean_dim}) " + f"does not match current feature count " + f"({1 + len(calibrator.columns)}). " + f"Skipping MLP inference and FDR for this configuration.\n" + f" To benchmark these stages, supply a model trained with the " + f"matching feature set via --model-path-no-prosit." + ) + + return BenchmarkRun( + config_label=config_label, + n_spectra=n_spectra, + n_features=n_features, + n_columns=n_columns, + stages=stages, + ) + + +# --------------------------------------------------------------------------- +# Output formatting +# --------------------------------------------------------------------------- + + +def format_run(run: BenchmarkRun) -> str: + """Format a single benchmark run as a human-readable table.""" + lines: List[str] = [] + header = ( + f"=== Configuration: {run.config_label} " + f"({run.n_features} features, {run.n_columns} columns) ===" + ) + lines.append("") + lines.append(header) + lines.append("") + + col_w = [38, 16, 15, 15] + hdr = ( + f"{'Stage':<{col_w[0]}}| {'Device':<{col_w[1]}}| " + f"{'Wall time (s)':>{col_w[2]}}| {'Peak mem (MB)':>{col_w[3]}}" + ) + sep = ( + "-" * col_w[0] + + "|" + + "-" * (col_w[1] + 1) + + "|" + + "-" * (col_w[2] + 1) + + "|" + + "-" * (col_w[3] + 1) + ) + + lines.append(hdr) + lines.append(sep) + + feat_time = 0.0 + feat_mem = 0.0 + total_time = 0.0 + + def _row(label: str, device: str, t: float, mem: float) -> str: + return ( + f"{label:<{col_w[0]}}| {device:<{col_w[1]}}| " + f"{t:>{col_w[2]}.2f}| {mem:>{col_w[3]}.1f}" + ) + + for i, s in enumerate(run.stages): + lines.append(_row(s.name, s.device, s.wall_time_s, s.peak_mem_mb)) + total_time += s.wall_time_s + if s.is_feature: + feat_time += s.wall_time_s + feat_mem = max(feat_mem, s.peak_mem_mb) + + is_last_feature = s.is_feature and not any( + st.is_feature for st in run.stages[i + 1 :] + ) + if is_last_feature: + lines.append( + _row( + " Feature computation subtotal", + "", + feat_time, + feat_mem, + ) + ) + + lines.append(sep) + total_mem = max(s.peak_mem_mb for s in run.stages) if run.stages else 0.0 + lines.append(_row("End-to-end total", "", total_time, total_mem)) + + return "\n".join(lines) + + +def format_full_report( + hw: Dict[str, str], + mlp_device: str, + runs: List[BenchmarkRun], + n_spectra: int, + spectrum_path: str, +) -> str: + """Assemble the complete benchmark report.""" + lines: List[str] = [] + lines.append("=" * 88) + lines.append(" Winnow Pipeline Runtime Benchmark") + lines.append("=" * 88) + lines.append("") + lines.append( + f"Hardware: {hw['cpu']} ({hw['cpu_cores']} cores), " + f"{hw['ram_gb']} GB RAM, GPU: {hw['gpu']}" + ) + lines.append(f"Dataset: {n_spectra:,} spectra from {spectrum_path}") + lines.append(f"Calibrator MLP device: {mlp_device}") + + for run in runs: + lines.append(format_run(run)) + + lines.append("") + lines.append("Notes:") + lines.append( + '- "CPU + network" = gRPC calls to a Koina/Triton server for' + " Prosit-derived spectral predictions." + ) + lines.append(" No local GPU is used by winnow during prediction.") + lines.append( + "- The MLP runs on CPU after loading. GPU is only used during training" + " (not benchmarked here)." + ) + lines.append( + "- Koina predictions are not cached to disk between runs; batching is" + " handled internally by koinapy." + ) + peak_rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + lines.append(f"- Process peak RSS: {peak_rss_mb:.1f} MB") + lines.append("") + return "\n".join(lines) + + +def build_json_report( + hw: Dict[str, str], + mlp_device: str, + runs: List[BenchmarkRun], + n_spectra: int, + spectrum_path: str, +) -> Dict[str, Any]: + """Build a structured dict suitable for JSON serialisation.""" + report: Dict[str, Any] = { + "hardware": hw, + "dataset": { + "spectrum_path": spectrum_path, + "n_spectra": n_spectra, + }, + "mlp_device": mlp_device, + "configurations": [], + } + for run in runs: + cfg: Dict[str, Any] = { + "label": run.config_label, + "n_features": run.n_features, + "n_columns": run.n_columns, + "total_wall_time_s": run.total_wall_time_s, + "feature_wall_time_s": run.feature_wall_time_s, + "stages": [asdict(s) for s in run.stages], + } + report["configurations"].append(cfg) + report["process_peak_rss_mb"] = ( + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + ) + report["notes"] = { + "koina_caching": ( + "Koina predictions are not cached to disk; each invocation " + "re-queries the server." + ), + "koina_batching": ( + "Batching is handled internally by koinapy " + "(gRPC streaming to the Koina/Triton server)." + ), + "gpu_usage": ( + "No local GPU is used during prediction. GPU is only used " + "during training (not benchmarked here)." + ), + } + return report + + +def _validate_runtime_cli_args( + *, + train_spectrum_path: Optional[Path], + train_predictions_path: Optional[Path], + val_spectrum_path: Optional[Path], + val_predictions_path: Optional[Path], + model_path_no_prosit: Optional[str], + force_retrain: bool, +) -> bool: + """Validate train/val CLI combinations. Returns whether all train paths were set.""" + train_args = ( + train_spectrum_path, + train_predictions_path, + val_spectrum_path, + val_predictions_path, + ) + train_provided = [p is not None for p in train_args] + if any(train_provided) and not all(train_provided): + raise typer.BadParameter( + "Provide all of --train-spectrum-path, --train-predictions-path, " + "--val-spectrum-path, and --val-predictions-path, or none of them" + ) + if all(train_provided) and not model_path_no_prosit: + raise typer.BadParameter( + "--model-path-no-prosit is required when training or reusing a " + "no-Prosit dummy (pass the directory to load/save)" + ) + if force_retrain and not all(train_provided): + raise typer.BadParameter( + "--force-retrain requires labelled train/val paths for the dummy" + ) + return all(train_provided) + + +def _ensure_no_prosit_model( + *, + train_all_provided: bool, + train_spectrum_path: Optional[Path], + train_predictions_path: Optional[Path], + val_spectrum_path: Optional[Path], + val_predictions_path: Optional[Path], + model_path_no_prosit: Optional[str], + model_path: str, + data_loader: str, + force_retrain: bool, +) -> str: + """Train or reuse a no-Prosit calibrator; return the path to load for benchmarking.""" + if train_all_provided: + assert train_spectrum_path is not None + assert train_predictions_path is not None + assert val_spectrum_path is not None + assert val_predictions_path is not None + assert model_path_no_prosit is not None + train_or_load_dummy_calibrator( + train_spectrum_path=train_spectrum_path, + train_predictions_path=train_predictions_path, + val_spectrum_path=val_spectrum_path, + val_predictions_path=val_predictions_path, + data_loader_name=data_loader, + model_output_dir=Path(model_path_no_prosit), + force_retrain=force_retrain, + ) + return model_path_no_prosit + return model_path_no_prosit or model_path + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +@app.command() +def main( + spectrum_path: Annotated[ + Optional[list[Path]], + typer.Option( + "--spectrum-path", + help=( + "Spectrum parquet; repeat to concatenate multiple splits " + "(diagonal_relaxed)." + ), + ), + ] = None, + predictions_path: Annotated[ + Optional[list[Path]], + typer.Option( + "--predictions-path", + help="Predictions CSV paired with each --spectrum-path.", + ), + ] = None, + model_path: Annotated[ + str, + typer.Option( + "--model-path", + help=( + "Path to a local calibrator directory or HuggingFace model " + "identifier (default: InstaDeepAI/winnow-general-model)." + ), + ), + ] = "InstaDeepAI/winnow-general-model", + model_path_no_prosit: Annotated[ + Optional[str], + typer.Option( + "--model-path-no-prosit", + help=( + "Directory for the no-Prosit calibrator. With labelled train/val " + "paths, reuse a checkpoint here or train and save one. Without " + "train/val paths, load an existing calibrator from this path " + "(falls back to --model-path)." + ), + ), + ] = None, + train_spectrum_path: Annotated[ + Optional[Path], + typer.Option( + "--train-spectrum-path", + help="Labelled train spectra for the no-Prosit dummy calibrator.", + ), + ] = None, + train_predictions_path: Annotated[ + Optional[Path], + typer.Option( + "--train-predictions-path", + help="Train predictions CSV for the no-Prosit dummy calibrator.", + ), + ] = None, + val_spectrum_path: Annotated[ + Optional[Path], + typer.Option( + "--val-spectrum-path", + help="Labelled val spectra for the no-Prosit dummy calibrator.", + ), + ] = None, + val_predictions_path: Annotated[ + Optional[Path], + typer.Option( + "--val-predictions-path", + help="Val predictions CSV for the no-Prosit dummy calibrator.", + ), + ] = None, + force_retrain: Annotated[ + bool, + typer.Option( + "--force-retrain/--reuse-model", + help="Retrain the no-Prosit dummy even if a checkpoint already exists.", + ), + ] = False, + data_loader: Annotated[ + str, + typer.Option("--data-loader", help="Data loader to use (default: instanovo)."), + ] = "instanovo", + no_prosit: Annotated[ + bool, + typer.Option( + "--no-prosit", + help=( + "Only benchmark without Prosit/Koina features. When omitted, " + "both configurations (full and no-Prosit) are benchmarked." + ), + ), + ] = False, + full_only: Annotated[ + bool, + typer.Option( + "--full-only", + help="Only benchmark the full feature set (skip no-Prosit run).", + ), + ] = False, + koina_input_constant: Annotated[ + Optional[List[str]], + typer.Option( + "--koina-input-constant", + help=( + "Koina model input constant as KEY=VALUE (repeatable). " + "E.g. --koina-input-constant collision_energies=27 " + "--koina-input-constant fragmentation_types=HCD" + ), + ), + ] = None, + output_json: Annotated[ + Optional[Path], + typer.Option( + "--output-json", + help="Save structured results to a JSON file.", + ), + ] = None, + output_text: Annotated[ + Optional[Path], + typer.Option( + "--output-text", + help="Save the human-readable report table to a text file.", + ), + ] = None, +) -> None: + """Run configured pipeline benchmarks and print (optionally save) results.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + hw = get_hardware_info() + koina_constants = _parse_koina_constants(koina_input_constant) + run_full = not no_prosit + run_no_prosit = not full_only + train_all_provided = _validate_runtime_cli_args( + train_spectrum_path=train_spectrum_path, + train_predictions_path=train_predictions_path, + val_spectrum_path=val_spectrum_path, + val_predictions_path=val_predictions_path, + model_path_no_prosit=model_path_no_prosit, + force_retrain=force_retrain, + ) + + spectrum_paths = spectrum_path or [Path("examples/example_data/spectra.ipc")] + predictions_paths = predictions_path or [ + Path("examples/example_data/predictions.csv") + ] + + tmpdir: Optional[Path] = None + try: + if len(spectrum_paths) > 1: + tmpdir = Path(tempfile.mkdtemp(prefix="winnow_runtime_")) + combined_dir = tmpdir + else: + combined_dir = Path(".") + spectrum_file, predictions_file, spectrum_label = write_combined_inputs( + spectrum_paths, + predictions_paths, + combined_dir, + ) + spectrum_path_str = str(spectrum_file) + predictions_path_str = str(predictions_file) + + probe_calibrator = ProbabilityCalibrator.load( + pretrained_model_name_or_path=model_path + ) + if probe_calibrator.network is None: + raise RuntimeError("Calibrator network is not loaded") + mlp_device = str(next(probe_calibrator.network.parameters()).device) + del probe_calibrator + + runs: List[BenchmarkRun] = [] + n_spectra = 0 + + if run_full: + logger.info(">>> Benchmarking: Full feature set (including Prosit) ...") + result = run_benchmark( + spectrum_path=spectrum_path_str, + predictions_path=predictions_path_str, + model_path=model_path, + data_loader_name=data_loader, + include_prosit=True, + koina_input_constants=koina_constants, + ) + runs.append(result) + n_spectra = result.n_spectra + + if run_no_prosit: + logger.info(">>> Benchmarking: Without Prosit features ...") + no_prosit_model = _ensure_no_prosit_model( + train_all_provided=train_all_provided, + train_spectrum_path=train_spectrum_path, + train_predictions_path=train_predictions_path, + val_spectrum_path=val_spectrum_path, + val_predictions_path=val_predictions_path, + model_path_no_prosit=model_path_no_prosit, + model_path=model_path, + data_loader=data_loader, + force_retrain=force_retrain, + ) + result = run_benchmark( + spectrum_path=spectrum_path_str, + predictions_path=predictions_path_str, + model_path=no_prosit_model, + data_loader_name=data_loader, + include_prosit=False, + koina_input_constants=koina_constants, + ) + runs.append(result) + n_spectra = n_spectra or result.n_spectra + + report = format_full_report(hw, mlp_device, runs, n_spectra, spectrum_label) + print(report) + + if output_text is not None: + output_text.parent.mkdir(parents=True, exist_ok=True) + output_text.write_text(report + "\n") + logger.info("Text report saved to %s", output_text) + + if output_json is not None: + json_report = build_json_report( + hw, mlp_device, runs, n_spectra, spectrum_label + ) + output_json.parent.mkdir(parents=True, exist_ok=True) + with open(output_json, "w") as f: + json.dump(json_report, f, indent=2) + logger.info("JSON results saved to %s", output_json) + finally: + if tmpdir is not None: + shutil.rmtree(tmpdir, ignore_errors=True) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/benchmark_scaling.py b/paper_scripts/benchmark_scaling.py new file mode 100644 index 00000000..452874c0 --- /dev/null +++ b/paper_scripts/benchmark_scaling.py @@ -0,0 +1,480 @@ +#!/usr/bin/env python3 +"""Measure pipeline scaling without Koina/Prosit feature stages. + +Trains a small dummy calibrator on Beam + Token Score + Mass Error (Da) using +HeLa (or any labelled) train/val inputs, then runs data loading, non-Koina +feature computation, MLP inference, and FDR at multiple subsampled sizes of the +benchmark spectrum set. Produces a JSON file and a matplotlib scaling figure. +""" + +from __future__ import annotations + +import json +import logging +import random +import shutil +import sys +import tempfile +import time +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Any, Dict, List, Optional, Sequence, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +import seaborn as sns +import typer + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.fdr.nonparametric import NonParametricFDRControl + +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + +from no_prosit_dummy import train_or_load_dummy_calibrator # noqa: E402 + +logger = logging.getLogger(__name__) +app = typer.Typer( + add_completion=False, pretty_exceptions_show_locals=False, no_args_is_help=True +) + +plt.switch_backend("Agg") + +# Paul Tol "bright" palette (colorblind-safe) — same as plot_eval_results.py +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +DEFAULT_FRACTIONS = [0.1, 0.5, 1.0] +_SEED = 42 +_DEFAULT_MODEL_OUTPUT_DIR = Path("paper_results/scaling/dummy_model") +_DEFAULT_RESULTS_DIR = Path("analysis") +_DEFAULT_PLOTS_DIR = Path("analysis") + + +@contextmanager +def measure(): + """Context manager that yields a dict with wall_time_s on exit.""" + result: Dict[str, float] = {} + t0 = time.perf_counter() + try: + yield result + finally: + result["wall_time_s"] = time.perf_counter() - t0 + + +def write_subsampled_files( + spectrum_paths: Sequence[Path], + predictions_paths: Sequence[Path], + fractions: Sequence[float], + output_dir: Path, + seed: int = _SEED, +) -> List[Tuple[float, Path, Path, int]]: + """Write nested subsamples of aligned spectrum and prediction files. + + All input splits are combined, spectrum rows are shuffled once, and each + fraction takes a prefix of that ordering. Predictions are filtered by + spectrum ID so every benchmark size has aligned inputs. + """ + if len(spectrum_paths) != len(predictions_paths): + raise ValueError( + "The number of --spectrum-path and --predictions-path values must match" + ) + + spectra = pl.concat( + [pl.read_parquet(path) for path in spectrum_paths], + how="diagonal_relaxed", + ) + preds = pl.concat( + [pl.read_csv(path) for path in predictions_paths], + how="diagonal_relaxed", + ) + n = len(spectra) + indices = list(range(n)) + random.Random(seed).shuffle(indices) + + outputs: List[Tuple[float, Path, Path, int]] = [] + for fraction in sorted(fractions): + k = n if fraction >= 1.0 else max(1, int(n * fraction)) + spectra_sub = spectra[indices[:k]] + keep_ids = set(spectra_sub["spectrum_id"].to_list()) + preds_sub = preds.filter(pl.col("spectrum_id").is_in(keep_ids)) + + spec_path = output_dir / f"spectra_{fraction:.2f}.parquet" + pred_path = output_dir / f"preds_{fraction:.2f}.csv" + spectra_sub.write_parquet(spec_path) + preds_sub.write_csv(pred_path) + outputs.append((fraction, spec_path, pred_path, len(spectra_sub))) + + return outputs + + +def load_dataset( + spectrum_path: str, + predictions_path: str, + data_loader_name: str, +) -> CalibrationDataset: + """Load a dataset through the package data loader and filter invalid rows.""" + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + from winnow.scripts.main import _filter_dataset + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + overrides = [f"data_loader={data_loader_name}"] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name="benchmark_scaling", + ): + cfg = compose(config_name="predict", overrides=overrides) + + data_loader = instantiate(cfg.data_loader) + dataset = data_loader.load( + data_path=spectrum_path, + predictions_path=predictions_path, + ) + return _filter_dataset(dataset) + + +@dataclass +class ScalingPoint: + """Timing measurements for one dataset-size fraction.""" + + fraction: float + n_spectra: int + stage_times: Dict[str, float] = field(default_factory=dict) + + +def run_at_size( + spec_path: Path, + pred_path: Path, + calibrator: ProbabilityCalibrator, + data_loader_name: str, + fraction: float, +) -> ScalingPoint: + """Run the full pipeline from disk at a given dataset size.""" + stage_times: Dict[str, float] = {} + + with measure() as m: + dataset = load_dataset(str(spec_path), str(pred_path), data_loader_name) + stage_times["Data loading"] = m["wall_time_s"] + n = len(dataset.metadata) + + feature_total = 0.0 + for name, feat in calibrator.feature_dict.items(): + with measure() as m: + feat.prepare(dataset=dataset) + feat.compute(dataset=dataset) + stage_times[f"Feature: {name}"] = m["wall_time_s"] + feature_total += m["wall_time_s"] + stage_times["Feature computation (total)"] = feature_total + + with measure() as m: + calibrator.predict(dataset) + stage_times["MLP inference"] = m["wall_time_s"] + + fdr = NonParametricFDRControl() + col = "calibrated_confidence" + with measure() as m: + fdr.fit(dataset=dataset.metadata[col]) + dataset.metadata = fdr.add_psm_pep(dataset.metadata, col) + dataset.metadata = fdr.add_psm_fdr(dataset.metadata, col) + dataset.metadata = fdr.add_psm_q_value(dataset.metadata, col) + stage_times["FDR / q-value"] = m["wall_time_s"] + + stage_times["End-to-end"] = ( + stage_times["Data loading"] + + feature_total + + stage_times["MLP inference"] + + stage_times["FDR / q-value"] + ) + + return ScalingPoint(fraction=fraction, n_spectra=n, stage_times=stage_times) + + +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + """Save figure as both PNG and PDF.""" + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def plot_scaling(points: List[ScalingPoint], output_path: Path) -> None: + """Plot per-stage wall time vs. dataset size.""" + sizes = [p.n_spectra for p in points] + + stages_to_plot = [ + ("Data loading", "Data loading", "v", "-", _PALETTE[0]), + ("Feature computation", "Feature computation (total)", "o", "-", _PALETTE[2]), + ("MLP inference", "MLP inference", "^", "-", _PALETTE[1]), + ("FDR / q-value", "FDR / q-value", "D", "-", _PALETTE[3]), + ("End-to-end", "End-to-end", "s", "--", _PALETTE[5]), + ] + + fig, ax = plt.subplots(figsize=(6, 4)) + for label, key, marker, linestyle, colour in stages_to_plot: + times = [p.stage_times[key] for p in points] + ax.plot( + sizes, + times, + marker=marker, + linestyle=linestyle, + color=colour, + label=label, + alpha=0.7, + ) + + max_size = max(sizes) + max_total = max(p.stage_times["End-to-end"] for p in points) + ref_sizes = np.linspace(0, max_size, 50) + ref_times = max_total * (ref_sizes / max_size) + ax.plot( + ref_sizes, + ref_times, + ls=":", + color=_PALETTE[6], + linewidth=1, + label="Linear reference", + ) + + ax.set_xlabel("Number of spectra") + ax.set_ylabel("Wall time (s)") + ax.set_ylim(top=max_total * 1.05) + ax.set_title("Pipeline scaling\nexcluding Koina-dependent features") + ax.legend(loc="upper left", fontsize=9) + fig.tight_layout() + base = Path(str(output_path).removesuffix(".png").removesuffix(".pdf")) + _save_fig(fig, base) + logger.info("Wrote %s.png/.pdf", base) + + +def load_points_from_json(json_path: Path) -> List[ScalingPoint]: + """Load previously saved scaling measurements from JSON.""" + with open(json_path) as f: + data = json.load(f) + return [ + ScalingPoint( + fraction=p["fraction"], + n_spectra=p["n_spectra"], + stage_times=p["stage_times"], + ) + for p in data["points"] + ] + + +@app.command() +def main( + replot_json: Annotated[ + Optional[Path], + typer.Option( + "--replot-json", + help=( + "Replot from a saved benchmark_scaling.json without rerunning " + "benchmarks." + ), + metavar="PATH", + ), + ] = None, + spectrum_path: Annotated[ + Optional[list[Path]], + typer.Option( + "--spectrum-path", + help="Benchmark spectrum parquet; repeat to combine multiple splits.", + ), + ] = None, + predictions_path: Annotated[ + Optional[list[Path]], + typer.Option( + "--predictions-path", + help="Predictions CSV paired with each --spectrum-path.", + ), + ] = None, + train_spectrum_path: Annotated[ + Optional[Path], + typer.Option( + "--train-spectrum-path", + help="Labelled train spectra for the no-Prosit dummy calibrator.", + ), + ] = None, + train_predictions_path: Annotated[ + Optional[Path], + typer.Option( + "--train-predictions-path", + help="Train predictions CSV for the dummy calibrator.", + ), + ] = None, + val_spectrum_path: Annotated[ + Optional[Path], + typer.Option( + "--val-spectrum-path", + help="Labelled val spectra for the dummy calibrator.", + ), + ] = None, + val_predictions_path: Annotated[ + Optional[Path], + typer.Option( + "--val-predictions-path", + help="Val predictions CSV for the dummy calibrator.", + ), + ] = None, + model_output_dir: Annotated[ + Path, + typer.Option( + "--model-output-dir", + help="Where to save/reuse the no-Prosit dummy calibrator.", + ), + ] = _DEFAULT_MODEL_OUTPUT_DIR, + force_retrain: Annotated[ + bool, + typer.Option( + "--force-retrain/--reuse-model", + help="Retrain the dummy even if model-output-dir already has a checkpoint.", + ), + ] = False, + train_only: Annotated[ + bool, + typer.Option( + "--train-only/--run-benchmark", + help=( + "Only train or reuse the no-Prosit dummy calibrator; skip the " + "scaling benchmark." + ), + ), + ] = False, + data_loader: Annotated[ + str, + typer.Option("--data-loader", help="Data loader to use (default: instanovo)."), + ] = "instanovo", + fractions: Annotated[ + Optional[list[float]], + typer.Option( + "--fractions", + help="Dataset fractions to benchmark (default: 0.1 0.5 1.0).", + ), + ] = None, + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for benchmark_scaling.json."), + ] = _DEFAULT_RESULTS_DIR, + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for benchmark_scaling.png."), + ] = _DEFAULT_PLOTS_DIR, +) -> None: + """Run scaling benchmarks or replot from a saved JSON file.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + if replot_json is not None: + plots_dir.mkdir(parents=True, exist_ok=True) + plot_scaling( + load_points_from_json(replot_json), + plots_dir / "benchmark_scaling.png", + ) + return + + required: list[tuple[str, Path | list[Path] | None]] = [ + ("--train-spectrum-path", train_spectrum_path), + ("--train-predictions-path", train_predictions_path), + ("--val-spectrum-path", val_spectrum_path), + ("--val-predictions-path", val_predictions_path), + ] + if not train_only: + required.extend( + ( + ("--spectrum-path", spectrum_path), + ("--predictions-path", predictions_path), + ) + ) + missing = [name for name, value in required if value is None] + if missing: + raise typer.BadParameter( + f"{', '.join(missing)} required unless --replot-json is set" + ) + + assert train_spectrum_path is not None + assert train_predictions_path is not None + assert val_spectrum_path is not None + assert val_predictions_path is not None + + calibrator = train_or_load_dummy_calibrator( + train_spectrum_path=train_spectrum_path, + train_predictions_path=train_predictions_path, + val_spectrum_path=val_spectrum_path, + val_predictions_path=val_predictions_path, + data_loader_name=data_loader, + model_output_dir=model_output_dir, + force_retrain=force_retrain, + ) + if train_only: + logger.info("Train-only complete; dummy at %s", model_output_dir) + return + + assert spectrum_path is not None + assert predictions_path is not None + + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + frac_list = sorted(fractions if fractions is not None else DEFAULT_FRACTIONS) + tmpdir = Path(tempfile.mkdtemp(prefix="winnow_scaling_")) + points: List[ScalingPoint] = [] + + try: + logger.info("Preparing subsampled files ...") + file_info = write_subsampled_files( + spectrum_path, + predictions_path, + frac_list, + tmpdir, + seed=_SEED, + ) + for frac, spec_p, pred_p, n in file_info: + logger.info( + " %.0f%%: %s spectra -> %s, %s", + frac * 100, + f"{n:,}", + spec_p.name, + pred_p.name, + ) + logger.info(">>> Running at %.0f%% (%s spectra) ...", frac * 100, f"{n:,}") + point = run_at_size(spec_p, pred_p, calibrator, data_loader, frac) + points.append(point) + logger.info( + " %s spectra -> %.2f s total", + f"{point.n_spectra:,}", + point.stage_times["End-to-end"], + ) + for stage, t in point.stage_times.items(): + if stage not in ("Feature computation (total)", "End-to-end"): + logger.info(" %s: %.3f s", stage, t) + + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + json_path = results_dir / "benchmark_scaling.json" + json_data: Dict[str, Any] = { + "points": [ + { + "fraction": p.fraction, + "n_spectra": p.n_spectra, + "stage_times": p.stage_times, + } + for p in points + ], + "dummy_model_dir": str(model_output_dir), + } + with open(json_path, "w") as handle: + json.dump(json_data, handle, indent=2) + logger.info("Wrote %s", json_path) + + plot_scaling(points, plots_dir / "benchmark_scaling.png") + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/calibrator_generalisation_utils.py b/paper_scripts/calibrator_generalisation_utils.py new file mode 100644 index 00000000..8847e3dd --- /dev/null +++ b/paper_scripts/calibrator_generalisation_utils.py @@ -0,0 +1,115 @@ +"""Shared helpers for calibrator generalisation analysis.""" + +from __future__ import annotations + +import logging +import re +from pathlib import Path + +import polars as pl + +logger = logging.getLogger(__name__) + +HEPG2_SOURCE = "PXD019483" + +SPECIES_NAME_MAPPING: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + HEPG2_SOURCE: "HepG2", + "hepg2": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD014877": "$\\it{C.\\;elegans}$", +} + + +def extract_project_name(parquet_path: Path) -> str: + """Extract project name from ``dataset-helaqc-annotated-0000-0001.parquet``.""" + match = re.match(r"dataset-(.+?)-annotated", parquet_path.stem) + if match: + return match.group(1) + return parquet_path.stem + + +def build_experiment_source_mapping(biological_validation_dir: Path) -> dict[str, str]: + """Map every experiment in biological validation parquets to its source label.""" + mapping: dict[str, str] = {} + parquet_files = sorted(biological_validation_dir.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError( + f"No parquet files found in biological validation directory: " + f"{biological_validation_dir}" + ) + + for parquet_path in parquet_files: + project = extract_project_name(parquet_path) + experiments = ( + pl.scan_parquet(parquet_path) + .select("experiment_name") + .unique() + .collect()["experiment_name"] + .to_list() + ) + for experiment_name in experiments: + mapping[experiment_name] = project + + logger.info( + "Built experiment->source mapping for %d experiments across %d projects", + len(mapping), + len(parquet_files), + ) + return mapping + + +def annotate_train_source_labels( + train_parquet: Path, + train_predictions: Path, + biological_validation_dir: Path, +) -> None: + """Add a ``source`` column to the train parquet and predictions CSV. + + Experiments found in ``biological_validation_dir`` inherit that project name. + All other experiments are labelled as HepG2 (``PXD019483``). + """ + experiment_to_source = build_experiment_source_mapping(biological_validation_dir) + lookup = pl.DataFrame( + { + "experiment_name": list(experiment_to_source.keys()), + "source": list(experiment_to_source.values()), + } + ) + + spectra = pl.read_parquet(train_parquet) + if "source" not in spectra.columns: + spectra = spectra.join(lookup, on="experiment_name", how="left").with_columns( + pl.col("source").fill_null(HEPG2_SOURCE) + ) + spectra.write_parquet(train_parquet) + logger.info("Wrote source labels to %s", train_parquet) + else: + logger.info( + "Parquet already has source column, leaving %s unchanged", train_parquet + ) + + predictions = pl.read_csv(train_predictions) + if "source" not in predictions.columns: + source_by_spectrum = spectra.select("spectrum_id", "source") + predictions = predictions.join(source_by_spectrum, on="spectrum_id", how="left") + missing = predictions.filter(pl.col("source").is_null()) + if len(missing) > 0: + raise ValueError( + f"{len(missing)} prediction rows in {train_predictions} have no matching " + "spectrum_id in the train parquet" + ) + predictions.write_csv(train_predictions) + logger.info("Wrote source labels to %s", train_predictions) + else: + logger.info( + "Predictions CSV already has source column, leaving %s unchanged", + train_predictions, + ) diff --git a/paper_scripts/download_figshare_article.py b/paper_scripts/download_figshare_article.py new file mode 100644 index 00000000..433e92d7 --- /dev/null +++ b/paper_scripts/download_figshare_article.py @@ -0,0 +1,365 @@ +#!/usr/bin/env python3 +"""Download Figshare article files reconstructing folder paths from the API. + +Fetches public article metadata (``files`` plus ``folder_structure``) and writes +each file under ``output_dir`` at ``{folder}/{name}``, matching the Figshare +folder layout used for paper reproduction artefacts. + +By default targets article ``30147601`` version ``7``. Use ``--check-only`` to +verify expected relative paths already exist locally without downloading. +""" + +from __future__ import annotations + +import fnmatch +import json +import logging +import urllib.error +import urllib.request +from pathlib import Path +from typing import Annotated, Any + +import typer + +logger = logging.getLogger(__name__) + +FIGSHARE_API = "https://api.figshare.com/v2" +_DEFAULT_OUTPUT_DIR = Path("paper_data") + +app = typer.Typer( + add_completion=False, + pretty_exceptions_show_locals=False, + no_args_is_help=True, +) + + +def _configure_logging() -> None: + if logger.handlers: + return + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s")) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + + +def _article_url(article_id: int, version: int | None) -> str: + if version is None: + return f"{FIGSHARE_API}/articles/{article_id}" + return f"{FIGSHARE_API}/articles/{article_id}/versions/{version}" + + +def _http_get_json(url: str) -> dict[str, Any]: + """GET a JSON object from ``url`` via urllib.""" + request = urllib.request.Request( + url, + headers={"Accept": "application/json", "User-Agent": "winnow-paper-scripts"}, + ) + try: + with urllib.request.urlopen(request, timeout=60) as response: + payload = response.read() + except urllib.error.HTTPError as exc: + raise RuntimeError( + f"Figshare API HTTP {exc.code} for {url}: {exc.reason}" + ) from exc + except urllib.error.URLError as exc: + raise RuntimeError( + f"Figshare API request failed for {url}: {exc.reason}" + ) from exc + data = json.loads(payload.decode("utf-8")) + if not isinstance(data, dict): + raise RuntimeError( + f"Expected JSON object from {url}, got {type(data).__name__}" + ) + return data + + +def _fetch_article(article_id: int, version: int | None) -> dict[str, Any]: + """Load article metadata, preferring the versioned endpoint when set.""" + if version is not None: + url = _article_url(article_id, version) + try: + return _http_get_json(url) + except RuntimeError as exc: + logger.warning( + "Versioned endpoint failed (%s); falling back to current article metadata.", + exc, + ) + return _http_get_json(_article_url(article_id, None)) + + +def _relative_path(file_info: dict[str, Any], folder_structure: dict[str, Any]) -> str: + """Build the reconstructed relative path for a Figshare file entry.""" + file_id = str(file_info.get("id", "")) + folder = str(folder_structure.get(file_id, "") or "").strip("/") + name = str(file_info.get("name", "")) + if not name: + raise ValueError(f"Figshare file {file_id!r} has no name") + if folder: + return f"{folder}/{name}" + return name + + +def _matches_include(rel_path: str, include: list[str] | None) -> bool: + if not include: + return True + return any(fnmatch.fnmatch(rel_path, pattern) for pattern in include) + + +def _list_article_entries( + article: dict[str, Any], + include: list[str] | None, +) -> list[tuple[str, dict[str, Any]]]: + """Return ``(relative_path, file_info)`` pairs filtered by ``include`` globs.""" + folder_structure = article.get("folder_structure") or {} + if not isinstance(folder_structure, dict): + raise RuntimeError("Article metadata folder_structure must be an object") + files = article.get("files") or [] + if not isinstance(files, list): + raise RuntimeError("Article metadata files must be a list") + + entries: list[tuple[str, dict[str, Any]]] = [] + for file_info in files: + if not isinstance(file_info, dict): + continue + rel_path = _relative_path(file_info, folder_structure) + if not _matches_include(rel_path, include): + continue + entries.append((rel_path, file_info)) + entries.sort(key=lambda item: item[0]) + return entries + + +def _local_size_matches(path: Path, expected_size: int | None) -> bool: + if expected_size is None or not path.is_file(): + return False + return path.stat().st_size == int(expected_size) + + +def _download_file(download_url: str, dest: Path) -> None: + """Download ``download_url`` to ``dest`` using urllib, writing via a temp file.""" + dest.parent.mkdir(parents=True, exist_ok=True) + tmp_path = dest.with_name(f".{dest.name}.partial") + request = urllib.request.Request( + download_url, + headers={"User-Agent": "winnow-paper-scripts"}, + ) + try: + with urllib.request.urlopen(request, timeout=300) as response: + with tmp_path.open("wb") as handle: + while True: + chunk = response.read(1024 * 1024) + if not chunk: + break + handle.write(chunk) + tmp_path.replace(dest) + except urllib.error.URLError as exc: + if tmp_path.exists(): + tmp_path.unlink() + raise RuntimeError(f"Download failed for {download_url}: {exc.reason}") from exc + except OSError: + if tmp_path.exists(): + tmp_path.unlink() + raise + + +def _check_paths(output_dir: Path, relative_paths: list[str]) -> int: + """Return the number of missing paths under ``output_dir``.""" + missing = 0 + for rel in relative_paths: + path = output_dir / rel + if path.is_file(): + logger.info("OK %s", rel) + else: + logger.error("Missing %s", rel) + missing += 1 + return missing + + +def _paths_for_check_only( + article_id: int, + version: int | None, + include: list[str] | None, + require: list[str] | None, +) -> list[str]: + """Resolve relative paths to verify under ``--check-only``.""" + if require: + to_check = list(require) + if include: + to_check = [path for path in to_check if _matches_include(path, include)] + return to_check + + article = _fetch_article(article_id, version) + to_check = [rel for rel, _ in _list_article_entries(article, include)] + if not to_check: + logger.error("No article files matched the selection for --check-only.") + raise typer.Exit(code=1) + return to_check + + +def _download_or_skip_entry( + rel_path: str, + file_info: dict[str, Any], + output_dir: Path, + *, + dry_run: bool, + force: bool, +) -> str: + """Download one article file or skip it. Returns ``downloaded``, ``skipped``, or ``listed``.""" + size = file_info.get("size") + expected_size = int(size) if size is not None else None + download_url = file_info.get("download_url") + dest = output_dir / rel_path + size_label = str(expected_size) if expected_size is not None else "?" + + if dry_run: + logger.info("DRY-RUN %s (%s bytes)", rel_path, size_label) + return "listed" + + if not force and _local_size_matches(dest, expected_size): + logger.info("Skip %s (size matches)", rel_path) + return "skipped" + + if not download_url: + raise RuntimeError(f"No download_url for {rel_path!r}") + + logger.info("Download %s (%s bytes)", rel_path, size_label) + _download_file(str(download_url), dest) + if expected_size is not None and not _local_size_matches(dest, expected_size): + raise RuntimeError( + f"Downloaded size mismatch for {rel_path!r}: " + f"expected {expected_size}, got {dest.stat().st_size}" + ) + return "downloaded" + + +@app.command() +def main( + article_id: Annotated[ + int, + typer.Option("--article-id", help="Figshare article id."), + ] = 30147601, + version: Annotated[ + int | None, + typer.Option( + "--version", + help="Article version to fetch (uses /versions/{n} when set).", + ), + ] = 7, + output_dir: Annotated[ + Path, + typer.Option( + "--output-dir", + "-o", + help="Root directory for reconstructed article paths.", + ), + ] = _DEFAULT_OUTPUT_DIR, + include: Annotated[ + list[str] | None, + typer.Option( + "--include", + help=( + "Glob pattern(s) matched against reconstructed relative paths. " + "If unset, all files are selected." + ), + ), + ] = None, + dry_run: Annotated[ + bool, + typer.Option("--dry-run", help="List selected files without downloading."), + ] = False, + force: Annotated[ + bool, + typer.Option( + "--force", + help="Re-download even when a local file already matches remote size.", + ), + ] = False, + check_only: Annotated[ + bool, + typer.Option( + "--check-only", + help=( + "Verify expected relative paths exist under output-dir and exit " + "(1 if any are missing). Does not download." + ), + ), + ] = False, + require: Annotated[ + list[str] | None, + typer.Option( + "--require", + help=( + "Relative path(s) that must exist when --check-only is set. " + "If unset with --check-only, all selected article files are checked." + ), + ), + ] = None, +) -> None: + """Download or verify Figshare article files with reconstructed folder paths.""" + _configure_logging() + + if check_only: + to_check = _paths_for_check_only(article_id, version, include, require) + missing = _check_paths(output_dir, to_check) + if missing: + logger.error( + "%d of %d required path(s) missing under %s", + missing, + len(to_check), + output_dir, + ) + raise typer.Exit(code=1) + logger.info( + "All %d required path(s) present under %s", len(to_check), output_dir + ) + return + + if require: + logger.warning("--require is only used with --check-only; ignoring.") + + article = _fetch_article(article_id, version) + title = article.get("title", "?") + article_version = article.get("version", version) + entries = _list_article_entries(article, include) + logger.info( + "Article %s (version %s): %r — %d file(s) selected", + article_id, + article_version, + title, + len(entries), + ) + if not entries: + logger.warning("No files matched; nothing to do.") + return + + downloaded = 0 + skipped = 0 + for rel_path, file_info in entries: + outcome = _download_or_skip_entry( + rel_path, + file_info, + output_dir, + dry_run=dry_run, + force=force, + ) + if outcome == "downloaded": + downloaded += 1 + elif outcome == "skipped": + skipped += 1 + + if dry_run: + logger.info("Dry run complete (%d file(s) listed).", len(entries)) + return + + logger.info( + "Done: downloaded %d, skipped %d, total selected %d under %s", + downloaded, + skipped, + len(entries), + output_dir, + ) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/evaluate_calibrator_generalisation.py b/paper_scripts/evaluate_calibrator_generalisation.py new file mode 100644 index 00000000..98259bc3 --- /dev/null +++ b/paper_scripts/evaluate_calibrator_generalisation.py @@ -0,0 +1,384 @@ +"""Evaluate calibrator generalisation by training on one source dataset and testing on all others. + +Uses a labelled training matrix (HF ``general_model_training_set`` or equivalent) +whose parquet metadata already has a ``source`` column. For each source, trains a +fresh calibrator, evaluates it in-distribution (held-out 20 %) and +out-of-distribution (every other source), then saves a combined results CSV. +""" + +import logging +import re +import sys +from pathlib import Path +from typing import Annotated, Dict, List, Optional + +import numpy as np +import pandas as pd +import yaml +from rich.logging import RichHandler +import typer + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.calibration.features import ( + BeamFeatures, + FragmentMatchFeatures, + MassErrorDaFeature, + RetentionTimeFeature, + TokenScoreFeatures, +) +from winnow.datasets.calibration_dataset import CalibrationDataset +from winnow.datasets.data_loaders import InstaNovoDatasetLoader + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +logger = logging.getLogger("winnow.evaluate_generalization") +logger.setLevel(logging.INFO) +logger.propagate = False +logger.addHandler(RichHandler()) + +# --------------------------------------------------------------------------- +# Constants — loaded from the canonical Winnow YAML configs +# --------------------------------------------------------------------------- +SEED = 42 +TEST_SIZE = 0.2 + +_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "winnow" / "configs" + +with open(_CONFIGS_DIR / "residues.yaml") as _f: + RESIDUE_MASSES: dict[str, float] = yaml.safe_load(_f)["residue_masses"] + +with open(_CONFIGS_DIR / "data_loader" / "instanovo.yaml") as _f: + _instanovo_cfg = yaml.safe_load(_f) + RESIDUE_REMAPPING: dict[str, str] = _instanovo_cfg.get("residue_remapping", {}) + BEAM_COLUMNS: dict[str, str] | None = _instanovo_cfg.get("beam_columns") + +with open(_CONFIGS_DIR / "calibrator.yaml") as _f: + _calibrator_cfg = yaml.safe_load(_f) +with open(_CONFIGS_DIR / "koina.yaml") as _f: + _KOINA_CFG = yaml.safe_load(_f)["koina"] +_KOINA_CONSTRAINTS = _KOINA_CFG["constraints"] +_KOINA_INPUT_CONSTANTS = _KOINA_CFG.get("input_constants") or { + "collision_energies": 27, + "fragmentation_types": "HCD", +} +_UNSUPPORTED_RESIDUES: list[str] = _KOINA_CONSTRAINTS.get("unsupported_residues") or [] +_MAX_PRECURSOR_CHARGE: int = _KOINA_CONSTRAINTS["max_precursor_charge"] +_MAX_PEPTIDE_LENGTH: int = _KOINA_CONSTRAINTS["max_peptide_length"] +_INTENSITY_MODEL: str = _KOINA_CFG["intensity_model"] +_IRT_MODEL: str = _KOINA_CFG["irt_model"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_IRT_TRAIN_FRACTION_OVERRIDES: Dict[str, float] = { + "herceptin": 0.15, +} + +# Mirrors Makefile train-extra-small-mass-error-da / EXTRA_SMALL_* overrides. +_EXTRA_SMALL_FRAGMENT_EXCLUDE = [ + "spectral_angle", + "xcorr", + "complementary_ion_count", + "max_ion_gap", +] +_EXTRA_SMALL_BEAM_EXCLUDE = ["edit_distance"] + + +def initialise_calibrator( + *, + train_project: Optional[str] = None, +) -> ProbabilityCalibrator: + """Create a fresh calibrator matching train-extra-small-mass-error-da.""" + irt_train_fraction = _IRT_TRAIN_FRACTION_OVERRIDES.get(train_project or "", 0.1) + + calibrator = ProbabilityCalibrator( + hidden_dims=(50, 50), + dropout=0.3, + learning_rate=0.0001, + weight_decay=0.001, + max_epochs=1000, + batch_size=1024, + n_iter_no_change=10, + tol=0.0001, + seed=SEED, + val_early_stopping_max_psms=None, + val_subsample_seed=None, + ) + calibrator.add_feature(MassErrorDaFeature(residue_masses=RESIDUE_MASSES)) + calibrator.add_feature( + FragmentMatchFeatures( + mz_tolerance=20, + mz_tolerance_unit="ppm", + learn_from_missing=False, + intensity_model_name=_INTENSITY_MODEL, + max_precursor_charge=_MAX_PRECURSOR_CHARGE, + max_peptide_length=_MAX_PEPTIDE_LENGTH, + unsupported_residues=_UNSUPPORTED_RESIDUES, + model_input_constants=_KOINA_INPUT_CONSTANTS, + ) + ) + calibrator.add_feature( + RetentionTimeFeature( + train_fraction=irt_train_fraction, + min_train_points=3, + learn_from_missing=False, + irt_model_name=_IRT_MODEL, + max_peptide_length=_MAX_PEPTIDE_LENGTH, + unsupported_residues=_UNSUPPORTED_RESIDUES, + ) + ) + calibrator.add_feature(BeamFeatures()) + calibrator.add_feature(TokenScoreFeatures()) + # Former excluded_columns behaviour: train on a reduced feature subset. + training_columns = [ + col + for col in calibrator.columns + if col not in _EXTRA_SMALL_FRAGMENT_EXCLUDE + and col not in _EXTRA_SMALL_BEAM_EXCLUDE + ] + calibrator.set_training_feature_columns(training_columns) + return calibrator + + +def load_dataset(data_path: Path, predictions_path: Path) -> CalibrationDataset: + """Load the combined train_extra_small dataset.""" + logger.info("Loading dataset from %s and %s", data_path, predictions_path) + loader = InstaNovoDatasetLoader( + residue_masses=RESIDUE_MASSES, + residue_remapping=RESIDUE_REMAPPING, + beam_columns=BEAM_COLUMNS, + ) + return loader.load(data_path=data_path, predictions_path=predictions_path) + + +def subset_dataset(dataset: CalibrationDataset, idx: np.ndarray) -> CalibrationDataset: + """Return a row subset of *dataset* with aligned beam predictions.""" + meta = dataset.metadata.iloc[idx].reset_index(drop=True) + preds = ( + [dataset.predictions[i] for i in idx.tolist()] + if dataset.predictions is not None + else None + ) + return CalibrationDataset(metadata=meta, predictions=preds) + + +def split_dataset_by_source( + dataset: CalibrationDataset, +) -> Dict[str, CalibrationDataset]: + """Split a combined dataset into one CalibrationDataset per ``source`` label.""" + if "source" not in dataset.metadata.columns: + raise ValueError( + "Expected a 'source' column in the train parquet metadata. " + "Use HF general_model_training_set or another pre-labelled matrix." + ) + + datasets: Dict[str, CalibrationDataset] = {} + for source in sorted(dataset.metadata["source"].unique()): + idx = np.where(dataset.metadata["source"].values == source)[0] + datasets[source] = subset_dataset(dataset, idx) + return datasets + + +_MOD_RE = re.compile(r"\[UNIMOD:\d+\]") + + +def _peptide_key(tokens: object) -> str: + """Normalise a tokenised peptide to a modification-free, I/L-collapsed key. + + Matches the strategy in ``scripts/split_annotated_raw_parquets.py``: + strip UNIMOD modifications, normalise I→L. + """ + if not isinstance(tokens, list): + return "__MISSING__" + stripped = [_MOD_RE.sub("", tok).replace("I", "L") for tok in tokens] + return "".join(stripped) + + +def create_train_test_split( + dataset: CalibrationDataset, +) -> tuple[CalibrationDataset, CalibrationDataset]: + """Split a dataset 80/20 by peptide so no peptide appears in both folds.""" + meta = dataset.metadata + n = len(meta) + if n <= 1: + return dataset, dataset + + pep_keys = meta["sequence"].apply(_peptide_key) + unique_peptides = pep_keys.unique() + + rng = np.random.default_rng(SEED) + perm = rng.permutation(len(unique_peptides)) + n_train = int(len(unique_peptides) * (1 - TEST_SIZE)) + + train_peptides = set(unique_peptides[perm[:n_train]]) + train_mask = pep_keys.isin(train_peptides).values + + train_idx = np.where(train_mask)[0] + test_idx = np.where(~train_mask)[0] + + return subset_dataset(dataset, train_idx), subset_dataset(dataset, test_idx) + + +def evaluate_model( + model: ProbabilityCalibrator, + test_dataset: CalibrationDataset, + train_project: str, + test_project: str, + evaluation_type: str, +) -> pd.DataFrame: + """Run prediction and tag the results.""" + model.compute_features(test_dataset) + model.predict(test_dataset) + + results = test_dataset.metadata.copy() + results["trained_on_dataset"] = train_project + results["test_dataset"] = test_project + results["evaluation_type"] = evaluation_type + return results + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +_DEFAULT_MODEL_OUTPUT_DIR = Path("models/generalisation") +_DEFAULT_RESULTS_OUTPUT_DIR = Path("results/generalisation") +_DEFAULT_TRAIN_PARQUET = Path("train_extra_small/train.parquet") +_DEFAULT_TRAIN_PREDS = Path("train_extra_small/train_preds.csv") + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +@app.command() +def main( + train_parquet: Annotated[ + Path, typer.Option(help="Combined train parquet with a source column.") + ] = _DEFAULT_TRAIN_PARQUET, + train_predictions: Annotated[ + Path, typer.Option(help="Combined train predictions CSV.") + ] = _DEFAULT_TRAIN_PREDS, + model_output_dir: Annotated[ + Path, typer.Option(help="Directory to save trained models.") + ] = _DEFAULT_MODEL_OUTPUT_DIR, + results_output_dir: Annotated[ + Path, typer.Option(help="Directory to save evaluation results.") + ] = _DEFAULT_RESULTS_OUTPUT_DIR, +) -> None: + """Evaluate calibrator generalisation across source-labelled training datasets.""" + model_output_dir.mkdir(parents=True, exist_ok=True) + results_output_dir.mkdir(parents=True, exist_ok=True) + + if not train_parquet.exists(): + logger.error("Train parquet not found: %s", train_parquet) + raise typer.Exit(1) + if not train_predictions.exists(): + logger.error("Train predictions CSV not found: %s", train_predictions) + raise typer.Exit(1) + + full_dataset = load_dataset(train_parquet, train_predictions) + if "source" not in full_dataset.metadata.columns: + logger.error( + "Train parquet metadata is missing a required 'source' column: %s", + train_parquet, + ) + raise typer.Exit(1) + + datasets = split_dataset_by_source(full_dataset) + logger.info("Found %d source datasets: %s", len(datasets), list(datasets.keys())) + for source, dataset in datasets.items(): + logger.info(" %s: %d samples", source, len(dataset.metadata)) + + # Train-on-each, evaluate-on-all + all_results: List[pd.DataFrame] = [] + for train_project in datasets: + logger.info("=== Training on %s ===", train_project) + + train_ds, in_dist_test_ds = create_train_test_split(datasets[train_project]) + logger.info( + " train: %d, in-dist test: %d", + len(train_ds.metadata), + len(in_dist_test_ds.metadata), + ) + + calibrator = initialise_calibrator(train_project=train_project) + calibrator.fit(train_ds) + + model_path = model_output_dir / f"trained_on_{train_project}" + ProbabilityCalibrator.save(calibrator, model_path) + + # In-distribution evaluation + logger.info( + " Evaluating in-distribution on %s (%d samples)", + train_project, + len(in_dist_test_ds.metadata), + ) + all_results.append( + evaluate_model( + calibrator, + in_dist_test_ds, + train_project, + train_project, + "in_distribution", + ) + ) + + # Out-of-distribution evaluation + for test_project in datasets: + if test_project == train_project: + continue + test_ds = datasets[test_project] + logger.info( + " Evaluating out-of-distribution on %s (%d samples)", + test_project, + len(test_ds.metadata), + ) + all_results.append( + evaluate_model( + calibrator, + test_ds, + train_project, + test_project, + "out_of_distribution", + ) + ) + + # Combine and save + combined = pd.concat(all_results, ignore_index=True) + + # Drop large array columns to save space + array_cols = [c for c in ["mz_array", "intensity_array"] if c in combined.columns] + if array_cols: + combined = combined.drop(columns=array_cols) + + results_path = results_output_dir / "calibrator_generalisation_results.csv" + combined.to_csv(results_path, index=False) + logger.info("Results saved to %s", results_path) + + # Summary + logger.info("Evaluation summary:") + summary = ( + combined.groupby(["trained_on_dataset", "test_dataset", "evaluation_type"]) + .size() + .reset_index(name="num_samples") + ) + for _, row in summary.iterrows(): + logger.info( + " Trained on %s, tested on %s (%s): %d samples", + row["trained_on_dataset"], + row["test_dataset"], + row["evaluation_type"], + row["num_samples"], + ) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/fdr_tool_comparison_preprocess.py b/paper_scripts/fdr_tool_comparison_preprocess.py new file mode 100644 index 00000000..2f936042 --- /dev/null +++ b/paper_scripts/fdr_tool_comparison_preprocess.py @@ -0,0 +1,809 @@ +"""Shared preprocessing helpers for FDR tool comparisons. + +PSM comparison and the external peptide score-mixture share: +1. Method-specific load / NovoBoard mass-delta → ProForma conversion. +2. Pair-gated NovoBoard target-decoy filters (equal twin counts). +3. :func:`filter_prediction_table` / :func:`filter_novoboard_prediction_table`. + +Labelled correctness uses Novor token matching; proteome-hit proxies use +PTM-stripped I→L substring search against an I→L FASTA haystack. + +Peptide score-mixture only then: +4. :func:`max_score_per_peptide` (no re-filtering). +5. NovoBoard max-target → twin-decoy peptide TDC helpers. +""" + +from __future__ import annotations + +import logging +import re +from functools import lru_cache +from pathlib import Path +from typing import Iterable + +import numpy as np +import pandas as pd +import polars as pl +import yaml +from instanovo.utils.metrics import Metrics +from instanovo.utils.residues import ResidueSet + +from winnow.utils.proteome import ( + _batch_peptide_substring_hits, + processed_peptide_for_match, + residue_token_count, +) + +logger = logging.getLogger(__name__) + +_REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_RESIDUES_YAML = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" + +MIN_PEPTIDE_LENGTH = 8 +# Labelled / reference sets use Novor agreement, so short peptides are valid. +# Keep a non-empty-key floor only. Unlabelled sets keep MIN_PEPTIDE_LENGTH +# because correctness is proteome substring membership. +LABELLED_MIN_PEPTIDE_LENGTH = 1 +_UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") +_MOD_SQUARE = re.compile(r"\[.*?\]") +_MOD_PAREN = re.compile(r"\(.*?\)") +_NOVOBOARD_TO_PROFORMA = { + "C(+57.02)": "C[UNIMOD:4]", + "M(+15.99)": "M[UNIMOD:35]", + "N(+0.98)": "N[UNIMOD:7]", + "Q(+0.98)": "Q[UNIMOD:7]", + "S(+79.97)": "S[UNIMOD:21]", + "T(+79.97)": "T[UNIMOD:21]", + "Y(+79.97)": "Y[UNIMOD:21]", +} + + +def normalize_peptide_key(peptide: object) -> str: + """Normalise sequence-only peptide identity (strip PTMs, I→L).""" + if isinstance(peptide, list): + peptide = "".join(str(token) for token in peptide) + if pd.isna(peptide) or not isinstance(peptide, str): + return "" + if len(peptide) > 4 and peptide[1] == "." and peptide[-2] == ".": + peptide = peptide[2:-2] + seq = _MOD_SQUARE.sub("", peptide) + seq = _MOD_PAREN.sub("", seq) + seq = "".join(c for c in seq if c.isalpha()) + return seq.replace("I", "L") + + +def has_unsupported_unimod(peptide: object) -> bool: + """Return True when *peptide* still contains an unsupported ``[UNIMOD:n]`` token.""" + if pd.isna(peptide) or not isinstance(peptide, str): + return True + return bool(_UNIMOD_RE.search(peptide)) + + +def novoboard_to_proforma(peptide: object) -> object: + """Convert NovoBoard's supported mass-delta notation to ProForma.""" + if pd.isna(peptide) or not isinstance(peptide, str): + return peptide + converted = peptide + for novoboard_mod, proforma_mod in _NOVOBOARD_TO_PROFORMA.items(): + converted = converted.replace(novoboard_mod, proforma_mod) + return converted + + +def sequence_only_correct_prediction(sequence: object, prediction: object) -> bool: + """Full sequence equality after PTM stripping and I/L normalisation.""" + sequence_key = normalize_peptide_key(sequence) + prediction_key = normalize_peptide_key(prediction) + return bool(sequence_key) and sequence_key == prediction_key + + +def sequence_only_correctness_mask( + sequences: pd.Series, predictions: pd.Series +) -> np.ndarray: + """Vectorized strip-PTM I→L equality (not for labelled Novor eval).""" + return np.array( + [ + sequence_only_correct_prediction(sequence, prediction) + for sequence, prediction in zip(sequences, predictions) + ], + dtype=bool, + ) + + +def load_residue_masses(residues_yaml: Path | None = None) -> dict[str, float]: + """Load residue masses from Winnow's residues YAML.""" + path = residues_yaml if residues_yaml is not None else DEFAULT_RESIDUES_YAML + with path.open(encoding="utf-8") as handle: + return yaml.safe_load(handle)["residue_masses"] + + +@lru_cache(maxsize=4) +def _metrics_from_residue_masses_frozen( + residues_items: tuple[tuple[str, float], ...], +) -> Metrics: + residue_masses = dict(residues_items) + return Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + +def metrics_from_residue_masses(residue_masses: dict[str, float]) -> Metrics: + """Build an InstaNovo ``Metrics`` instance for Novor matching.""" + items = tuple(sorted((str(k), float(v)) for k, v in residue_masses.items())) + return _metrics_from_residue_masses_frozen(items) + + +def novor_correct_prediction( + sequence: object, + prediction: object, + metrics: Metrics, +) -> bool: + """Winnow/InstaNovo Novor correctness: full residue-token match.""" + if isinstance(sequence, list): + gt = sequence + elif pd.isna(sequence) or not isinstance(sequence, str) or not sequence: + return False + else: + gt = metrics._split_peptide(sequence) + + if isinstance(prediction, list): + pred = prediction + elif pd.isna(prediction) or not isinstance(prediction, str) or not prediction: + return False + else: + pred = metrics._split_peptide(prediction) + + if not gt or not pred: + return False + num_matches = metrics._novor_match(gt, pred) + return bool(num_matches == len(gt) == len(pred)) + + +def novor_correctness_mask( + sequences: pd.Series | Iterable[object], + predictions: pd.Series | Iterable[object], + *, + residue_masses: dict[str, float] | None = None, + metrics: Metrics | None = None, +) -> np.ndarray: + """Vectorized Novor correctness (same rule as ``DatabaseGroundedFDRControl.fit``).""" + if metrics is None: + masses = residue_masses if residue_masses is not None else load_residue_masses() + metrics = metrics_from_residue_masses(masses) + return np.array( + [ + novor_correct_prediction(sequence, prediction, metrics) + for sequence, prediction in zip(sequences, predictions) + ], + dtype=bool, + ) + + +def dedupe_best_score_per_peptide( + df: pd.DataFrame, peptide_col: str, score_col: str +) -> pd.DataFrame: + """Keep the highest-scoring row per peptide key.""" + return ( + df.sort_values(score_col, ascending=False) + .groupby(peptide_col, as_index=False) + .first() + ) + + +def compute_q_values(fdr: np.ndarray) -> np.ndarray: + """Convert ranked FDR estimates to q-values using suffix minima.""" + values = np.asarray(fdr, dtype=float) + q_values = np.empty_like(values) + fdr_min = np.inf + for i in range(len(values) - 1, -1, -1): + fdr_min = min(fdr_min, values[i]) + q_values[i] = fdr_min + return q_values + + +def monotonize_q_by_confidence( + confidence: np.ndarray, q_value: np.ndarray +) -> np.ndarray: + """Enforce non-increasing q-values when confidence increases.""" + order = np.argsort(-np.asarray(confidence, dtype=float)) + q_sorted = np.asarray(q_value, dtype=float)[order] + q_mono = np.empty_like(q_sorted) + q_min = np.inf + for i in range(len(q_sorted) - 1, -1, -1): + if q_sorted[i] > q_min: + q_mono[i] = q_min + else: + q_mono[i] = q_sorted[i] + q_min = q_sorted[i] + out = np.empty_like(q_mono) + out[order] = q_mono + return out + + +def add_peptide_key( + df: pd.DataFrame, + peptide_col: str, + *, + key_col: str = "peptide_key", +) -> pd.DataFrame: + """Append normalised peptide keys.""" + work = df.copy() + work[key_col] = work[peptide_col].map(normalize_peptide_key) + return work + + +def filter_prediction_table( + df: pd.DataFrame, + peptide_col: str, + *, + min_length: int = MIN_PEPTIDE_LENGTH, + key_col: str = "peptide_key", + drop_unsupported_mods: bool = True, + log: bool = True, +) -> pd.DataFrame: + """Drop unsupported mods and peptides shorter than *min_length* (normalised).""" + work = add_peptide_key(df, peptide_col, key_col=key_col) + before = len(work) + if drop_unsupported_mods: + work = work[~work[peptide_col].map(has_unsupported_unimod)].copy() + work = work[work[key_col].str.len() >= min_length].copy() + dropped = before - len(work) + if log and dropped: + logger.info( + "Filtered %d/%d rows (unsupported mods and/or length < %d) on %s", + dropped, + before, + min_length, + peptide_col, + ) + return work.reset_index(drop=True) + + +def filter_novoboard_prediction_table( + df: pd.DataFrame, + *, + peptide_col: str = "Peptide", + min_length: int = MIN_PEPTIDE_LENGTH, + key_col: str = "_peptide_key", + log: bool = True, +) -> pd.DataFrame: + """Convert NovoBoard modifications to ProForma, then apply shared filters.""" + work = df.copy() + work[peptide_col] = work[peptide_col].map(novoboard_to_proforma) + return filter_prediction_table( + work, + peptide_col, + min_length=min_length, + key_col=key_col, + log=log, + ) + + +def filter_novoboard_target_decoy_pairs( + target: pd.DataFrame, + decoy: pd.DataFrame, + *, + peptide_col: str = "Peptide", + min_length: int = MIN_PEPTIDE_LENGTH, + key_col: str = "_peptide_key", + log: bool = True, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Filter NovoBoard target/decoy as spectrum twins so pair counts stay equal. + + Both sides are converted to ProForma and passed through the shared + mod/length filter. Only ``_pair_key`` values present in **both** filtered + tables are kept, so dropping an unsupported-mod decoy also drops its + target (and vice versa). + """ + if "_pair_key" not in target.columns or "_pair_key" not in decoy.columns: + raise ValueError( + "filter_novoboard_target_decoy_pairs requires '_pair_key' on both tables" + ) + + target_f = filter_novoboard_prediction_table( + target, + peptide_col=peptide_col, + min_length=min_length, + key_col=key_col, + log=log, + ) + decoy_f = filter_novoboard_prediction_table( + decoy, + peptide_col=peptide_col, + min_length=min_length, + key_col=key_col, + log=log, + ) + + def _valid_pair_keys(series: pd.Series) -> set[str]: + keys = series.astype(str) + return {k for k in keys if k and k != "nan"} + + shared_keys = _valid_pair_keys(target_f["_pair_key"]) & _valid_pair_keys( + decoy_f["_pair_key"] + ) + target_out = target_f[target_f["_pair_key"].astype(str).isin(shared_keys)].copy() + decoy_out = decoy_f[decoy_f["_pair_key"].astype(str).isin(shared_keys)].copy() + n_target = target_out["_pair_key"].astype(str).nunique() + n_decoy = decoy_out["_pair_key"].astype(str).nunique() + if n_target != n_decoy: + raise AssertionError( + f"Pair filter left unequal twin counts: targets={n_target} decoys={n_decoy}" + ) + if log: + before_pairs = len( + _valid_pair_keys(target_f["_pair_key"]) + | _valid_pair_keys(decoy_f["_pair_key"]) + ) + logger.info( + "NovoBoard pair filter: %s → %s twin spectra " + "(target rows %s → %s, decoy rows %s → %s)", + before_pairs, + n_target, + len(target_f), + len(target_out), + len(decoy_f), + len(decoy_out), + ) + return target_out, decoy_out + + +def restrict_winnow_to_novoboard_spectra( + winnow: pd.DataFrame, novoboard: pd.DataFrame +) -> pd.DataFrame: + """Trim Winnow to NovoBoard twin-valid spectra under the subset invariant. + + After shared peptide filters and NovoBoard pair-gating, NovoBoard targets are + expected to be a subset of Winnow-filtered spectra (same InstaNovo + predictions; NovoBoard additionally drops pairs whose decoy fails). The + shared pool is therefore the NovoBoard spectrum set: only Winnow is trimmed. + + Raises: + AssertionError: If any NovoBoard spectrum is missing from Winnow. + """ + winnow_ids = set(winnow["spectrum_id"].astype(str)) + novoboard_ids = set(novoboard["spectrum_id"].astype(str)) + only_novoboard = novoboard_ids - winnow_ids + if only_novoboard: + examples = sorted(only_novoboard)[:5] + raise AssertionError( + "NovoBoard twin-valid spectra are not a subset of Winnow-filtered " + f"spectra ({len(only_novoboard)} missing); examples={examples}. " + "Expected identical InstaNovo predictions after ProForma remapping." + ) + winnow_shared = winnow[winnow["spectrum_id"].astype(str).isin(novoboard_ids)].copy() + logger.info( + "Shared spectrum pool: %d spectra (trimmed Winnow=%d; NovoBoard unchanged)", + len(novoboard_ids), + len(winnow) - len(winnow_shared), + ) + return winnow_shared + + +def assert_shared_prediction_keys( + winnow: pd.DataFrame, + novoboard: pd.DataFrame, + *, + winnow_peptide_col: str = "prediction", + novoboard_peptide_col: str = "Peptide", +) -> None: + """Require I/L-normalised prediction identity on the shared spectrum pool.""" + w_ids = winnow["spectrum_id"].astype(str) + nb_ids = novoboard["spectrum_id"].astype(str) + if w_ids.nunique() != len(winnow) or nb_ids.nunique() != len(novoboard): + raise AssertionError( + "Shared-pool tables must have one row per spectrum_id before " + f"prediction-key assert (winnow rows={len(winnow)} unique={w_ids.nunique()}, " + f"novoboard rows={len(novoboard)} unique={nb_ids.nunique()})" + ) + merged = ( + winnow[["spectrum_id", winnow_peptide_col]] + .assign(spectrum_id=w_ids) + .merge( + novoboard[["spectrum_id", novoboard_peptide_col]].assign( + spectrum_id=nb_ids + ), + on="spectrum_id", + how="inner", + validate="one_to_one", + ) + ) + if len(merged) != len(winnow) or len(merged) != len(novoboard): + raise AssertionError( + "Shared-pool spectrum_id join is not 1:1 " + f"(winnow={len(winnow)} novoboard={len(novoboard)} inner={len(merged)})" + ) + w_keys = merged[winnow_peptide_col].map(normalize_peptide_key) + nb_keys = merged[novoboard_peptide_col].map(normalize_peptide_key) + mismatch = w_keys != nb_keys + if bool(mismatch.any()): + bad = merged.loc[ + mismatch, ["spectrum_id", winnow_peptide_col, novoboard_peptide_col] + ] + examples = bad.head(5).to_dict(orient="records") + raise AssertionError( + "Winnow and NovoBoard predictions disagree after I/L-normalised " + f"peptide keys ({int(mismatch.sum())} spectra); examples={examples}" + ) + + +def label_series_by_spectrum_id(winnow: pd.DataFrame, label_col: str) -> pd.Series: + """Map ``spectrum_id`` → boolean label from a Winnow table.""" + if label_col not in winnow.columns: + raise KeyError(f"Missing label column {label_col!r}") + ids = winnow["spectrum_id"].astype(str) + if ids.duplicated().any(): + raise AssertionError( + f"Duplicate spectrum_id values when building {label_col} label map" + ) + return pd.Series( + winnow[label_col].astype(bool).to_numpy(), + index=ids, + name=label_col, + ) + + +def attach_labels_by_spectrum_id( + novoboard: pd.DataFrame, + label_by_id: pd.Series, + *, + label_col: str, +) -> pd.DataFrame: + """Attach a shared label column to NovoBoard rows by ``spectrum_id``.""" + out = novoboard.copy() + mapped = out["spectrum_id"].astype(str).map(label_by_id) + if mapped.isna().any(): + missing = out.loc[mapped.isna(), "spectrum_id"].astype(str).head(5).tolist() + raise AssertionError( + f"NovoBoard rows missing shared {label_col} labels; examples={missing}" + ) + out[label_col] = mapped.astype(bool) + return out + + +def _best_alc_per_pair_key(df: pd.DataFrame) -> pd.DataFrame: + """One highest-ALC row per ``_pair_key``.""" + work = df.dropna(subset=["ALC (%)", "_pair_key"]) + work = work[work["_pair_key"].astype(str) != "nan"] + return ( + work.sort_values("ALC (%)", ascending=False) + .groupby("_pair_key", as_index=False) + .first() + ) + + +def proteome_hit_mask( + peptides: pd.Series | Iterable[str], + haystack: str, + *, + min_length: int = MIN_PEPTIDE_LENGTH, +) -> np.ndarray: + """True when normalised peptide key (length ≥ *min_length*) hits the proteome.""" + keys = [normalize_peptide_key(p) for p in peptides] + eligible = [bool(k) and len(k) >= min_length for k in keys] + unique_keys = sorted({k for k, ok in zip(keys, eligible) if ok}) + hit_map: dict[str, bool] = {} + if unique_keys: + hits = _batch_peptide_substring_hits(unique_keys, haystack) + hit_map = dict(zip(unique_keys, hits)) + return np.array( + [bool(eligible[i] and hit_map.get(keys[i], False)) for i in range(len(keys))], + dtype=bool, + ) + + +def filter_and_annotate_preds( + preds: pl.DataFrame, + haystack: str, + metrics: Metrics, + min_residue_length: int, +) -> pl.DataFrame: + """Filter short peptides and annotate ``proteome_hit`` via ``winnow.utils.proteome``. + + Args: + preds: Polars frame with a ``prediction`` column. + haystack: I/L-normalised FASTA haystack from ``load_proteome_haystack``. + metrics: InstaNovo ``Metrics`` (uses ``metrics.residue_set`` for length). + min_residue_length: Drop PSMs with fewer than this many residue tokens. + """ + residue_set = metrics.residue_set + n_tok = preds["prediction"].map_elements( + lambda x: residue_token_count(x, residue_set), + return_dtype=pl.Int32, + ) + filtered = preds.with_columns(n_tok.alias("_n_residue_tokens")).filter( + pl.col("_n_residue_tokens") >= min_residue_length + ) + processed = filtered["prediction"].map_elements( + lambda x: processed_peptide_for_match(x) if isinstance(x, str) else "", + return_dtype=pl.Utf8, + ) + hits = _batch_peptide_substring_hits(processed.to_list(), haystack) + return filtered.drop("_n_residue_tokens").with_columns( + pl.Series("proteome_hit", hits, dtype=pl.Boolean) + ) + + +def max_score_per_peptide( + df: pd.DataFrame, + key_col: str, + score_col: str, +) -> pd.DataFrame: + """Keep the max-scoring row per peptide key (no filtering). + + Call after :func:`filter_prediction_table` or + :func:`filter_novoboard_prediction_table` so all methods share the same + filter → max-dedupe sequence in the peptide score-mixture benchmark. + """ + work = df.dropna(subset=[score_col, key_col]) + work = work[work[key_col].astype(str) != ""] + return dedupe_best_score_per_peptide(work, key_col, score_col).reset_index( + drop=True + ) + + +def confidence_to_log_prob(confidence: pd.Series | np.ndarray) -> np.ndarray: + """Map raw InstaNovo confidence in (0, 1] to Glissade-style log probabilities.""" + conf = np.asarray(confidence, dtype=float) + return np.log(np.clip(conf, 1e-300, 1.0)) + + +def _load_mgf_title_to_scan(mgf_path: Path) -> dict[str, str]: + """Parse TITLE→SCANS mapping from an MGF file.""" + mapping: dict[str, str] = {} + title: str | None = None + scan: str | None = None + with open(mgf_path, encoding="utf-8", errors="replace") as handle: + for line in handle: + value = line.strip() + if value.startswith("TITLE="): + title = value.removeprefix("TITLE=") + elif value.startswith("SCANS="): + scan = value.removeprefix("SCANS=") + elif value == "END IONS" and title is not None and scan is not None: + mapping[title] = scan + title = None + scan = None + return mapping + + +def attach_novoboard_pair_keys( + target: pd.DataFrame, + decoy: pd.DataFrame, + *, + novoboard_dir: Path, + split_prefix: str, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Attach ``_pair_key`` using Scan identity or decoy-MGF TITLE→SCANS mapping.""" + target_out = target.copy() + decoy_out = decoy.copy() + if "Scan" not in target_out.columns or "Scan" not in decoy_out.columns: + raise ValueError("NovoBoard tables require a 'Scan' column for twin pairing") + + target_key = target_out["Scan"].astype(str) + decoy_key = decoy_out["Scan"].astype(str) + best_decoy_key = decoy_key + best_overlap = len(set(target_key) & set(decoy_key)) + + mgf_path = novoboard_dir.parent / f"{split_prefix}.mgf" + if mgf_path.is_file(): + title_to_scan = _load_mgf_title_to_scan(mgf_path) + if title_to_scan: + mapped = decoy_key.map(title_to_scan) + mapped_overlap = len(set(target_key) & set(mapped.dropna())) + if mapped_overlap > best_overlap: + best_decoy_key = mapped + best_overlap = mapped_overlap + + target_out["_pair_key"] = target_key + decoy_out["_pair_key"] = best_decoy_key + logger.info( + "NovoBoard %s pair-key overlap: target=%d decoy=%d overlap=%d", + split_prefix, + target_key.nunique(dropna=True), + pd.Series(best_decoy_key).nunique(dropna=True), + best_overlap, + ) + return target_out, decoy_out + + +def prepare_novoboard_decoy_by_pair( + decoy_df: pd.DataFrame, + *, + min_length: int = MIN_PEPTIDE_LENGTH, + already_filtered: bool = False, +) -> pd.DataFrame: + """Index the best decoy row per ``_pair_key``. + + Args: + decoy_df: Decoy table with ``_pair_key``. Prefer pair-gated output from + :func:`filter_novoboard_target_decoy_pairs`. + already_filtered: When True, skip ProForma/mod/length filtering (caller + already pair-filtered). + """ + if "_pair_key" not in decoy_df.columns: + raise ValueError("NovoBoard twin TDC requires '_pair_key' on decoy") + if already_filtered: + decoy = decoy_df.copy() + if "_peptide_key" not in decoy.columns: + if "peptide_key" in decoy.columns: + decoy["_peptide_key"] = decoy["peptide_key"] + else: + decoy = add_peptide_key(decoy, "Peptide", key_col="_peptide_key") + elif "_peptide_key" in decoy_df.columns and decoy_df["_peptide_key"].notna().all(): + decoy = decoy_df.copy() + else: + decoy = filter_novoboard_prediction_table( + decoy_df, min_length=min_length, key_col="_peptide_key" + ) + decoy = decoy.dropna(subset=["ALC (%)", "_pair_key"]) + decoy = decoy[ + (decoy["_pair_key"].astype(str) != "nan") & (decoy["_peptide_key"] != "") + ] + return _best_alc_per_pair_key(decoy).set_index("_pair_key", drop=False) + + +def novoboard_psm_tdc( + target_df: pd.DataFrame, + decoy_df: pd.DataFrame, + *, + min_length: int = MIN_PEPTIDE_LENGTH, +) -> pd.DataFrame: + """Recompute NovoBoard's pooled PSM TDC after pair-gated filtering. + + Target and decoy are filtered as spectrum twins so unsupported-mod drops + remove the pair. Competition uses one best-ALC row per twin on each side, + guaranteeing ``sum(is_target) == sum(~is_target)``. + """ + target, decoy = filter_novoboard_target_decoy_pairs( + target_df, decoy_df, min_length=min_length + ) + target = _best_alc_per_pair_key(target).assign(is_target=True) + decoy = _best_alc_per_pair_key(decoy).assign(is_target=False) + n_target = int(target["_pair_key"].nunique()) + n_decoy = int(decoy["_pair_key"].nunique()) + if n_target != n_decoy or len(target) != len(decoy): + raise AssertionError( + f"PSM TDC unbalanced after pair gate: " + f"target_rows={len(target)} decoy_rows={len(decoy)} " + f"target_pairs={n_target} decoy_pairs={n_decoy}" + ) + + combined = pd.concat([target, decoy], ignore_index=True, sort=False) + combined = combined.sort_values( + ["ALC (%)", "is_target"], ascending=[False, False] + ).reset_index(drop=True) + return _assign_cumulative_tdc_fdr(combined) + + +def _prepare_targets_for_twin_tdc( + target_df: pd.DataFrame, + decoy_df: pd.DataFrame, + *, + min_length: int, + decoy_by_pair: pd.DataFrame | None, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Pair-gate or twin-filter targets and return ``(target, decoy_by_pair)``.""" + if "_pair_key" not in target_df.columns: + raise ValueError("NovoBoard twin TDC requires '_pair_key' on target") + if decoy_by_pair is None and "_pair_key" not in decoy_df.columns: + raise ValueError("NovoBoard twin TDC requires '_pair_key' on decoy") + + if decoy_by_pair is None: + target, decoy = filter_novoboard_target_decoy_pairs( + target_df, decoy_df, min_length=min_length, log=False + ) + decoy_by_pair = prepare_novoboard_decoy_by_pair( + decoy, min_length=min_length, already_filtered=True + ) + return target, decoy_by_pair + + target = target_df.copy() + if "_peptide_key" not in target.columns: + if "peptide_key" in target.columns: + target["_peptide_key"] = target["peptide_key"] + else: + target = filter_novoboard_prediction_table( + target, + min_length=min_length, + key_col="_peptide_key", + log=False, + ) + twin_keys = set(decoy_by_pair.index.astype(str)) + target = target[target["_pair_key"].astype(str).isin(twin_keys)] + return target, decoy_by_pair + + +def _assign_cumulative_tdc_fdr(combined: pd.DataFrame) -> pd.DataFrame: + """Add estimated FDR / q-value columns for a balanced target-decoy table.""" + out = combined.copy() + n_target = out["is_target"].astype(int).cumsum() + n_decoy = (~out["is_target"]).astype(int).cumsum() + out["estimated_fdr"] = np.divide( + n_decoy, + n_target, + out=np.ones(len(out), dtype=float), + where=n_target > 0, + ) + out["estimated_q_value"] = np.nan + target_mask = out["is_target"].to_numpy() + out.loc[target_mask, "estimated_q_value"] = compute_q_values( + out.loc[target_mask, "estimated_fdr"].to_numpy() + ) + return out + + +def novoboard_max_target_twin_decoy_tdc( + target_df: pd.DataFrame, + decoy_df: pd.DataFrame, + *, + min_length: int = MIN_PEPTIDE_LENGTH, + target_peptide_keys: set[str] | None = None, + decoy_by_pair: pd.DataFrame | None = None, + log_missing_twins: bool = True, +) -> pd.DataFrame: + """Peptide TDC: pair-gate, max ALC per target peptide, twin decoy by ``_pair_key``. + + Returns the combined ranked table with ``is_target``, ``estimated_fdr``, and + ``estimated_q_value`` (targets only). Targets without a twin-valid decoy are + dropped; the competition table is always 1:1. + + Args: + decoy_by_pair: Optional precomputed output of + :func:`prepare_novoboard_decoy_by_pair` from an already pair-gated + decoy table. When omitted, target/decoy are pair-filtered together. + """ + target, decoy_by_pair = _prepare_targets_for_twin_tdc( + target_df, + decoy_df, + min_length=min_length, + decoy_by_pair=decoy_by_pair, + ) + target = target.dropna(subset=["ALC (%)", "_pair_key"]) + target = target[ + (target["_pair_key"].astype(str) != "nan") & (target["_peptide_key"] != "") + ] + + if target_peptide_keys is not None: + target = target[target["_peptide_key"].isin(target_peptide_keys)] + + # Max-score only among twin-valid targets. + target_best = max_score_per_peptide(target, "_peptide_key", "ALC (%)") + pair_keys = target_best["_pair_key"].astype(str) + has_twin = pair_keys.isin(decoy_by_pair.index.astype(str)) + n_missing_twin = int((~has_twin).sum()) + if log_missing_twins and n_missing_twin: + logger.warning( + "NovoBoard twin-decoy TDC dropped %d/%d max-target peptides without twin", + n_missing_twin, + len(target_best), + ) + target_keep = target_best.loc[has_twin].copy() + if target_keep.empty: + return pd.DataFrame( + columns=[ + "spectrum_id", + "Peptide", + "ALC (%)", + "_peptide_key", + "_pair_key", + "is_target", + "estimated_fdr", + "estimated_q_value", + ] + ) + + decoy_keep = decoy_by_pair.loc[target_keep["_pair_key"].astype(str)].copy() + decoy_keep = decoy_keep.reset_index(drop=True) + target_keep = target_keep.assign(is_target=True).reset_index(drop=True) + decoy_keep = decoy_keep.assign(is_target=False) + if len(target_keep) != len(decoy_keep): + raise AssertionError( + f"Peptide TDC unbalanced: targets={len(target_keep)} decoys={len(decoy_keep)}" + ) + # Preserve 1:1 balance: one decoy row per retained target (do not dedupe decoys). + combined = pd.concat([target_keep, decoy_keep], ignore_index=True, sort=False) + combined = combined.sort_values( + ["ALC (%)", "is_target"], ascending=[False, False] + ).reset_index(drop=True) + return _assign_cumulative_tdc_fdr(combined) diff --git a/paper_scripts/fdr_tool_comparison_summaries.py b/paper_scripts/fdr_tool_comparison_summaries.py new file mode 100644 index 00000000..e6d27b5d --- /dev/null +++ b/paper_scripts/fdr_tool_comparison_summaries.py @@ -0,0 +1,538 @@ +"""Summary tables for Winnow / NovoBoard / Glissade FDR tool comparisons. + +Produces two long-form CSVs: + +- ``*_acceptance.csv``: accepted counts and recovery at q-value thresholds. +- ``*_error_gain.csv``: observed FDP, excess over nominal FDR, optional mean + absolute q-value deviation vs a database-grounded reference (calibrated-score + DBG for Winnow; raw-score / ALC DBG for NovoBoard and Glissade), and relative + gain/loss of a primary method vs each comparator. +""" + +from __future__ import annotations + +import logging +import sys +from pathlib import Path +from typing import Iterable, Sequence, cast + +import numpy as np +import pandas as pd + +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + +from fdr_tool_comparison_preprocess import compute_q_values # noqa: E402 + +logger = logging.getLogger(__name__) + +SUMMARY_THRESHOLDS: list[float] = [0.01, 0.05, 0.10] +# Match ``DatabaseGroundedFDRControl`` / PSM comparison default. +_DB_GROUNDED_DROP = 10 + +_KEY_COLS = ["dataset", "panel", "level", "method", "q_value_threshold"] + + +def _slug_comparator(name: str) -> str: + """Map a method label to a filesystem-/column-safe slug.""" + return ( + name.lower() + .replace(" ", "_") + .replace("(", "") + .replace(")", "") + .replace("-", "_") + ) + + +def acceptance_rows_from_q( + *, + dataset: str, + panel: str, + level: str, + method: str, + q_value: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, + label_mask: np.ndarray | None = None, + recovery_denom: int | None = None, +) -> list[dict[str, object]]: + """Build acceptance/yield rows for one method at each q-value threshold. + + Args: + dataset: Dataset key (e.g. ``helaqc``). + panel: Evaluation panel (e.g. ``labelled_test``, ``unlabelled``, ``external``). + level: ``psm`` or ``peptide``. + method: Method display label. + q_value: Per-row estimated q-values. + thresholds: Nominal FDR thresholds. + label_mask: Optional boolean correctness / proteome-hit labels aligned with + ``q_value``. When provided, ``n_correct`` is filled. + recovery_denom: Denominator for recovery percentage. Defaults to the number + of True labels when ``label_mask`` is given. + + Returns: + One dict per threshold. + """ + q = np.asarray(q_value, dtype=float) + labels = None if label_mask is None else np.asarray(label_mask, dtype=bool) + if labels is not None and len(labels) != len(q): + raise ValueError( + f"label_mask length {len(labels)} does not match q_value length {len(q)}" + ) + if recovery_denom is None and labels is not None: + recovery_denom = int(labels.sum()) + + rows: list[dict[str, object]] = [] + for threshold in thresholds: + valid = ~np.isnan(q) + accepted = valid & (q <= threshold) + n_accepted = int(accepted.sum()) + n_correct: float | int = np.nan + recovery_pct: float = np.nan + if labels is not None: + n_correct = int((accepted & labels).sum()) + if recovery_denom and recovery_denom > 0: + recovery_pct = 100.0 * float(n_correct) / float(recovery_denom) + rows.append( + { + "dataset": dataset, + "panel": panel, + "level": level, + "method": method, + "q_value_threshold": float(threshold), + "n_accepted": n_accepted, + "n_correct": n_correct, + "recovery_pct": recovery_pct, + } + ) + return rows + + +def observed_fdp_at_thresholds( + q_value: np.ndarray, + label_mask: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, +) -> list[float]: + """Return observed false-discovery proportion among accepted rows at each threshold. + + ``label_mask`` is True for correct (or proteome-hit) rows. Observed FDP is + ``1 - n_correct / n_accepted`` when any rows are accepted, else NaN. + """ + q = np.asarray(q_value, dtype=float) + labels = np.asarray(label_mask, dtype=bool) + if len(q) != len(labels): + raise ValueError( + f"label_mask length {len(labels)} does not match q_value length {len(q)}" + ) + out: list[float] = [] + for threshold in thresholds: + valid = ~np.isnan(q) + accepted = valid & (q <= threshold) + n_accepted = int(accepted.sum()) + if n_accepted == 0: + out.append(float("nan")) + continue + n_correct = int((accepted & labels).sum()) + out.append(1.0 - n_correct / n_accepted) + return out + + +def database_grounded_q_from_labels( + scores: np.ndarray, + labels: np.ndarray, + *, + drop: int = _DB_GROUNDED_DROP, +) -> np.ndarray: + """In-sample database-grounded q-values from ranked scores and boolean labels. + + Builds the empirical precision curve ``1 - cumsum(correct) / rank`` on scores + sorted descending (same construction as the proteome-hit shortcut in the PSM + comparison), drops the first *drop* ranks from the FDR map, assigns FDR by + score lookup, then converts to q-values. + + Args: + scores: Ranking scores (higher = more confident). + labels: Boolean correctness / hit labels aligned with *scores*. + drop: Leading ranks excluded from the FDR map (default 10). + + Returns: + q-value array aligned with *scores*. + """ + scores_a = np.asarray(scores, dtype=float) + labels_a = np.asarray(labels, dtype=bool) + n = len(scores_a) + if n == 0: + return np.asarray([], dtype=float) + if len(labels_a) != n: + raise ValueError( + f"labels length {len(labels_a)} does not match scores length {n}" + ) + + order = np.argsort(-scores_a, kind="mergesort") + precision = np.cumsum(labels_a[order].astype(float)) / np.arange(1, n + 1) + fdr_ranked = 1.0 - precision + drop_eff = min(drop, max(0, n - 1)) + fit_scores = scores_a[order][drop_eff:] + fit_fdr = fdr_ranked[drop_eff:] + n_fit = len(fit_scores) + + idx = np.searchsorted(-fit_scores, -scores_a, side="left") + fdr = np.empty(n, dtype=float) + below = (idx == n_fit) & (scores_a < fit_scores[-1]) + above = (idx == 0) & (scores_a > fit_scores[0]) + normal = ~(below | above) + fdr[below] = 1.0 + fdr[above] = float(fit_fdr[0]) + fdr[normal] = fit_fdr[np.clip(idx[normal], 0, n_fit - 1)] + + q_sorted = compute_q_values(fdr[order]) + q = np.empty(n, dtype=float) + q[order] = q_sorted + return q + + +def mean_abs_q_dev_vs_reference( + q_method: np.ndarray, + q_ref: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, +) -> list[float]: + """Mean absolute q-value deviation vs a row-aligned reference at each threshold. + + For each threshold, restrict to rows accepted by either method + (``q_method <= t`` or ``q_ref <= t``) with finite q for both, then report + ``mean(|q_method - q_ref|)``. + """ + q_m = np.asarray(q_method, dtype=float) + q_r = np.asarray(q_ref, dtype=float) + if len(q_m) != len(q_r): + raise ValueError( + f"q_ref length {len(q_r)} does not match q_method length {len(q_m)}" + ) + out: list[float] = [] + for threshold in thresholds: + both_finite = ~np.isnan(q_m) & ~np.isnan(q_r) + either_accepted = both_finite & ((q_m <= threshold) | (q_r <= threshold)) + if not np.any(either_accepted): + out.append(float("nan")) + continue + out.append(float(np.mean(np.abs(q_m[either_accepted] - q_r[either_accepted])))) + return out + + +def error_rows_from_q( + *, + dataset: str, + panel: str, + level: str, + method: str, + q_value: np.ndarray, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, + label_mask: np.ndarray | None = None, + q_ref: np.ndarray | None = None, + observed_fdp: Sequence[float] | None = None, +) -> list[dict[str, object]]: + """Build error-metric rows for one method (without relative-gain columns). + + Args: + observed_fdp: Optional precomputed FDP values (e.g. from a mixture + benchmark). When omitted, FDP is derived from ``label_mask`` if given. + """ + if observed_fdp is not None and len(observed_fdp) != len(thresholds): + raise ValueError("observed_fdp length must match thresholds") + if observed_fdp is None and label_mask is not None: + fdp_values = observed_fdp_at_thresholds(q_value, label_mask, thresholds) + elif observed_fdp is not None: + fdp_values = [float(x) for x in observed_fdp] + else: + fdp_values = [float("nan")] * len(thresholds) + + if q_ref is not None: + q_dev = mean_abs_q_dev_vs_reference(q_value, q_ref, thresholds) + else: + q_dev = [float("nan")] * len(thresholds) + + rows: list[dict[str, object]] = [] + for threshold, fdp, dev in zip(thresholds, fdp_values, q_dev): + fdp_f = float(fdp) + excess = fdp_f - float(threshold) if np.isfinite(fdp_f) else float("nan") + rows.append( + { + "dataset": dataset, + "panel": panel, + "level": level, + "method": method, + "q_value_threshold": float(threshold), + "observed_fdp": fdp_f, + "fdp_excess": excess, + "mean_abs_q_dev_vs_db": float(dev), + } + ) + return rows + + +def _relative_gain_column(value_col: str, comparator: str) -> str: + """Return the plan-specified relative-gain column name for *value_col*.""" + slug = _slug_comparator(comparator) + if value_col == "n_accepted": + return f"accepted_pct_vs_{slug}" + if value_col == "recovery_pct": + return f"recovery_pct_vs_{slug}" + if value_col == "observed_fdp": + return f"fdp_delta_vs_{slug}" + if "fdp" in value_col: + return f"{value_col}_delta_vs_{slug}" + return f"{value_col}_pct_vs_{slug}" + + +def _relative_gain_value( + value_col: str, primary_val: object, comparator_val: object +) -> float: + """Compute primary-vs-comparator gain for one metric cell.""" + if pd.isna(primary_val) or pd.isna(comparator_val): + return float("nan") + primary = float(cast("float | int | str", primary_val)) + comparator = float(cast("float | int | str", comparator_val)) + if value_col == "observed_fdp" or "fdp" in value_col: + return primary - comparator + if comparator == 0: + return float("nan") + return 100.0 * (primary - comparator) / comparator + + +def _fill_primary_relative_gains( + work: pd.DataFrame, + group: pd.DataFrame, + *, + primary_method: str, + comparators: Sequence[str], + value_cols: Sequence[str], + group_cols: Sequence[str], + group_keys: tuple[object, ...], +) -> None: + """Write relative-gain columns onto the primary-method row for one group.""" + primary_rows = group[group["method"] == primary_method] + if primary_rows.empty: + return + primary = primary_rows.iloc[0] + mask = pd.Series(True, index=work.index) + for col, val in zip(group_cols, group_keys): + mask &= work[col] == val + primary_idx = work.index[mask & (work["method"] == primary_method)] + if len(primary_idx) == 0: + return + idx = primary_idx[0] + + for comparator in comparators: + comp_rows = group[group["method"] == comparator] + if comp_rows.empty: + continue + comp = comp_rows.iloc[0] + for col in value_cols: + out_col = _relative_gain_column(col, comparator) + gain = _relative_gain_value(col, primary[col], comp[col]) + if np.isfinite(gain): + work.at[idx, out_col] = gain + + +def add_relative_gain_columns( + df: pd.DataFrame, + *, + primary_method: str, + comparators: Iterable[str], + value_cols: Sequence[str], + group_cols: Sequence[str] = ("dataset", "panel", "level", "q_value_threshold"), +) -> pd.DataFrame: + """Attach primary-vs-comparator relative columns onto a long-form metrics table. + + For ``n_accepted`` / ``recovery_pct``, writes + ``100 * (primary - comparator) / comparator``. + For ``observed_fdp``, writes the signed difference ``primary - comparator``. + Relative columns are filled only on primary-method rows. + """ + if df.empty: + return df.copy() + + work = df.copy() + comparator_list = list(comparators) + for col in value_cols: + for comparator in comparator_list: + work[_relative_gain_column(col, comparator)] = np.nan + + group_list = list(group_cols) + for keys, group in work.groupby(group_list, dropna=False, sort=False): + if not isinstance(keys, tuple): + keys = (keys,) + _fill_primary_relative_gains( + work, + group, + primary_method=primary_method, + comparators=comparator_list, + value_cols=value_cols, + group_cols=group_list, + group_keys=keys, + ) + return work + + +def merge_acceptance_and_error( + acceptance: pd.DataFrame, + error: pd.DataFrame, + *, + key_cols: Sequence[str] | None = None, +) -> pd.DataFrame: + """Join acceptance counts onto error rows for relative-gain construction.""" + if acceptance.empty or error.empty: + return error.copy() + keys = list(key_cols) if key_cols is not None else list(_KEY_COLS) + keys = [c for c in keys if c in acceptance.columns and c in error.columns] + cols = [ + c + for c in ("n_accepted", "n_correct", "recovery_pct") + if c in acceptance.columns + ] + return error.merge( + acceptance[keys + cols], + on=keys, + how="left", + ) + + +def finalise_error_gain_table( + acceptance: pd.DataFrame, + error: pd.DataFrame, + *, + primary_method: str, + comparators: Sequence[str], + key_cols: Sequence[str] | None = None, + group_cols: Sequence[str] | None = None, +) -> pd.DataFrame: + """Merge counts into error rows and add primary-vs-comparator relative columns.""" + merged = merge_acceptance_and_error(acceptance, error, key_cols=key_cols) + value_cols = [ + c for c in ("n_accepted", "recovery_pct", "observed_fdp") if c in merged.columns + ] + gain_groups = ( + tuple(group_cols) + if group_cols is not None + else ("dataset", "panel", "level", "q_value_threshold") + ) + with_gain = add_relative_gain_columns( + merged, + primary_method=primary_method, + comparators=comparators, + value_cols=value_cols, + group_cols=gain_groups, + ) + # Keep error-table identity columns first; drop helper count cols that duplicate + # the acceptance table except when used only for gain calculation. + drop_helpers = [c for c in ("n_accepted", "n_correct") if c in with_gain.columns] + return with_gain.drop(columns=drop_helpers, errors="ignore") + + +def write_summary_tables( + acceptance_df: pd.DataFrame, + error_df: pd.DataFrame, + output_dir: Path, + stem: str, +) -> tuple[Path, Path]: + """Write acceptance and error/gain CSVs under *output_dir*.""" + output_dir.mkdir(parents=True, exist_ok=True) + acceptance_path = output_dir / f"{stem}_acceptance.csv" + error_path = output_dir / f"{stem}_error_gain.csv" + acceptance_df.to_csv(acceptance_path, index=False) + error_df.to_csv(error_path, index=False) + logger.info("Wrote %s", acceptance_path) + logger.info("Wrote %s", error_path) + return acceptance_path, error_path + + +def summarise_holdout_results( + raw: pd.DataFrame, + *, + thresholds: Sequence[float] = SUMMARY_THRESHOLDS, + primary_method: str = "Winnow", + comparators: Sequence[str] = ("NovoBoard", "Glissade"), + panel: str = "score_mixture", + level: str = "peptide", + group_extra: Sequence[str] = (), +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Aggregate mixture-benchmark iterations into the two summary tables. + + Args: + raw: Per-iteration rows from ``external_peptide_holdout_results.csv``. + thresholds: Nominal FDR thresholds to retain. + primary_method: Method used for relative gain/loss columns. + comparators: Comparator method labels. + panel: Panel name written into the summary tables. + level: Identification level written into the summary tables. + group_extra: Extra columns to group by (e.g. ``pi0_target``). + + Returns: + ``(acceptance_df, error_gain_df)``. + """ + if raw.empty: + return pd.DataFrame(), pd.DataFrame() + + filtered = raw[raw["q_value_threshold"].isin(thresholds)].copy() + if filtered.empty: + return pd.DataFrame(), pd.DataFrame() + + extra = [c for c in group_extra if c in filtered.columns] + group_cols = ["dataset", *extra, "method", "q_value_threshold"] + gain_group_cols = ("dataset", *extra, "panel", "level", "q_value_threshold") + agg_kwargs: dict[str, tuple[str, str]] = { + "n_accepted": ("accepted_peptides", "mean"), + "n_correct": ("true_correct_peptides", "mean"), + "recovery_pct": ("correct_discovery_pct", "mean"), + "observed_fdp": ("observed_fdp", "mean"), + "n_accepted_std": ("accepted_peptides", "std"), + "observed_fdp_std": ("observed_fdp", "std"), + "recovery_pct_std": ("correct_discovery_pct", "std"), + } + if "mean_abs_q_dev_vs_db" in filtered.columns: + agg_kwargs["mean_abs_q_dev_vs_db"] = ("mean_abs_q_dev_vs_db", "mean") + agg = ( + filtered.groupby(group_cols, as_index=False) + .agg(**agg_kwargs) + .sort_values(group_cols) + .reset_index(drop=True) + ) + agg["panel"] = panel + agg["level"] = level + agg["fdp_excess"] = agg["observed_fdp"] - agg["q_value_threshold"] + if "mean_abs_q_dev_vs_db" not in agg.columns: + agg["mean_abs_q_dev_vs_db"] = np.nan + + id_cols = ["dataset", *extra, "panel", "level", "method", "q_value_threshold"] + acceptance = agg[ + id_cols + + [ + "n_accepted", + "n_correct", + "recovery_pct", + "n_accepted_std", + "recovery_pct_std", + ] + ].copy() + + error = agg[ + id_cols + + [ + "observed_fdp", + "fdp_excess", + "mean_abs_q_dev_vs_db", + "observed_fdp_std", + ] + ].copy() + + error_gain = finalise_error_gain_table( + acceptance.drop( + columns=["n_accepted_std", "recovery_pct_std"], errors="ignore" + ), + error, + primary_method=primary_method, + comparators=comparators, + key_cols=id_cols, + group_cols=gain_group_cols, + ) + return acceptance, error_gain diff --git a/paper_scripts/feature_subsets.py b/paper_scripts/feature_subsets.py new file mode 100644 index 00000000..34712b4d --- /dev/null +++ b/paper_scripts/feature_subsets.py @@ -0,0 +1,74 @@ +"""Feature column sets for subset calibrator training and evaluation.""" + +from __future__ import annotations + +from typing import TypedDict + + +class FeatureSubsetSpec(TypedDict): + """Metadata and column list for one feature-subset experiment.""" + + description: str + from_parquet: bool + columns: list[str] + + +# Full feature matrix columns (``confidence`` + features + ``correct`` label). +FULL_FEATURE_COLUMNS: list[str] = [ + "confidence", + "mass_error_ppm", + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", + "irt_error", + "margin", + "median_margin", + "entropy", + "z-score", + "edit_distance", + "min_token_probability", + "std_token_probability", +] + +_NO_XCORR_SPECTRAL = {"spectral_angle", "xcorr"} +_NO_FRAGMENT_SIMILARITY = _NO_XCORR_SPECTRAL | { + "complementary_ion_count", + "max_ion_gap", + "edit_distance", +} + + +def _columns_excluding(*, drop: set[str]) -> list[str]: + return [c for c in FULL_FEATURE_COLUMNS if c not in drop] + + +FEATURE_SUBSETS: dict[str, FeatureSubsetSpec] = { + "no_xcorr_spectral": { + "description": "Exclude spectral_angle and xcorr only.", + "from_parquet": True, + "columns": _columns_excluding(drop=_NO_XCORR_SPECTRAL), + }, + "no_fragment_similarity": { + "description": ( + "Exclude spectral_angle, xcorr, complementary_ion_count, " + "max_ion_gap, and edit_distance." + ), + "from_parquet": True, + "columns": _columns_excluding(drop=_NO_FRAGMENT_SIMILARITY), + }, + "mass_error_da_no_similarity": { + "description": ( + "Exclude mass_error_ppm and fragment-similarity features; " + "use mass_error_da (Daltons) instead. Requires full winnow train " + "(recomputes features from raw spectra)." + ), + "from_parquet": False, + "columns": list( + _columns_excluding(drop=_NO_FRAGMENT_SIMILARITY | {"mass_error_ppm"}) + ) + + ["mass_error_da"], + }, +} diff --git a/paper_scripts/no_prosit_dummy.py b/paper_scripts/no_prosit_dummy.py new file mode 100644 index 00000000..0dca98f4 --- /dev/null +++ b/paper_scripts/no_prosit_dummy.py @@ -0,0 +1,125 @@ +"""Shared no-Prosit dummy calibrator for paper timing benchmarks. + +Used by ``benchmark_scaling.py`` and ``benchmark_runtime.py`` so either can +reuse a checkpoint under the same directory or train and save one if missing. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import yaml + +from winnow.calibration.calibrator import ProbabilityCalibrator +from winnow.calibration.features import ( + BeamFeatures, + MassErrorDaFeature, + TokenScoreFeatures, +) +from winnow.datasets.calibration_dataset import CalibrationDataset + +logger = logging.getLogger(__name__) + +_SEED = 42 +_REPO_ROOT = Path(__file__).resolve().parent.parent +_CONFIGS_DIR = _REPO_ROOT / "winnow" / "configs" + + +def _load_residue_masses() -> dict[str, float]: + with open(_CONFIGS_DIR / "residues.yaml") as handle: + return yaml.safe_load(handle)["residue_masses"] + + +def load_dataset( + spectrum_path: str, + predictions_path: str, + data_loader_name: str, +) -> CalibrationDataset: + """Load a labelled dataset through the package data loader and filter rows.""" + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + from winnow.scripts.main import _filter_dataset + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + overrides = [f"data_loader={data_loader_name}"] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name="no_prosit_dummy", + ): + cfg = compose(config_name="predict", overrides=overrides) + + data_loader = instantiate(cfg.data_loader) + dataset = data_loader.load( + data_path=spectrum_path, + predictions_path=predictions_path, + ) + return _filter_dataset(dataset) + + +def build_no_prosit_calibrator() -> ProbabilityCalibrator: + """Build a small calibrator with only non-Koina features.""" + calibrator = ProbabilityCalibrator( + hidden_dims=(50, 50), + dropout=0.3, + learning_rate=0.0001, + weight_decay=0.001, + max_epochs=50, + batch_size=1024, + n_iter_no_change=5, + tol=0.0001, + seed=_SEED, + val_early_stopping_max_psms=None, + val_subsample_seed=None, + ) + calibrator.add_feature(MassErrorDaFeature(residue_masses=_load_residue_masses())) + calibrator.add_feature(BeamFeatures()) + calibrator.add_feature(TokenScoreFeatures()) + return calibrator + + +def checkpoint_exists(model_dir: Path) -> bool: + """Return True if ``model_dir`` has a loadable calibrator checkpoint.""" + return (model_dir / "config.json").is_file() and ( + model_dir / "model.safetensors" + ).is_file() + + +def train_or_load_dummy_calibrator( + *, + train_spectrum_path: Path, + train_predictions_path: Path, + val_spectrum_path: Path, + val_predictions_path: Path, + data_loader_name: str, + model_output_dir: Path, + force_retrain: bool = False, +) -> ProbabilityCalibrator: + """Train a no-Prosit dummy on labelled data, or reuse a checkpoint.""" + model_output_dir.mkdir(parents=True, exist_ok=True) + if checkpoint_exists(model_output_dir) and not force_retrain: + logger.info("Reusing dummy calibrator at %s", model_output_dir) + return ProbabilityCalibrator.load( + pretrained_model_name_or_path=str(model_output_dir) + ) + + logger.info( + "Training no-Prosit dummy calibrator (Beam + Token Score + Mass Error Da)" + ) + calibrator = build_no_prosit_calibrator() + train_ds = load_dataset( + str(train_spectrum_path), str(train_predictions_path), data_loader_name + ) + val_ds = load_dataset( + str(val_spectrum_path), str(val_predictions_path), data_loader_name + ) + logger.info( + " train=%d rows, val=%d rows", len(train_ds.metadata), len(val_ds.metadata) + ) + calibrator.fit(train_ds, val_ds, progress_bar=True) + ProbabilityCalibrator.save(calibrator, model_output_dir) + logger.info("Saved dummy calibrator to %s", model_output_dir) + return calibrator diff --git a/paper_scripts/plot_ablation_summary.py b/paper_scripts/plot_ablation_summary.py new file mode 100644 index 00000000..20626b52 --- /dev/null +++ b/paper_scripts/plot_ablation_summary.py @@ -0,0 +1,674 @@ +#!/usr/bin/env python3 +"""Bar charts of ablation calibration metrics from ``ablation_summary.csv``. + +The bar charts use ECE in the top 10% of PSMs ranked by calibrated +score. Tabular metrics use tail ECEat 5% and 10% FDR. +""" + +from __future__ import annotations + +import json +import logging +import sys +from pathlib import Path +from typing import Annotated, Literal + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +from plot_eval_results import ( # noqa: E402 + _display_name, + _fit_database_grounded_fdr, + _save_fig, + _style_ax, +) +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +# Paul Tol qualitative palette (colour-blind safe) — canonical ablation colours. +_ABLATION_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#EE7733", + "#0077BB", + "#33BBEE", + "#CC3311", +] + +# Ablation summary keys → ``plot_eval_results.DATASET_DISPLAY_NAMES`` keys. +_ABLATION_DATASET_KEYS: dict[str, str] = { + "Arabidopsis": "01747_C01_P018218_S00_I00_N03_R1", + "Astral": "astral", + "HCT116": "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46", +} + +logger = logging.getLogger(__name__) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +ABLATION_CONFIG_ORDER: list[str] = [ + "Confidence only", + "Confidence + mass error", + "Confidence + iRT error", + "Confidence + token-level", + "Confidence + beam search", + "Confidence + fragment matching", + "All features", +] + + +def assign_ablation_colors(config_names: list[str]) -> dict[str, str]: + """Assign a unique colour per ablation config (no palette wrap).""" + if len(config_names) > len(_ABLATION_PALETTE): + raise ValueError( + f"Need {len(config_names)} ablation colours but only " + f"{len(_ABLATION_PALETTE)} defined." + ) + return {name: _ABLATION_PALETTE[i] for i, name in enumerate(config_names)} + + +def ordered_ablation_configs(present: set[str]) -> list[str]: + """Canonical config order for ablation plots and colour assignment.""" + ordered = [c for c in ABLATION_CONFIG_ORDER if c in present] + extra = sorted(present - set(ordered)) + return ordered + extra + + +_CONFIG_SHORT_LABELS: dict[str, str] = { + "Confidence only": "Confidence", + "Confidence + mass error": "+ Mass", + "Confidence + iRT error": "+ iRT", + "Confidence + token-level": "+ Token", + "Confidence + beam search": "+ Beam", + "Confidence + fragment matching": "+ Fragment", + "All features": "All features", +} + +MetricName = Literal[ + "tail_ECE", + "tail_ECE@5%FDR", + "tail_ECE@10%FDR", + "ECE", + "Brier", + "PR_AUC", + "fdr_bias@5%FDR", + "fdr_bias@10%FDR", + "q_dev@5%FDR", + "q_dev@10%FDR", +] + +TAIL_ECE_FRACTION = 0.1 +FDR_TAIL_THRESHOLDS: tuple[float, ...] = (0.05, 0.10) +TAIL_ECE_COLUMN_BY_THRESHOLD: dict[float, str] = { + 0.05: "tail_ECE@5%FDR", + 0.10: "tail_ECE@10%FDR", +} +Q_DEV_COLUMN_BY_THRESHOLD: dict[float, str] = { + 0.05: "q_dev@5%FDR", + 0.10: "q_dev@10%FDR", +} +FDR_BIAS_COLUMN_BY_THRESHOLD: dict[float, str] = { + 0.05: "fdr_bias@5%FDR", + 0.10: "fdr_bias@10%FDR", +} + +_DEFAULT_SUMMARY = Path("results/ablations/ablation_summary.csv") +_DEFAULT_OUTPUT_DIR = Path("results/ablations/plots") + + +def load_ablation_summary(path: Path) -> pd.DataFrame: + """Load ``ablation_summary.csv`` or ``.json``.""" + if not path.is_file(): + raise FileNotFoundError(path) + if path.suffix == ".json": + with open(path) as f: + return pd.DataFrame(json.load(f)) + return pd.read_csv(path) + + +def compute_ece( + pred: np.ndarray, + labels: np.ndarray, + n_bins: int = 10, +) -> float: + """Expected calibration error.""" + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_indices = np.digitize(pred, bins) - 1 + bin_indices = np.clip(bin_indices, 0, n_bins - 1) + ece = 0.0 + for b in range(n_bins): + mask = bin_indices == b + if mask.sum() == 0: + continue + avg_conf = pred[mask].mean() + avg_acc = labels[mask].mean() + ece += mask.sum() / len(pred) * abs(avg_conf - avg_acc) + return float(ece) + + +def compute_tail_ece( + pred: np.ndarray, + labels: np.ndarray, + tail_fraction: float = TAIL_ECE_FRACTION, + n_bins: int = 10, +) -> float: + """ECE in the top score-ranked fraction of PSMs.""" + threshold_idx = max(1, int(len(pred) * (1 - tail_fraction))) + sorted_indices = np.argsort(pred) + tail_mask = np.zeros(len(pred), dtype=bool) + tail_mask[sorted_indices[threshold_idx:]] = True + if tail_mask.sum() == 0: + return 0.0 + return compute_ece(pred[tail_mask], labels[tail_mask], n_bins=n_bins) + + +def compute_tail_ece_at_fdr( + pred: np.ndarray, + labels: np.ndarray, + fdr_threshold: float, + *, + fdr_ctrl: NonParametricFDRControl | None = None, + n_bins: int = 10, +) -> float: + """ECE among PSMs accepted at a non-parametric FDR threshold.""" + if len(pred) == 0: + return float("nan") + + if fdr_ctrl is None: + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=pd.Series(pred, name="score")) + + cutoff = fdr_ctrl.get_confidence_cutoff(threshold=fdr_threshold) + if np.isnan(cutoff): + return float("nan") + + mask = pred >= cutoff + if not mask.any(): + return float("nan") + + return compute_ece(pred[mask], labels[mask], n_bins=n_bins) + + +def compute_tail_ece_at_fdr_thresholds( + df: pd.DataFrame, + *, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", + fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, +) -> dict[float, float]: + """Tail ECE at each non-parametric FDR operating point.""" + work = df[[confidence_col, label_col]].dropna() + if work.empty: + return {threshold: float("nan") for threshold in fdr_thresholds} + + pred = work[confidence_col].to_numpy(dtype=float) + labels = work[label_col].to_numpy(dtype=float) + + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=work[confidence_col]) + + return { + threshold: compute_tail_ece_at_fdr( + pred, + labels, + threshold, + fdr_ctrl=fdr_ctrl, + ) + for threshold in fdr_thresholds + } + + +def compute_fdr_bias_at_fdr_thresholds( + df: pd.DataFrame, + *, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", + fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, +) -> dict[float, float]: + """Signed FDR bias at each NP-FDR cutoff; equal to empirical sTECE.""" + work = df[[confidence_col, label_col]].dropna() + if work.empty: + return {threshold: float("nan") for threshold in fdr_thresholds} + + scores = work[confidence_col].to_numpy(dtype=float) + labels = work[label_col].to_numpy(dtype=float) + + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=work[confidence_col]) + + results: dict[float, float] = {} + for threshold in fdr_thresholds: + cutoff = fdr_ctrl.get_confidence_cutoff(threshold=threshold) + if np.isnan(cutoff): + results[threshold] = float("nan") + continue + mask = scores >= cutoff + if not mask.any(): + results[threshold] = float("nan") + continue + + # E[1-S | S>=tau] - E[1-Y | S>=tau] = E[Y-S | S>=tau]. + results[threshold] = float(np.mean(labels[mask] - scores[mask])) + return results + + +def compute_pr_auc( + df: pd.DataFrame, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", +) -> float: + """Area under the ablation PR curve (matches ``run_feature_ablations`` plots).""" + work = df[[confidence_col, label_col]].dropna() + if work.empty: + return float("nan") + + sorted_data = work.sort_values(by=confidence_col, ascending=False) + cum_correct = np.cumsum(sorted_data[label_col].values) + precision = cum_correct / np.arange(1, len(sorted_data) + 1) + total_correct = cum_correct[-1] if len(cum_correct) else 0 + if total_correct <= 0 or len(precision) < 2: + return 0.0 + + recall = cum_correct / total_correct + from sklearn.metrics import auc + + return float(auc(recall, precision)) + + +def _vectorized_psm_fdr( + scores: np.ndarray, + ctrl: NonParametricFDRControl, +) -> np.ndarray: + """Map confidence scores to PSM FDR using a fitted controller.""" + conf = np.asarray(ctrl._confidence_scores, dtype=float) + fdr = np.asarray(ctrl._fdr_values, dtype=float) + scores = np.asarray(scores, dtype=float) + idx = np.searchsorted(-conf, -scores, side="left") + idx = np.clip(idx, 0, max(len(fdr) - 1, 0)) + if len(fdr) == 0: + return np.ones_like(scores) + + out = fdr[idx] + below = (idx == len(conf)) & (scores < conf[-1]) + above = (idx == 0) & (scores > conf[0]) + out[below] = 1.0 + out[above] = fdr[0] + return out + + +def _vectorized_psm_q_values( + scores: np.ndarray, + ctrl: NonParametricFDRControl, +) -> np.ndarray: + """Assign PSM q-values without per-row ``compute_fdr`` calls.""" + row_fdr = _vectorized_psm_fdr(scores, ctrl) + order = np.argsort(-scores) + sorted_fdr = row_fdr[order] + q_sorted = np.empty_like(sorted_fdr) + fdr_min = np.inf + for i in range(len(sorted_fdr) - 1, -1, -1): + current = sorted_fdr[i] + if current > fdr_min: + q_sorted[i] = fdr_min + else: + q_sorted[i] = current + fdr_min = current + q_values = np.empty_like(q_sorted) + q_values[order] = q_sorted + return q_values + + +def compute_q_value_deviations( + df: pd.DataFrame, + *, + confidence_col: str = "calibrated_confidence", + label_col: str = "correct", + fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, +) -> dict[float, float]: + """Mean absolute q-value deviation among NP-accepted PSMs at each FDR level.""" + work = df[[confidence_col, label_col]].dropna().copy() + if work.empty or label_col not in work.columns: + return {threshold: float("nan") for threshold in fdr_thresholds} + + np_fdr = NonParametricFDRControl() + np_fdr.fit(dataset=work[confidence_col]) + + dbg_ctrl = _fit_database_grounded_fdr( + work, + confidence_col=confidence_col, + correct_col=label_col, + drop=0 if len(work) <= 10 else 10, + ) + + scores = work[confidence_col].to_numpy(dtype=float) + est_q = _vectorized_psm_q_values(scores, np_fdr) + true_q = _vectorized_psm_q_values(scores, dbg_ctrl) + deviations = np.abs(est_q - true_q) + + results: dict[float, float] = {} + for threshold in fdr_thresholds: + mask = est_q <= threshold + if not mask.any(): + results[threshold] = float("nan") + else: + results[threshold] = float(np.mean(deviations[mask])) + return results + + +def metrics_from_eval_parquet(path: Path) -> dict[str, float | str]: + """Compute top-decile/FDR-tail ECE and metrics from one eval-results Parquet.""" + df = pd.read_parquet(path) + config_name = str(df["config_name"].iloc[0]) + dataset_name = str(df["dataset_name"].iloc[0]) + meta = df.drop(columns=["config_name", "dataset_name"], errors="ignore") + + pred = meta["calibrated_confidence"].to_numpy(dtype=float) + labels = meta["correct"].to_numpy(dtype=float) + tail_ece = compute_tail_ece(pred, labels) + tail_ece_at_fdr = compute_tail_ece_at_fdr_thresholds(meta) + fdr_bias = compute_fdr_bias_at_fdr_thresholds(meta) + pr_auc = compute_pr_auc(meta) + q_dev = compute_q_value_deviations(meta) + + return { + "config": config_name, + "dataset": dataset_name, + "tail_ECE": round(tail_ece, 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.05]: round(tail_ece_at_fdr[0.05], 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.10]: round(tail_ece_at_fdr[0.10], 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.05]: round(fdr_bias[0.05], 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.10]: round(fdr_bias[0.10], 5), + "PR_AUC": round(pr_auc, 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.05]: round(q_dev[0.05], 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.10]: round(q_dev[0.10], 5), + } + + +def enrich_summary_from_eval_results( + summary: pd.DataFrame, + eval_results_dir: Path, + *, + datasets: list[str] | None = None, +) -> pd.DataFrame: + """Add PR-AUC and q-value deviation columns using saved eval Parquets.""" + if not eval_results_dir.is_dir(): + raise FileNotFoundError(eval_results_dir) + + metric_rows: list[dict[str, float | str]] = [] + for path in sorted(eval_results_dir.glob("*.parquet")): + dataset_name = path.name.split("_", 1)[0] + if datasets is not None and dataset_name not in datasets: + continue + metric_rows.append(metrics_from_eval_parquet(path)) + + if not metric_rows: + raise FileNotFoundError( + f"No eval Parquets found under {eval_results_dir}" + + (f" for datasets {datasets!r}" if datasets else "") + ) + + metrics_df = pd.DataFrame(metric_rows) + merge_cols = ["config", "dataset"] + extra_cols = [ + *TAIL_ECE_COLUMN_BY_THRESHOLD.values(), + *FDR_BIAS_COLUMN_BY_THRESHOLD.values(), + "PR_AUC", + *Q_DEV_COLUMN_BY_THRESHOLD.values(), + "tail_ECE", + ] + summary = summary.drop(columns=extra_cols, errors="ignore") + return summary.merge(metrics_df, on=merge_cols, how="left") + + +def _ablation_dataset_display(dataset: str) -> str: + """Publication label via ``plot_eval_results._display_name``.""" + return _display_name(_ABLATION_DATASET_KEYS.get(dataset, dataset)) + + +def _wrap_title_before_dataset(title: str, *, max_line: int = 52) -> str: + """Break before ``on `` when the title would be too wide.""" + marker = " on " + if marker not in title or len(title) <= max_line: + return title + split = title.index(marker) + return f"{title[:split]}\n{title[split + 1 :]}" + + +def _metric_axis_label(metric: MetricName) -> str: + if metric == "tail_ECE": + return "Tail ECE" + if metric == "tail_ECE@5%FDR": + return "Tail ECE at 5% FDR" + if metric == "tail_ECE@10%FDR": + return "Tail ECE at 10% FDR" + if metric == "ECE": + return "ECE" + if metric == "Brier": + return "Brier score" + if metric == "PR_AUC": + return "PR-AUC" + if metric == "fdr_bias@5%FDR": + return "FDR bias (= sTECE) at 5% FDR" + if metric == "fdr_bias@10%FDR": + return "FDR bias (= sTECE) at 10% FDR" + if metric == "q_dev@5%FDR": + return "Mean |q-value deviation| at 5% FDR" + return "Mean |q-value deviation| at 10% FDR" + + +def _metric_plot_title(metric: MetricName, dataset_display: str) -> str: + """Publication title: full sentence, ECE capitalised.""" + if metric == "tail_ECE": + pct = int(TAIL_ECE_FRACTION * 100) + title = ( + f"Tail ECE in the top {pct}% of PSMs ranked by calibrated score " + f"on {dataset_display}" + ) + elif metric == "tail_ECE@5%FDR": + title = ( + f"Tail expected calibration error among PSMs accepted at 5% FDR " + f"on {dataset_display}" + ) + elif metric == "tail_ECE@10%FDR": + title = ( + f"Tail expected calibration error among PSMs accepted at 10% FDR " + f"on {dataset_display}" + ) + elif metric == "ECE": + title = f"Expected calibration error (ECE) on {dataset_display}" + elif metric == "Brier": + title = f"Brier score on {dataset_display}." + elif metric == "PR_AUC": + title = f"Precision-recall AUC on {dataset_display}" + elif metric == "fdr_bias@5%FDR": + title = ( + f"FDR bias, equal to signed tail calibration error, " + f"at 5% FDR on {dataset_display}" + ) + elif metric == "fdr_bias@10%FDR": + title = ( + f"FDR bias, equal to signed tail calibration error, " + f"at 10% FDR on {dataset_display}" + ) + elif metric == "q_dev@5%FDR": + title = ( + f"Non-parametric q-value deviation from database-grounded q-values " + f"at 5% FDR on {dataset_display}" + ) + else: + title = ( + f"Non-parametric q-value deviation from database-grounded q-values " + f"at 10% FDR on {dataset_display}" + ) + return _wrap_title_before_dataset(title) + + +def plot_ablation_calibration_bars( + summary: pd.DataFrame, + dataset: str, + *, + metric: MetricName = "tail_ECE", + output_path: Path, + figsize: tuple[float, float] = (7.5, 4), +) -> pd.DataFrame: + """Bar chart of *metric* for one dataset; returns the plotted slice.""" + ds = summary.loc[summary["dataset"] == dataset].copy() + if ds.empty: + available = sorted(summary["dataset"].unique()) + raise ValueError(f"No rows for dataset {dataset!r}. Available: {available}") + + configs = ordered_ablation_configs(set(ds["config"])) + ds = ds.set_index("config").loc[configs].reset_index() + if metric not in ds.columns: + raise ValueError(f"Metric {metric!r} not in summary columns: {ds.columns}") + + values = ds[metric].to_numpy(dtype=float) + all_features_value = float(ds.loc[ds["config"] == "All features", metric].iloc[0]) + + colors = assign_ablation_colors(configs) + short_labels = [_CONFIG_SHORT_LABELS.get(c, c) for c in configs] + + fig, ax = plt.subplots(figsize=figsize) + x = np.arange(len(configs)) + bar_colors = [colors[c] for c in configs] + ax.bar(x, values, color=bar_colors, edgecolor="black", linewidth=0.6, zorder=2) + ax.axhline( + all_features_value, + color="#333333", + linestyle="--", + linewidth=1.2, + zorder=1, + label="All features", + ) + + display = _ablation_dataset_display(dataset) + ax.set_ylabel(_metric_axis_label(metric)) + ax.set_xlabel("Calibrator feature groups") + ax.set_title(_metric_plot_title(metric, display)) + ax.set_xticks(x) + ax.set_xticklabels(short_labels, rotation=35, ha="right") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s.png and %s.pdf", output_path, output_path) + return ds[["config", metric]] + + +@app.command() +def main( + summary: Annotated[ + Path, + typer.Option("--summary", help="ablation_summary.csv or .json"), + ] = _DEFAULT_SUMMARY, + dataset: Annotated[ + str, + typer.Option("--dataset", help="Dataset key in the summary table"), + ] = "Arabidopsis", + metric: Annotated[ + MetricName, + typer.Option("--metric", help="Calibration metric to plot"), + ] = "tail_ECE", + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for figure outputs"), + ] = _DEFAULT_OUTPUT_DIR, + eval_results_dir: Annotated[ + Path | None, + typer.Option( + "--eval-results-dir", + help="Optional eval_results/ directory to enrich summary before plotting", + ), + ] = None, +) -> None: + """Plot ablation calibration bars for one dataset.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + sns.set_theme(style="white", context="paper", font_scale=1.5) + + summary_df = load_ablation_summary(summary) + if eval_results_dir is not None: + summary_df = enrich_summary_from_eval_results( + summary_df, + eval_results_dir, + datasets=[dataset], + ) + plots_dir.mkdir(parents=True, exist_ok=True) + slug = dataset.lower().replace(" ", "_") + metric_slug = metric.lower().replace("%", "pct").replace("@", "_at_") + out_base = plots_dir / f"ablation_{metric_slug}_{slug}" + table = plot_ablation_calibration_bars( + summary_df, dataset, metric=metric, output_path=out_base + ) + print(table.to_string(index=False)) + + +@app.command("recompute-summary") +def recompute_summary( + eval_results_dir: Annotated[ + Path, + typer.Option("--eval-results-dir", help="Directory of eval_results Parquets"), + ], + summary: Annotated[ + Path | None, + typer.Option( + "--summary", + help="Existing ablation_summary.csv to merge with (optional)", + ), + ] = None, + datasets: Annotated[ + list[str] | None, + typer.Option( + "--datasets", + help="Restrict to these dataset keys (repeatable)", + ), + ] = None, + output: Annotated[ + Path, + typer.Option("--output", help="Output CSV path"), + ] = _DEFAULT_SUMMARY, +) -> None: + """Recompute PR-AUC and q-value deviation columns from eval Parquets.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + if summary is not None: + base = load_ablation_summary(summary) + if datasets is not None: + base = base.loc[base["dataset"].isin(datasets)].copy() + enriched = enrich_summary_from_eval_results( + base, + eval_results_dir, + datasets=datasets, + ) + else: + rows = [] + for path in sorted(eval_results_dir.glob("*.parquet")): + dataset_name = path.name.split("_", 1)[0] + if datasets is not None and dataset_name not in datasets: + continue + rows.append(metrics_from_eval_parquet(path)) + if not rows: + raise typer.BadParameter(f"No eval Parquets found under {eval_results_dir}") + enriched = pd.DataFrame(rows) + enriched = enriched.sort_values(["dataset", "config"]).reset_index(drop=True) + + output.parent.mkdir(parents=True, exist_ok=True) + enriched.to_csv(output, index=False) + logger.info("Wrote %s", output) + print(enriched.to_string(index=False)) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/plot_analysis.py b/paper_scripts/plot_analysis.py new file mode 100644 index 00000000..1b290531 --- /dev/null +++ b/paper_scripts/plot_analysis.py @@ -0,0 +1,1201 @@ +"""Generate analysis plots from Winnow predict outputs. + +Produces calibration, PR, confidence, FDR-accuracy and labelled-diagnostic +figures from a predict directory, annotated against a FASTA proteome. +""" + +from __future__ import annotations + +import json +import logging +import sys +import warnings +from pathlib import Path +from typing import Annotated, Literal, Optional + +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +import seaborn as sns +import typer +import yaml +from matplotlib.patches import Patch +from scipy.stats import gaussian_kde +from sklearn.calibration import calibration_curve +from sklearn.decomposition import PCA +from sklearn.preprocessing import StandardScaler + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +from fdr_tool_comparison_preprocess import ( # noqa: E402 + filter_and_annotate_preds, +) +from winnow.utils.proteome import load_proteome_haystack # noqa: E402 +from winnow.calibration.calibrator import TrainingHistory # noqa: E402 +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 + +logger = logging.getLogger(__name__) +app = typer.Typer( + add_completion=False, pretty_exceptions_show_locals=False, no_args_is_help=True +) + +# ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_MAIN_LINE_COLOUR = _PALETTE[0] +_RAW_LINE_COLOUR = _PALETTE[5] +_IDEAL_LINE_COLOUR = _PALETTE[6] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) +warnings.filterwarnings("ignore", module="winnow") + + +def _spine_fmt(ax: plt.Axes) -> None: + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _save(fig: plt.Figure, out_dir: Path, name: str) -> None: + base = out_dir / name + fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + print(f" saved {name}") + + +# ── Plot functions ──────────────────────────────────────────────────── + + +def _df_for_raw_confidence_plots(df: pl.DataFrame) -> pl.DataFrame: + """Drop PSMs with negative raw confidence (Casanovo mass-mismatch penalty scores).""" + if "confidence" not in df.columns: + return df + n_neg = int((df["confidence"] < 0).sum()) + if n_neg == 0: + return df + print( + f" excluding {n_neg} PSMs with negative raw confidence " + "from raw-confidence plots" + ) + return df.filter(pl.col("confidence") >= 0) + + +def plot_calibration_curves( + df: pl.DataFrame, + label_col: str, + title: str, + bins: int = 10, + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Plot reliability curves for calibrated and (optionally) raw confidence.""" + fig, ax = plt.subplots(figsize=(8, 6)) + + frac_pos, mean_pred = calibration_curve( + df[label_col].to_numpy(), + df["calibrated_confidence"].to_numpy(), + n_bins=bins, + strategy="uniform", + ) + ax.plot( + mean_pred, + frac_pos, + marker="o", + color=_MAIN_LINE_COLOUR, + label="Calibrated confidence", + linewidth=1.5, + markersize=6, + zorder=3, + ) + + if df_raw is not None and len(df_raw) > 0 and "confidence" in df_raw.columns: + frac_pos, mean_pred = calibration_curve( + df_raw[label_col].to_numpy(), + df_raw["confidence"].to_numpy(), + n_bins=bins, + strategy="uniform", + ) + ax.plot( + mean_pred, + frac_pos, + marker="D", + color=_RAW_LINE_COLOUR, + label="Raw confidence", + linewidth=1.5, + markersize=6, + zorder=3, + ) + + ax.plot( + [0, 1], + [0, 1], + "--", + color=_IDEAL_LINE_COLOUR, + label="Perfectly calibrated", + alpha=0.7, + zorder=2, + ) + ax.set_xlabel("Mean predicted probability") + ax.set_ylabel("Fraction of positives") + ax.set_title(title) + ax.legend(loc="lower right") + ax.set_xlim([0, 1.05]) + ax.set_ylim([0, 1.05]) + ax.grid(False) + _spine_fmt(ax) + return fig + + +def plot_pr_curves( + df: pl.DataFrame, + label_col: str, + title: str, + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Plot precision–recall curves for calibrated and (optionally) raw confidence.""" + fig, ax = plt.subplots(figsize=(8, 6)) + + sorted_cal = df.sort("calibrated_confidence", descending=True) + labels = sorted_cal[label_col].to_numpy() + cum = np.cumsum(labels) + precision = cum / np.arange(1, len(labels) + 1) + recall = cum / len(labels) + ax.plot( + recall, + precision, + color=_MAIN_LINE_COLOUR, + label="Calibrated confidence", + linewidth=1.5, + ) + + if df_raw is not None and len(df_raw) > 0 and "confidence" in df_raw.columns: + sorted_raw = df_raw.sort("confidence", descending=True) + labels = sorted_raw[label_col].to_numpy() + cum = np.cumsum(labels) + precision = cum / np.arange(1, len(labels) + 1) + recall = cum / len(labels) + ax.plot( + recall, + precision, + color=_RAW_LINE_COLOUR, + label="Raw confidence", + linewidth=1.5, + ) + + ax.set_xlabel("Recall") + ax.set_ylabel("Precision") + ax.set_title(title) + ax.set_xlim(0, 1.05) + ax.set_ylim(0, 1.05) + ax.legend(loc="lower left") + ax.grid(False) + _spine_fmt(ax) + return fig + + +def plot_confidence_histogram( + df: pl.DataFrame, + label_col: str, + conf_col: str, + col_label: str, + title: str, + bins: int = 50, +) -> plt.Figure: + """Plot confidence histograms with KDE overlays for correct vs incorrect PSMs.""" + fig, ax = plt.subplots(1, 1, figsize=(7, 5)) + pos = df.filter(pl.col(label_col)) + neg = df.filter(~pl.col(label_col)) + n_data = neg[conf_col].to_numpy() + p_data = pos[conf_col].to_numpy() + + ax.hist( + n_data, + bins=bins, + alpha=0.6, + label="Incorrect", + density=False, + edgecolor="black", + color=_INCORRECT_COLOUR, + ) + ax.hist( + p_data, + bins=bins, + alpha=0.6, + label="Correct", + density=False, + edgecolor="black", + color=_CORRECT_COLOUR, + ) + + x_min = min(n_data.min(), p_data.min()) + x_max = max(n_data.max(), p_data.max()) + x_grid = np.linspace(x_min, x_max, 300) + bin_width = (x_max - x_min) / bins if bins > 1 else 1.0 + + if len(n_data) > 1: + y_neg = gaussian_kde(n_data)(x_grid) * len(n_data) * bin_width + ax.plot(x_grid, y_neg, color=_INCORRECT_COLOUR, lw=1.5) + if len(p_data) > 1: + y_pos = gaussian_kde(p_data)(x_grid) * len(p_data) * bin_width + ax.plot(x_grid, y_pos, color=_CORRECT_COLOUR, lw=1.5) + + ax.set_xlabel(col_label) + ax.set_ylabel("Frequency") + ax.legend(loc="upper center") + ax.grid(False) + ax.set_title(title) + _spine_fmt(ax) + fig.tight_layout() + return fig + + +def _fit_db_fdr( + df: pl.DataFrame, + correct_col: str, + confidence_feature: str = "calibrated_confidence", + drop: int = 10, +) -> DatabaseGroundedFDRControl: + """Fit a DatabaseGroundedFDRControl from per-row correctness labels.""" + ctrl = DatabaseGroundedFDRControl( + confidence_feature=confidence_feature, + drop=drop, + ) + sorted_df = df.sort(confidence_feature, descending=True) + correct_vals = sorted_df[correct_col].to_numpy().astype(float) + confidence_vals = sorted_df[confidence_feature].to_numpy() + precision = np.cumsum(correct_vals) / np.arange(1, len(sorted_df) + 1) + ctrl._fdr_values = np.array(1 - precision[drop:]) + ctrl._confidence_scores = confidence_vals[drop:] + return ctrl + + +def plot_fdr_accuracy( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, + metric: str = "fdr", + use_proteome_shortcut: bool = False, +) -> plt.Figure: + """Compare non-parametric vs database-grounded FDR or q-value vs confidence.""" + del residue_masses, use_proteome_shortcut # retained for call-site compatibility + fig, ax = plt.subplots(figsize=(8, 6)) + col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" + winnow_col = col_name + + ctrl = _fit_db_fdr(df, correct_col) + + if metric == "fdr": + db_pd = ctrl.add_psm_fdr(df.to_pandas(), "calibrated_confidence") + else: + df_pd = df.to_pandas() + if "psm_q_value" in df_pd.columns: + df_pd = df_pd.drop(columns=["psm_q_value"]) + db_pd = ctrl.add_psm_q_value(df_pd, "calibrated_confidence") + db_df = pl.from_pandas(db_pd).select(["spectrum_id", col_name]) + + merged = ( + df.select(["spectrum_id", "calibrated_confidence", winnow_col]) + .join(db_df, on="spectrum_id", how="inner", suffix="_db") + .sort("calibrated_confidence") + ) + + conf = merged["calibrated_confidence"].to_numpy() + ax.plot( + conf, + merged[winnow_col].to_numpy(), + color=_MAIN_LINE_COLOUR, + label="Non-parametric", + linewidth=1.5, + ) + ax.plot( + conf, + merged[f"{col_name}_db"].to_numpy(), + color=_RAW_LINE_COLOUR, + label="Database-grounded", + linewidth=1.5, + ) + + ax.set_xlabel("Calibrated confidence") + ylabel = "FDR" if metric == "fdr" else "Q-value" + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.legend(loc="upper right") + ax.grid(False) + _spine_fmt(ax) + return fig + + +def plot_ranked_qvalue( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, +) -> plt.Figure: + """Ranked predictions vs q-value (non-parametric & database-grounded).""" + ctrl = _fit_db_fdr(df, correct_col) + test_pd = df.to_pandas() + test_pd_no_q = test_pd.drop(columns=["psm_q_value"], errors="ignore") + db_q = ctrl.add_psm_q_value(test_pd_no_q, "calibrated_confidence") + + sorted_np = test_pd.sort_values( + "calibrated_confidence", ascending=False + ).reset_index(drop=True) + sorted_db = db_q.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + ranks = np.arange(1, len(sorted_np) + 1) + + fig, ax = plt.subplots(figsize=(8, 6)) + ax.plot( + ranks, + sorted_np["psm_q_value"].values, + color=_MAIN_LINE_COLOUR, + label="Non-parametric", + linewidth=1.5, + ) + ax.plot( + ranks, + sorted_db["psm_q_value"].values, + color=_RAW_LINE_COLOUR, + label="Database-grounded", + linewidth=1.5, + ) + ax.set_xlabel("Ranked predictions") + ax.set_ylabel("Q-value") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_ranked_fdr_pep(df: pl.DataFrame, title: str) -> plt.Figure: + """Ranked predictions vs non-parametric FDR and PEP.""" + sorted_df = df.sort("calibrated_confidence", descending=True) + ranks = np.arange(1, len(sorted_df) + 1) + fig, ax = plt.subplots(figsize=(8, 6)) + ax.plot( + ranks, + sorted_df["psm_fdr"].to_numpy(), + color=_MAIN_LINE_COLOUR, + label="FDR", + linewidth=1.5, + ) + if "psm_pep" in sorted_df.columns: + ax.plot( + ranks, + sorted_df["psm_pep"].to_numpy(), + color=_PALETTE[3], + label="PEP", + linewidth=1.5, + ) + ax.set_xlabel("Ranked predictions") + ax.set_ylabel("Error rate") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_ranked_fdr_raw_vs_cal( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, + metric: str = "fdr", + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Ranked predictions vs FDR/q-value for non-parametric and database-grounded on raw+calibrated.""" + test_pd = df.to_pandas() + test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") + + np_cal = test_pd.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + + db_cal_ctrl = _fit_db_fdr( + df, correct_col, confidence_feature="calibrated_confidence" + ) + raw_df = df_raw if df_raw is not None else df + db_raw_ctrl = _fit_db_fdr(raw_df, correct_col, confidence_feature="confidence") + + col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" + add_fn = "add_psm_fdr" if metric == "fdr" else "add_psm_q_value" + sort_col_cal = "calibrated_confidence" + sort_col_raw = "confidence" + + db_cal = getattr(db_cal_ctrl, add_fn)(test_pd_no_q.copy(), sort_col_cal) + db_cal = db_cal.sort_values(sort_col_cal, ascending=False).reset_index(drop=True) + + raw_pd_no_q = raw_df.to_pandas().drop( + columns=["psm_q_value", "psm_fdr"], errors="ignore" + ) + db_raw = getattr(db_raw_ctrl, add_fn)(raw_pd_no_q.copy(), sort_col_raw) + db_raw = db_raw.sort_values(sort_col_raw, ascending=False).reset_index(drop=True) + + ranks_cal = np.arange(1, len(np_cal) + 1) + ranks_raw = np.arange(1, len(db_raw) + 1) + + fig, ax = plt.subplots(figsize=(8, 6)) + ylabel = "FDR" if metric == "fdr" else "Q-value" + np_col = "psm_fdr" if metric == "fdr" else "psm_q_value" + ax.plot( + ranks_cal, + np_cal[np_col].values, + color=_MAIN_LINE_COLOUR, + label="Non-parametric (calibrated)", + linewidth=1.5, + ) + ax.plot( + ranks_cal, + db_cal[col_name].values, + color=_RAW_LINE_COLOUR, + label="Database-grounded (calibrated)", + linewidth=1.5, + ) + ax.plot( + ranks_raw, + db_raw[col_name].values, + color=_PALETTE[3], + label="Database-grounded (raw)", + linewidth=1.5, + ) + ax.set_xlabel("Ranked predictions") + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +def plot_bar_psms_fdr( + df: pl.DataFrame, + correct_col: str, + residue_masses: dict, + title: str, + df_raw: pl.DataFrame | None = None, +) -> plt.Figure: + """Bar plot of PSMs at q-value thresholds (calibrated vs raw, database-grounded).""" + test_pd = df.to_pandas() + test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") + + db_cal_ctrl = _fit_db_fdr( + df, correct_col, confidence_feature="calibrated_confidence" + ) + raw_df = df_raw if df_raw is not None else df + db_raw_ctrl = _fit_db_fdr(raw_df, correct_col, confidence_feature="confidence") + + db_cal = db_cal_ctrl.add_psm_q_value(test_pd_no_q.copy(), "calibrated_confidence") + raw_pd_no_q = raw_df.to_pandas().drop( + columns=["psm_q_value", "psm_fdr"], errors="ignore" + ) + db_raw = db_raw_ctrl.add_psm_q_value(raw_pd_no_q.copy(), "confidence") + + thresholds = [0.001, 0.01, 0.05, 0.1] + counts_cal = [int((db_cal["psm_q_value"] <= t).sum()) for t in thresholds] + counts_raw = [int((db_raw["psm_q_value"] <= t).sum()) for t in thresholds] + + x = np.arange(len(thresholds)) + width, gap = 0.32, 0.04 + + fig, ax = plt.subplots(figsize=(8, 6)) + bars_cal = ax.bar( + x - width / 2 - gap / 2, + counts_cal, + width, + label="Calibrated confidence", + color=_MAIN_LINE_COLOUR, + edgecolor="black", + linewidth=1, + ) + bars_raw = ax.bar( + x + width / 2 + gap / 2, + counts_raw, + width, + label="Raw confidence", + color=_RAW_LINE_COLOUR, + edgecolor="black", + linewidth=1, + ) + ax.set_xlabel("FDR threshold") + ax.set_ylabel("Peptide-spectrum matches") + ax.set_title(title) + ax.set_xticks(x) + ax.set_xticklabels([str(t) for t in thresholds]) + ax.legend(loc="upper left") + + for bar_group in [bars_cal, bars_raw]: + for bar in bar_group: + h = bar.get_height() + ax.annotate( + f"{h:,}", + xy=(bar.get_x() + bar.get_width() / 2, h), + xytext=(0, 3), + textcoords="offset points", + ha="center", + va="bottom", + fontsize=10, + ) + _spine_fmt(ax) + return fig + + +def plot_raw_vs_cal_scatter( + df: pl.DataFrame, + label_col: str, + title: str, +) -> plt.Figure: + """Raw confidence vs calibrated confidence coloured by correctness.""" + fig, ax = plt.subplots(figsize=(8, 7)) + inc = df.filter(~pl.col(label_col)) + cor = df.filter(pl.col(label_col)) + ax.scatter( + inc["confidence"].to_numpy(), + inc["calibrated_confidence"].to_numpy(), + c=_INCORRECT_COLOUR, + label="Incorrect", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.scatter( + cor["confidence"].to_numpy(), + cor["calibrated_confidence"].to_numpy(), + c=_CORRECT_COLOUR, + label="Correct", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.plot( + [0, 1], + [0, 1], + color=_IDEAL_LINE_COLOUR, + linestyle="--", + linewidth=1, + label="Identity", + ) + ax.set_xlabel("Raw confidence") + ax.set_ylabel("Calibrated confidence") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +# Features used by the paper HeLa InstantNovo / Casanovo calibrators (and the +# general model). Prefer ``config.json`` ``feature_columns`` via ``--model-dir``. +_DEFAULT_TRAINED_FEATURE_COLUMNS = [ + "confidence", + "ion_matches", + "ion_match_intensity", + "irt_error", + "margin", + "median_margin", + "entropy", + "z-score", + "min_token_probability", + "std_token_probability", + "mass_error_da", +] + + +def _resolve_pca_feature_columns(model_dir: Path | None) -> list[str]: + """Return MLP input columns: confidence + checkpoint ``feature_columns``.""" + if model_dir is None: + return list(_DEFAULT_TRAINED_FEATURE_COLUMNS) + config_path = model_dir / "config.json" + if not config_path.is_file(): + logger.warning( + "No config.json under %s; using default trained feature columns for PCA", + model_dir, + ) + return list(_DEFAULT_TRAINED_FEATURE_COLUMNS) + cfg = json.loads(config_path.read_text()) + feature_columns = cfg.get("feature_columns") + if not isinstance(feature_columns, list) or not feature_columns: + logger.warning( + "config.json under %s has no feature_columns; using default PCA set", + model_dir, + ) + return list(_DEFAULT_TRAINED_FEATURE_COLUMNS) + return ["confidence", *[str(c) for c in feature_columns]] + + +def plot_pca_features( + df: pl.DataFrame, + label_col: str, + title: str, + feature_cols: list[str] | None = None, +) -> tuple[plt.Figure, PCA, list[str]]: + """PCA of trained calibrator features coloured by correctness.""" + if feature_cols is None: + feature_cols = list(_DEFAULT_TRAINED_FEATURE_COLUMNS) + available = [c for c in feature_cols if c in df.columns] + missing = [c for c in feature_cols if c not in df.columns] + if missing: + logger.warning("PCA skipped missing trained feature columns: %s", missing) + if len(available) < 2: + raise ValueError( + f"Need at least 2 trained feature columns for PCA; found {available}" + ) + feat_df = df.select(available).to_pandas().dropna() + labels = df.filter(pl.all_horizontal([pl.col(c).is_not_null() for c in available]))[ + label_col + ].to_numpy() + + scaler = StandardScaler() + features_scaled = scaler.fit_transform(feat_df.values) + pca = PCA(n_components=2) + coords = pca.fit_transform(features_scaled) + + fig, ax = plt.subplots(figsize=(8, 7)) + mask_inc, mask_cor = ~labels, labels + ax.scatter( + coords[mask_inc, 0], + coords[mask_inc, 1], + c=_INCORRECT_COLOUR, + label="Incorrect", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.scatter( + coords[mask_cor, 0], + coords[mask_cor, 1], + c=_CORRECT_COLOUR, + label="Correct", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.set_xlabel(f"PC 1 ({pca.explained_variance_ratio_[0]:.1%} variance)") + ax.set_ylabel(f"PC 2 ({pca.explained_variance_ratio_[1]:.1%} variance)") + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig, pca, available + + +def plot_pca_loadings( + pca: PCA, + feature_names: list[str], + title: str, +) -> plt.Figure: + """PCA loadings for PC1 and PC2, ordered by |PC1|.""" + pretty = { + "confidence": "Raw confidence", + "mass_error_ppm": "Log absolute mass error (ppm)", + "mass_error_da": "Mass error (Da)", + "ion_matches": "Ion matches", + "ion_match_intensity": "Ion match intensity", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Maximum ion gap", + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation", + "irt_error": "Retention time error", + "margin": "Margin", + "median_margin": "Median margin", + "entropy": "Entropy", + "z-score": "Z-score", + "edit_distance": "Edit distance", + "min_token_probability": "Minimum token probability", + "std_token_probability": "Token probability std. dev.", + } + pc1 = pca.components_[0] + pc2 = pca.components_[1] + names = [pretty.get(c, c) for c in feature_names] + order = np.argsort(np.abs(pc1))[::-1] + + y = np.arange(len(names)) + fig, ax = plt.subplots(figsize=(10, 7)) + ax.barh(y, pc1[order], color=_MAIN_LINE_COLOUR, alpha=0.6, edgecolor="black") + ax.barh(y, pc2[order], color=_RAW_LINE_COLOUR, alpha=0.4, edgecolor="black") + ax.set_yticks(y) + ax.set_yticklabels([names[i] for i in order]) + ax.invert_yaxis() + ax.set_xlabel("Loading value") + ax.set_title(title) + ax.axvline(0, color="black", linewidth=0.5) + ax.legend( + handles=[ + Patch(facecolor=_MAIN_LINE_COLOUR, alpha=0.6, label="PC 1 loading"), + Patch(facecolor=_RAW_LINE_COLOUR, alpha=0.4, label="PC 2 loading"), + ], + loc="lower right", + ) + _spine_fmt(ax) + return fig + + +def plot_scatter_feature_vs_conf( + df: pl.DataFrame, + label_col: str, + x_col: str, + y_col: str, + title: str, + x_label: str | None = None, + y_label: str | None = None, +) -> plt.Figure: + """Scatter of x_col vs y_col coloured by correctness.""" + fig, ax = plt.subplots(figsize=(8, 7)) + inc = df.filter(~pl.col(label_col)) + cor = df.filter(pl.col(label_col)) + ax.scatter( + inc[x_col].to_numpy(), + inc[y_col].to_numpy(), + c=_INCORRECT_COLOUR, + label="Incorrect", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.scatter( + cor[x_col].to_numpy(), + cor[y_col].to_numpy(), + c=_CORRECT_COLOUR, + label="Correct", + s=10, + alpha=0.3, + rasterized=True, + ) + ax.set_xlabel(x_label or x_col.replace("_", " ").title()) + ax.set_ylabel(y_label or y_col.replace("_", " ").title()) + ax.set_title(title) + ax.legend(loc="upper left") + _spine_fmt(ax) + return fig + + +# ── Main logic ──────────────────────────────────────────────────────── + + +def _load_residue_masses() -> dict: + cfg_path = REPO_ROOT / "winnow" / "configs" / "residues.yaml" + with open(cfg_path) as f: + return yaml.safe_load(f)["residue_masses"] + + +def _load_data(predictions_dir: Path) -> pl.DataFrame: + preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") + meta_path = predictions_dir / "metadata.csv" + if meta_path.exists(): + meta = pl.read_csv(meta_path) + preds = preds.join(meta, on="spectrum_id", how="inner") + return preds + + +_SPLIT_DISPLAY_NAMES = { + "test": "test set", + "unlabelled": "unlabelled space", + "raw_less_train": "full search space", +} + + +def _split_display(split: str) -> str: + key = split.strip().replace("-", "_") + return _SPLIT_DISPLAY_NAMES.get(key, split) + + +def _dns_model_tag(dns_model: str | None) -> str: + return f" ({dns_model})" if dns_model else "" + + +def _split_title(split_label: str, dns_model: str | None = None) -> str: + return f"{split_label}{_dns_model_tag(dns_model)}" + + +def _eval_title( + prefix: str, + split_label: str, + eval_kind: str, + dns_model: str | None = None, +) -> str: + return f"{prefix} {_split_title(split_label, dns_model)}\nusing {eval_kind}" + + +def _save_training_history_plot(model_dir: Path, out_dir: Path) -> None: + hist_path = model_dir / "training_history.json" + if not hist_path.exists(): + return + print("Plotting training history") + th = TrainingHistory.load(str(hist_path)) + th.plot(output_path=out_dir / "training_history.png", show=False) + print(" saved training_history") + + +def _plot_calibration_and_pr( + df: pl.DataFrame, + df_raw_conf: pl.DataFrame, + split: str, + labelled: bool, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + print("Plotting calibration curves") + if labelled: + fig = plot_calibration_curves( + df, + "correct", + _eval_title( + "Calibration curves for", split_label, "database search", dns_model + ), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"calibration_{split}_db_search") + + fig = plot_calibration_curves( + df, + "proteome_hit", + _eval_title( + "Calibration curves for", split_label, "proteome mapping", dns_model + ), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"calibration_{split}_proteome") + + print("Plotting PR curves") + if labelled: + fig = plot_pr_curves( + df, + "correct", + _eval_title("PR curves for", split_label, "database search", dns_model), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"pr_{split}_db_search") + + fig = plot_pr_curves( + df, + "proteome_hit", + _eval_title("PR curves for", split_label, "proteome mapping", dns_model), + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"pr_{split}_proteome") + + +def _plot_confidence_histograms( + df: pl.DataFrame, + df_raw_conf: pl.DataFrame, + split: str, + labelled: bool, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + print("Plotting confidence histograms") + for conf_col, conf_label, tag in [ + ("confidence", "Raw confidence", "raw"), + ("calibrated_confidence", "Calibrated confidence", "cal"), + ]: + hist_df = df_raw_conf if conf_col == "confidence" else df + if labelled: + fig = plot_confidence_histogram( + hist_df, + "correct", + conf_col, + conf_label, + _eval_title( + f"{conf_label} for", split_label, "database search", dns_model + ), + ) + _save(fig, out_dir, f"hist_{tag}_{split}_db_search") + + fig = plot_confidence_histogram( + hist_df, + "proteome_hit", + conf_col, + conf_label, + _eval_title( + f"{conf_label} for", split_label, "proteome mapping", dns_model + ), + ) + _save(fig, out_dir, f"hist_{tag}_{split}_proteome") + + +def _plot_fdr_accuracy_plots( + df: pl.DataFrame, + split: str, + labelled: bool, + residue_masses: dict, + out_dir: Path, + dns_model: str | None = None, +) -> None: + split_label = _split_display(split) + print("Plotting FDR accuracy") + use_shortcut = not labelled + for metric, tag in [("fdr", "fdr"), ("q_value", "qvalue")]: + metric_name = "FDR" if metric == "fdr" else "Q-value" + if labelled: + fig = plot_fdr_accuracy( + df, + "correct", + residue_masses, + _eval_title( + f"{metric_name} accuracy for", + split_label, + "database search", + dns_model, + ), + metric="fdr" if metric == "fdr" else "q_value", + use_proteome_shortcut=False, + ) + _save(fig, out_dir, f"{tag}_{split}_db_search") + + fig = plot_fdr_accuracy( + df, + "proteome_hit", + residue_masses, + _eval_title( + f"{metric_name} accuracy for", + split_label, + "proteome mapping", + dns_model, + ), + metric="fdr" if metric == "fdr" else "q_value", + use_proteome_shortcut=use_shortcut, + ) + _save(fig, out_dir, f"{tag}_{split}_proteome") + + +def _plot_labelled_diagnostics( + df: pl.DataFrame, + df_raw_conf: pl.DataFrame, + split: str, + residue_masses: dict, + out_dir: Path, + dns_model: str | None = None, + model_dir: Path | None = None, +) -> None: + split_label = _split_display(split) + title_split = _split_title(split_label, dns_model) + label_col = "correct" + print("Plotting labelled-only diagnostics") + + fig = plot_ranked_qvalue( + df, + label_col, + residue_masses, + _eval_title( + "Ranked predictions vs q-value for", + split_label, + "database search", + dns_model, + ), + ) + _save(fig, out_dir, f"ranked_qvalue_{split}_db_search") + + fig = plot_ranked_fdr_pep( + df, f"Ranked predictions vs FDR and PEP for {title_split}" + ) + _save(fig, out_dir, f"ranked_fdr_pep_{split}_nonparametric") + + for metric, file_tag, metric_name in [ + ("fdr", "fdr", "FDR"), + ("q_value", "qvalue", "q-value"), + ]: + fig = plot_ranked_fdr_raw_vs_cal( + df, + label_col, + residue_masses, + _eval_title( + f"Ranked predictions vs {metric_name} for", + split_label, + "database search", + dns_model, + ), + metric=metric, + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"ranked_{file_tag}_raw_vs_cal_{split}_db_search") + + fig = plot_bar_psms_fdr( + df, + label_col, + residue_masses, + f"PSMs at database-grounded FDR thresholds for {title_split}", + df_raw=df_raw_conf, + ) + _save(fig, out_dir, f"bar_psms_fdr_thresholds_{split}_db_search") + + if len(df_raw_conf) > 0: + fig = plot_raw_vs_cal_scatter( + df_raw_conf, + label_col, + f"Raw vs calibrated confidence for {title_split}", + ) + _save(fig, out_dir, f"scatter_raw_vs_cal_confidence_{split}") + + if "margin" in df.columns and len(df_raw_conf) > 0: + fig = plot_scatter_feature_vs_conf( + df_raw_conf, + label_col, + "margin", + "confidence", + f"Raw confidence vs margin for {title_split}", + x_label="Margin", + y_label="Raw confidence", + ) + _save(fig, out_dir, f"scatter_raw_confidence_vs_margin_{split}") + + fig = plot_scatter_feature_vs_conf( + df, + label_col, + "margin", + "calibrated_confidence", + f"Calibrated confidence vs margin for {title_split}", + x_label="Margin", + y_label="Calibrated confidence", + ) + _save(fig, out_dir, f"scatter_cal_confidence_vs_margin_{split}") + + pca_feature_cols = _resolve_pca_feature_columns(model_dir) + logger.info("PCA using trained feature columns: %s", pca_feature_cols) + fig, pca_model, feat_names = plot_pca_features( + df, + label_col, + f"PCA of calibrator features for {title_split}", + feature_cols=pca_feature_cols, + ) + _save(fig, out_dir, f"pca_features_{split}") + + fig = plot_pca_loadings( + pca_model, feat_names, "PCA loadings for first two principal components" + ) + _save(fig, out_dir, f"pca_loadings_pc1_pc2_{split}") + + +@app.command() +def main( + predictions_dir: Annotated[ + Path, + typer.Option( + "--predictions-dir", + help="Winnow predict output dir with preds_and_fdr_metrics.csv", + ), + ], + split: Annotated[ + str, + typer.Option( + "--split", + help=( + "Split id for filenames; titles use test set / unlabelled space / " + "full search space" + ), + ), + ], + label_mode: Annotated[ + Literal["labelled", "unlabelled"], + typer.Option( + "--label-mode", + help="labelled = has 'correct' column; unlabelled = proteome mapping only", + ), + ], + fasta: Annotated[ + Path, + typer.Option("--fasta", help="FASTA file for proteome annotation"), + ], + plots_dir: Annotated[ + Optional[Path], + typer.Option( + "--plots-dir", + help="Directory for plots (defaults to predictions-dir/plots/).", + ), + ] = None, + model_dir: Annotated[ + Optional[Path], + typer.Option( + "--model-dir", + help=( + "Calibrator directory (config.json + optional training_history). " + "PCA uses checkpoint feature_columns (+ confidence)." + ), + ), + ] = None, + dns_model: Annotated[ + Optional[str], + typer.Option( + "--dns-model", + help=( + "Upstream DNS model name for plot titles " + "(e.g. InstaNovo, Casanovo, $\\pi$-PrimeNovo)" + ), + ), + ] = None, +) -> None: + """Load predictions, annotate proteome hits, and write analysis plots.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + out_dir = plots_dir if plots_dir is not None else predictions_dir / "plots" + out_dir.mkdir(parents=True, exist_ok=True) + + labelled = label_mode == "labelled" + residue_masses = _load_residue_masses() + + logger.info("Loading predictions from %s", predictions_dir) + df = _load_data(predictions_dir) + + from instanovo.utils.metrics import Metrics + from instanovo.utils.residues import ResidueSet + + metrics = Metrics( + residue_set=ResidueSet(residue_masses=residue_masses), + isotope_error_range=(0, 1), + ) + + logger.info("Annotating with proteome hits from %s", fasta) + haystack = load_proteome_haystack(str(fasta)) + df = filter_and_annotate_preds(df, haystack, metrics, min_residue_length=7) + df_raw_conf = _df_for_raw_confidence_plots(df) + + if model_dir is not None: + _save_training_history_plot(model_dir, out_dir) + + _plot_calibration_and_pr( + df, df_raw_conf, split, labelled, out_dir, dns_model=dns_model + ) + _plot_confidence_histograms( + df, df_raw_conf, split, labelled, out_dir, dns_model=dns_model + ) + _plot_fdr_accuracy_plots( + df, split, labelled, residue_masses, out_dir, dns_model=dns_model + ) + + if labelled: + _plot_labelled_diagnostics( + df, + df_raw_conf, + split, + residue_masses, + out_dir, + dns_model=dns_model, + model_dir=model_dir, + ) + else: + # Unlabelled / full-search: ranked q-value against proteome mapping. + split_label = _split_display(split) + print("Plotting ranked q-value (proteome mapping)") + fig = plot_ranked_qvalue( + df, + "proteome_hit", + residue_masses, + _eval_title( + "Ranked predictions vs q-value for", + split_label, + "proteome mapping", + dns_model, + ), + ) + _save(fig, out_dir, f"ranked_qvalue_{split}_proteome") + + logger.info("All plots saved to %s", out_dir) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/plot_calibrator_generalisation_heatmap.py b/paper_scripts/plot_calibrator_generalisation_heatmap.py new file mode 100644 index 00000000..b8f292f0 --- /dev/null +++ b/paper_scripts/plot_calibrator_generalisation_heatmap.py @@ -0,0 +1,270 @@ +"""Plot PR-AUC heatmaps for calibrator generalisation results. + +Reads the combined CSV produced by ``evaluate_calibrator_generalisation.py`` +and creates heatmaps comparing raw vs calibrated confidence PR-AUC values. +""" + +import logging +import sys +from pathlib import Path +from typing import Annotated + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +from matplotlib.colors import LinearSegmentedColormap +from rich.logging import RichHandler +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +from calibrator_generalisation_utils import SPECIES_NAME_MAPPING # noqa: E402 + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- +logger = logging.getLogger("winnow.plot_generalisation_heatmap") +logger.setLevel(logging.INFO) +logger.propagate = False +logger.addHandler(RichHandler()) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette + "sunset" diverging colourmap +# --------------------------------------------------------------------------- +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +_SUNSET_COLORS = [ + "#364B9A", + "#4A7BB7", + "#6EA6CD", + "#98CAE1", + "#C2E4EF", + "#EAECCC", + "#FEDA8B", + "#FDB366", + "#F67E4B", + "#DD3D2D", + "#A50026", +] +_BAD_COLOUR = "#FFFFFF" + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=2) + + +def _diverging_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) + cmap.set_bad(color=_BAD_COLOUR) + return cmap + + +def _sequential_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list( + "tol_sunset_seq", _SUNSET_COLORS[5:], N=256 + ) + cmap.set_bad(color=_BAD_COLOUR) + return cmap + + +# --------------------------------------------------------------------------- +# PR-AUC computation +# --------------------------------------------------------------------------- +def compute_pr_auc( + input_dataset: pd.DataFrame, + confidence_column: str, + label_column: str, +) -> float: + """Compute Area Under Curve for precision-recall curve.""" + if len(input_dataset) == 0: + return 0.0 + + sorted_data = input_dataset[[confidence_column, label_column]].sort_values( + by=confidence_column, ascending=False + ) + + cum_correct = np.cumsum(sorted_data[label_column]) + precision = cum_correct / np.arange(1, len(sorted_data) + 1) + recall = ( + cum_correct / cum_correct.iloc[-1] + if cum_correct.iloc[-1] > 0 + else np.zeros_like(cum_correct) + ) + + if len(precision) < 2: + return 0.0 + + from sklearn.metrics import auc + + return auc(recall, precision) + + +# --------------------------------------------------------------------------- +# Heatmap creation +# --------------------------------------------------------------------------- +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + """Save figure as both PNG and PDF.""" + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def create_auc_heatmap( + auc_df: pd.DataFrame, + output_path: Path, + title: str = "Calibrator generalisation PR-AUC heatmap", +) -> None: + """Create and save a heatmap of PR-AUC values.""" + fig, ax = plt.subplots(figsize=(12, 10)) + + sns.heatmap( + auc_df, + annot=True, + fmt=".3f", + cmap=_sequential_cmap(), + cbar_kws={"label": "PR-AUC"}, + square=True, + linewidths=0.5, + ax=ax, + ) + + ax.set_title(title) + ax.set_xlabel("Test dataset") + ax.set_ylabel("Train dataset") + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") + ax.set_yticklabels(ax.get_yticklabels(), rotation=0) + + base = str(output_path).removesuffix(".png") + _save_fig(fig, Path(base)) + logger.info("Heatmap saved to %s", output_path) + + +def create_comparison_heatmaps(results_path: Path, output_dir: Path) -> None: + """Create heatmaps comparing raw vs calibrated confidence PR-AUC values.""" + logger.info("Scanning results from %s", results_path) + results = pl.scan_csv(results_path) + + trained_datasets = sorted( + results.select(pl.col("trained_on_dataset")) + .unique() + .collect() + .to_series() + .to_list() + ) + test_datasets = sorted( + results.select(pl.col("test_dataset")).unique().collect().to_series().to_list() + ) + logger.info("Trained datasets: %s", trained_datasets) + logger.info("Test datasets: %s", test_datasets) + + trained_labels = [SPECIES_NAME_MAPPING.get(ds, ds) for ds in trained_datasets] + test_labels = [SPECIES_NAME_MAPPING.get(ds, ds) for ds in test_datasets] + + # Compute PR-AUC matrices for both confidence types + auc_matrices = {} + for conf_type in ["confidence", "calibrated_confidence"]: + auc_matrix = [] + for trained_dataset in trained_datasets: + auc_row = [] + for test_dataset in test_datasets: + logger.info( + "Computing PR-AUC (%s) for trained=%s, test=%s", + conf_type, + trained_dataset, + test_dataset, + ) + subset = ( + results.filter( + (pl.col("trained_on_dataset") == trained_dataset) + & (pl.col("test_dataset") == test_dataset) + ) + .collect() + .to_pandas() + ) + + if len(subset) > 0: + auc_row.append(compute_pr_auc(subset, conf_type, "correct")) + else: + auc_row.append(np.nan) + auc_matrix.append(auc_row) + + auc_matrices[conf_type] = pd.DataFrame( + auc_matrix, index=trained_labels, columns=test_labels + ) + + # Individual heatmaps + for conf_type, auc_df in auc_matrices.items(): + conf_name = conf_type.replace("_", " ") + output_path = ( + output_dir / f"calibrator_generalisation_{conf_type}_auc_heatmap.png" + ) + create_auc_heatmap( + auc_df, + output_path, + f"Calibrator generalisation {conf_name} PR-AUC", + ) + + # Difference heatmap (calibrated - raw) + diff_matrix = auc_matrices["calibrated_confidence"] - auc_matrices["confidence"] + + fig, ax = plt.subplots(figsize=(12, 10)) + sns.heatmap( + diff_matrix, + annot=True, + fmt=".3f", + cmap=_diverging_cmap(), + center=0, + cbar_kws={"label": r"PR-AUC difference $(\mathrm{calibrated} - \mathrm{raw})$"}, + square=True, + linewidths=0.5, + ax=ax, + ) + ax.set_title("Calibrator generalisation PR-AUC improvement") + ax.set_xlabel("Test dataset") + ax.set_ylabel("Train dataset") + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") + ax.set_yticklabels(ax.get_yticklabels(), rotation=0) + + diff_base = output_dir / "calibrator_generalisation_auc_difference_heatmap" + _save_fig(fig, diff_base) + logger.info("Difference heatmap saved to %s", diff_base) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +_DEFAULT_OUTPUT_DIR = Path("results/generalisation/plots") + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + + +@app.command() +def main( + results_path: Annotated[ + Path, typer.Option(help="Path to calibrator generalisation results CSV.") + ], + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory to save plots."), + ] = _DEFAULT_OUTPUT_DIR, +) -> None: + """Create PR-AUC heatmaps for calibrator generalisation results.""" + plots_dir.mkdir(parents=True, exist_ok=True) + + if not results_path.exists(): + logger.error("Results file not found: %s", results_path) + raise typer.Exit(1) + + logger.info("Loading results from: %s", results_path) + logger.info("Saving plots to: %s", plots_dir) + + create_comparison_heatmaps(results_path, plots_dir) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/plot_eval_results.py b/paper_scripts/plot_eval_results.py new file mode 100644 index 00000000..3addab87 --- /dev/null +++ b/paper_scripts/plot_eval_results.py @@ -0,0 +1,995 @@ +"""Generate publication-quality evaluation plots from ``winnow predict`` outputs. + +Supports both annotated (database-grounded) and raw (proteome-hit) evaluation +modes, producing six plots per project: precision-recall, FDR run, true vs +estimated FDR (full + zoomed), probability calibration, and before/after score +histograms. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Annotated +import warnings + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import typer +from rich.logging import RichHandler + +from winnow.fdr.nonparametric import NonParametricFDRControl + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# Filter by the specific message or category +warnings.filterwarnings("ignore", message=".*range of fitted confidence scores.*") +warnings.filterwarnings("ignore", message=".*range of fitted FDR thresholds.*") + +# --------------------------------------------------------------------------- +# Dataset display names +# --------------------------------------------------------------------------- +DATASET_DISPLAY_NAMES: dict[str, str] = { + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD004732": "ProteomeTools-1", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", +} + +_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") +_DISPLAY_NAME_LOOKUP = {k.lower(): v for k, v in DATASET_DISPLAY_NAMES.items()} + +# Paul Tol "bright" palette (colour-blind safe) +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_MAIN_LINE_COLOUR = _PALETTE[0] +_RAW_LINE_COLOUR = _PALETTE[5] +_IDEAL_LINE_COLOUR = _PALETTE[6] +_BAND_COLOUR = _PALETTE[0] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_DIAGNOSTIC_ALPHAS = (0.01, 0.05, 0.10) +_HOEFFDING_DELTA = 0.05 + + +def _safe_stem(key: str) -> str: + """Return a flat, filesystem-safe filename stem for a project key.""" + return key.strip().replace("/", "_") + + +def _normalize_project_key(key: str) -> str: + """Strip eval suffixes and nested path segments for display-name lookup.""" + key = key.strip() + if "/" in key: + key = key.rsplit("/", 1)[-1] + for suffix in _FOLDER_SUFFIXES: + if key.endswith(suffix): + return key[: -len(suffix)] + return key + + +def _display_name(key: str) -> str: + """Look up the publication-ready display name for a dataset key.""" + normalized = _normalize_project_key(key) + if normalized in DATASET_DISPLAY_NAMES: + return DATASET_DISPLAY_NAMES[normalized] + return _DISPLAY_NAME_LOOKUP.get(normalized.lower(), normalized) + + +def _ground_truth_qualifier(eval_type: str) -> str: + """Return the title-friendly ground truth qualifier for plot titles.""" + if eval_type in ("annotated", "labelled"): + return "using database search" + return "using proteome mapping" + + +def _save_fig(fig: plt.Figure, base_path: Path) -> None: + """Save figure as both PNG and PDF.""" + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +# --------------------------------------------------------------------------- +# PR curve (non-standard cumulative definition) +# --------------------------------------------------------------------------- +def _compute_precision_recall( + df: pd.DataFrame, confidence_col: str = "calibrated_confidence" +) -> pd.DataFrame: + """Non-standard cumulative PR curve matching the codebase convention.""" + sorted_df = df.sort_values(confidence_col, ascending=False) + labels = sorted_df["correct"].values + cum_correct = np.cumsum(labels) + n = len(labels) + precision = cum_correct / np.arange(1, n + 1) + recall = cum_correct / n + return pd.DataFrame({"precision": precision, "recall": recall}) + + +def plot_precision_recall( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot precision-recall curve.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + pr_cal = _compute_precision_recall(df, "calibrated_confidence") + pr_raw = _compute_precision_recall(df, "confidence") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + pr_raw["recall"], + pr_raw["precision"], + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Raw confidence", + ) + ax.plot( + pr_cal["recall"], + pr_cal["precision"], + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Calibrated confidence", + ) + ax.set_xlabel("Recall") + ax.set_ylabel("Precision") + ax.set_title(f"{display} precision-recall {qualifier}") + ax.set_xlim(0, 1) + ax.set_ylim(0, 1.02) + ax.legend(loc="lower right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"pr_curve_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# FDR run plot +# --------------------------------------------------------------------------- +def plot_fdr_run( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot calibrated confidence vs estimated and true PSM FDR.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + true_fdr_df = true_fdr_ctrl.add_psm_fdr( + df.copy(), confidence_col="calibrated_confidence" + ) + true_fdr_df = true_fdr_df.sort_values("calibrated_confidence") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + df["calibrated_confidence"].values, + df["psm_fdr"].values, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_fdr_df["calibrated_confidence"].values, + true_fdr_df["psm_fdr"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM FDR") + ax.set_title(f"{display} FDR run {qualifier}") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_run_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# Q-value run plot +# --------------------------------------------------------------------------- +def _fit_database_grounded_fdr( + df: pd.DataFrame, + confidence_col: str = "calibrated_confidence", + correct_col: str = "correct", + drop: int = 10, +) -> NonParametricFDRControl: + """Fit an FDR controller using ground-truth labels. + + Replicates the fitting logic of ``DatabaseGroundedFDRControl`` (computing + FDR as 1 − precision over sorted predictions, with the first *drop* entries + removed) without pulling in the instanovo dependency. + """ + sorted_desc = df.sort_values(confidence_col, ascending=False) + labels = sorted_desc[correct_col].values.astype(float) + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + confidence = sorted_desc[confidence_col].values + + ctrl = NonParametricFDRControl() + ctrl._fdr_values = (1.0 - precision)[drop:] + ctrl._confidence_scores = confidence[drop:] + return ctrl + + +def plot_q_value_run( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot calibrated confidence vs estimated and true PSM q-values.""" + if "psm_q_value" not in df.columns: + logger.warning( + "Skipping q-value run plot for %s: psm_q_value column missing", project + ) + return + + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + sorted_df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + qval_input = df[["calibrated_confidence"]].copy() + true_q_df = true_fdr_ctrl.add_psm_q_value( + qval_input, confidence_col="calibrated_confidence" + ) + true_q_df = true_q_df.sort_values("calibrated_confidence") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + sorted_df["calibrated_confidence"].values, + sorted_df["psm_q_value"].values, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_q_df["calibrated_confidence"].values, + true_q_df["psm_q_value"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM q-value") + ax.set_title(f"{display} q-value run {qualifier}") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"qvalue_run_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# FDR / q-value run plots with Hoeffding confidence bands +# --------------------------------------------------------------------------- +def _hoeffding_band_arrays(n: int) -> np.ndarray: + """Compute pointwise Hoeffding half-widths for ranks 1..n (descending confidence).""" + ranks = np.arange(1, n + 1) + return np.sqrt(np.log(2.0 / _HOEFFDING_DELTA) / (2.0 * ranks)) + + +def plot_fdr_run_with_bands( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """FDR run plot with Hoeffding 95% confidence band on the non-parametric curve.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + true_fdr_df = true_fdr_ctrl.add_psm_fdr( + df.copy(), confidence_col="calibrated_confidence" + ) + true_fdr_df = true_fdr_df.sort_values("calibrated_confidence") + + fdr_vals = df["psm_fdr"].values + conf_vals = df["calibrated_confidence"].values + n = len(fdr_vals) + hw = _hoeffding_band_arrays(n)[::-1] + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.fill_between( + conf_vals, + np.clip(fdr_vals - hw, 0, None), + np.clip(fdr_vals + hw, None, 1), + color=_BAND_COLOUR, + alpha=0.2, + label="95% Hoeffding bound", + ) + ax.plot( + conf_vals, + fdr_vals, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_fdr_df["calibrated_confidence"].values, + true_fdr_df["psm_fdr"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM FDR") + ax.set_title(f"{display} FDR run with sampling error bounds {qualifier}") + ax.legend(loc="upper right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_run_bands_{_safe_stem(project)}") + + +def plot_q_value_run_with_bands( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Q-value run plot with Hoeffding 95% confidence band on the non-parametric curve.""" + if "psm_q_value" not in df.columns: + logger.warning( + "Skipping banded q-value run plot for %s: psm_q_value column missing", + project, + ) + return + + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + sorted_df = df.sort_values("calibrated_confidence") + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + qval_input = df[["calibrated_confidence"]].copy() + true_q_df = true_fdr_ctrl.add_psm_q_value( + qval_input, confidence_col="calibrated_confidence" + ) + true_q_df = true_q_df.sort_values("calibrated_confidence") + + qvals = sorted_df["psm_q_value"].values + conf_vals = sorted_df["calibrated_confidence"].values + n = len(qvals) + hw = _hoeffding_band_arrays(n)[::-1] + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.fill_between( + conf_vals, + np.clip(qvals - hw, 0, None), + np.clip(qvals + hw, None, 1), + color=_BAND_COLOUR, + alpha=0.2, + label="95% Hoeffding bound", + ) + ax.plot( + conf_vals, + qvals, + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Non-parametric", + ) + ax.plot( + true_q_df["calibrated_confidence"].values, + true_q_df["psm_q_value"].values, + color=_RAW_LINE_COLOUR, + lw=1.5, + label="Database-grounded", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM q-value") + ax.set_title(f"{display} q-value run with sampling error bounds {qualifier}") + ax.legend(loc="upper center") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"qvalue_run_bands_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# True FDR vs estimated FDR +# --------------------------------------------------------------------------- +def _compute_true_vs_estimated_fdr(df: pd.DataFrame) -> pd.DataFrame: + """Compute true and estimated FDR arrays, sorted by confidence descending.""" + sorted_df = df.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + + true_fdr_ctrl = _fit_database_grounded_fdr(sorted_df) + with_true_fdr = true_fdr_ctrl.add_psm_fdr( + sorted_df, confidence_col="calibrated_confidence" + ) + + return pd.DataFrame( + { + "estimated_fdr": sorted_df["psm_fdr"].values, + "true_fdr": with_true_fdr["psm_fdr"].values, + } + ) + + +def plot_true_vs_estimated_fdr( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, + *, + zoomed: bool = False, +) -> None: + """Plot true FDR vs estimated FDR.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + fdr_data = _compute_true_vs_estimated_fdr(df) + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + fdr_data["estimated_fdr"], + fdr_data["true_fdr"], + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Observed", + ) + + # Only plot the ideal line up to the max extent of the observed line + max_x = float(fdr_data["estimated_fdr"].max()) + max_y = float(fdr_data["true_fdr"].max()) + lim = 0.1 if zoomed else 1.0 + ideal_end = min(lim, max(max_x, max_y)) + + ax.plot( + [0, ideal_end], + [0, ideal_end], + ls="--", + color=_IDEAL_LINE_COLOUR, + lw=1, + label="Perfectly calibrated", + ) + ax.set_xlabel("Non-parametric estimated FDR") + ax.set_ylabel("Database-grounded FDR") + zoom_suffix = " (0 to 0.1)" if zoomed else "" + ax.set_title(f"{display} true vs estimated FDR{zoom_suffix} {qualifier}") + if zoomed: + ax.set_xlim(0, 0.1) + ax.set_ylim(0, 0.1) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + tag = "fdr_true_vs_est_zoom" if zoomed else "fdr_true_vs_est" + _save_fig(fig, output_dir / f"{tag}_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# True q-values vs estimated q-values +# --------------------------------------------------------------------------- +def _compute_true_vs_estimated_q_values(df: pd.DataFrame) -> pd.DataFrame: + """Compute true and estimated q-value arrays, sorted by confidence descending.""" + sorted_df = df.sort_values("calibrated_confidence", ascending=False).reset_index( + drop=True + ) + + true_q_val_ctrl = _fit_database_grounded_fdr(sorted_df) + qval_input = sorted_df[["calibrated_confidence"]].copy() + with_true_q_df = true_q_val_ctrl.add_psm_q_value( + qval_input, confidence_col="calibrated_confidence" + ) + return pd.DataFrame( + { + "estimated_q_value": sorted_df["psm_q_value"].values, + "true_q_value": with_true_q_df["psm_q_value"].values, + } + ) + + +def plot_true_vs_estimated_q_values( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, + *, + zoomed: bool = False, +) -> None: + """Plot true q-values vs estimated q-values.""" + if "psm_q_value" not in df.columns: + logger.warning( + "Skipping true vs estimated q-value plot for %s: psm_q_value column missing", + project, + ) + return + + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + q_value_data = _compute_true_vs_estimated_q_values(df) + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + q_value_data["estimated_q_value"], + q_value_data["true_q_value"], + color=_MAIN_LINE_COLOUR, + lw=1.5, + label="Observed", + ) + + # Only plot the ideal line up to the max extent of the observed line + max_x = float(q_value_data["estimated_q_value"].max()) + max_y = float(q_value_data["true_q_value"].max()) + lim = 0.1 if zoomed else 1.0 + ideal_end = min(lim, max(max_x, max_y)) + + ax.plot( + [0, ideal_end], + [0, ideal_end], + ls="--", + color=_IDEAL_LINE_COLOUR, + lw=1, + label="Perfectly calibrated", + ) + ax.set_xlabel("Non-parametric estimated q-values") + ax.set_ylabel("Database-grounded q-values") + zoom_suffix = " (0 to 0.1)" if zoomed else "" + ax.set_title(f"{display} true vs estimated q-values{zoom_suffix} {qualifier}") + if zoomed: + ax.set_xlim(0, 0.1) + ax.set_ylim(0, 0.1) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + tag = "qvalue_true_vs_est_zoom" if zoomed else "qvalue_true_vs_est" + _save_fig(fig, output_dir / f"{tag}_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# Probability calibration (reliability diagram) +# --------------------------------------------------------------------------- +def _compute_calibration_curve( + df: pd.DataFrame, + pred_col: str, + label_col: str, + n_bins: int = 10, +) -> pd.DataFrame: + """Fixed-width bin calibration curve.""" + data = df[[pred_col, label_col]].dropna().copy() + data[pred_col] = data[pred_col].clip(0.0, 1.0) + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_cats = pd.cut(data[pred_col], bins=bins, include_lowest=True) + bin_cats.name = "bin" + grouped = ( + data.groupby(bin_cats, observed=True) + .agg( + pred_mean=(pred_col, "mean"), + empirical=(label_col, "mean"), + count=(label_col, "size"), + ) + .reset_index() + ) + grouped = grouped[grouped["count"] > 0] + grouped["bin_center"] = grouped["bin"].apply(lambda iv: (iv.left + iv.right) / 2) + return grouped[["pred_mean", "empirical", "count", "bin_center"]] + + +def _estimate_calibration_values( + df: pd.DataFrame, + pred_col: str, + label_col: str, + n_bins: int = 20, +) -> np.ndarray: + """Estimate c(s) for each PSM via binned calibration. + + Returns an array of the same length as *df* where each entry is the + empirical accuracy of the bin that PSM falls into. + """ + scores = df[pred_col].values.clip(0.0, 1.0) + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_idx = np.digitize(scores, bins) - 1 + bin_idx = np.clip(bin_idx, 0, n_bins - 1) + labels = df[label_col].values.astype(float) + bin_sums = np.bincount(bin_idx, weights=labels, minlength=n_bins) + bin_counts = np.bincount(bin_idx, minlength=n_bins).astype(float) + bin_counts[bin_counts == 0] = 1.0 + bin_means = bin_sums / bin_counts + return bin_means[bin_idx] + + +def _hoeffding_halfwidth(k: int, delta: float = _HOEFFDING_DELTA) -> float: + """Hoeffding 95% confidence half-width for a mean of *k* bounded [0,1] r.v.s.""" + if k <= 0: + return float("nan") + return float(np.sqrt(np.log(2.0 / delta) / (2.0 * k))) + + +def plot_calibration( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot probability calibration (reliability diagram).""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + cal_calibrated = _compute_calibration_curve(df, "calibrated_confidence", "correct") + cal_raw = _compute_calibration_curve(df, "confidence", "correct") + + fig, ax = plt.subplots(figsize=(6, 4)) + ax.plot( + cal_raw["pred_mean"], + cal_raw["empirical"], + marker="D", + color=_RAW_LINE_COLOUR, + label="Raw confidence", + ) + ax.plot( + cal_calibrated["pred_mean"], + cal_calibrated["empirical"], + marker="o", + color=_MAIN_LINE_COLOUR, + label="Calibrated confidence", + ) + ax.plot( + [0, 1], + [0, 1], + ls="--", + color=_IDEAL_LINE_COLOUR, + lw=1, + label="Perfectly calibrated", + ) + ax.set_xlabel("Mean predicted probability") + ax.set_ylabel("Empirical accuracy") + ax.set_title(f"{display} probability calibration {qualifier}") + ax.set_xlim(0, 1) + ax.set_ylim(0, 1) + ax.legend(loc="lower right") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"calibration_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# Before/after score histograms +# --------------------------------------------------------------------------- +def plot_score_histograms( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Plot before/after score histograms with correct/incorrect overlays.""" + display = _display_name(project) + qualifier = _ground_truth_qualifier(eval_type) + + correct_mask = df["correct"].astype(bool) + + fig, axes = plt.subplots(2, 1, figsize=(7, 6), sharex=False) + + # Before calibration + ax = axes[0] + bins_before = np.linspace(0, 1, 51) + ax.hist( + df.loc[correct_mask, "confidence"], + bins=bins_before, + alpha=0.5, + color=_CORRECT_COLOUR, + edgecolor="black", + label="Correct", + ) + ax.hist( + df.loc[~correct_mask, "confidence"], + bins=bins_before, + alpha=0.5, + color=_INCORRECT_COLOUR, + edgecolor="black", + label="Incorrect", + ) + ax.set_xlabel("Raw confidence") + ax.set_ylabel("Count") + ax.set_title("Before calibration") + ax.legend(loc="upper center") + _style_ax(ax) + + # After calibration + ax = axes[1] + bins_after = np.linspace(0, 1, 51) + ax.hist( + df.loc[correct_mask, "calibrated_confidence"], + bins=bins_after, + alpha=0.5, + color=_CORRECT_COLOUR, + edgecolor="black", + label="Correct", + ) + ax.hist( + df.loc[~correct_mask, "calibrated_confidence"], + bins=bins_after, + alpha=0.5, + color=_INCORRECT_COLOUR, + edgecolor="black", + label="Incorrect", + ) + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("Count") + ax.set_title("After calibration") + ax.legend(loc="upper center") + _style_ax(ax) + + fig.suptitle(f"{display} score distributions {qualifier}", fontsize=13) + fig.tight_layout() + _save_fig(fig, output_dir / f"score_histograms_{_safe_stem(project)}") + + +# --------------------------------------------------------------------------- +# Diagnostics CSV +# --------------------------------------------------------------------------- +def _is_labelled(eval_type: str) -> bool: + return eval_type in ("annotated", "labelled") + + +def _compute_diagnostics( + df: pd.DataFrame, + eval_type: str, + alphas: tuple[float, ...] = _DIAGNOSTIC_ALPHAS, +) -> pd.DataFrame: + """Compute FDR diagnostics at each target alpha. + + Label-dependent metrics (sTECE, TECE, realised FDR, etc.) are only + populated for annotated/labelled eval types. + """ + labelled = _is_labelled(eval_type) + + np_ctrl = NonParametricFDRControl() + np_ctrl.fit(dataset=df["calibrated_confidence"]) + + if labelled: + db_ctrl = _fit_database_grounded_fdr(df) + c_hat = _estimate_calibration_values(df, "calibrated_confidence", "correct") + scores = df["calibrated_confidence"].values.clip(0.0, 1.0) + + rows: list[dict] = [] + for alpha in alphas: + tau_hat = np_ctrl.get_confidence_cutoff(threshold=alpha) + if np.isnan(tau_hat): + rows.append({"alpha": alpha, "tau_hat": float("nan")}) + continue + + mask_hat = df["calibrated_confidence"].values >= tau_hat + k = int(mask_hat.sum()) + est_fdr = float(np_ctrl.compute_fdr(tau_hat)) + eps = _hoeffding_halfwidth(k) + + row: dict = { + "alpha": alpha, + "tau_hat": float(tau_hat), + "k_accepted": k, + "estimated_fdr": est_fdr, + "hoeffding_halfwidth": eps, + } + + if labelled: + residuals = c_hat[mask_hat] - scores[mask_hat] + row["stece"] = float(np.mean(residuals)) + row["tece"] = float(np.mean(np.abs(residuals))) + row["tece_2"] = float(np.sqrt(np.mean(residuals**2))) + + realised_fdr = float(db_ctrl.compute_fdr(tau_hat)) + row["realised_fdr"] = realised_fdr + row["fdr_bias"] = est_fdr - realised_fdr + + tau_star = db_ctrl.get_confidence_cutoff(threshold=alpha) + row["tau_star"] = float(tau_star) + if not np.isnan(tau_star): + k_star = int((df["calibrated_confidence"].values >= tau_star).sum()) + row["discovery_count_shift"] = k - k_star + else: + row["discovery_count_shift"] = float("nan") + + rows.append(row) + + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- +def _load_project_data( + predictions_root: Path, + project: str, + suffix: str, + eval_type: str, +) -> pd.DataFrame: + """Load and merge metadata.csv and preds_and_fdr_metrics.csv for a project.""" + # folder = predictions_root / f"{project}_{suffix}" + folder = predictions_root / f"{project}" + preds_path = folder / "preds_and_fdr_metrics.csv" + meta_path = folder / "metadata.csv" + if not preds_path.is_file(): + raise FileNotFoundError(f"Missing predictions file: {preds_path}") + + preds_df = pd.read_csv(preds_path) + + if meta_path.is_file(): + meta_df = pd.read_csv(meta_path) + # Drop columns already present in preds to avoid duplicates on merge + overlap = [ + c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" + ] + if overlap: + meta_df = meta_df.drop(columns=overlap) + df = preds_df.merge(meta_df, on="spectrum_id", how="left") + else: + df = preds_df + + if eval_type in ("raw", "unlabelled"): + if "proteome_hit" not in df.columns: + raise ValueError( + f"Expected 'proteome_hit' column for eval-type={eval_type} in {preds_path}" + ) + df["correct"] = df["proteome_hit"].astype(float) + + required = ["confidence", "calibrated_confidence", "correct"] + missing = [c for c in required if c not in df.columns] + if missing: + raise ValueError(f"Missing columns {missing} in {preds_path}") + + return df + + +def generate_all_plots( + df: pd.DataFrame, + project: str, + eval_type: str, + output_dir: Path, +) -> None: + """Generate all plots for a single project.""" + output_dir.mkdir(parents=True, exist_ok=True) + + plot_precision_recall(df, project, eval_type, output_dir) + plot_fdr_run(df, project, eval_type, output_dir) + plot_fdr_run_with_bands(df, project, eval_type, output_dir) + plot_q_value_run(df, project, eval_type, output_dir) + plot_q_value_run_with_bands(df, project, eval_type, output_dir) + plot_true_vs_estimated_fdr(df, project, eval_type, output_dir, zoomed=False) + plot_true_vs_estimated_fdr(df, project, eval_type, output_dir, zoomed=True) + plot_true_vs_estimated_q_values(df, project, eval_type, output_dir, zoomed=False) + plot_true_vs_estimated_q_values(df, project, eval_type, output_dir, zoomed=True) + plot_calibration(df, project, eval_type, output_dir) + plot_score_histograms(df, project, eval_type, output_dir) + + +_EVAL_TYPE_SUFFIX: dict[str, str] = { + "annotated": "annotated", + "raw": "raw", + "labelled": "labelled", + "unlabelled": "unlabelled", +} + + +@app.command() +def main( + predictions_root: Annotated[ + Path, + typer.Option( + "--predictions-root", + help="Root directory containing per-project prediction folders.", + ), + ], + projects: Annotated[ + str, + typer.Option( + "--projects", + help="Space- or comma-separated project keys (e.g. 'helaqc,gluc' or 'helaqc gluc').", + ), + ], + eval_type: Annotated[ + str, + typer.Option( + "--eval-type", + help="Evaluation type: annotated, raw, labelled, or unlabelled.", + ), + ], + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for summary/diagnostics CSVs."), + ], + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for png/pdf figures."), + ], +) -> None: + """Generate evaluation plots from winnow predict outputs.""" + logging.basicConfig(level=logging.INFO, format="%(message)s", datefmt="%H:%M:%S") + + if eval_type not in _EVAL_TYPE_SUFFIX: + raise typer.BadParameter( + f"Unknown eval-type {eval_type!r}. Expected one of: {list(_EVAL_TYPE_SUFFIX)}" + ) + + project_list = [p.strip() for p in projects.replace(",", " ").split() if p.strip()] + if not project_list: + raise typer.BadParameter("No projects specified.") + + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + suffix = _EVAL_TYPE_SUFFIX[eval_type] + + for project in project_list: + display = _display_name(project) + logger.info("Processing %s (%s, eval-type=%s)...", project, display, eval_type) + + df = _load_project_data(predictions_root, project, suffix, eval_type) + logger.info(" Loaded %d rows", len(df)) + + true_fdr_ctrl = _fit_database_grounded_fdr(df) + db_fdr = true_fdr_ctrl.add_psm_fdr( + df[["calibrated_confidence"]].copy(), confidence_col="calibrated_confidence" + ) + df["db_grounded_psm_fdr"] = db_fdr["psm_fdr"] + db_qval = true_fdr_ctrl.add_psm_q_value( + df[["calibrated_confidence"]].copy(), confidence_col="calibrated_confidence" + ) + df["db_grounded_psm_q_value"] = db_qval["psm_q_value"] + + summary_cols = ["confidence", "calibrated_confidence", "correct"] + if "psm_fdr" in df.columns: + summary_cols.append("psm_fdr") + summary_cols.append("db_grounded_psm_fdr") + if "psm_q_value" in df.columns: + summary_cols.append("psm_q_value") + summary_cols.append("db_grounded_psm_q_value") + stem = _safe_stem(project) + df[summary_cols].to_csv(results_dir / f"{stem}_summary.csv", index=False) + + diag = _compute_diagnostics(df, eval_type) + diag.to_csv(results_dir / f"{stem}_diagnostics.csv", index=False) + logger.info(" Diagnostics saved (%d alpha levels)", len(diag)) + + generate_all_plots(df, project, eval_type, plots_dir) + logger.info(" Plots saved to %s", plots_dir) + + logger.info("Done. Tables in %s; plots in %s", results_dir, plots_dir) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/plot_fdr_method_comparison.py b/paper_scripts/plot_fdr_method_comparison.py new file mode 100644 index 00000000..ff6f443a --- /dev/null +++ b/paper_scripts/plot_fdr_method_comparison.py @@ -0,0 +1,1121 @@ +#!/usr/bin/env python3 +"""Compare PSM-level FDR estimates from Winnow and NovoBoard. + +Plots and summaries cover labelled-test Novor correctness and unlabelled +reference-proteome membership at 1 %, 5 %, and 10 % FDR on a shared filtered +spectrum pool (NovoBoard mass-deltas converted to ProForma; unsupported +modifications dropped; NovoBoard target-decoy pairs gated). Unlabelled panels +also drop normalised peptides shorter than 8 residues (proteome-substring +proxy); labelled panels keep short peptides because correctness is Novor +agreement. The shared pool is the twin-valid NovoBoard set; Winnow is trimmed +to match under the invariant that NovoBoard ⊆ Winnow after identical InstaNovo +filters. + +A long-form ``fdr_method_comparison_curves.csv`` (per spectrum x method) is +written so plots and summary tables can be regenerated with ``--summarise-only``. + +External NovoBoard inputs (``--novoboard-root``) must follow +``{root}/{dataset}/novoboard/`` with target/decoy CSVs such as +``annotated_test.csv``, ``annotated_test_decoy_{rate}.csv``, +``raw_unlabelled.csv`` and ``raw_unlabelled_decoy_{rate}.csv``. Point +``--novoboard-root`` at the ``datasets`` directory of a NovoBoard checkout. +Local results were produced from the fork +``git@github.com:JemmaLDaniel/NovoBoard.git``, branch +``feat/adapt-to-instanovo`` at commit +``a9faab3ef1af06987599c2f01e6ba96072c80172``. +""" + +from __future__ import annotations + +import logging +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Annotated, Literal, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +from winnow.utils.proteome import load_proteome_haystack # noqa: E402 +from fdr_tool_comparison_preprocess import ( # noqa: E402 + LABELLED_MIN_PEPTIDE_LENGTH, + MIN_PEPTIDE_LENGTH, + assert_shared_prediction_keys, + attach_labels_by_spectrum_id, + attach_novoboard_pair_keys, + compute_q_values, + filter_novoboard_target_decoy_pairs, + filter_prediction_table, + label_series_by_spectrum_id, + load_residue_masses, + novoboard_psm_tdc, + novor_correctness_mask, + proteome_hit_mask, + restrict_winnow_to_novoboard_spectra, +) +from fdr_tool_comparison_summaries import ( # noqa: E402 + SUMMARY_THRESHOLDS, + acceptance_rows_from_q, + error_rows_from_q, + finalise_error_gain_table, + write_summary_tables, +) +from plot_eval_results import ( # noqa: E402 + _MAIN_LINE_COLOUR, + _PALETTE, + _RAW_LINE_COLOUR, + _display_name, + _ground_truth_qualifier, + _save_fig, + _style_ax, +) +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +PRIMARY_METHOD = "Winnow (non-parametric)" +DB_CAL_METHOD = "Database-grounded (calibrated confidence)" +DB_RAW_METHOD = "Database-grounded (raw confidence)" +NOVOBOARD_METHOD = "NovoBoard" +CURVES_CSV_NAME = "fdr_method_comparison_curves.csv" +_WINNOW_METHODS = (PRIMARY_METHOD, DB_CAL_METHOD, DB_RAW_METHOD) +_METHOD_ORDER = (*_WINNOW_METHODS, NOVOBOARD_METHOD) + +logger = logging.getLogger(__name__) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +FDR_THRESHOLDS = [0.01, 0.05, 0.10] +_DB_GROUNDED_DROP = 10 + +DEFAULT_WINNOW_RESULTS = _REPO_ROOT / "results" +DEFAULT_MODEL_ROOT = _REPO_ROOT / "models" +DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/fdr_method_comparison_psm" +DEFAULT_DATASETS = ["helaqc", "celegans"] +_METHOD_COLOURS = { + PRIMARY_METHOD: _MAIN_LINE_COLOUR, + DB_CAL_METHOD: _RAW_LINE_COLOUR, + DB_RAW_METHOD: _PALETTE[3], + NOVOBOARD_METHOD: _PALETTE[2], +} + +EvalType = Literal["labelled", "unlabelled"] + +_DATASET_META = { + "helaqc": { + "fasta": "fasta/human.fasta", + "novoboard_decoy": "0.50", + "winnow_suffix": "helaqc", + }, + "celegans": { + "fasta": "fasta/celegans.fasta", + "novoboard_decoy": "0.70", + "winnow_suffix": "celegans", + }, + "sbrodae": { + "fasta": "fasta/Sb_proteome.fasta", + "novoboard_decoy": "0.50", + "winnow_suffix": "sbrodae", + }, + "PXD019483": { + "fasta": "fasta/human.fasta", + "novoboard_decoy": "0.70", + "winnow_suffix": "pxd019483", + }, +} + + +@dataclass(frozen=True) +class DatasetConfig: + """Paths and metadata for one evaluation dataset.""" + + key: str + fasta: Path + winnow_unlabelled: Path + winnow_test: Path + novoboard_dir: Path + novoboard_decoy_rate: str + calibrator_train_metadata: Path + + +def build_dataset_configs( + winnow_results: Path = DEFAULT_WINNOW_RESULTS, + *, + novoboard_root: Path, + model_root: Path = DEFAULT_MODEL_ROOT, +) -> dict[str, DatasetConfig]: + """Build per-dataset path bundles from repo roots.""" + configs: dict[str, DatasetConfig] = {} + for key, meta in _DATASET_META.items(): + suffix = meta["winnow_suffix"] + configs[key] = DatasetConfig( + key=key, + fasta=_REPO_ROOT / meta["fasta"], + winnow_unlabelled=winnow_results + / f"instanovo_{suffix}_predictions_unlabelled", + winnow_test=winnow_results / f"instanovo_{suffix}_predictions_test", + novoboard_dir=novoboard_root / f"{key}/novoboard", + novoboard_decoy_rate=meta["novoboard_decoy"], + calibrator_train_metadata=model_root + / f"instanovo_{suffix}/metadata_train.parquet", + ) + return configs + + +@dataclass +class MethodCurve: + """One method's confidence and q-value arrays for curve plotting.""" + + label: str + color: str + confidence: np.ndarray + q_value: np.ndarray + + +@dataclass +class MethodCounts: + """PSM counts per q-value threshold for one method.""" + + label: str + color: str + counts: list[int] + + +@dataclass +class MethodRecovery: + """Correct labelled identifications recovered at q-value thresholds.""" + + label: str + color: str + q_value: np.ndarray + correct: np.ndarray + + +def _load_residue_masses() -> dict[str, float]: + return load_residue_masses() + + +def _fit_database_grounded_fdr( + df: pd.DataFrame, + correct_col: str, + confidence_col: str, + *, + drop: int = _DB_GROUNDED_DROP, +) -> DatabaseGroundedFDRControl: + """Fit ``DatabaseGroundedFDRControl`` from per-row correctness labels.""" + ctrl = DatabaseGroundedFDRControl( + confidence_feature=confidence_col, + drop=drop, + ) + sorted_df = df.sort_values(confidence_col, ascending=False) + labels = sorted_df[correct_col].astype(float).to_numpy() + conf = sorted_df[confidence_col].to_numpy() + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + ctrl._fdr_values = np.array(1.0 - precision)[drop:] + ctrl._confidence_scores = conf[drop:] + return ctrl + + +def load_winnow( + predictions_dir: Path, fasta: Path, eval_type: EvalType +) -> pd.DataFrame: + """Load Winnow preds + metadata; annotate proteome hits or use labelled ``correct``. + + Always drops unsupported ``[UNIMOD:n]`` tokens. Unlabelled panels also require + normalised peptide length ≥ :data:`MIN_PEPTIDE_LENGTH` (proteome-substring + proxy). Labelled panels only require a non-empty normalised key + (:data:`LABELLED_MIN_PEPTIDE_LENGTH`), because correctness is Novor agreement. + + Labelled ``correct`` keeps Winnow predict-time Novor labels when present; + otherwise recomputes Novor from ``sequence`` / ``prediction``. + """ + preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") + meta_path = predictions_dir / "metadata.csv" + if meta_path.exists(): + meta = pl.read_csv(meta_path, columns=["spectrum_id", "confidence"]) + preds = preds.join(meta, on="spectrum_id", how="inner") + + df = preds.to_pandas() + min_length = ( + LABELLED_MIN_PEPTIDE_LENGTH if eval_type == "labelled" else MIN_PEPTIDE_LENGTH + ) + df = filter_prediction_table( + df, "prediction", min_length=min_length, key_col="peptide_key" + ) + + if eval_type == "labelled": + if "correct" in df.columns: + pass + elif {"sequence", "prediction"}.issubset(df.columns): + df["correct"] = novor_correctness_mask(df["sequence"], df["prediction"]) + else: + raise ValueError( + f"Missing 'correct' (and sequence/prediction) in " + f"{predictions_dir}/preds_and_fdr_metrics.csv" + ) + return df + + haystack = load_proteome_haystack(fasta) + df["proteome_hit"] = proteome_hit_mask( + df["prediction"], haystack, min_length=MIN_PEPTIDE_LENGTH + ) + return df + + +def _effective_db_grounded_drop(n_rows: int, drop: int = _DB_GROUNDED_DROP) -> int: + """Cap drop so FDR fit retains at least one score when *n_rows* is small.""" + return min(drop, max(0, n_rows - 1)) + + +def _vectorized_fdr_from_control( + confidence: np.ndarray, ctrl: DatabaseGroundedFDRControl | NonParametricFDRControl +) -> np.ndarray: + """Vectorized equivalent of ``FDRControl.compute_fdr`` for an array of scores.""" + if ctrl._confidence_scores is None or ctrl._fdr_values is None: + raise AttributeError("FDR method not fitted, please call `fit()` first") + conf = np.asarray(confidence, dtype=float) + scores = np.asarray(ctrl._confidence_scores, dtype=float) + fdr_values = np.asarray(ctrl._fdr_values, dtype=float) + n = len(scores) + idx = np.searchsorted(-scores, -conf, side="left") + fdr = np.empty(len(conf), dtype=float) + below = (idx == n) & (conf < scores[-1]) + above = (idx == 0) & (conf > scores[0]) + normal = ~(below | above) + fdr[below] = 1.0 + fdr[above] = float(fdr_values[0]) + clipped = np.clip(idx[normal], 0, n - 1) + fdr[normal] = fdr_values[clipped] + return fdr + + +def _assign_q_values_fast( + df: pd.DataFrame, + confidence_col: str, + ctrl: DatabaseGroundedFDRControl | NonParametricFDRControl, + out_col: str, +) -> pd.DataFrame: + """Assign q-values without per-row ``compute_fdr`` applies (needed for large tables).""" + work = df.copy() + conf = work[confidence_col].to_numpy(dtype=float) + fdr = _vectorized_fdr_from_control(conf, ctrl) + order = np.argsort(-conf, kind="mergesort") + q_sorted = compute_q_values(fdr[order]) + q = np.empty_like(q_sorted) + q[order] = q_sorted + work[out_col] = q + return work + + +def _add_database_grounded_qvalues( + df: pd.DataFrame, + correct_col: str, + confidence_col: str, + out_col: str, + residue_masses: dict[str, float], + *, + fit_df: pd.DataFrame | None = None, + drop: int = _DB_GROUNDED_DROP, +) -> pd.DataFrame: + """Append database-grounded PSM q-values; fit on *fit_df* (defaults to *df*).""" + reference = fit_df if fit_df is not None else df + work = df.drop(columns=[out_col], errors="ignore").copy() + ctrl = _fit_database_grounded_fdr( + reference, + correct_col, + confidence_col, + drop=_effective_db_grounded_drop(len(reference), drop), + ) + return _assign_q_values_fast(work, confidence_col, ctrl, out_col) + + +def _prepare_winnow_psm_table( + df: pd.DataFrame, + correct_col: str, + residue_masses: dict[str, float], + *, + fit_df: pd.DataFrame | None = None, +) -> pd.DataFrame: + """Append Winnow PSM-level q-value columns while retaining labels.""" + reference = fit_df if fit_df is not None else df + db_cal = _add_database_grounded_qvalues( + df, + correct_col, + "calibrated_confidence", + "psm_q_value_db_cal", + residue_masses, + fit_df=reference, + ) + db_raw = _add_database_grounded_qvalues( + db_cal, + correct_col, + "confidence", + "psm_q_value_db_raw", + residue_masses, + fit_df=reference, + ) + return db_raw + + +def _curves_df_from_winnow_table( + table: pd.DataFrame, + *, + dataset: str, + panel: str, + label_col: str, +) -> pd.DataFrame: + """Long-form curve rows for the three Winnow PSM q-value methods.""" + if "spectrum_id" not in table.columns: + raise KeyError("Winnow curve export requires spectrum_id") + if label_col not in table.columns: + raise KeyError(f"Missing label column {label_col!r}") + label = table[label_col].astype(bool).to_numpy() + spectrum_id = table["spectrum_id"].astype(str) + specs = ( + (PRIMARY_METHOD, "calibrated_confidence", "psm_q_value"), + (DB_CAL_METHOD, "calibrated_confidence", "psm_q_value_db_cal"), + (DB_RAW_METHOD, "confidence", "psm_q_value_db_raw"), + ) + parts: list[pd.DataFrame] = [] + for method, score_col, q_col in specs: + if score_col not in table.columns or q_col not in table.columns: + raise KeyError(f"Missing {score_col!r} / {q_col!r} for {method}") + parts.append( + pd.DataFrame( + { + "dataset": dataset, + "panel": panel, + "method": method, + "spectrum_id": spectrum_id, + "score": table[score_col].to_numpy(dtype=float), + "q_value": table[q_col].to_numpy(dtype=float), + "label": label, + } + ) + ) + return pd.concat(parts, ignore_index=True) + + +def _curves_df_from_novoboard( + df: pd.DataFrame, + *, + dataset: str, + panel: str, + label_col: str, +) -> pd.DataFrame: + """Long-form curve rows for NovoBoard PSM TDC targets.""" + if "spectrum_id" not in df.columns: + raise KeyError("NovoBoard curve export requires spectrum_id") + if label_col not in df.columns: + raise KeyError(f"Missing label column {label_col!r}") + return pd.DataFrame( + { + "dataset": dataset, + "panel": panel, + "method": NOVOBOARD_METHOD, + "spectrum_id": df["spectrum_id"].astype(str), + "score": df["ALC (%)"].to_numpy(dtype=float), + "q_value": df["estimated_q_value"].to_numpy(dtype=float), + "label": df[label_col].astype(bool).to_numpy(), + } + ) + + +def load_novoboard_target_decoy( + novoboard_dir: Path, split: Literal["unlabelled", "test"], decoy_rate: str +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Load NovoBoard target/decoy tables and attach twin ``_pair_key`` values.""" + prefix = "raw_unlabelled" if split == "unlabelled" else "annotated_test" + target_path = novoboard_dir / f"{prefix}.csv" + decoy_path = novoboard_dir / f"{prefix}_decoy_{decoy_rate}.csv" + if not target_path.is_file(): + raise FileNotFoundError(target_path) + if not decoy_path.is_file(): + raise FileNotFoundError(decoy_path) + target = pd.read_csv(target_path) + decoy = pd.read_csv(decoy_path) + return attach_novoboard_pair_keys( + target, decoy, novoboard_dir=novoboard_dir, split_prefix=prefix + ) + + +def _restrict_winnow_to_novoboard_spectra( + winnow: pd.DataFrame, novoboard: pd.DataFrame +) -> pd.DataFrame: + """Trim Winnow to NovoBoard twin-valid spectra under the subset invariant.""" + return restrict_winnow_to_novoboard_spectra(winnow, novoboard) + + +def _assert_shared_prediction_keys( + winnow: pd.DataFrame, + novoboard: pd.DataFrame, + *, + winnow_peptide_col: str = "prediction", + novoboard_peptide_col: str = "Peptide", +) -> None: + """Require I/L-normalised prediction identity on the shared spectrum pool.""" + assert_shared_prediction_keys( + winnow, + novoboard, + winnow_peptide_col=winnow_peptide_col, + novoboard_peptide_col=novoboard_peptide_col, + ) + + +def _label_series_by_spectrum_id(winnow: pd.DataFrame, label_col: str) -> pd.Series: + """Map ``spectrum_id`` → boolean label from a Winnow table.""" + return label_series_by_spectrum_id(winnow, label_col) + + +def _attach_labels_by_spectrum_id( + novoboard: pd.DataFrame, + label_by_id: pd.Series, + *, + label_col: str, +) -> pd.DataFrame: + """Attach a shared label column to NovoBoard rows by ``spectrum_id``.""" + return attach_labels_by_spectrum_id(novoboard, label_by_id, label_col=label_col) + + +def _method_curves_from_panel(panel_df: pd.DataFrame) -> list[MethodCurve]: + """Rebuild plot curves from long-form curve rows for one panel.""" + curves: list[MethodCurve] = [] + for method in _METHOD_ORDER: + sub = panel_df.loc[panel_df["method"] == method] + if sub.empty: + continue + colour = _METHOD_COLOURS.get(str(method), _PALETTE[0]) + curves.append( + MethodCurve( + str(method), + colour, + sub["score"].to_numpy(dtype=float), + sub["q_value"].to_numpy(dtype=float), + ) + ) + return curves + + +def _recovery_series_from_panel(panel_df: pd.DataFrame) -> list[MethodRecovery]: + """Rebuild labelled recovery series from long-form curve rows.""" + series: list[MethodRecovery] = [] + for method in _METHOD_ORDER: + sub = panel_df.loc[panel_df["method"] == method] + if sub.empty: + continue + colour = _METHOD_COLOURS.get(str(method), _PALETTE[0]) + series.append( + MethodRecovery( + str(method), + colour, + sub["q_value"].to_numpy(dtype=float), + sub["label"].astype(bool).to_numpy(), + ) + ) + return series + + +def plot_dataset_from_curves( + curves: pd.DataFrame, dataset_key: str, output_dir: Path +) -> None: + """Write PSM comparison plots for one dataset from a curves table.""" + out = output_dir / dataset_key + out.mkdir(parents=True, exist_ok=True) + ds = curves.loc[curves["dataset"] == dataset_key] + if ds.empty: + raise ValueError(f"No curve rows for dataset {dataset_key!r}") + + panel_specs: tuple[tuple[str, EvalType, str], ...] = ( + ("unlabelled", "unlabelled", "unlabelled"), + ("labelled_test", "labelled", "test"), + ) + for panel, eval_type, stem in panel_specs: + panel_df = ds.loc[ds["panel"] == panel] + method_curves = _method_curves_from_panel(panel_df) + if not method_curves: + continue + plot_qvalue_by_rank( + method_curves, + dataset_key, + eval_type, + out / f"psm_qvalue_by_rank_{stem}_{dataset_key}", + ) + plot_threshold_barplot( + _bar_series_from_curves(method_curves), + dataset_key, + eval_type, + out / f"psm_counts_{stem}_{dataset_key}", + ) + + labelled = ds.loc[ds["panel"] == "labelled_test"] + recovery = _recovery_series_from_panel(labelled) + if recovery: + plot_recovery_curves( + recovery, + dataset_key, + out / f"psm_recovery_test_{dataset_key}", + ) + + +def write_curves_csv(curves: pd.DataFrame, output_dir: Path) -> Path: + """Write the long-form replot curves table.""" + output_dir.mkdir(parents=True, exist_ok=True) + path = output_dir / CURVES_CSV_NAME + # Preserve float64 q/score values so threshold edge cases survive round-trip. + curves.to_csv(path, index=False, float_format="%.17g") + logger.info("Wrote %s (%d rows)", path, len(curves)) + return path + + +def plot_qvalue_by_rank( + curves: list[MethodCurve], + dataset_key: str, + eval_type: EvalType, + output_path: Path, + *, + title_suffix: str = "", +) -> None: + """Plot q-value against native-score rank/accepted count.""" + display = _display_name(dataset_key) + qualifier = _ground_truth_qualifier( + "labelled" if eval_type == "labelled" else "unlabelled" + ) + title = f"{display} PSM q-value by rank {qualifier}{title_suffix}" + + fig, ax = plt.subplots(figsize=(8, 6)) + q_max = 0.0 + for curve in curves: + order = np.argsort(-np.asarray(curve.confidence, dtype=float)) + y = np.asarray(curve.q_value, dtype=float)[order] + rank = np.arange(1, len(y) + 1) + valid = ~np.isnan(y) + if not np.any(valid): + continue + q_max = max(q_max, float(np.nanmax(y[valid]))) + ax.plot( + rank[valid], + y[valid], + color=curve.color, + lw=1.5, + label=curve.label, + ) + + ax.set_xlabel("Accepted PSMs by native-score rank") + ax.set_ylabel("PSM q-value") + ax.set_title(title) + y_top = min(max(q_max * 1.15, 0.05), 1.0) + ax.set_ylim(0, y_top) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s", output_path) + + +def _count_at_thresholds( + q: np.ndarray, thresholds: list[float] = FDR_THRESHOLDS +) -> list[int]: + q = np.asarray(q, dtype=float) + valid = q[~np.isnan(q)] + return [int((valid <= t).sum()) for t in thresholds] + + +def _bar_series_from_curves(curves: list[MethodCurve]) -> list[MethodCounts]: + return [ + MethodCounts( + label=c.label, + color=c.color, + counts=_count_at_thresholds(c.q_value), + ) + for c in curves + ] + + +def _recovery_at_thresholds( + q: np.ndarray, + correct: np.ndarray, + thresholds: list[float] = FDR_THRESHOLDS, +) -> list[float]: + """Return correct-identification recovery percentage at each q-value threshold.""" + q = np.asarray(q, dtype=float) + correct = np.asarray(correct, dtype=bool) + denom = int(correct.sum()) + if denom == 0: + return [np.nan for _ in thresholds] + valid = ~np.isnan(q) + return [100.0 * int((valid & correct & (q <= t)).sum()) / denom for t in thresholds] + + +def plot_recovery_curves( + series: list[MethodRecovery], + dataset_key: str, + output_path: Path, +) -> None: + """Plot correct-identification recovery versus q-value threshold.""" + display = _display_name(dataset_key) + fig, ax = plt.subplots(figsize=(8, 6)) + for item in series: + y = _recovery_at_thresholds(item.q_value, item.correct) + ax.plot( + FDR_THRESHOLDS, + y, + marker="o", + lw=1.5, + label=item.label, + color=item.color, + ) + + ax.set_xlim(0, max(FDR_THRESHOLDS)) + ax.set_ylim(0, 100) + ax.set_xlabel("Estimated q-value threshold") + ax.set_ylabel("Correct PSM recovery\n(% of labelled correct PSMs)") + ax.set_title(f"{display} labelled PSM recovery by q-value threshold") + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s", output_path) + + +def plot_threshold_barplot( + series: list[MethodCounts], + dataset_key: str, + eval_type: EvalType, + output_path: Path, +) -> None: + """Bar chart of identifications retained at each q-value threshold.""" + display = _display_name(dataset_key) + if eval_type == "labelled": + split_label = "labelled test set" + else: + split_label = "unlabelled set" + title = f"{display}: accepted PSMs on the {split_label} at q-value thresholds" + ylabel = "Peptide-spectrum matches" + + n_methods = len(series) + n_thresh = len(FDR_THRESHOLDS) + group_spacing = 0.825 + cluster_width = min(0.75, group_spacing * 0.92) + width = cluster_width / n_methods + x = np.arange(n_thresh) * group_spacing + + fig_w = max(10.0, 2.2 * n_thresh * group_spacing) + fig, ax = plt.subplots(figsize=(fig_w, 7)) + for i, item in enumerate(series): + offset = (i - (n_methods - 1) / 2) * width + bars = ax.bar( + x + offset, + item.counts, + width, + label=item.label, + color=item.color, + edgecolor="black", + linewidth=1, + ) + for bar in bars: + h = bar.get_height() + ax.annotate( + f"{int(h):,}", + xy=(bar.get_x() + bar.get_width() / 2, h), + xytext=(0, 3), + textcoords="offset points", + ha="center", + va="bottom", + fontsize=9, + ) + + max_count = max((c for item in series for c in item.counts), default=1) + y_headroom = (1.55 + 0.06 * n_methods) * (2 / 3) + ax.set_ylim(0, max_count * y_headroom) + + half_cluster = cluster_width / 2 + ax.set_xlim( + -half_cluster - 0.25, + (n_thresh - 1) * group_spacing + half_cluster + 0.25, + ) + + ax.set_xlabel("Q-value threshold") + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.set_xticks(x) + ax.set_xticklabels([str(t) for t in FDR_THRESHOLDS]) + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, output_path) + logger.info("Wrote %s", output_path) + + +def _append_method_summary_rows( + acceptance_rows: list[dict[str, object]], + error_rows: list[dict[str, object]], + *, + dataset: str, + panel: str, + level: str, + method: str, + q_value: np.ndarray, + label_mask: np.ndarray | None = None, + recovery_denom: int | None = None, + q_ref: np.ndarray | None = None, +) -> None: + """Append acceptance and error rows for one method.""" + acceptance_rows.extend( + acceptance_rows_from_q( + dataset=dataset, + panel=panel, + level=level, + method=method, + q_value=q_value, + thresholds=SUMMARY_THRESHOLDS, + label_mask=label_mask, + recovery_denom=recovery_denom, + ) + ) + error_rows.extend( + error_rows_from_q( + dataset=dataset, + panel=panel, + level=level, + method=method, + q_value=q_value, + thresholds=SUMMARY_THRESHOLDS, + label_mask=label_mask, + q_ref=q_ref, + ) + ) + + +def summary_rows_from_curves( + curves: pd.DataFrame, +) -> tuple[list[dict[str, object]], list[dict[str, object]]]: + """Build acceptance and error summary rows from a long-form curves table.""" + acceptance_rows: list[dict[str, object]] = [] + error_rows: list[dict[str, object]] = [] + required = {"dataset", "panel", "method", "spectrum_id", "q_value", "label"} + missing = required - set(curves.columns) + if missing: + raise ValueError(f"Curves table missing columns: {sorted(missing)}") + + for (dataset, panel), group in curves.groupby(["dataset", "panel"], sort=False): + cal = group.loc[group["method"] == DB_CAL_METHOD, ["spectrum_id", "q_value"]] + q_ref_cal_by_id = cal.drop_duplicates("spectrum_id").set_index("spectrum_id")[ + "q_value" + ] + raw = group.loc[group["method"] == DB_RAW_METHOD, ["spectrum_id", "q_value"]] + q_ref_raw_by_id = raw.drop_duplicates("spectrum_id").set_index("spectrum_id")[ + "q_value" + ] + for method, mdf in group.groupby("method", sort=False): + method_s = str(method) + q_value = mdf["q_value"].to_numpy(dtype=float) + label_mask = mdf["label"].astype(bool).to_numpy() + recovery_denom = int(label_mask.sum()) + q_ref: np.ndarray | None = None + # Winnow-family methods: deviation vs calibrated-confidence DBG. + # NovoBoard / Glissade: deviation vs raw-confidence DBG. + if method_s in _WINNOW_METHODS: + q_ref = ( + mdf["spectrum_id"] + .astype(str) + .map(q_ref_cal_by_id) + .to_numpy(dtype=float) + ) + elif method_s in (NOVOBOARD_METHOD, "Glissade"): + q_ref = ( + mdf["spectrum_id"] + .astype(str) + .map(q_ref_raw_by_id) + .to_numpy(dtype=float) + ) + _append_method_summary_rows( + acceptance_rows, + error_rows, + dataset=str(dataset), + panel=str(panel), + level="psm", + method=method_s, + q_value=q_value, + label_mask=label_mask, + recovery_denom=recovery_denom, + q_ref=q_ref, + ) + return acceptance_rows, error_rows + + +def _comparators_for_panel(panel: str, level: str) -> list[str]: + """Comparator method labels present in a given summary panel.""" + del panel, level + return ["NovoBoard"] + + +def _finalise_method_comparison_tables( + acceptance_rows: list[dict[str, object]], + error_rows: list[dict[str, object]], +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Build acceptance and error/gain DataFrames with per-panel relative columns.""" + acceptance = pd.DataFrame(acceptance_rows) + error = pd.DataFrame(error_rows) + if acceptance.empty: + return acceptance, error + + gain_parts: list[pd.DataFrame] = [] + group_cols = ["dataset", "panel", "level"] + for keys, acc_group in acceptance.groupby(group_cols, sort=False): + if not isinstance(keys, tuple): + keys = (keys,) + _, panel, level = keys + err_mask = True + for col, val in zip(group_cols, keys): + err_mask = err_mask & (error[col] == val) + err_group = error.loc[err_mask] + gain_parts.append( + finalise_error_gain_table( + acc_group, + err_group, + primary_method=PRIMARY_METHOD, + comparators=_comparators_for_panel(str(panel), str(level)), + ) + ) + error_gain = pd.concat(gain_parts, ignore_index=True) if gain_parts else error + return acceptance, error_gain + + +def process_dataset(cfg: DatasetConfig, plots_dir: Path) -> pd.DataFrame: + """Generate comparison plots and return long-form curve rows for one dataset.""" + residue_masses = _load_residue_masses() + + winnow_unlabelled = load_winnow(cfg.winnow_unlabelled, cfg.fasta, "unlabelled") + winnow_test = load_winnow(cfg.winnow_test, cfg.fasta, "labelled") + nb_u_target, nb_u_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "unlabelled", cfg.novoboard_decoy_rate + ) + nb_t_target, nb_t_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "test", cfg.novoboard_decoy_rate + ) + nb_u_target, nb_u_decoy = filter_novoboard_target_decoy_pairs( + nb_u_target, nb_u_decoy, min_length=MIN_PEPTIDE_LENGTH + ) + nb_t_target, nb_t_decoy = filter_novoboard_target_decoy_pairs( + nb_t_target, nb_t_decoy, min_length=LABELLED_MIN_PEPTIDE_LENGTH + ) + # Pair-gated NovoBoard ⊆ Winnow after identical InstaNovo filters; only trim Winnow. + winnow_unlabelled = _restrict_winnow_to_novoboard_spectra( + winnow_unlabelled, nb_u_target + ) + winnow_test = _restrict_winnow_to_novoboard_spectra(winnow_test, nb_t_target) + _assert_shared_prediction_keys(winnow_test, nb_t_target) + _assert_shared_prediction_keys(winnow_unlabelled, nb_u_target) + + # Shared labels once on Winnow; NovoBoard reuses them by spectrum_id. + winnow_test = winnow_test.copy() + winnow_test["correct"] = novor_correctness_mask( + winnow_test["sequence"], + winnow_test["prediction"], + residue_masses=residue_masses, + ) + if "proteome_hit" not in winnow_unlabelled.columns: + raise KeyError("Expected proteome_hit on unlabelled Winnow table") + correct_by_id = _label_series_by_spectrum_id(winnow_test, "correct") + hit_by_id = _label_series_by_spectrum_id(winnow_unlabelled, "proteome_hit") + + novoboard_unlabelled = novoboard_psm_tdc( + nb_u_target, nb_u_decoy, min_length=MIN_PEPTIDE_LENGTH + ) + n_u_tgt = int(novoboard_unlabelled["is_target"].sum()) + n_u_dec = int((~novoboard_unlabelled["is_target"]).sum()) + if n_u_tgt != n_u_dec: + raise AssertionError( + f"Unlabelled PSM TDC unbalanced: targets={n_u_tgt} decoys={n_u_dec}" + ) + novoboard_unlabelled = novoboard_unlabelled[ + novoboard_unlabelled["is_target"] + ].copy() + novoboard_test = novoboard_psm_tdc( + nb_t_target, nb_t_decoy, min_length=LABELLED_MIN_PEPTIDE_LENGTH + ) + n_t_tgt = int(novoboard_test["is_target"].sum()) + n_t_dec = int((~novoboard_test["is_target"]).sum()) + if n_t_tgt != n_t_dec: + raise AssertionError( + f"Labelled PSM TDC unbalanced: targets={n_t_tgt} decoys={n_t_dec}" + ) + novoboard_test = novoboard_test[novoboard_test["is_target"]].copy() + novoboard_test = _attach_labels_by_spectrum_id( + novoboard_test, correct_by_id, label_col="correct" + ) + novoboard_unlabelled = _attach_labels_by_spectrum_id( + novoboard_unlabelled, hit_by_id, label_col="proteome_hit" + ) + + n_w_correct = int(winnow_test["correct"].sum()) + n_nb_correct = int(novoboard_test["correct"].sum()) + if n_w_correct != n_nb_correct: + raise AssertionError( + f"Shared labelled correct counts disagree: Winnow={n_w_correct} " + f"NovoBoard={n_nb_correct}" + ) + n_w_hit = int(winnow_unlabelled["proteome_hit"].sum()) + n_nb_hit = int(novoboard_unlabelled["proteome_hit"].sum()) + if n_w_hit != n_nb_hit: + raise AssertionError( + f"Shared proteome-hit counts disagree: Winnow={n_w_hit} NovoBoard={n_nb_hit}" + ) + + logger.info( + "%s shared PSM pools: unlabelled=%d labelled=%d " + "(NovoBoard twin-valid; Winnow trimmed; shared labels correct=%d hits=%d)", + cfg.key, + len(winnow_unlabelled), + len(winnow_test), + n_w_correct, + n_w_hit, + ) + + winnow_u_psm_table = _prepare_winnow_psm_table( + winnow_unlabelled, "proteome_hit", residue_masses + ) + winnow_t_psm_table = _prepare_winnow_psm_table( + winnow_test, "correct", residue_masses + ) + curves = pd.concat( + [ + _curves_df_from_winnow_table( + winnow_t_psm_table, + dataset=cfg.key, + panel="labelled_test", + label_col="correct", + ), + _curves_df_from_novoboard( + novoboard_test, + dataset=cfg.key, + panel="labelled_test", + label_col="correct", + ), + _curves_df_from_winnow_table( + winnow_u_psm_table, + dataset=cfg.key, + panel="unlabelled", + label_col="proteome_hit", + ), + _curves_df_from_novoboard( + novoboard_unlabelled, + dataset=cfg.key, + panel="unlabelled", + label_col="proteome_hit", + ), + ], + ignore_index=True, + ) + plot_dataset_from_curves(curves, cfg.key, plots_dir) + return curves + + +@app.command() +def main( + novoboard_root: Annotated[ + Optional[Path], + typer.Option( + "--novoboard-root", + help=( + "Root of NovoBoard per-dataset tables: " + "{root}/{dataset}/novoboard/ with annotated_test*.csv and " + "raw_unlabelled*.csv target/decoy pairs (the datasets/ dir of " + "a NovoBoard checkout). Required unless --summarise-only is set. " + "Local runs used fork JemmaLDaniel/NovoBoard, branch " + "feat/adapt-to-instanovo " + "(commit a9faab3ef1af06987599c2f01e6ba96072c80172)." + ), + ), + ] = None, + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for curves/summary CSVs."), + ] = DEFAULT_OUTPUT_DIR, + plots_dir: Annotated[ + Path, + typer.Option("--plots-dir", help="Directory for png/pdf figures."), + ] = DEFAULT_OUTPUT_DIR, + datasets: Annotated[ + Optional[list[str]], + typer.Option("--datasets", help="Dataset keys to plot."), + ] = None, + winnow_results: Annotated[ + Path, + typer.Option("--winnow-results", help="Winnow results directory."), + ] = DEFAULT_WINNOW_RESULTS, + summarise_only: Annotated[ + Optional[Path], + typer.Option( + "--summarise-only", + help="Only write plots and summary CSVs from an existing curves CSV.", + ), + ] = None, +) -> None: + """Generate FDR method comparison plots and summary CSVs.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + + if summarise_only is not None: + curves = pd.read_csv(summarise_only, float_precision="round_trip") + if "label" in curves.columns: + curves["label"] = curves["label"].astype(bool) + dataset_keys = ( + datasets + if datasets is not None + else sorted(curves["dataset"].astype(str).unique()) + ) + for key in dataset_keys: + logger.info("Replotting %s from curves CSV", key) + plot_dataset_from_curves(curves, str(key), plots_dir) + curves_out = curves.loc[curves["dataset"].astype(str).isin(dataset_keys)] + acceptance_rows, error_rows = summary_rows_from_curves(curves_out) + acceptance, error_gain = _finalise_method_comparison_tables( + acceptance_rows, error_rows + ) + write_summary_tables( + acceptance, error_gain, results_dir, "fdr_method_comparison" + ) + if summarise_only.resolve() != (results_dir / CURVES_CSV_NAME).resolve(): + write_curves_csv(curves_out, results_dir) + return + + if novoboard_root is None: + raise typer.BadParameter( + "--novoboard-root is required unless --summarise-only is set." + ) + + dataset_keys = datasets if datasets is not None else list(DEFAULT_DATASETS) + configs = build_dataset_configs(winnow_results, novoboard_root=novoboard_root) + + curve_parts: list[pd.DataFrame] = [] + for key in dataset_keys: + if key not in configs: + raise typer.BadParameter(f"Unknown dataset {key!r}") + logger.info("Processing %s", key) + curve_parts.append(process_dataset(configs[key], plots_dir)) + + curves = pd.concat(curve_parts, ignore_index=True) + write_curves_csv(curves, results_dir) + acceptance_rows, error_rows = summary_rows_from_curves(curves) + acceptance, error_gain = _finalise_method_comparison_tables( + acceptance_rows, error_rows + ) + write_summary_tables(acceptance, error_gain, results_dir, "fdr_method_comparison") + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/plot_feature_investigation.py b/paper_scripts/plot_feature_investigation.py new file mode 100644 index 00000000..0bb8cea9 --- /dev/null +++ b/paper_scripts/plot_feature_investigation.py @@ -0,0 +1,1349 @@ +"""Generate feature investigation plots from calibrator training feature matrices. + +Produces KDE, scatter, violin, correlation, discriminative-power, pairplot, +mirror-spectrum, retention-time, token-stem, and beam-stem figures matching the +style of ``analysis/feature_investigation_new.ipynb``. +""" + +from __future__ import annotations + +import ast +import logging +import warnings +from pathlib import Path +from typing import Annotated, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import typer +from matplotlib.colors import LinearSegmentedColormap +from matplotlib.patches import Patch +from scipy import stats as sp_stats +from scipy.stats import gaussian_kde +from sklearn.decomposition import PCA +from sklearn.metrics import roc_auc_score +from sklearn.preprocessing import StandardScaler + +logger = logging.getLogger(__name__) +app = typer.Typer( + add_completion=False, pretty_exceptions_show_locals=False, no_args_is_help=True +) + +# --------------------------------------------------------------------------- +# Style — Paul Tol "bright" palette (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = [ + "#4477AA", + "#EE6677", + "#228833", + "#CCBB44", + "#66CCEE", + "#AA3377", + "#BBBBBB", +] +_CORRECT_COLOUR = _PALETTE[0] +_INCORRECT_COLOUR = _PALETTE[1] +_NEUTRAL_COLOUR = _PALETTE[6] + +_HIGH_CONF_BEAM_COLOUR = _PALETTE[2] # green +_LOW_CONF_BEAM_COLOUR = _PALETTE[5] # purple +_MED_CONF_BEAM_COLOUR = _PALETTE[3] # yellow + +_OBS_COLOUR = _PALETTE[3] # yellow (observed spectrum) +_THEO_COLOUR = _PALETTE[5] # purple (predicted spectrum) + +HUE_LABEL_CORRECT = "Correct" +HUE_LABEL_INCORRECT = "Incorrect" +HUE_ORDER = [HUE_LABEL_CORRECT, HUE_LABEL_INCORRECT] +HUE_PALETTE = { + HUE_LABEL_CORRECT: _CORRECT_COLOUR, + HUE_LABEL_INCORRECT: _INCORRECT_COLOUR, +} + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +_SUNSET_COLORS = [ + "#364B9A", + "#4A7BB7", + "#6EA6CD", + "#98CAE1", + "#C2E4EF", + "#EAECCC", + "#FEDA8B", + "#FDB366", + "#F67E4B", + "#DD3D2D", + "#A50026", +] + + +def _diverging_cmap() -> LinearSegmentedColormap: + cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) + cmap.set_bad(color="#FFFFFF") + return cmap + + +# --------------------------------------------------------------------------- +# Feature definitions +# --------------------------------------------------------------------------- +FEATURE_COLUMNS = [ + "confidence", + "mass_error_ppm", + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", + "irt_error", + "margin", + "median_margin", + "entropy", + "z-score", + "edit_distance", + "min_token_probability", + "std_token_probability", +] + +FRAGMENT_FEATURES = [ + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", +] + +BEAM_FEATURES = ["margin", "median_margin", "entropy", "z-score", "edit_distance"] + +TOKEN_FEATURES = ["min_token_probability", "std_token_probability"] + +SKEWED_FEATURES = {"irt_error"} + +_NICE_LABELS: dict[str, str] = { + "ion_matches": "Ion match rate", + "ion_match_intensity": "Ion match intensity", + "complementary_ion_count": "Complementary ion count", + "max_ion_gap": "Max ion gap", + "spectral_angle": "Spectral angle", + "xcorr": "Cross-correlation (XCorr)", + "mass_error_ppm": "Precursor mass error", + "log_abs_mass_error_ppm": "Log-absolute mass error", + "log_abs_mass_error_da": "Log-absolute mass error", + "mass_error_da": "Precursor mass error", + "irt_error": "iRT prediction error", + "confidence": "Model confidence", + "margin": "Beam margin", + "median_margin": "Beam median margin", + "entropy": "Beam entropy", + "z-score": "Beam z-score", + "edit_distance": "Runner-up edit distance", + "min_token_probability": "Min. token probability", + "std_token_probability": "Std. token probability", + "predicted_irt": "Regressor-predicted iRT", + "irt": "Koina-predicted iRT", + "retention_time": "Retention time (s)", +} + + +def _nice_label(col: str) -> str: + return _NICE_LABELS.get(col, col.replace("_", " ").capitalize()) + + +# --------------------------------------------------------------------------- +# Plotting helpers +# --------------------------------------------------------------------------- +def _save_fig(fig: plt.Figure, name: str, output_dir: Path) -> None: + base = output_dir / name + fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) + fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _style_ax(ax: plt.Axes) -> None: + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _auto_ylim(df: pd.DataFrame, feature: str): + if feature in SKEWED_FEATURES: + q99 = df[feature].quantile(0.99) + q01 = df[feature].quantile(0.01) + margin = (q99 - q01) * 0.1 + return (q01 - margin, q99 + margin) + return None + + +def plot_feature_vs_confidence( + df: pd.DataFrame, + feature: str, + title: str | None = None, + ylim: tuple[float, float] | None = None, +) -> tuple[plt.Figure, plt.Axes]: + """Scatter plot of a feature against model confidence, coloured by class.""" + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label] + ax.scatter( + subset["confidence"], + subset[feature], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("confidence")) + ax.set_ylabel(_nice_label(feature)) + if ylim is not None: + ax.set_ylim(ylim) + if title: + ax.set_title(title) + ax.legend() + _style_ax(ax) + fig.tight_layout() + return fig, ax + + +def _plot_peak_normalised_kde( + ax: plt.Axes, + subset: pd.Series, + colour: str, + label: str, + fill: bool, + clip: tuple[float, float] | None, +) -> None: + """Plot a peak-normalised KDE curve on *ax*.""" + kde = gaussian_kde(subset) + lo = subset.min() if clip is None else clip[0] + hi = subset.max() if clip is None else clip[1] + xs = np.linspace(lo, hi, 500) + ys = kde(xs) + ys /= ys.max() + ax.plot(xs, ys, color=colour, label=label, linewidth=1.5) + if fill: + ax.fill_between(xs, ys, alpha=0.3, color=colour) + + +def plot_kde_by_class( + df: pd.DataFrame, + feature: str, + title: str | None = None, + fill: bool = True, + clip: tuple[float, float] | None = None, + peak_normalise: bool = False, +) -> tuple[plt.Figure, plt.Axes]: + """KDE density plot of a feature split by correct/incorrect class.""" + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label][feature].dropna() + if len(subset) < 2: + continue + if peak_normalise: + _plot_peak_normalised_kde(ax, subset, HUE_PALETTE[label], label, fill, clip) + else: + kw: dict = {} + if clip is not None: + kw["clip"] = clip + sns.kdeplot( + subset, + ax=ax, + color=HUE_PALETTE[label], + label=label, + fill=fill, + alpha=0.3, + linewidth=1.5, + **kw, + ) + ax.set_xlabel(_nice_label(feature)) + ax.set_ylabel("Peak-normalised density" if peak_normalise else "Density") + if title: + ax.set_title(title) + ax.legend(loc="upper center") + _style_ax(ax) + fig.tight_layout() + return fig, ax + + +def plot_mirror_spectrum( + obs_mz, + obs_int, + theo_mz, + theo_int, + annotations, + title: str, + ax: plt.Axes | None = None, +) -> tuple[plt.Figure, plt.Axes]: + """Mirror plot comparing observed vs predicted spectra.""" + own_fig = ax is None + if own_fig: + fig, ax = plt.subplots(figsize=(10, 4)) + else: + assert ax is not None + fig = ax.get_figure() + + obs_int_norm = np.array(obs_int) / max(obs_int) * 100 + theo_int_norm = np.array(theo_int) / max(theo_int) * 100 + + ax.vlines(obs_mz, 0, obs_int_norm, color=_OBS_COLOUR, linewidth=1.8) + ax.vlines(theo_mz, 0, -theo_int_norm, color=_THEO_COLOUR, linewidth=1.8) + + if annotations is not None: + for mz_val, intensity_val, ann in zip(theo_mz, theo_int_norm, annotations): + if intensity_val > 10: + label_text = ann.decode() if isinstance(ann, bytes) else str(ann) + ax.annotate( + label_text, + (mz_val, -intensity_val), + fontsize=8, + ha="center", + va="top", + rotation=90, + color=_THEO_COLOUR, + ) + + ax.axhline(0, color="black", linewidth=0.5) + ax.set_xlabel("m/z") + ax.set_ylabel("Relative intensity (%)") + ax.set_title(title) + + ax.text( + 0.99, + 0.95, + "Observed", + transform=ax.transAxes, + ha="right", + va="top", + fontsize=9, + color=_OBS_COLOUR, + fontweight="bold", + ) + ax.text( + 0.99, + 0.05, + "Predicted", + transform=ax.transAxes, + ha="right", + va="bottom", + fontsize=9, + color=_THEO_COLOUR, + fontweight="bold", + ) + + _style_ax(ax) + if own_fig: + fig.tight_layout() + return fig, ax + + +def compute_discriminative_stats(df: pd.DataFrame, features: list[str]) -> pd.DataFrame: + """Compute AUROC, KS statistic, and Cohen's d for each feature.""" + results = [] + labels = df["correct"].astype(int) + for feat in features: + vals = df[feat].dropna() + valid_mask = df[feat].notna() + valid_labels = labels[valid_mask] + valid_vals = vals + if len(valid_vals) < 10 or valid_labels.nunique() < 2: + results.append( + { + "feature": feat, + "auroc": np.nan, + "ks_stat": np.nan, + "cohens_d": np.nan, + } + ) + continue + try: + auroc = roc_auc_score(valid_labels, valid_vals) + auroc = max(auroc, 1 - auroc) + except ValueError: + auroc = np.nan + correct_vals = valid_vals[valid_labels == 1] + incorrect_vals = valid_vals[valid_labels == 0] + ks_stat, _ = sp_stats.ks_2samp(correct_vals, incorrect_vals) + pooled_std = np.sqrt( + ( + (len(correct_vals) - 1) * correct_vals.std() ** 2 + + (len(incorrect_vals) - 1) * incorrect_vals.std() ** 2 + ) + / (len(correct_vals) + len(incorrect_vals) - 2) + ) + cohens_d = ( + abs(correct_vals.mean() - incorrect_vals.mean()) / pooled_std + if pooled_std > 0 + else np.nan + ) + results.append( + {"feature": feat, "auroc": auroc, "ks_stat": ks_stat, "cohens_d": cohens_d} + ) + return ( + pd.DataFrame(results) + .sort_values("auroc", ascending=False) + .reset_index(drop=True) + ) + + +# --------------------------------------------------------------------------- +# Token / beam stem helpers +# --------------------------------------------------------------------------- +_STEM_FIGSIZE = (10, 4.5) + + +def _parse_token_probs(row): + """Extract token probabilities and residue labels from a row.""" + try: + token_probs = np.exp(np.array(ast.literal_eval(row["token_log_probs"]))) + except (ValueError, SyntaxError): + return np.array([]), [] + seq_str = row["prediction"] + residues: list[str] = [] + j = 0 + while j < len(seq_str): + if j + 1 < len(seq_str) and seq_str[j + 1] == "[": + end = seq_str.index("]", j + 1) + 1 + residues.append(seq_str[j:end]) + j = end + else: + residues.append(seq_str[j]) + j += 1 + n_tokens = min(len(token_probs), len(residues)) + return token_probs[:n_tokens], residues[:n_tokens] + + +def _plot_token_stem(row, beam_colour: str, title_suffix: str = ""): + """Stem plot of per-residue token probabilities for one PSM.""" + token_probs, residues = _parse_token_probs(row) + if len(token_probs) == 0: + return None + + n = len(token_probs) + fig, ax = plt.subplots(figsize=_STEM_FIGSIZE) + markerline, stemlines, baseline = ax.stem( + range(n), + token_probs, + linefmt="-", + markerfmt="o", + basefmt="k-", + ) + plt.setp(stemlines, color=beam_colour, linewidth=2.5) + plt.setp(markerline, color=beam_colour, markersize=7, zorder=5) + + ax.set_xticks(range(n)) + ax.set_xticklabels( + residues, + fontsize=11, + rotation=45, + ha="right", + rotation_mode="anchor", + ) + ax.set_xlim(-0.5, n - 0.5) + ax.set_ylim(-0.03, 1.05) + ax.set_ylabel("Token probability") + ax.set_xlabel("Residue") + + charge = int(row["precursor_charge"]) if "precursor_charge" in row.index else "?" + ax.set_title( + f"Token probabilities for {row['prediction']}, " + f"+{charge}, confidence={row['confidence']:.3f}{title_suffix}", + ) + _style_ax(ax) + fig.tight_layout() + return fig + + +def _infer_charge(row) -> str: + """Best-effort charge extraction from a beam CSV row.""" + for col in ("precursor_charge", "charge"): + if col in row.index and pd.notna(row[col]): + return str(int(row[col])) + return "?" + + +def _plot_beam_stem(row, beam_log_prob_cols, beam_seq_cols, colour: str): + """Stem plot of per-beam confidence for one spectrum.""" + probs: list[float] = [] + labels: list[str] = [] + for i, (lp_col, seq_col) in enumerate(zip(beam_log_prob_cols, beam_seq_cols)): + lp = row[lp_col] + seq = row[seq_col] + if pd.isna(lp) or np.isinf(lp): + continue + probs.append(np.exp(float(lp))) + label = str(seq) if pd.notna(seq) else f"beam {i}" + labels.append(label) + + if len(probs) < 2: + return None + + n = len(probs) + fig, ax = plt.subplots(figsize=_STEM_FIGSIZE) + markerline, stemlines, baseline = ax.stem( + range(n), + probs, + linefmt="-", + markerfmt="o", + basefmt="k-", + ) + plt.setp(stemlines, color=colour, linewidth=2.5) + plt.setp(markerline, color=colour, markersize=7, zorder=5) + + ax.set_xticks(range(n)) + ax.set_xlim(-0.5, n - 0.5) + ax.set_ylim(-max(probs) * 0.03, max(probs) * 1.15) + ax.set_ylabel("Beam confidence") + ax.set_xlabel("Beam prediction index") + ax.set_title(f"Beam confidence for {labels[0]}, +{_infer_charge(row)}") + _style_ax(ax) + fig.tight_layout() + return fig + + +# --------------------------------------------------------------------------- +# Section generators — each mirrors a notebook section +# --------------------------------------------------------------------------- +def plot_confidence(df: pd.DataFrame, output_dir: Path) -> None: + """Section 1: confidence distribution.""" + fig, ax = plot_kde_by_class(df, "confidence", title="Model confidence distribution") + _save_fig(fig, "01a_confidence_kde", output_dir) + + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label] + ax.hist( + subset["confidence"], + bins=50, + alpha=0.5, + color=HUE_PALETTE[label], + edgecolor="black", + label=label, + density=True, + ) + ax.set_xlabel(_nice_label("confidence")) + ax.set_ylabel("Density") + ax.set_title("Model confidence histogram") + ax.legend(loc="upper center") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "01b_confidence_histogram", output_dir) + + +def _mass_error_log_column(df: pd.DataFrame) -> tuple[str, str]: + """Return (raw mass error column, log-absolute column) for plotting.""" + if "mass_error_da" in df.columns: + return "mass_error_da", "log_abs_mass_error_da" + if "mass_error_ppm" in df.columns: + return "mass_error_ppm", "log_abs_mass_error_ppm" + raise KeyError( + "Feature matrix must contain 'mass_error_da' or 'mass_error_ppm' for mass error plots" + ) + + +def plot_mass_error(df: pd.DataFrame, output_dir: Path) -> None: + """Section 2: mass error vs confidence (Da or ppm).""" + raw_col, log_col = _mass_error_log_column(df) + work = df.copy() + work[log_col] = np.log(work[raw_col].abs().clip(lower=1e-12)) + + fig, _ = plot_kde_by_class( + work, + raw_col, + title=f"{_nice_label(raw_col)} distribution", + ) + _save_fig(fig, f"02a_{raw_col}_kde", output_dir) + + fig, _ = plot_feature_vs_confidence( + work, + raw_col, + title=f"{_nice_label(raw_col)} vs model confidence", + ) + _save_fig(fig, f"02b_{raw_col}_vs_confidence", output_dir) + + fig, _ = plot_kde_by_class( + work, + log_col, + title="Log-absolute precursor mass error distribution", + ) + _save_fig(fig, "02c_mass_error_log_kde", output_dir) + + fig, _ = plot_feature_vs_confidence( + work, + log_col, + title="Log-absolute precursor mass error vs model confidence", + ) + _save_fig(fig, "02d_mass_error_log_vs_confidence", output_dir) + + if raw_col == "mass_error_da": + da_ylim = (-0.2, 0.2) + fig, _ = plot_kde_by_class( + work, + raw_col, + title=f"{_nice_label(raw_col)} distribution", + clip=da_ylim, + ) + _save_fig(fig, "02e_mass_error_da_kde_within_0.2da", output_dir) + + fig, _ = plot_feature_vs_confidence( + work, + raw_col, + title=f"{_nice_label(raw_col)} vs model confidence", + ylim=da_ylim, + ) + _save_fig(fig, "02f_mass_error_da_vs_confidence_within_0.2da", output_dir) + + +def plot_mirror_spectra(df_meta: pd.DataFrame, output_dir: Path) -> None: + """Section 3: mirror plots of observed vs predicted spectra.""" + required = { + "theoretical_mz", + "mz_array", + "intensity_array", + "theoretical_intensity", + } + if not required.issubset(df_meta.columns): + print(" Skipping mirror plots — missing spectrum columns in metadata.") + return + + valid_mirror = df_meta[ + df_meta["theoretical_mz"].apply(lambda x: x is not None and len(x) > 0) + & df_meta["mz_array"].apply(lambda x: x is not None and len(x) > 0) + ].copy() + + if len(valid_mirror) == 0: + print(" Skipping mirror plots — no rows with valid spectrum arrays.") + return + + has_annotations = "theoretical_annotation" in valid_mirror.columns + + def _add_mirror_margin(ax, y_frac=0.11): + ymin, ymax = ax.get_ylim() + y_pad = (ymax - ymin) * y_frac + ax.set_ylim(ymin - y_pad, ymax + y_pad) + + def _mirror_title(row) -> str: + pred = row.get("prediction", "?") + charge = ( + int(row["precursor_charge"]) if "precursor_charge" in row.index else "?" + ) + return f"Observed vs predicted spectrum for {pred}, +{charge}" + + correct_high = valid_mirror[valid_mirror["correct"]].nlargest(3, "confidence") + incorrect_low = valid_mirror[~valid_mirror["correct"]].nsmallest(3, "confidence") + + conf_middle_lo, conf_middle_hi = 0.45, 0.55 + middle_mask = valid_mirror["confidence"].between(conf_middle_lo, conf_middle_hi) + n_correct_mid = (valid_mirror["correct"] & middle_mask).sum() + n_incorrect_mid = (~valid_mirror["correct"] & middle_mask).sum() + correct_middle = valid_mirror[valid_mirror["correct"] & middle_mask].sample( + n=min(3, n_correct_mid), random_state=42 + ) + incorrect_middle = valid_mirror[~valid_mirror["correct"] & middle_mask].sample( + n=min(3, n_incorrect_mid), random_state=42 + ) + + groups = [ + (correct_high, "correct", "03a_mirror_high_conf"), + (incorrect_low, "incorrect", "03b_mirror_low_conf"), + (correct_middle, "correct", "03c_mirror_middle_conf_correct"), + (incorrect_middle, "incorrect", "03d_mirror_middle_conf_incorrect"), + ] + + for subset, _status, prefix in groups: + for i, (_, row) in enumerate(subset.iterrows()): + fig, ax = plt.subplots(figsize=(8, 5)) + annotations = row.get("theoretical_annotation") if has_annotations else None + plot_mirror_spectrum( + obs_mz=row["mz_array"], + obs_int=row["intensity_array"], + theo_mz=row["theoretical_mz"], + theo_int=row["theoretical_intensity"], + annotations=annotations, + title=_mirror_title(row), + ax=ax, + ) + _add_mirror_margin(ax) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, f"{prefix}_{i}", output_dir) + + +def plot_fragment_features(df: pd.DataFrame, output_dir: Path) -> None: + """Section 4: fragment ion match features vs confidence.""" + available = [f for f in FRAGMENT_FEATURES if f in df.columns] + for feat in available: + fig, _ = plot_feature_vs_confidence( + df, + feat, + title=f"{_nice_label(feat)} vs model confidence", + ylim=_auto_ylim(df, feat), + ) + _save_fig(fig, f"04a_fragment_{feat}_vs_confidence", output_dir) + + fig, _ = plot_kde_by_class(df, feat, title=f"{_nice_label(feat)} distribution") + _save_fig(fig, f"04b_fragment_{feat}_kde", output_dir) + + +def plot_irt(df: pd.DataFrame, df_meta: pd.DataFrame | None, output_dir: Path) -> None: + """Section 7: iRT error plots + RT scatter when metadata is available.""" + if df_meta is not None: + has_rt = ( + "retention_time" in df_meta.columns and "predicted_irt" in df_meta.columns + ) + has_koina_irt = "irt" in df_meta.columns + + if has_rt and has_koina_irt: + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df_meta[df_meta["hue"] == label] + ax.scatter( + subset["retention_time"], + subset["irt"], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("retention_time")) + ax.set_ylabel(_nice_label("irt")) + ax.set_title("Retention time vs Koina-predicted iRT") + ax.legend(markerscale=3, frameon=True) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "07a_rt_vs_koina_irt", output_dir) + + if has_koina_irt and has_rt: + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df_meta[df_meta["hue"] == label] + ax.scatter( + subset["predicted_irt"], + subset["irt"], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("predicted_irt")) + ax.set_ylabel(_nice_label("irt")) + ax.set_title("Koina-predicted iRT vs regressor-predicted iRT") + ax.legend(markerscale=3, frameon=True) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "07b_predicted_vs_koina_irt", output_dir) + + if "irt_error" not in df.columns: + return + + irt_ylim = _auto_ylim(df, "irt_error") + fig, _ = plot_feature_vs_confidence( + df, "irt_error", title="iRT prediction error vs model confidence", ylim=irt_ylim + ) + _save_fig(fig, "07c_irt_error_vs_confidence", output_dir) + + fig, _ = plot_kde_by_class( + df, + "irt_error", + title="iRT prediction error distribution", + clip=(0, df["irt_error"].quantile(0.99)), + ) + _save_fig(fig, "07d_irt_error_kde", output_dir) + + +def plot_token_stems(df_meta: pd.DataFrame, output_dir: Path) -> None: + """Section 8: token-level probability stem plots from metadata.""" + if "token_log_probs" not in df_meta.columns: + print(" Skipping token stem plots — no token_log_probs column in metadata.") + return + if "prediction" not in df_meta.columns: + print(" Skipping token stem plots — no prediction column in metadata.") + return + + # High confidence + high_conf_pool = df_meta[df_meta["confidence"] >= 0.9] + high_samples = ( + high_conf_pool.sample(3, random_state=42) + if len(high_conf_pool) >= 3 + else high_conf_pool + ) + for i, (_, row) in enumerate(high_samples.iterrows()): + fig = _plot_token_stem(row, _HIGH_CONF_BEAM_COLOUR) + if fig is not None: + _save_fig(fig, f"08a_token_stem_high_conf_{i}", output_dir) + + # Medium confidence + med_conf_pool = df_meta[ + (df_meta["confidence"] >= 0.4) & (df_meta["confidence"] <= 0.7) + ] + med_samples = ( + med_conf_pool.sample(3, random_state=42) + if len(med_conf_pool) >= 3 + else med_conf_pool + ) + for i, (_, row) in enumerate(med_samples.iterrows()): + fig = _plot_token_stem(row, _MED_CONF_BEAM_COLOUR) + if fig is not None: + _save_fig(fig, f"08b_token_stem_med_conf_{i}", output_dir) + + # Low confidence + low_conf_pool = df_meta[df_meta["confidence"] <= 0.2] + low_samples = ( + low_conf_pool.sample(3, random_state=42) + if len(low_conf_pool) >= 3 + else low_conf_pool + ) + for i, (_, row) in enumerate(low_samples.iterrows()): + fig = _plot_token_stem(row, _LOW_CONF_BEAM_COLOUR) + if fig is not None: + _save_fig(fig, f"08c_token_stem_low_conf_{i}", output_dir) + + +def plot_beam_stems(predictions_csv: Path, output_dir: Path) -> None: + """Section 8b: beam confidence stem plots from the predictions CSV.""" + beam_csv = pd.read_csv(predictions_csv) + + beam_log_prob_cols = sorted( + [ + c + for c in beam_csv.columns + if c.startswith("predictions_log_probability_beam_") + ], + key=lambda c: int(c.rsplit("_", 1)[1]), + ) + beam_seq_cols = sorted( + [ + c + for c in beam_csv.columns + if c.startswith("predictions_beam_") + and "log_probability" not in c + and "token" not in c + ], + key=lambda c: int(c.rsplit("_", 1)[1]), + ) + + if not beam_log_prob_cols or not beam_seq_cols: + print(" Skipping beam stem plots — no beam columns in predictions CSV.") + return + + beam_csv["top_confidence"] = np.exp(beam_csv[beam_log_prob_cols[0]].astype(float)) + + # Filter rows where all beams are -inf or NaN + valid_beams = beam_csv.dropna(subset=beam_log_prob_cols, how="all").copy() + for col in beam_log_prob_cols: + valid_beams[col] = pd.to_numeric(valid_beams[col], errors="coerce") + valid_beams = valid_beams[ + valid_beams[beam_log_prob_cols].apply( + lambda row: not all(np.isinf(row) | row.isna()), axis=1 + ) + ] + valid_beams = valid_beams[ + valid_beams[beam_log_prob_cols].apply( + lambda row: any(np.exp(row.dropna()) > 1e-15), axis=1 + ) + ] + + if len(valid_beams) == 0: + print(" Skipping beam stem plots — no valid beam rows after filtering.") + return + + # High confidence beams + high_beam = valid_beams[valid_beams["top_confidence"] >= 0.9] + high_beam_samples = ( + high_beam.sample(3, random_state=42) if len(high_beam) >= 3 else high_beam + ) + for i, (_, row) in enumerate(high_beam_samples.iterrows()): + fig = _plot_beam_stem( + row, beam_log_prob_cols, beam_seq_cols, _HIGH_CONF_BEAM_COLOUR + ) + if fig is not None: + _save_fig(fig, f"08d_beam_conf_high_{i}", output_dir) + + # Low confidence beams + low_beam = valid_beams[valid_beams["top_confidence"] <= 0.2] + low_beam_samples = ( + low_beam.sample(3, random_state=42) if len(low_beam) >= 3 else low_beam + ) + for i, (_, row) in enumerate(low_beam_samples.iterrows()): + fig = _plot_beam_stem( + row, beam_log_prob_cols, beam_seq_cols, _LOW_CONF_BEAM_COLOUR + ) + if fig is not None: + _save_fig(fig, f"08e_beam_conf_low_{i}", output_dir) + + +def plot_beam_features(df: pd.DataFrame, output_dir: Path) -> None: + """Section 9: beam search features vs confidence.""" + available = [f for f in BEAM_FEATURES if f in df.columns] + for feat in available: + fig, _ = plot_feature_vs_confidence( + df, feat, title=f"{_nice_label(feat)} vs model confidence" + ) + _save_fig(fig, f"09a_beam_{feat}_scatter", output_dir) + + fig, _ = plot_kde_by_class(df, feat, title=f"{_nice_label(feat)} distribution") + _save_fig(fig, f"09b_beam_{feat}_kde", output_dir) + + +def plot_token_features(df: pd.DataFrame, output_dir: Path) -> None: + """Section 11: token-level features.""" + if "min_token_probability" not in df.columns: + return + + fig, _ = plot_kde_by_class( + df, "min_token_probability", title="Min. token probability distribution" + ) + _save_fig(fig, "11a_min_token_prob_kde", output_dir) + + fig, _ = plot_kde_by_class( + df, "std_token_probability", title="Std. token probability distribution" + ) + _save_fig(fig, "11b_std_token_prob_kde", output_dir) + + fig, _ = plot_feature_vs_confidence( + df, "min_token_probability", title="Min. token probability vs confidence" + ) + _save_fig(fig, "11c_min_token_prob_scatter", output_dir) + + fig, _ = plot_feature_vs_confidence( + df, "std_token_probability", title="Std. token probability vs confidence" + ) + _save_fig(fig, "11d_std_token_prob_scatter", output_dir) + + fig, ax = plt.subplots(figsize=(8, 6)) + for label in HUE_ORDER: + subset = df[df["hue"] == label] + ax.scatter( + subset["min_token_probability"], + subset["std_token_probability"], + c=HUE_PALETTE[label], + label=label, + alpha=0.3, + s=10, + rasterized=True, + ) + ax.set_xlabel(_nice_label("min_token_probability")) + ax.set_ylabel(_nice_label("std_token_probability")) + ax.set_title("Token-level feature space") + ax.legend(markerscale=3, frameon=True) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "11e_token_feature_2d", output_dir) + + +def _plot_pca(df: pd.DataFrame, available: list[str], output_dir: Path) -> None: + """12f/12g — PCA scatter and loadings for calibrator features.""" + feat_df = df[available].dropna() + if len(feat_df) < 10: + return + + hue_pca = df.loc[feat_df.index, "hue"].values + scaler = StandardScaler() + x_scaled = scaler.fit_transform(feat_df.values) + pca = PCA(n_components=2) + z_pca = pca.fit_transform(x_scaled) + + fig, ax = plt.subplots(figsize=(8, 7)) + for label, colour in zip(reversed(HUE_ORDER), [_INCORRECT_COLOUR, _CORRECT_COLOUR]): + mask = hue_pca == label + ax.scatter( + z_pca[mask, 0], + z_pca[mask, 1], + c=colour, + label=label, + s=10, + alpha=0.3, + rasterized=True, + ) + ax.set_xlabel(f"PC 1 ({pca.explained_variance_ratio_[0]:.1%} variance)") + ax.set_ylabel(f"PC 2 ({pca.explained_variance_ratio_[1]:.1%} variance)") + ax.set_title("PCA of calibrator features") + ax.legend(loc="upper left") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "12f_pca_features", output_dir) + + pc1 = pca.components_[0] + pc2 = pca.components_[1] + names = [_nice_label(c) for c in available] + order = np.argsort(np.abs(pc1))[::-1] + + y = np.arange(len(names)) + fig, ax = plt.subplots(figsize=(10, 7)) + ax.barh( + y, + pc1[order], + color=_CORRECT_COLOUR, + alpha=0.6, + edgecolor="black", + linewidth=0.4, + ) + ax.barh( + y, pc2[order], color=_PALETTE[5], alpha=0.4, edgecolor="black", linewidth=0.4 + ) + ax.set_yticks(y) + ax.set_yticklabels([names[i] for i in order]) + ax.invert_yaxis() + ax.set_xlabel("Loading value") + ax.set_title("PCA loadings for first two principal components") + ax.axvline(0, color="black", linewidth=0.5) + ax.legend( + handles=[ + Patch(facecolor=_CORRECT_COLOUR, alpha=0.6, label="PC 1 loading"), + Patch(facecolor=_PALETTE[5], alpha=0.4, label="PC 2 loading"), + ], + loc="lower right", + ) + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "12g_pca_loadings", output_dir) + + +def plot_discriminative_analysis(df: pd.DataFrame, output_dir: Path) -> None: + """Section 12: discriminative stats, correlation, violins, pairplot.""" + available = [f for f in FEATURE_COLUMNS if f in df.columns] + disc_stats = compute_discriminative_stats(df, available) + + # 12a — AUROC bar chart + fig, ax = plt.subplots(figsize=(8, 7)) + colours = [ + _PALETTE[0] if v >= 0.7 else _NEUTRAL_COLOUR for v in disc_stats["auroc"] + ] + ax.barh(range(len(disc_stats)), disc_stats["auroc"], color=colours) + ax.set_yticks(range(len(disc_stats))) + ax.set_yticklabels([_nice_label(f) for f in disc_stats["feature"]], fontsize=9) + ax.set_xlabel("AUROC") + ax.set_title("Per-feature AUROC for separating correct vs incorrect") + ax.axvline(0.5, color="grey", linestyle="--", linewidth=0.8) + ax.invert_yaxis() + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, "12a_feature_auroc_ranking", output_dir) + + # 12b — correlation matrix + corr = df[available].corr() + fig, ax = plt.subplots(figsize=(14, 12)) + mask = np.triu(np.ones_like(corr, dtype=bool), k=1) + sns.heatmap( + corr, + mask=mask, + cmap=_diverging_cmap(), + center=0, + ax=ax, + xticklabels=[_nice_label(c) for c in available], + yticklabels=[_nice_label(c) for c in available], + annot=True, + fmt=".2f", + annot_kws={"size": 6}, + linewidths=0.5, + square=True, + vmin=-1, + vmax=1, + cbar_kws={"label": "Pearson r"}, + ) + ax.set_title("Feature correlation matrix") + ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right", fontsize=7) + ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=7) + fig.tight_layout() + _save_fig(fig, "12b_correlation_matrix", output_dir) + + # 12c — violin plots per feature group + feature_groups = { + "Fragment match": [f for f in FRAGMENT_FEATURES if f in df.columns], + "Beam search": [f for f in BEAM_FEATURES if f in df.columns], + "Token-level": [f for f in TOKEN_FEATURES if f in df.columns], + } + for _group_name, group_feats in feature_groups.items(): + for feat in group_feats: + fig, ax = plt.subplots(figsize=(6, 5)) + sns.violinplot( + data=df, + x="hue", + y=feat, + hue="hue", + ax=ax, + palette=HUE_PALETTE, + order=HUE_ORDER, + hue_order=HUE_ORDER, + inner="quartile", + cut=0, + linewidth=0.8, + legend=False, + ) + ax.set_xlabel("") + ax.set_ylabel(_nice_label(feat)) + ax.set_title(f"{_nice_label(feat)} by identification status") + _style_ax(ax) + fig.tight_layout() + _save_fig(fig, f"12c_violin_{feat}", output_dir) + + # 12e — pairplot of top-5 features + top5 = disc_stats.head(5)["feature"].tolist() + sample_size = min(2000, len(df)) + df_sample = df[top5 + ["hue"]].sample(n=sample_size, random_state=42) + + g = sns.pairplot( + df_sample, + vars=top5, + hue="hue", + palette=HUE_PALETTE, + hue_order=HUE_ORDER, + diag_kind="kde", + plot_kws={"alpha": 0.25, "s": 8, "rasterized": True}, + diag_kws={"fill": True, "alpha": 0.3}, + height=2.2, + ) + g.figure.suptitle("Pairplot of top-5 discriminative features", y=1.01, fontsize=13) + g._legend.set_title("Identification") + for ax_row in g.axes: + for ax_item in ax_row: + xl = ax_item.get_xlabel() + yl = ax_item.get_ylabel() + if xl: + ax_item.set_xlabel(_nice_label(xl), fontsize=7) + if yl: + ax_item.set_ylabel(_nice_label(yl), fontsize=7) + _style_ax(ax_item) + _save_fig(g.figure, "12e_pairplot_top5", output_dir) + + _plot_pca(df, available, output_dir) + + # Save discriminative stats as CSV for reference + disc_stats.to_csv(output_dir / "discriminative_stats.csv", index=False) + + +def print_summary(df: pd.DataFrame) -> None: + """Section 13: summary statistics printed to stdout.""" + available = [f for f in FEATURE_COLUMNS if f in df.columns] + disc_stats = compute_discriminative_stats(df, available) + + print("=" * 80) + print("DISCRIMINATIVE STATISTICS SUMMARY") + print("=" * 80) + n_correct = df["correct"].sum() + n_incorrect = (~df["correct"]).sum() + print( + f"\nDataset: {len(df):,} spectra | {n_correct:,} correct | {n_incorrect:,} incorrect" + ) + print(f"Class balance: {df['correct'].mean():.1%} correct\n") + + print("Per-feature discriminative power (sorted by AUROC):") + print("-" * 80) + print(disc_stats.to_string(index=False, float_format="%.3f")) + + print("\nTop-5 features by AUROC:") + for _, row in disc_stats.head(5).iterrows(): + print( + f" {_nice_label(row['feature']):40s} AUROC={row['auroc']:.3f} " + f"KS={row['ks_stat']:.3f} d={row['cohens_d']:.3f}" + ) + + print("\nBottom-5 features by AUROC:") + for _, row in disc_stats.tail(5).iterrows(): + print( + f" {_nice_label(row['feature']):40s} AUROC={row['auroc']:.3f} " + f"KS={row['ks_stat']:.3f} d={row['cohens_d']:.3f}" + ) + + print("\nConfidence statistics:") + print(f" Overall mean confidence: {df['confidence'].mean():.3f}") + print(f" Correct mean confidence: {df[df['correct']]['confidence'].mean():.3f}") + print(f" Incorrect mean confidence: {df[~df['correct']]['confidence'].mean():.3f}") + print( + f" Confidence AUROC: " + f"{roc_auc_score(df['correct'].astype(int), df['confidence']):.3f}" + ) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +def _load_metadata( + metadata_train: Path | None, + metadata_val: Path | None, +) -> pd.DataFrame | None: + """Load and concatenate metadata parquets when provided.""" + parts: list[pd.DataFrame] = [] + for path in (metadata_train, metadata_val): + if path is not None: + logger.info("Loading metadata from %s", path) + parts.append(pl.read_parquet(path).to_pandas()) + if not parts: + return None + df_meta = pd.concat(parts, ignore_index=True) + logger.info( + "Combined metadata: %s rows, %s columns", + f"{len(df_meta):,}", + len(df_meta.columns), + ) + df_meta["hue"] = df_meta["correct"].map( + {True: HUE_LABEL_CORRECT, False: HUE_LABEL_INCORRECT} + ) + return df_meta + + +@app.command() +def main( + features_train: Annotated[ + Path, + typer.Option( + "--features-train", + help=( + "Path to the training features parquet produced by " + "`winnow compute-features`." + ), + ), + ], + features_val: Annotated[ + Optional[Path], + typer.Option( + "--features-val", + help=( + "Optional path to validation features parquet. When provided the " + "train and val splits are concatenated for richer plots." + ), + ), + ] = None, + metadata_train: Annotated[ + Optional[Path], + typer.Option( + "--metadata-train", + help=( + "Optional path to full training metadata parquet (produced by " + "compute-features with metadata_output_path set). Enables mirror " + "plots, RT scatter, and token stems." + ), + ), + ] = None, + metadata_val: Annotated[ + Optional[Path], + typer.Option( + "--metadata-val", + help="Optional path to full validation metadata parquet.", + ), + ] = None, + predictions_csv: Annotated[ + Optional[Path], + typer.Option( + "--predictions-csv", + help=( + "Optional path to InstaNovo-style predictions CSV with beam columns. " + "Enables beam confidence stem plots." + ), + ), + ] = None, + plots_dir: Annotated[ + Optional[Path], + typer.Option( + "--plots-dir", + help=( + "Directory for output plots. Defaults to a " + "`feature_investigation_plots` subdirectory next to --features-train." + ), + ), + ] = None, +) -> None: + """Generate all feature investigation plots.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + warnings.filterwarnings("ignore", category=FutureWarning) + + # -- Load features (always required) -- + logger.info("Loading training features from %s", features_train) + df = pl.read_parquet(features_train).to_pandas() + + if features_val is not None: + logger.info("Loading validation features from %s", features_val) + df_val = pl.read_parquet(features_val).to_pandas() + df = pd.concat([df, df_val], ignore_index=True) + logger.info("Combined dataset: %s spectra", f"{len(df):,}") + + df["hue"] = df["correct"].map({True: HUE_LABEL_CORRECT, False: HUE_LABEL_INCORRECT}) + + # -- Load metadata (optional) -- + df_meta = _load_metadata(metadata_train, metadata_val) + + has_metadata = df_meta is not None + has_predictions_csv = predictions_csv is not None + + if plots_dir is None: + resolved_plots = features_train.parent / "feature_investigation_plots" + else: + resolved_plots = plots_dir + resolved_plots.mkdir(parents=True, exist_ok=True) + logger.info("Saving plots to %s/", resolved_plots) + + logger.info("Dataset shape: %s", df.shape) + n_correct = df["correct"].sum() + n_incorrect = (~df["correct"]).sum() + logger.info( + "Correct: %s | Incorrect: %s | Total: %s", + f"{n_correct:,}", + f"{n_incorrect:,}", + f"{len(df):,}", + ) + logger.info("Class balance: %.1f%% correct", df["correct"].mean() * 100) + + n_steps = 7 + has_metadata * 2 + has_predictions_csv + step = 0 + + step += 1 + logger.info("[%s/%s] Confidence distribution...", step, n_steps) + plot_confidence(df, resolved_plots) + + step += 1 + logger.info("[%s/%s] Mass error vs confidence...", step, n_steps) + plot_mass_error(df, resolved_plots) + + if has_metadata: + step += 1 + logger.info("[%s/%s] Mirror spectrum plots...", step, n_steps) + plot_mirror_spectra(df_meta, resolved_plots) + + step += 1 + logger.info("[%s/%s] Fragment ion match features...", step, n_steps) + plot_fragment_features(df, resolved_plots) + + step += 1 + logger.info("[%s/%s] iRT error...", step, n_steps) + plot_irt(df, df_meta, resolved_plots) + + if has_metadata: + step += 1 + logger.info("[%s/%s] Token-level stem plots...", step, n_steps) + plot_token_stems(df_meta, resolved_plots) + + if has_predictions_csv: + step += 1 + logger.info("[%s/%s] Beam confidence stem plots...", step, n_steps) + assert predictions_csv is not None + plot_beam_stems(predictions_csv, resolved_plots) + + step += 1 + logger.info("[%s/%s] Beam search features...", step, n_steps) + plot_beam_features(df, resolved_plots) + + step += 1 + logger.info("[%s/%s] Token-level features...", step, n_steps) + plot_token_features(df, resolved_plots) + + step += 1 + logger.info( + "[%s/%s] Discriminative analysis (AUROC, correlation, violins, pairplot)...", + step, + n_steps, + ) + plot_discriminative_analysis(df, resolved_plots) + + print_summary(df) + + logger.info("Done — plots saved to %s/", resolved_plots) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/run_external_peptide_holdout_benchmark.py b/paper_scripts/run_external_peptide_holdout_benchmark.py new file mode 100644 index 00000000..e791d80b --- /dev/null +++ b/paper_scripts/run_external_peptide_holdout_benchmark.py @@ -0,0 +1,1160 @@ +#!/usr/bin/env python3 +"""Glissade-style external peptide score-mixture benchmark. + +Builds a **shared** matched pool S_m (labelled-test Novor-correct peptides) and +external pool S_e (unlabelled proteome-external peptides) after filter → +max-score-per-peptide (NovoBoard mass-deltas converted to ProForma; unsupported +mods dropped; NovoBoard target-decoy pairs gated so every retained key has a +twin). Unlabelled / external peptides require normalised length ≥ 8 (proteome +substring proxy); labelled matched peptides and Glissade's training-split +reference keep short peptides (Novor agreement). Novor and proteome-hit labels +are computed once on Winnow and reused for NovoBoard by ``spectrum_id``. +Method-specific scores are attached to the same peptide keys. Mixtures control +π₀ explicitly. + +Mixtures are drawn **without replacement** so every peptide key is unique. All +three methods therefore score the identical mixture and realise the same π₀; +NovoBoard's max-score-per-peptide step is a no-op, which is asserted per +mixture. + +Each tool draws its null/reference information from the same place, the +annotated *training* split of its own organism: Winnow through the pretrained +per-dataset calibrator, Glissade through the training-split matched score +distribution, NovoBoard through its training-tuned decoy masking rate. No tool +fits on the evaluation labels. + +NovoBoard peptide FDR uses max-target → twin-decoy TDC. Winnow uses max +calibrated confidence then nonparametric FDR (PSM-calibrator proxy). Glissade +uses native bootstrap FDR with NumPy seeded from the benchmark RNG. + +External tool checkouts for local results: + +- ``--novoboard-root``: ``{root}/{dataset}/novoboard/`` target/decoy CSVs (the + ``datasets`` dir of a NovoBoard checkout). Local runs used fork + ``git@github.com:JemmaLDaniel/NovoBoard.git``, branch + ``feat/adapt-to-instanovo`` at + ``a9faab3ef1af06987599c2f01e6ba96072c80172``. +- Glissade is installed from the ``paper`` dependency group + (``uv sync --group paper``), which pins + ``https://github.com/JemmaLDaniel/glissade.git`` branch ``winnow-benchmark`` + at ``7c723a2af4a88fda84a6bd4f223b351179bd36da``. +""" + +from __future__ import annotations + +import importlib +import logging +import sys +from pathlib import Path +from typing import Annotated, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import typer + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +from fdr_tool_comparison_preprocess import ( # noqa: E402 + LABELLED_MIN_PEPTIDE_LENGTH, + MIN_PEPTIDE_LENGTH, + assert_shared_prediction_keys, + attach_labels_by_spectrum_id, + confidence_to_log_prob, + filter_novoboard_target_decoy_pairs, + filter_prediction_table, + label_series_by_spectrum_id, + max_score_per_peptide, + novoboard_max_target_twin_decoy_tdc, + novor_correctness_mask, + prepare_novoboard_decoy_by_pair, + restrict_winnow_to_novoboard_spectra, +) +from fdr_tool_comparison_summaries import ( # noqa: E402 + SUMMARY_THRESHOLDS, + database_grounded_q_from_labels, + mean_abs_q_dev_vs_reference, + summarise_holdout_results, + write_summary_tables, +) +from plot_eval_results import _PALETTE, _display_name, _save_fig, _style_ax # noqa: E402 +from plot_fdr_method_comparison import ( # noqa: E402 + DEFAULT_MODEL_ROOT, + DEFAULT_WINNOW_RESULTS, + build_dataset_configs, + load_novoboard_target_decoy, + load_winnow, +) +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +logger = logging.getLogger(__name__) +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark" +DEFAULT_DATASETS = ["helaqc", "celegans"] +DEFAULT_Q_THRESHOLDS = [round(float(x), 2) for x in np.linspace(0.0, 0.25, 26)] +DEFAULT_PI0_GRID = [0.5, 0.6, 0.7, 0.8, 0.9] +DEFAULT_N_ITERATIONS = 20 +# Matches the Glissade package default for run_bootstraps. +DEFAULT_N_BOOTSTRAPS = 10 +# Prefer the full matched pool; |S_c| is then capped so the highest π₀ remains +# drawable without replacement from S_e. +DEFAULT_HOLDOUT_FRAC = 1.0 +DEFAULT_SEED = 42 +METHODS = ("Winnow", "NovoBoard", "Glissade") + + +def max_correct_pool_for_pi0_grid(n_external_pool: int, pi0_grid: list[float]) -> int: + """Largest |S_c| such that every π₀ in ``pi0_grid`` fits without replacement. + + Requires ``round(π₀ / (1 - π₀) · |S_c|) ≤ |S_e|`` for each target π₀. + """ + if n_external_pool < 1: + return 0 + limit = n_external_pool + for pi0 in pi0_grid: + if not 0.0 < pi0 < 1.0: + continue + ratio = pi0 / (1.0 - pi0) + # Largest n with round(ratio * n) <= n_external_pool. + # For ratio = k integer (e.g. 0.9 → 9), this is floor(pool / k). + hi = int(n_external_pool / ratio) + 2 + n_ok = 0 + for n in range(1, hi + 1): + if int(round(ratio * n)) <= n_external_pool: + n_ok = n + else: + break + limit = min(limit, n_ok) + return max(0, limit) + + +def _load_glissade_functions(): + """Import Glissade FDR helpers from the installed ``paper`` dependency group.""" + try: + module = importlib.import_module("glissade.glissade") + except ImportError as exc: + raise ImportError( + "Glissade is required for this benchmark. Install it with " + "`uv sync --group paper` (pins JemmaLDaniel/glissade @ " + "7c723a2af4a88fda84a6bd4f223b351179bd36da)." + ) from exc + return module.run_bootstraps, module.annotate_results, module.compute_fdr_transform + + +def _load_winnow_with_raw_confidence( + predictions_dir: Path, fasta: Path, eval_type: str +) -> pd.DataFrame: + """Load Winnow preds and ensure raw ``confidence`` is present.""" + df = load_winnow(predictions_dir, fasta, eval_type) # type: ignore[arg-type] + if "confidence" not in df.columns: + meta_path = predictions_dir / "metadata.csv" + if not meta_path.is_file(): + raise FileNotFoundError(meta_path) + meta = pd.read_csv(meta_path, usecols=["spectrum_id", "confidence"]) + df = df.merge(meta, on="spectrum_id", how="inner") + return df + + +def build_glissade_training_reference( + train_metadata: Path, + *, + min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, +) -> pd.DataFrame: + """Matched reference score distribution for Glissade, from the training split. + + Glissade anchors its null-fraction estimate on a database-matched score + distribution. Taking that anchor from the annotated training split puts it on + the same data the Winnow calibrator was trained on and the NovoBoard decoy + masking rate was tuned on, and keeps it disjoint from the evaluation spectra. + Short peptides are retained: the reference is labelled (Novor-correct). + + Args: + train_metadata: Calibrator training metadata parquet. + min_length: Minimum normalised peptide length (default: labelled floor). + + Returns: + One row per Novor-correct training peptide with ``raw_confidence`` and + ``score_glissade``. + """ + if not train_metadata.is_file(): + raise FileNotFoundError(train_metadata) + train = pl.read_parquet( + train_metadata, + columns=["spectrum_id", "prediction", "sequence", "confidence"], + ).to_pandas() + train = filter_prediction_table( + train, "prediction", min_length=min_length, key_col="peptide_key" + ) + train["correct"] = novor_correctness_mask(train["sequence"], train["prediction"]) + matched = max_score_per_peptide( + train.loc[train["correct"]], "peptide_key", "confidence" + ) + reference = matched[["peptide_key", "confidence"]].rename( + columns={"confidence": "raw_confidence"} + ) + reference["score_glissade"] = confidence_to_log_prob(reference["raw_confidence"]) + logger.info( + "Glissade training reference: %d matched peptides from %s", + len(reference), + train_metadata, + ) + return reference + + +def _namespace_pair_keys(df: pd.DataFrame, namespace: str) -> pd.DataFrame: + """Prefix ``_pair_key`` to avoid collisions across splits.""" + work = df.copy() + if "_pair_key" not in work.columns: + raise ValueError("Missing '_pair_key'") + work["_pair_key"] = namespace + ":" + work["_pair_key"].astype(str) + return work + + +def build_shared_score_tables( + *, + dataset: str, + winnow_results: Path, + novoboard_root: Path, + model_root: Path = DEFAULT_MODEL_ROOT, + unlabelled_min_length: int = MIN_PEPTIDE_LENGTH, + labelled_min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, +) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """Return shared matched/external keys with per-method scores. + + All methods follow filter → :func:`max_score_per_peptide`. NovoBoard tables + are pair-gated (equal target/decoy twins) before max-dedupe so mixture keys + always have a twin. Labelled S_m and Glissade's training reference keep short + peptides (``labelled_min_length``); unlabelled S_e uses + ``unlabelled_min_length`` for the proteome-substring proxy. Labelled S_m uses + Novor correctness computed once on Winnow and reused for NovoBoard by + ``spectrum_id``; S_e uses proteome-hit the same way. Glissade scores are + ``log`` raw InstaNovo confidence on the shared keys. + + Returns: + matched_scores: one row per peptide_key in S_m with method scores. + external_scores: same for S_e. + nb_decoy_combined: namespaced twin-valid decoys for twin TDC. + glissade_reference: training-split matched scores for Glissade's anchor. + """ + cfg = build_dataset_configs( + winnow_results, novoboard_root=novoboard_root, model_root=model_root + )[dataset] + + winnow_test = _load_winnow_with_raw_confidence( + cfg.winnow_test, cfg.fasta, "labelled" + ) + winnow_unlab = _load_winnow_with_raw_confidence( + cfg.winnow_unlabelled, cfg.fasta, "unlabelled" + ) + # load_winnow already applied the labelled/unlabelled length floors; re-apply + # explicitly so callers can override without depending on that path. + winnow_test = filter_prediction_table( + winnow_test, + "prediction", + min_length=labelled_min_length, + key_col="peptide_key", + ) + w_unlab = filter_prediction_table( + winnow_unlab, + "prediction", + min_length=unlabelled_min_length, + key_col="peptide_key", + ) + if "proteome_hit" not in w_unlab.columns: + raise KeyError("Expected proteome_hit on unlabelled Winnow table") + + nb_test_target, nb_test_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "test", cfg.novoboard_decoy_rate + ) + nb_unlab_target, nb_unlab_decoy = load_novoboard_target_decoy( + cfg.novoboard_dir, "unlabelled", cfg.novoboard_decoy_rate + ) + nb_test_target = _namespace_pair_keys(nb_test_target, "test") + nb_test_decoy = _namespace_pair_keys(nb_test_decoy, "test") + nb_unlab_target = _namespace_pair_keys(nb_unlab_target, "unlabelled") + nb_unlab_decoy = _namespace_pair_keys(nb_unlab_decoy, "unlabelled") + + nb_test_target, nb_test_decoy = filter_novoboard_target_decoy_pairs( + nb_test_target, + nb_test_decoy, + min_length=labelled_min_length, + key_col="peptide_key", + ) + nb_unlab_target, nb_unlab_decoy = filter_novoboard_target_decoy_pairs( + nb_unlab_target, + nb_unlab_decoy, + min_length=unlabelled_min_length, + key_col="peptide_key", + ) + nb_decoy_combined = pd.concat( + [nb_test_decoy, nb_unlab_decoy], ignore_index=True, sort=False + ) + + # Twin-valid NovoBoard ⊆ Winnow; shared labels once on Winnow. + winnow_test = restrict_winnow_to_novoboard_spectra(winnow_test, nb_test_target) + w_unlab = restrict_winnow_to_novoboard_spectra(w_unlab, nb_unlab_target) + assert_shared_prediction_keys(winnow_test, nb_test_target) + assert_shared_prediction_keys(w_unlab, nb_unlab_target) + + winnow_test = winnow_test.copy() + winnow_test["correct"] = novor_correctness_mask( + winnow_test["sequence"], winnow_test["prediction"] + ) + correct_by_id = label_series_by_spectrum_id(winnow_test, "correct") + hit_by_id = label_series_by_spectrum_id(w_unlab, "proteome_hit") + nb_test_target = attach_labels_by_spectrum_id( + nb_test_target, correct_by_id, label_col="correct" + ) + nb_unlab_f = attach_labels_by_spectrum_id( + nb_unlab_target, hit_by_id, label_col="proteome_hit" + ) + + w_matched = max_score_per_peptide( + winnow_test.loc[winnow_test["correct"]], + "peptide_key", + "calibrated_confidence", + ) + w_matched_raw = max_score_per_peptide( + winnow_test.loc[winnow_test["correct"]], + "peptide_key", + "confidence", + ) + nb_correct = nb_test_target.loc[nb_test_target["correct"]].copy() + nb_matched = max_score_per_peptide(nb_correct, "peptide_key", "ALC (%)") + + w_external = max_score_per_peptide( + w_unlab.loc[~w_unlab["proteome_hit"].astype(bool)], + "peptide_key", + "calibrated_confidence", + ) + w_external_raw = max_score_per_peptide( + w_unlab.loc[~w_unlab["proteome_hit"].astype(bool)], + "peptide_key", + "confidence", + ) + nb_external = max_score_per_peptide( + nb_unlab_f.loc[~nb_unlab_f["proteome_hit"].astype(bool)], + "peptide_key", + "ALC (%)", + ) + + # Shared membership = Winnow ∩ twin-valid NovoBoard keys. + sm_keys = ( + set(w_matched["peptide_key"]) + & set(nb_matched["peptide_key"]) + & set(w_matched_raw["peptide_key"]) + ) + se_keys = ( + set(w_external["peptide_key"]) + & set(nb_external["peptide_key"]) + & set(w_external_raw["peptide_key"]) + ) + if sm_keys != set(w_matched["peptide_key"]) or sm_keys != set( + nb_matched["peptide_key"] + ): + raise AssertionError( + f"{dataset} shared S_m keys disagree after shared Novor labels: " + f"winnow={len(w_matched)} novoboard={len(nb_matched)} " + f"intersection={len(sm_keys)}" + ) + if se_keys != set(w_external["peptide_key"]) or se_keys != set( + nb_external["peptide_key"] + ): + raise AssertionError( + f"{dataset} shared S_e keys disagree after shared proteome-hit labels: " + f"winnow={len(w_external)} novoboard={len(nb_external)} " + f"intersection={len(se_keys)}" + ) + if not sm_keys: + raise ValueError(f"Empty shared matched pool for {dataset}") + if not se_keys: + raise ValueError(f"Empty shared external pool for {dataset}") + + matched = pd.DataFrame({"peptide_key": sorted(sm_keys)}) + matched = matched.merge( + w_matched[["peptide_key", "calibrated_confidence"]].rename( + columns={"calibrated_confidence": "score_winnow"} + ), + on="peptide_key", + how="left", + ) + matched = matched.merge( + nb_matched[["peptide_key", "ALC (%)"]].rename( + columns={"ALC (%)": "score_novoboard"} + ), + on="peptide_key", + how="left", + ) + matched = matched.merge( + w_matched_raw[["peptide_key", "confidence"]].rename( + columns={"confidence": "raw_confidence"} + ), + on="peptide_key", + how="left", + ) + matched["score_glissade"] = confidence_to_log_prob(matched["raw_confidence"]) + + external = pd.DataFrame({"peptide_key": sorted(se_keys)}) + external = external.merge( + w_external[["peptide_key", "calibrated_confidence"]].rename( + columns={"calibrated_confidence": "score_winnow"} + ), + on="peptide_key", + how="left", + ) + external = external.merge( + nb_external[["peptide_key", "ALC (%)"]].rename( + columns={"ALC (%)": "score_novoboard"} + ), + on="peptide_key", + how="left", + ) + external = external.merge( + w_external_raw[["peptide_key", "confidence"]].rename( + columns={"confidence": "raw_confidence"} + ), + on="peptide_key", + how="left", + ) + external["score_glissade"] = confidence_to_log_prob(external["raw_confidence"]) + + # Attach NovoBoard pair metadata for twin TDC on mixture subsets. + nb_ext_meta = nb_unlab_f.loc[ + ~nb_unlab_f["proteome_hit"].astype(bool), + ["peptide_key", "Peptide", "ALC (%)", "spectrum_id", "_pair_key", "Scan"], + ].copy() + nb_ext_meta = ( + nb_ext_meta.sort_values("ALC (%)", ascending=False) + .groupby("peptide_key", as_index=False) + .first() + ) + external = external.merge( + nb_ext_meta.rename( + columns={ + "Peptide": "peptide_novoboard", + "spectrum_id": "spectrum_id_novoboard", + } + ), + on="peptide_key", + how="left", + ) + + nb_match_meta = ( + nb_correct[ + ["peptide_key", "Peptide", "ALC (%)", "spectrum_id", "_pair_key", "Scan"] + ] + .sort_values("ALC (%)", ascending=False) + .groupby("peptide_key", as_index=False) + .first() + .rename( + columns={ + "Peptide": "peptide_novoboard", + "spectrum_id": "spectrum_id_novoboard", + } + ) + ) + matched = matched.merge(nb_match_meta, on="peptide_key", how="left") + + twin_coverage_m = ( + float(matched["_pair_key"].notna().mean()) if len(matched) else 0.0 + ) + twin_coverage_e = ( + float(external["_pair_key"].notna().mean()) if len(external) else 0.0 + ) + if twin_coverage_m < 1.0 or twin_coverage_e < 1.0: + raise AssertionError( + f"{dataset} NovoBoard twin coverage incomplete: " + f"Sm={twin_coverage_m:.3f} Se={twin_coverage_e:.3f}" + ) + + glissade_reference = build_glissade_training_reference( + cfg.calibrator_train_metadata, min_length=labelled_min_length + ) + + logger.info( + "%s shared pools: matched=%d external=%d glissade_reference=%d " + "(shared Novor/proteome-hit labels; NB twin coverage 100%%)", + dataset, + len(matched), + len(external), + len(glissade_reference), + ) + return matched, external, nb_decoy_combined, glissade_reference + + +def _estimate_winnow_q_values(mixed: pd.DataFrame) -> pd.DataFrame: + work = mixed[["peptide_key", "score_winnow", "source"]].copy() + work = work.rename(columns={"score_winnow": "score"}) + work = work.dropna(subset=["score"]) + ctrl = NonParametricFDRControl() + ctrl.fit(work["score"]) + q_table = ctrl.add_psm_q_value(work.copy(), "score") + return pd.DataFrame( + { + "peptide_key": q_table["peptide_key"], + "score": q_table["score"], + "source": q_table["source"], + "q_value": q_table["psm_q_value"], + "method": "Winnow", + } + ) + + +def _estimate_novoboard_q_values( + mixed: pd.DataFrame, + decoy_by_pair: pd.DataFrame, +) -> pd.DataFrame: + target = pd.DataFrame( + { + "Peptide": mixed["peptide_novoboard"].fillna(mixed["peptide_key"]), + "ALC (%)": mixed["score_novoboard"], + "spectrum_id": mixed.get( + "spectrum_id_novoboard", + pd.Series(np.arange(len(mixed)), dtype=str), + ), + "_pair_key": mixed["_pair_key"], + "Scan": mixed.get("Scan", mixed["_pair_key"]), + "source": mixed["source"], + "peptide_key": mixed["peptide_key"], + } + ) + target = target.dropna(subset=["ALC (%)", "_pair_key"]) + # Restrict twin TDC to the mixture peptide keys only. + table = novoboard_max_target_twin_decoy_tdc( + target, + pd.DataFrame(), + target_peptide_keys=set(target["peptide_key"].astype(str)), + decoy_by_pair=decoy_by_pair, + log_missing_twins=False, + ) + targets = table[table["is_target"]].copy() + # Map back sources for FDP. + source_map = mixed.set_index("peptide_key")["source"].to_dict() + key_col = "_peptide_key" if "_peptide_key" in targets.columns else "peptide_key" + targets["peptide_key"] = targets[key_col] + targets["source"] = targets["peptide_key"].map(source_map) + return pd.DataFrame( + { + "peptide_key": targets["peptide_key"], + "score": targets["ALC (%)"], + "source": targets["source"], + "q_value": targets["estimated_q_value"], + "method": "NovoBoard", + } + ) + + +def _estimate_glissade_q_values( + mixed: pd.DataFrame, + matched_reference: pd.DataFrame, + *, + n_bootstraps: int, + rng: np.random.Generator, +) -> pd.DataFrame: + run_bootstraps, annotate_results, compute_fdr_transform = _load_glissade_functions() + # Glissade FDR is defined on the mixture scores; the reference is the + # training-split matched distribution and is never scored itself. + mixed_scores = mixed["score_glissade"].astype(float).to_numpy() + matched_scores = matched_reference["score_glissade"].astype(float).to_numpy() + if len(mixed_scores) < 10 or len(matched_scores) < 10: + raise ValueError("Glissade FDR requires ≥10 reference and mixture scores") + + peptides = mixed["peptide_key"].astype(str).tolist() + np.random.seed(int(rng.integers(0, 2**32 - 1))) + fdrs, grid, _ = run_bootstraps( + matched_scores, + mixed_scores, + n_bootstraps=n_bootstraps, + ) + out_peptides, peptide_fdrs, scores = annotate_results( + peptides, mixed_scores, fdrs, grid + ) + peptide_fdrs = compute_fdr_transform(peptide_fdrs) + source_map = mixed.set_index("peptide_key")["source"].to_dict() + return pd.DataFrame( + { + "peptide_key": out_peptides, + "score": scores, + "source": [source_map.get(p) for p in out_peptides], + "q_value": peptide_fdrs, + "method": "Glissade", + } + ) + + +def _mixture_result_rows( + *, + dataset: str, + pi0: float, + true_pi0: float, + holdout_frac: float, + seed: int, + iteration: int, + method: str, + q_table: pd.DataFrame, + correct_keys: set[str], + n_external: int, + thresholds: list[float], + q_ref: np.ndarray | None = None, +) -> list[dict[str, object]]: + """Build per-threshold result rows for one method on one mixture.""" + work = q_table.dropna(subset=["q_value"]) + if q_ref is not None: + if len(q_ref) != len(q_table): + raise ValueError( + f"q_ref length {len(q_ref)} does not match q_table length {len(q_table)}" + ) + q_ref_aligned = pd.Series(np.asarray(q_ref, dtype=float), index=q_table.index) + q_ref_work = q_ref_aligned.loc[work.index].to_numpy(dtype=float) + q_devs = mean_abs_q_dev_vs_reference( + work["q_value"].to_numpy(dtype=float), q_ref_work, thresholds + ) + else: + q_devs = [float("nan")] * len(thresholds) + + rows: list[dict[str, object]] = [] + for threshold, q_dev in zip(thresholds, q_devs): + accepted = work[work["q_value"] <= threshold] + n_accepted = len(accepted) + n_true = int(accepted["peptide_key"].isin(correct_keys).sum()) + n_false = n_accepted - n_true + rows.append( + { + "dataset": dataset, + "pi0_target": float(pi0), + "true_pi0": float(true_pi0), + "holdout_frac": float(holdout_frac), + "seed": seed, + "iteration": iteration, + "method": method, + "q_value_threshold": float(threshold), + "mixed_external_peptides": n_external, + "correct_peptides": len(correct_keys), + "accepted_peptides": n_accepted, + "true_correct_peptides": n_true, + "false_external_peptides": n_false, + "observed_fdp": (n_false / n_accepted if n_accepted else np.nan), + "correct_discovery_pct": ( + 100.0 * n_true / len(correct_keys) if correct_keys else np.nan + ), + "mean_abs_q_dev_vs_db": float(q_dev), + } + ) + return rows + + +def _evaluate_mixture_methods( + *, + mixed: pd.DataFrame, + estimator_reference: pd.DataFrame, + decoy_by_pair: pd.DataFrame, + n_bootstraps: int, + iter_rng: np.random.Generator, + dataset: str, + pi0: float, + true_pi0: float, + holdout_frac: float, + seed: int, + iteration: int, + correct_keys: set[str], + n_external: int, + thresholds: list[float], +) -> list[dict[str, object]]: + """Run Winnow / NovoBoard / Glissade FDR on one mixture and collect rows.""" + estimators: dict[str, object] = { + "Winnow": lambda m, _r: _estimate_winnow_q_values(m), + "NovoBoard": lambda m, _r: _estimate_novoboard_q_values(m, decoy_by_pair), + "Glissade": lambda m, r, rng=iter_rng: _estimate_glissade_q_values( + m, + r, + n_bootstraps=n_bootstraps, + rng=rng, + ), + } + # Raw-confidence DBG on the mixture used for NovoBoard and Glissade q-deviation. Winnow keeps calibrated-score DBG unset here. + is_correct = mixed["peptide_key"].astype(str).isin(correct_keys).to_numpy() + q_db_raw = database_grounded_q_from_labels( + mixed["score_novoboard"].to_numpy(dtype=float), is_correct + ) + q_db_raw_by_key = dict(zip(mixed["peptide_key"].astype(str), q_db_raw, strict=True)) + + rows: list[dict[str, object]] = [] + mixture_keys = set(mixed["peptide_key"].astype(str)) + for method, estimator in estimators.items(): + try: + q_table = estimator(mixed, estimator_reference) # type: ignore[operator] + except Exception as exc: # noqa: BLE001 - boundary around external tool + logger.warning( + "%s FDR failed dataset=%s pi0=%.3g iter=%d: %s", + method, + dataset, + pi0, + iteration, + exc, + ) + continue + scored_keys = set(q_table["peptide_key"].astype(str)) + if scored_keys != mixture_keys: + raise AssertionError( + f"{method} scored a different mixture on dataset={dataset} " + f"pi0={pi0:.3g} iter={iteration}: mixture={len(mixture_keys)} " + f"scored={len(scored_keys)} " + f"missing={len(mixture_keys - scored_keys)} " + f"extra={len(scored_keys - mixture_keys)}" + ) + q_ref: np.ndarray | None = None + if method in ("NovoBoard", "Glissade"): + q_ref = ( + q_table["peptide_key"] + .astype(str) + .map(q_db_raw_by_key) + .to_numpy(dtype=float) + ) + rows.extend( + _mixture_result_rows( + dataset=dataset, + pi0=pi0, + true_pi0=true_pi0, + holdout_frac=holdout_frac, + seed=seed, + iteration=iteration, + method=method, + q_table=q_table, + correct_keys=correct_keys, + n_external=n_external, + thresholds=thresholds, + q_ref=q_ref, + ) + ) + return rows + + +def _sample_correct_component( + *, + dataset: str, + matched: pd.DataFrame, + external: pd.DataFrame, + pi0_grid: list[float], + holdout_frac: float, + rng: np.random.Generator, +) -> tuple[pd.DataFrame, set[str], float]: + """Sample S_c from matched, capped so the π₀ grid fits in S_e.""" + n_from_frac = max(1, int(round(len(matched) * holdout_frac))) + n_from_frac = min(n_from_frac, len(matched)) + n_cap = max_correct_pool_for_pi0_grid(len(external), pi0_grid) + if n_cap < 1: + raise ValueError( + f"{dataset}: external pool of {len(external)} cannot support any " + f"π₀ in {pi0_grid} without replacement" + ) + n_correct = min(n_from_frac, n_cap) + if n_correct < n_from_frac: + logger.info( + "%s capping |S_c| from %d to %d so π₀ grid %s fits in |S_e|=%d " + "without replacement", + dataset, + n_from_frac, + n_correct, + pi0_grid, + len(external), + ) + correct = matched.sample(n=n_correct, random_state=int(rng.integers(0, 2**32 - 1))) + correct = correct.copy() + correct["source"] = "correct" + correct_keys = set(correct["peptide_key"]) + effective_holdout_frac = n_correct / len(matched) if len(matched) else 0.0 + overlap = correct_keys & set(external["peptide_key"]) + if overlap: + raise AssertionError( + f"{dataset} S_c and S_e share {len(overlap)} peptide keys; " + "mixture sources would be ambiguous" + ) + return correct, correct_keys, effective_holdout_frac + + +def _external_draw_size( + *, + dataset: str, + pi0: float, + n_correct: int, + n_external_pool: int, +) -> int | None: + """Return |S_e'| for π₀, or None to skip; raise if the pool is too small.""" + if not 0.0 < pi0 < 1.0: + return None + n_external = int(round(pi0 / (1.0 - pi0) * n_correct)) + if n_external < 1: + logger.warning("Skipping pi0=%.3g: requested |S_e'|=%d", pi0, n_external) + return None + if n_external > n_external_pool: + raise AssertionError( + f"{dataset} pi0={pi0:.3g}: |S_e'|={n_external} exceeds pool " + f"{n_external_pool} after |S_c| cap {n_correct}" + ) + return n_external + + +def evaluate_controlled_mixtures( + *, + dataset: str, + matched: pd.DataFrame, + external: pd.DataFrame, + nb_decoy: pd.DataFrame, + glissade_reference: pd.DataFrame, + pi0_grid: list[float], + holdout_frac: float, + seed: int, + n_iterations: int, + thresholds: list[float], + n_bootstraps: int, +) -> list[dict[str, object]]: + """Evaluate all methods on shared mixtures with controlled π₀. + + Uses as much of the matched pool as possible (up to holdout_frac), capped + so every π₀ in pi0_grid can draw its null component from S_e without + replacement. The null component is drawn without replacement, so mixture + peptide keys are unique and every method realises the same π₀ on the same + rows. + """ + rng = np.random.default_rng(seed) + correct, correct_keys, effective_holdout_frac = _sample_correct_component( + dataset=dataset, + matched=matched, + external=external, + pi0_grid=pi0_grid, + holdout_frac=holdout_frac, + rng=rng, + ) + decoy_by_pair = prepare_novoboard_decoy_by_pair(nb_decoy, already_filtered=True) + + rows: list[dict[str, object]] = [] + for pi0 in pi0_grid: + n_external = _external_draw_size( + dataset=dataset, + pi0=pi0, + n_correct=len(correct_keys), + n_external_pool=len(external), + ) + if n_external is None: + continue + + for iteration in range(n_iterations): + iter_rng = np.random.default_rng( + seed + 10_000 * iteration + int(1000 * pi0) + ) + ext_sample = external.sample( + n=n_external, + replace=False, + random_state=int(iter_rng.integers(0, 2**32 - 1)), + ).copy() + ext_sample["source"] = "external" + mixed = pd.concat([ext_sample, correct], ignore_index=True, sort=False) + if mixed["peptide_key"].duplicated().any(): + raise AssertionError( + f"{dataset} mixture has duplicate peptide keys at " + f"pi0={pi0:.3g} iter={iteration}" + ) + true_pi0 = len(ext_sample) / (len(ext_sample) + len(correct_keys)) + rows.extend( + _evaluate_mixture_methods( + mixed=mixed, + estimator_reference=glissade_reference, + decoy_by_pair=decoy_by_pair, + n_bootstraps=n_bootstraps, + iter_rng=iter_rng, + dataset=dataset, + pi0=pi0, + true_pi0=true_pi0, + holdout_frac=effective_holdout_frac, + seed=seed, + iteration=iteration, + correct_keys=correct_keys, + n_external=len(ext_sample), + thresholds=thresholds, + ) + ) + return rows + + +def _plot_metric_by_pi0( + dataset_results: pd.DataFrame, + *, + dataset_label: str, + dataset_slug: str, + metric: str, + ylabel: str, + base_name: str, + percent_axis: bool, + output_dir: Path, +) -> None: + method_order = list(METHODS) + colors = {"Winnow": _PALETTE[0], "NovoBoard": _PALETTE[2], "Glissade": _PALETTE[4]} + pi0_values = sorted(dataset_results["pi0_target"].dropna().unique()) + n_pi0 = max(1, len(pi0_values)) + fig, axes = plt.subplots( + 1, n_pi0, figsize=(4.2 * n_pi0, 5.5), sharey=True, squeeze=False + ) + for ax, pi0 in zip(axes[0], pi0_values): + sub = dataset_results[dataset_results["pi0_target"] == pi0] + for method in method_order: + msub = sub[sub["method"] == method] + if msub.empty: + continue + summary = ( + msub.groupby("q_value_threshold", as_index=False)[metric] + .mean(numeric_only=True) + .sort_values("q_value_threshold") + ) + ax.plot( + summary["q_value_threshold"], + summary[metric], + lw=1.5, + label=method, + color=colors[method], + ) + if metric == "observed_fdp": + max_threshold = float(sub["q_value_threshold"].max()) + ax.plot( + [0.0, max_threshold], + [0.0, max_threshold], + color="#666666", + lw=1, + ls="--", + label="Nominal FDR", + ) + ax.set_ylim(bottom=0) + if percent_axis: + ax.set_ylim(0, 100) + ax.set_xlim(0, float(sub["q_value_threshold"].max())) + ax.set_xlabel("Estimated q-value threshold") + ax.set_title(f"π₀={pi0:g}") + _style_ax(ax) + axes[0][0].set_ylabel(ylabel) + axes[0][0].legend(loc="best", fontsize=9) + fig.suptitle(f"{dataset_label} external peptide score-mixture benchmark", y=1.02) + fig.tight_layout() + _save_fig(fig, output_dir / f"{base_name}_{dataset_slug}") + + +def plot_benchmark_results(results: pd.DataFrame, output_dir: Path) -> None: + """Save FDP and recovery plots faceted by π₀.""" + if results.empty: + return + output_dir.mkdir(parents=True, exist_ok=True) + specs = [ + ("observed_fdp", "Observed FDP", "external_peptide_score_mixture_fdp", False), + ( + "correct_discovery_pct", + "Correct peptide recovery\n(% of held-out correct peptides)", + "external_peptide_score_mixture_correct_discovery_pct", + True, + ), + ] + for dataset, dataset_results in results.groupby("dataset", sort=False): + for metric, ylabel, base_name, percent_axis in specs: + _plot_metric_by_pi0( + dataset_results, + dataset_label=_display_name(str(dataset)), + dataset_slug=str(dataset).replace("/", "_"), + metric=metric, + ylabel=ylabel, + base_name=base_name, + percent_axis=percent_axis, + output_dir=output_dir, + ) + + +def write_holdout_summary_tables( + results: pd.DataFrame, + output_dir: Path, + *, + thresholds: list[float] | None = None, +) -> tuple[Path, Path]: + """Aggregate raw mixture rows into acceptance and error/gain CSVs.""" + acceptance, error_gain = summarise_holdout_results( + results, + thresholds=thresholds if thresholds is not None else SUMMARY_THRESHOLDS, + group_extra=("pi0_target",), + ) + return write_summary_tables( + acceptance, error_gain, output_dir, "external_peptide_holdout" + ) + + +@app.command() +def main( + novoboard_root: Annotated[ + Optional[Path], + typer.Option( + "--novoboard-root", + help=( + "Root of NovoBoard per-dataset tables: " + "{root}/{dataset}/novoboard/ with annotated_test*.csv and " + "raw_unlabelled*.csv target/decoy pairs (the datasets/ dir of " + "a NovoBoard checkout). Required unless --summarise-only is set. " + "Local runs used fork JemmaLDaniel/NovoBoard, branch " + "feat/adapt-to-instanovo " + "(commit a9faab3ef1af06987599c2f01e6ba96072c80172)." + ), + ), + ] = None, + results_dir: Annotated[ + Path, + typer.Option("--results-dir", help="Directory for results/summary CSVs."), + ] = DEFAULT_OUTPUT_DIR, + plots_dir: Annotated[ + Optional[Path], + typer.Option( + "--plots-dir", + help="Directory for png/pdf figures (required when plotting).", + ), + ] = None, + datasets: Annotated[ + Optional[list[str]], + typer.Option("--datasets", help="Dataset keys to benchmark."), + ] = None, + pi0_grid: Annotated[ + Optional[list[float]], + typer.Option("--pi0-grid", help="Target mixture null fractions."), + ] = None, + holdout_frac: Annotated[ + float, + typer.Option( + "--holdout-frac", + help=( + "Maximum fraction of shared matched peptides used as S_c " + "(default 1 = prefer the full pool). |S_c| is further capped so " + "every π₀ in --pi0-grid fits in S_e without replacement." + ), + ), + ] = DEFAULT_HOLDOUT_FRAC, + q_thresholds: Annotated[ + Optional[list[float]], + typer.Option( + "--q-thresholds", help="Estimated q-value thresholds to evaluate." + ), + ] = None, + seed: Annotated[int, typer.Option("--seed", help="Random seed.")] = DEFAULT_SEED, + n_iterations: Annotated[ + int, + typer.Option( + "--n-iterations", help="Number of external-score resampling iterations." + ), + ] = DEFAULT_N_ITERATIONS, + winnow_results: Annotated[ + Path, + typer.Option("--winnow-results", help="Winnow results directory."), + ] = DEFAULT_WINNOW_RESULTS, + model_root: Annotated[ + Path, + typer.Option( + "--model-root", + help="Per-dataset calibrator directories, used for Glissade's anchor.", + ), + ] = DEFAULT_MODEL_ROOT, + n_bootstraps: Annotated[ + int, + typer.Option( + "--n-bootstraps", help="Glissade bootstraps per mixture iteration." + ), + ] = DEFAULT_N_BOOTSTRAPS, + min_peptide_length: Annotated[ + int, + typer.Option( + "--min-peptide-length", + help=( + "Minimum normalised peptide length for unlabelled / external " + "pools. Labelled matched pools and Glissade's training reference " + "use the labelled floor (non-empty key only)." + ), + ), + ] = MIN_PEPTIDE_LENGTH, + plot: Annotated[bool, typer.Option(help="Create summary plots.")] = True, + summarise_only: Annotated[ + Optional[Path], + typer.Option( + "--summarise-only", + help="Only write summary CSVs/plots from an existing results CSV.", + ), + ] = None, +) -> None: + """Run the controlled-π₀ external peptide score-mixture benchmark.""" + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + if plot: + if plots_dir is None: + plots_dir = DEFAULT_OUTPUT_DIR / "plots" + results_dir.mkdir(parents=True, exist_ok=True) + plots_dir.mkdir(parents=True, exist_ok=True) + else: + if plots_dir is not None: + raise typer.BadParameter( + "Do not pass --plots-dir with --no-plot; this run writes tables only." + ) + results_dir.mkdir(parents=True, exist_ok=True) + + if summarise_only is not None: + results = pd.read_csv(summarise_only) + write_holdout_summary_tables(results, results_dir) + if plot: + assert plots_dir is not None + plot_benchmark_results(results, plots_dir) + return + + if novoboard_root is None: + raise typer.BadParameter( + "--novoboard-root is required unless --summarise-only is set." + ) + + dataset_keys = datasets if datasets is not None else DEFAULT_DATASETS + pi0s = pi0_grid if pi0_grid is not None else list(DEFAULT_PI0_GRID) + thresholds = ( + q_thresholds if q_thresholds is not None else list(DEFAULT_Q_THRESHOLDS) + ) + + rows: list[dict[str, object]] = [] + for dataset in dataset_keys: + logger.info("Building shared score tables for %s", _display_name(dataset)) + matched, external, nb_decoy, glissade_reference = build_shared_score_tables( + dataset=dataset, + winnow_results=winnow_results, + novoboard_root=novoboard_root, + model_root=model_root, + unlabelled_min_length=min_peptide_length, + labelled_min_length=LABELLED_MIN_PEPTIDE_LENGTH, + ) + rows.extend( + evaluate_controlled_mixtures( + dataset=dataset, + matched=matched, + external=external, + nb_decoy=nb_decoy, + glissade_reference=glissade_reference, + pi0_grid=pi0s, + holdout_frac=holdout_frac, + seed=seed, + n_iterations=n_iterations, + thresholds=thresholds, + n_bootstraps=n_bootstraps, + ) + ) + + results = pd.DataFrame(rows) + results_path = results_dir / "external_peptide_holdout_results.csv" + results.to_csv(results_path, index=False) + logger.info("Wrote %s (%d rows)", results_path, len(results)) + write_holdout_summary_tables(results, results_dir) + if plot: + assert plots_dir is not None + plot_benchmark_results(results, plots_dir) + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/run_feature_ablations.py b/paper_scripts/run_feature_ablations.py new file mode 100644 index 00000000..9d4640f7 --- /dev/null +++ b/paper_scripts/run_feature_ablations.py @@ -0,0 +1,1542 @@ +"""Feature ablation study for Winnow calibrator. + +Trains MLP calibrators on subsets of pre-computed training feature matrices, +computes features from raw spectra for evaluation datasets, and produces +publication-quality plots of calibration, discrimination, and FDR behavior. +""" + +from __future__ import annotations + +import json +import logging +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Iterable, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import polars as pl +import seaborn as sns +import torch +import typer +from rich.logging import RichHandler + +_REPO_ROOT = Path(__file__).resolve().parent.parent +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +_PAPER_SCRIPTS = Path(__file__).resolve().parent +if str(_PAPER_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_PAPER_SCRIPTS)) + + +from feature_subsets import FEATURE_SUBSETS # noqa: E402 +from plot_ablation_summary import ( # noqa: E402 + FDR_BIAS_COLUMN_BY_THRESHOLD, + Q_DEV_COLUMN_BY_THRESHOLD, + TAIL_ECE_COLUMN_BY_THRESHOLD, + assign_ablation_colors, + compute_ece, + compute_fdr_bias_at_fdr_thresholds, + compute_pr_auc, + compute_q_value_deviations, + compute_tail_ece, + compute_tail_ece_at_fdr, + ordered_ablation_configs, +) + +from winnow.calibration.calibrator import ProbabilityCalibrator # noqa: E402 +from winnow.datasets.feature_dataset import FeatureDataset # noqa: E402 +from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 +from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +if not logger.handlers: + logger.addHandler(RichHandler()) + +app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) + +# --------------------------------------------------------------------------- +# Plot theme — Paul Tol qualitative (colour-blind safe) +# --------------------------------------------------------------------------- +_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] + +sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) + +DATASET_DISPLAY_NAMES: dict[str, str] = { + "HCT116": "Human colon", + "gluc": "HeLa degradome", + "helaqc": "HeLa single shot", + "herceptin": "Herceptin", + "immuno": "Immunopeptidomics-1", + "celegans": "$\\it{C.\\;elegans}$", + "sbrodae": "$\\it{Scalindua\\;brodae}$", + "PXD019483": "HepG2", + "snakevenoms": "Snake venomics", + "tplantibodies": "Therapeutic nanobodies", + "woundfluids": "Wound exudates", + "PXD004732": "ProteomeTools-1", + "PXD014877": "$\\it{C.\\;elegans}$", + "PXD023064": "Immunopeptidomics-2", + "astral": "Astral $\\it{E.\\;coli}$", + "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", + "Arabidopsis": "$\\it{Arabidopsis\\;thaliana}$", + "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", + "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", + "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", + "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", + "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", +} + +# --------------------------------------------------------------------------- +# Feature group definitions (reduced set: no xcorr, spectral_angle, gap/similarity, edit_distance) +# --------------------------------------------------------------------------- +_EXCLUDED_REDUCED = frozenset( + { + "xcorr", + "spectral_angle", + "complementary_ion_count", + "max_ion_gap", + "edit_distance", + } +) + +BEAM_COLUMNS = ["margin", "median_margin", "entropy", "z-score", "edit_distance"] +TOKEN_COLUMNS = ["min_token_probability", "std_token_probability"] +FRAGMENT_MATCH_COLUMNS = [ + "ion_matches", + "ion_match_intensity", + "complementary_ion_count", + "max_ion_gap", + "spectral_angle", + "xcorr", +] +RETENTION_TIME_COLUMNS = ["irt_error"] +MASS_ERROR_PPM = "mass_error_ppm" +MASS_ERROR_DA = "mass_error_da" + +REDUCED_BEAM_COLUMNS = [c for c in BEAM_COLUMNS if c not in _EXCLUDED_REDUCED] +REDUCED_FRAGMENT_COLUMNS = [ + c for c in FRAGMENT_MATCH_COLUMNS if c not in _EXCLUDED_REDUCED +] + +# Default training matrix columns (train_extra_small_matrix.parquet). +REDUCED_TRAIN_COLUMNS: list[str] = FEATURE_SUBSETS["no_fragment_similarity"]["columns"] + +# Hydra overrides aligned with Makefile ANALYSIS_REDUCED_FEATURE_OVERRIDES (mass_error_da model). +REDUCED_FEATURE_COMPUTE_OVERRIDES: list[str] = [ + "~calibrator.features.mass_error", + "+calibrator.features.mass_error_da._target_=winnow.calibration.calibration_features.MassErrorDaFeature", + "+calibrator.features.mass_error_da.residue_masses=${residue_masses}", +] + + +def _reference_model_columns(model_dir: Path | None) -> list[str] | None: + """Return ``feature_columns`` from a saved calibrator, if present.""" + if model_dir is None: + return None + config_path = model_dir / "config.json" + if not config_path.is_file(): + return None + with open(config_path) as f: + config = json.load(f) + cols = config.get("feature_columns") + return list(cols) if cols else None + + +def _resolve_mass_error_column( + df: pl.DataFrame, + reference_model_dir: Path | None, +) -> str: + """Pick mass-error column present in *df*, preferring the reference model.""" + ref_cols = _reference_model_columns(reference_model_dir) + if MASS_ERROR_DA in df.columns: + return MASS_ERROR_DA + if MASS_ERROR_PPM in df.columns: + if ref_cols and MASS_ERROR_DA in ref_cols: + logger.warning( + "Reference model uses %s but data has %s; using %s for ablations.", + MASS_ERROR_DA, + MASS_ERROR_PPM, + MASS_ERROR_PPM, + ) + return MASS_ERROR_PPM + raise ValueError( + f"No mass error column in data (tried {MASS_ERROR_DA}, {MASS_ERROR_PPM})" + ) + + +def _columns_available(df: pl.DataFrame, columns: list[str]) -> list[str]: + missing = [c for c in columns if c not in df.columns] + if missing: + raise ValueError(f"Missing columns: {missing}. Available: {df.columns}") + return columns + + +def resolve_all_feature_columns( + df: pl.DataFrame, + reference_model_dir: Path | None, +) -> list[str]: + """Full reduced feature set for the 'All features' ablation config.""" + ref_cols = _reference_model_columns(reference_model_dir) + if ref_cols: + cols = ["confidence"] + for col in ref_cols: + if ( + col == MASS_ERROR_DA + and col not in df.columns + and MASS_ERROR_PPM in df.columns + ): + cols.append(MASS_ERROR_PPM) + elif col in df.columns: + cols.append(col) + else: + cols = [c for c in REDUCED_TRAIN_COLUMNS if c in df.columns] + return _columns_available(df, cols) + + +def build_ablation_configs( + df: pl.DataFrame, + reference_model_dir: Path | None, +) -> dict[str, list[str]]: + """Build ablation configs using columns available in *df*.""" + mass_col = _resolve_mass_error_column(df, reference_model_dir) + all_features = resolve_all_feature_columns(df, reference_model_dir) + return { + "Confidence only": ["confidence"], + "Confidence + mass error": ["confidence", mass_col], + "Confidence + iRT error": ["confidence", *RETENTION_TIME_COLUMNS], + "Confidence + token-level": ["confidence", *TOKEN_COLUMNS], + "Confidence + beam search": ["confidence", *REDUCED_BEAM_COLUMNS], + "Confidence + fragment matching": ["confidence", *REDUCED_FRAGMENT_COLUMNS], + "All features": all_features, + } + + +ABLATION_CONFIGS: dict[str, list[str]] = {} + +ABLATION_COLORS: dict[str, str] = {} + + +def _dataset_display_name(key: str) -> str: + """Publication-ready dataset label for plot titles.""" + if key in EVAL_DATASETS: + return str(EVAL_DATASETS[key]["label"]) + return DATASET_DISPLAY_NAMES.get(key, key) + + +def _configure_ablation_colors(config_names: Iterable[str]) -> None: + global ABLATION_COLORS + ABLATION_COLORS = assign_ablation_colors( + ordered_ablation_configs(set(config_names)) + ) + + +# Default training hyperparameters (overridden by --hyperparams-from-model). +TRAIN_HYPERPARAMS = { + "hidden_dims": [128, 64], + "learning_rate": 0.0001, + "weight_decay": 0.0001, + "batch_size": 4096, + "max_epochs": 200, + "n_iter_no_change": 10, + "tol": 1e-4, +} + + +def train_hyperparams_from_model(model_dir: Path) -> dict[str, object]: + """Load MLP training hyperparameters from a saved calibrator ``config.json``.""" + config_path = model_dir / "config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"No config.json at {model_dir}") + with open(config_path) as f: + config = json.load(f) + return { + "hidden_dims": tuple(config["hidden_dims"]), + "dropout": config["dropout"], + "learning_rate": config["learning_rate"], + "weight_decay": config["weight_decay"], + "batch_size": config["batch_size"], + "max_epochs": config["max_epochs"], + "n_iter_no_change": config["n_iter_no_change"], + "tol": config["tol"], + } + + +EVAL_DATASETS = { + "HCT116": { + "label": "Human colon", + "spectra": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.parquet", + "predictions": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.csv", + "koina_mode": "columns", + }, + "Arabidopsis": { + "label": "Arabidopsis", + "spectra": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.parquet", + "predictions": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.csv", + "koina_mode": "columns", + }, + "PXD023064": { + "label": "Immunopeptidomics-2", + "spectra": "held_out_projects/lcfm/PXD023064/", + "predictions": "held_out_projects/lcfm/PXD023064_predictions/PXD023064.csv", + "koina_mode": "columns", + }, +} + +# Residue masses for DatabaseGroundedFDRControl (loaded from config at runtime) +_RESIDUE_MASSES: dict[str, float] | None = None + + +def _get_residue_masses() -> dict[str, float]: + """Load residue masses from the winnow residues config.""" + global _RESIDUE_MASSES + if _RESIDUE_MASSES is None: + import yaml + + config_path = ( + Path(__file__).resolve().parent.parent + / "winnow" + / "configs" + / "residues.yaml" + ) + with open(config_path) as f: + cfg = yaml.safe_load(f) + _RESIDUE_MASSES = cfg["residue_masses"] + return _RESIDUE_MASSES + + +# --------------------------------------------------------------------------- +# Eval feature computation +# --------------------------------------------------------------------------- +def _feature_compute_overrides(reference_model_dir: Path | None) -> list[str]: + """Hydra overrides so eval features match a mass_error_da / reduced-feature model.""" + ref_cols = _reference_model_columns(reference_model_dir) + if ref_cols and MASS_ERROR_DA in ref_cols: + return list(REDUCED_FEATURE_COMPUTE_OVERRIDES) + return [] + + +def _compute_eval_features_for_dataset( + name: str, + spectra_path: str, + predictions_path: str, + cache_dir: Path, + koina_url: str, + koina_ssl: bool, + koina_mode: str = "columns", + feature_overrides: list[str] | None = None, +) -> Path: + """Compute the full feature matrix for an eval dataset and cache as Parquet. + + Args: + koina_mode: ``"columns"`` to read collision_energy / frag_type from + per-row metadata columns, or ``"constants"`` to use fixed values + (CE=27, HCD). + """ + cache_path = cache_dir / f"{name}.parquet" + if cache_path.exists(): + logger.info("Using cached eval features for %s at %s", name, cache_path) + return cache_path + + from hydra import compose, initialize_config_dir + from hydra.utils import instantiate + + from winnow.utils.config_path import get_primary_config_dir + + primary_config_dir = get_primary_config_dir(None) + + logger.info("Computing features for eval dataset %s ...", name) + + if koina_mode == "columns": + koina_overrides = [ + "+koina.input_columns.collision_energies=collision_energy", + "+koina.input_columns.fragmentation_types=frag_type", + "+calibrator.features.fragment_match_features.model_input_columns.collision_energies=collision_energy", + "+calibrator.features.fragment_match_features.model_input_columns.fragmentation_types=frag_type", + ] + else: + koina_overrides = [ + "+koina.input_constants.collision_energies=27", + "+koina.input_constants.fragmentation_types=HCD", + "+calibrator.features.fragment_match_features.model_input_constants.collision_energies=27", + "+calibrator.features.fragment_match_features.model_input_constants.fragmentation_types=HCD", + ] + + with initialize_config_dir( + config_dir=str(primary_config_dir), + version_base="1.3", + job_name=f"winnow_ablation_features_{name}", + ): + cfg = compose( + config_name="compute_features", + overrides=[ + f"dataset.spectrum_path_or_directory={spectra_path}", + f"dataset.predictions_path={predictions_path}", + f"koina.server_url={koina_url}", + f"koina.ssl={koina_ssl}", + *koina_overrides, + *(feature_overrides or []), + "labelled=true", + "filter_empty_predictions=true", + ], + ) + + data_loader = instantiate(cfg.data_loader) + calibrator = instantiate(cfg.calibrator) + + from winnow.scripts.main import ( + _compute_features_batched_metadata, + ) + + spectrum_path = Path(spectra_path) + preds_path = cfg.dataset.get("predictions_path", predictions_path) + + all_metadata = _compute_features_batched_metadata( + spectrum_path, + preds_path, + data_loader, + calibrator, + labelled=True, + ) + + combined_metadata = pd.concat(all_metadata, ignore_index=True) + logger.info( + " %s: %d spectra after feature computation", name, len(combined_metadata) + ) + + # Write the training matrix parquet with all feature columns + correct + extra cols for FDR + feature_columns = ["confidence"] + calibrator.columns + keep_cols = list(feature_columns) + if "correct" in combined_metadata.columns: + keep_cols.append("correct") + if "sequence" in combined_metadata.columns: + keep_cols.append("sequence") + if "prediction" in combined_metadata.columns: + keep_cols.append("prediction") + if "precursor_mz" in combined_metadata.columns: + keep_cols.append("precursor_mz") + if "precursor_charge" in combined_metadata.columns: + keep_cols.append("precursor_charge") + + # Deduplicate while preserving order + seen = set() + unique_cols = [] + for c in keep_cols: + if c not in seen and c in combined_metadata.columns: + seen.add(c) + unique_cols.append(c) + + training_df = pl.from_pandas(combined_metadata[unique_cols]) + cache_dir.mkdir(parents=True, exist_ok=True) + training_df.write_parquet(cache_path) + logger.info( + " Cached eval features to %s (%d rows, %d cols)", + cache_path, + len(training_df), + len(training_df.columns), + ) + return cache_path + + +def compute_all_eval_features( + output_dir: Path, + koina_url: str, + koina_ssl: bool, + astral_spectra: str | None, + astral_predictions: str | None, + skip_feature_compute: bool, + reference_model_dir: Path | None = None, +) -> dict[str, Path]: + """Compute (or locate cached) eval feature Parquets for all datasets.""" + cache_dir = output_dir / "eval_feature_cache" + result: dict[str, Path] = {} + feature_overrides = _feature_compute_overrides(reference_model_dir) + + for name, info in EVAL_DATASETS.items(): + if skip_feature_compute: + cache_path = cache_dir / f"{name}.parquet" + if not cache_path.exists(): + raise FileNotFoundError( + f"--skip-feature-compute set but cache not found: {cache_path}" + ) + result[name] = cache_path + else: + result[name] = _compute_eval_features_for_dataset( + name, + info["spectra"], + info["predictions"], + cache_dir, + koina_url, + koina_ssl, + koina_mode=info.get("koina_mode", "columns"), + feature_overrides=feature_overrides, + ) + + if astral_spectra and astral_predictions: + name = "Astral" + if skip_feature_compute: + cache_path = cache_dir / f"{name}.parquet" + if not cache_path.exists(): + raise FileNotFoundError( + f"--skip-feature-compute set but cache not found: {cache_path}" + ) + result[name] = cache_path + else: + result[name] = _compute_eval_features_for_dataset( + name, + astral_spectra, + astral_predictions, + cache_dir, + koina_url, + koina_ssl, + feature_overrides=feature_overrides, + ) + + return result + + +# --------------------------------------------------------------------------- +# Training +# --------------------------------------------------------------------------- +def _load_parquet_as_polars(path: str | Path) -> pl.DataFrame: + """Load a Parquet file or directory of Parquets into a single Polars DataFrame.""" + path = Path(path) + if path.is_dir(): + parquet_files = sorted(path.glob("*.parquet")) + if not parquet_files: + raise FileNotFoundError(f"No .parquet files in {path}") + return pl.concat([pl.read_parquet(f) for f in parquet_files]) + return pl.read_parquet(path) + + +def split_train_val_frames( + df: pl.DataFrame, + validation_fraction: float, + seed: int, +) -> tuple[pl.DataFrame, pl.DataFrame]: + """Random train/validation split with a fixed index permutation. + + Uses the same scheme as ``winnow.scripts.main._maybe_split_calibration_dataset``: + shuffle all row indices with *seed*, then assign the last ``validation_fraction`` + fraction to validation. The same split is reused for every ablation config. + """ + if "correct" not in df.columns: + raise ValueError("Training Parquet must contain a 'correct' column") + if not 0 < validation_fraction < 1: + raise ValueError( + f"validation_fraction must be in (0, 1), got {validation_fraction}" + ) + + n = len(df) + n_val = max(1, int(n * validation_fraction)) + rng = np.random.default_rng(seed) + perm = rng.permutation(n) + train_df = df[perm[: n - n_val].tolist()] + val_df = df[perm[n - n_val :].tolist()] + logger.info( + "Train/val split: %d train, %d val (fraction=%.2f, seed=%d)", + len(train_df), + len(val_df), + validation_fraction, + seed, + ) + return train_df, val_df + + +def _column_slice_to_feature_dataset( + df: pl.DataFrame, columns: list[str] +) -> FeatureDataset: + """Select columns from a Polars DataFrame and build a FeatureDataset.""" + if "correct" not in df.columns: + raise ValueError("Parquet must contain a 'correct' column") + missing = [c for c in columns if c not in df.columns] + if missing: + raise ValueError( + f"Missing columns in Parquet: {missing}. Available: {df.columns}" + ) + features = df.select(columns).to_numpy().astype(np.float32) + labels = df["correct"].to_numpy().astype(np.float32) + non_confidence = [c for c in columns if c != "confidence"] + return FeatureDataset(features=features, labels=labels, columns=non_confidence) + + +def _config_dir_name(config_name: str) -> str: + """Derive the on-disk directory name for an ablation config.""" + return config_name.lower().replace(" ", "_").replace("+", "and") + + +_LEGACY_DIR_NAMES: dict[str, list[str]] = { + "Confidence only": ["confidence_only"], + "Confidence + mass error": [ + "confidence_and_mass_error", + "confidence_and_mass_error_and_rt", + ], + "Confidence + iRT error": [ + "confidence_and_irt_error", + "confidence_and_mass_error_and_rt", + "confidence_and_fragment_matching", + "prosit", + ], + "Confidence + token-level": [ + "confidence_and_token_level", + "confidence_and_beam_search", + "beam_and_token", + ], + "Confidence + beam search": ["confidence_and_beam_search", "beam_and_token"], + "Confidence + fragment matching": [ + "confidence_and_fragment_matching", + "prosit", + ], + "All features": ["all_features", "full_model"], +} + + +def _resolve_model_dir(output_dir: Path, config_name: str) -> Path: + """Find the model directory for a config, falling back to legacy names.""" + candidates = _LEGACY_DIR_NAMES.get(config_name, [_config_dir_name(config_name)]) + for candidate in candidates: + model_dir = output_dir / "models" / candidate + if model_dir.exists(): + return model_dir + raise FileNotFoundError( + f"No saved model found for '{config_name}'. " + f"Searched: {[str(output_dir / 'models' / c) for c in candidates]}" + ) + + +def train_ablation_models( + train_df: pl.DataFrame, + val_df: pl.DataFrame, + output_dir: Path, + seed: int, + train_hyperparams: dict[str, object] | None = None, +) -> dict[str, ProbabilityCalibrator]: + """Train one calibrator per ablation config, return dict of fitted calibrators.""" + models: dict[str, ProbabilityCalibrator] = {} + hp = {**TRAIN_HYPERPARAMS, **(train_hyperparams or {})} + + for config_name, columns in ABLATION_CONFIGS.items(): + logger.info( + "Training ablation config: %s (%d features)", config_name, len(columns) + ) + + train_ds = _column_slice_to_feature_dataset(train_df, columns) + val_ds = _column_slice_to_feature_dataset(val_df, columns) + + calibrator = ProbabilityCalibrator( + seed=seed, + **hp, # type: ignore[arg-type] + ) + history = calibrator.fit_from_features(train_ds, val_ds) + + model_dir = output_dir / "models" / _config_dir_name(config_name) + ProbabilityCalibrator.save(calibrator, model_dir) + logger.info( + " Trained %s: %d epochs, best_epoch=%d", + config_name, + history.epochs_trained, + history.best_epoch, + ) + + models[config_name] = calibrator + + return models + + +def load_ablation_models( + output_dir: Path, +) -> dict[str, ProbabilityCalibrator]: + """Load pre-trained ablation calibrators from ``{output_dir}/models/``.""" + models: dict[str, ProbabilityCalibrator] = {} + + for config_name in ABLATION_CONFIGS: + model_dir = _resolve_model_dir(output_dir, config_name) + calibrator = ProbabilityCalibrator.load(model_dir) + logger.info(" Loaded %s from %s", config_name, model_dir) + models[config_name] = calibrator + + return models + + +# --------------------------------------------------------------------------- +# Evaluation helpers +# --------------------------------------------------------------------------- +def _predict_calibrated_scores( + calibrator: ProbabilityCalibrator, + features: np.ndarray, +) -> np.ndarray: + """Run forward pass through a fitted calibrator and return calibrated probabilities.""" + assert calibrator.network is not None + assert calibrator.feature_mean is not None + assert calibrator.feature_std is not None + + device = next(calibrator.network.parameters()).device + x = torch.as_tensor(features, dtype=torch.float32, device=device) + x = (x - calibrator.feature_mean) / calibrator.feature_std + + calibrator.network.eval() + with torch.no_grad(): + logits = calibrator.network(x) + probs = torch.sigmoid(logits).cpu().numpy().flatten() + + return probs + + +def compute_precision_recall_curve( + dataset: pd.DataFrame, + confidence_column: str, + label_column: str, + name: str, +) -> pd.DataFrame: + """Non-standard cumulative PR curve matching the casanovo notebook.""" + original = dataset[[confidence_column, label_column]] + original = original.sort_values(by=confidence_column, ascending=False) + cum_correct = np.cumsum(original[label_column].values) + precision = cum_correct / np.arange(1, len(original) + 1) + recall = cum_correct / len(original) + metrics = pd.DataFrame({"precision": precision, "recall": recall}).reset_index( + drop=True + ) + metrics["name"] = name + return metrics + + +def compute_calibration_curve( + df: pd.DataFrame, + pred_col: str, + label_col: str, + name: str, + n_bins: int = 10, +) -> pd.DataFrame: + """Fixed-width bin calibration curve matching the casanovo notebook.""" + data = df[[pred_col, label_col]].dropna().copy(deep=True) + data[pred_col] = data[pred_col].clip(0.0, 1.0) + bins = np.linspace(0.0, 1.0, n_bins + 1) + bin_cats = pd.cut(data[pred_col], bins=bins, include_lowest=True) + bin_cats.name = "bin" + grouped = ( + data.groupby(bin_cats, observed=True) + .agg( + pred_mean=(pred_col, "mean"), + empirical=(label_col, "mean"), + count=(label_col, "size"), + ) + .reset_index() + ) + grouped = grouped[grouped["count"] > 0] + grouped["bin_center"] = grouped["bin"].apply(lambda iv: (iv.left + iv.right) / 2) + grouped["name"] = name + return grouped[["pred_mean", "empirical", "count", "bin_center", "name"]] + + +def compute_brier_score(pred: np.ndarray, labels: np.ndarray) -> float: + """Brier score.""" + return float(np.mean((pred - labels) ** 2)) + + +def compute_ids_at_fdr( + calibrated_scores: np.ndarray, + labels: np.ndarray, + fdr_threshold: float, +) -> int: + """Count PSMs accepted at a given FDR threshold using NonParametricFDRControl.""" + fdr_ctrl = NonParametricFDRControl() + scores_series = pd.Series(calibrated_scores, name="score") + fdr_ctrl.fit(dataset=scores_series) + cutoff = fdr_ctrl.get_confidence_cutoff(threshold=fdr_threshold) + if np.isnan(cutoff): + return 0 + return int((calibrated_scores >= cutoff).sum()) + + +@dataclass +class EvalResult: + """Metrics and curves for a single ablation config evaluated on one dataset.""" + + config_name: str + dataset_name: str + ece: float + tail_ece: float + tail_ece_at_5pct: float + tail_ece_at_10pct: float + brier: float + ids_at_1pct: int + ids_at_5pct: int + ids_at_10pct: int + pr_auc: float + fdr_bias_at_5pct: float + fdr_bias_at_10pct: float + q_dev_at_5pct: float + q_dev_at_10pct: float + pr_curve: pd.DataFrame = field(repr=False) + calibration_curve: pd.DataFrame = field(repr=False) + calibrated_scores: np.ndarray = field(repr=False) + labels: np.ndarray = field(repr=False) + raw_confidence: np.ndarray = field(repr=False) + eval_df: pd.DataFrame = field(repr=False) + + +def evaluate_single( + config_name: str, + calibrator: ProbabilityCalibrator, + columns: list[str], + eval_df: pl.DataFrame, + dataset_name: str, +) -> EvalResult: + """Evaluate a single ablation config on a single eval dataset.""" + features = eval_df.select(columns).to_numpy().astype(np.float32) + labels = eval_df["correct"].to_numpy().astype(np.float32) + raw_confidence = eval_df["confidence"].to_numpy().astype(np.float64) + + calibrated = _predict_calibrated_scores(calibrator, features) + + # Build a pandas DataFrame for PR / calibration / FDR computations + meta = pd.DataFrame( + { + "confidence": raw_confidence, + "calibrated_confidence": calibrated, + "correct": labels, + } + ) + + # Carry over sequence and prediction for database-grounded FDR if available + if "sequence" in eval_df.columns: + meta["sequence"] = eval_df["sequence"].to_pandas() + if "prediction" in eval_df.columns: + meta["prediction"] = eval_df["prediction"].to_pandas() + + pr = compute_precision_recall_curve( + meta, "calibrated_confidence", "correct", config_name + ) + + cal = compute_calibration_curve( + meta, "calibrated_confidence", "correct", config_name + ) + + ece = compute_ece(calibrated, labels) + tail_ece = compute_tail_ece(calibrated, labels) + fdr_ctrl = NonParametricFDRControl() + fdr_ctrl.fit(dataset=pd.Series(calibrated, name="score")) + tail_ece_5 = compute_tail_ece_at_fdr(calibrated, labels, 0.05, fdr_ctrl=fdr_ctrl) + tail_ece_10 = compute_tail_ece_at_fdr(calibrated, labels, 0.10, fdr_ctrl=fdr_ctrl) + brier = compute_brier_score(calibrated, labels) + + ids_1 = compute_ids_at_fdr(calibrated, labels, 0.01) + ids_5 = compute_ids_at_fdr(calibrated, labels, 0.05) + ids_10 = compute_ids_at_fdr(calibrated, labels, 0.10) + + pr_auc = compute_pr_auc(meta) + fdr_bias = compute_fdr_bias_at_fdr_thresholds(meta) + q_dev = compute_q_value_deviations(meta) + + return EvalResult( + config_name=config_name, + dataset_name=dataset_name, + ece=ece, + tail_ece=tail_ece, + tail_ece_at_5pct=tail_ece_5, + tail_ece_at_10pct=tail_ece_10, + brier=brier, + ids_at_1pct=ids_1, + ids_at_5pct=ids_5, + ids_at_10pct=ids_10, + pr_auc=pr_auc, + fdr_bias_at_5pct=fdr_bias[0.05], + fdr_bias_at_10pct=fdr_bias[0.10], + q_dev_at_5pct=q_dev[0.05], + q_dev_at_10pct=q_dev[0.10], + pr_curve=pr, + calibration_curve=cal, + calibrated_scores=calibrated, + labels=labels, + raw_confidence=raw_confidence, + eval_df=meta, + ) + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- +def _style_axes(ax: plt.Axes) -> None: + """Apply standard axes formatting: no grid, black spines.""" + ax.set_axisbelow(True) + ax.grid(False) + for spine in ax.spines.values(): + spine.set_edgecolor("black") + spine.set_linewidth(0.8) + + +def _save_fig(fig: plt.Figure, base_path: Path, plot_format: str) -> None: + """Save figure in the requested format(s).""" + if plot_format in ("pdf", "both"): + fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) + if plot_format in ("png", "both"): + fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) + plt.close(fig) + + +def _lineplot( + ax: plt.Axes, + data: pd.DataFrame, + *, + x: str, + y: str, + label: str, + color: str, + linestyle: str = "-", + linewidth: float = 0.5, + marker: str | None = None, +) -> None: + """Line plot with consistent linewidth (seaborn, no auto legend).""" + kwargs: dict = { + "data": data, + "x": x, + "y": y, + "label": label, + "color": color, + "linestyle": linestyle, + "linewidth": linewidth, + "ax": ax, + "legend": False, + } + if marker is not None: + kwargs["marker"] = marker + sns.lineplot(**kwargs) + + +def _generate_plots_for_dataset( + ds_results: list[EvalResult], + ds_name: str, + plots_dir: Path, + plot_format: str, +) -> None: + """Generate all ablation figures for one dataset.""" + plot_precision_recall(ds_results, ds_name, plots_dir, plot_format) + plot_calibration(ds_results, ds_name, plots_dir, plot_format) + plot_fdr_vs_confidence(ds_results, ds_name, plots_dir, plot_format) + plot_fdr_accepted_psms(ds_results, ds_name, plots_dir, plot_format) + + +def plot_precision_recall( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """PR curve: one line per ablation config.""" + fig, ax = plt.subplots(figsize=(6, 4)) + + for r in results: + _lineplot( + ax, + r.pr_curve, + x="recall", + y="precision", + label=r.config_name, + color=ABLATION_COLORS[r.config_name], + ) + + display = _dataset_display_name(dataset_name) + ax.set( + xlabel="Recall", + ylabel="Precision", + title=f"{display} precision-recall by feature set", + ) + ax.legend(loc="lower left", fontsize=7) + _style_axes(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"pr_curve_{dataset_name}", plot_format) + + +def plot_calibration( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Calibration diagram: reliability curves + diagonal.""" + fig, ax = plt.subplots(figsize=(6, 4)) + + for r in results: + _lineplot( + ax, + r.calibration_curve, + x="pred_mean", + y="empirical", + label=r.config_name, + color=ABLATION_COLORS[r.config_name], + marker="o", + ) + + display = _dataset_display_name(dataset_name) + ax.plot([0, 1], [0, 1], ls="--", color="gray", lw=0.5) + ax.set( + xlabel="Mean predicted probability", + ylabel="Empirical accuracy\n(database label)", + title=f"{display} probability calibration by feature set", + ) + ax.legend(loc="lower right", fontsize=7) + _style_axes(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"calibration_{dataset_name}", plot_format) + + +def plot_fdr_vs_confidence( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """PSM FDR vs calibrated confidence: non-parametric vs database-grounded per config.""" + n_configs = len(results) + fig, axes = plt.subplots(1, n_configs, figsize=(5 * n_configs, 4), squeeze=False) + + for i, r in enumerate(results): + ax = axes[0, i] + + np_fdr = NonParametricFDRControl() + np_fdr.fit(dataset=r.eval_df["calibrated_confidence"]) + winnow_metrics = np_fdr.add_psm_fdr( + r.eval_df.copy(), confidence_col="calibrated_confidence" + ) + + has_sequence = ( + "sequence" in r.eval_df.columns and "prediction" in r.eval_df.columns + ) + + if has_sequence: + dbg_fdr = DatabaseGroundedFDRControl( + confidence_feature="calibrated_confidence", + ) + try: + sorted_df = r.eval_df.sort_values( + "calibrated_confidence", ascending=False + ) + labels = sorted_df["correct"].astype(float).to_numpy() + conf = sorted_df["calibrated_confidence"].to_numpy() + drop = 10 + precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) + dbg_fdr._fdr_values = np.array(1.0 - precision)[drop:] + dbg_fdr._confidence_scores = conf[drop:] + dbg_metrics = dbg_fdr.add_psm_fdr( + r.eval_df.copy(), confidence_col="calibrated_confidence" + ) + + sns.lineplot( + x=np.asarray(dbg_metrics["calibrated_confidence"], dtype=float), + y=np.asarray(dbg_metrics["psm_fdr"], dtype=float), + label="Database-grounded", + ax=ax, + color=_PALETTE[3], + linewidth=0.5, + legend=False, + ) + except Exception as e: + logger.warning( + "Database-grounded FDR failed for %s/%s: %s", + r.config_name, + dataset_name, + e, + ) + + sns.lineplot( + x=np.asarray(winnow_metrics["calibrated_confidence"], dtype=float), + y=np.asarray(winnow_metrics["psm_fdr"], dtype=float), + label="Winnow (non-parametric)", + ax=ax, + color=_PALETTE[0], + linewidth=0.5, + legend=False, + ) + + ax.set_xlabel("Calibrated confidence") + ax.set_ylabel("PSM FDR") + ax.set_title(r.config_name) + ax.legend(fontsize=7) + _style_axes(ax) + + display = _dataset_display_name(dataset_name) + fig.suptitle( + f"{display} PSM FDR vs calibrated confidence by feature set", fontsize=12 + ) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_vs_confidence_{dataset_name}", plot_format) + + +def plot_fdr_accepted_psms( + results: list[EvalResult], + dataset_name: str, + output_dir: Path, + plot_format: str, +) -> None: + """Number of accepted PSMs vs q-value threshold.""" + fig, ax = plt.subplots(figsize=(6, 4)) + + thresholds = np.linspace(0.001, 0.10, 200) + + for r in results: + np_fdr = NonParametricFDRControl() + scores_series = pd.Series(r.calibrated_scores, name="score") + np_fdr.fit(dataset=scores_series) + + meta_with_q = np_fdr.add_psm_q_value( + pd.DataFrame({"calibrated_confidence": r.calibrated_scores}), + confidence_col="calibrated_confidence", + ) + + q_values = meta_with_q["psm_q_value"].values + counts = [] + for t in thresholds: + counts.append(int((q_values <= t).sum())) + + ax.plot( + thresholds, + counts, + label=r.config_name, + color=ABLATION_COLORS[r.config_name], + linewidth=0.5, + ) + + for fdr_line in [0.01, 0.05, 0.10]: + ax.axvline(fdr_line, ls="--", color="gray", lw=0.5, alpha=0.7) + + ax.relim() + ax.autoscale_view() + y_text = ax.get_ylim()[1] * 0.02 + for fdr_line in [0.01, 0.05, 0.10]: + ax.text( + fdr_line - 0.002, + y_text, + f"{fdr_line:.0%}", + ha="right", + va="bottom", + fontsize=7, + color="gray", + ) + + display = _dataset_display_name(dataset_name) + ax.set_xlabel("Non-parametric q-value threshold") + ax.set_ylabel("Accepted PSMs") + ax.set_title(f"{display} accepted PSMs at non-parametric q-value threshold") + ax.legend(loc="upper left", fontsize=7) + _style_axes(ax) + fig.tight_layout() + _save_fig(fig, output_dir / f"fdr_accepted_psms_{dataset_name}", plot_format) + + +# --------------------------------------------------------------------------- +# Saving eval results +# --------------------------------------------------------------------------- +def save_eval_results(all_results: list[EvalResult], output_dir: Path) -> None: + """Persist per-PSM eval DataFrames so plots can be reproduced without re-inference.""" + results_dir = output_dir / "eval_results" + results_dir.mkdir(parents=True, exist_ok=True) + + for r in all_results: + safe_config = r.config_name.lower().replace(" ", "_").replace("+", "and") + path = results_dir / f"{r.dataset_name}_{safe_config}.parquet" + df = r.eval_df.copy() + df["config_name"] = r.config_name + df["dataset_name"] = r.dataset_name + df.to_parquet(path, index=False) + + logger.info("Saved %d eval result Parquets to %s", len(all_results), results_dir) + + +def load_eval_results_for_plotting( + output_dir: Path, +) -> dict[str, list[EvalResult]]: + """Load saved eval Parquets and rebuild curve data for plotting.""" + results_dir = output_dir / "eval_results" + if not results_dir.is_dir(): + raise FileNotFoundError(f"No eval_results directory at {results_dir}") + + paths = sorted(results_dir.glob("*.parquet")) + if not paths: + raise FileNotFoundError(f"No eval result Parquets in {results_dir}") + + grouped: dict[str, list[EvalResult]] = defaultdict(list) + for path in paths: + df = pd.read_parquet(path) + config_name = str(df["config_name"].iloc[0]) + dataset_name = str(df["dataset_name"].iloc[0]) + meta = df.drop(columns=["config_name", "dataset_name"], errors="ignore") + calibrated = meta["calibrated_confidence"].to_numpy(dtype=np.float64) + labels = meta["correct"].to_numpy(dtype=np.float32) + pr = compute_precision_recall_curve( + meta, "calibrated_confidence", "correct", config_name + ) + cal = compute_calibration_curve( + meta, "calibrated_confidence", "correct", config_name + ) + grouped[dataset_name].append( + EvalResult( + config_name=config_name, + dataset_name=dataset_name, + ece=0.0, + tail_ece=compute_tail_ece(calibrated, labels), + tail_ece_at_5pct=float("nan"), + tail_ece_at_10pct=float("nan"), + brier=0.0, + ids_at_1pct=0, + ids_at_5pct=0, + ids_at_10pct=0, + pr_auc=0.0, + fdr_bias_at_5pct=float("nan"), + fdr_bias_at_10pct=float("nan"), + q_dev_at_5pct=float("nan"), + q_dev_at_10pct=float("nan"), + pr_curve=pr, + calibration_curve=cal, + calibrated_scores=calibrated, + labels=labels, + raw_confidence=meta["confidence"].to_numpy(dtype=np.float64), + eval_df=meta, + ) + ) + + for ds_name in grouped: + grouped[ds_name].sort(key=lambda r: r.config_name) + + return dict(grouped) + + +def _run_plots_only(output_dir: Path, plot_format: str) -> None: + """Regenerate plots from ``{output_dir}/eval_results`` without inference.""" + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + grouped = load_eval_results_for_plotting(output_dir) + config_names = [r.config_name for results in grouped.values() for r in results] + _configure_ablation_colors(config_names) + + for ds_name in sorted(grouped): + ds_results = grouped[ds_name] + logger.info("Generating plots for %s (%d configs)...", ds_name, len(ds_results)) + _generate_plots_for_dataset(ds_results, ds_name, plots_dir, plot_format) + + logger.info("Plots saved to %s", plots_dir) + + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- +def build_summary_table(all_results: list[EvalResult]) -> pd.DataFrame: + """Aggregate all EvalResults into a single summary DataFrame.""" + rows = [] + for r in all_results: + rows.append( + { + "config": r.config_name, + "dataset": r.dataset_name, + "ECE": round(r.ece, 5), + "tail_ECE": round(r.tail_ece, 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.05]: round(r.tail_ece_at_5pct, 5), + TAIL_ECE_COLUMN_BY_THRESHOLD[0.10]: round(r.tail_ece_at_10pct, 5), + "Brier": round(r.brier, 5), + "PR_AUC": round(r.pr_auc, 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.05]: round(r.fdr_bias_at_5pct, 5), + FDR_BIAS_COLUMN_BY_THRESHOLD[0.10]: round(r.fdr_bias_at_10pct, 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.05]: round(r.q_dev_at_5pct, 5), + Q_DEV_COLUMN_BY_THRESHOLD[0.10]: round(r.q_dev_at_10pct, 5), + "IDs@1%FDR": r.ids_at_1pct, + "IDs@5%FDR": r.ids_at_5pct, + "IDs@10%FDR": r.ids_at_10pct, + } + ) + return pd.DataFrame(rows) + + +_DEFAULT_OUTPUT_DIR = Path("analysis/hpo_ablation") + + +def _validate_training_inputs( + *, + skip_training: bool, + train_features: Path | None, + val_features: Path | None, + validation_fraction: float | None, +) -> None: + if skip_training: + return + if train_features is None: + raise typer.BadParameter( + "--train-features is required unless --skip-training is set." + ) + if val_features is None and validation_fraction is None: + raise typer.BadParameter( + "Provide --val-features or --validation-fraction when training." + ) + if val_features is not None and validation_fraction is not None: + logger.warning( + "Both --val-features and --validation-fraction set; using --val-features." + ) + + +def _configure_ablation_configs( + *, + skip_training: bool, + train_features: Path | None, + eval_dfs: dict[str, pl.DataFrame], + hyperparams_from_model: Path | None, +) -> None: + global ABLATION_CONFIGS + + if not skip_training: + assert train_features is not None + train_schema_df = _load_parquet_as_polars(train_features) + ABLATION_CONFIGS = build_ablation_configs( + train_schema_df, hyperparams_from_model + ) + else: + first_eval = next(iter(eval_dfs.values())) + ABLATION_CONFIGS = build_ablation_configs(first_eval, hyperparams_from_model) + _configure_ablation_colors(ABLATION_CONFIGS.keys()) + logger.info("Ablation configs: %s", list(ABLATION_CONFIGS.keys())) + + +def _train_or_load_ablation_models( + *, + skip_training: bool, + train_features: Path | None, + val_features: Path | None, + validation_fraction: float | None, + output_dir: Path, + seed: int, + hyperparams_from_model: Path | None, +) -> dict[str, ProbabilityCalibrator]: + if skip_training: + logger.info("Step 3: Loading pre-trained ablation models...") + return load_ablation_models(output_dir) + + logger.info("Step 3: Training ablation models...") + assert train_features is not None + full_train_df = _load_parquet_as_polars(train_features) + if val_features is not None: + train_df = full_train_df + val_df = _load_parquet_as_polars(val_features) + else: + assert validation_fraction is not None + train_df, val_df = split_train_val_frames( + full_train_df, validation_fraction, seed + ) + train_hp = None + if hyperparams_from_model is not None: + train_hp = train_hyperparams_from_model(hyperparams_from_model) + logger.info( + "Using training hyperparameters from %s: %s", + hyperparams_from_model, + train_hp, + ) + return train_ablation_models( + train_df, val_df, output_dir, seed, train_hyperparams=train_hp + ) + + +def _evaluate_ablations( + *, + eval_dfs: dict[str, pl.DataFrame], + models: dict[str, ProbabilityCalibrator], + plots_dir: Path, + plot_format: str, +) -> list[EvalResult]: + logger.info("Step 4: Evaluating ablation models...") + all_results: list[EvalResult] = [] + + for ds_name, ds_df in eval_dfs.items(): + ds_results: list[EvalResult] = [] + for config_name, columns in ABLATION_CONFIGS.items(): + result = evaluate_single( + config_name, models[config_name], columns, ds_df, ds_name + ) + ds_results.append(result) + all_results.append(result) + logger.info( + " %s / %s: ECE=%.4f, Brier=%.4f, IDs@1%%=%d, IDs@5%%=%d, IDs@10%%=%d", + ds_name, + config_name, + result.ece, + result.brier, + result.ids_at_1pct, + result.ids_at_5pct, + result.ids_at_10pct, + ) + + logger.info("Step 5: Generating plots for %s...", ds_name) + _generate_plots_for_dataset(ds_results, ds_name, plots_dir, plot_format) + + return all_results + + +def _write_ablation_summary(output_dir: Path, all_results: list[EvalResult]) -> None: + logger.info("Step 7: Writing summary...") + summary = build_summary_table(all_results) + summary.to_csv(output_dir / "ablation_summary.csv", index=False) + + summary_json = summary.to_dict(orient="records") + with open(output_dir / "ablation_summary.json", "w") as f: + json.dump(summary_json, f, indent=2) + + logger.info("Summary table:\n%s", summary.to_string(index=False)) + + +# --------------------------------------------------------------------------- +# Main CLI +# --------------------------------------------------------------------------- +@app.command() +def main( + train_features: Annotated[ + Optional[Path], + typer.Option( + help="Path to pre-computed training Parquet file or directory. " + "Required unless --skip-training is set.", + ), + ] = None, + val_features: Annotated[ + Optional[Path], + typer.Option( + help="Pre-computed validation Parquet. Omit if using --validation-fraction.", + ), + ] = None, + validation_fraction: Annotated[ + Optional[float], + typer.Option( + "--validation-fraction", + min=0.0, + max=1.0, + help=( + "Hold out this fraction of --train-features for validation " + "(same row split for every ablation model). Alternative to --val-features." + ), + ), + ] = None, + output_dir: Annotated[ + Path, + typer.Option(help="Directory for cached features, models, metrics, and plots."), + ] = _DEFAULT_OUTPUT_DIR, + astral_spectra: Annotated[ + Optional[str], + typer.Option(help="Optional: path to Astral spectra directory."), + ] = None, + astral_predictions: Annotated[ + Optional[str], + typer.Option(help="Optional: path to Astral predictions CSV."), + ] = None, + plot_format: Annotated[ + str, + typer.Option(help="Plot format: 'pdf', 'png', or 'both'."), + ] = "both", + seed: Annotated[ + int, + typer.Option(help="Random seed."), + ] = 42, + koina_url: Annotated[ + str, + typer.Option(help="Koina server URL for eval feature computation."), + ] = "koina.wilhelmlab.org:443", + koina_ssl: Annotated[ + bool, + typer.Option(help="Use SSL for Koina server."), + ] = True, + skip_feature_compute: Annotated[ + bool, + typer.Option( + "--skip-feature-compute", + help="Skip eval feature computation; assume cache exists.", + ), + ] = False, + skip_training: Annotated[ + bool, + typer.Option( + "--skip-training", + help="Load pre-trained ablation models from {output-dir}/models/ " + "instead of training from scratch.", + ), + ] = False, + hyperparams_from_model: Annotated[ + Optional[Path], + typer.Option( + help="Use training hyperparameters from this saved calibrator directory " + "(e.g. HPO best model). Reads config.json.", + ), + ] = None, + plots_only: Annotated[ + bool, + typer.Option( + "--plots-only", + help="Regenerate plots from {output-dir}/eval_results only " + "(no feature compute, training, or evaluation).", + ), + ] = False, +) -> None: + """Run feature ablation study for the Winnow calibrator.""" + output_dir.mkdir(parents=True, exist_ok=True) + + if plots_only: + _run_plots_only(output_dir, plot_format) + logger.info("Feature ablation plots complete.") + return + + _validate_training_inputs( + skip_training=skip_training, + train_features=train_features, + val_features=val_features, + validation_fraction=validation_fraction, + ) + + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + logger.info("Step 1: Computing eval features...") + eval_parquets = compute_all_eval_features( + output_dir, + koina_url, + koina_ssl, + astral_spectra, + astral_predictions, + skip_feature_compute, + reference_model_dir=hyperparams_from_model, + ) + + logger.info("Step 2: Loading Parquets...") + eval_dfs: dict[str, pl.DataFrame] = {} + for name, path in eval_parquets.items(): + eval_dfs[name] = _load_parquet_as_polars(path) + logger.info(" Loaded eval %s: %d rows", name, len(eval_dfs[name])) + + _configure_ablation_configs( + skip_training=skip_training, + train_features=train_features, + eval_dfs=eval_dfs, + hyperparams_from_model=hyperparams_from_model, + ) + models = _train_or_load_ablation_models( + skip_training=skip_training, + train_features=train_features, + val_features=val_features, + validation_fraction=validation_fraction, + output_dir=output_dir, + seed=seed, + hyperparams_from_model=hyperparams_from_model, + ) + all_results = _evaluate_ablations( + eval_dfs=eval_dfs, + models=models, + plots_dir=plots_dir, + plot_format=plot_format, + ) + + logger.info("Step 6: Saving eval results...") + save_eval_results(all_results, output_dir) + _write_ablation_summary(output_dir, all_results) + logger.info("Results saved to %s", output_dir) + logger.info("Feature ablation study complete.") + + +if __name__ == "__main__": + app() diff --git a/paper_scripts/subset_eval_by_experiment.py b/paper_scripts/subset_eval_by_experiment.py new file mode 100644 index 00000000..9178a80d --- /dev/null +++ b/paper_scripts/subset_eval_by_experiment.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Subset general-model eval parquet + preds CSV by experiment name. + +Used to reproduce the revisions-era PXD023064 / immuno2 cohort, which only +kept a fixed list of RAW runs (``PXD023064_FILES`` in ``Makefile.paper``). +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Annotated + +import pandas as pd +import typer + +logger = logging.getLogger(__name__) + +app = typer.Typer( + add_completion=False, + pretty_exceptions_show_locals=False, + no_args_is_help=True, +) + + +def _configure_logging() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + +def _experiment_mask(spectrum_ids: pd.Series, experiments: set[str]) -> pd.Series: + prefixes = spectrum_ids.astype("string").str.split(":", n=1).str[0] + return prefixes.isin(experiments) + + +@app.command() +def main( + spectra: Annotated[ + Path, + typer.Option("--spectra", help="Input spectra parquet."), + ], + preds: Annotated[ + Path, + typer.Option("--preds", help="Input InstantNovo predictions CSV."), + ], + output_dir: Annotated[ + Path, + typer.Option("--output-dir", help="Directory for subset parquet + preds CSV."), + ], + stem: Annotated[ + str, + typer.Option("--stem", help="Output basename stem (e.g. immuno2)."), + ], + experiment: Annotated[ + list[str], + typer.Option( + "--experiment", + help="Keep rows whose spectrum_id prefix matches this experiment. Repeatable.", + ), + ], +) -> None: + """Write ``{stem}.parquet`` and ``{stem}_preds.csv`` restricted to experiments.""" + _configure_logging() + if not experiment: + raise typer.BadParameter("Pass at least one --experiment") + if not spectra.is_file(): + raise FileNotFoundError(f"Missing spectra parquet: {spectra}") + if not preds.is_file(): + raise FileNotFoundError(f"Missing preds CSV: {preds}") + + wanted = set(experiment) + spectra_df = pd.read_parquet(spectra) + preds_df = pd.read_csv(preds) + + if "spectrum_id" not in spectra_df.columns: + raise ValueError(f"{spectra} missing spectrum_id") + if "spectrum_id" not in preds_df.columns: + raise ValueError(f"{preds} missing spectrum_id") + + spectra_mask = _experiment_mask(spectra_df["spectrum_id"], wanted) + preds_mask = _experiment_mask(preds_df["spectrum_id"], wanted) + spectra_out = spectra_df.loc[spectra_mask].reset_index(drop=True) + preds_out = preds_df.loc[preds_mask].reset_index(drop=True) + + present = set( + spectra_out["spectrum_id"].astype("string").str.split(":", n=1).str[0].unique() + ) + missing = sorted(wanted - present) + if missing: + logger.warning("Requested experiments absent from spectra: %s", missing) + + output_dir.mkdir(parents=True, exist_ok=True) + spectra_path = output_dir / f"{stem}.parquet" + preds_path = output_dir / f"{stem}_preds.csv" + spectra_out.to_parquet(spectra_path, index=False) + preds_out.to_csv(preds_path, index=False) + logger.info( + "Wrote %s (%d rows) and %s (%d rows) for experiments %s", + spectra_path, + len(spectra_out), + preds_path, + len(preds_out), + sorted(wanted), + ) + + +if __name__ == "__main__": + app() From a539d9a67ea30a5a91bdc28dea1266816ca4722a Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:51:51 +0100 Subject: [PATCH 08/26] chore: uncomment training matrix --- winnow/configs/compute_features.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/winnow/configs/compute_features.yaml b/winnow/configs/compute_features.yaml index e0445bae..daff80e4 100644 --- a/winnow/configs/compute_features.yaml +++ b/winnow/configs/compute_features.yaml @@ -12,6 +12,6 @@ dataset: metadata_output_path: results/metadata.csv # Optional: write a lean numeric Parquet for model training. -# training_matrix_output_path: results/training_matrix.parquet +training_matrix_output_path: null # If true, the dataset must include ground-truth sequence labels. labelled: true From 1b45beb6d5c6c1069266a15607919efa95d183de Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:52:14 +0100 Subject: [PATCH 09/26] docs: correct InstaNovo misspelling --- docs/cli.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cli.md b/docs/cli.md index b68978a4..60660cf5 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -222,7 +222,7 @@ This produces the `proteome_hit` column consumed by `diagnose-calibration` with **Configuration:** see [Proteome-hit annotation configuration](configuration.md#proteome-hit-annotation-configuration). ```bash -# Annotate an InstantNovo holdout and save a Winnow dataset directory +# Annotate an InstaNovo holdout and save a Winnow dataset directory winnow annotate-proteome-hits \ data_loader=instanovo \ dataset.spectrum_path_or_directory=holdout/spectra.mgf \ From 27ed01f4ad0def6b986228c170bf2bdf5852071d Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:52:37 +0100 Subject: [PATCH 10/26] chore: ignore paper reproduction artefacts --- .gitignore | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.gitignore b/.gitignore index 22b9d6cd..f6bcc7f5 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,8 @@ htmlcov/ .coverage coverage.xml pytest.xml + +# Paper reproduction local downloads / outputs +paper_data/ +paper_results/ +paper_plots/ From f1e2f3724378f6c8acb70bad1c4d62f65ccfa6bd Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:26:23 +0100 Subject: [PATCH 11/26] feat: first Make commands --- Makefile.paper | 740 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 740 insertions(+) create mode 100644 Makefile.paper diff --git a/Makefile.paper b/Makefile.paper new file mode 100644 index 00000000..cfa4ebed --- /dev/null +++ b/Makefile.paper @@ -0,0 +1,740 @@ +# Paper results / plots reproduction suite (Figshare 30147601 v7). +# Usage: make -f Makefile.paper +# +# Prefer replot from deposited artefacts. Missing pieces (feature matrices, +# scaling JSON) are recomputed from HF / original inputs. + +UV ?= uv run +PYTHON := $(UV) python +WINNOW := $(UV) winnow +HF := $(UV) hf +PAPER_SCRIPTS := paper_scripts + +# Keep all PSMs in predict outputs (default CLI filters at 5% FDR). Paper plots +# and deposits need FDR metrics on the full table, not a filtered subset. +PREDICT_FDR_THRESHOLD ?= 1.0 + +# Koina CE / fragmentation inputs (mirrors revisions intent + docs/configuration.md). +# HeLa Single Shot tiles CE=27 + HCD. Most general-model stems null those package +# defaults so predict resolves to metadata columns collision_energy / frag_type +# (see resolve_feature_model_inputs in winnow.utils.koina_intensity_config). +# immuno2 / PXD023064 is the exception: the Figshare deposit used tiled CE=27 + HCD +# (verified by recompute match); override via GENERAL_KOINA_immuno2 below. +KOINA_FRAGMENT_MATCH_CONSTANTS = \ + koina.input_constants.collision_energies=27 \ + koina.input_constants.fragmentation_types=HCD +KOINA_FRAGMENT_MATCH_COLUMNS = \ + koina.input_constants.collision_energies=null \ + koina.input_constants.fragmentation_types=null + +PAPER_DATA_DIR ?= paper_data +PAPER_RESULTS_DIR ?= paper_results +PAPER_PLOTS_DIR ?= paper_plots +HF_DATASETS_DIR ?= $(PAPER_DATA_DIR)/winnow-ms-datasets +HF_MODELS_DIR ?= $(PAPER_DATA_DIR)/models + +FIGSHARE_ARTICLE_ID ?= 30147601 +FIGSHARE_ARTICLE_VERSION ?= 7 +FIGSHARE_HELA_MODELS_ID ?= 32744946 +FIGSHARE_HELA_MODELS_VERSION ?= 2 + +HF_DATASETS_REV ?= 659802319d618a359de5ab90ec6b0195681e94a6 +HF_GENERAL_MODEL_REV ?= e2089330dd59adb9685e5b3d7d61f0cd69a3bbb0 +HF_HELAQC_MODEL_REV ?= d56542b961eac7d896e51bf0716a242fc394ab1f +GLISSADE_GIT_SHA ?= 7c723a2af4a88fda84a6bd4f223b351179bd36da +NOVOBOARD_GIT_SHA ?= a9faab3ef1af06987599c2f01e6ba96072c80172 + +FASTA_HUMAN ?= $(HF_DATASETS_DIR)/fasta/human.fasta +FASTA_CELEGANS ?= $(HF_DATASETS_DIR)/fasta/celegans.fasta +FASTA_ATHALIANA ?= $(HF_DATASETS_DIR)/fasta/athaliana.fasta +FASTA_ARABIDOPSIS ?= $(FASTA_ATHALIANA) +FASTA_ECOLI ?= $(HF_DATASETS_DIR)/fasta/ecoli_zorya.fasta + +GENERAL_EVAL_DIR ?= $(HF_DATASETS_DIR)/general_model_evaluation + +# Leave-one-out calibrator generalisation (full retrain; needs Koina). +GENERALISATION_TRAIN_PARQUET ?= $(HF_DATASETS_DIR)/general_model_training_set/train.parquet +GENERALISATION_TRAIN_PREDS ?= $(HF_DATASETS_DIR)/general_model_training_set/train_preds.csv +GENERALISATION_MODELS_DIR ?= $(PAPER_RESULTS_DIR)/generalisation/models +GENERALISATION_RESULTS_DIR ?= $(PAPER_RESULTS_DIR)/generalisation + +FDR_INPUTS ?= $(PAPER_DATA_DIR)/fdr_benchmark_inputs +NOVOBOARD_ROOT ?= $(FDR_INPUTS)/novoboard +WINNOW_FDR_RESULTS ?= $(FDR_INPUTS)/winnow_results +FDR_MODEL_ROOT ?= $(FDR_INPUTS)/models + +HELAQC_MODEL ?= $(HF_MODELS_DIR)/instanovo_helaqc +GENERAL_MODEL ?= $(HF_MODELS_DIR)/winnow-general-model +CASANOVO_HELAQC_MODEL ?= $(HF_MODELS_DIR)/casanovo_helaqc +PRIMENOVO_HELAQC_MODEL ?= $(HF_MODELS_DIR)/primenovo_helaqc + +HELAQC_DATA ?= $(HF_DATASETS_DIR)/helaqc +FEATURE_MATRIX_DIR ?= $(PAPER_RESULTS_DIR)/models/instanovo_helaqc +SCALING_DUMMY_MODEL ?= $(PAPER_RESULTS_DIR)/scaling/dummy_model +RUNTIME_RESULTS_DIR ?= $(PAPER_RESULTS_DIR)/runtime +FEATURE_IMPORTANCE_RESULTS_DIR ?= $(PAPER_RESULTS_DIR)/feature_importance/PXD014877 +FEATURE_IMPORTANCE_PLOTS_DIR ?= $(PAPER_PLOTS_DIR)/feature_importance/PXD014877 +CELEGANS_LABELLED_DIR ?= $(GENERAL_EVAL_DIR)/celegans/labelled + +# Nine general_results projects. +GENERAL_PROJECTS := \ + PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ + PXD004452/20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 \ + PXD004732 \ + PXD006939/20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2 \ + PXD006939/20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1 \ + PXD013868/01747_C01_P018218_S00_I00_N03_R1 \ + PXD014877 \ + PXD023064 \ + astral + +GENERAL_PROJECTS_FLAG := --projects "$(strip $(GENERAL_PROJECTS))" + +# HF general_model_evaluation/ stems (Make-safe suffixes for per-dataset targets). +GENERAL_STEMS := \ + hela_chymotrypsin \ + human_lung \ + proteometools1 \ + HLA_I \ + HLA_II \ + athaliana \ + celegans \ + immuno2 \ + astral + +# Local labelled recompute default (override with GENERAL_RECOMPUTE_STEMS=...). +GENERAL_RECOMPUTE_STEMS ?= immuno2 hela_chymotrypsin human_lung + +GENERAL_FULL_SMALL_STEMS ?= \ + hela_chymotrypsin \ + human_lung \ + HLA_I \ + HLA_II \ + athaliana \ + immuno2 + +GENERAL_FULL_LARGE_STEMS ?= astral celegans proteometools1 + +# Match revisions Makefile: PXD023064 / immuno2 only used these RAW runs. +PXD023064_FILES := MSB33410B MSB33411A MSB33659A MSB33663B MSB37876A MSB37878A MSB37880A MSB37884A +GENERAL_EXPERIMENT_FILTER_immuno2 := $(PXD023064_FILES) +# Deposit used tiled CE=27 / HCD, not per-row collision_energy / frag_type. +GENERAL_KOINA_immuno2 := $(KOINA_FRAGMENT_MATCH_CONSTANTS) +GENERAL_SUBSET_DIR ?= $(PAPER_RESULTS_DIR)/_subset + +HELAQC_SPLITS := test unlabelled raw_less_train +HELAQC_DNS := instanovo casanovo primenovo + +.PHONY: help \ + download-paper-artefacts download-paper-datasets download-paper-models paper-check-artefacts \ + paper-sync-group-paper \ + paper-plot paper-recompute \ + paper-plot-helaqc-analysis paper-plot-general-labelled paper-plot-general-full \ + paper-plot-feature-investigation paper-plot-scaling \ + paper-plot-fdr-method-comparison paper-plot-external-peptide-holdout \ + paper-plot-ablations paper-plot-generalisation \ + paper-plot-feature-importance \ + paper-plot-upscored-fps paper-plot-fdr-overlap paper-plot-novelty \ + paper-recompute-helaqc paper-recompute-general-labelled \ + paper-recompute-general-full-small paper-recompute-general-full-large \ + paper-recompute-general-print \ + paper-recompute-feature-investigation paper-recompute-scaling \ + paper-recompute-runtime \ + paper-recompute-feature-importance \ + paper-recompute-fdr-method-comparison paper-recompute-external-peptide-holdout \ + paper-recompute-novelty paper-recompute-fdr-overlap paper-recompute-upscored-fps \ + paper-recompute-generalisation \ + paper-smoke-cli paper-smoke-recompute-fdr paper-smoke-recompute-holdout \ + paper-smoke-recompute-compare + +help: + @echo "Paper reproduction (make -f Makefile.paper )" + @echo " download-paper-artefacts | download-paper-datasets | download-paper-models" + @echo " paper-check-artefacts" + @echo " paper-plot / paper-recompute (umbrellas)" + @echo " paper-plot- / paper-recompute- (one analysis at a time)" + @echo " paper-recompute-general-{labelled,full}- | paper-recompute-general-" + @echo " paper-recompute-general-print (echo all predict/annotate commands)" + @echo " stems: $(GENERAL_STEMS)" + @echo " paper-recompute-generalisation (full retrain; requires Koina)" + @echo " paper-recompute-runtime (HeLa full+no-Prosit table; requires Koina)" + @echo " paper-recompute-feature-importance (general model × C. elegans; requires Koina)" + +################################################################################ +# Downloads +################################################################################ + +download-paper-artefacts: + $(PYTHON) $(PAPER_SCRIPTS)/download_figshare_article.py \ + --article-id $(FIGSHARE_ARTICLE_ID) \ + --version $(FIGSHARE_ARTICLE_VERSION) \ + --output-dir $(PAPER_DATA_DIR) + +download-paper-datasets: + mkdir -p $(HF_DATASETS_DIR) + $(HF) download InstaDeepAI/winnow-ms-datasets \ + --repo-type dataset \ + --revision $(HF_DATASETS_REV) \ + --local-dir $(HF_DATASETS_DIR) + +download-paper-models: + mkdir -p $(HF_MODELS_DIR)/instanovo_helaqc $(HF_MODELS_DIR)/winnow-general-model + $(HF) download InstaDeepAI/winnow-helaqc-model \ + --revision $(HF_HELAQC_MODEL_REV) \ + --local-dir $(HELAQC_MODEL) + $(HF) download InstaDeepAI/winnow-general-model \ + --revision $(HF_GENERAL_MODEL_REV) \ + --local-dir $(GENERAL_MODEL) + $(PYTHON) $(PAPER_SCRIPTS)/download_figshare_article.py \ + --article-id $(FIGSHARE_HELA_MODELS_ID) \ + --version $(FIGSHARE_HELA_MODELS_VERSION) \ + --output-dir $(HF_MODELS_DIR) + +# Soft check of key deposited paths (extend --require as needed). +paper-check-artefacts: + $(PYTHON) $(PAPER_SCRIPTS)/download_figshare_article.py \ + --article-id $(FIGSHARE_ARTICLE_ID) \ + --version $(FIGSHARE_ARTICLE_VERSION) \ + --output-dir $(PAPER_DATA_DIR) \ + --check-only \ + --require helaqc_results/instanovo/test/preds_and_fdr_metrics.csv \ + --require general_results/labelled/PXD023064/preds_and_fdr_metrics.csv \ + --require fdr_tool_comparison/fdr_method_comparison_curves.csv \ + --require external_peptide_holdout_benchmark/external_peptide_holdout_results.csv \ + --require ablations/ablation_summary_arabidopsis.csv \ + --require ablations/ablation_summary_astral.csv \ + --require fdr_benchmark_inputs/novoboard/helaqc/novoboard/annotated_test.csv + +paper-sync-group-paper: + uv sync --group paper --group notebook --group docs --group dev + +################################################################################ +# Replot umbrellas / per-group plot targets +################################################################################ + +paper-plot: paper-plot-helaqc-analysis paper-plot-general-labelled paper-plot-general-full \ + paper-plot-fdr-method-comparison paper-plot-external-peptide-holdout \ + paper-plot-generalisation paper-plot-feature-importance \ + paper-plot-upscored-fps paper-plot-fdr-overlap paper-plot-novelty \ + paper-plot-feature-investigation paper-plot-scaling + +# HeLa analysis plots: test set + full search space (raw_less_train) only. +paper-plot-helaqc-analysis: + @test -f $(FASTA_HUMAN) || (echo "Missing $(FASTA_HUMAN); run download-paper-datasets" && exit 1) + @mkdir -p $(PAPER_PLOTS_DIR)/helaqc_analysis + $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/instanovo/test \ + --split test --label-mode labelled --fasta $(FASTA_HUMAN) --dns-model InstaNovo \ + --model-dir $(HELAQC_MODEL) \ + --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/instanovo_test + $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/instanovo/raw_less_train \ + --split raw_less_train --label-mode unlabelled --fasta $(FASTA_HUMAN) --dns-model InstaNovo \ + --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/instanovo_raw_less_train + $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/casanovo/test \ + --split test --label-mode labelled --fasta $(FASTA_HUMAN) --dns-model Casanovo \ + --model-dir $(CASANOVO_HELAQC_MODEL) \ + --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/casanovo_test + $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/casanovo/raw_less_train \ + --split raw_less_train --label-mode unlabelled --fasta $(FASTA_HUMAN) --dns-model Casanovo \ + --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/casanovo_raw_less_train + $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/primenovo/test \ + --split test --label-mode labelled --fasta $(FASTA_HUMAN) --dns-model 'pi-PrimeNovo' \ + --model-dir $(PRIMENOVO_HELAQC_MODEL) \ + --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/primenovo_test + $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/primenovo/raw_less_train \ + --split raw_less_train --label-mode unlabelled --fasta $(FASTA_HUMAN) --dns-model 'pi-PrimeNovo' \ + --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/primenovo_raw_less_train + +paper-plot-general-labelled: + mkdir -p $(PAPER_RESULTS_DIR)/general_labelled $(PAPER_PLOTS_DIR)/general_labelled + $(PYTHON) $(PAPER_SCRIPTS)/plot_eval_results.py \ + --predictions-root $(PAPER_DATA_DIR)/general_results/labelled \ + --eval-type labelled \ + $(GENERAL_PROJECTS_FLAG) \ + --results-dir $(PAPER_RESULTS_DIR)/general_labelled \ + --plots-dir $(PAPER_PLOTS_DIR)/general_labelled + +paper-plot-general-full: + mkdir -p $(PAPER_RESULTS_DIR)/general_full $(PAPER_PLOTS_DIR)/general_full + $(PYTHON) $(PAPER_SCRIPTS)/plot_eval_results.py \ + --predictions-root $(PAPER_DATA_DIR)/general_results/full \ + --eval-type unlabelled \ + $(GENERAL_PROJECTS_FLAG) \ + --results-dir $(PAPER_RESULTS_DIR)/general_full \ + --plots-dir $(PAPER_PLOTS_DIR)/general_full + +# No Figshare feature matrices → recompute then plot. +paper-plot-feature-investigation: paper-recompute-feature-investigation + mkdir -p $(PAPER_PLOTS_DIR)/feature_investigation + $(PYTHON) $(PAPER_SCRIPTS)/plot_feature_investigation.py \ + --features-train $(FEATURE_MATRIX_DIR)/features_train.parquet \ + --features-val $(FEATURE_MATRIX_DIR)/features_val.parquet \ + --metadata-train $(FEATURE_MATRIX_DIR)/metadata_train.parquet \ + --metadata-val $(FEATURE_MATRIX_DIR)/metadata_val.parquet \ + --plots-dir $(PAPER_PLOTS_DIR)/feature_investigation + +# No deposited scaling JSON → live recompute/plot. +paper-plot-scaling: paper-recompute-scaling + +paper-plot-fdr-method-comparison: + mkdir -p $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + $(PAPER_PLOTS_DIR)/fdr_method_comparison + $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ + --summarise-only $(PAPER_DATA_DIR)/fdr_tool_comparison/fdr_method_comparison_curves.csv \ + --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison + +paper-plot-external-peptide-holdout: + mkdir -p $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + $(PAPER_PLOTS_DIR)/external_peptide_holdout + $(PYTHON) $(PAPER_SCRIPTS)/run_external_peptide_holdout_benchmark.py \ + --summarise-only $(PAPER_DATA_DIR)/external_peptide_holdout_benchmark/external_peptide_holdout_results.csv \ + --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout + +paper-plot-ablations: + @echo "Cannot replot the original ablation bars from the deposited summaries." + @echo "They contain reviewer metrics at 5%/10% FDR, but not top-decile tail_ECE." + @echo "Recompute the ablation evaluation outputs, then run plot_ablation_summary.py." + @exit 1 + +paper-plot-generalisation: + mkdir -p $(PAPER_PLOTS_DIR)/generalisation + $(PYTHON) $(PAPER_SCRIPTS)/plot_calibrator_generalisation_heatmap.py \ + --results-path $(PAPER_DATA_DIR)/generalisation/calibrator_generalisation_results.csv \ + --plots-dir $(PAPER_PLOTS_DIR)/generalisation + +# Replot general-model feature importance for PXD014877 (C. elegans) from Figshare pickles. +# Skips SHAP bar / correlation matrix (need feature matrices not in the deposit). +paper-plot-feature-importance: + @test -f $(PAPER_DATA_DIR)/feature_importance/PXD014877/perm_importance.pkl || \ + (echo "Missing feature_importance pickles; run download-paper-artefacts" && exit 1) + @test -f $(PAPER_DATA_DIR)/feature_importance/PXD014877/shap_values.pkl || \ + (echo "Missing feature_importance pickles; run download-paper-artefacts" && exit 1) + mkdir -p $(FEATURE_IMPORTANCE_PLOTS_DIR) + $(PYTHON) $(PAPER_SCRIPTS)/analyze_features.py \ + --replot-dir $(PAPER_DATA_DIR)/feature_importance/PXD014877 \ + --output-dir $(FEATURE_IMPORTANCE_PLOTS_DIR) + +# Recompute general-model feature importance on labelled C. elegans (PXD014877). +# Heavy to recompute +paper-recompute-feature-importance: + @echo "Recomputing feature importance (Koina required)." + @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) + @test -f $(CELEGANS_LABELLED_DIR)/celegans.parquet || \ + (echo "Missing $(CELEGANS_LABELLED_DIR)/celegans.parquet; run download-paper-datasets" && exit 1) + @test -f $(CELEGANS_LABELLED_DIR)/celegans_preds.csv || \ + (echo "Missing $(CELEGANS_LABELLED_DIR)/celegans_preds.csv; run download-paper-datasets" && exit 1) + mkdir -p $(FEATURE_IMPORTANCE_RESULTS_DIR) $(FEATURE_IMPORTANCE_PLOTS_DIR) + $(PYTHON) $(PAPER_SCRIPTS)/analyze_features.py \ + --model-path $(GENERAL_MODEL) \ + --data-dir $(CELEGANS_LABELLED_DIR) \ + --train-spectra celegans.parquet \ + --train-preds celegans_preds.csv \ + --test-spectra celegans.parquet \ + --test-preds celegans_preds.csv \ + --koina-input-constant collision_energies=27 \ + --koina-input-constant fragmentation_types=HCD \ + --n-background-samples 200 \ + --n-test-samples 500 \ + --output-dir $(FEATURE_IMPORTANCE_RESULTS_DIR) + @echo "Pickles+plots under $(FEATURE_IMPORTANCE_RESULTS_DIR)." + +# Full retrain of leave-one-source-out calibrators on HF general_model_training_set. +# Heavy (Koina + nine train/eval loops). +# Outputs: $(GENERALISATION_RESULTS_DIR)/calibrator_generalisation_results.csv and +# $(GENERALISATION_MODELS_DIR)/trained_on_/. +paper-recompute-generalisation: + @echo "Retraining leave-one-out calibrators (Koina required; not Figshare summaries alone)." + @test -f $(GENERALISATION_TRAIN_PARQUET) || \ + (echo "Missing $(GENERALISATION_TRAIN_PARQUET); run download-paper-datasets" && exit 1) + @test -f $(GENERALISATION_TRAIN_PREDS) || \ + (echo "Missing $(GENERALISATION_TRAIN_PREDS); run download-paper-datasets" && exit 1) + mkdir -p $(GENERALISATION_MODELS_DIR) $(GENERALISATION_RESULTS_DIR) + $(PYTHON) $(PAPER_SCRIPTS)/evaluate_calibrator_generalisation.py \ + --train-parquet $(GENERALISATION_TRAIN_PARQUET) \ + --train-predictions $(GENERALISATION_TRAIN_PREDS) \ + --model-output-dir $(GENERALISATION_MODELS_DIR) \ + --results-output-dir $(GENERALISATION_RESULTS_DIR) + +# Figshare analysis CSVs alone cannot rebuild these figures; recompute from +# deposited general_results PSM trees, writing tables and plots separately. +paper-plot-upscored-fps: paper-recompute-upscored-fps + +paper-plot-fdr-overlap: paper-recompute-fdr-overlap + +# Novelty plots need a full recompute (deposited novelty trees are CSV-only). +paper-plot-novelty: paper-recompute-novelty + +paper-recompute-novelty: + @echo "Recomputing novelty tables+plots from deposited general_results (not Figshare novelty CSVs alone)." + mkdir -p $(PAPER_RESULTS_DIR)/novelty/chymotrypsin \ + $(PAPER_RESULTS_DIR)/novelty/proteometools \ + $(PAPER_PLOTS_DIR)/novelty/chymotrypsin \ + $(PAPER_PLOTS_DIR)/novelty/proteometools + @test -f $(FASTA_HUMAN) || (echo "Missing $(FASTA_HUMAN); run download-paper-datasets" && exit 1) + $(PYTHON) $(PAPER_SCRIPTS)/analyze_novelty.py nontryptic-digest \ + --predictions-dir $(PAPER_DATA_DIR)/general_results/full/PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ + --fasta $(FASTA_HUMAN) \ + --results-dir $(PAPER_RESULTS_DIR)/novelty/chymotrypsin \ + --plots-dir $(PAPER_PLOTS_DIR)/novelty/chymotrypsin \ + --file-prefix chymotrypsin \ + --dataset-label "HeLa chymotrypsin" + $(PYTHON) $(PAPER_SCRIPTS)/analyze_novelty.py proteometools \ + --lcfm-predictions-dir $(PAPER_DATA_DIR)/general_results/labelled/PXD004732 \ + --acfm-predictions-dir $(PAPER_DATA_DIR)/general_results/full/PXD004732 \ + --results-dir $(PAPER_RESULTS_DIR)/novelty/proteometools \ + --plots-dir $(PAPER_PLOTS_DIR)/novelty/proteometools + +# Leaf keys under general_results/{labelled,full}/ (nested runs use the run folder name). +# Astral / PXD014877 are slow; subset for a smoke pass, e.g.: +# FDR_OVERLAP_PROJECTS='$(FDR_OVERLAP_immuno2)' make -f Makefile.paper paper-recompute-fdr-overlap +FDR_OVERLAP_immuno2 := PXD023064 +FDR_OVERLAP_PROJECTS ?= \ + 20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ + 20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 \ + PXD004732 \ + 20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2 \ + 20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1 \ + 01747_C01_P018218_S00_I00_N03_R1 \ + PXD014877 \ + PXD023064 \ + astral + +paper-recompute-fdr-overlap: + @echo "Recomputing FDR overlap from deposited general_results (not Figshare overlap CSVs alone)." + @echo "Projects: $(strip $(FDR_OVERLAP_PROJECTS))" + @echo "Note: astral and PXD014877 are slow; option to subset with FDR_OVERLAP_PROJECTS='$$(FDR_OVERLAP_immuno2)'." + mkdir -p $(PAPER_RESULTS_DIR)/fdr_overlap $(PAPER_PLOTS_DIR)/fdr_overlap + $(PYTHON) $(PAPER_SCRIPTS)/analyze_fdr_overlap.py \ + --labelled-dir $(PAPER_DATA_DIR)/general_results/labelled \ + --unlabelled-dir $(PAPER_DATA_DIR)/general_results/full \ + --results-dir $(PAPER_RESULTS_DIR)/fdr_overlap \ + --plots-dir $(PAPER_PLOTS_DIR)/fdr_overlap \ + --projects "$(strip $(FDR_OVERLAP_PROJECTS))" + +paper-recompute-upscored-fps: + @echo "Recomputing upscored-FP tables+plots from deposited general_results (not Figshare upscored CSVs alone)." + mkdir -p $(PAPER_RESULTS_DIR)/upscored_fps $(PAPER_PLOTS_DIR)/upscored_fps + $(PYTHON) $(PAPER_SCRIPTS)/analyze_upscored_fps.py \ + --predictions-root $(PAPER_DATA_DIR)/general_results/labelled \ + --results-dir $(PAPER_RESULTS_DIR)/upscored_fps \ + --plots-dir $(PAPER_PLOTS_DIR)/upscored_fps + +################################################################################ +# Recompute +################################################################################ + +paper-recompute: paper-recompute-helaqc paper-recompute-general-labelled \ + paper-recompute-feature-investigation paper-recompute-scaling \ + paper-recompute-fdr-method-comparison paper-recompute-external-peptide-holdout + @echo "Skipping general-full-large (cluster-only). Optional: paper-recompute-general-full-small" + +# InstaNovo × 3 splits by default; set HELAQC_DNS_TOOLS=instanovo casanovo primenovo for all. +HELAQC_DNS_TOOLS ?= instanovo + +paper-recompute-helaqc: + @test -d $(HELAQC_DATA) || (echo "Missing $(HELAQC_DATA); run download-paper-datasets" && exit 1) + @test -d $(HELAQC_MODEL) || (echo "Missing $(HELAQC_MODEL); run download-paper-models" && exit 1) + @mkdir -p $(PAPER_RESULTS_DIR) + @for tool in $(HELAQC_DNS_TOOLS); do \ + for split in $(HELAQC_SPLITS); do \ + out=$(PAPER_RESULTS_DIR)/$${tool}_helaqc_predictions_$${split}; \ + echo "Predict $${tool} $${split} -> $${out}"; \ + case $${tool} in \ + instanovo) loader=instanovo; model=$(HELAQC_MODEL); preds=$(HELAQC_DATA)/instanovo/$${split}_preds.csv ;; \ + casanovo) loader=mztab; model=$(CASANOVO_HELAQC_MODEL); preds=$(HELAQC_DATA)/casanovo/$${split}_preds.csv ;; \ + primenovo) loader=primenovo; model=$(PRIMENOVO_HELAQC_MODEL); preds=$(HELAQC_DATA)/primenovo/$${split}_preds.csv ;; \ + *) echo "Unknown tool $${tool}"; exit 1 ;; \ + esac; \ + spectra=$(HELAQC_DATA)/$${split}.parquet; \ + $(WINNOW) predict \ + data_loader=$${loader} \ + dataset.spectrum_path_or_directory=$${spectra} \ + dataset.predictions_path=$${preds} \ + calibrator.pretrained_model_name_or_path=$${model} \ + fdr_control.fdr_threshold=$(PREDICT_FDR_THRESHOLD) \ + $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ + output_folder=$${out}; \ + done; \ + done + +paper-recompute-general-labelled: + @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) + @test -d $(GENERAL_EVAL_DIR) || (echo "Missing $(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) + @mkdir -p $(PAPER_RESULTS_DIR)/general_results/labelled + @for stem in $(GENERAL_RECOMPUTE_STEMS); do \ + $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-labelled-$${stem}; \ + done + +paper-recompute-general-full-small: + @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) + @test -d $(GENERAL_EVAL_DIR) || (echo "Missing $(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) + @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full + @for stem in $(GENERAL_FULL_SMALL_STEMS); do \ + $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ + done + +paper-recompute-general-full-large: + @echo "CLUSTER ONLY: $(GENERAL_FULL_LARGE_STEMS) — not part of paper-recompute umbrella." + @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) + @test -d $(GENERAL_EVAL_DIR) || (echo "Missing $(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) + @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full + @for stem in $(GENERAL_FULL_LARGE_STEMS); do \ + $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ + done + +# Figshare project key and proteome FASTA per HF stem. +GENERAL_PROJECT_hela_chymotrypsin := PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin +GENERAL_FASTA_hela_chymotrypsin := $(FASTA_HUMAN) +GENERAL_PROJECT_human_lung := PXD004452/20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 +GENERAL_FASTA_human_lung := $(FASTA_HUMAN) +GENERAL_PROJECT_proteometools1 := PXD004732 +GENERAL_FASTA_proteometools1 := $(FASTA_HUMAN) +GENERAL_PROJECT_HLA_I := PXD006939/20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2 +GENERAL_FASTA_HLA_I := $(FASTA_HUMAN) +GENERAL_PROJECT_HLA_II := PXD006939/20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1 +GENERAL_FASTA_HLA_II := $(FASTA_HUMAN) +GENERAL_PROJECT_athaliana := PXD013868/01747_C01_P018218_S00_I00_N03_R1 +GENERAL_FASTA_athaliana := $(FASTA_ATHALIANA) +GENERAL_PROJECT_celegans := PXD014877 +GENERAL_FASTA_celegans := $(FASTA_CELEGANS) +GENERAL_PROJECT_immuno2 := PXD023064 +GENERAL_FASTA_immuno2 := $(FASTA_HUMAN) +GENERAL_PROJECT_astral := astral +GENERAL_FASTA_astral := $(FASTA_ECOLI) + +# If GENERAL_EXPERIMENT_FILTER_ is set, materialise a filtered parquet/CSV +# under GENERAL_SUBSET_DIR (revisions-era PXD023064 cohort for immuno2). +# Expanded via $(call ...) from stem recipes (single $$ → shell $). +define GENERAL_MAYBE_SUBSET +filter="$(strip $(GENERAL_EXPERIMENT_FILTER_$(1)))"; \ +spectra_src=$(GENERAL_EVAL_DIR)/$(1)/$(2)/$(1).parquet; \ +preds_src=$(GENERAL_EVAL_DIR)/$(1)/$(2)/$(1)_preds.csv; \ +test -f "$$spectra_src" || (echo "Missing $$spectra_src" && exit 1); \ +test -f "$$preds_src" || (echo "Missing $$preds_src" && exit 1); \ +if [ -n "$$filter" ]; then \ + subset_dir=$(GENERAL_SUBSET_DIR)/$(1)/$(2); \ + mkdir -p "$$subset_dir"; \ + echo "Subset $(1)/$(2) to experiments: $$filter"; \ + exp_args=""; \ + for exp in $$filter; do exp_args="$$exp_args --experiment $$exp"; done; \ + $(PYTHON) $(PAPER_SCRIPTS)/subset_eval_by_experiment.py \ + --spectra "$$spectra_src" \ + --preds "$$preds_src" \ + --output-dir "$$subset_dir" \ + --stem $(1) \ + $$exp_args; \ + spectra="$$subset_dir/$(1).parquet"; \ + preds="$$subset_dir/$(1)_preds.csv"; \ +else \ + spectra="$$spectra_src"; \ + preds="$$preds_src"; \ +fi +endef + +# Explicit per-stem rules (empty .PHONY + pattern recipes do not mix). +define GENERAL_LABELLED_STEM_RULE +.PHONY: paper-recompute-general-labelled-$(1) +paper-recompute-general-labelled-$(1): + @test -d $$(GENERAL_MODEL) || (echo "Missing $$(GENERAL_MODEL); run download-paper-models" && exit 1) + @test -d $$(GENERAL_EVAL_DIR) || (echo "Missing $$(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) + @$$(call GENERAL_MAYBE_SUBSET,$(1),labelled); \ + out=$$(PAPER_RESULTS_DIR)/general_results/labelled/$$(GENERAL_PROJECT_$(1)); \ + mkdir -p "$$$$out"; \ + echo "Predict labelled $(1) -> $$$$out"; \ + $$(WINNOW) predict \ + data_loader=instanovo \ + dataset.spectrum_path_or_directory=$$$$spectra \ + dataset.predictions_path=$$$$preds \ + calibrator.pretrained_model_name_or_path=$$(GENERAL_MODEL) \ + fdr_control.fdr_threshold=$$(PREDICT_FDR_THRESHOLD) \ + $$(if $$(GENERAL_KOINA_$(1)),$$(GENERAL_KOINA_$(1)),$$(KOINA_FRAGMENT_MATCH_COLUMNS)) \ + output_folder=$$$$out +endef + +define GENERAL_FULL_STEM_RULE +.PHONY: paper-recompute-general-full-$(1) +paper-recompute-general-full-$(1): + @test -d $$(GENERAL_MODEL) || (echo "Missing $$(GENERAL_MODEL); run download-paper-models" && exit 1) + @test -d $$(GENERAL_EVAL_DIR) || (echo "Missing $$(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) + @$$(call GENERAL_MAYBE_SUBSET,$(1),full); \ + out=$$(PAPER_RESULTS_DIR)/general_results/full/$$(GENERAL_PROJECT_$(1)); \ + fasta=$$(GENERAL_FASTA_$(1)); \ + test -f $$$$fasta || (echo "Missing $$$$fasta; run download-paper-datasets" && exit 1); \ + mkdir -p "$$$$out"; \ + echo "Predict full $(1) -> $$$$out"; \ + $$(WINNOW) predict \ + data_loader=instanovo \ + dataset.spectrum_path_or_directory=$$$$spectra \ + dataset.predictions_path=$$$$preds \ + calibrator.pretrained_model_name_or_path=$$(GENERAL_MODEL) \ + fdr_control.fdr_threshold=$$(PREDICT_FDR_THRESHOLD) \ + $$(if $$(GENERAL_KOINA_$(1)),$$(GENERAL_KOINA_$(1)),$$(KOINA_FRAGMENT_MATCH_COLUMNS)) \ + output_folder=$$$$out; \ + echo "Annotate proteome hits $$(GENERAL_PROJECT_$(1)) (fasta=$$$$fasta)"; \ + $$(PYTHON) $$(PAPER_SCRIPTS)/annotate_preds_proteome_hits.py \ + "$$(GENERAL_PROJECT_$(1))" \ + --predictions-root $$(PAPER_RESULTS_DIR)/general_results/full \ + --fasta $$$$fasta +endef + +define GENERAL_BOTH_STEM_RULE +.PHONY: paper-recompute-general-$(1) +paper-recompute-general-$(1): paper-recompute-general-labelled-$(1) paper-recompute-general-full-$(1) +endef + +$(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_LABELLED_STEM_RULE,$(s)))) +$(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_FULL_STEM_RULE,$(s)))) +$(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_BOTH_STEM_RULE,$(s)))) + +# Echo exact recipes for all nine stems × labelled/full (safe on a laptop). +paper-recompute-general-print: + @echo "# General-model recompute commands (pins from Makefile.paper)" + @echo "# Requires: download-paper-datasets download-paper-models; Koina network for predict" + @echo "" + @$(foreach s,$(GENERAL_STEMS), \ + echo "# --- $(s) (labelled) ---"; \ + echo "$(WINNOW) predict \\"; \ + echo " data_loader=instanovo \\"; \ + echo " dataset.spectrum_path_or_directory=$(GENERAL_EVAL_DIR)/$(s)/labelled/$(s).parquet \\"; \ + echo " dataset.predictions_path=$(GENERAL_EVAL_DIR)/$(s)/labelled/$(s)_preds.csv \\"; \ + echo " calibrator.pretrained_model_name_or_path=$(GENERAL_MODEL) \\"; \ + echo " fdr_control.fdr_threshold=$(PREDICT_FDR_THRESHOLD) \\"; \ + echo " $(if $(GENERAL_KOINA_$(s)),$(GENERAL_KOINA_$(s)),$(KOINA_FRAGMENT_MATCH_COLUMNS)) \\"; \ + echo " output_folder=$(PAPER_RESULTS_DIR)/general_results/labelled/$(GENERAL_PROJECT_$(s))"; \ + echo ""; \ + echo "# --- $(s) (full: predict then post-annotate) ---"; \ + echo "$(WINNOW) predict \\"; \ + echo " data_loader=instanovo \\"; \ + echo " dataset.spectrum_path_or_directory=$(GENERAL_EVAL_DIR)/$(s)/full/$(s).parquet \\"; \ + echo " dataset.predictions_path=$(GENERAL_EVAL_DIR)/$(s)/full/$(s)_preds.csv \\"; \ + echo " calibrator.pretrained_model_name_or_path=$(GENERAL_MODEL) \\"; \ + echo " fdr_control.fdr_threshold=$(PREDICT_FDR_THRESHOLD) \\"; \ + echo " $(if $(GENERAL_KOINA_$(s)),$(GENERAL_KOINA_$(s)),$(KOINA_FRAGMENT_MATCH_COLUMNS)) \\"; \ + echo " output_folder=$(PAPER_RESULTS_DIR)/general_results/full/$(GENERAL_PROJECT_$(s))"; \ + echo "$(PYTHON) $(PAPER_SCRIPTS)/annotate_preds_proteome_hits.py \\"; \ + echo " $(GENERAL_PROJECT_$(s)) \\"; \ + echo " --predictions-root $(PAPER_RESULTS_DIR)/general_results/full \\"; \ + echo " --fasta $(GENERAL_FASTA_$(s))"; \ + echo ""; \ + ) + +# Feature investigation needs the InstaNovo HeLa feature stack (beam + token + +# mass_error_da). Package compute-features has no pretrained_model_name_or_path; +# narrow the lean training matrix to the checkpoint feature_columns schema. +HELAQC_FEATURE_INVESTIGATION_OVERRIDES = \ + +calibrator.features.beam_features._target_=winnow.calibration.calibration_features.BeamFeatures \ + +calibrator.features.token_score_features._target_=winnow.calibration.calibration_features.TokenScoreFeatures \ + calibrator.features.retention_time_feature.min_train_points=5 \ + 'calibrator.training_feature_columns=[ion_matches,ion_match_intensity,irt_error,margin,median_margin,entropy,z-score,min_token_probability,std_token_probability,mass_error_da]' + +paper-recompute-feature-investigation: + @echo "Recomputing feature matrices (not deposited on Figshare)." + @test -d $(HELAQC_DATA) || (echo "Missing $(HELAQC_DATA); run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/train.parquet || (echo "Missing $(HELAQC_DATA)/train.parquet" && exit 1) + @test -f $(HELAQC_DATA)/val.parquet || (echo "Missing $(HELAQC_DATA)/val.parquet" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/train_preds.csv || (echo "Missing InstaNovo train preds" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/val_preds.csv || (echo "Missing InstaNovo val preds" && exit 1) + mkdir -p $(FEATURE_MATRIX_DIR) + $(WINNOW) compute-features \ + data_loader=instanovo \ + labelled=true \ + dataset.spectrum_path_or_directory=$(HELAQC_DATA)/train.parquet \ + dataset.predictions_path=$(HELAQC_DATA)/instanovo/train_preds.csv \ + $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ + $(HELAQC_FEATURE_INVESTIGATION_OVERRIDES) \ + metadata_output_path=$(FEATURE_MATRIX_DIR)/metadata_train.parquet \ + training_matrix_output_path=$(FEATURE_MATRIX_DIR)/features_train.parquet + $(WINNOW) compute-features \ + data_loader=instanovo \ + labelled=true \ + dataset.spectrum_path_or_directory=$(HELAQC_DATA)/val.parquet \ + dataset.predictions_path=$(HELAQC_DATA)/instanovo/val_preds.csv \ + $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ + $(HELAQC_FEATURE_INVESTIGATION_OVERRIDES) \ + metadata_output_path=$(FEATURE_MATRIX_DIR)/metadata_val.parquet \ + training_matrix_output_path=$(FEATURE_MATRIX_DIR)/features_val.parquet + +paper-recompute-scaling: + @echo "Recomputing scaling benchmark (trains a no-Prosit dummy; no Figshare JSON)." + @test -f $(HELAQC_DATA)/train.parquet || (echo "Missing helaqc train parquet; run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/val.parquet || (echo "Missing helaqc val parquet; run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/raw_less_train.parquet || (echo "Missing helaqc raw_less_train parquet; run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/train_preds.csv || (echo "Missing InstaNovo train preds" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/val_preds.csv || (echo "Missing InstaNovo val preds" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv || (echo "Missing InstaNovo raw_less_train preds" && exit 1) + mkdir -p $(PAPER_RESULTS_DIR)/scaling $(PAPER_PLOTS_DIR)/scaling + $(PYTHON) $(PAPER_SCRIPTS)/benchmark_scaling.py \ + --spectrum-path $(HELAQC_DATA)/train.parquet \ + --predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ + --spectrum-path $(HELAQC_DATA)/raw_less_train.parquet \ + --predictions-path $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv \ + --train-spectrum-path $(HELAQC_DATA)/train.parquet \ + --train-predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ + --val-spectrum-path $(HELAQC_DATA)/val.parquet \ + --val-predictions-path $(HELAQC_DATA)/instanovo/val_preds.csv \ + --model-output-dir $(SCALING_DUMMY_MODEL) \ + --data-loader instanovo \ + --fractions 0.1 --fractions 0.5 --fractions 1.0 \ + --results-dir $(PAPER_RESULTS_DIR)/scaling \ + --plots-dir $(PAPER_PLOTS_DIR)/scaling + +# Stage-wise runtime table (full Prosit + no-Prosit) on the full HeLa search space. +# Prints the table to stdout and writes JSON + a text copy under paper_results/runtime/. +# Not part of paper-recompute (Koina-heavy). Requires download-paper-datasets/models. +# Reuses or trains the no-Prosit dummy under SCALING_DUMMY_MODEL (same as scaling). +paper-recompute-runtime: + @echo "Recomputing HeLa runtime table (full + no-Prosit; requires Koina for full)." + @test -d $(HELAQC_MODEL) || (echo "Missing $(HELAQC_MODEL); run download-paper-models" && exit 1) + @test -f $(HELAQC_DATA)/train.parquet || (echo "Missing helaqc train parquet; run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/val.parquet || (echo "Missing helaqc val parquet; run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/raw_less_train.parquet || (echo "Missing helaqc raw_less_train parquet; run download-paper-datasets" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/train_preds.csv || (echo "Missing InstaNovo train preds" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/val_preds.csv || (echo "Missing InstaNovo val preds" && exit 1) + @test -f $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv || (echo "Missing InstaNovo raw_less_train preds" && exit 1) + mkdir -p $(RUNTIME_RESULTS_DIR) $(SCALING_DUMMY_MODEL) + $(PYTHON) $(PAPER_SCRIPTS)/benchmark_runtime.py \ + --spectrum-path $(HELAQC_DATA)/train.parquet \ + --predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ + --spectrum-path $(HELAQC_DATA)/raw_less_train.parquet \ + --predictions-path $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv \ + --model-path $(HELAQC_MODEL) \ + --model-path-no-prosit $(SCALING_DUMMY_MODEL) \ + --train-spectrum-path $(HELAQC_DATA)/train.parquet \ + --train-predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ + --val-spectrum-path $(HELAQC_DATA)/val.parquet \ + --val-predictions-path $(HELAQC_DATA)/instanovo/val_preds.csv \ + --data-loader instanovo \ + --koina-input-constant collision_energies=27 \ + --koina-input-constant fragmentation_types=HCD \ + --output-json $(RUNTIME_RESULTS_DIR)/benchmark_results.json \ + --output-text $(RUNTIME_RESULTS_DIR)/benchmark_results.txt + +paper-recompute-fdr-method-comparison: + @echo "Recomputing FDR method comparison from fdr_benchmark_inputs (not Figshare curves alone)." + mkdir -p $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + $(PAPER_PLOTS_DIR)/fdr_method_comparison + $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ + --novoboard-root $(NOVOBOARD_ROOT) \ + --winnow-results $(WINNOW_FDR_RESULTS) \ + --datasets helaqc --datasets celegans \ + --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison + +paper-recompute-external-peptide-holdout: paper-sync-group-paper + @echo "Recomputing external peptide holdout from fdr_benchmark_inputs (not Figshare results CSV alone)." + mkdir -p $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + $(PAPER_PLOTS_DIR)/external_peptide_holdout + $(PYTHON) $(PAPER_SCRIPTS)/run_external_peptide_holdout_benchmark.py \ + --novoboard-root $(NOVOBOARD_ROOT) \ + --winnow-results $(WINNOW_FDR_RESULTS) \ + --model-root $(FDR_MODEL_ROOT) \ + --datasets helaqc --datasets celegans \ + --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout From 87ff0173c8b716e53a61f66140d92b876406ed58 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:59:11 +0100 Subject: [PATCH 12/26] chore: remove unused dataset names and add checks for missing files --- paper_scripts/analyze_fdr_overlap.py | 1 - paper_scripts/analyze_novelty.py | 9 ++ paper_scripts/analyze_upscored_fps.py | 1 - paper_scripts/plot_ablation_summary.py | 1 - paper_scripts/plot_analysis.py | 5 + paper_scripts/plot_eval_results.py | 35 ++++++- paper_scripts/run_feature_ablations.py | 121 +++++++++++++++++++------ 7 files changed, 137 insertions(+), 36 deletions(-) diff --git a/paper_scripts/analyze_fdr_overlap.py b/paper_scripts/analyze_fdr_overlap.py index fd433c7b..151cc3a1 100644 --- a/paper_scripts/analyze_fdr_overlap.py +++ b/paper_scripts/analyze_fdr_overlap.py @@ -96,7 +96,6 @@ "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", "PXD004732": "ProteomeTools-1", diff --git a/paper_scripts/analyze_novelty.py b/paper_scripts/analyze_novelty.py index b366d65d..2bfdfd25 100644 --- a/paper_scripts/analyze_novelty.py +++ b/paper_scripts/analyze_novelty.py @@ -1112,6 +1112,15 @@ def _nontryptic_digest_analysis( plots_dir.mkdir(parents=True, exist_ok=True) print(f"Loading predictions from {predictions_dir}") + if not (predictions_dir / "preds_and_fdr_metrics.csv").is_file(): + raise FileNotFoundError( + f"Missing preds_and_fdr_metrics.csv under {predictions_dir}" + ) + if not fasta.is_file(): + raise FileNotFoundError( + f"Missing FASTA {fasta}; run download-paper-datasets " + "(or set --fasta to an existing proteome)." + ) df_pl = _load_data(predictions_dir) print(f" {df_pl.height:,} rows loaded") diff --git a/paper_scripts/analyze_upscored_fps.py b/paper_scripts/analyze_upscored_fps.py index 3ea56dbf..77faecce 100644 --- a/paper_scripts/analyze_upscored_fps.py +++ b/paper_scripts/analyze_upscored_fps.py @@ -88,7 +88,6 @@ "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", } diff --git a/paper_scripts/plot_ablation_summary.py b/paper_scripts/plot_ablation_summary.py index 20626b52..7494ec1b 100644 --- a/paper_scripts/plot_ablation_summary.py +++ b/paper_scripts/plot_ablation_summary.py @@ -52,7 +52,6 @@ _ABLATION_DATASET_KEYS: dict[str, str] = { "Arabidopsis": "01747_C01_P018218_S00_I00_N03_R1", "Astral": "astral", - "HCT116": "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46", } logger = logging.getLogger(__name__) diff --git a/paper_scripts/plot_analysis.py b/paper_scripts/plot_analysis.py index 1b290531..67f9b9bd 100644 --- a/paper_scripts/plot_analysis.py +++ b/paper_scripts/plot_analysis.py @@ -1150,6 +1150,11 @@ def main( ) logger.info("Annotating with proteome hits from %s", fasta) + if not fasta.is_file(): + raise FileNotFoundError( + f"Missing FASTA {fasta}; run download-paper-datasets " + "(or set --fasta to an existing proteome)." + ) haystack = load_proteome_haystack(str(fasta)) df = filter_and_annotate_preds(df, haystack, metrics, min_residue_length=7) df_raw_conf = _df_for_raw_confidence_plots(df) diff --git a/paper_scripts/plot_eval_results.py b/paper_scripts/plot_eval_results.py index 3addab87..8da257a6 100644 --- a/paper_scripts/plot_eval_results.py +++ b/paper_scripts/plot_eval_results.py @@ -55,7 +55,6 @@ "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", } @@ -833,6 +832,37 @@ def _compute_diagnostics( # --------------------------------------------------------------------------- # Orchestration # --------------------------------------------------------------------------- +_PXD_ACCESSION_PREFIX = "PXD" + + +def _resolve_project_folder(predictions_root: Path, project: str) -> Path: + """Resolve a project key to a preds folder under ``predictions_root``. + + Accepts nested keys (``PXD004452/``), flat leaf folders + (``PXD004732``), or leaf run names under ``PXD*//``. + """ + direct = predictions_root / project + if (direct / "preds_and_fdr_metrics.csv").is_file(): + return direct + + leaf = project.rsplit("/", 1)[-1] + leaf_direct = predictions_root / leaf + if (leaf_direct / "preds_and_fdr_metrics.csv").is_file(): + return leaf_direct + + if predictions_root.is_dir(): + for child in sorted(predictions_root.iterdir()): + if not child.is_dir() or not child.name.startswith(_PXD_ACCESSION_PREFIX): + continue + candidate = child / leaf + if (candidate / "preds_and_fdr_metrics.csv").is_file(): + return candidate + + raise FileNotFoundError( + f"No preds folder for project {project!r} under {predictions_root}" + ) + + def _load_project_data( predictions_root: Path, project: str, @@ -840,8 +870,7 @@ def _load_project_data( eval_type: str, ) -> pd.DataFrame: """Load and merge metadata.csv and preds_and_fdr_metrics.csv for a project.""" - # folder = predictions_root / f"{project}_{suffix}" - folder = predictions_root / f"{project}" + folder = _resolve_project_folder(predictions_root, project) preds_path = folder / "preds_and_fdr_metrics.csv" meta_path = folder / "metadata.csv" if not preds_path.is_file(): diff --git a/paper_scripts/run_feature_ablations.py b/paper_scripts/run_feature_ablations.py index 9d4640f7..37bf606d 100644 --- a/paper_scripts/run_feature_ablations.py +++ b/paper_scripts/run_feature_ablations.py @@ -68,11 +68,11 @@ sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) DATASET_DISPLAY_NAMES: dict[str, str] = { - "HCT116": "Human colon", "gluc": "HeLa degradome", "helaqc": "HeLa single shot", "herceptin": "Herceptin", "immuno": "Immunopeptidomics-1", + "immuno2": "Immunopeptidomics-2", "celegans": "$\\it{C.\\;elegans}$", "sbrodae": "$\\it{Scalindua\\;brodae}$", "PXD019483": "HepG2", @@ -85,9 +85,9 @@ "astral": "Astral $\\it{E.\\;coli}$", "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", "Arabidopsis": "$\\it{Arabidopsis\\;thaliana}$", + "athaliana": "$\\it{Arabidopsis\\;thaliana}$", "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", } @@ -268,26 +268,20 @@ def train_hyperparams_from_model(model_dir: Path) -> dict[str, object]: EVAL_DATASETS = { - "HCT116": { - "label": "Human colon", - "spectra": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.parquet", - "predictions": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.csv", - "koina_mode": "columns", - }, - "Arabidopsis": { + "athaliana": { "label": "Arabidopsis", - "spectra": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.parquet", - "predictions": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.csv", + "stem": "athaliana", "koina_mode": "columns", }, - "PXD023064": { + "immuno2": { "label": "Immunopeptidomics-2", - "spectra": "held_out_projects/lcfm/PXD023064/", - "predictions": "held_out_projects/lcfm/PXD023064_predictions/PXD023064.csv", - "koina_mode": "columns", + "stem": "immuno2", + "koina_mode": "constants", }, } +_DEFAULT_EVAL_ROOT = Path("paper_data/winnow-ms-datasets/general_model_evaluation") + # Residue masses for DatabaseGroundedFDRControl (loaded from config at runtime) _RESIDUE_MASSES: dict[str, float] | None = None @@ -443,6 +437,39 @@ def _compute_eval_features_for_dataset( return cache_path +def _resolve_eval_dataset_paths( + name: str, + info: dict[str, str], + eval_root: Path, + overrides: dict[str, Path], +) -> tuple[str, str]: + """Return ``(spectra_path, predictions_path)`` for an eval dataset.""" + stem = info["stem"] + if name in overrides: + labelled_dir = overrides[name] + spectra = labelled_dir / f"{stem}.parquet" + preds = labelled_dir / f"{stem}_preds.csv" + else: + labelled_dir = eval_root / stem / "labelled" + spectra = labelled_dir / f"{stem}.parquet" + preds = labelled_dir / f"{stem}_preds.csv" + if not spectra.exists(): + raise FileNotFoundError(f"Missing eval spectra for {name}: {spectra}") + if not preds.is_file(): + raise FileNotFoundError(f"Missing eval predictions for {name}: {preds}") + return str(spectra), str(preds) + + +def _require_cached_eval_features(cache_dir: Path, name: str) -> Path: + """Return the cached eval Parquet, or raise if ``--skip-feature-compute``.""" + cache_path = cache_dir / f"{name}.parquet" + if not cache_path.exists(): + raise FileNotFoundError( + f"--skip-feature-compute set but cache not found: {cache_path}" + ) + return cache_path + + def compute_all_eval_features( output_dir: Path, koina_url: str, @@ -451,25 +478,27 @@ def compute_all_eval_features( astral_predictions: str | None, skip_feature_compute: bool, reference_model_dir: Path | None = None, + eval_root: Path | None = None, + eval_dir_overrides: dict[str, Path] | None = None, ) -> dict[str, Path]: """Compute (or locate cached) eval feature Parquets for all datasets.""" cache_dir = output_dir / "eval_feature_cache" result: dict[str, Path] = {} feature_overrides = _feature_compute_overrides(reference_model_dir) + root = eval_root if eval_root is not None else _DEFAULT_EVAL_ROOT + overrides = eval_dir_overrides or {} for name, info in EVAL_DATASETS.items(): if skip_feature_compute: - cache_path = cache_dir / f"{name}.parquet" - if not cache_path.exists(): - raise FileNotFoundError( - f"--skip-feature-compute set but cache not found: {cache_path}" - ) - result[name] = cache_path + result[name] = _require_cached_eval_features(cache_dir, name) else: + spectra_path, predictions_path = _resolve_eval_dataset_paths( + name, info, root, overrides + ) result[name] = _compute_eval_features_for_dataset( name, - info["spectra"], - info["predictions"], + spectra_path, + predictions_path, cache_dir, koina_url, koina_ssl, @@ -480,12 +509,7 @@ def compute_all_eval_features( if astral_spectra and astral_predictions: name = "Astral" if skip_feature_compute: - cache_path = cache_dir / f"{name}.parquet" - if not cache_path.exists(): - raise FileNotFoundError( - f"--skip-feature-compute set but cache not found: {cache_path}" - ) - result[name] = cache_path + result[name] = _require_cached_eval_features(cache_dir, name) else: result[name] = _compute_eval_features_for_dataset( name, @@ -1462,7 +1486,7 @@ def main( Optional[Path], typer.Option( help="Use training hyperparameters from this saved calibrator directory " - "(e.g. HPO best model). Reads config.json.", + "Reads config.json.", ), ] = None, plots_only: Annotated[ @@ -1473,6 +1497,26 @@ def main( "(no feature compute, training, or evaluation).", ), ] = False, + eval_root: Annotated[ + Path, + typer.Option( + "--eval-root", + help=( + "Root of HF general_model_evaluation trees " + "(each stem under //labelled/)." + ), + ), + ] = _DEFAULT_EVAL_ROOT, + override_eval_dir: Annotated[ + Optional[list[str]], + typer.Option( + "--override-eval-dir", + help=( + "Override labelled dir for a dataset as name=/path/to/labelled " + "(repeatable). Path must contain .parquet and _preds.csv." + ), + ), + ] = None, ) -> None: """Run feature ablation study for the Winnow calibrator.""" output_dir.mkdir(parents=True, exist_ok=True) @@ -1492,6 +1536,21 @@ def main( plots_dir = output_dir / "plots" plots_dir.mkdir(parents=True, exist_ok=True) + eval_dir_overrides: dict[str, Path] = {} + for item in override_eval_dir or []: + if "=" not in item: + raise typer.BadParameter( + f"--override-eval-dir expects name=path, got {item!r}" + ) + name, path_str = item.split("=", 1) + name = name.strip() + if name not in EVAL_DATASETS: + raise typer.BadParameter( + f"Unknown eval dataset {name!r}; expected one of " + f"{sorted(EVAL_DATASETS)}" + ) + eval_dir_overrides[name] = Path(path_str.strip()) + logger.info("Step 1: Computing eval features...") eval_parquets = compute_all_eval_features( output_dir, @@ -1501,6 +1560,8 @@ def main( astral_predictions, skip_feature_compute, reference_model_dir=hyperparams_from_model, + eval_root=eval_root, + eval_dir_overrides=eval_dir_overrides, ) logger.info("Step 2: Loading Parquets...") From c860c42568f22707b21fa194a7c658cb24c4148e Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:59:41 +0100 Subject: [PATCH 13/26] refactor: organise paper repro commands around analyses --- Makefile.paper | 748 +++++++++++++++++++++++++------------------------ 1 file changed, 376 insertions(+), 372 deletions(-) diff --git a/Makefile.paper b/Makefile.paper index cfa4ebed..9b732b5f 100644 --- a/Makefile.paper +++ b/Makefile.paper @@ -1,8 +1,9 @@ # Paper results / plots reproduction suite (Figshare 30147601 v7). # Usage: make -f Makefile.paper # -# Prefer replot from deposited artefacts. Missing pieces (feature matrices, -# scaling JSON) are recomputed from HF / original inputs. +# Naming: paper-plot-* = Figshare deposits (or a deterministic local transform). +# paper-recompute-* = redo predict / train / benchmark. +# See paper_scripts/README.md for the full guide. UV ?= uv run PYTHON := $(UV) python @@ -52,7 +53,7 @@ FASTA_ECOLI ?= $(HF_DATASETS_DIR)/fasta/ecoli_zorya.fasta GENERAL_EVAL_DIR ?= $(HF_DATASETS_DIR)/general_model_evaluation -# Leave-one-out calibrator generalisation (full retrain; needs Koina). +# Leave-one-out calibrator generalisation (full retrain; needs Koina + GPU). GENERALISATION_TRAIN_PARQUET ?= $(HF_DATASETS_DIR)/general_model_training_set/train.parquet GENERALISATION_TRAIN_PREDS ?= $(HF_DATASETS_DIR)/general_model_training_set/train_preds.csv GENERALISATION_MODELS_DIR ?= $(PAPER_RESULTS_DIR)/generalisation/models @@ -74,16 +75,18 @@ SCALING_DUMMY_MODEL ?= $(PAPER_RESULTS_DIR)/scaling/dummy_model RUNTIME_RESULTS_DIR ?= $(PAPER_RESULTS_DIR)/runtime FEATURE_IMPORTANCE_RESULTS_DIR ?= $(PAPER_RESULTS_DIR)/feature_importance/PXD014877 FEATURE_IMPORTANCE_PLOTS_DIR ?= $(PAPER_PLOTS_DIR)/feature_importance/PXD014877 +ABLATIONS_RESULTS_DIR ?= $(PAPER_RESULTS_DIR)/ablations CELEGANS_LABELLED_DIR ?= $(GENERAL_EVAL_DIR)/celegans/labelled -# Nine general_results projects. +# Leaf folder names under general_results/{labelled,full}/ (nested runs use the +# run folder name). Used by plot_eval_results and FDR overlap. GENERAL_PROJECTS := \ - PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ - PXD004452/20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 \ + 20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ + 20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 \ PXD004732 \ - PXD006939/20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2 \ - PXD006939/20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1 \ - PXD013868/01747_C01_P018218_S00_I00_N03_R1 \ + 20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2 \ + 20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1 \ + 01747_C01_P018218_S00_I00_N03_R1 \ PXD014877 \ PXD023064 \ astral @@ -124,17 +127,33 @@ GENERAL_SUBSET_DIR ?= $(PAPER_RESULTS_DIR)/_subset HELAQC_SPLITS := test unlabelled raw_less_train HELAQC_DNS := instanovo casanovo primenovo +# Default: all three DNS tools (matches paper-plot-helaqc-analysis). +HELAQC_DNS_TOOLS ?= instanovo casanovo primenovo + +HELAQC_DNS_LABEL_instanovo := InstaNovo +HELAQC_DNS_MODEL_instanovo := $(HELAQC_MODEL) +HELAQC_DNS_LABEL_casanovo := Casanovo +HELAQC_DNS_MODEL_casanovo := $(CASANOVO_HELAQC_MODEL) +HELAQC_DNS_LABEL_primenovo := pi-PrimeNovo +HELAQC_DNS_MODEL_primenovo := $(PRIMENOVO_HELAQC_MODEL) + +# FDR tool recompute datasets. Laptop keeps helaqc only (celegans holdout >5 min). +FDR_RECOMPUTE_DATASETS_LAPTOP ?= helaqc +FDR_RECOMPUTE_DATASETS_HEAVY ?= helaqc celegans +FDR_RECOMPUTE_DATASETS ?= $(FDR_RECOMPUTE_DATASETS_LAPTOP) + +# FDR overlap: default is all projects. Light umbrella uses immuno2 only. +FDR_OVERLAP_immuno2 := PXD023064 +FDR_OVERLAP_PROJECTS ?= $(GENERAL_PROJECTS) .PHONY: help \ download-paper-artefacts download-paper-datasets download-paper-models paper-check-artefacts \ paper-sync-group-paper \ - paper-plot paper-recompute \ + paper-setup paper-plot-light paper-recompute-laptop paper-recompute-heavy \ paper-plot-helaqc-analysis paper-plot-general-labelled paper-plot-general-full \ - paper-plot-feature-investigation paper-plot-scaling \ paper-plot-fdr-method-comparison paper-plot-external-peptide-holdout \ - paper-plot-ablations paper-plot-generalisation \ - paper-plot-feature-importance \ - paper-plot-upscored-fps paper-plot-fdr-overlap paper-plot-novelty \ + paper-plot-generalisation paper-plot-feature-importance \ + paper-plot-novelty paper-plot-upscored-fps paper-plot-fdr-overlap \ paper-recompute-helaqc paper-recompute-general-labelled \ paper-recompute-general-full-small paper-recompute-general-full-large \ paper-recompute-general-print \ @@ -142,26 +161,36 @@ HELAQC_DNS := instanovo casanovo primenovo paper-recompute-runtime \ paper-recompute-feature-importance \ paper-recompute-fdr-method-comparison paper-recompute-external-peptide-holdout \ - paper-recompute-novelty paper-recompute-fdr-overlap paper-recompute-upscored-fps \ - paper-recompute-generalisation \ - paper-smoke-cli paper-smoke-recompute-fdr paper-smoke-recompute-holdout \ - paper-smoke-recompute-compare + paper-recompute-generalisation paper-recompute-ablations help: @echo "Paper reproduction (make -f Makefile.paper )" - @echo " download-paper-artefacts | download-paper-datasets | download-paper-models" - @echo " paper-check-artefacts" - @echo " paper-plot / paper-recompute (umbrellas)" - @echo " paper-plot- / paper-recompute- (one analysis at a time)" - @echo " paper-recompute-general-{labelled,full}- | paper-recompute-general-" - @echo " paper-recompute-general-print (echo all predict/annotate commands)" + @echo "See paper_scripts/README.md for the full guide." + @echo "" + @echo "Umbrellas:" + @echo " paper-setup download Figshare + HF datasets + models, then check" + @echo " paper-plot-light replot from deposits" + @echo " paper-recompute-laptop no GPU / no large datasets" + @echo " paper-recompute-heavy large datasets, may include GPU usage" + @echo "" + @echo "Downloads: download-paper-artefacts | download-paper-datasets | download-paper-models" + @echo " paper-check-artefacts | paper-sync-group-paper" + @echo "" + @echo "Plot (from deposits): paper-plot-" + @echo " helaqc-analysis, general-labelled, general-full (slow)," + @echo " fdr-method-comparison, external-peptide-holdout, generalisation," + @echo " feature-importance, novelty, upscored-fps, fdr-overlap" + @echo "" + @echo "Recompute (predict / Koina / train / benchmark): paper-recompute-" + @echo " helaqc | general-labelled | general-full-small | general-full-large" + @echo " general-{labelled,full}- | general- | general-print" + @echo " feature-investigation | fdr-method-comparison | external-peptide-holdout" + @echo " feature-importance | generalisation [GPU] | ablations [GPU]" + @echo " runtime [GPU] | scaling [GPU]" @echo " stems: $(GENERAL_STEMS)" - @echo " paper-recompute-generalisation (full retrain; requires Koina)" - @echo " paper-recompute-runtime (HeLa full+no-Prosit table; requires Koina)" - @echo " paper-recompute-feature-importance (general model × C. elegans; requires Koina)" ################################################################################ -# Downloads +# Downloads / setup / umbrellas ################################################################################ download-paper-artefacts: @@ -208,68 +237,62 @@ paper-check-artefacts: paper-sync-group-paper: uv sync --group paper --group notebook --group docs --group dev -################################################################################ -# Replot umbrellas / per-group plot targets -################################################################################ +paper-setup: download-paper-artefacts download-paper-datasets download-paper-models \ + paper-check-artefacts -paper-plot: paper-plot-helaqc-analysis paper-plot-general-labelled paper-plot-general-full \ - paper-plot-fdr-method-comparison paper-plot-external-peptide-holdout \ - paper-plot-generalisation paper-plot-feature-importance \ - paper-plot-upscored-fps paper-plot-fdr-overlap paper-plot-novelty \ - paper-plot-feature-investigation paper-plot-scaling - -# HeLa analysis plots: test set + full search space (raw_less_train) only. -paper-plot-helaqc-analysis: - @test -f $(FASTA_HUMAN) || (echo "Missing $(FASTA_HUMAN); run download-paper-datasets" && exit 1) - @mkdir -p $(PAPER_PLOTS_DIR)/helaqc_analysis - $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ - --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/instanovo/test \ - --split test --label-mode labelled --fasta $(FASTA_HUMAN) --dns-model InstaNovo \ - --model-dir $(HELAQC_MODEL) \ - --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/instanovo_test - $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ - --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/instanovo/raw_less_train \ - --split raw_less_train --label-mode unlabelled --fasta $(FASTA_HUMAN) --dns-model InstaNovo \ - --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/instanovo_raw_less_train - $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ - --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/casanovo/test \ - --split test --label-mode labelled --fasta $(FASTA_HUMAN) --dns-model Casanovo \ - --model-dir $(CASANOVO_HELAQC_MODEL) \ - --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/casanovo_test - $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ - --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/casanovo/raw_less_train \ - --split raw_less_train --label-mode unlabelled --fasta $(FASTA_HUMAN) --dns-model Casanovo \ - --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/casanovo_raw_less_train - $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ - --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/primenovo/test \ - --split test --label-mode labelled --fasta $(FASTA_HUMAN) --dns-model 'pi-PrimeNovo' \ - --model-dir $(PRIMENOVO_HELAQC_MODEL) \ - --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/primenovo_test - $(PYTHON) $(PAPER_SCRIPTS)/plot_analysis.py \ - --predictions-dir $(PAPER_DATA_DIR)/helaqc_results/primenovo/raw_less_train \ - --split raw_less_train --label-mode unlabelled --fasta $(FASTA_HUMAN) --dns-model 'pi-PrimeNovo' \ - --plots-dir $(PAPER_PLOTS_DIR)/helaqc_analysis/primenovo_raw_less_train +paper-plot-light: paper-plot-helaqc-analysis paper-plot-general-labelled \ + paper-plot-fdr-method-comparison \ + paper-plot-external-peptide-holdout paper-plot-generalisation \ + paper-plot-feature-importance \ + paper-plot-novelty paper-plot-upscored-fps + $(MAKE) -f Makefile.paper --no-print-directory \ + FDR_OVERLAP_PROJECTS='$(FDR_OVERLAP_immuno2)' paper-plot-fdr-overlap + +paper-recompute-laptop: paper-recompute-helaqc \ + paper-recompute-general-labelled \ + paper-recompute-feature-investigation \ + paper-recompute-fdr-method-comparison \ + paper-recompute-external-peptide-holdout + +paper-recompute-heavy: paper-recompute-general-full-small \ + paper-recompute-general-full-large \ + paper-recompute-feature-importance paper-recompute-generalisation \ + paper-recompute-ablations paper-recompute-runtime paper-recompute-scaling \ + paper-plot-general-full paper-plot-fdr-overlap + $(MAKE) -f Makefile.paper --no-print-directory \ + FDR_RECOMPUTE_DATASETS='$(FDR_RECOMPUTE_DATASETS_HEAVY)' \ + paper-recompute-fdr-method-comparison paper-recompute-external-peptide-holdout -paper-plot-general-labelled: - mkdir -p $(PAPER_RESULTS_DIR)/general_labelled $(PAPER_PLOTS_DIR)/general_labelled - $(PYTHON) $(PAPER_SCRIPTS)/plot_eval_results.py \ - --predictions-root $(PAPER_DATA_DIR)/general_results/labelled \ - --eval-type labelled \ - $(GENERAL_PROJECTS_FLAG) \ - --results-dir $(PAPER_RESULTS_DIR)/general_labelled \ - --plots-dir $(PAPER_PLOTS_DIR)/general_labelled +################################################################################ +# 1. Feature investigation +################################################################################ -paper-plot-general-full: - mkdir -p $(PAPER_RESULTS_DIR)/general_full $(PAPER_PLOTS_DIR)/general_full - $(PYTHON) $(PAPER_SCRIPTS)/plot_eval_results.py \ - --predictions-root $(PAPER_DATA_DIR)/general_results/full \ - --eval-type unlabelled \ - $(GENERAL_PROJECTS_FLAG) \ - --results-dir $(PAPER_RESULTS_DIR)/general_full \ - --plots-dir $(PAPER_PLOTS_DIR)/general_full +# Feature investigation needs the InstaNovo HeLa feature stack. +HELAQC_FEATURE_INVESTIGATION_OVERRIDES = \ + +calibrator.features.beam_features._target_=winnow.calibration.calibration_features.BeamFeatures \ + +calibrator.features.token_score_features._target_=winnow.calibration.calibration_features.TokenScoreFeatures \ + calibrator.features.retention_time_feature.min_train_points=5 -# No Figshare feature matrices → recompute then plot. -paper-plot-feature-investigation: paper-recompute-feature-investigation +paper-recompute-feature-investigation: + mkdir -p $(FEATURE_MATRIX_DIR) + $(WINNOW) compute-features \ + data_loader=instanovo \ + labelled=true \ + dataset.spectrum_path_or_directory=$(HELAQC_DATA)/train.parquet \ + dataset.predictions_path=$(HELAQC_DATA)/instanovo/train_preds.csv \ + $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ + $(HELAQC_FEATURE_INVESTIGATION_OVERRIDES) \ + metadata_output_path=$(FEATURE_MATRIX_DIR)/metadata_train.parquet \ + training_matrix_output_path=$(FEATURE_MATRIX_DIR)/features_train.parquet + $(WINNOW) compute-features \ + data_loader=instanovo \ + labelled=true \ + dataset.spectrum_path_or_directory=$(HELAQC_DATA)/val.parquet \ + dataset.predictions_path=$(HELAQC_DATA)/instanovo/val_preds.csv \ + $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ + $(HELAQC_FEATURE_INVESTIGATION_OVERRIDES) \ + metadata_output_path=$(FEATURE_MATRIX_DIR)/metadata_val.parquet \ + training_matrix_output_path=$(FEATURE_MATRIX_DIR)/features_val.parquet mkdir -p $(PAPER_PLOTS_DIR)/feature_investigation $(PYTHON) $(PAPER_SCRIPTS)/plot_feature_investigation.py \ --features-train $(FEATURE_MATRIX_DIR)/features_train.parquet \ @@ -278,169 +301,32 @@ paper-plot-feature-investigation: paper-recompute-feature-investigation --metadata-val $(FEATURE_MATRIX_DIR)/metadata_val.parquet \ --plots-dir $(PAPER_PLOTS_DIR)/feature_investigation -# No deposited scaling JSON → live recompute/plot. -paper-plot-scaling: paper-recompute-scaling - -paper-plot-fdr-method-comparison: - mkdir -p $(PAPER_RESULTS_DIR)/fdr_method_comparison \ - $(PAPER_PLOTS_DIR)/fdr_method_comparison - $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ - --summarise-only $(PAPER_DATA_DIR)/fdr_tool_comparison/fdr_method_comparison_curves.csv \ - --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ - --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison - -paper-plot-external-peptide-holdout: - mkdir -p $(PAPER_RESULTS_DIR)/external_peptide_holdout \ - $(PAPER_PLOTS_DIR)/external_peptide_holdout - $(PYTHON) $(PAPER_SCRIPTS)/run_external_peptide_holdout_benchmark.py \ - --summarise-only $(PAPER_DATA_DIR)/external_peptide_holdout_benchmark/external_peptide_holdout_results.csv \ - --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ - --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout - -paper-plot-ablations: - @echo "Cannot replot the original ablation bars from the deposited summaries." - @echo "They contain reviewer metrics at 5%/10% FDR, but not top-decile tail_ECE." - @echo "Recompute the ablation evaluation outputs, then run plot_ablation_summary.py." - @exit 1 - -paper-plot-generalisation: - mkdir -p $(PAPER_PLOTS_DIR)/generalisation - $(PYTHON) $(PAPER_SCRIPTS)/plot_calibrator_generalisation_heatmap.py \ - --results-path $(PAPER_DATA_DIR)/generalisation/calibrator_generalisation_results.csv \ - --plots-dir $(PAPER_PLOTS_DIR)/generalisation - -# Replot general-model feature importance for PXD014877 (C. elegans) from Figshare pickles. -# Skips SHAP bar / correlation matrix (need feature matrices not in the deposit). -paper-plot-feature-importance: - @test -f $(PAPER_DATA_DIR)/feature_importance/PXD014877/perm_importance.pkl || \ - (echo "Missing feature_importance pickles; run download-paper-artefacts" && exit 1) - @test -f $(PAPER_DATA_DIR)/feature_importance/PXD014877/shap_values.pkl || \ - (echo "Missing feature_importance pickles; run download-paper-artefacts" && exit 1) - mkdir -p $(FEATURE_IMPORTANCE_PLOTS_DIR) - $(PYTHON) $(PAPER_SCRIPTS)/analyze_features.py \ - --replot-dir $(PAPER_DATA_DIR)/feature_importance/PXD014877 \ - --output-dir $(FEATURE_IMPORTANCE_PLOTS_DIR) - -# Recompute general-model feature importance on labelled C. elegans (PXD014877). -# Heavy to recompute -paper-recompute-feature-importance: - @echo "Recomputing feature importance (Koina required)." - @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) - @test -f $(CELEGANS_LABELLED_DIR)/celegans.parquet || \ - (echo "Missing $(CELEGANS_LABELLED_DIR)/celegans.parquet; run download-paper-datasets" && exit 1) - @test -f $(CELEGANS_LABELLED_DIR)/celegans_preds.csv || \ - (echo "Missing $(CELEGANS_LABELLED_DIR)/celegans_preds.csv; run download-paper-datasets" && exit 1) - mkdir -p $(FEATURE_IMPORTANCE_RESULTS_DIR) $(FEATURE_IMPORTANCE_PLOTS_DIR) - $(PYTHON) $(PAPER_SCRIPTS)/analyze_features.py \ - --model-path $(GENERAL_MODEL) \ - --data-dir $(CELEGANS_LABELLED_DIR) \ - --train-spectra celegans.parquet \ - --train-preds celegans_preds.csv \ - --test-spectra celegans.parquet \ - --test-preds celegans_preds.csv \ - --koina-input-constant collision_energies=27 \ - --koina-input-constant fragmentation_types=HCD \ - --n-background-samples 200 \ - --n-test-samples 500 \ - --output-dir $(FEATURE_IMPORTANCE_RESULTS_DIR) - @echo "Pickles+plots under $(FEATURE_IMPORTANCE_RESULTS_DIR)." - -# Full retrain of leave-one-source-out calibrators on HF general_model_training_set. -# Heavy (Koina + nine train/eval loops). -# Outputs: $(GENERALISATION_RESULTS_DIR)/calibrator_generalisation_results.csv and -# $(GENERALISATION_MODELS_DIR)/trained_on_/. -paper-recompute-generalisation: - @echo "Retraining leave-one-out calibrators (Koina required; not Figshare summaries alone)." - @test -f $(GENERALISATION_TRAIN_PARQUET) || \ - (echo "Missing $(GENERALISATION_TRAIN_PARQUET); run download-paper-datasets" && exit 1) - @test -f $(GENERALISATION_TRAIN_PREDS) || \ - (echo "Missing $(GENERALISATION_TRAIN_PREDS); run download-paper-datasets" && exit 1) - mkdir -p $(GENERALISATION_MODELS_DIR) $(GENERALISATION_RESULTS_DIR) - $(PYTHON) $(PAPER_SCRIPTS)/evaluate_calibrator_generalisation.py \ - --train-parquet $(GENERALISATION_TRAIN_PARQUET) \ - --train-predictions $(GENERALISATION_TRAIN_PREDS) \ - --model-output-dir $(GENERALISATION_MODELS_DIR) \ - --results-output-dir $(GENERALISATION_RESULTS_DIR) - -# Figshare analysis CSVs alone cannot rebuild these figures; recompute from -# deposited general_results PSM trees, writing tables and plots separately. -paper-plot-upscored-fps: paper-recompute-upscored-fps - -paper-plot-fdr-overlap: paper-recompute-fdr-overlap - -# Novelty plots need a full recompute (deposited novelty trees are CSV-only). -paper-plot-novelty: paper-recompute-novelty - -paper-recompute-novelty: - @echo "Recomputing novelty tables+plots from deposited general_results (not Figshare novelty CSVs alone)." - mkdir -p $(PAPER_RESULTS_DIR)/novelty/chymotrypsin \ - $(PAPER_RESULTS_DIR)/novelty/proteometools \ - $(PAPER_PLOTS_DIR)/novelty/chymotrypsin \ - $(PAPER_PLOTS_DIR)/novelty/proteometools - @test -f $(FASTA_HUMAN) || (echo "Missing $(FASTA_HUMAN); run download-paper-datasets" && exit 1) - $(PYTHON) $(PAPER_SCRIPTS)/analyze_novelty.py nontryptic-digest \ - --predictions-dir $(PAPER_DATA_DIR)/general_results/full/PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ - --fasta $(FASTA_HUMAN) \ - --results-dir $(PAPER_RESULTS_DIR)/novelty/chymotrypsin \ - --plots-dir $(PAPER_PLOTS_DIR)/novelty/chymotrypsin \ - --file-prefix chymotrypsin \ - --dataset-label "HeLa chymotrypsin" - $(PYTHON) $(PAPER_SCRIPTS)/analyze_novelty.py proteometools \ - --lcfm-predictions-dir $(PAPER_DATA_DIR)/general_results/labelled/PXD004732 \ - --acfm-predictions-dir $(PAPER_DATA_DIR)/general_results/full/PXD004732 \ - --results-dir $(PAPER_RESULTS_DIR)/novelty/proteometools \ - --plots-dir $(PAPER_PLOTS_DIR)/novelty/proteometools - -# Leaf keys under general_results/{labelled,full}/ (nested runs use the run folder name). -# Astral / PXD014877 are slow; subset for a smoke pass, e.g.: -# FDR_OVERLAP_PROJECTS='$(FDR_OVERLAP_immuno2)' make -f Makefile.paper paper-recompute-fdr-overlap -FDR_OVERLAP_immuno2 := PXD023064 -FDR_OVERLAP_PROJECTS ?= \ - 20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ - 20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 \ - PXD004732 \ - 20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2 \ - 20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1 \ - 01747_C01_P018218_S00_I00_N03_R1 \ - PXD014877 \ - PXD023064 \ - astral - -paper-recompute-fdr-overlap: - @echo "Recomputing FDR overlap from deposited general_results (not Figshare overlap CSVs alone)." - @echo "Projects: $(strip $(FDR_OVERLAP_PROJECTS))" - @echo "Note: astral and PXD014877 are slow; option to subset with FDR_OVERLAP_PROJECTS='$$(FDR_OVERLAP_immuno2)'." - mkdir -p $(PAPER_RESULTS_DIR)/fdr_overlap $(PAPER_PLOTS_DIR)/fdr_overlap - $(PYTHON) $(PAPER_SCRIPTS)/analyze_fdr_overlap.py \ - --labelled-dir $(PAPER_DATA_DIR)/general_results/labelled \ - --unlabelled-dir $(PAPER_DATA_DIR)/general_results/full \ - --results-dir $(PAPER_RESULTS_DIR)/fdr_overlap \ - --plots-dir $(PAPER_PLOTS_DIR)/fdr_overlap \ - --projects "$(strip $(FDR_OVERLAP_PROJECTS))" - -paper-recompute-upscored-fps: - @echo "Recomputing upscored-FP tables+plots from deposited general_results (not Figshare upscored CSVs alone)." - mkdir -p $(PAPER_RESULTS_DIR)/upscored_fps $(PAPER_PLOTS_DIR)/upscored_fps - $(PYTHON) $(PAPER_SCRIPTS)/analyze_upscored_fps.py \ - --predictions-root $(PAPER_DATA_DIR)/general_results/labelled \ - --results-dir $(PAPER_RESULTS_DIR)/upscored_fps \ - --plots-dir $(PAPER_PLOTS_DIR)/upscored_fps - ################################################################################ -# Recompute +# 2. HeLa DNS ################################################################################ -paper-recompute: paper-recompute-helaqc paper-recompute-general-labelled \ - paper-recompute-feature-investigation paper-recompute-scaling \ - paper-recompute-fdr-method-comparison paper-recompute-external-peptide-holdout - @echo "Skipping general-full-large (cluster-only). Optional: paper-recompute-general-full-small" +define HELAQC_PLOT_TOOL_RULE +.PHONY: paper-plot-helaqc-$(1) +paper-plot-helaqc-$(1): + @mkdir -p $$(PAPER_PLOTS_DIR)/helaqc_analysis + $$(PYTHON) $$(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $$(PAPER_DATA_DIR)/helaqc_results/$(1)/test \ + --split test --label-mode labelled --fasta $$(FASTA_HUMAN) \ + --dns-model '$$(HELAQC_DNS_LABEL_$(1))' \ + --model-dir $$(HELAQC_DNS_MODEL_$(1)) \ + --plots-dir $$(PAPER_PLOTS_DIR)/helaqc_analysis/$(1)_test + $$(PYTHON) $$(PAPER_SCRIPTS)/plot_analysis.py \ + --predictions-dir $$(PAPER_DATA_DIR)/helaqc_results/$(1)/raw_less_train \ + --split raw_less_train --label-mode unlabelled --fasta $$(FASTA_HUMAN) \ + --dns-model '$$(HELAQC_DNS_LABEL_$(1))' \ + --plots-dir $$(PAPER_PLOTS_DIR)/helaqc_analysis/$(1)_raw_less_train +endef + +$(foreach t,$(HELAQC_DNS),$(eval $(call HELAQC_PLOT_TOOL_RULE,$(t)))) -# InstaNovo × 3 splits by default; set HELAQC_DNS_TOOLS=instanovo casanovo primenovo for all. -HELAQC_DNS_TOOLS ?= instanovo +paper-plot-helaqc-analysis: $(foreach t,$(HELAQC_DNS),paper-plot-helaqc-$(t)) paper-recompute-helaqc: - @test -d $(HELAQC_DATA) || (echo "Missing $(HELAQC_DATA); run download-paper-datasets" && exit 1) - @test -d $(HELAQC_MODEL) || (echo "Missing $(HELAQC_MODEL); run download-paper-models" && exit 1) @mkdir -p $(PAPER_RESULTS_DIR) @for tool in $(HELAQC_DNS_TOOLS); do \ for split in $(HELAQC_SPLITS); do \ @@ -464,32 +350,74 @@ paper-recompute-helaqc: done; \ done -paper-recompute-general-labelled: - @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) - @test -d $(GENERAL_EVAL_DIR) || (echo "Missing $(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) - @mkdir -p $(PAPER_RESULTS_DIR)/general_results/labelled - @for stem in $(GENERAL_RECOMPUTE_STEMS); do \ - $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-labelled-$${stem}; \ - done +################################################################################ +# 3. FDR method comparison +################################################################################ -paper-recompute-general-full-small: - @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) - @test -d $(GENERAL_EVAL_DIR) || (echo "Missing $(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) - @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full - @for stem in $(GENERAL_FULL_SMALL_STEMS); do \ - $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ - done +paper-plot-fdr-method-comparison: + mkdir -p $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + $(PAPER_PLOTS_DIR)/fdr_method_comparison + $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ + --summarise-only $(PAPER_DATA_DIR)/fdr_tool_comparison/fdr_method_comparison_curves.csv \ + --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison -paper-recompute-general-full-large: - @echo "CLUSTER ONLY: $(GENERAL_FULL_LARGE_STEMS) — not part of paper-recompute umbrella." - @test -d $(GENERAL_MODEL) || (echo "Missing $(GENERAL_MODEL); run download-paper-models" && exit 1) - @test -d $(GENERAL_EVAL_DIR) || (echo "Missing $(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) - @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full - @for stem in $(GENERAL_FULL_LARGE_STEMS); do \ - $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ - done +paper-recompute-fdr-method-comparison: + mkdir -p $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + $(PAPER_PLOTS_DIR)/fdr_method_comparison + $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ + --novoboard-root $(NOVOBOARD_ROOT) \ + --winnow-results $(WINNOW_FDR_RESULTS) \ + $(foreach d,$(FDR_RECOMPUTE_DATASETS),--datasets $(d)) \ + --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ + --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison + +################################################################################ +# 4. External peptide holdout +################################################################################ + +paper-plot-external-peptide-holdout: + mkdir -p $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + $(PAPER_PLOTS_DIR)/external_peptide_holdout + $(PYTHON) $(PAPER_SCRIPTS)/run_external_peptide_holdout_benchmark.py \ + --summarise-only $(PAPER_DATA_DIR)/external_peptide_holdout_benchmark/external_peptide_holdout_results.csv \ + --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout + +paper-recompute-external-peptide-holdout: + mkdir -p $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + $(PAPER_PLOTS_DIR)/external_peptide_holdout + $(PYTHON) $(PAPER_SCRIPTS)/run_external_peptide_holdout_benchmark.py \ + --novoboard-root $(NOVOBOARD_ROOT) \ + --winnow-results $(WINNOW_FDR_RESULTS) \ + --model-root $(FDR_MODEL_ROOT) \ + $(foreach d,$(FDR_RECOMPUTE_DATASETS),--datasets $(d)) \ + --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ + --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout + +################################################################################ +# 5. General-model evaluation +################################################################################ + +paper-plot-general-labelled: + mkdir -p $(PAPER_RESULTS_DIR)/general_labelled $(PAPER_PLOTS_DIR)/general_labelled + $(PYTHON) $(PAPER_SCRIPTS)/plot_eval_results.py \ + --predictions-root $(PAPER_DATA_DIR)/general_results/labelled \ + --eval-type labelled \ + $(GENERAL_PROJECTS_FLAG) \ + --results-dir $(PAPER_RESULTS_DIR)/general_labelled \ + --plots-dir $(PAPER_PLOTS_DIR)/general_labelled + +paper-plot-general-full: + mkdir -p $(PAPER_RESULTS_DIR)/general_full $(PAPER_PLOTS_DIR)/general_full + $(PYTHON) $(PAPER_SCRIPTS)/plot_eval_results.py \ + --predictions-root $(PAPER_DATA_DIR)/general_results/full \ + --eval-type unlabelled \ + $(GENERAL_PROJECTS_FLAG) \ + --results-dir $(PAPER_RESULTS_DIR)/general_full \ + --plots-dir $(PAPER_PLOTS_DIR)/general_full -# Figshare project key and proteome FASTA per HF stem. +# Nested Figshare output paths (and proteome FASTA) per HF stem. GENERAL_PROJECT_hela_chymotrypsin := PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin GENERAL_FASTA_hela_chymotrypsin := $(FASTA_HUMAN) GENERAL_PROJECT_human_lung := PXD004452/20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46 @@ -510,7 +438,7 @@ GENERAL_PROJECT_astral := astral GENERAL_FASTA_astral := $(FASTA_ECOLI) # If GENERAL_EXPERIMENT_FILTER_ is set, materialise a filtered parquet/CSV -# under GENERAL_SUBSET_DIR (revisions-era PXD023064 cohort for immuno2). +# under GENERAL_SUBSET_DIR (subset for immuno2). # Expanded via $(call ...) from stem recipes (single $$ → shell $). define GENERAL_MAYBE_SUBSET filter="$(strip $(GENERAL_EXPERIMENT_FILTER_$(1)))"; \ @@ -538,12 +466,9 @@ else \ fi endef -# Explicit per-stem rules (empty .PHONY + pattern recipes do not mix). define GENERAL_LABELLED_STEM_RULE .PHONY: paper-recompute-general-labelled-$(1) paper-recompute-general-labelled-$(1): - @test -d $$(GENERAL_MODEL) || (echo "Missing $$(GENERAL_MODEL); run download-paper-models" && exit 1) - @test -d $$(GENERAL_EVAL_DIR) || (echo "Missing $$(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) @$$(call GENERAL_MAYBE_SUBSET,$(1),labelled); \ out=$$(PAPER_RESULTS_DIR)/general_results/labelled/$$(GENERAL_PROJECT_$(1)); \ mkdir -p "$$$$out"; \ @@ -561,12 +486,9 @@ endef define GENERAL_FULL_STEM_RULE .PHONY: paper-recompute-general-full-$(1) paper-recompute-general-full-$(1): - @test -d $$(GENERAL_MODEL) || (echo "Missing $$(GENERAL_MODEL); run download-paper-models" && exit 1) - @test -d $$(GENERAL_EVAL_DIR) || (echo "Missing $$(GENERAL_EVAL_DIR); run download-paper-datasets" && exit 1) @$$(call GENERAL_MAYBE_SUBSET,$(1),full); \ out=$$(PAPER_RESULTS_DIR)/general_results/full/$$(GENERAL_PROJECT_$(1)); \ fasta=$$(GENERAL_FASTA_$(1)); \ - test -f $$$$fasta || (echo "Missing $$$$fasta; run download-paper-datasets" && exit 1); \ mkdir -p "$$$$out"; \ echo "Predict full $(1) -> $$$$out"; \ $$(WINNOW) predict \ @@ -593,27 +515,70 @@ $(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_LABELLED_STEM_RULE,$(s)))) $(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_FULL_STEM_RULE,$(s)))) $(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_BOTH_STEM_RULE,$(s)))) -# Echo exact recipes for all nine stems × labelled/full (safe on a laptop). +paper-recompute-general-labelled: + @mkdir -p $(PAPER_RESULTS_DIR)/general_results/labelled + @for stem in $(GENERAL_RECOMPUTE_STEMS); do \ + $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-labelled-$${stem}; \ + done + +paper-recompute-general-full-small: + @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full + @for stem in $(GENERAL_FULL_SMALL_STEMS); do \ + $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ + done + +paper-recompute-general-full-large: + @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full + @for stem in $(GENERAL_FULL_LARGE_STEMS); do \ + $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ + done + +# Echo exact recipes for all nine stems × labelled/full (immuno2 always subsets). paper-recompute-general-print: @echo "# General-model recompute commands (pins from Makefile.paper)" @echo "# Requires: download-paper-datasets download-paper-models; Koina network for predict" + @echo "# immuno2 inputs always subsetted to the Figshare cohort before predict." @echo "" @$(foreach s,$(GENERAL_STEMS), \ echo "# --- $(s) (labelled) ---"; \ - echo "$(WINNOW) predict \\"; \ - echo " data_loader=instanovo \\"; \ - echo " dataset.spectrum_path_or_directory=$(GENERAL_EVAL_DIR)/$(s)/labelled/$(s).parquet \\"; \ - echo " dataset.predictions_path=$(GENERAL_EVAL_DIR)/$(s)/labelled/$(s)_preds.csv \\"; \ + $(if $(strip $(GENERAL_EXPERIMENT_FILTER_$(s))), \ + echo "$(PYTHON) $(PAPER_SCRIPTS)/subset_eval_by_experiment.py \\"; \ + echo " --spectra $(GENERAL_EVAL_DIR)/$(s)/labelled/$(s).parquet \\"; \ + echo " --preds $(GENERAL_EVAL_DIR)/$(s)/labelled/$(s)_preds.csv \\"; \ + echo " --output-dir $(GENERAL_SUBSET_DIR)/$(s)/labelled \\"; \ + echo " --stem $(s) $(foreach exp,$(GENERAL_EXPERIMENT_FILTER_$(s)),--experiment $(exp))"; \ + echo "$(WINNOW) predict \\"; \ + echo " data_loader=instanovo \\"; \ + echo " dataset.spectrum_path_or_directory=$(GENERAL_SUBSET_DIR)/$(s)/labelled/$(s).parquet \\"; \ + echo " dataset.predictions_path=$(GENERAL_SUBSET_DIR)/$(s)/labelled/$(s)_preds.csv \\"; \ + , \ + echo "$(WINNOW) predict \\"; \ + echo " data_loader=instanovo \\"; \ + echo " dataset.spectrum_path_or_directory=$(GENERAL_EVAL_DIR)/$(s)/labelled/$(s).parquet \\"; \ + echo " dataset.predictions_path=$(GENERAL_EVAL_DIR)/$(s)/labelled/$(s)_preds.csv \\"; \ + ) \ echo " calibrator.pretrained_model_name_or_path=$(GENERAL_MODEL) \\"; \ echo " fdr_control.fdr_threshold=$(PREDICT_FDR_THRESHOLD) \\"; \ echo " $(if $(GENERAL_KOINA_$(s)),$(GENERAL_KOINA_$(s)),$(KOINA_FRAGMENT_MATCH_COLUMNS)) \\"; \ echo " output_folder=$(PAPER_RESULTS_DIR)/general_results/labelled/$(GENERAL_PROJECT_$(s))"; \ echo ""; \ echo "# --- $(s) (full: predict then post-annotate) ---"; \ - echo "$(WINNOW) predict \\"; \ - echo " data_loader=instanovo \\"; \ - echo " dataset.spectrum_path_or_directory=$(GENERAL_EVAL_DIR)/$(s)/full/$(s).parquet \\"; \ - echo " dataset.predictions_path=$(GENERAL_EVAL_DIR)/$(s)/full/$(s)_preds.csv \\"; \ + $(if $(strip $(GENERAL_EXPERIMENT_FILTER_$(s))), \ + echo "$(PYTHON) $(PAPER_SCRIPTS)/subset_eval_by_experiment.py \\"; \ + echo " --spectra $(GENERAL_EVAL_DIR)/$(s)/full/$(s).parquet \\"; \ + echo " --preds $(GENERAL_EVAL_DIR)/$(s)/full/$(s)_preds.csv \\"; \ + echo " --output-dir $(GENERAL_SUBSET_DIR)/$(s)/full \\"; \ + echo " --stem $(s) $(foreach exp,$(GENERAL_EXPERIMENT_FILTER_$(s)),--experiment $(exp))"; \ + echo "$(WINNOW) predict \\"; \ + echo " data_loader=instanovo \\"; \ + echo " dataset.spectrum_path_or_directory=$(GENERAL_SUBSET_DIR)/$(s)/full/$(s).parquet \\"; \ + echo " dataset.predictions_path=$(GENERAL_SUBSET_DIR)/$(s)/full/$(s)_preds.csv \\"; \ + , \ + echo "$(WINNOW) predict \\"; \ + echo " data_loader=instanovo \\"; \ + echo " dataset.spectrum_path_or_directory=$(GENERAL_EVAL_DIR)/$(s)/full/$(s).parquet \\"; \ + echo " dataset.predictions_path=$(GENERAL_EVAL_DIR)/$(s)/full/$(s)_preds.csv \\"; \ + ) \ echo " calibrator.pretrained_model_name_or_path=$(GENERAL_MODEL) \\"; \ echo " fdr_control.fdr_threshold=$(PREDICT_FDR_THRESHOLD) \\"; \ echo " $(if $(GENERAL_KOINA_$(s)),$(GENERAL_KOINA_$(s)),$(KOINA_FRAGMENT_MATCH_COLUMNS)) \\"; \ @@ -625,79 +590,119 @@ paper-recompute-general-print: echo ""; \ ) -# Feature investigation needs the InstaNovo HeLa feature stack (beam + token + -# mass_error_da). Package compute-features has no pretrained_model_name_or_path; -# narrow the lean training matrix to the checkpoint feature_columns schema. -HELAQC_FEATURE_INVESTIGATION_OVERRIDES = \ - +calibrator.features.beam_features._target_=winnow.calibration.calibration_features.BeamFeatures \ - +calibrator.features.token_score_features._target_=winnow.calibration.calibration_features.TokenScoreFeatures \ - calibrator.features.retention_time_feature.min_train_points=5 \ - 'calibrator.training_feature_columns=[ion_matches,ion_match_intensity,irt_error,margin,median_margin,entropy,z-score,min_token_probability,std_token_probability,mass_error_da]' +################################################################################ +# 6. Feature importance +################################################################################ -paper-recompute-feature-investigation: - @echo "Recomputing feature matrices (not deposited on Figshare)." - @test -d $(HELAQC_DATA) || (echo "Missing $(HELAQC_DATA); run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/train.parquet || (echo "Missing $(HELAQC_DATA)/train.parquet" && exit 1) - @test -f $(HELAQC_DATA)/val.parquet || (echo "Missing $(HELAQC_DATA)/val.parquet" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/train_preds.csv || (echo "Missing InstaNovo train preds" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/val_preds.csv || (echo "Missing InstaNovo val preds" && exit 1) - mkdir -p $(FEATURE_MATRIX_DIR) - $(WINNOW) compute-features \ - data_loader=instanovo \ - labelled=true \ - dataset.spectrum_path_or_directory=$(HELAQC_DATA)/train.parquet \ - dataset.predictions_path=$(HELAQC_DATA)/instanovo/train_preds.csv \ - $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ - $(HELAQC_FEATURE_INVESTIGATION_OVERRIDES) \ - metadata_output_path=$(FEATURE_MATRIX_DIR)/metadata_train.parquet \ - training_matrix_output_path=$(FEATURE_MATRIX_DIR)/features_train.parquet - $(WINNOW) compute-features \ - data_loader=instanovo \ - labelled=true \ - dataset.spectrum_path_or_directory=$(HELAQC_DATA)/val.parquet \ - dataset.predictions_path=$(HELAQC_DATA)/instanovo/val_preds.csv \ - $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ - $(HELAQC_FEATURE_INVESTIGATION_OVERRIDES) \ - metadata_output_path=$(FEATURE_MATRIX_DIR)/metadata_val.parquet \ - training_matrix_output_path=$(FEATURE_MATRIX_DIR)/features_val.parquet +# Replot from Figshare pickles (skips SHAP bar / correlation matrix). +paper-plot-feature-importance: + mkdir -p $(FEATURE_IMPORTANCE_PLOTS_DIR) + $(PYTHON) $(PAPER_SCRIPTS)/analyze_features.py \ + --replot-dir $(PAPER_DATA_DIR)/feature_importance/PXD014877 \ + --output-dir $(FEATURE_IMPORTANCE_PLOTS_DIR) -paper-recompute-scaling: - @echo "Recomputing scaling benchmark (trains a no-Prosit dummy; no Figshare JSON)." - @test -f $(HELAQC_DATA)/train.parquet || (echo "Missing helaqc train parquet; run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/val.parquet || (echo "Missing helaqc val parquet; run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/raw_less_train.parquet || (echo "Missing helaqc raw_less_train parquet; run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/train_preds.csv || (echo "Missing InstaNovo train preds" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/val_preds.csv || (echo "Missing InstaNovo val preds" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv || (echo "Missing InstaNovo raw_less_train preds" && exit 1) - mkdir -p $(PAPER_RESULTS_DIR)/scaling $(PAPER_PLOTS_DIR)/scaling - $(PYTHON) $(PAPER_SCRIPTS)/benchmark_scaling.py \ - --spectrum-path $(HELAQC_DATA)/train.parquet \ - --predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ - --spectrum-path $(HELAQC_DATA)/raw_less_train.parquet \ - --predictions-path $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv \ - --train-spectrum-path $(HELAQC_DATA)/train.parquet \ - --train-predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ - --val-spectrum-path $(HELAQC_DATA)/val.parquet \ - --val-predictions-path $(HELAQC_DATA)/instanovo/val_preds.csv \ - --model-output-dir $(SCALING_DUMMY_MODEL) \ - --data-loader instanovo \ - --fractions 0.1 --fractions 0.5 --fractions 1.0 \ - --results-dir $(PAPER_RESULTS_DIR)/scaling \ - --plots-dir $(PAPER_PLOTS_DIR)/scaling +paper-recompute-feature-importance: + mkdir -p $(FEATURE_IMPORTANCE_RESULTS_DIR) $(FEATURE_IMPORTANCE_PLOTS_DIR) + $(PYTHON) $(PAPER_SCRIPTS)/analyze_features.py \ + --model-path $(GENERAL_MODEL) \ + --data-dir $(CELEGANS_LABELLED_DIR) \ + --train-spectra celegans.parquet \ + --train-preds celegans_preds.csv \ + --test-spectra celegans.parquet \ + --test-preds celegans_preds.csv \ + --koina-input-constant collision_energies=27 \ + --koina-input-constant fragmentation_types=HCD \ + --n-background-samples 200 \ + --n-test-samples 500 \ + --output-dir $(FEATURE_IMPORTANCE_RESULTS_DIR) + +################################################################################ +# 7. Generalisation [GPU] +################################################################################ + +paper-plot-generalisation: + mkdir -p $(PAPER_PLOTS_DIR)/generalisation + $(PYTHON) $(PAPER_SCRIPTS)/plot_calibrator_generalisation_heatmap.py \ + --results-path $(PAPER_DATA_DIR)/generalisation/calibrator_generalisation_results.csv \ + --plots-dir $(PAPER_PLOTS_DIR)/generalisation + +# Full retrain of leave-one-source-out calibrators on HF general_model_training_set. +paper-recompute-generalisation: + mkdir -p $(GENERALISATION_MODELS_DIR) $(GENERALISATION_RESULTS_DIR) + $(PYTHON) $(PAPER_SCRIPTS)/evaluate_calibrator_generalisation.py \ + --train-parquet $(GENERALISATION_TRAIN_PARQUET) \ + --train-predictions $(GENERALISATION_TRAIN_PREDS) \ + --model-output-dir $(GENERALISATION_MODELS_DIR) \ + --results-output-dir $(GENERALISATION_RESULTS_DIR) + +################################################################################ +# 8. Feature ablations [GPU] (recompute only; plots written by the script) +################################################################################ + +paper-recompute-ablations: paper-recompute-feature-investigation + @$(call GENERAL_MAYBE_SUBSET,immuno2,labelled); \ + mkdir -p $(ABLATIONS_RESULTS_DIR); \ + $(PYTHON) $(PAPER_SCRIPTS)/run_feature_ablations.py \ + --train-features $(FEATURE_MATRIX_DIR)/features_train.parquet \ + --val-features $(FEATURE_MATRIX_DIR)/features_val.parquet \ + --eval-root $(GENERAL_EVAL_DIR) \ + --override-eval-dir immuno2=$(GENERAL_SUBSET_DIR)/immuno2/labelled \ + --output-dir $(ABLATIONS_RESULTS_DIR) + +################################################################################ +# 9. Upscored FPs (plot from deposited general_results/labelled) +################################################################################ + +paper-plot-upscored-fps: + mkdir -p $(PAPER_RESULTS_DIR)/upscored_fps $(PAPER_PLOTS_DIR)/upscored_fps + $(PYTHON) $(PAPER_SCRIPTS)/analyze_upscored_fps.py \ + --predictions-root $(PAPER_DATA_DIR)/general_results/labelled \ + --results-dir $(PAPER_RESULTS_DIR)/upscored_fps \ + --plots-dir $(PAPER_PLOTS_DIR)/upscored_fps + +################################################################################ +# 10. Novelty (plot from deposited general_results) +################################################################################ + +paper-plot-novelty: + mkdir -p $(PAPER_RESULTS_DIR)/novelty/chymotrypsin \ + $(PAPER_RESULTS_DIR)/novelty/proteometools \ + $(PAPER_PLOTS_DIR)/novelty/chymotrypsin \ + $(PAPER_PLOTS_DIR)/novelty/proteometools + $(PYTHON) $(PAPER_SCRIPTS)/analyze_novelty.py nontryptic-digest \ + --predictions-dir $(PAPER_DATA_DIR)/general_results/full/PXD004452/20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin \ + --fasta $(FASTA_HUMAN) \ + --results-dir $(PAPER_RESULTS_DIR)/novelty/chymotrypsin \ + --plots-dir $(PAPER_PLOTS_DIR)/novelty/chymotrypsin \ + --file-prefix chymotrypsin \ + --dataset-label "HeLa chymotrypsin" + $(PYTHON) $(PAPER_SCRIPTS)/analyze_novelty.py proteometools \ + --lcfm-predictions-dir $(PAPER_DATA_DIR)/general_results/labelled/PXD004732 \ + --acfm-predictions-dir $(PAPER_DATA_DIR)/general_results/full/PXD004732 \ + --results-dir $(PAPER_RESULTS_DIR)/novelty/proteometools \ + --plots-dir $(PAPER_PLOTS_DIR)/novelty/proteometools + +################################################################################ +# 11. FDR overlap (plot from deposited general_results) +################################################################################ + +# Default FDR_OVERLAP_PROJECTS is all projects, but this is slow. Overrides to lighter immuno2: +# FDR_OVERLAP_PROJECTS='$(FDR_OVERLAP_immuno2)' make -f Makefile.paper paper-plot-fdr-overlap +paper-plot-fdr-overlap: + mkdir -p $(PAPER_RESULTS_DIR)/fdr_overlap $(PAPER_PLOTS_DIR)/fdr_overlap + $(PYTHON) $(PAPER_SCRIPTS)/analyze_fdr_overlap.py \ + --labelled-dir $(PAPER_DATA_DIR)/general_results/labelled \ + --unlabelled-dir $(PAPER_DATA_DIR)/general_results/full \ + --results-dir $(PAPER_RESULTS_DIR)/fdr_overlap \ + --plots-dir $(PAPER_PLOTS_DIR)/fdr_overlap \ + --projects "$(strip $(FDR_OVERLAP_PROJECTS))" + +################################################################################ +# 12. Runtime [GPU] +################################################################################ -# Stage-wise runtime table (full Prosit + no-Prosit) on the full HeLa search space. -# Prints the table to stdout and writes JSON + a text copy under paper_results/runtime/. -# Not part of paper-recompute (Koina-heavy). Requires download-paper-datasets/models. -# Reuses or trains the no-Prosit dummy under SCALING_DUMMY_MODEL (same as scaling). +# Trains a dummy model with GPU and requires Koina for full feature benchmark. paper-recompute-runtime: - @echo "Recomputing HeLa runtime table (full + no-Prosit; requires Koina for full)." - @test -d $(HELAQC_MODEL) || (echo "Missing $(HELAQC_MODEL); run download-paper-models" && exit 1) - @test -f $(HELAQC_DATA)/train.parquet || (echo "Missing helaqc train parquet; run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/val.parquet || (echo "Missing helaqc val parquet; run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/raw_less_train.parquet || (echo "Missing helaqc raw_less_train parquet; run download-paper-datasets" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/train_preds.csv || (echo "Missing InstaNovo train preds" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/val_preds.csv || (echo "Missing InstaNovo val preds" && exit 1) - @test -f $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv || (echo "Missing InstaNovo raw_less_train preds" && exit 1) mkdir -p $(RUNTIME_RESULTS_DIR) $(SCALING_DUMMY_MODEL) $(PYTHON) $(PAPER_SCRIPTS)/benchmark_runtime.py \ --spectrum-path $(HELAQC_DATA)/train.parquet \ @@ -716,25 +721,24 @@ paper-recompute-runtime: --output-json $(RUNTIME_RESULTS_DIR)/benchmark_results.json \ --output-text $(RUNTIME_RESULTS_DIR)/benchmark_results.txt -paper-recompute-fdr-method-comparison: - @echo "Recomputing FDR method comparison from fdr_benchmark_inputs (not Figshare curves alone)." - mkdir -p $(PAPER_RESULTS_DIR)/fdr_method_comparison \ - $(PAPER_PLOTS_DIR)/fdr_method_comparison - $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ - --novoboard-root $(NOVOBOARD_ROOT) \ - --winnow-results $(WINNOW_FDR_RESULTS) \ - --datasets helaqc --datasets celegans \ - --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ - --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison +################################################################################ +# 13. Scaling [GPU] +################################################################################ -paper-recompute-external-peptide-holdout: paper-sync-group-paper - @echo "Recomputing external peptide holdout from fdr_benchmark_inputs (not Figshare results CSV alone)." - mkdir -p $(PAPER_RESULTS_DIR)/external_peptide_holdout \ - $(PAPER_PLOTS_DIR)/external_peptide_holdout - $(PYTHON) $(PAPER_SCRIPTS)/run_external_peptide_holdout_benchmark.py \ - --novoboard-root $(NOVOBOARD_ROOT) \ - --winnow-results $(WINNOW_FDR_RESULTS) \ - --model-root $(FDR_MODEL_ROOT) \ - --datasets helaqc --datasets celegans \ - --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ - --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout +# Trains a dummy model with GPU +paper-recompute-scaling: + mkdir -p $(PAPER_RESULTS_DIR)/scaling $(PAPER_PLOTS_DIR)/scaling + $(PYTHON) $(PAPER_SCRIPTS)/benchmark_scaling.py \ + --spectrum-path $(HELAQC_DATA)/train.parquet \ + --predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ + --spectrum-path $(HELAQC_DATA)/raw_less_train.parquet \ + --predictions-path $(HELAQC_DATA)/instanovo/raw_less_train_preds.csv \ + --train-spectrum-path $(HELAQC_DATA)/train.parquet \ + --train-predictions-path $(HELAQC_DATA)/instanovo/train_preds.csv \ + --val-spectrum-path $(HELAQC_DATA)/val.parquet \ + --val-predictions-path $(HELAQC_DATA)/instanovo/val_preds.csv \ + --model-output-dir $(SCALING_DUMMY_MODEL) \ + --data-loader instanovo \ + --fractions 0.1 --fractions 0.5 --fractions 1.0 \ + --results-dir $(PAPER_RESULTS_DIR)/scaling \ + --plots-dir $(PAPER_PLOTS_DIR)/scaling From 4805961f35821867329c20a8a81792da08e8102b Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:59:57 +0100 Subject: [PATCH 14/26] docs: add paper reproduction guide --- paper_scripts/README.md | 242 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 paper_scripts/README.md diff --git a/paper_scripts/README.md b/paper_scripts/README.md new file mode 100644 index 00000000..d6fb592b --- /dev/null +++ b/paper_scripts/README.md @@ -0,0 +1,242 @@ +# Reproducing Winnow paper results and plots + +This folder contains the analysis scripts used in the Winnow paper ([arXiv:2509.24952](https://arxiv.org/abs/2509.24952)). +Together with [`Makefile.paper`](../Makefile.paper) in the repository root, they let a reviewer regenerate the reported figures and summary tables from public artefacts. + +## 1. Setup + +```bash +git clone https://github.com/instadeepai/winnow.git +cd winnow +make -f Makefile.paper paper-sync-group-paper +make -f Makefile.paper help +make -f Makefile.paper paper-setup # Figshare + HF datasets + models, then artefact check +make -f Makefile.paper paper-plot-light +``` + +`paper-setup` is the one-shot acquire path. +Be aware that the HF dataset pin is large (Astral and full-search inputs making up the bulk of the size). +The Hugging Face CLI can be used to download specific datasets. +Alternatively, selective downloads are available: + +```bash +make -f Makefile.paper download-paper-artefacts +make -f Makefile.paper download-paper-datasets +make -f Makefile.paper download-paper-models +make -f Makefile.paper paper-check-artefacts +``` + +Recompute umbrellas (after setup): + +```bash +make -f Makefile.paper paper-recompute-laptop # no large datasets / no GPU +make -f Makefile.paper paper-recompute-heavy # large datasets, GPU +``` + +Or run one analysis at a time (`paper-plot-helaqc-analysis`, +`paper-recompute-external-peptide-holdout`, …). +See `make -f Makefile.paper help`. + +`winnow predict` for general-model and HeLa recomputes calls Koina (public server via `koinapy`) for fragment-intensity and iRT features. +Outbound network access is required and runtime is likely to be dominated by these calls. + +### Command naming: `plot` vs `recompute` + +| Prefix | Means | May call | +| --- | --- | --- | +| `paper-plot-*` | Build figures / summary tables from Figshare deposits (or a deterministic local transform of them). | CPU only; never Koina-dependent, never uses calibrator training or inference | +| `paper-recompute-*` | Rebuild intermediates that the deposit does not fully pin: inference, Koina feature computation, calibrator training, or timed benchmarks. Scripts often write plots at the end of the same target. | Koina and/or GPU as noted | + +If a paper figure has no deposited artefact that can be plotted without those steps, there is no `paper-plot-*` target (table cell says “n/a”). The `paper-recompute-*` target still produces the plots when it finishes. + +### Umbrella membership + +Laptop vs heavy is about workload size and hardware requirements. + +| Umbrella | Rule | Contents | +| --- | --- | --- | +| `paper-setup` | acquire | `download-paper-artefacts`, `download-paper-datasets`, `download-paper-models`, `paper-check-artefacts` | +| `paper-plot-light` | deposit replot without the largest datasets | HeLa DNS, general labelled, FDR method / holdout, generalisation heatmap, feature importance, novelty, upscored FPs, FDR overlap (**immuno2** only) | +| `paper-recompute-laptop` | no GPU; no unlabelled / large cohorts | HeLa DNS predict, general **labelled** (default stems), feature investigation, FDR method / holdout for **helaqc** | +| `paper-recompute-heavy` | large cohorts, GPU | `general-full-*`, feature-importance, generalisation, ablations, runtime, scaling, FDR tools including **celegans**, `paper-plot-general-full`, FDR overlap (**all** projects) | + +## 2. Public pins + +| Resource | Pin | +| --- | --- | +| Figshare [Analysis outputs](https://figshare.com/articles/dataset/Analysis_outputs/30147601) | article `30147601` **v7** (`10.6084/m9.figshare.30147601.v7`) | +| Hugging Face [winnow-ms-datasets](https://huggingface.co/datasets/InstaDeepAI/winnow-ms-datasets) | `659802319d618a359de5ab90ec6b0195681e94a6` | +| Hugging Face [winnow-general-model](https://huggingface.co/InstaDeepAI/winnow-general-model) | `e2089330dd59adb9685e5b3d7d61f0cd69a3bbb0` | +| Hugging Face [winnow-helaqc-model](https://huggingface.co/InstaDeepAI/winnow-helaqc-model) | `d56542b961eac7d896e51bf0716a242fc394ab1f` | +| Figshare [Additional HeLa Single Shot models](https://doi.org/10.6084/m9.figshare.32744946.v2) | `32744946` v2 (Casanovo / π-PrimeNovo calibrators) | +| Glissade ([JemmaLDaniel/glissade](https://github.com/JemmaLDaniel/glissade), branch `winnow-benchmark`) | `7c723a2af4a88fda84a6bd4f223b351179bd36da` via `uv sync --group paper` | +| NovoBoard (only if regenerating decoy CSVs from MGFs) | `a9faab3ef1af06987599c2f01e6ba96072c80172` | + +## 3. What each analysis does + +| Paper analysis | Plot (from Figshare) | Recompute | Needs | Script | +| --- | --- | --- | --- | --- | +| HeLa Single Shot DNS (InstaNovo, Casanovo, π-PrimeNovo × test / raw_less_train) | `paper-plot-helaqc-analysis` | `paper-recompute-helaqc` | Koina | `plot_analysis.py` | +| General-model evaluation (nine projects, labelled and full-search) | `paper-plot-general-labelled`, `paper-plot-general-full` | tiered / per-stem `paper-recompute-general-*` | Koina | `plot_eval_results.py` | +| Feature investigation (InstaNovo HeLa) | n/a (matrices not deposited) | `paper-recompute-feature-investigation` (compute + plot) | Koina | `plot_feature_investigation.py` | +| General-model feature importance (*C. elegans*) | `paper-plot-feature-importance` (no SHAP bar / correlations) | `paper-recompute-feature-importance` | Koina | `analyze_features.py` | +| Pipeline scaling excluding Koina features | n/a (JSON not deposited; script writes plots) | `paper-recompute-scaling` | GPU + Koina | `benchmark_scaling.py` | +| Pipeline runtime table (full + no-Prosit) | n/a (JSON not deposited) | `paper-recompute-runtime` | GPU + Koina | `benchmark_runtime.py` | +| PSM-level FDR vs NovoBoard | `paper-plot-fdr-method-comparison` (`--summarise-only`) | `paper-recompute-fdr-method-comparison` | CPU | `plot_fdr_method_comparison.py` | +| External peptide score-mixture (Winnow / NovoBoard / Glissade) | `paper-plot-external-peptide-holdout` | `paper-recompute-external-peptide-holdout` | CPU | `run_external_peptide_holdout_benchmark.py` | +| Feature ablations | n/a (deposit lacks tail ECE for top 10% PSMs; script writes plots) | `paper-recompute-ablations` | GPU + Koina | `run_feature_ablations.py` | +| Calibrator generalisation heatmap | `paper-plot-generalisation` | `paper-recompute-generalisation` | GPU + Koina | `evaluate_calibrator_generalisation.py`, `plot_calibrator_generalisation_heatmap.py` | +| Upscored FPs | `paper-plot-upscored-fps` | n/a (uses deposited labelled `general_results/`) | CPU | `analyze_upscored_fps.py` | +| FDR overlap | `paper-plot-fdr-overlap` | n/a (uses deposited labelled + full trees) | CPU | `analyze_fdr_overlap.py` | +| Novelty | `paper-plot-novelty` | n/a (uses deposited chymotrypsin / ProteomeTools trees) | CPU | `analyze_novelty.py` | + +### General-model projects (`plot_eval_results`) + +Make / CLI `--projects` use leaf folder names. +Nested Figshare trees are resolved automatically (`PXD004452//` or flat `PXD004732/`): + +- `20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin` +- `20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46` +- `PXD004732` +- `20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2` +- `20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1` +- `01747_C01_P018218_S00_I00_N03_R1` +- `PXD014877` +- `PXD023064` +- `astral` + +### General-model recompute (HF map and Make targets) + +Inputs live under `paper_data/winnow-ms-datasets/general_model_evaluation//{labelled,full}/`. +Outputs mirror Figshare under `paper_results/general_results/{labelled,full}//`. + +| HF stem (Make suffix) | Figshare project key | FASTA | +| --- | --- | --- | +| `hela_chymotrypsin` | `PXD004452/20150708_…_Chymotrypsin` | `fasta/human.fasta` | +| `human_lung` | `PXD004452/20151020_…_A549_Rep2_46` | `fasta/human.fasta` | +| `proteometools1` | `PXD004732` | `fasta/human.fasta` | +| `HLA_I` | `PXD006939/…_HLApI_…` | `fasta/human.fasta` | +| `HLA_II` | `PXD006939/…_HLAIIp_…` | `fasta/human.fasta` | +| `athaliana` | `PXD013868/01747_…` | `fasta/athaliana.fasta` | +| `celegans` | `PXD014877` | `fasta/celegans.fasta` | +| `immuno2` | `PXD023064` | `fasta/human.fasta` | +| `astral` | `astral` | `fasta/ecoli_zorya.fasta` | + +**Labelled:** `winnow predict` only. + +**Full-search:** `winnow predict`, then `annotate_preds_proteome_hits.py` (post-predict; adds `proteome_hit`, drops peptides shorter than 7 residues). + +Note that **immuno2** always subsets to the Figshare cohort (`PXD023064_FILES`) before predict. + +All paper `winnow predict` calls pass `fdr_control.fdr_threshold=1.0` (`PREDICT_FDR_THRESHOLD`) so outputs keep every PSM with FDR columns. + +Koina collision energy / fragmentation: + +- **HeLa Single Shot** (`paper-recompute-helaqc`, feature investigation): tiled constants `collision_energies=27`, `fragmentation_types=HCD` (`KOINA_FRAGMENT_MATCH_CONSTANTS`). +- **General-model evaluation** (most HF stems): null those constants (`KOINA_FRAGMENT_MATCH_COLUMNS`) so runtime resolution uses the default metadata columns `collision_energy` / `frag_type` (see `docs/configuration.md` and `resolve_feature_model_inputs`). +- **immuno2 / PXD023064**: tiled `CE=27` / `HCD`. + +Batch / print: + +```bash +make -f Makefile.paper paper-recompute-general-labelled # default stems: immuno2 hela_chymotrypsin human_lung +make -f Makefile.paper paper-recompute-general-full-small # in paper-recompute-heavy +make -f Makefile.paper paper-recompute-general-full-large # in paper-recompute-heavy +make -f Makefile.paper paper-recompute-general-print # echo all nine × labelled/full commands +``` + +Per dataset (HF stem): + +```bash +make -f Makefile.paper paper-recompute-general-labelled-immuno2 +make -f Makefile.paper paper-recompute-general-full-immuno2 +make -f Makefile.paper paper-recompute-general-immuno2 # labelled + full +``` + +Override labelled batch: `GENERAL_RECOMPUTE_STEMS='immuno2 athaliana' make -f Makefile.paper paper-recompute-general-labelled`. + +HeLa DNS tools default to all three (`HELAQC_DNS_TOOLS=instanovo casanovo primenovo`). Subset with e.g. `HELAQC_DNS_TOOLS=instanovo`. + +Selective HF download (skip Astral when you only need a small project): + +```bash +uv run hf download InstaDeepAI/winnow-ms-datasets \ + --repo-type dataset \ + --revision 659802319d618a359de5ab90ec6b0195681e94a6 \ + --include 'general_model_evaluation/immuno2/**' \ + --include 'fasta/**' \ + --local-dir paper_data/winnow-ms-datasets +``` + +### How full-search proteome annotation differs from the package CLI + +Paper full recompute intentionally keeps the old post-predict path so outputs stay aligned with deposited `general_results/full/` artefacts. +The package CLI `winnow annotate-proteome-hits` is the supported product path going forward (annotate-then-predict / diagnose). +Reviewers who follow only `docs/cli.md` will get a different order and can differ slightly from deposits. + +| | Paper helper (this suite) | Package CLI | +| --- | --- | --- | +| Command | `paper_scripts/annotate_preds_proteome_hits.py` | `winnow annotate-proteome-hits` | +| When | **After** `winnow predict` | **Before** predict (or for diagnose) | +| Input | Predict output folder | Spectra + de novo preds via `DatasetLoader` | +| Output | In-place CSVs with `proteome_hit`; short peptides removed | New Winnow dataset | +| FDR | Estimated on all PSMs, then short peptides dropped | Short peptides dropped before feature/FDR if you annotate then predict | +| Used for | Reproducing Figshare `general_results/full/` | Diagnose / annotate-then-predict workflows in main docs | + +## 4. Approximate deposit sizes + +Orders of magnitude for Figshare v7 outputs (metadata + prediction CSVs unless noted): + +| Tree | Size | +| --- | --- | +| `general_results/labelled/` (nine projects) | ~2 GB | +| `general_results/full/` | ~5 GB (astral ~2 GB; *C. elegans* ~1.8 GB) | +| `fdr_benchmark_inputs/` | ~1 GB+ (includes HeLa MGFs for twin pairing) | +| Generalisation results CSV | ~6.8 GB | +| Remaining analysis CSVs / HeLa result trees | much smaller | + +Approximate HF input sizes (parquet + InstaNovo preds) for recompute: + +| Tier | Stems | Notes | +| --- | --- | --- | +| Labelled defaults | `immuno2`, `hela_chymotrypsin`, `human_lung` | Tens–hundreds of MB each | +| Full small | `hela_chymotrypsin`, `human_lung`, `HLA_I/II`, `athaliana`, `immuno2` | immuno2 full ~1.4 GB inputs | +| Cluster | `astral` full and labelled, `celegans` full, `proteometools1` full | Astral labelled ~2.7 GB; Astral full ~6.7 GB inputs | + +Replotting labelled deposits and `paper-plot-light` is the most accessible review path. +`paper-plot-general-full` (and full FDR overlap for large datasets) are slow, so they are listed under `paper-recompute-heavy` even though they are still `paper-plot-*` targets. +Use `paper-recompute-general-print` to obtain the exact cluster predict commands without running them. + +## 5. FDR tool benchmarking + +For both PSM-level comparison and the peptide holdout: + +1. Download `fdr_benchmark_inputs/` with the Figshare article (or `paper-setup` / `download-paper-artefacts`). +2. Replot with `--summarise-only` on the deposited curves / results CSVs, or recompute with `--novoboard-root paper_data/fdr_benchmark_inputs/novoboard` and the matching Winnow prediction folders. + +NovoBoard’s twin-decoy competition is reimplemented here against those CSVs. Glissade’s bootstrap FDR is imported from the package installed by `uv sync --group paper`. +We do not cover reproducing decoy spectra here, but this can be done using standard NovoBoard decoy generation commands and the per-dataset decoy generation strategies described in the paper. + +## 6. Entrypoints + +| Script | Role | +| --- | --- | +| `plot_analysis.py` | HeLa DNS evaluation plots for InstaNovo, Casanovo and π-PrimeNovo | +| `plot_eval_results.py` | General model evaluation figures | +| `plot_feature_investigation.py` | Feature distributions / investigation (after recompute matrices) | +| `benchmark_scaling.py` | Runtime vs dataset size (trains/reuses a no-Prosit dummy, then times the pipeline) **[GPU]** | +| `benchmark_runtime.py` | Stage-wise wall-time / memory table (full Prosit + no-Prosit; same dummy as scaling) **[GPU]** | +| `no_prosit_dummy.py` | Shared train-or-reuse helper for the no-Prosit dummy calibrator | +| `plot_fdr_method_comparison.py` | Winnow vs NovoBoard PSM FDR | +| `run_external_peptide_holdout_benchmark.py` | Controlled-π₀ peptide mixture benchmark | +| `run_feature_ablations.py` | Feature-subset calibrator training + eval **[GPU]** | +| `plot_ablation_summary.py` | Ablation bar summaries (from recompute outputs, not Figshare alone) | +| `plot_calibrator_generalisation_heatmap.py` | Hold-one-out generalisation heatmap | +| `evaluate_calibrator_generalisation.py` | Full retrain leave-one-source-out results CSV **[GPU]** | +| `analyze_features.py` | Feature importance / SHAP (supports `--replot-dir` from Figshare pickles) | +| `analyze_upscored_fps.py`, `analyze_fdr_overlap.py`, `analyze_novelty.py` | Downstream analyses of deposited `general_results/` | +| `annotate_preds_proteome_hits.py` | Post-predict proteome hits for full-search recompute | +| `download_figshare_article.py` | Figshare download with folder layout | +| `subset_eval_by_experiment.py` | Filter immuno2 (and similar) to the deposited cohort | + +Every entrypoint supports `--help`. From e3d10aa2e6248dd4b514f7d038a107d7d389764d Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:23:57 +0100 Subject: [PATCH 15/26] chore: delete old scripts --- scripts/analyze_fdr_overlap.py | 1293 ------------ scripts/analyze_features.py | 804 -------- scripts/analyze_novelty.py | 1812 ----------------- scripts/analyze_upscored_fps.py | 709 ------- scripts/benchmark_runtime.py | 690 ------- scripts/benchmark_scaling.py | 421 ---- scripts/calibrator_generalisation_utils.py | 114 -- scripts/evaluate_calibrator_generalisation.py | 415 ---- scripts/fdr_tool_comparison_preprocess.py | 809 -------- scripts/fdr_tool_comparison_summaries.py | 533 ----- scripts/feature_subsets.py | 74 - scripts/plot_ablation_summary.py | 641 ------ scripts/plot_acfm_minus_lcfm_fdr.py | 470 ----- scripts/plot_analysis.py | 1122 ---------- .../plot_calibrator_generalisation_heatmap.py | 265 --- scripts/plot_eval_results.py | 984 --------- scripts/plot_fdr_method_comparison.py | 1106 ---------- scripts/plot_feature_investigation.py | 1323 ------------ .../run_external_peptide_holdout_benchmark.py | 1143 ----------- scripts/run_feature_ablations.py | 1532 -------------- 20 files changed, 16260 deletions(-) delete mode 100644 scripts/analyze_fdr_overlap.py delete mode 100644 scripts/analyze_features.py delete mode 100644 scripts/analyze_novelty.py delete mode 100644 scripts/analyze_upscored_fps.py delete mode 100644 scripts/benchmark_runtime.py delete mode 100644 scripts/benchmark_scaling.py delete mode 100644 scripts/calibrator_generalisation_utils.py delete mode 100644 scripts/evaluate_calibrator_generalisation.py delete mode 100644 scripts/fdr_tool_comparison_preprocess.py delete mode 100644 scripts/fdr_tool_comparison_summaries.py delete mode 100644 scripts/feature_subsets.py delete mode 100644 scripts/plot_ablation_summary.py delete mode 100644 scripts/plot_acfm_minus_lcfm_fdr.py delete mode 100644 scripts/plot_analysis.py delete mode 100644 scripts/plot_calibrator_generalisation_heatmap.py delete mode 100644 scripts/plot_eval_results.py delete mode 100644 scripts/plot_fdr_method_comparison.py delete mode 100644 scripts/plot_feature_investigation.py delete mode 100644 scripts/run_external_peptide_holdout_benchmark.py delete mode 100644 scripts/run_feature_ablations.py diff --git a/scripts/analyze_fdr_overlap.py b/scripts/analyze_fdr_overlap.py deleted file mode 100644 index 6e51fc2d..00000000 --- a/scripts/analyze_fdr_overlap.py +++ /dev/null @@ -1,1293 +0,0 @@ -#!/usr/bin/env python3 -"""Post-FDR overlap analysis: Winnow-filtered identifications vs database search. - -For each project (see ``plot_eval_results.py`` CLI pattern), at 1 %, 5 %, and 10 % -nominal FDR: - - * Count retained PSMs / unique peptides vs database-search reference peptides at - the same nominal FDR (Winnow: non-parametric on calibrated confidence; - database search: database-grounded on raw confidence). - * Match rule: exact ProForma sequence after I/L equivalence; PTM differences - are not a match. - * Categorise discordant calls (partial match, PTM candidate, single-AA variant, - near-miss edit distance 2-3, fully discordant). - * Full-search Venns: Winnow (non-parametric calibrated confidence) vs database - search unique peptides at the same nominal FDR (database-grounded raw confidence). - * Violin plots comparing database-matched vs fully novel retained PSMs at - matched FDR (same retention rules as overlap summaries). - -Inputs are ``winnow predict`` output folders arranged as subdirectories under two -roots: an **unlabelled** tree (full-search Winnow predictions) and a **labelled** -tree (database-search reference with ``sequence``). Each project folder (flat or -``PXD*//`` nested) must contain ``preds_and_fdr_metrics.csv``; -``metadata.csv`` is merged when present for violin plots. -""" - -from __future__ import annotations - -import json -import logging -import re -from collections import defaultdict -from pathlib import Path -from typing import Annotated, Callable, cast - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import seaborn as sns -import typer -import yaml -from matplotlib_venn import venn2 -from rich.logging import RichHandler - -from winnow.fdr.database_grounded import DatabaseGroundedFDRControl -from winnow.fdr.nonparametric import NonParametricFDRControl - -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -logger.propagate = False -if not logger.handlers: - logger.addHandler(RichHandler()) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -# --------------------------------------------------------------------------- -# Style — Paul Tol "bright" palette (colour-blind safe) -# --------------------------------------------------------------------------- -_PALETTE = [ - "#4477AA", - "#EE6677", - "#228833", - "#CCBB44", - "#66CCEE", - "#AA3377", - "#BBBBBB", -] -_CORRECT_COLOUR = _PALETTE[0] -_INCORRECT_COLOUR = _PALETTE[1] -_NOVEL_COLOUR = _PALETTE[2] - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -_REPO_ROOT = Path(__file__).resolve().parent.parent -_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)-?") -_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") - -_PXD_ACCESSION_PREFIX = "PXD" - -FDR_THRESHOLDS = [0.01, 0.05, 0.10] -_DB_GROUNDED_DROP = 10 -RAW_CONFIDENCE_COL = "confidence" -DB_Q_VALUE_COL = "db_psm_q_value" - -DATASET_DISPLAY_NAMES: dict[str, str] = { - "gluc": "HeLa degradome", - "helaqc": "HeLa single shot", - "herceptin": "Herceptin", - "immuno": "Immunopeptidomics-1", - "celegans": "$\\it{C.\\;elegans}$", - "sbrodae": "$\\it{Scalindua\\;brodae}$", - "PXD019483": "HepG2", - "snakevenoms": "Snake venomics", - "tplantibodies": "Therapeutic nanobodies", - "woundfluids": "Wound exudates", - "PXD014877": "$\\it{C.\\;elegans}$", - "PXD023064": "Immunopeptidomics-2", - "astral": "Astral $\\it{E.\\;coli}$", - "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", - "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", - "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", - "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", - "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", - "PXD004732": "ProteomeTools-1", -} - -_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") -_UNLABELLED_FOLDER_SUFFIXES = ("_raw", "_unlabelled") -_LABELLED_FOLDER_SUFFIXES = ("_annotated", "_labelled") - -NOVEL_FEATURE_COLUMNS: list[tuple[str, str]] = [ - ("spectral_angle", "Spectral angle"), - ("ion_matches", "Ion match rate"), - ("ion_match_intensity", "Ion match intensity"), - ("precursor_charge", "Precursor charge"), - ("mass_error_da", "Precursor mass error (Da)"), - ("irt_error", "iRT error"), - ("confidence", "Raw confidence"), - ("margin", "Beam margin"), -] - -_DISCORDANCE_COUNT_COLS = [ - "n_partial_match", - "n_ptm_candidate", - "n_single_aa_variant", - "n_near_miss_edit_dist", - "n_fully_discordant", -] - -_PTM_DELTAS = { - "oxidation": 15.995, - "phosphorylation": 79.966, - "deamidation": 0.984, - "acetylation": 42.011, - "methylation": 14.016, - "carbamidomethyl": 57.021, -} -_PTM_TOLERANCE_DA = 0.02 - -_MIN_VIOLIN_GROUP_SIZE = 5 -_MAX_VIOLIN_PSMs_PER_GROUP = 5000 -_VIOLIN_SUBSAMPLE_SEED = 42 - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- -def _display_name(key: str) -> str: - return DATASET_DISPLAY_NAMES.get(key, key) - - -def _project_key_from_folder(folder_name: str) -> str: - """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" - for suffix in _FOLDER_SUFFIXES: - if folder_name.endswith(suffix): - return folder_name[: -len(suffix)] - return folder_name - - -def _search_space_tag_from_folder(folder_name: str) -> str: - """Infer eval-type label for tables/plots from the unlabelled subfolder name.""" - for suffix in _UNLABELLED_FOLDER_SUFFIXES: - if folder_name.endswith(suffix): - return suffix[1:] # raw | unlabelled - return "full_search" - - -def _eval_type_display(eval_type: str) -> str: - """Human-readable eval-type label for plot titles.""" - return { - "full_search": "full search space", - "raw": "raw", - "unlabelled": "unlabelled", - }.get(eval_type, eval_type) - - -def _get_residue_masses() -> dict[str, float]: - config_path = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" - with open(config_path) as f: - cfg = yaml.safe_load(f) - return cfg["residue_masses"] - - -def _save_fig(fig: plt.Figure, base_path: Path) -> None: - fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def _style_ax(ax: plt.Axes) -> None: - ax.grid(False) - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -def _sequence_match_key(seq: str) -> str: - """Exact match key: ProForma with mods preserved, I/L equivalent.""" - if not seq or not isinstance(seq, str): - return "" - return seq.replace("I", "L") - - -def _strip_mods(seq: str) -> str: - """Strip PTM annotations and normalise I -> L (discordance subtyping only).""" - if not seq or not isinstance(seq, str): - return "" - s = _MOD_PLUS.sub("", seq) - s = _MOD_UNIMOD.sub("", s) - return s.replace("I", "L") - - -def _levenshtein(s: str, t: str) -> int: - n, m = len(s), len(t) - if n == 0: - return m - if m == 0: - return n - prev = list(range(m + 1)) - for i in range(1, n + 1): - curr = [i] + [0] * m - for j in range(1, m + 1): - cost = 0 if s[i - 1] == t[j - 1] else 1 - curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost) - prev = curr - return prev[m] - - -def _mass_from_sequence(seq_str: str, residue_masses: dict[str, float]) -> float | None: - key = _strip_mods(seq_str) - if not key: - return None - total = 18.010565 - for aa in key: - m = residue_masses.get(aa) - if m is None: - return None - total += m - return total - - -def _db_stripped_by_length(db_stripped_list: list[str]) -> dict[int, list[str]]: - by_len: dict[int, list[str]] = defaultdict(list) - for s in db_stripped_list: - by_len[len(s)].append(s) - return by_len - - -def _min_edit_distance_to_db( - pred_stripped: str, db_stripped_by_len: dict[int, list[str]] -) -> int: - if not pred_stripped: - return 999 - plen = len(pred_stripped) - best = 999 - for length in range(max(0, plen - 3), plen + 4): - for db in db_stripped_by_len.get(length, []): - d = _levenshtein(pred_stripped, db) - if d < best: - best = d - if best == 0: - return 0 - return best - - -def _build_db_reference_sets( - db_df: pd.DataFrame, -) -> tuple[set[str], set[str], list[str], dict[int, list[str]]]: - sequences = db_df["sequence"].dropna().astype(str) - db_keys = {_sequence_match_key(s) for s in sequences if _sequence_match_key(s)} - db_stripped = [_strip_mods(s) for s in sequences if _strip_mods(s)] - db_stripped_set = set(db_stripped) - db_stripped_unique = list(dict.fromkeys(db_stripped)) - return ( - db_keys, - db_stripped_set, - db_stripped_unique, - _db_stripped_by_length(db_stripped_unique), - ) - - -def _is_db_match(pred: str, db_keys: set[str]) -> bool: - return _sequence_match_key(pred) in db_keys - - -# --------------------------------------------------------------------------- -# Discordance classification -# --------------------------------------------------------------------------- -LabelledDiscordanceKey = tuple[str, str, int] -LabelledDiscordanceCache = dict[LabelledDiscordanceKey, str] -DbDiscordanceCache = dict[str, str] - - -def _discordance_from_edit_distance(seq_norm: str, pred_norm: str) -> str | None: - if not (seq_norm and pred_norm): - return None - ed = _levenshtein(seq_norm, pred_norm) - if ed == 1: - return "single_aa_variant" - if ed in (2, 3): - return "near_miss_edit_dist" - return None - - -def _discordance_from_ptm_mass( - seq_str: str, - pred_str: str, - residue_masses: dict[str, float], -) -> str | None: - seq_mass = _mass_from_sequence(seq_str, residue_masses) - pred_mass = _mass_from_sequence(pred_str, residue_masses) - if seq_mass is None or pred_mass is None: - return None - delta = abs(seq_mass - pred_mass) - for ptm_delta in _PTM_DELTAS.values(): - if abs(delta - ptm_delta) < _PTM_TOLERANCE_DA: - return "ptm_candidate" - return None - - -def _labelled_discordance_key(row: pd.Series) -> LabelledDiscordanceKey: - return ( - str(row.get("sequence", "")), - str(row.get("prediction", "")), - int(row.get("num_matches", 0)), - ) - - -def _lookup_labelled_discordance( - row: pd.Series, cache: LabelledDiscordanceCache -) -> str: - return cache[_labelled_discordance_key(row)] - - -def classify_discordance_labelled( - seq_str: str, - pred_str: str, - num_matches: int, - residue_masses: dict[str, float], -) -> str: - """Classify a non-matching PSM on a labelled spectrum.""" - if num_matches > 0: - return "partial_match" - - seq_norm = _strip_mods(seq_str) - pred_norm = _strip_mods(pred_str) - if seq_norm == pred_norm and _sequence_match_key(seq_str) != _sequence_match_key( - pred_str - ): - return "ptm_candidate" - - edit_label = _discordance_from_edit_distance(seq_norm, pred_norm) - if edit_label is not None: - return edit_label - - ptm_label = _discordance_from_ptm_mass(seq_str, pred_str, residue_masses) - if ptm_label is not None: - return ptm_label - - return "fully_discordant" - - -def classify_discordance_vs_db( - pred_str: str, - db_keys: set[str], - db_stripped_set: set[str], - db_stripped_by_len: dict[int, list[str]], -) -> str: - """Classify a discordant full-search PSM vs the database peptide reference set.""" - pred_key = _sequence_match_key(pred_str) - if pred_key in db_keys: - return "db_match" - - pred_stripped = _strip_mods(pred_str) - if pred_stripped in db_stripped_set: - return "ptm_candidate" - - ed = _min_edit_distance_to_db(pred_stripped, db_stripped_by_len) - if ed == 1: - return "single_aa_variant" - if ed in (2, 3): - return "near_miss_edit_dist" - return "fully_discordant" - - -def _build_discordance_cache( - predictions: pd.Series, - db_keys: set[str], - db_stripped_set: set[str], - db_stripped_by_len: dict[int, list[str]], -) -> dict[str, str]: - """Classify each unique discordant prediction string once.""" - cache: dict[str, str] = {} - for pred in predictions.dropna().unique(): - key = str(pred) - if _is_db_match(key, db_keys): - cache[key] = "db_match" - elif key not in cache: - cache[key] = classify_discordance_vs_db( - key, db_keys, db_stripped_set, db_stripped_by_len - ) - return cache - - -def _classify_predictions_vs_db( - predictions: pd.Series, - cache: dict[str, str], -) -> pd.Series: - return predictions.map(lambda p: cache.get(str(p), "fully_discordant")) - - -# --------------------------------------------------------------------------- -# Data loading -# --------------------------------------------------------------------------- -def _preds_header(folder: Path) -> set[str] | None: - preds_path = folder / "preds_and_fdr_metrics.csv" - if not preds_path.is_file(): - return None - return set(pd.read_csv(preds_path, nrows=0).columns.tolist()) - - -def _is_unlabelled_preds_folder(folder: Path) -> bool: - header = _preds_header(folder) - return header is not None and "sequence" not in header - - -def _is_labelled_preds_folder(folder: Path) -> bool: - header = _preds_header(folder) - return header is not None and "sequence" in header - - -def _load_from_folder(folder: Path) -> pd.DataFrame: - """Load preds_and_fdr_metrics.csv merged with metadata.csv from a project folder.""" - preds_path = folder / "preds_and_fdr_metrics.csv" - if not preds_path.is_file(): - raise FileNotFoundError(f"Missing predictions file: {preds_path}") - - preds_df = pd.read_csv(preds_path) - meta_path = folder / "metadata.csv" - if meta_path.is_file(): - meta_df = pd.read_csv(meta_path) - overlap_cols = [ - c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" - ] - if overlap_cols: - meta_df = meta_df.drop(columns=overlap_cols) - return preds_df.merge(meta_df, on="spectrum_id", how="left") - return preds_df - - -def _register_project_folder(projects: dict[str, Path], key: str, folder: Path) -> None: - """Register *folder* under *key*, warning on duplicate keys.""" - if key in projects: - logger.warning( - "Duplicate project key %r: %s and %s", - key, - projects[key], - folder, - ) - return - projects[key] = folder - - -def _discover_project_folders( - root: Path, - *, - is_preds_folder: Callable[[Path], bool], -) -> dict[str, Path]: - """Map project key -> preds folder under *root* (flat or ``PXD*//``).""" - projects: dict[str, Path] = {} - if not root.is_dir(): - return projects - - for child in sorted(root.iterdir()): - if not child.is_dir(): - continue - if is_preds_folder(child): - _register_project_folder( - projects, _project_key_from_folder(child.name), child - ) - continue - if not child.name.startswith(_PXD_ACCESSION_PREFIX): - continue - for run_dir in sorted(child.iterdir()): - if run_dir.is_dir() and is_preds_folder(run_dir): - _register_project_folder(projects, run_dir.name, run_dir) - return projects - - -def _discover_unlabelled_folders(root: Path) -> dict[str, Path]: - """Map project key -> full-search folder under ``root``. - - Supports flat project folders (``{root}/PXD004732/``) and nested per-run - layouts (``{root}/PXD004452//``) used by new eval sets. - """ - return _discover_project_folders(root, is_preds_folder=_is_unlabelled_preds_folder) - - -def _discover_labelled_folders(root: Path) -> dict[str, Path]: - """Map project key -> database-reference folder under ``root``. - - Supports flat project folders (``{root}/PXD004732/``) and nested per-run - layouts (``{root}/PXD004452//``) used by new eval sets. - """ - return _discover_project_folders(root, is_preds_folder=_is_labelled_preds_folder) - - -def discover_project_pairs( - unlabelled_dir: Path, - labelled_dir: Path, - *, - projects_filter: set[str] | None = None, -) -> list[tuple[str, Path, Path, str]]: - """Return ``(project, unlabelled_folder, labelled_folder, search_space_tag)`` pairs.""" - unlabelled = _discover_unlabelled_folders(unlabelled_dir) - labelled = _discover_labelled_folders(labelled_dir) - - keys = sorted(unlabelled.keys() & labelled.keys()) - if projects_filter is not None: - keys = [k for k in keys if k in projects_filter] - - pairs: list[tuple[str, Path, Path, str]] = [] - for key in keys: - pairs.append( - ( - key, - unlabelled[key], - labelled[key], - _search_space_tag_from_folder(unlabelled[key].name), - ) - ) - - for key in sorted(unlabelled.keys() - labelled.keys()): - if projects_filter is None or key in projects_filter: - logger.warning("No labelled folder for unlabelled project %r", key) - for key in sorted(labelled.keys() - unlabelled.keys()): - if projects_filter is None or key in projects_filter: - logger.warning("No unlabelled folder for labelled project %r", key) - - return pairs - - -def _effective_db_grounded_drop(n_rows: int, drop: int = _DB_GROUNDED_DROP) -> int: - return min(drop, max(0, n_rows - 1)) - - -def _add_q_values( - df: pd.DataFrame, - conf_col: str = "calibrated_confidence", - *, - q_col: str = "psm_q_value", -) -> pd.DataFrame: - """Attach PSM q-values from a non-parametric FDR fit on *conf_col*.""" - if q_col in df.columns: - return df - if conf_col not in df.columns: - raise ValueError(f"Missing confidence column {conf_col!r}") - - existing_q = df["psm_q_value"] if "psm_q_value" in df.columns else None - work = df.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") - - fdr = NonParametricFDRControl() - fdr.fit(dataset=work[conf_col]) - out = fdr.add_psm_q_value(work, confidence_col=conf_col) - if q_col != "psm_q_value": - out = out.rename(columns={"psm_q_value": q_col}) - if existing_q is not None and q_col != "psm_q_value": - out["psm_q_value"] = existing_q - return out - - -def _add_database_grounded_q_values( - df: pd.DataFrame, - confidence_col: str = RAW_CONFIDENCE_COL, - *, - q_col: str = DB_Q_VALUE_COL, - correct_col: str = "correct", -) -> pd.DataFrame: - """Attach PSM q-values from database-grounded FDR on *confidence_col*.""" - if q_col in df.columns: - return df - if confidence_col not in df.columns: - raise ValueError(f"Missing confidence column {confidence_col!r}") - if "sequence" not in df.columns or "prediction" not in df.columns: - raise ValueError( - "Database-grounded FDR requires 'sequence' and 'prediction' columns" - ) - - # Drop Winnow NP q-values before fitting; merge-based add_psm_q_value can - # leave duplicate psm_q_value_* columns and skip the rename to db_psm_q_value. - work = df.drop( - columns=[q_col, "psm_q_value", "psm_fdr", "fdr"], - errors="ignore", - ).copy() - drop = _effective_db_grounded_drop(len(work)) - ctrl = DatabaseGroundedFDRControl( - confidence_feature=confidence_col, - drop=drop, - ) - # DataFrame path: mirror DatabaseGroundedFDRControl.fit without CalibrationDataset. - sorted_df = work.sort_values(confidence_col, ascending=False) - labels = sorted_df[correct_col].astype(float).to_numpy() - conf = sorted_df[confidence_col].to_numpy() - precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) - ctrl._fdr_values = np.array(1.0 - precision)[drop:] - ctrl._confidence_scores = conf[drop:] - q_df = ctrl.add_psm_q_value( - work[[confidence_col]].copy(), confidence_col=confidence_col - ) - work[q_col] = q_df["psm_q_value"].values - return work - - -def _unique_peptides_at_fdr( - df: pd.DataFrame, - sequence_col: str, - q_col: str, - fdr_t: float, -) -> set[str]: - retained = df[df[q_col] <= fdr_t] - return set(retained[sequence_col].dropna().map(_sequence_match_key)) - {""} - - -def _empty_overlap_row( - project: str, - eval_type: str, - fdr_t: float, - n_db_peptides: int, - labelled_subset: bool, -) -> dict: - row: dict = { - "project": project, - "eval_type": eval_type, - "fdr_threshold": fdr_t, - "n_psms_retained": 0, - "n_unique_peptides_retained": 0, - "n_db_search_peptides": n_db_peptides, - "n_matching": 0, - "pct_matching": 0.0, - "n_discordant": 0, - "pct_discordant": 0.0, - } - for col in _DISCORDANCE_COUNT_COLS: - if col == "n_partial_match" and not labelled_subset: - continue - row[col] = 0 - return row - - -def _discordance_cache_for_fdr_retained( - df: pd.DataFrame, - db_keys: set[str], - db_stripped_set: set[str], - db_stripped_by_len: dict[int, list[str]], - *, - labelled_subset: bool, - residue_masses: dict[str, float] | None, -) -> LabelledDiscordanceCache | DbDiscordanceCache: - """Build discordance lookup for all predictions retained at any FDR threshold.""" - df = _add_q_values(df) - retained = df[df["psm_q_value"] <= max(FDR_THRESHOLDS)] - if labelled_subset and residue_masses is not None: - cache: LabelledDiscordanceCache = {} - disc = retained[ - retained["prediction"].map(_sequence_match_key) - != retained["sequence"].map(_sequence_match_key) - ] - for _, row in disc.drop_duplicates( - subset=["sequence", "prediction"] - ).iterrows(): - trip = ( - str(row.get("sequence", "")), - str(row.get("prediction", "")), - int(row.get("num_matches", 0)), - ) - if trip not in cache: - cache[trip] = classify_discordance_labelled( - trip[0], trip[1], trip[2], residue_masses - ) - return cache - - return _build_discordance_cache( - retained["prediction"], db_keys, db_stripped_set, db_stripped_by_len - ) - - -def compute_overlap_table( - df: pd.DataFrame, - project: str, - eval_type: str, - db_df: pd.DataFrame, - discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, - residue_masses: dict[str, float], - *, - labelled_subset: bool = False, -) -> pd.DataFrame: - """Overlap summary at each FDR threshold. - - Winnow uses non-parametric FDR on calibrated confidence (full-search run). - Database reference peptides use database-grounded FDR on raw confidence in the - database-labelled run (see ``plot_full_search_venn``). - """ - df = _add_q_values(df.copy()) - db_scored = _add_database_grounded_q_values( - db_df.copy(), - confidence_col=RAW_CONFIDENCE_COL, - q_col=DB_Q_VALUE_COL, - ) - - rows: list[dict] = [] - for fdr_t in FDR_THRESHOLDS: - db_keys_at_fdr = _unique_peptides_at_fdr( - db_scored, "sequence", DB_Q_VALUE_COL, fdr_t - ) - n_db_peptides = len(db_keys_at_fdr) - - retained = df[df["psm_q_value"] <= fdr_t].copy() - n_retained = len(retained) - if n_retained == 0: - rows.append( - _empty_overlap_row( - project, eval_type, fdr_t, n_db_peptides, labelled_subset - ) - ) - continue - - if labelled_subset: - retained["pred_key"] = retained["prediction"].map(_sequence_match_key) - retained["seq_key"] = retained["sequence"].map(_sequence_match_key) - is_match = retained["pred_key"] == retained["seq_key"] - else: - is_match = retained["prediction"].map( - lambda p, keys=db_keys_at_fdr: _is_db_match(str(p), keys) - ) - - n_matching = int(is_match.sum()) - n_discordant = n_retained - n_matching - n_unique_peptides = int( - retained["prediction"].map(_sequence_match_key).replace("", pd.NA).nunique() - ) - - disc = retained[~is_match] - cat_counts: dict[str, int] = {} - if len(disc) > 0: - if labelled_subset: - labelled_cache = cast(LabelledDiscordanceCache, discordance_cache) - cats = disc.apply( - _lookup_labelled_discordance, axis=1, cache=labelled_cache - ) - else: - cats = _classify_predictions_vs_db( - disc["prediction"], - cast(DbDiscordanceCache, discordance_cache), - ) - cat_counts = cats.value_counts().to_dict() - - row: dict = { - "project": project, - "eval_type": eval_type, - "fdr_threshold": fdr_t, - "n_psms_retained": n_retained, - "n_unique_peptides_retained": n_unique_peptides, - "n_db_search_peptides": n_db_peptides, - "n_matching": n_matching, - "pct_matching": round(n_matching / n_retained * 100, 2), - "n_discordant": n_discordant, - "pct_discordant": round(n_discordant / n_retained * 100, 2), - "n_ptm_candidate": cat_counts.get("ptm_candidate", 0), - "n_single_aa_variant": cat_counts.get("single_aa_variant", 0), - "n_near_miss_edit_dist": cat_counts.get("near_miss_edit_dist", 0), - "n_fully_discordant": cat_counts.get("fully_discordant", 0), - } - if labelled_subset: - row["n_partial_match"] = cat_counts.get("partial_match", 0) - rows.append(row) - - return pd.DataFrame(rows) - - -# --------------------------------------------------------------------------- -# Plots -# --------------------------------------------------------------------------- -def _draw_venn_panel( - ax, - *, - winnow_peptides: set[str], - db_peptides: set[str], - winnow_label: str, - fdr_t: float, -) -> None: - """Draw one FDR-threshold Venn panel onto *ax*.""" - pct = int(fdr_t * 100) - if not winnow_peptides and not db_peptides: - ax.set_title(f"No peptides retained at {pct}% FDR") - ax.axis("off") - return - - if not winnow_peptides or not db_peptides: - missing = "Winnow" if not winnow_peptides else "Database search" - ax.text( - 0.5, - 0.5, - f"No {missing} peptides at {pct}% FDR", - ha="center", - va="center", - transform=ax.transAxes, - ) - ax.set_title(f"{pct}% FDR") - ax.axis("off") - return - - venn2( - [db_peptides, winnow_peptides], - set_labels=("Database search", winnow_label), - set_colors=(_CORRECT_COLOUR, _INCORRECT_COLOUR), - alpha=0.6, - ax=ax, - ) - ax.set_title(f"Unique peptides at {pct}% FDR") - _style_ax(ax) - - -def _plot_venn_panels( - winnow_df: pd.DataFrame, - project: str, - output_path: Path, - *, - winnow_label: str, - suptitle: str, - residue_masses: dict[str, float] | None = None, - db_df: pd.DataFrame | None = None, - db_peptides_static: set[str] | None = None, -) -> None: - del residue_masses # kept for call-site compatibility - del project - if db_df is None and db_peptides_static is None: - raise ValueError("Provide db_df or db_peptides_static for Venn panels") - - winnow_scored = _add_q_values(winnow_df.copy()) - db_scored: pd.DataFrame | None = None - if db_df is not None: - db_scored = _add_database_grounded_q_values( - db_df.copy(), - confidence_col=RAW_CONFIDENCE_COL, - q_col=DB_Q_VALUE_COL, - ) - - n_thresholds = len(FDR_THRESHOLDS) - fig, axes = plt.subplots(1, n_thresholds, figsize=(5 * n_thresholds, 5)) - if n_thresholds == 1: - axes = [axes] - - for ax, fdr_t in zip(axes, FDR_THRESHOLDS): - winnow_peptides = _unique_peptides_at_fdr( - winnow_scored, "prediction", "psm_q_value", fdr_t - ) - if db_scored is not None: - db_peptides = _unique_peptides_at_fdr( - db_scored, "sequence", DB_Q_VALUE_COL, fdr_t - ) - else: - assert db_peptides_static is not None - db_peptides = db_peptides_static - - _draw_venn_panel( - ax, - winnow_peptides=winnow_peptides, - db_peptides=db_peptides, - winnow_label=winnow_label, - fdr_t=fdr_t, - ) - - fig.suptitle(suptitle, fontsize=12) - fig.tight_layout() - _save_fig(fig, output_path) - - -def plot_full_search_venn( - df: pd.DataFrame, - db_df: pd.DataFrame, - project: str, - plots_dir: Path, - residue_masses: dict[str, float], -) -> None: - """Venn diagrams of FDR-filtered DB vs Winnow full-search unique peptides. - - Database peptides use database-grounded q-values on raw ``confidence`` in the - database-labelled run. Winnow peptides use non-parametric q-values on - ``calibrated_confidence`` in the full-search run. - """ - display = _display_name(project) - _plot_venn_panels( - df, - project, - plots_dir / f"venn_{project}_full_search", - winnow_label="Winnow", - suptitle=f"Database search vs Winnow full search at matched FDR for {display}", - residue_masses=residue_masses, - db_df=db_df, - ) - - -def plot_labelled_subset_venn( - df: pd.DataFrame, - project: str, - plots_dir: Path, -) -> None: - """Venn diagrams of labelled reference peptides vs Winnow predictions per FDR.""" - db_peptides = { - _sequence_match_key(s) - for s in df["sequence"].dropna() - if _sequence_match_key(s) - } - display = _display_name(project) - _plot_venn_panels( - df, - project, - plots_dir / f"venn_{project}_labelled_subset", - winnow_label="Winnow", - suptitle=f"Peptide overlap on labelled spectra for {display}", - db_peptides_static=db_peptides, - ) - - -def _assign_retained_groups( - retained: pd.DataFrame, - db_keys: set[str] | None, - discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, - *, - labelled_subset: bool, -) -> pd.Series: - if labelled_subset and "sequence" in retained.columns: - labelled_cache = cast(LabelledDiscordanceCache, discordance_cache) - is_match = retained["prediction"].map(_sequence_match_key) == retained[ - "sequence" - ].map(_sequence_match_key) - groups = pd.Series("Database match", index=retained.index) - disc_mask = ~is_match - if disc_mask.any(): - disc = retained.loc[disc_mask] - groups.loc[disc_mask] = disc.apply( - _lookup_labelled_discordance, axis=1, cache=labelled_cache - ).values - return groups - - db_cache = cast(DbDiscordanceCache, discordance_cache) - if db_keys is None: - raise ValueError("db_keys required for full-search retained-group assignment") - is_match = retained["prediction"].map(lambda p: _is_db_match(str(p), db_keys)) - groups = pd.Series("Database match", index=retained.index) - disc_mask = ~is_match - if disc_mask.any(): - groups.loc[disc_mask] = _classify_predictions_vs_db( - retained.loc[disc_mask, "prediction"], - db_cache, - ).values - return groups - - -def _subsample_violin_groups(df: pd.DataFrame, category_col: str) -> pd.DataFrame: - """Limit points per category so violin plots stay responsive.""" - parts: list[pd.DataFrame] = [] - for _cat, group in df.groupby(category_col, observed=True): - if len(group) > _MAX_VIOLIN_PSMs_PER_GROUP: - group = group.sample( - n=_MAX_VIOLIN_PSMs_PER_GROUP, - random_state=_VIOLIN_SUBSAMPLE_SEED, - ) - parts.append(group) - return pd.concat(parts, ignore_index=True) if parts else df - - -def _plot_novel_violins_at_fdr( - df: pd.DataFrame, - db_scored: pd.DataFrame | None, - discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, - available: list[tuple[str, str]], - *, - project: str, - display: str, - eval_type: str, - plots_dir: Path, - fdr_t: float, - labelled_subset: bool, -) -> None: - """Render one FDR-threshold novel-feature violin figure, or skip if too sparse.""" - retained = df[df["psm_q_value"] <= fdr_t].copy() - if len(retained) < _MIN_VIOLIN_GROUP_SIZE: - logger.info( - "%s: skip violins at %d%% FDR (n=%d retained)", - project, - int(fdr_t * 100), - len(retained), - ) - return - - match_keys: set[str] | None = None - if not labelled_subset: - assert db_scored is not None - match_keys = _unique_peptides_at_fdr( - db_scored, "sequence", DB_Q_VALUE_COL, fdr_t - ) - - groups = _assign_retained_groups( - retained, - match_keys, - discordance_cache, - labelled_subset=labelled_subset, - ) - retained = retained.assign(_overlap_group=groups) - plot_df = retained[ - retained["_overlap_group"].isin(["Database match", "fully_discordant"]) - ].copy() - plot_df["Category"] = plot_df["_overlap_group"].map( - { - "Database match": "Database match", - "fully_discordant": "Novel", - } - ) - plot_df = _subsample_violin_groups(plot_df, "Category") - - n_match = (plot_df["Category"] == "Database match").sum() - n_novel = (plot_df["Category"] == "Novel").sum() - if n_match < _MIN_VIOLIN_GROUP_SIZE or n_novel < _MIN_VIOLIN_GROUP_SIZE: - logger.info( - "%s: skip violins at %d%% FDR (match=%d, novel=%d)", - project, - int(fdr_t * 100), - n_match, - n_novel, - ) - return - - n_feats = len(available) - n_cols = 4 - n_rows = int(np.ceil(n_feats / n_cols)) - fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows)) - axes_flat = np.atleast_1d(axes).flatten() - palette = {"Database match": _CORRECT_COLOUR, "Novel": _NOVEL_COLOUR} - cat_order = ["Database match", "Novel"] - - for ax, (col, label) in zip(axes_flat, available): - sub = plot_df[[col, "Category"]].dropna() - if sub["Category"].nunique() < 2: - ax.set_visible(False) - continue - sns.violinplot( - data=sub, - x="Category", - y=col, - order=cat_order, - palette=palette, - ax=ax, - inner="quartile", - cut=0, - linewidth=0.8, - ) - ax.set_xlabel("") - ax.set_ylabel(label) - ax.tick_params(axis="x", rotation=15) - _style_ax(ax) - - for ax in axes_flat[len(available) :]: - ax.set_visible(False) - - pct = int(fdr_t * 100) - fig.suptitle( - f"{display} ({_eval_type_display(eval_type)}): " - f"database-matched vs novel features at {pct}% FDR", - fontsize=12, - ) - fig.tight_layout() - _save_fig(fig, plots_dir / f"novel_feature_violins_{project}_fdr{pct}") - - -def plot_novel_feature_violins( - df: pd.DataFrame, - db_df: pd.DataFrame | None, - residue_masses: dict[str, float] | None, - discordance_cache: LabelledDiscordanceCache | DbDiscordanceCache, - project: str, - eval_type: str, - plots_dir: Path, - *, - labelled_subset: bool = False, -) -> None: - """Violin plots: database-matched vs fully discordant retained PSMs. - - Winnow retention uses non-parametric FDR on calibrated confidence. For - full-search comparisons, database match uses peptides retained at the same - nominal FDR via database-grounded FDR on raw confidence (see overlap table). - """ - del residue_masses # kept for call-site compatibility - available = [ - (col, label) for col, label in NOVEL_FEATURE_COLUMNS if col in df.columns - ] - if not available: - logger.warning( - "%s: no feature columns for violin plots (metadata merge missing?)", - project, - ) - return - - df = _add_q_values(df.copy()) - display = _display_name(project) - - db_scored: pd.DataFrame | None = None - if not labelled_subset: - if db_df is None: - raise ValueError("Full-search violin plots require db_df") - db_scored = _add_database_grounded_q_values( - db_df.copy(), - confidence_col=RAW_CONFIDENCE_COL, - q_col=DB_Q_VALUE_COL, - ) - - for fdr_t in FDR_THRESHOLDS: - _plot_novel_violins_at_fdr( - df, - db_scored, - discordance_cache, - available, - project=project, - display=display, - eval_type=eval_type, - plots_dir=plots_dir, - fdr_t=fdr_t, - labelled_subset=labelled_subset, - ) - - -# --------------------------------------------------------------------------- -# Per-project orchestration -# --------------------------------------------------------------------------- -def generate_all_analyses( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, - db_df: pd.DataFrame | None, - residue_masses: dict[str, float], -) -> pd.DataFrame: - """Tables and plots for one project.""" - output_dir.mkdir(parents=True, exist_ok=True) - plots_dir = output_dir / "plots" - plots_dir.mkdir(parents=True, exist_ok=True) - - if db_df is None: - raise ValueError(f"Full-search analysis requires DB reference for {project}") - db_keys, db_stripped_set, _, db_by_len = _build_db_reference_sets(db_df) - disc_cache = _discordance_cache_for_fdr_retained( - df, - db_keys, - db_stripped_set, - db_by_len, - labelled_subset=False, - residue_masses=None, - ) - overlap = compute_overlap_table( - df, - project, - eval_type, - db_df, - disc_cache, - residue_masses, - labelled_subset=False, - ) - plot_full_search_venn(df, db_df, project, plots_dir, residue_masses) - plot_novel_feature_violins( - df, - db_df, - residue_masses, - disc_cache, - project, - eval_type, - plots_dir, - labelled_subset=False, - ) - - overlap.to_csv(output_dir / f"{project}_overlap_summary.csv", index=False) - with open(output_dir / f"{project}_overlap_summary.json", "w") as f: - json.dump(overlap.to_dict(orient="records"), f, indent=2) - - logger.info("\n%s", overlap.to_string(index=False)) - return overlap - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- -@app.command() -def main( - unlabelled_dir: Annotated[ - Path, - typer.Option( - "--unlabelled-dir", - help="Root with per-project full-search folders (e.g. gluc_raw/, PXD014877_unlabelled/).", - ), - ], - labelled_dir: Annotated[ - Path, - typer.Option( - "--labelled-dir", - help="Root with per-project database-search folders (e.g. gluc_annotated/, PXD014877_labelled/).", - ), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for tables and plots."), - ], - projects: Annotated[ - str | None, - typer.Option( - "--projects", - help="Optional space- or comma-separated project keys to restrict analysis.", - ), - ] = None, -) -> None: - """Post-FDR overlap: full-search Winnow identifications vs database search.""" - logging.basicConfig(level=logging.INFO, format="%(message)s", datefmt="%H:%M:%S") - - projects_filter: set[str] | None = None - if projects is not None: - project_list = [ - p.strip() for p in projects.replace(",", " ").split() if p.strip() - ] - if not project_list: - raise typer.BadParameter("No projects specified in --projects.") - projects_filter = set(project_list) - - pairs = discover_project_pairs( - unlabelled_dir, labelled_dir, projects_filter=projects_filter - ) - if not pairs: - logger.error( - "No paired projects found under unlabelled-dir=%s and labelled-dir=%s", - unlabelled_dir, - labelled_dir, - ) - raise typer.Exit(code=1) - - output_dir.mkdir(parents=True, exist_ok=True) - residue_masses = _get_residue_masses() - - all_tables: list[pd.DataFrame] = [] - for project, unlabelled_folder, labelled_folder, search_tag in pairs: - display = _display_name(project) - logger.info( - "Processing %s (%s): %s vs %s", - project, - display, - unlabelled_folder.name, - labelled_folder.name, - ) - - try: - df = _load_from_folder(unlabelled_folder) - db_df = _load_from_folder(labelled_folder) - except FileNotFoundError as exc: - logger.warning("Skipping %s: %s", project, exc) - continue - - logger.info( - " Full search: %d rows; DB reference: %d rows", len(df), len(db_df) - ) - - try: - table = generate_all_analyses( - df, - project, - search_tag, - output_dir, - db_df, - residue_masses, - ) - all_tables.append(table) - except ValueError as exc: - logger.warning("Skipping %s: %s", project, exc) - - if not all_tables: - logger.error("No projects produced overlap output under %s", output_dir) - raise typer.Exit(code=1) - - combined = pd.concat(all_tables, ignore_index=True) - combined.to_csv(output_dir / "all_projects_overlap_summary.csv", index=False) - with open(output_dir / "all_projects_overlap_summary.json", "w") as f: - json.dump(combined.to_dict(orient="records"), f, indent=2) - - logger.info("FDR overlap analysis complete. Output in %s", output_dir) - - -if __name__ == "__main__": - app() diff --git a/scripts/analyze_features.py b/scripts/analyze_features.py deleted file mode 100644 index 6fc4fec7..00000000 --- a/scripts/analyze_features.py +++ /dev/null @@ -1,804 +0,0 @@ -"""Analyze feature importance and correlations for a pretrained calibrator. - -This script provides comprehensive analysis of feature importance: - - Permutation importance on test set - - SHAP values with training background on test set - - Feature correlation analysis on training data - - Optional visualization of results -""" - -import logging -import pickle -from pathlib import Path -from typing import Annotated, Any, Dict, List, Optional - -import matplotlib.pyplot as plt -from matplotlib.colors import LinearSegmentedColormap -import numpy as np -import pandas as pd -import seaborn as sns -import shap -import torch -import typer -import yaml -from rich.console import Console -from rich.theme import Theme -from sklearn.inspection import permutation_importance - -from winnow.calibration.calibrator import ProbabilityCalibrator -from winnow.datasets.calibration_dataset import CalibrationDataset -from winnow.datasets.data_loaders import InstaNovoDatasetLoader - -# --------------------------------------------------------------------------- -# Style — Paul Tol "bright" palette + "sunset" diverging colourmap -# --------------------------------------------------------------------------- -_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] - -_SUNSET_COLORS = [ - "#364B9A", - "#4A7BB7", - "#6EA6CD", - "#98CAE1", - "#C2E4EF", - "#EAECCC", - "#FEDA8B", - "#FDB366", - "#F67E4B", - "#DD3D2D", - "#A50026", -] -_BAD_COLOUR = "#FFFFFF" - - -def _sunset_cmap() -> LinearSegmentedColormap: - cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) - cmap.set_bad(color=_BAD_COLOUR) - return cmap - - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - - -def _style_axes(ax: plt.Axes) -> None: - """Apply standard axes formatting: no grid, black spines.""" - ax.set_axisbelow(True) - ax.grid(False) - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -# --------------------------------------------------------------------------- -# Logging -# --------------------------------------------------------------------------- -logger = logging.getLogger("winnow") -logger.setLevel(logging.INFO) - -logging.getLogger("shap").setLevel(logging.WARNING) - -# --------------------------------------------------------------------------- -# Constants — loaded from the canonical Winnow YAML configs -# --------------------------------------------------------------------------- -SEED = 42 - -_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "winnow" / "configs" - -with open(_CONFIGS_DIR / "residues.yaml") as _f: - RESIDUE_MASSES: dict[str, float] = yaml.safe_load(_f)["residue_masses"] - -with open(_CONFIGS_DIR / "data_loader" / "instanovo.yaml") as _f: - _instanovo_cfg = yaml.safe_load(_f) - RESIDUE_REMAPPING: dict[str, str] = _instanovo_cfg.get("residue_remapping", {}) - BEAM_COLUMNS: dict[str, str] | None = _instanovo_cfg.get("beam_columns") - -COLUMN_DISPLAY_NAMES = { - "confidence": "Raw confidence", - "mass_error": "Mass error", - "mass_error_ppm": "Mass error (ppm)", - "mass_error_da": "Mass error (Da)", - "spectral_angle": "Spectral angle", - "ion_matches": "Ion matches", - "ion_match_intensity": "Ion match intensity", - "chimeric_ion_matches": "Chimeric ion matches", - "chimeric_ion_match_intensity": "Chimeric ion match intensity", - "irt_error": "iRT error", - "margin": "Margin", - "median_margin": "Median margin", - "entropy": "Entropy", - "z-score": "Z-score", - "edit_distance": "Edit distance", - "xcorr": "XCorr", - "chimeric_xcorr": "Chimeric XCorr", - "longest_ion_series": "Longest ion series", - "complementary_ion_count": "Complementary ion count", - "max_ion_gap": "Max ion gap", - "chimeric_longest_ion_series": "Chimeric longest ion series", - "chimeric_complementary_ion_count": "Chimeric complementary ion count", - "chimeric_max_ion_gap": "Chimeric max ion gap", - "is_missing_fragment_match_features": "Missing fragment match", - "is_missing_chimeric_features": "Missing chimeric", - "is_missing_irt_error": "Missing iRT", - "sequence_length": "Sequence length", - "precursor_charge": "Precursor charge", - "min_token_probability": "Min token probability", - "std_token_probability": "Std token probability", -} - -error_theme = Theme({"error": "red bold", "error_highlight": "red bold underline"}) -console = Console(theme=error_theme) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - - -def feature_display_name(column: str) -> str: - """Human-readable label for a feature column.""" - if column in COLUMN_DISPLAY_NAMES: - return COLUMN_DISPLAY_NAMES[column] - return column.replace("_", " ").replace("-", " ").strip().title() - - -def to_sentence_case(name: str) -> str: - """Convert a feature display name to sentence case.""" - return name.lower() - - -_SUPPORTED_EXTENSIONS = {".parquet", ".ipc", ".mgf"} - - -def _load_and_compute_features( - spectrum_path: Path, - predictions_path: Path, - loader: InstaNovoDatasetLoader, - calibrator: ProbabilityCalibrator, -) -> CalibrationDataset: - """Load spectra (single file or directory) and compute calibration features. - - When *spectrum_path* is a directory the contained spectrum files are - processed one at a time and the resulting metadata frames are - concatenated, mirroring the batched logic in - ``winnow.scripts.main._compute_features_directory``. - """ - if spectrum_path.is_dir(): - files = sorted( - f for f in spectrum_path.iterdir() if f.suffix in _SUPPORTED_EXTENSIONS - ) - if not files: - raise FileNotFoundError( - f"No spectrum files found in {spectrum_path}. " - f"Supported extensions: {', '.join(sorted(_SUPPORTED_EXTENSIONS))}" - ) - all_metadata: list[pd.DataFrame] = [] - for file_path in files: - logger.info(" Processing experiment file: %s", file_path.name) - ds = loader.load(data_path=file_path, predictions_path=predictions_path) - calibrator.compute_features(ds) - all_metadata.append(ds.metadata) - combined = pd.concat(all_metadata, ignore_index=True) - return CalibrationDataset(metadata=combined, predictions=None) - - dataset = loader.load(data_path=spectrum_path, predictions_path=predictions_path) - calibrator.compute_features(dataset) - return dataset - - -# --------------------------------------------------------------------------- -# Model wrapper for sklearn-compatible predict_proba -# --------------------------------------------------------------------------- -class _CalibratorPredictor: - """Wraps a fitted ProbabilityCalibrator as an sklearn-style estimator. - - Provides ``predict_proba`` and ``predict`` on *pre-normalised* feature - arrays so that permutation importance and SHAP can treat it like a - classifier. The ``classes_`` attribute is set to ``[0, 1]``. - """ - - def __init__(self, calibrator: ProbabilityCalibrator) -> None: - assert calibrator.network is not None - assert calibrator.feature_mean is not None - assert calibrator.feature_std is not None - - self.network = calibrator.network - self.feature_mean = calibrator.feature_mean.cpu() - self.feature_std = calibrator.feature_std.cpu() - self.classes_ = np.array([0, 1]) - - def fit(self, x_input: np.ndarray, y: np.ndarray) -> "_CalibratorPredictor": - """No-op fit to satisfy sklearn estimator interface.""" - return self - - def score(self, x_input: np.ndarray, y: np.ndarray) -> float: - """Return accuracy to satisfy sklearn estimator interface.""" - return float(np.mean(self.predict(x_input) == y)) - - def predict_proba(self, x_input: np.ndarray) -> np.ndarray: # noqa: N803 - """Return class probabilities for each sample.""" - x = torch.as_tensor(x_input, dtype=torch.float32) - self.network.eval() - with torch.no_grad(): - logits = self.network(x) - probs = torch.sigmoid(logits).numpy().ravel() - return np.column_stack([1 - probs, probs]) - - def predict(self, x_input: np.ndarray) -> np.ndarray: # noqa: N803 - """Return binary predictions for each sample.""" - return (self.predict_proba(x_input)[:, 1] >= 0.5).astype(int) - - -# --------------------------------------------------------------------------- -# Plotting functions -# --------------------------------------------------------------------------- -def plot_feature_importance( - importance_scores: Dict[str, float], - title: str, - output_path_base: Path, -) -> None: - """Plot horizontal bar chart of permutation feature importance scores.""" - plt.figure(figsize=(8, 6)) - features = list(importance_scores.keys()) - scores = list(importance_scores.values()) - - sorted_idx = np.argsort(scores) - features = [features[i] for i in sorted_idx] - scores = [scores[i] for i in sorted_idx] - - plt.barh(range(len(features)), scores, color=_PALETTE[0]) - plt.yticks(range(len(features)), features) - plt.xlabel("Importance score") - plt.title(title) - _style_axes(plt.gca()) - - plt.savefig(f"{output_path_base}.pdf", bbox_inches="tight", dpi=300) - plt.savefig(f"{output_path_base}.png", bbox_inches="tight", dpi=300) - plt.close() - - -def plot_feature_correlations(features: pd.DataFrame, output_path_base: Path) -> None: - """Plot lower-triangle feature correlation heatmap.""" - plt.figure(figsize=(12, 10)) - corr_matrix = features.corr() - mask = np.triu(np.ones_like(corr_matrix, dtype=bool)) - - sns.heatmap( - corr_matrix, - mask=mask, - cmap=_sunset_cmap(), - vmin=-1, - vmax=1, - center=0, - square=True, - annot=True, - fmt=".2f", - cbar_kws={"shrink": 0.5}, - ) - plt.title("Feature correlation matrix") - _style_axes(plt.gca()) - - plt.savefig(f"{output_path_base}.pdf", bbox_inches="tight", dpi=300) - plt.savefig(f"{output_path_base}.png", bbox_inches="tight", dpi=300) - plt.close() - - -def plot_shap_summary(shap_values, correct_class_idx: int, output_dir: Path) -> None: - """Plot SHAP beeswarm summary for the correct class.""" - plt.figure(figsize=(8, 6)) - shap.plots.beeswarm( - shap_values[:, :, correct_class_idx], - show=False, - max_display=12, - color=_sunset_cmap(), - ) - plt.title(r"SHAP feature impact on $P(\text{correct})$") - _style_axes(plt.gca()) - - plt.savefig(output_dir / "shap_summary.pdf", bbox_inches="tight", dpi=300) - plt.savefig(output_dir / "shap_summary.png", bbox_inches="tight", dpi=300) - plt.close() - - -def plot_shap_bar( - shap_values, - test_features_scaled, - test_labels, - correct_class_idx: int, - output_dir: Path, -) -> None: - """Plot SHAP bar chart with hierarchical clustering.""" - plt.figure(figsize=(8, 6)) - clustering = shap.utils.hclust(test_features_scaled, test_labels) - shap.plots.bar( - shap_values[:, :, correct_class_idx], - clustering=clustering, - show=False, - clustering_cutoff=0.5, - max_display=12, - ) - ax = plt.gca() - for patch in ax.patches: - patch.set_facecolor(_PALETTE[1]) - plt.title(r"SHAP feature importance for $P(\text{correct})$") - _style_axes(plt.gca()) - - plt.savefig(output_dir / "shap_importance.pdf", bbox_inches="tight", dpi=300) - plt.savefig(output_dir / "shap_importance.png", bbox_inches="tight", dpi=300) - plt.close() - - -def plot_shap_dependence( - shap_values, - feature_names: list, - display_feature_names: list, - correct_class_idx: int, - output_dir: Path, - top_n: int = 3, -) -> None: - """Plot SHAP dependence scatter for the top-N most important features.""" - mean_abs_shap = np.abs(shap_values.values[:, :, correct_class_idx]).mean(axis=0) - top_features_idx = np.argsort(mean_abs_shap)[-top_n:][::-1] - - for idx in top_features_idx: - feature_name = display_feature_names[idx] - original_feature_name = feature_names[idx] - plt.figure(figsize=(8, 6)) - shap.plots.scatter( - shap_values[:, idx, correct_class_idx], - show=False, - color=_PALETTE[2], - ) - plt.title( - "SHAP dependence plot for " - + to_sentence_case(feature_name) - + "\n" - + r"(impact on $P(\text{correct})$)" - ) - - ax = plt.gca() - ylabel = ax.get_ylabel() - if "SHAP value for" in ylabel: - ax.set_ylabel(ylabel.replace(feature_name, to_sentence_case(feature_name))) - - _style_axes(ax) - plt.savefig( - output_dir / f"shap_dependence_{original_feature_name}.pdf", - bbox_inches="tight", - dpi=300, - ) - plt.savefig( - output_dir / f"shap_dependence_{original_feature_name}.png", - bbox_inches="tight", - dpi=300, - ) - plt.close() - - -def plot_shap_interactions( - shap_values, - feature_names: list, - display_feature_names: list, - correct_class_idx: int, - output_dir: Path, - top_n: int = 3, -) -> None: - """Plot pairwise SHAP interaction scatter plots for top-N features.""" - mean_abs_shap = np.abs(shap_values.values[:, :, correct_class_idx]).mean(axis=0) - top_features_idx = np.argsort(mean_abs_shap)[-top_n:][::-1] - - for i, idx1 in enumerate(top_features_idx): - f1_display = display_feature_names[idx1] - f1_orig = feature_names[idx1] - for j, idx2 in enumerate(top_features_idx): - if i == j: - continue - f2_display = display_feature_names[idx2] - f2_orig = feature_names[idx2] - - plt.figure(figsize=(8, 6)) - shap.plots.scatter( - shap_values[:, f1_display, correct_class_idx], - color=shap_values[:, f2_display, correct_class_idx], - show=False, - cmap=_sunset_cmap(), - ) - plt.title( - "SHAP interaction plot for " - + to_sentence_case(f1_display) - + " vs " - + to_sentence_case(f2_display) - + "\n" - + r" (impact on $P(\text{correct})$)" - ) - - ax = plt.gca() - ylabel = ax.get_ylabel() - if "SHAP value for" in ylabel: - ax.set_ylabel(ylabel.replace(f1_display, to_sentence_case(f1_display))) - - _style_axes(ax) - plt.savefig( - output_dir / f"shap_interaction_{f1_orig}_vs_{f2_orig}.pdf", - bbox_inches="tight", - dpi=300, - ) - plt.savefig( - output_dir / f"shap_interaction_{f1_orig}_vs_{f2_orig}.png", - bbox_inches="tight", - dpi=300, - ) - plt.close() - - -def plot_shap_heatmap(shap_values, correct_class_idx: int, output_dir: Path) -> None: - """Plot SHAP heatmap showing per-sample feature contributions.""" - plt.figure(figsize=(8, 6)) - shap.plots.heatmap( - shap_values[:, :, correct_class_idx], - max_display=12, - show=False, - cmap=_sunset_cmap(), - ) - plt.title("SHAP feature impact heatmap\n" + r"(impact on $P(\text{correct})$)") - _style_axes(plt.gca()) - - plt.savefig(output_dir / "shap_heatmap.pdf", bbox_inches="tight", dpi=300) - plt.savefig(output_dir / "shap_heatmap.png", bbox_inches="tight", dpi=300) - plt.close() - - -def create_all_plots( - perm_importance_dict: Dict[str, float], - shap_values, - train_features_scaled_df: pd.DataFrame, - test_features_scaled: np.ndarray, - test_labels: np.ndarray, - feature_names: list, - display_feature_names: list, - correct_class_idx: int, - output_dir: Path, -) -> None: - """Generate all analysis plots (importance, correlations, SHAP).""" - logger.info("Creating plots...") - - plot_feature_importance( - perm_importance_dict, - "Permutation feature importance", - output_dir / "permutation_importance", - ) - plot_shap_summary(shap_values, correct_class_idx, output_dir) - plot_shap_bar( - shap_values, test_features_scaled, test_labels, correct_class_idx, output_dir - ) - plot_shap_dependence( - shap_values, feature_names, display_feature_names, correct_class_idx, output_dir - ) - plot_shap_interactions( - shap_values, feature_names, display_feature_names, correct_class_idx, output_dir - ) - plot_shap_heatmap(shap_values, correct_class_idx, output_dir) - plot_feature_correlations( - train_features_scaled_df, output_dir / "feature_correlations" - ) - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- - - -def _load_features_from_parquet( - path: Path, - feature_columns: list[str], -) -> tuple[np.ndarray, np.ndarray]: - """Load feature matrix and labels from a Parquet file or directory.""" - import polars as pl - - p = Path(path) - if p.is_dir(): - parquet_files = sorted(p.glob("*.parquet")) - if not parquet_files: - raise FileNotFoundError(f"No .parquet files found in {p}") - df = pl.concat([pl.read_parquet(f) for f in parquet_files]) - else: - df = pl.read_parquet(p) - - if "correct" not in df.columns: - raise ValueError(f"Parquet at {path} must contain a 'correct' column") - missing = [c for c in feature_columns if c not in df.columns] - if missing: - raise ValueError(f"Missing feature columns in Parquet: {missing}") - - features = df.select(feature_columns).to_numpy().astype(np.float32) - labels = df["correct"].to_numpy().astype(np.float32) - return features, labels - - -def _parse_koina_constants(raw: Optional[List[str]]) -> Optional[Dict[str, Any]]: - """Parse ``KEY=VALUE`` pairs into a dict, casting numeric strings.""" - if not raw: - return None - out: Dict[str, Any] = {} - for item in raw: - if "=" not in item: - raise typer.BadParameter( - f"Invalid --koina-input-constant format: '{item}'. Expected KEY=VALUE." - ) - key, value = item.split("=", 1) - try: - out[key] = int(value) - except ValueError: - try: - out[key] = float(value) - except ValueError: - out[key] = value - return out - - -@app.command() -def main( - model_path: Annotated[ - Path, typer.Option(help="Path to pretrained calibrator model directory.") - ], - output_dir: Annotated[ - Path, typer.Option(help="Directory to save analysis results and plots.") - ], - data_dir: Annotated[ - Optional[Path], - typer.Option( - help="Directory containing train and test data files (raw spectra path)." - ), - ] = None, - train_features_path: Annotated[ - Optional[Path], - typer.Option( - help="Path to pre-computed training feature Parquet (alternative to --data-dir)." - ), - ] = None, - test_features_path: Annotated[ - Optional[Path], - typer.Option( - help="Path to pre-computed test feature Parquet (alternative to --data-dir)." - ), - ] = None, - train_spectra: Annotated[ - str, typer.Option(help="Filename of training spectra parquet inside data-dir.") - ] = "general_train.parquet", - train_preds: Annotated[ - str, typer.Option(help="Filename of training predictions CSV inside data-dir.") - ] = "general_train_beams.csv", - test_spectra: Annotated[ - str, typer.Option(help="Filename of test spectra parquet inside data-dir.") - ] = "general_test.parquet", - test_preds: Annotated[ - str, typer.Option(help="Filename of test predictions CSV inside data-dir.") - ] = "general_test_beams.csv", - koina_url: Annotated[ - Optional[str], - typer.Option( - "--koina-url", - help="Override Koina server URL on loaded features (e.g. localhost:8500).", - ), - ] = None, - koina_ssl: Annotated[ - bool, - typer.Option( - "--koina-ssl/--no-koina-ssl", - help="Use SSL for Koina (disable for in-pod Triton).", - ), - ] = True, - koina_input_constant: Annotated[ - Optional[List[str]], - typer.Option( - help="Koina model input constant as KEY=VALUE (repeatable). " - "E.g. --koina-input-constant collision_energies=27 " - "--koina-input-constant fragmentation_types=HCD", - ), - ] = None, - n_background_samples: Annotated[ - int, typer.Option(help="Background samples for SHAP.", min=1, max=10000) - ] = 500, - n_test_samples: Annotated[ - int, typer.Option(help="Test samples for SHAP.", min=1, max=10000) - ] = 1000, - create_plots: Annotated[ - bool, typer.Option("--create-plots/--no-plots", help="Whether to create plots.") - ] = True, -) -> None: - """Analyze feature importance and correlations for a pretrained calibrator.""" - use_parquet = train_features_path is not None or test_features_path is not None - if use_parquet and (train_features_path is None or test_features_path is None): - raise typer.BadParameter( - "--train-features-path and --test-features-path must both be provided." - ) - if not use_parquet and data_dir is None: - raise typer.BadParameter( - "Either --data-dir or --train-features-path/--test-features-path must be provided." - ) - - output_dir.mkdir(parents=True, exist_ok=True) - - # Load calibrator - logger.info("Loading pretrained calibrator from %s", model_path) - calibrator = ProbabilityCalibrator.load(model_path) - - if koina_url is not None or not koina_ssl: - logger.warning( - "Ignoring --koina-url/--koina-ssl; Koina server overrides are no longer " - "supported on ProbabilityCalibrator (url=%s, ssl=%s).", - koina_url, - koina_ssl, - ) - - koina_constants = _parse_koina_constants(koina_input_constant) - if koina_constants: - logger.info("Applying Koina input constant overrides: %s", koina_constants) - calibrator.apply_koina_model_input_overrides( - model_input_constants=koina_constants, - ) - - # Build sklearn-compatible predictor wrapper - predictor = _CalibratorPredictor(calibrator) - - if use_parquet: - assert train_features_path is not None - assert test_features_path is not None - - feature_columns = ["confidence"] + calibrator.columns - - logger.info("Loading training features from Parquet: %s", train_features_path) - train_features, train_labels = _load_features_from_parquet( - train_features_path, - feature_columns, - ) - logger.info( - " %d training samples, %d features", - len(train_labels), - train_features.shape[1], - ) - - logger.info("Loading test features from Parquet: %s", test_features_path) - test_features, test_labels = _load_features_from_parquet( - test_features_path, - feature_columns, - ) - logger.info( - " %d test samples, %d features", len(test_labels), test_features.shape[1] - ) - - feature_names = feature_columns - else: - assert data_dir is not None - loader = InstaNovoDatasetLoader( - residue_masses=RESIDUE_MASSES, - residue_remapping=RESIDUE_REMAPPING, - beam_columns=BEAM_COLUMNS, - add_index_cols=True, - ) - - logger.info("Loading and computing features for training set...") - train_dataset = _load_and_compute_features( - data_dir / train_spectra, - data_dir / train_preds, - loader, - calibrator, - ) - train_features, train_labels = calibrator._extract_feature_matrix( - train_dataset, labelled=True - ) - - logger.info("Loading and computing features for test set...") - test_dataset = _load_and_compute_features( - data_dir / test_spectra, - data_dir / test_preds, - loader, - calibrator, - ) - test_features, test_labels = calibrator._extract_feature_matrix( - test_dataset, labelled=True - ) - - feature_names = [train_dataset.confidence_column] + calibrator.columns - - display_feature_names = [feature_display_name(name) for name in feature_names] - - assert calibrator.feature_mean is not None - assert calibrator.feature_std is not None - feature_mean = calibrator.feature_mean.cpu().numpy() - feature_std = calibrator.feature_std.cpu().numpy() - train_features_scaled = (train_features - feature_mean) / feature_std - test_features_scaled = (test_features - feature_mean) / feature_std - - correct_class_idx = 1 # class 1 = correct - - # 1. Permutation importance on test set - logger.info("Computing permutation importance on test set...") - perm_importance = permutation_importance( - predictor, - test_features_scaled, - test_labels, - n_repeats=10, - random_state=SEED, - n_jobs=-1, - ) - perm_importance_dict = dict( - zip(display_feature_names, perm_importance.importances_mean) - ) - - # 2. SHAP values - logger.info("Computing SHAP values...") - background = shap.sample( - train_features_scaled, - min(n_background_samples, len(train_features_scaled)), - random_state=SEED, - ) - - explainer = shap.KernelExplainer( - model=predictor.predict_proba, - data=background, - seed=SEED, - link="identity", - ) - - np.random.seed(SEED) - n_samples = min(n_test_samples, test_features_scaled.shape[0]) - indices = np.random.choice( - test_features_scaled.shape[0], size=n_samples, replace=False - ) - - shap_values = explainer(test_features_scaled[indices]) - - # Switch to original feature space for visualisation - shap_values.data = test_features[indices] - shap_values.feature_names = display_feature_names - - # 3. Feature correlations on training data - logger.info("Computing feature correlations on training data...") - train_features_scaled_df = pd.DataFrame( - train_features_scaled, columns=display_feature_names - ) - - if create_plots: - create_all_plots( - perm_importance_dict=perm_importance_dict, - shap_values=shap_values, - train_features_scaled_df=train_features_scaled_df, - test_features_scaled=test_features_scaled, - test_labels=test_labels, - feature_names=feature_names, - display_feature_names=display_feature_names, - correct_class_idx=correct_class_idx, - output_dir=output_dir, - ) - - # Save raw objects - logger.info("Saving raw analysis objects...") - - with open(output_dir / "perm_importance.pkl", "wb") as f: - pickle.dump(perm_importance, f) - - with open(output_dir / "shap_values.pkl", "wb") as f: - pickle.dump(shap_values, f) - - logger.info("Analysis complete!") - logger.info("Results saved to %s", output_dir) - logger.info( - "Permutation Feature Importance: computed on %d test samples", len(test_labels) - ) - logger.info( - "SHAP values: computed on %d test samples with %d training samples as background", - n_samples, - len(background), - ) - logger.info( - "Correlation matrix: computed on %d training samples", len(train_labels) - ) - - saved_files = ["perm_importance.pkl", "shap_values.pkl"] - if create_plots: - saved_files.append("All plots in PDF and PNG formats") - else: - logger.info("Plots were skipped (--no-plots flag used)") - - logger.info("Saved files: %s", ", ".join(saved_files)) - print(f"\nResults saved to {output_dir}") - - -if __name__ == "__main__": - app() diff --git a/scripts/analyze_novelty.py b/scripts/analyze_novelty.py deleted file mode 100644 index 81a619b5..00000000 --- a/scripts/analyze_novelty.py +++ /dev/null @@ -1,1812 +0,0 @@ -#!/usr/bin/env python3 -"""Analyse Winnow calibrator behaviour on out-of-distribution / novel peptides. - -Two analyses demonstrate that the calibrator does not penalise peptides absent -from the standard tryptic database-search training distribution: - -1. **Non-tryptic enzyme digest (``nontryptic_digest`` subcommand)** -- The model - was trained on tryptic data. Enzymes such as GluC, AspN, LysC or chymotrypsin cleave at - non-K/R sites, so retained peptides often have a C-terminus that is *not* - K or R. We classify predictions by whether their C-terminal residue - is tryptic (K/R) or non-tryptic, report terminus proportions before and - after FDR, compare raw InstaNovo versus Winnow calibrated scores, and - quantify calibration shifts (``calibrated_confidence - confidence``). - - Inputs are **full search space** Winnow predictions (acfm / unlabelled eval: - all candidate spectra, not the labelled database-search subset and not - acfm-minus-lcfm). The ``proteome_hit`` column flags predictions whose - stripped sequence occurs in the reference proteome FASTA; plots and tables - label this cohort **full search space**. - - *Non-tryptic* is defined solely by the C-terminal residue of the - mod-stripped, I/L-normalised prediction. N-terminal context is not - checked because positional information is lost in the substring proteome - match. - -2. **ProteomeTools-1 PXD004732 (``proteometools`` subcommand)** -- Synthetic - peptide library. The *lcfm* set contains database-search-confirmed - peptides; the *acfm* set contains all candidates. For each acfm - prediction we check whether it exactly matches, is a subsequence of, or - shares no overlap with any lcfm peptide. Subsequence matches are - validated novel identifications the search engine missed. - -3. **Summary (``summary`` subcommand)** -- Combines tables from both analyses - into a single grouped bar chart. -""" - -from __future__ import annotations - -import re -import warnings -from pathlib import Path -from typing import Annotated - -import ahocorasick -import matplotlib.pyplot as plt -from matplotlib.lines import Line2D -import numpy as np -import pandas as pd -import polars as pl -import seaborn as sns -import typer -from Bio import SeqIO -from scipy.stats import gaussian_kde - -from winnow.fdr.nonparametric import NonParametricFDRControl - -warnings.filterwarnings("ignore", module="winnow") - -# ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── -_PALETTE = [ - "#4477AA", - "#EE6677", - "#228833", - "#CCBB44", - "#66CCEE", - "#AA3377", - "#BBBBBB", -] -_CORRECT_COLOUR = _PALETTE[0] -_INCORRECT_COLOUR = _PALETTE[1] -_NOVEL_COLOUR = _PALETTE[2] -_MAIN_LINE_COLOUR = _PALETTE[3] -_RAW_LINE_COLOUR = _PALETTE[5] -_IDEAL_LINE_COLOUR = _PALETTE[6] - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -_REPO_ROOT = Path(__file__).resolve().parent.parent -_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)") -_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") -_PROTEOME_JOIN_SEP = "\x1f" - -FDR_THRESHOLDS = [0.01, 0.05, 0.10] - -# Display label for the ``proteome_hit`` cohort (full-search / acfm evaluation). -FULL_SEARCH_SPACE_LABEL = "full search space" -COHORT_FULL_SEARCH = "full_search_space" -COHORT_FULL_SEARCH_AT_FDR = "full_search_space_at_fdr" - -_NONTRYPTIC_CALIBRATION_SCATTER_MAX_POINTS = 10_000 -_NONTRYPTIC_CALIBRATION_SCATTER_RANDOM_STATE = 42 -_C_TERMINUS_ORDER = ("Tryptic (K/R)", "Non-tryptic") -_NONTRYPTIC_GROUP_ALPHA = 0.6 -_TRYPTIC_GROUP_ALPHA = 0.7 - -FEATURE_COLUMNS = [ - "spectral_angle", - "xcorr", - "ion_matches", - "ion_match_intensity", - "irt_error", - "mass_error_ppm", -] - -DATASET_DISPLAY_NAMES: dict[str, str] = { - "PXD004732": "ProteomeTools-1", -} - -app = typer.Typer( - add_completion=False, - no_args_is_help=True, - pretty_exceptions_show_locals=False, -) - - -# ── Shared helpers ──────────────────────────────────────────────────── - - -def _spine_fmt(ax: plt.Axes) -> None: - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -def _save(fig: plt.Figure, out_dir: Path, name: str) -> None: - base = out_dir / name - fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - print(f" saved {name}") - - -def _subsample_psms( - df: pd.DataFrame, - max_points: int, - random_state: int = _NONTRYPTIC_CALIBRATION_SCATTER_RANDOM_STATE, -) -> pd.DataFrame: - """Return up to ``max_points`` rows without replacement.""" - if len(df) <= max_points: - return df - return df.sample(n=max_points, random_state=random_state) - - -def _strip_mods(seq: str) -> str: - """Strip PTM annotations and normalise I -> L.""" - if not seq or not isinstance(seq, str): - return "" - s = _MOD_PLUS.sub("", seq) - s = _MOD_UNIMOD.sub("", s) - return s.replace("I", "L") - - -def _load_data(predictions_dir: Path) -> pl.DataFrame: - """Load and join ``preds_and_fdr_metrics.csv`` + ``metadata.csv``.""" - preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") - meta_path = predictions_dir / "metadata.csv" - if meta_path.exists(): - meta = pl.read_csv(meta_path) - join_cols = ["spectrum_id"] + [ - c for c in meta.columns if c != "spectrum_id" and c not in preds.columns - ] - if len(join_cols) > 1: - preds = preds.join(meta.select(join_cols), on="spectrum_id", how="inner") - return preds - - -def _add_q_values( - df: pd.DataFrame, - conf_col: str = "calibrated_confidence", -) -> pd.DataFrame: - """Fit non-parametric FDR and append ``psm_q_value`` if missing.""" - if "psm_q_value" in df.columns: - return df - fdr = NonParametricFDRControl() - fdr.fit(dataset=df[conf_col]) - return fdr.add_psm_q_value(df, confidence_col=conf_col) - - -def _load_proteome_haystack(fasta_file: Path) -> str: - """Load a FASTA proteome into a single string for substring matching.""" - parts: list[str] = [] - for record in SeqIO.parse(fasta_file, "fasta"): - s = str(record.seq).replace("I", "L") - if s: - parts.append(s) - return _PROTEOME_JOIN_SEP.join(parts) - - -def _batch_substring_hits( - needles: list[str], - haystack: str, -) -> list[bool]: - """Aho-Corasick batch substring matching.""" - n = len(needles) - out = [False] * n - if not haystack: - return out - - by_needle: dict[str, list[int]] = {} - for i, p in enumerate(needles): - if not p: - continue - by_needle.setdefault(p, []).append(i) - if not by_needle: - return out - - auto = ahocorasick.Automaton() - needle_for_pid: list[str] = [] - for pid, needle in enumerate(by_needle): - auto.add_word(needle, pid) - needle_for_pid.append(needle) - auto.make_automaton() - - matched_pids: set[int] = set() - for _end_idx, pid in auto.iter(haystack): - matched_pids.add(pid) - - for pid in matched_pids: - needle = needle_for_pid[pid] - for row_i in by_needle[needle]: - out[row_i] = True - return out - - -def _nice_feature_label(col: str) -> str: - labels = { - "spectral_angle": "Spectral angle", - "xcorr": "Cross-correlation", - "ion_matches": "Ion match rate", - "ion_match_intensity": "Ion match intensity", - "irt_error": "iRT prediction error", - "mass_error_ppm": "Mass error (ppm)", - } - return labels.get(col, col.replace("_", " ").capitalize()) - - -# ── Non-tryptic digest analysis ───────────────────────────────────────── - - -def _is_tryptic_cterm(seq: str) -> bool: - """Return True if the mod-stripped C-terminal residue is K or R.""" - stripped = _strip_mods(seq) - if not stripped: - return False - return stripped[-1] in ("K", "R") - - -def _nontryptic_full_search_panel_title(fdr_t: float, n_psms: int) -> str: - """Subplot title for FDR-filtered full-search-space PSMs.""" - pct = int(fdr_t * 100) - return ( - f"{FULL_SEARCH_SPACE_LABEL.title()} identifications at {pct}% FDR\n" - f"(n={n_psms:,})" - ) - - -def _nontryptic_annotate( - df: pl.DataFrame, - fasta_path: Path, -) -> pl.DataFrame: - """Annotate full-search predictions with proteome and tryptic-terminus flags. - - ``proteome_hit`` is True when the mod-stripped prediction is a substring of - the reference proteome FASTA (same rule as ``winnow.utils.proteome``). - """ - haystack = _load_proteome_haystack(fasta_path) - - processed = df["prediction"].map_elements( - lambda x: _strip_mods(x) if isinstance(x, str) else "", - return_dtype=pl.Utf8, - ) - hits = _batch_substring_hits(processed.to_list(), haystack) - tryptic = df["prediction"].map_elements( - lambda x: _is_tryptic_cterm(x) if isinstance(x, str) else False, - return_dtype=pl.Boolean, - ) - return df.with_columns( - pl.Series("proteome_hit", hits, dtype=pl.Boolean), - tryptic.alias("tryptic_cterm"), - ) - - -def _terminus_count_row( - sub: pd.DataFrame, - *, - cohort: str, - fdr_threshold: float | None, -) -> dict: - """Return one row of tryptic / non-tryptic counts for *sub*.""" - n = len(sub) - n_tryp = int(sub["tryptic_cterm"].sum()) if n > 0 else 0 - n_non = n - n_tryp - return { - "cohort": cohort, - "fdr_threshold": fdr_threshold, - "n": n, - "n_tryptic": n_tryp, - "n_non_tryptic": n_non, - "pct_tryptic": round(n_tryp / n * 100, 2) if n > 0 else 0.0, - "pct_non_tryptic": round(n_non / n * 100, 2) if n > 0 else 0.0, - } - - -def _nontryptic_terminus_proportions_table(df: pd.DataFrame) -> pd.DataFrame: - """Tryptic versus non-tryptic counts across cohorts and FDR cutoffs.""" - df = _add_q_values(df) - rows: list[dict] = [ - _terminus_count_row(df, cohort="all_predictions", fdr_threshold=None), - _terminus_count_row( - df[df["proteome_hit"]], - cohort=COHORT_FULL_SEARCH, - fdr_threshold=None, - ), - ] - for fdr_t in FDR_THRESHOLDS: - retained = df[df["psm_q_value"] <= fdr_t] - rows.append( - _terminus_count_row( - retained, - cohort="retained_at_fdr", - fdr_threshold=fdr_t, - ) - ) - rows.append( - _terminus_count_row( - retained[retained["proteome_hit"]], - cohort=COHORT_FULL_SEARCH_AT_FDR, - fdr_threshold=fdr_t, - ) - ) - return pd.DataFrame(rows) - - -def _nontryptic_calibration_delta_rows( - sub: pd.DataFrame, - *, - cohort: str, - fdr_threshold: float | None, -) -> list[dict]: - """Build calibration-shift summary rows for tryptic and non-tryptic groups.""" - out: list[dict] = [] - for tryptic, label in ((True, "tryptic"), (False, "non_tryptic")): - grp = sub[sub["tryptic_cterm"] == tryptic] - n = len(grp) - out.append( - { - "cohort": cohort, - "fdr_threshold": fdr_threshold, - "terminus_group": label, - "n": n, - "mean_confidence": ( - round(float(grp["confidence"].mean()), 4) if n > 0 else float("nan") - ), - "mean_calibrated_confidence": ( - round(float(grp["calibrated_confidence"].mean()), 4) - if n > 0 - else float("nan") - ), - "mean_delta_confidence": ( - round(float(grp["delta_confidence"].mean()), 4) - if n > 0 - else float("nan") - ), - "median_delta_confidence": ( - round(float(grp["delta_confidence"].median()), 4) - if n > 0 - else float("nan") - ), - } - ) - return out - - -def _nontryptic_calibration_delta_table(df: pd.DataFrame) -> pd.DataFrame: - """Mean calibration shift by terminus and cohort.""" - if "confidence" not in df.columns: - return pd.DataFrame() - - work = _add_q_values(df.copy()) - work["delta_confidence"] = work["calibrated_confidence"] - work["confidence"] - work = work.dropna( - subset=["confidence", "calibrated_confidence", "delta_confidence"] - ) - - rows: list[dict] = [] - rows.extend( - _nontryptic_calibration_delta_rows( - work, cohort="all_predictions", fdr_threshold=None - ) - ) - rows.extend( - _nontryptic_calibration_delta_rows( - work[work["proteome_hit"]], - cohort=COHORT_FULL_SEARCH, - fdr_threshold=None, - ) - ) - for fdr_t in FDR_THRESHOLDS: - retained = work[(work["psm_q_value"] <= fdr_t) & work["proteome_hit"]] - rows.extend( - _nontryptic_calibration_delta_rows( - retained, - cohort=COHORT_FULL_SEARCH_AT_FDR, - fdr_threshold=fdr_t, - ) - ) - return pd.DataFrame(rows) - - -def _nontryptic_summary_table(df: pd.DataFrame) -> pd.DataFrame: - """Build the tryptic-summary table at each FDR threshold.""" - df = _add_q_values(df) - rows: list[dict] = [] - for fdr_t in FDR_THRESHOLDS: - retained = df[df["psm_q_value"] <= fdr_t] - n_retained = len(retained) - hits = retained[retained["proteome_hit"]] - n_hit = len(hits) - tryptic_hits = hits[hits["tryptic_cterm"]] - non_tryptic_hits = hits[~hits["tryptic_cterm"]] - n_tryp = len(tryptic_hits) - n_non = len(non_tryptic_hits) - rows.append( - { - "fdr_threshold": fdr_t, - "n_retained": n_retained, - "n_full_search_space": n_hit, - "n_tryptic_hit": n_tryp, - "n_non_tryptic_hit": n_non, - "pct_non_tryptic_among_hits": ( - round(n_non / n_hit * 100, 2) if n_hit > 0 else 0.0 - ), - "mean_cal_conf_tryptic": ( - round(float(tryptic_hits["calibrated_confidence"].mean()), 4) - if n_tryp > 0 - else float("nan") - ), - "mean_cal_conf_non_tryptic": ( - round(float(non_tryptic_hits["calibrated_confidence"].mean()), 4) - if n_non > 0 - else float("nan") - ), - } - ) - return pd.DataFrame(rows) - - -def _nontryptic_score_label(score_col: str) -> str: - if score_col == "confidence": - return "Raw InstaNovo confidence" - return "Calibrated confidence" - - -def _finalize_nontryptic_violin_figure( - fig: plt.Figure, - axes: list[plt.Axes], - *, - y_label: str, - suptitle: str, -) -> None: - """Apply shared y-axis label and suptitle after multi-panel violin plots.""" - if len(axes) > 0: - axes[0].set_ylabel(y_label) - fig.tight_layout(rect=[0, 0, 1, 0.92]) - fig.suptitle(suptitle, fontsize=13, y=0.98) - - -def _plot_nontryptic_score_by_terminus( - df: pd.DataFrame, - out_dir: Path, - *, - score_col: str, - save_name: str, - suptitle: str, -) -> None: - """Violin plot of a score column split by C-terminal residue at each FDR.""" - if score_col not in df.columns: - print(f" skipping {save_name} (missing {score_col})") - return - - df = _add_q_values(df) - y_label = _nontryptic_score_label(score_col) - palette = {"Tryptic (K/R)": _MAIN_LINE_COLOUR, "Non-tryptic": _NOVEL_COLOUR} - - panels: list[tuple[float, pd.DataFrame]] = [] - for fdr_t in FDR_THRESHOLDS: - retained = df[(df["psm_q_value"] <= fdr_t) & df["proteome_hit"]] - retained = retained.dropna(subset=[score_col]) - if len(retained) < 5: - print( - f" skipping {int(fdr_t * 100)}% FDR panel in {save_name} " - f"(n={len(retained):,})" - ) - continue - retained = retained.copy() - retained["C-terminus"] = retained["tryptic_cterm"].map( - {True: "Tryptic (K/R)", False: "Non-tryptic"}, - ) - panels.append((fdr_t, retained)) - - if not panels: - print( - f" skipping {save_name} (no FDR panels with enough " - f"{FULL_SEARCH_SPACE_LABEL} PSMs)" - ) - return - - n_cols = len(panels) - fig, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 5), sharey=True) - if n_cols == 1: - axes = [axes] - - for ax, (fdr_t, retained) in zip(axes, panels): - sns.violinplot( - data=retained, - x="C-terminus", - y=score_col, - order=list(_C_TERMINUS_ORDER), - palette=palette, - ax=ax, - inner="quartile", - cut=0, - linewidth=0.8, - ) - ax.set_xlabel("") - ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) - ax.grid(False) - _spine_fmt(ax) - - _finalize_nontryptic_violin_figure(fig, axes, y_label=y_label, suptitle=suptitle) - _save(fig, out_dir, save_name) - - -def _plot_nontryptic_conf_by_terminus( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Violin plot of calibrated confidence split by C-terminal residue.""" - _plot_nontryptic_score_by_terminus( - df, - out_dir, - score_col="calibrated_confidence", - save_name=f"{prefix}_conf_by_terminus", - suptitle=( - f"Calibrated confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" - "by C-terminal residue" - ), - ) - - -def _plot_nontryptic_raw_conf_by_terminus( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Violin plot of raw InstaNovo confidence split by C-terminal residue.""" - _plot_nontryptic_score_by_terminus( - df, - out_dir, - score_col="confidence", - save_name=f"{prefix}_raw_conf_by_terminus", - suptitle=( - f"Raw InstaNovo confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" - "by C-terminal residue" - ), - ) - - -def _plot_nontryptic_overlapping_score_histogram( - tryp: np.ndarray, - non_tryp: np.ndarray, - *, - score_col: str, - title: str, - out_dir: Path, - save_name: str, -) -> None: - """Overlapping tryptic / non-tryptic histogram with KDE overlays.""" - if len(tryp) + len(non_tryp) < 2: - print(f" skipping {save_name} (too few PSMs)") - return - - fig, ax = plt.subplots(figsize=(7, 5)) - bins = 50 - - ax.hist( - non_tryp, - bins=bins, - alpha=_NONTRYPTIC_GROUP_ALPHA, - label=f"Non-tryptic (n={len(non_tryp):,})", - density=False, - edgecolor="black", - color=_NOVEL_COLOUR, - ) - ax.hist( - tryp, - bins=bins, - alpha=_TRYPTIC_GROUP_ALPHA, - label=f"Tryptic (K/R, n={len(tryp):,})", - density=False, - edgecolor="black", - color=_MAIN_LINE_COLOUR, - ) - - all_vals = np.concatenate([tryp, non_tryp]) if len(non_tryp) else tryp - x_min, x_max = float(all_vals.min()), float(all_vals.max()) - if x_max <= x_min: - x_max = x_min + 1e-6 - x_grid = np.linspace(x_min, x_max, 300) - bin_width = (x_max - x_min) / bins if bins > 1 else 1.0 - - if len(non_tryp) > 1: - y_non = gaussian_kde(non_tryp)(x_grid) * len(non_tryp) * bin_width - ax.plot(x_grid, y_non, color=_NOVEL_COLOUR, lw=1.5) - if len(tryp) > 1: - y_tryp = gaussian_kde(tryp)(x_grid) * len(tryp) * bin_width - ax.plot(x_grid, y_tryp, color=_MAIN_LINE_COLOUR, lw=1.5) - - ax.set_xlabel(_nontryptic_score_label(score_col)) - ax.set_ylabel("Frequency") - ax.set_title(title) - ax.legend(loc="upper center") - ax.grid(False) - _spine_fmt(ax) - fig.tight_layout() - _save(fig, out_dir, save_name) - - -def _plot_nontryptic_score_histograms( - df: pd.DataFrame, - out_dir: Path, - *, - score_col: str, - save_name: str, - title: str, - subset: pd.DataFrame | None = None, - min_psms: int = 10, -) -> None: - """Overlapping tryptic / non-tryptic histograms for a score column.""" - if score_col not in df.columns: - print(f" skipping {save_name} (missing {score_col})") - return - - retained = df if subset is None else subset - retained = retained.dropna(subset=[score_col]) - if len(retained) < min_psms: - print(f" skipping {save_name} (too few PSMs)") - return - - tryp = retained.loc[retained["tryptic_cterm"], score_col].to_numpy() - non_tryp = retained.loc[~retained["tryptic_cterm"], score_col].to_numpy() - _plot_nontryptic_overlapping_score_histogram( - tryp, - non_tryp, - score_col=score_col, - title=title, - out_dir=out_dir, - save_name=save_name, - ) - - -def _plot_nontryptic_score_histograms_at_fdr( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Calibrated confidence histogram at 5% FDR for full search space PSMs.""" - df = _add_q_values(df) - retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] - _plot_nontryptic_score_histograms( - df, - out_dir, - score_col="calibrated_confidence", - save_name=f"{prefix}_score_histograms", - title=( - f"Calibrated confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} " - "at 5% FDR" - ), - subset=retained, - ) - - -def _plot_nontryptic_raw_score_histograms_at_fdr( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Raw InstaNovo confidence histogram at 5% FDR for full search space PSMs.""" - df = _add_q_values(df) - retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] - _plot_nontryptic_score_histograms( - df, - out_dir, - score_col="confidence", - save_name=f"{prefix}_raw_score_histograms", - title=( - f"Raw InstaNovo confidence for {dataset_label} {FULL_SEARCH_SPACE_LABEL} " - "at 5% FDR" - ), - subset=retained, - ) - - -def _plot_nontryptic_full_score_histograms( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Full-dataset raw and calibrated histograms by C-terminal residue.""" - all_preds = df.dropna(subset=["calibrated_confidence"]) - _plot_nontryptic_score_histograms( - df, - out_dir, - score_col="calibrated_confidence", - save_name=f"{prefix}_calibrated_score_histogram_full", - title=( - f"Calibrated confidence for all {dataset_label} predictions\n" - "by C-terminal residue" - ), - subset=all_preds, - min_psms=2, - ) - if "confidence" not in df.columns: - print(f" skipping {prefix}_raw_score_histogram_full (missing confidence)") - return - raw_preds = df.dropna(subset=["confidence"]) - _plot_nontryptic_score_histograms( - df, - out_dir, - score_col="confidence", - save_name=f"{prefix}_raw_score_histogram_full", - title=( - f"Raw InstaNovo confidence for all {dataset_label} predictions\n" - "by C-terminal residue" - ), - subset=raw_preds, - min_psms=2, - ) - - has_raw = "confidence" in df.columns - panels: list[tuple[str, str, pd.DataFrame]] = [ - ( - "calibrated_confidence", - "Calibrated confidence", - all_preds, - ), - ] - if has_raw: - panels.append(("confidence", "Raw InstaNovo confidence", raw_preds)) - - n_cols = len(panels) - fig, axes = plt.subplots(1, n_cols, figsize=(7 * n_cols, 5), sharey=True) - if n_cols == 1: - axes = [axes] - - bins = 50 - for ax, (score_col, y_label, subset) in zip(axes, panels): - tryp = subset.loc[subset["tryptic_cterm"], score_col].to_numpy() - non_tryp = subset.loc[~subset["tryptic_cterm"], score_col].to_numpy() - ax.hist( - non_tryp, - bins=bins, - alpha=_NONTRYPTIC_GROUP_ALPHA, - label=f"Non-tryptic (n={len(non_tryp):,})", - density=False, - edgecolor="black", - color=_NOVEL_COLOUR, - ) - ax.hist( - tryp, - bins=bins, - alpha=_TRYPTIC_GROUP_ALPHA, - label=f"Tryptic (K/R, n={len(tryp):,})", - density=False, - edgecolor="black", - color=_MAIN_LINE_COLOUR, - ) - ax.set_xlabel(y_label) - ax.set_ylabel("Frequency") - ax.set_title(f"All predictions (n={len(subset):,})") - ax.legend(loc="upper center", fontsize=9) - ax.grid(False) - _spine_fmt(ax) - - fig.tight_layout(rect=[0, 0, 1, 0.92]) - fig.suptitle( - f"Score distributions for all {dataset_label} predictions by C-terminal residue", - fontsize=13, - y=0.98, - ) - _save(fig, out_dir, f"{prefix}_score_histogram_full_panel") - - -def _plot_nontryptic_calibration_scatter( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Subsampled scatter of raw versus calibrated confidence by C-terminus.""" - if "confidence" not in df.columns: - print(f" skipping {prefix}_calibration_scatter (missing confidence)") - return - - work = df.dropna(subset=["confidence", "calibrated_confidence"]) - n_total = len(work) - if n_total < 10: - print(f" skipping {prefix}_calibration_scatter (too few PSMs)") - return - - plot_df = _subsample_psms(work, _NONTRYPTIC_CALIBRATION_SCATTER_MAX_POINTS) - n_show = len(plot_df) - fig, ax = plt.subplots(figsize=(7.5, 7)) - panels = [ - ("Non-tryptic", _NOVEL_COLOUR, False, _NONTRYPTIC_GROUP_ALPHA), - ("Tryptic (K/R)", _MAIN_LINE_COLOUR, True, _TRYPTIC_GROUP_ALPHA), - ] - for label, colour, tryptic, alpha in panels: - sub = plot_df.loc[plot_df["tryptic_cterm"] == tryptic] - if len(sub) == 0: - continue - ax.scatter( - sub["confidence"], - sub["calibrated_confidence"], - c=colour, - s=12, - alpha=alpha, - rasterized=True, - label=f"{label}", - ) - - ax.plot( - [-0.01, 1.01], - [-0.01, 1.01], - ls="--", - color="black", - lw=1, - label="No recalibration", - zorder=5, - ) - ax.set_xlim(-0.01, 1.01) - ax.set_ylim(-0.01, 1.01) - ax.set_xlabel("Raw InstaNovo confidence") - ax.set_ylabel("Calibrated confidence") - if n_show < n_total: - ax.set_title( - f"Raw vs calibrated confidence for all {dataset_label} predictions" - ) - else: - ax.set_title(f"All {dataset_label} predictions") - ax.legend(loc="lower right", fontsize=9) - ax.grid(False) - _spine_fmt(ax) - fig.tight_layout() - _save(fig, out_dir, f"{prefix}_calibration_scatter") - - -def _plot_nontryptic_delta_by_terminus( - df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Calibration shift by C-terminal residue.""" - if "confidence" not in df.columns: - print(f" skipping {prefix}_delta_by_terminus (missing confidence)") - return - - df = _add_q_values(df) - y_label = "Calibration shift" - palette = {"Tryptic (K/R)": _MAIN_LINE_COLOUR, "Non-tryptic": _NOVEL_COLOUR} - work = df.copy() - work["delta_confidence"] = work["calibrated_confidence"] - work["confidence"] - - panels: list[tuple[float, pd.DataFrame]] = [] - for fdr_t in FDR_THRESHOLDS: - retained = work[(work["psm_q_value"] <= fdr_t) & work["proteome_hit"]] - retained = retained.dropna(subset=["delta_confidence"]) - if len(retained) < 5: - print( - f" skipping {int(fdr_t * 100)}% FDR panel in " - f"{prefix}_delta_by_terminus (n={len(retained):,})" - ) - continue - retained = retained.copy() - retained["C-terminus"] = retained["tryptic_cterm"].map( - {True: "Tryptic (K/R)", False: "Non-tryptic"}, - ) - panels.append((fdr_t, retained)) - - if not panels: - print(f" skipping {prefix}_delta_by_terminus (no FDR panels with enough PSMs)") - return - - n_cols = len(panels) - fig, axes = plt.subplots(1, n_cols, figsize=(5 * n_cols, 5), sharey=True) - if n_cols == 1: - axes = [axes] - - for ax, (fdr_t, retained) in zip(axes, panels): - sns.violinplot( - data=retained, - x="C-terminus", - y="delta_confidence", - order=list(_C_TERMINUS_ORDER), - palette=palette, - ax=ax, - inner="quartile", - cut=0, - linewidth=0.8, - ) - ax.axhline(0.0, ls="--", color=_IDEAL_LINE_COLOUR, lw=1) - ax.set_xlabel("") - ax.set_title(_nontryptic_full_search_panel_title(fdr_t, len(retained))) - ax.grid(False) - _spine_fmt(ax) - - _finalize_nontryptic_violin_figure( - fig, - axes, - y_label=y_label, - suptitle=( - f"Winnow calibration shift for {dataset_label} {FULL_SEARCH_SPACE_LABEL} PSMs\n" - "by C-terminal residue" - ), - ) - _save(fig, out_dir, f"{prefix}_delta_by_terminus") - - -def _pooled_feature_mean_std(retained: pd.DataFrame, col: str) -> tuple[float, float]: - vals = retained[col].dropna() - if len(vals) == 0: - return float("nan"), float("nan") - if len(vals) == 1: - return float(vals.iloc[0]), float("nan") - return float(vals.mean()), float(vals.std()) - - -def _feature_group_median_z_row( - sub: pd.DataFrame, - available: list[str], - pooled: dict[str, tuple[float, float]], -) -> dict[str, float]: - row: dict[str, float] = {} - for col in available: - vals = sub[col].dropna() - if len(vals) == 0: - row[f"median_{col}"] = float("nan") - row[f"z_median_{col}"] = float("nan") - continue - med = float(vals.median()) - row[f"median_{col}"] = round(med, 4) - mu, std = pooled[col] - if np.isnan(std) or std == 0: - row[f"z_median_{col}"] = float("nan") - else: - row[f"z_median_{col}"] = round((med - mu) / std, 4) - return row - - -def _feature_median_z_score_table( - retained: pd.DataFrame, - available: list[str], - groups: list[tuple[str, str, pd.Series]], - *, - group_col: str, - reference: pd.DataFrame | None = None, -) -> pd.DataFrame: - """Per-group feature medians and z-scores relative to *reference* or *retained* PSMs.""" - pool_from = reference if reference is not None else retained - pooled = {col: _pooled_feature_mean_std(pool_from, col) for col in available} - - rows: list[dict] = [] - for group_key, _label, mask in groups: - sub = retained[mask] - row: dict = {group_col: group_key, "n": len(sub)} - row.update(_feature_group_median_z_row(sub, available, pooled)) - rows.append(row) - return pd.DataFrame(rows) - - -def _nontryptic_feature_table(df: pd.DataFrame) -> pd.DataFrame: - """Median feature values for tryptic vs non-tryptic full search space PSMs at 5% FDR.""" - df = _add_q_values(df) - retained = df[(df["psm_q_value"] <= 0.05) & df["proteome_hit"]] - available = [c for c in FEATURE_COLUMNS if c in retained.columns] - if not available: - return pd.DataFrame() - - return _feature_median_z_score_table( - retained, - available, - [ - ("tryptic", "Tryptic (K/R)", retained["tryptic_cterm"]), - ("non_tryptic", "Non-tryptic", ~retained["tryptic_cterm"]), - ], - group_col="group", - ) - - -def _plot_grouped_feature_z_scores( - feat_df: pd.DataFrame, - *, - group_col: str, - out_dir: Path, - save_name: str, - title: str, - group_style: list[tuple[str, str, str]], - z_score_ylabel: str = "Median z-score", -) -> None: - """Grouped bar chart of pooled z-scored feature medians.""" - if feat_df.empty: - print(f" skipping {save_name} (no feature data)") - return - - z_cols = [c for c in feat_df.columns if c.startswith("z_median_")] - if not z_cols: - print(f" skipping {save_name} (no z-scored feature columns)") - return - - plot_df = feat_df.set_index(group_col) - feature_labels = [_nice_feature_label(c.replace("z_median_", "")) for c in z_cols] - x = np.arange(len(z_cols)) - - present = [ - (key, label, colour) - for key, label, colour in group_style - if key in plot_df.index - ] - n_groups = len(present) - total_width = 0.7 - bar_w = total_width / max(n_groups, 1) - - fig, ax = plt.subplots(figsize=(9, 6.5)) - ax.axhline(0.0, color=_IDEAL_LINE_COLOUR, lw=0.8, zorder=0) - bar_groups = [] - for plot_i, (group_key, label, colour) in enumerate(present): - offset = (plot_i - (n_groups - 1) / 2) * bar_w - vals = plot_df.loc[group_key, z_cols].to_numpy(dtype=float) - bars = ax.bar( - x + offset, - vals, - bar_w, - label=label, - color=colour, - edgecolor="black", - linewidth=1, - ) - bar_groups.append(bars) - - if not bar_groups: - print(f" skipping {save_name} (no groups to plot)") - plt.close(fig) - return - - ax.set_xticks(x) - ax.set_xticklabels(feature_labels, rotation=30, ha="right") - ax.set_ylabel(z_score_ylabel) - ax.set_title(title) - ax.legend(loc="upper right", fontsize=9) - ax.grid(False) - _spine_fmt(ax) - fig.tight_layout() - _save(fig, out_dir, save_name) - - -def _plot_nontryptic_feature_comparison( - feat_df: pd.DataFrame, - out_dir: Path, - *, - prefix: str = "nontryptic_digest", - dataset_label: str = "Non-tryptic digest", -) -> None: - """Grouped bar chart of median features, tryptic vs non-tryptic.""" - _plot_grouped_feature_z_scores( - feat_df, - group_col="group", - out_dir=out_dir, - save_name=f"{prefix}_feature_comparison", - title=( - f"Median feature values for {dataset_label} tryptic versus " - f"non-tryptic {FULL_SEARCH_SPACE_LABEL} PSMs at 5% FDR" - ), - group_style=[ - ("tryptic", "Tryptic (K/R)", _MAIN_LINE_COLOUR), - ("non_tryptic", "Non-tryptic", _NOVEL_COLOUR), - ], - ) - - -def _nontryptic_digest_analysis( - predictions_dir: Path, - fasta: Path, - output_dir: Path, - *, - file_prefix: str, - dataset_label: str, -) -> None: - """Shared tryptic vs non-tryptic digest analysis (any non-tryptic enzyme dataset).""" - output_dir.mkdir(parents=True, exist_ok=True) - - print(f"Loading predictions from {predictions_dir}") - df_pl = _load_data(predictions_dir) - print(f" {df_pl.height:,} rows loaded") - - print(f"Annotating {FULL_SEARCH_SPACE_LABEL} predictions against {fasta}") - df_pl = _nontryptic_annotate(df_pl, fasta) - n_hits = df_pl.filter(pl.col("proteome_hit")).height - print(f" {n_hits:,} PSMs in {FULL_SEARCH_SPACE_LABEL} (proteome substring match)") - - df = df_pl.to_pandas() - - print("Building tryptic summary table") - summary = _nontryptic_summary_table(df) - summary.to_csv(output_dir / f"{file_prefix}_tryptic_summary.csv", index=False) - print(summary.to_string(index=False)) - - print("Building terminus proportion table") - prop_df = _nontryptic_terminus_proportions_table(df) - prop_df.to_csv(output_dir / f"{file_prefix}_terminus_proportions.csv", index=False) - print(prop_df.to_string(index=False)) - - if "confidence" in df.columns: - print("Building calibration shift table") - delta_df = _nontryptic_calibration_delta_table(df) - delta_df.to_csv( - output_dir / f"{file_prefix}_calibration_delta_summary.csv", index=False - ) - print(delta_df.to_string(index=False)) - else: - print(" skipping calibration shift table (missing raw confidence)") - - print("Building feature comparison table") - feat_df = _nontryptic_feature_table(df) - if not feat_df.empty: - feat_df.to_csv( - output_dir / f"{file_prefix}_feature_comparison.csv", index=False - ) - - plot_kw = {"prefix": file_prefix, "dataset_label": dataset_label} - print("Plotting") - _plot_nontryptic_conf_by_terminus(df, output_dir, **plot_kw) - _plot_nontryptic_raw_conf_by_terminus(df, output_dir, **plot_kw) - _plot_nontryptic_score_histograms_at_fdr(df, output_dir, **plot_kw) - _plot_nontryptic_raw_score_histograms_at_fdr(df, output_dir, **plot_kw) - _plot_nontryptic_full_score_histograms(df, output_dir, **plot_kw) - _plot_nontryptic_calibration_scatter(df, output_dir, **plot_kw) - _plot_nontryptic_delta_by_terminus(df, output_dir, **plot_kw) - _plot_nontryptic_feature_comparison(feat_df, output_dir, **plot_kw) - - print(f"\n{dataset_label} analysis complete. Output in {output_dir}") - - -@app.command() -def nontryptic_digest( - predictions_dir: Annotated[ - Path, - typer.Option( - "--predictions-dir", - help=( - "winnow predict output folder for the non-tryptic enzyme digest " - "(full search space / acfm)." - ), - ), - ], - fasta: Annotated[ - Path, - typer.Option("--fasta", help="Proteome FASTA for substring matching."), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for output tables and plots."), - ], - file_prefix: Annotated[ - str, - typer.Option( - "--file-prefix", - help="Prefix for output CSV and plot filenames (e.g. chymotrypsin).", - ), - ] = "nontryptic_digest", - dataset_label: Annotated[ - str, - typer.Option( - "--dataset-label", - help="Human-readable dataset name used in plot titles.", - ), - ] = "Non-tryptic digest", -) -> None: - """Analyse calibrator behaviour on non-tryptic enzyme digest peptides.""" - _nontryptic_digest_analysis( - predictions_dir, - fasta, - output_dir, - file_prefix=file_prefix, - dataset_label=dataset_label, - ) - - -@app.command() -def chymotrypsin( - predictions_dir: Annotated[ - Path, - typer.Option( - "--predictions-dir", - help="winnow predict output folder for HeLa chymotrypsin (full search space / acfm).", - ), - ], - fasta: Annotated[ - Path, - typer.Option("--fasta", help="Human proteome FASTA for substring matching."), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for output tables and plots."), - ], -) -> None: - """Convenience wrapper for ``nontryptic_digest`` on HeLa chymotrypsin data.""" - _nontryptic_digest_analysis( - predictions_dir, - fasta, - output_dir, - file_prefix="chymotrypsin", - dataset_label="HeLa chymotrypsin", - ) - - -@app.command() -def gluc( - predictions_dir: Annotated[ - Path, - typer.Option( - "--predictions-dir", - help="[Deprecated] winnow predict output folder for GluC raw.", - ), - ], - fasta: Annotated[ - Path, - typer.Option("--fasta", help="Human proteome FASTA for substring matching."), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for output tables and plots."), - ], -) -> None: - """[Deprecated] Use ``nontryptic_digest`` with ``--file-prefix gluc``.""" - typer.secho( - "Warning: gluc is deprecated; use nontryptic_digest --file-prefix gluc", - fg=typer.colors.YELLOW, - err=True, - ) - _nontryptic_digest_analysis( - predictions_dir, - fasta, - output_dir, - file_prefix="gluc", - dataset_label="HeLa degradome", - ) - - -# ── ProteomeTools-1 analysis ──────────────────────────────────────────── - - -def _exact_match_category(fits_precursor: bool) -> str: - if fits_precursor: - return "exact_match_and_fits_precursor" - return "exact_match_and_no_precursor_fit" - - -def _subsequence_category(fits_precursor: bool) -> str: - if fits_precursor: - return "subsequence_and_fits_precursor" - return "subsequence_and_no_precursor_fit" - - -def _neither_category(fits_precursor: bool) -> str: - if fits_precursor: - return "neither_and_fits_precursor" - return "neither_and_no_precursor_fit" - - -def _classify_exact_matches( - categories: list[str], - predictions: list[str], - fits_precursor: list[bool], - lcfm_peptide_set: set[str], -) -> None: - for i, (peptide, fit) in enumerate(zip(predictions, fits_precursor)): - if peptide in lcfm_peptide_set: - categories[i] = _exact_match_category(fit) - - -def _classify_subsequence_matches( - categories: list[str], - predictions: list[str], - fits_precursor: list[bool], - lcfm_haystack: str, -) -> None: - remaining_indices = [ - i for i, category in enumerate(categories) if category == "neither" - ] - remaining_peps = [predictions[i] for i in remaining_indices] - remaining_fits = [fits_precursor[i] for i in remaining_indices] - if not remaining_peps: - return - - hits = _batch_substring_hits(remaining_peps, lcfm_haystack) - for idx, fit, hit in zip(remaining_indices, remaining_fits, hits): - if hit: - categories[idx] = _subsequence_category(fit) - - -def _classify_remaining_neither( - categories: list[str], - fits_precursor: list[bool], -) -> None: - for i, fit in enumerate(fits_precursor): - if categories[i] == "neither": - categories[i] = _neither_category(fit) - - -def _classify_predictions( - predictions: list[str], - fits_precursor: list[bool], - lcfm_peptide_set: set[str], - lcfm_haystack: str, -) -> list[str]: - """Classify each prediction by lcfm overlap and precursor mass fit (<20 ppm). - - Uses Aho-Corasick to find which predictions are substrings of at least - one lcfm peptide (the haystack is built by joining all lcfm peptides - with a separator). Unmatched predictions are split by precursor fit. - """ - categories = ["neither"] * len(predictions) - _classify_exact_matches(categories, predictions, fits_precursor, lcfm_peptide_set) - _classify_subsequence_matches( - categories, predictions, fits_precursor, lcfm_haystack - ) - _classify_remaining_neither(categories, fits_precursor) - return categories - - -def _proteometools_summary_table(df: pd.DataFrame) -> pd.DataFrame: - """Build novelty summary table at each FDR threshold.""" - df = _add_q_values(df) - rows: list[dict] = [] - for fdr_t in FDR_THRESHOLDS: - retained = df[df["psm_q_value"] <= fdr_t] - n = len(retained) - if n == 0: - rows.append( - { - "fdr_threshold": fdr_t, - "n_retained": 0, - "n_exact_match_and_no_precursor_fit": 0, - "n_exact_match_and_fits_precursor": 0, - "n_subsequence_and_no_precursor_fit": 0, - "n_subsequence_and_fits_precursor": 0, - "n_neither_and_no_precursor_fit": 0, - "n_neither_and_fits_precursor": 0, - "pct_exact_or_sub_and_fit_among_retained": 0.0, - } - ) - continue - - cats = retained["novelty_category"] - n_exact_and_no_fit = int((cats == "exact_match_and_no_precursor_fit").sum()) - n_exact_and_fits_precursor = int( - (cats == "exact_match_and_fits_precursor").sum() - ) - n_sub_and_no_fit = int((cats == "subsequence_and_no_precursor_fit").sum()) - n_sub_and_fits_precursor = int((cats == "subsequence_and_fits_precursor").sum()) - n_neither_and_no_fit = int((cats == "neither_and_no_precursor_fit").sum()) - n_neither_and_fits_precursor = int((cats == "neither_and_fits_precursor").sum()) - - rows.append( - { - "fdr_threshold": fdr_t, - "n_retained": n, - "n_exact_match_and_no_precursor_fit": n_exact_and_no_fit, - "n_exact_match_and_fits_precursor": n_exact_and_fits_precursor, - "n_subsequence_and_no_precursor_fit": n_sub_and_no_fit, - "n_subsequence_and_fits_precursor": n_sub_and_fits_precursor, - "n_neither_and_no_precursor_fit": n_neither_and_no_fit, - "n_neither_and_fits_precursor": n_neither_and_fits_precursor, - "pct_exact_or_sub_and_fit_among_retained": round( - (n_exact_and_fits_precursor + n_sub_and_fits_precursor) / n * 100, 2 - ) - if n > 0 - else 0.0, - } - ) - return pd.DataFrame(rows) - - -_PROTEOMETOOLS_CONF_PLOT_LABELS: dict[str, str] = { - "exact_match_and_no_precursor_fit": "ID-", - "exact_match_and_fits_precursor": "ID+", - "subsequence_and_no_precursor_fit": "Sub-", - "subsequence_and_fits_precursor": "Sub+", - "neither_and_no_precursor_fit": "Novel-", - "neither_and_fits_precursor": "Novel+", -} - -_PROTEOMETOOLS_CONF_CATEGORY_ORDER = list(_PROTEOMETOOLS_CONF_PLOT_LABELS.keys()) - - -def _proteometools_conf_category_legend(ax: plt.Axes) -> None: - handles = [ - Line2D([], [], color="none", label="ID: Exact sequence match to labelled set."), - Line2D([], [], color="none", label="Sub: Subsequence of labelled set peptide."), - Line2D([], [], color="none", label="Novel: No sequence match to labelled set."), - Line2D( - [], - [], - color="none", - label="+: Matches precursor mass within 20 ppm.", - ), - ] - ( - Line2D( - [], - [], - color="none", - label="-: Does not match precursor mass within 20 ppm.", - ), - ) - ax.legend( - handles=handles, - loc="upper left", - bbox_to_anchor=(1.02, 1.0), - borderaxespad=0, - frameon=True, - fontsize=9, - ) - - -def _plot_proteometools_conf_by_category( - df: pd.DataFrame, - out_dir: Path, -) -> None: - """Violin plot of calibrated confidence by novelty category (all unlabelled PSMs).""" - plot_df = df.dropna(subset=["calibrated_confidence"]).copy() - if len(plot_df) < 5: - print(" skipping proteometools_conf_by_category (too few PSMs)") - return - - palette = { - _PROTEOMETOOLS_CONF_PLOT_LABELS[k]: _PALETTE[i] - for i, k in enumerate(_PROTEOMETOOLS_CONF_CATEGORY_ORDER) - } - - plot_df["Category"] = plot_df["novelty_category"].map( - _PROTEOMETOOLS_CONF_PLOT_LABELS - ) - present_cats = [ - _PROTEOMETOOLS_CONF_PLOT_LABELS[c] - for c in _PROTEOMETOOLS_CONF_CATEGORY_ORDER - if _PROTEOMETOOLS_CONF_PLOT_LABELS[c] in plot_df["Category"].values - ] - if not present_cats: - print(" skipping proteometools_conf_by_category (no categories present)") - return - - conf = plot_df["calibrated_confidence"] - y_min, y_max = float(conf.min()), float(conf.max()) - - fig, ax = plt.subplots(figsize=(10, 5)) - sns.violinplot( - data=plot_df, - x="Category", - y="calibrated_confidence", - order=present_cats, - palette=palette, - ax=ax, - inner="quartile", - cut=0, - linewidth=0.8, - ) - ax.set_ylim(y_min, y_max) - ax.set_xlabel("") - ax.set_ylabel("Calibrated confidence") - ax.set_title( - "Calibrated confidence for ProteomeTools-1 predictions\nby novelty category" - ) - ax.grid(False) - _spine_fmt(ax) - _proteometools_conf_category_legend(ax) - - fig.tight_layout() - _save(fig, out_dir, "proteometools_conf_by_category") - - -def _plot_proteometools_hit_rate( - df: pd.DataFrame, - out_dir: Path, -) -> None: - """Line plot: validated hit rate (exact or subsequence match and fits precursor mass within 20ppm) by calibrated confidence decile.""" - if len(df) < 20: - print(" skipping proteometools_hit_rate_vs_conf (too few PSMs)") - return - - df = df.copy() - df["is_validated"] = df["novelty_category"].isin( - ["exact_match_and_fits_precursor", "subsequence_and_fits_precursor"] - ) - df["conf_decile"] = pd.qcut( - df["calibrated_confidence"], - q=10, - duplicates="drop", - ) - grouped = ( - df.groupby("conf_decile", observed=True) - .agg( - hit_rate=("is_validated", "mean"), - mid=("calibrated_confidence", "mean"), - ) - .sort_values("mid") - ) - - fig, ax = plt.subplots(figsize=(8, 6)) - ax.plot( - grouped["mid"], - grouped["hit_rate"], - color=_MAIN_LINE_COLOUR, - linewidth=1.5, - marker="o", - markersize=6, - label="Validated hit rate", - ) - overall = float(df["is_validated"].mean()) - ax.axhline( - overall, - color=_IDEAL_LINE_COLOUR, - lw=1, - linestyle="--", - label=f"Overall mean ({overall:.2%})", - ) - ax.set_xlabel("Mean calibrated confidence per decile") - ax.set_ylabel( - "Fraction validated\n(exact match or subsequence fitting precursor mass)" - ) - ax.set_title( - "Validated hit rate by calibrated confidence decile for ProteomeTools-1" - ) - ax.legend(loc="lower right") - ax.grid(False) - _spine_fmt(ax) - fig.tight_layout() - _save(fig, out_dir, "proteometools_hit_rate_vs_conf") - - -def _proteometools_feature_table( - df: pd.DataFrame, - labelled_df: pd.DataFrame, -) -> pd.DataFrame: - """Median features for exact / novel / neither at 5% FDR.""" - df = _add_q_values(df) - labelled_df = _add_q_values(labelled_df) - retained = df[df["psm_q_value"] <= 0.05] - labelled_ref = labelled_df[labelled_df["psm_q_value"] <= 0.05] - available = [c for c in FEATURE_COLUMNS if c in retained.columns] - if not available: - return pd.DataFrame() - - return _feature_median_z_score_table( - retained, - available, - [ - ( - "exact_match_and_fits_precursor", - "Exact match, fits precursor mass", - retained["novelty_category"] == "exact_match_and_fits_precursor", - ), - ( - "exact_match_and_no_precursor_fit", - "Exact match, no precursor mass fit", - retained["novelty_category"] == "exact_match_and_no_precursor_fit", - ), - ( - "subsequence_and_fits_precursor", - "Subsequence, precursor mass fit", - retained["novelty_category"] == "subsequence_and_fits_precursor", - ), - ( - "subsequence_and_no_precursor_fit", - "Subsequence, no precursor mass fit", - retained["novelty_category"] == "subsequence_and_no_precursor_fit", - ), - ( - "neither_and_fits_precursor", - "Novel, precursor mass fit", - retained["novelty_category"] == "neither_and_fits_precursor", - ), - ( - "neither_and_no_precursor_fit", - "Novel, no precursor mass fit", - retained["novelty_category"] == "neither_and_no_precursor_fit", - ), - ], - group_col="category", - reference=labelled_ref, - ) - - -def _plot_proteometools_feature_comparison( - feat_df: pd.DataFrame, - out_dir: Path, -) -> None: - """Grouped bar chart of median features by category.""" - _plot_grouped_feature_z_scores( - feat_df, - group_col="category", - out_dir=out_dir, - save_name="proteometools_feature_comparison", - title=( - "Median feature values for ProteomeTools-1 predictions by novelty " - "category at 5% FDR" - ), - group_style=[ - ( - "exact_match_and_fits_precursor", - "Exact match, precursor mass fit", - _PALETTE[1], - ), - ( - "exact_match_and_no_precursor_fit", - "Exact match, no precursor mass fit", - _PALETTE[0], - ), - ( - "subsequence_and_fits_precursor", - "Subsequence, precursor mass fit", - _PALETTE[3], - ), - ( - "subsequence_and_no_precursor_fit", - "Subsequence, no precursor mass fit", - _PALETTE[2], - ), - ("neither_and_fits_precursor", "Novel, precursor mass fit", _PALETTE[5]), - ( - "neither_and_no_precursor_fit", - "Novel, no precursor mass fit", - _PALETTE[4], - ), - ], - z_score_ylabel="Median z-score (vs labelled PSMs at 5% FDR)", - ) - - -@app.command() -def proteometools( - lcfm_predictions_dir: Annotated[ - Path, - typer.Option( - "--lcfm-predictions-dir", - help="winnow predict output for PXD004732 lcfm (labelled).", - ), - ], - acfm_predictions_dir: Annotated[ - Path, - typer.Option( - "--acfm-predictions-dir", - help="winnow predict output for PXD004732 acfm (unlabelled).", - ), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for output tables and plots."), - ], -) -> None: - """Analyse calibrator behaviour on ProteomeTools-1 novel identifications.""" - output_dir.mkdir(parents=True, exist_ok=True) - - print(f"Loading lcfm predictions from {lcfm_predictions_dir}") - lcfm_pl = _load_data(lcfm_predictions_dir) - print(f" {lcfm_pl.height:,} lcfm rows") - - print(f"Loading acfm predictions from {acfm_predictions_dir}") - acfm_pl = _load_data(acfm_predictions_dir) - print(f" {acfm_pl.height:,} acfm rows") - - print("Building lcfm peptide set for subsequence matching") - lcfm_sequences = lcfm_pl["sequence"].drop_nulls().to_list() - lcfm_peptide_set: set[str] = set() - for seq in lcfm_sequences: - stripped = _strip_mods(seq) - if stripped: - lcfm_peptide_set.add(stripped) - print(f" {len(lcfm_peptide_set):,} unique lcfm peptides") - - lcfm_haystack = _PROTEOME_JOIN_SEP.join(sorted(lcfm_peptide_set)) - - print("Classifying unlabelled predictions") - unlabelled_pl = acfm_pl.join( - lcfm_pl.select("spectrum_id"), on="spectrum_id", how="anti" - ) - unlabelled_pl = unlabelled_pl.with_columns( - (pl.col("delta_mass_ppm").abs() < 20).alias("fits_precursor") - ) - unlabelled_preds_raw = unlabelled_pl["prediction"].to_list() - unlabelled_preds_stripped = [ - _strip_mods(p) if isinstance(p, str) else "" for p in unlabelled_preds_raw - ] - fits_precursor = unlabelled_pl["fits_precursor"].to_list() - - categories = _classify_predictions( - unlabelled_preds_stripped, - fits_precursor, - lcfm_peptide_set, - lcfm_haystack, - ) - unlabelled_pl = unlabelled_pl.with_columns( - pl.Series("novelty_category", categories, dtype=pl.Utf8), - ) - df = unlabelled_pl.to_pandas() - labelled_df = lcfm_pl.to_pandas() - - print("Building novelty summary table") - summary = _proteometools_summary_table(df) - summary.to_csv(output_dir / "proteometools_novelty_summary.csv", index=False) - print(summary.to_string(index=False)) - - print("Building feature comparison table") - feat_df = _proteometools_feature_table(df, labelled_df) - if not feat_df.empty: - feat_df.to_csv(output_dir / "proteometools_feature_comparison.csv", index=False) - - print("Plotting") - _plot_proteometools_conf_by_category(df, output_dir) - _plot_proteometools_hit_rate(df, output_dir) - _plot_proteometools_feature_comparison(feat_df, output_dir) - - print(f"\nProteomeTools-1 analysis complete. Output in {output_dir}") - - -# ── Summary figure ──────────────────────────────────────────────────── - - -@app.command() -def summary( - proteometools_dir: Annotated[ - Path, - typer.Option( - "--proteometools-dir", - help="Output directory from the proteometools subcommand.", - ), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for the combined summary figure."), - ], - nontryptic_digest_dir: Annotated[ - Path | None, - typer.Option( - "--nontryptic-digest-dir", - help="Output directory from the nontryptic_digest subcommand.", - ), - ] = None, - chymotrypsin_dir: Annotated[ - Path | None, - typer.Option( - "--chymotrypsin-dir", - help="Legacy alias for --nontryptic-digest-dir (chymotrypsin outputs).", - ), - ] = None, - gluc_dir: Annotated[ - Path | None, - typer.Option( - "--gluc-dir", - help="Legacy alias for --nontryptic-digest-dir (gluc outputs).", - ), - ] = None, -) -> None: - """Produce a combined summary bar chart from both analyses.""" - output_dir.mkdir(parents=True, exist_ok=True) - - digest_dir = nontryptic_digest_dir or chymotrypsin_dir or gluc_dir - if digest_dir is None: - raise typer.BadParameter( - "Provide --nontryptic-digest-dir (or --chymotrypsin-dir / --gluc-dir)." - ) - - digest_csvs = sorted(digest_dir.glob("*_tryptic_summary.csv")) - if len(digest_csvs) != 1: - raise typer.BadParameter( - f"Expected exactly one *_tryptic_summary.csv under {digest_dir}, " - f"found {len(digest_csvs)}" - ) - digest_csv = digest_csvs[0] - pt_csv = proteometools_dir / "proteometools_novelty_summary.csv" - if not digest_csv.is_file(): - raise typer.BadParameter(f"Missing tryptic summary under {digest_dir}") - if not pt_csv.is_file(): - raise typer.BadParameter(f"Missing {pt_csv}") - - -if __name__ == "__main__": - app() diff --git a/scripts/analyze_upscored_fps.py b/scripts/analyze_upscored_fps.py deleted file mode 100644 index 184ae2ae..00000000 --- a/scripts/analyze_upscored_fps.py +++ /dev/null @@ -1,709 +0,0 @@ -#!/usr/bin/env python3 -"""Characterise up-scored false positives from Winnow calibration. - -For each labelled evaluation dataset (where both ``sequence`` and ``prediction`` -are available), this script quantifies the false positives that calibration -"rescues" into high-confidence regions and compares their feature profiles to -true positives. - -Inputs are ``winnow predict`` output folders, each containing -``preds_and_fdr_metrics.csv`` and ``metadata.csv``. -""" - -from __future__ import annotations - -import json -import logging -import re -from pathlib import Path -from typing import Annotated - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import seaborn as sns -import typer -import yaml -from instanovo.utils.metrics import Metrics -from instanovo.utils.residues import ResidueSet -from rich.logging import RichHandler - -from winnow.fdr.nonparametric import NonParametricFDRControl - -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -logger.propagate = False -if not logger.handlers: - logger.addHandler(RichHandler()) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -# --------------------------------------------------------------------------- -# Style — Paul Tol "bright" palette (colour-blind safe) -# --------------------------------------------------------------------------- -_PALETTE = [ - "#4477AA", - "#EE6677", - "#228833", - "#CCBB44", - "#66CCEE", - "#AA3377", - "#BBBBBB", -] -_CORRECT_COLOUR = _PALETTE[0] -_INCORRECT_COLOUR = _PALETTE[1] - -TP_COLOR = _CORRECT_COLOUR -FP_COLOR = _INCORRECT_COLOUR - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -_REPO_ROOT = Path(__file__).resolve().parent.parent -_MOD_PLUS = re.compile(r"\(\+\d+\.?\d*\)-?") -_MOD_UNIMOD = re.compile(r"\[UNIMOD:\d+\]-?") - -FDR_THRESHOLDS = [0.01, 0.05, 0.10] -_FEATURE_VIOLIN_FDR_THRESHOLDS = [0.05, 0.10] -_MIN_FEATURE_VIOLIN_PSMs = 20 - -# Max PSMs per correctness panel in the raw-vs-calibrated confidence scatter. -_CONFIDENCE_SCATTER_MAX_POINTS = 10_000 -_CONFIDENCE_SCATTER_RANDOM_STATE = 42 - -DATASET_DISPLAY_NAMES: dict[str, str] = { - "gluc": "HeLa degradome", - "helaqc": "HeLa single shot", - "herceptin": "Herceptin", - "immuno": "Immunopeptidomics-1", - "celegans": "$\\it{C.\\;elegans}$", - "sbrodae": "$\\it{Scalindua\\;brodae}$", - "PXD019483": "HepG2", - "snakevenoms": "Snake venomics", - "tplantibodies": "Therapeutic nanobodies", - "woundfluids": "Wound exudates", - "PXD004732": "ProteomeTools-1", - "PXD014877": "$\\it{C.\\;elegans}$", - "PXD023064": "Immunopeptidomics-2", - "astral": "Astral $\\it{E.\\;coli}$", - "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", - "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", - "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", - "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", - "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", -} - -_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") - -# new_eval_sets_results layout: lcfm/PXD004452//preds_and_fdr_metrics.csv -_PXD_ACCESSION_PREFIX = "PXD" - -FEATURE_COLUMNS_OF_INTEREST = [ - "spectral_angle", - "xcorr", - "ion_matches", - "ion_match_intensity", - "irt_error", - "mass_error_ppm", - "margin", - "entropy", - "confidence", -] - -_MASS_ERROR_COLUMNS = ("mass_error_da", "mass_error_ppm") - -_NICE_LABELS: dict[str, str] = { - "ion_matches": "Ion match rate", - "ion_match_intensity": "Ion match intensity", - "complementary_ion_count": "Complementary ion count", - "max_ion_gap": "Max ion gap", - "spectral_angle": "Spectral angle", - "xcorr": "Cross-correlation (XCorr)", - "mass_error_ppm": "Precursor mass error (ppm)", - "mass_error_da": "Precursor mass error (Da)", - "irt_error": "iRT prediction error", - "confidence": "Model confidence", - "margin": "Beam margin", - "median_margin": "Beam median margin", - "entropy": "Beam entropy", - "z-score": "Beam z-score", - "edit_distance": "Runner-up edit distance", - "min_token_probability": "Min. token probability", - "std_token_probability": "Std. token probability", -} - - -def _nice_label(col: str) -> str: - return _NICE_LABELS.get(col, col.replace("_", " ").capitalize()) - - -def _mass_error_column(df: pd.DataFrame, *, min_count: int = 10) -> str | None: - """Return ``mass_error_da`` or ``mass_error_ppm`` when present with enough data.""" - for col in _MASS_ERROR_COLUMNS: - if col in df.columns and df[col].notna().sum() > min_count: - return col - return None - - -def _violin_feature_columns(df: pd.DataFrame) -> list[str]: - """Feature list for violin plots, resolving Da vs ppm mass error.""" - cols: list[str] = [] - for col in FEATURE_COLUMNS_OF_INTEREST: - if col == "mass_error_ppm": - mass_col = _mass_error_column(df) - if mass_col is not None: - cols.append(mass_col) - else: - cols.append(col) - return cols - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- -def _get_residue_masses() -> dict[str, float]: - config_path = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" - with open(config_path) as f: - cfg = yaml.safe_load(f) - return cfg["residue_masses"] - - -def _save_fig(fig: plt.Figure, base_path: Path, fmt: str = "both") -> None: - if fmt in ("pdf", "both"): - fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) - if fmt in ("png", "both"): - fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def _style_ax(ax: plt.Axes) -> None: - ax.grid(False) - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -def _folder_display_name(folder_name: str) -> str: - """Map an evaluation folder name to a publication-ready dataset label.""" - key = folder_name - for suffix in _FOLDER_SUFFIXES: - if key.endswith(suffix): - key = key[: -len(suffix)] - break - return DATASET_DISPLAY_NAMES.get(key, key) - - -def _subsample_psms( - df: pd.DataFrame, - max_points: int, - random_state: int = _CONFIDENCE_SCATTER_RANDOM_STATE, -) -> pd.DataFrame: - """Return up to ``max_points`` rows without replacement.""" - if len(df) <= max_points: - return df - return df.sample(n=max_points, random_state=random_state) - - -def _strip_mods(seq: str) -> str: - if not seq or not isinstance(seq, str): - return "" - s = _MOD_PLUS.sub("", seq) - s = _MOD_UNIMOD.sub("", s) - return s.replace("I", "L") - - -def _project_key_from_folder(folder_name: str) -> str: - """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" - for suffix in _FOLDER_SUFFIXES: - if folder_name.endswith(suffix): - return folder_name[: -len(suffix)] - return folder_name - - -def _is_labelled_preds_folder(folder: Path) -> bool: - preds_csv = folder / "preds_and_fdr_metrics.csv" - if not preds_csv.is_file(): - return False - header = pd.read_csv(preds_csv, nrows=0).columns.tolist() - required = {"sequence", "prediction", "calibrated_confidence"} - return required.issubset(header) - - -def _register_labelled_folder(results: dict[str, Path], key: str, folder: Path) -> None: - """Register *folder* under *key*, warning on duplicate keys.""" - if key in results: - logger.warning( - "Duplicate labelled project key %r: %s and %s", - key, - results[key], - folder, - ) - return - results[key] = folder - - -def _discover_labelled_folders(root: Path) -> dict[str, Path]: - """Find folders with labelled ``preds_and_fdr_metrics.csv``. - - Supports flat project folders (``{root}/PXD004732/``) and nested per-run - layouts used by new eval sets (``{root}/PXD004452//``). - """ - results: dict[str, Path] = {} - if not root.is_dir(): - return results - - for child in sorted(root.iterdir()): - if not child.is_dir(): - continue - if _is_labelled_preds_folder(child): - _register_labelled_folder( - results, _project_key_from_folder(child.name), child - ) - continue - if not child.name.startswith(_PXD_ACCESSION_PREFIX): - continue - for run_dir in sorted(child.iterdir()): - if run_dir.is_dir() and _is_labelled_preds_folder(run_dir): - _register_labelled_folder(results, run_dir.name, run_dir) - return results - - -def _load_dataset(folder: Path) -> pd.DataFrame: - """Load and merge preds + metadata CSVs for a single evaluation folder.""" - preds = pd.read_csv(folder / "preds_and_fdr_metrics.csv") - meta_path = folder / "metadata.csv" - if meta_path.is_file(): - meta = pd.read_csv(meta_path) - join_cols = ["spectrum_id"] + [ - c for c in meta.columns if c != "spectrum_id" and c not in preds.columns - ] - if len(join_cols) > 1: - preds = preds.merge( - meta[join_cols].drop_duplicates(subset=["spectrum_id"]), - on="spectrum_id", - how="left", - ) - if "correct" not in preds.columns and {"sequence", "prediction"}.issubset( - preds.columns - ): - preds = preds.copy() - preds["correct"] = preds["sequence"] == preds["prediction"] - return preds - - -def _add_q_values( - df: pd.DataFrame, conf_col: str = "calibrated_confidence" -) -> pd.DataFrame: - """Fit non-parametric FDR and append ``psm_q_value`` if missing.""" - if "psm_q_value" in df.columns: - return df - fdr = NonParametricFDRControl() - fdr.fit(dataset=df[conf_col]) - df = fdr.add_psm_q_value(df, confidence_col=conf_col) - return df - - -# --------------------------------------------------------------------------- -# Analysis -# --------------------------------------------------------------------------- -def _upscored_summary_table( - df: pd.DataFrame, - delta_threshold: float, - dataset_name: str, -) -> pd.DataFrame: - """Build per-FDR-threshold summary of up-scored TP / FP counts.""" - df = _add_q_values(df) - upscored = df["delta_confidence"] > delta_threshold - - rows = [] - for fdr_t in FDR_THRESHOLDS: - passing = df["psm_q_value"] <= fdr_t - for label, mask in [ - ("all", pd.Series(True, index=df.index)), - ("up-scored", upscored), - ("not up-scored", ~upscored), - ]: - sub = df[mask & passing] - n = len(sub) - n_correct = int(sub["correct"].sum()) if "correct" in sub.columns else 0 - n_incorrect = n - n_correct - rows.append( - { - "dataset": dataset_name, - "fdr_threshold": fdr_t, - "subset": label, - "n_passing": n, - "n_correct": n_correct, - "n_incorrect": n_incorrect, - "pct_correct": round(n_correct / n * 100, 2) if n > 0 else 0.0, - } - ) - return pd.DataFrame(rows) - - -def _plot_confidence_scatter( - df: pd.DataFrame, - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """Subsampled scatter of raw vs calibrated confidence, colored by correctness.""" - display = _folder_display_name(dataset_name) - fig, ax = plt.subplots(figsize=(8, 6)) - - # Define panels as before - panels = [ - ("Correct", TP_COLOR, df["correct"].astype(bool)), - ("Incorrect", FP_COLOR, ~df["correct"].astype(bool)), - ] - - handles = [] - for label, colour, mask in panels: - sub = df.loc[mask, ["confidence", "calibrated_confidence"]].dropna() - n_total = len(sub) - if n_total < 2: - # Only skip plotting, no data for this class - continue - - plot_df = _subsample_psms(sub, _CONFIDENCE_SCATTER_MAX_POINTS) - handle = ax.scatter( - plot_df["confidence"], - plot_df["calibrated_confidence"], - c=colour, - s=10, - alpha=0.3, - rasterized=True, - label=f"{label}", - ) - handles.append(handle) - - ax.plot( - [-0.01, 1.01], - [-0.01, 1.01], - ls="--", - color="black", - lw=1, - label="No recalibration", - zorder=5, - ) - ax.set_xlim(-0.01, 1.01) - ax.set_ylim(-0.01, 1.01) - ax.set_xlabel("Raw confidence") - ax.set_ylabel("Calibrated confidence") - ax.legend(loc="lower right", fontsize=9) - _style_ax(ax) - - fig.suptitle( - f"Raw versus calibrated confidence for {display}", - fontsize=13, - ) - fig.tight_layout() - _save_fig(fig, output_dir / f"confidence_scatter_{dataset_name}", plot_format) - - -def _plot_feature_distributions_at_fdr( - upscored: pd.DataFrame, - *, - fdr_t: float, - delta_threshold: float, - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """Violin plots of features for up-scored TPs vs FPs retained at one FDR cutoff.""" - n = len(upscored) - pct = int(fdr_t * 100) - if n < _MIN_FEATURE_VIOLIN_PSMs: - logger.info( - "Skipping feature violins for %s at %d%% FDR (n=%d up-scored retained)", - dataset_name, - pct, - n, - ) - return - - available = [ - c - for c in _violin_feature_columns(upscored) - if c in upscored.columns - and upscored[c].notna().sum() >= _MIN_FEATURE_VIOLIN_PSMs - ] - if not available: - logger.warning( - "No feature columns with >=%d values for %s at %d%% FDR", - _MIN_FEATURE_VIOLIN_PSMs, - dataset_name, - pct, - ) - return - - plot_df = upscored.copy() - plot_df["label"] = plot_df["correct"].map({True: "Correct", False: "Incorrect"}) - n_features = len(available) - n_cols = min(3, n_features) - n_rows = (n_features + n_cols - 1) // n_cols - fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows)) - axes = np.atleast_1d(axes).flatten() - - palette = {"Correct": TP_COLOR, "Incorrect": FP_COLOR} - n_plotted = 0 - - for i, col in enumerate(available): - ax = axes[i] - sub = plot_df[[col, "label"]].dropna(subset=[col]) - if len(sub) < _MIN_FEATURE_VIOLIN_PSMs: - ax.set_visible(False) - continue - sns.violinplot( - data=sub, - x="label", - y=col, - palette=palette, - ax=ax, - inner="quartile", - cut=0, - linewidth=0.8, - ) - ax.set_xlabel("") - ax.set_ylabel(_nice_label(col)) - ax.set_title(_nice_label(col)) - _style_ax(ax) - n_plotted += 1 - - for i in range(len(available), len(axes)): - axes[i].set_visible(False) - - if n_plotted == 0: - plt.close(fig) - logger.info( - "Skipping feature violins for %s at %d%% FDR (no feature panels with n>=%d)", - dataset_name, - pct, - _MIN_FEATURE_VIOLIN_PSMs, - ) - return - - display = _folder_display_name(dataset_name) - fig.suptitle( - f"Feature distributions for up-scored PSMs retained at {pct}% FDR " - f"(calibration increase > {delta_threshold:.2f}) on {display}\n" - f"(n={n:,})", - fontsize=13, - ) - fig.tight_layout() - _save_fig( - fig, - output_dir / f"upscored_features_{dataset_name}_fdr{pct}", - plot_format, - ) - - -def _plot_feature_distributions( - df: pd.DataFrame, - delta_threshold: float, - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """Violin plots at 5% and 10% FDR for up-scored correct vs incorrect PSMs.""" - df = _add_q_values(df) - upscored_mask = df["delta_confidence"] > delta_threshold - for fdr_t in _FEATURE_VIOLIN_FDR_THRESHOLDS: - retained = df[upscored_mask & (df["psm_q_value"] <= fdr_t)].copy() - _plot_feature_distributions_at_fdr( - retained, - fdr_t=fdr_t, - delta_threshold=delta_threshold, - dataset_name=dataset_name, - output_dir=output_dir, - plot_format=plot_format, - ) - - -def _upscored_fp_detail( - df: pd.DataFrame, - delta_threshold: float, - dataset_name: str, - metrics: Metrics, -) -> pd.DataFrame: - """Detailed characterisation of up-scored FPs that pass FDR thresholds.""" - df = _add_q_values(df) - upscored_fps = df[ - (df["delta_confidence"] > delta_threshold) & (~df["correct"].astype(bool)) - ].copy() - - if len(upscored_fps) == 0: - return pd.DataFrame() - - def _match_fraction(row: pd.Series) -> float: - nm = row.get("num_matches", 0) - seq = row.get("sequence", "") - if isinstance(seq, str): - tokens = metrics._split_peptide(seq) - else: - tokens = seq if seq else [] - return nm / len(tokens) if tokens else 0.0 - - upscored_fps["match_fraction"] = upscored_fps.apply(_match_fraction, axis=1) - - def _edit_dist(row: pd.Series) -> int: - s = _strip_mods(str(row.get("sequence", ""))) - p = _strip_mods(str(row.get("prediction", ""))) - if not s or not p: - return -1 - return _levenshtein(s, p) - - upscored_fps["edit_distance_norm"] = upscored_fps.apply(_edit_dist, axis=1) - - rows = [] - for fdr_t in FDR_THRESHOLDS: - sub = upscored_fps[upscored_fps["psm_q_value"] <= fdr_t] - if len(sub) == 0: - rows.append( - {"dataset": dataset_name, "fdr_threshold": fdr_t, "n_upscored_fps": 0} - ) - continue - rows.append( - { - "dataset": dataset_name, - "fdr_threshold": fdr_t, - "n_upscored_fps": len(sub), - "mean_match_fraction": round(float(sub["match_fraction"].mean()), 4), - "median_edit_distance": int(sub["edit_distance_norm"].median()), - "n_edit_dist_le2": int((sub["edit_distance_norm"] <= 2).sum()), - "n_partial_match": int((sub["match_fraction"] > 0).sum()), - } - ) - return pd.DataFrame(rows) - - -def _levenshtein(s: str, t: str) -> int: - """Simple Levenshtein distance for short peptide strings.""" - n, m = len(s), len(t) - if n == 0: - return m - if m == 0: - return n - prev = list(range(m + 1)) - for i in range(1, n + 1): - curr = [i] + [0] * m - for j in range(1, m + 1): - cost = 0 if s[i - 1] == t[j - 1] else 1 - curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost) - prev = curr - return prev[m] - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- -_DEFAULT_PREDICTIONS_ROOT = Path("predictions/general_model") -_DEFAULT_OUTPUT_DIR = Path("analysis/upscored_fps") - - -@app.command() -def main( - predictions_root: Annotated[ - Path, - typer.Option( - help="Root directory containing winnow predict output folders.", - ), - ] = _DEFAULT_PREDICTIONS_ROOT, - output_dir: Annotated[ - Path, - typer.Option(help="Directory for output tables and plots."), - ] = _DEFAULT_OUTPUT_DIR, - delta_threshold: Annotated[ - float, - typer.Option( - help="Minimum delta (calibrated - raw) to classify a PSM as up-scored. " - "Default 0.2 (20 percentage-point increase).", - ), - ] = 0.2, - plot_format: Annotated[ - str, - typer.Option(help="Plot format: 'pdf', 'png', or 'both'."), - ] = "both", - residues_config: Annotated[ - Path, - typer.Option(help="Path to residues.yaml for InstaNovo Metrics."), - ] = _REPO_ROOT / "winnow" / "configs" / "residues.yaml", -) -> None: - """Characterise false positives that calibration up-scores into high-confidence regions.""" - output_dir.mkdir(parents=True, exist_ok=True) - plots_dir = output_dir / "plots" - plots_dir.mkdir(parents=True, exist_ok=True) - - residue_masses = _get_residue_masses() - metrics = Metrics( - residue_set=ResidueSet(residue_masses=residue_masses), - isotope_error_range=(0, 1), - ) - - folders = _discover_labelled_folders(predictions_root) - if not folders: - logger.error("No labelled output folders found under %s", predictions_root) - raise typer.Exit(code=1) - - logger.info("Found %d labelled folder(s): %s", len(folders), list(folders.keys())) - - all_summary: list[pd.DataFrame] = [] - all_detail: list[pd.DataFrame] = [] - - for name, folder in folders.items(): - logger.info("Processing %s ...", name) - df = _load_dataset(folder) - - missing_conf = [ - c for c in ("confidence", "calibrated_confidence") if c not in df.columns - ] - if missing_conf: - logger.warning( - "Skipping %s: missing confidence columns %s", - name, - missing_conf, - ) - continue - if "correct" not in df.columns: - logger.warning("Skipping %s: missing 'correct' column", name) - continue - - df["delta_confidence"] = df["calibrated_confidence"] - df["confidence"] - - logger.info( - " %s: %d PSMs, %d correct, delta stats: mean=%.3f, q75=%.3f", - name, - len(df), - int(df["correct"].sum()), - df["delta_confidence"].mean(), - df["delta_confidence"].quantile(0.75), - ) - - summary = _upscored_summary_table(df, delta_threshold, name) - all_summary.append(summary) - - _plot_confidence_scatter(df, name, plots_dir, plot_format) - _plot_feature_distributions(df, delta_threshold, name, plots_dir, plot_format) - - detail = _upscored_fp_detail(df, delta_threshold, name, metrics) - if len(detail) > 0: - all_detail.append(detail) - - if all_summary: - combined = pd.concat(all_summary, ignore_index=True) - combined.to_csv(output_dir / "upscored_summary.csv", index=False) - logger.info("Summary table:\n%s", combined.to_string(index=False)) - - with open(output_dir / "upscored_summary.json", "w") as f: - json.dump(combined.to_dict(orient="records"), f, indent=2) - - if all_detail: - detail_df = pd.concat(all_detail, ignore_index=True) - detail_df.to_csv(output_dir / "upscored_fp_detail.csv", index=False) - logger.info("FP detail table:\n%s", detail_df.to_string(index=False)) - - logger.info("Up-scored FP analysis complete. Output in %s", output_dir) - - -if __name__ == "__main__": - app() diff --git a/scripts/benchmark_runtime.py b/scripts/benchmark_runtime.py deleted file mode 100644 index a287959f..00000000 --- a/scripts/benchmark_runtime.py +++ /dev/null @@ -1,690 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark wall-clock time and memory for the winnow prediction pipeline. - -Measures end-to-end processing time and peak memory, broken down by: - (i) data loading, - (ii) per-feature computation (individually timed), - (iii) MLP calibration inference, and - (iv) FDR / q-value computation. - -Two configurations are benchmarked by default: - 1. Full feature set (including Prosit/Koina-derived features). - 2. Without Prosit features (fragment match + iRT removed). - -This directly addresses the modularity claim: Prosit-based features can be -omitted without disrupting the pipeline. - -Usage examples: - # Both configurations on sample data (requires Koina server for full run) - python scripts/benchmark_runtime.py - - # Only the no-Prosit configuration (no Koina server required) - python scripts/benchmark_runtime.py --no-prosit - - # Custom dataset and locally-trained model - python scripts/benchmark_runtime.py \ - --spectrum-path data/spectra.ipc \ - --predictions-path data/predictions.csv \ - --model-path models/my_model \ - --data-loader instanovo - - # Save structured results to JSON - python scripts/benchmark_runtime.py --output-json results/benchmark.json -""" - -from __future__ import annotations - -import argparse -import json -import os -import platform -import resource -import time -import tracemalloc -from contextlib import contextmanager -from dataclasses import asdict, dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional - -import torch - -from winnow.calibration.calibrator import ProbabilityCalibrator -from winnow.calibration.features.fragment_match import FragmentMatchFeatures -from winnow.calibration.features.retention_time import RetentionTimeFeature -from winnow.datasets.calibration_dataset import CalibrationDataset -from winnow.fdr.nonparametric import NonParametricFDRControl - - -PROSIT_FEATURE_CLASSES = (FragmentMatchFeatures, RetentionTimeFeature) - - -# --------------------------------------------------------------------------- -# Measurement helpers -# --------------------------------------------------------------------------- - - -@dataclass -class StageResult: - """Timing and memory result for a single pipeline stage.""" - - name: str - device: str - wall_time_s: float - peak_mem_mb: float - is_feature: bool = False - is_prosit: bool = False - columns: List[str] = field(default_factory=list) - - -@contextmanager -def measure(): - """Context manager that yields a dict populated with wall_time_s and peak_mem_mb on exit.""" - result: Dict[str, float] = {} - tracemalloc.start() - # Reset the peak so we measure only this block - tracemalloc.reset_peak() - t0 = time.perf_counter() - try: - yield result - finally: - result["wall_time_s"] = time.perf_counter() - t0 - _, peak = tracemalloc.get_traced_memory() - tracemalloc.stop() - result["peak_mem_mb"] = peak / (1024 * 1024) - - -# --------------------------------------------------------------------------- -# Hardware info -# --------------------------------------------------------------------------- - - -def get_hardware_info() -> Dict[str, str]: - """Collect CPU, RAM, and GPU identifiers for the benchmark report.""" - info: Dict[str, str] = {} - - # CPU - cpu_name = None - try: - with open("/proc/cpuinfo") as f: - for line in f: - if line.startswith("model name"): - cpu_name = line.split(":", 1)[1].strip() - break - except OSError: - pass - if not cpu_name: - cpu_name = platform.processor() or "unknown" - info["cpu"] = cpu_name - info["cpu_cores"] = str(os.cpu_count() or "unknown") - - # RAM - try: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemTotal"): - kb = int(line.split()[1]) - info["ram_gb"] = f"{kb / (1024**2):.0f}" - break - except OSError: - info["ram_gb"] = "unknown" - - # GPU - if torch.cuda.is_available(): - info["gpu"] = torch.cuda.get_device_name(0) - else: - info["gpu"] = "none" - - return info - - -# --------------------------------------------------------------------------- -# Pipeline stages -# --------------------------------------------------------------------------- - - -def load_dataset( - spectrum_path: str, - predictions_path: Optional[str], - data_loader_name: str, -) -> CalibrationDataset: - """Load and filter the dataset, returning a CalibrationDataset.""" - from hydra import compose, initialize_config_dir - from hydra.utils import instantiate - from winnow.utils.config_path import get_primary_config_dir - - primary_config_dir = get_primary_config_dir(None) - overrides = [f"data_loader={data_loader_name}"] - - with initialize_config_dir( - config_dir=str(primary_config_dir), - version_base="1.3", - job_name="benchmark", - ): - cfg = compose(config_name="predict", overrides=overrides) - - data_loader = instantiate(cfg.data_loader) - dataset = data_loader.load( - data_path=spectrum_path, - predictions_path=predictions_path, - ) - - from winnow.scripts.main import _filter_dataset - - dataset = _filter_dataset(dataset) - return dataset - - -def compute_features_individually( - calibrator: ProbabilityCalibrator, - dataset: CalibrationDataset, -) -> List[StageResult]: - """Run each feature's prepare+compute with individual timing.""" - results: List[StageResult] = [] - - # Dependencies (currently all features return [], but measure for completeness) - for dep in calibrator.dependencies.values(): - with measure() as m: - dep.compute(dataset=dataset) - results.append( - StageResult( - name=f"Dependency: {dep.name}", - device="CPU", - wall_time_s=m["wall_time_s"], - peak_mem_mb=m["peak_mem_mb"], - ) - ) - - for name, feat in calibrator.feature_dict.items(): - is_prosit = isinstance(feat, PROSIT_FEATURE_CLASSES) - device = "CPU + network" if is_prosit else "CPU" - - with measure() as m: - feat.prepare(dataset=dataset) - feat.compute(dataset=dataset) - - results.append( - StageResult( - name=f"Feature: {name}", - device=device, - wall_time_s=m["wall_time_s"], - peak_mem_mb=m["peak_mem_mb"], - is_feature=True, - is_prosit=is_prosit, - columns=list(feat.columns), - ) - ) - - return results - - -def run_mlp_inference( - calibrator: ProbabilityCalibrator, - dataset: CalibrationDataset, -) -> StageResult: - """Run MLP calibration inference.""" - if calibrator.network is None: - raise RuntimeError("Calibrator network is not loaded") - device = str(next(calibrator.network.parameters()).device) - with measure() as m: - calibrator.predict(dataset) - return StageResult( - name="MLP calibration inference", - device=device.upper() if device == "cpu" else device, - wall_time_s=m["wall_time_s"], - peak_mem_mb=m["peak_mem_mb"], - ) - - -def run_fdr( - dataset: CalibrationDataset, - confidence_column: str = "calibrated_confidence", - fdr_threshold: float = 0.05, -) -> StageResult: - """Run FDR / q-value computation.""" - fdr_control = NonParametricFDRControl() - - with measure() as m: - fdr_control.fit(dataset=dataset.metadata[confidence_column]) - dataset.metadata = fdr_control.add_psm_pep(dataset.metadata, confidence_column) - dataset.metadata = fdr_control.add_psm_fdr(dataset.metadata, confidence_column) - dataset.metadata = fdr_control.add_psm_q_value( - dataset.metadata, confidence_column - ) - confidence_cutoff = fdr_control.get_confidence_cutoff(threshold=fdr_threshold) - _ = dataset.metadata[dataset.metadata[confidence_column] >= confidence_cutoff] - - return StageResult( - name="FDR / q-value computation", - device="CPU", - wall_time_s=m["wall_time_s"], - peak_mem_mb=m["peak_mem_mb"], - ) - - -# --------------------------------------------------------------------------- -# Single benchmark run -# --------------------------------------------------------------------------- - - -@dataclass -class BenchmarkRun: - """Results from a single pipeline configuration.""" - - config_label: str - n_spectra: int - n_features: int - n_columns: int - stages: List[StageResult] - - @property - def total_wall_time_s(self) -> float: - """Sum of wall times across all recorded stages.""" - return sum(s.wall_time_s for s in self.stages) - - @property - def feature_wall_time_s(self) -> float: - """Sum of wall times for feature computation stages only.""" - return sum(s.wall_time_s for s in self.stages if s.is_feature) - - @property - def peak_mem_mb(self) -> float: - """Maximum peak memory across stages, in megabytes.""" - return max(s.peak_mem_mb for s in self.stages) if self.stages else 0.0 - - -def _model_matches_features(calibrator: ProbabilityCalibrator) -> bool: - """Check whether the loaded MLP input dim matches the current feature set.""" - if calibrator.network is None or calibrator.feature_mean is None: - return False - expected_dim = calibrator.feature_mean.shape[0] - actual_dim = 1 + len(calibrator.columns) # confidence + feature columns - return expected_dim == actual_dim - - -def run_benchmark( - spectrum_path: str, - predictions_path: Optional[str], - model_path: str, - data_loader_name: str, - include_prosit: bool, - koina_url: Optional[str] = None, - koina_ssl: Optional[bool] = None, -) -> BenchmarkRun: - """Execute the full prediction pipeline with per-stage timing.""" - config_label = "Full feature set" if include_prosit else "Without Prosit features" - - # Load calibrator - calibrator = ProbabilityCalibrator.load(pretrained_model_name_or_path=model_path) - - # Remove Prosit features if requested - if not include_prosit: - to_remove = [ - name - for name, feat in calibrator.feature_dict.items() - if isinstance(feat, PROSIT_FEATURE_CLASSES) - ] - for name in to_remove: - calibrator.remove_feature(name) - - stages: List[StageResult] = [] - - # Stage 1: Data loading - with measure() as m: - dataset = load_dataset(spectrum_path, predictions_path, data_loader_name) - n_spectra = len(dataset.metadata) - stages.append( - StageResult( - name="Data loading", - device="CPU", - wall_time_s=m["wall_time_s"], - peak_mem_mb=m["peak_mem_mb"], - ) - ) - - # Stage 2: Per-feature computation - feature_results = compute_features_individually(calibrator, dataset) - stages.extend(feature_results) - - n_features = len(calibrator.feature_dict) - n_columns = len(calibrator.columns) - - # Stage 3: MLP calibration inference - # The MLP input dimension must match the feature set. If features were - # removed (e.g. Prosit features dropped) but the model was trained with - # the full set, the dimensions won't match. In that case we skip MLP + - # FDR and note the mismatch -- these stages are sub-millisecond anyway - # and their cost is independent of the feature set used. - can_infer = _model_matches_features(calibrator) - if can_infer: - stages.append(run_mlp_inference(calibrator, dataset)) - # Stage 4: FDR / q-value - stages.append(run_fdr(dataset)) - else: - mean_dim = ( - calibrator.feature_mean.shape[0] - if calibrator.feature_mean is not None - else "unknown" - ) - print( - f" [note] MLP input dim ({mean_dim}) " - f"does not match current feature count " - f"({1 + len(calibrator.columns)}). " - f"Skipping MLP inference and FDR for this configuration.\n" - f" To benchmark these stages, supply a model trained with the " - f"matching feature set via --model-path-no-prosit." - ) - - return BenchmarkRun( - config_label=config_label, - n_spectra=n_spectra, - n_features=n_features, - n_columns=n_columns, - stages=stages, - ) - - -# --------------------------------------------------------------------------- -# Output formatting -# --------------------------------------------------------------------------- - - -def format_run(run: BenchmarkRun) -> str: - """Format a single benchmark run as a human-readable table.""" - lines: List[str] = [] - header = ( - f"=== Configuration: {run.config_label} " - f"({run.n_features} features, {run.n_columns} columns) ===" - ) - lines.append("") - lines.append(header) - lines.append("") - - col_w = [38, 16, 15, 15] - hdr = ( - f"{'Stage':<{col_w[0]}}| {'Device':<{col_w[1]}}| " - f"{'Wall time (s)':>{col_w[2]}}| {'Peak mem (MB)':>{col_w[3]}}" - ) - sep = ( - "-" * col_w[0] - + "|" - + "-" * (col_w[1] + 1) - + "|" - + "-" * (col_w[2] + 1) - + "|" - + "-" * (col_w[3] + 1) - ) - - lines.append(hdr) - lines.append(sep) - - feat_time = 0.0 - feat_mem = 0.0 - total_time = 0.0 - - def _row(label: str, device: str, t: float, mem: float) -> str: - return ( - f"{label:<{col_w[0]}}| {device:<{col_w[1]}}| " - f"{t:>{col_w[2]}.2f}| {mem:>{col_w[3]}.1f}" - ) - - for i, s in enumerate(run.stages): - lines.append(_row(s.name, s.device, s.wall_time_s, s.peak_mem_mb)) - total_time += s.wall_time_s - if s.is_feature: - feat_time += s.wall_time_s - feat_mem = max(feat_mem, s.peak_mem_mb) - - is_last_feature = s.is_feature and not any( - st.is_feature for st in run.stages[i + 1 :] - ) - if is_last_feature: - lines.append( - _row( - " Feature computation subtotal", - "", - feat_time, - feat_mem, - ) - ) - - lines.append(sep) - total_mem = max(s.peak_mem_mb for s in run.stages) if run.stages else 0.0 - lines.append(_row("End-to-end total", "", total_time, total_mem)) - - return "\n".join(lines) - - -def format_full_report( - hw: Dict[str, str], - mlp_device: str, - runs: List[BenchmarkRun], - n_spectra: int, - spectrum_path: str, -) -> str: - """Assemble the complete benchmark report.""" - lines: List[str] = [] - lines.append("=" * 88) - lines.append(" Winnow Pipeline Runtime Benchmark") - lines.append("=" * 88) - lines.append("") - lines.append( - f"Hardware: {hw['cpu']} ({hw['cpu_cores']} cores), " - f"{hw['ram_gb']} GB RAM, GPU: {hw['gpu']}" - ) - lines.append(f"Dataset: {n_spectra:,} spectra from {spectrum_path}") - lines.append(f"Calibrator MLP device: {mlp_device}") - - for run in runs: - lines.append(format_run(run)) - - lines.append("") - lines.append("Notes:") - lines.append( - '- "CPU + network" = gRPC calls to a Koina/Triton server for' - " Prosit-derived spectral predictions." - ) - lines.append(" No local GPU is used by winnow during prediction.") - lines.append( - "- The MLP runs on CPU after loading. GPU is only used during training" - " (not benchmarked here)." - ) - lines.append( - "- Koina predictions are not cached to disk between runs; batching is" - " handled internally by koinapy." - ) - peak_rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 - lines.append(f"- Process peak RSS: {peak_rss_mb:.1f} MB") - lines.append("") - return "\n".join(lines) - - -def build_json_report( - hw: Dict[str, str], - mlp_device: str, - runs: List[BenchmarkRun], - n_spectra: int, - spectrum_path: str, -) -> Dict[str, Any]: - """Build a structured dict suitable for JSON serialisation.""" - report: Dict[str, Any] = { - "hardware": hw, - "dataset": { - "spectrum_path": spectrum_path, - "n_spectra": n_spectra, - }, - "mlp_device": mlp_device, - "configurations": [], - } - for run in runs: - cfg: Dict[str, Any] = { - "label": run.config_label, - "n_features": run.n_features, - "n_columns": run.n_columns, - "total_wall_time_s": run.total_wall_time_s, - "feature_wall_time_s": run.feature_wall_time_s, - "stages": [asdict(s) for s in run.stages], - } - report["configurations"].append(cfg) - report["process_peak_rss_mb"] = ( - resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 - ) - report["notes"] = { - "koina_caching": ( - "Koina predictions are not cached to disk; each invocation " - "re-queries the server." - ), - "koina_batching": ( - "Batching is handled internally by koinapy " - "(gRPC streaming to the Koina/Triton server)." - ), - "gpu_usage": ( - "No local GPU is used during prediction. GPU is only used " - "during training (not benchmarked here)." - ), - } - return report - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- - - -def parse_args() -> argparse.Namespace: - """Parse command-line arguments for the runtime benchmark.""" - parser = argparse.ArgumentParser( - description="Benchmark winnow prediction pipeline runtime and memory.", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=__doc__, - ) - parser.add_argument( - "--spectrum-path", - default="examples/example_data/spectra.ipc", - help="Path to the spectrum data file (default: example data).", - ) - parser.add_argument( - "--predictions-path", - default="examples/example_data/predictions.csv", - help="Path to predictions file (default: example data).", - ) - parser.add_argument( - "--model-path", - default="InstaDeepAI/winnow-general-model", - help=( - "Path to a local calibrator directory or HuggingFace model " - "identifier (default: InstaDeepAI/winnow-general-model)." - ), - ) - parser.add_argument( - "--model-path-no-prosit", - default=None, - help=( - "Path to a calibrator trained without Prosit features, used for " - "the no-Prosit benchmark. If omitted, --model-path is used and " - "MLP/FDR stages are skipped when the input dimension mismatches." - ), - ) - parser.add_argument( - "--data-loader", - default="instanovo", - help="Data loader to use (default: instanovo).", - ) - parser.add_argument( - "--no-prosit", - action="store_true", - help=( - "Only benchmark without Prosit/Koina features. When omitted, " - "both configurations (full and no-Prosit) are benchmarked." - ), - ) - parser.add_argument( - "--full-only", - action="store_true", - help="Only benchmark the full feature set (skip no-Prosit run).", - ) - parser.add_argument( - "--koina-url", - default=None, - help="Koina server URL (e.g. localhost:8500).", - ) - parser.add_argument( - "--koina-ssl", - default=None, - type=lambda x: x.lower() in ("true", "1", "yes"), - help="Use SSL for Koina (true/false).", - ) - parser.add_argument( - "--output-json", - default=None, - help="Save structured results to a JSON file.", - ) - return parser.parse_args() - - -def main() -> None: - """Run configured pipeline benchmarks and print (optionally save) results.""" - args = parse_args() - - hw = get_hardware_info() - - # Determine which configs to run - run_full = not args.no_prosit - run_no_prosit = not args.full_only - - # Detect MLP device from a probe load - probe_calibrator = ProbabilityCalibrator.load( - pretrained_model_name_or_path=args.model_path - ) - if probe_calibrator.network is None: - raise RuntimeError("Calibrator network is not loaded") - mlp_device = str(next(probe_calibrator.network.parameters()).device) - del probe_calibrator - - runs: List[BenchmarkRun] = [] - n_spectra = 0 - - if run_full: - print("\n>>> Benchmarking: Full feature set (including Prosit) ...") - result = run_benchmark( - spectrum_path=args.spectrum_path, - predictions_path=args.predictions_path, - model_path=args.model_path, - data_loader_name=args.data_loader, - include_prosit=True, - koina_url=args.koina_url, - koina_ssl=args.koina_ssl, - ) - runs.append(result) - n_spectra = result.n_spectra - - if run_no_prosit: - print("\n>>> Benchmarking: Without Prosit features ...") - no_prosit_model = args.model_path_no_prosit or args.model_path - result = run_benchmark( - spectrum_path=args.spectrum_path, - predictions_path=args.predictions_path, - model_path=no_prosit_model, - data_loader_name=args.data_loader, - include_prosit=False, - koina_url=args.koina_url, - koina_ssl=args.koina_ssl, - ) - runs.append(result) - n_spectra = n_spectra or result.n_spectra - - report = format_full_report(hw, mlp_device, runs, n_spectra, args.spectrum_path) - print(report) - - if args.output_json: - json_report = build_json_report( - hw, mlp_device, runs, n_spectra, args.spectrum_path - ) - out_path = Path(args.output_json) - out_path.parent.mkdir(parents=True, exist_ok=True) - with open(out_path, "w") as f: - json.dump(json_report, f, indent=2) - print(f"JSON results saved to {out_path}") - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmark_scaling.py b/scripts/benchmark_scaling.py deleted file mode 100644 index 63d07852..00000000 --- a/scripts/benchmark_scaling.py +++ /dev/null @@ -1,421 +0,0 @@ -#!/usr/bin/env python3 -"""Measure pipeline scaling by running the no-Prosit benchmark at multiple dataset sizes. - -Writes subsampled spectrum and prediction files to a temporary directory, then -runs the full pipeline (data loading, feature computation, MLP inference, FDR) -from scratch at each size. Produces a JSON file with the raw measurements and -a matplotlib figure showing per-stage scaling. - -Usage: - python scripts/benchmark_scaling.py \ - --spectrum-path held_out_projects/.../dataset-helaqc-raw-0000-0001.parquet \ - --predictions-path held_out_projects/.../dataset-helaqc-raw-0000-0001.csv \ - --model-path models/benchmark_model_no_prosit \ - --output-dir analysis -""" - -from __future__ import annotations - -import argparse -import json -import random -import shutil -import tempfile -import time -import tracemalloc -from contextlib import contextmanager -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Tuple - -import matplotlib.pyplot as plt -import numpy as np -import polars as pl -import seaborn as sns - -from winnow.calibration.calibrator import ProbabilityCalibrator -from winnow.calibration.features.fragment_match import FragmentMatchFeatures -from winnow.calibration.features.retention_time import RetentionTimeFeature -from winnow.datasets.calibration_dataset import CalibrationDataset -from winnow.fdr.nonparametric import NonParametricFDRControl - -plt.switch_backend("Agg") - -# Paul Tol "bright" palette (colorblind-safe) — same as plot_eval_results.py -_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -PROSIT_FEATURE_CLASSES = (FragmentMatchFeatures, RetentionTimeFeature) - -DEFAULT_FRACTIONS = [0.1, 0.5, 1.0] - - -@contextmanager -def measure(): - """Context manager that yields a dict populated with wall_time_s and peak_mem_mb on exit.""" - result: Dict[str, float] = {} - tracemalloc.start() - tracemalloc.reset_peak() - t0 = time.perf_counter() - try: - yield result - finally: - result["wall_time_s"] = time.perf_counter() - t0 - _, peak = tracemalloc.get_traced_memory() - tracemalloc.stop() - result["peak_mem_mb"] = peak / (1024 * 1024) - - -def write_subsampled_files( - spectrum_path: str, - predictions_path: str, - fraction: float, - output_dir: Path, - seed: int = 42, -) -> Tuple[Path, Path, int]: - """Write subsampled spectrum and prediction files, returning paths and row count. - - Both files are joined on spectrum_id, subsampled together, and written out - so the data loader sees consistent, smaller files. - """ - spectra = pl.read_parquet(spectrum_path) - preds = pl.read_csv(predictions_path) - - n = len(spectra) - if fraction >= 1.0: - k = n - indices = list(range(n)) - else: - k = max(1, int(n * fraction)) - rng = random.Random(seed) - indices = sorted(rng.sample(range(n), k)) - - spectra_sub = spectra[indices] - # Match predictions to the subsampled spectra by spectrum_id - keep_ids = set(spectra_sub["spectrum_id"].to_list()) - preds_sub = preds.filter(pl.col("spectrum_id").is_in(keep_ids)) - - spec_path = output_dir / f"spectra_{fraction:.2f}.parquet" - pred_path = output_dir / f"preds_{fraction:.2f}.csv" - spectra_sub.write_parquet(spec_path) - preds_sub.write_csv(pred_path) - - return spec_path, pred_path, len(spectra_sub) - - -def load_dataset_timed( - spectrum_path: str, - predictions_path: str, - data_loader_name: str, -) -> Tuple[CalibrationDataset, float]: - """Load a dataset through the full data loader and return it with timing.""" - from hydra import compose, initialize_config_dir - from hydra.utils import instantiate - from winnow.utils.config_path import get_primary_config_dir - - primary_config_dir = get_primary_config_dir(None) - overrides = [f"data_loader={data_loader_name}"] - - with initialize_config_dir( - config_dir=str(primary_config_dir), - version_base="1.3", - job_name="benchmark_scaling", - ): - cfg = compose(config_name="predict", overrides=overrides) - - data_loader = instantiate(cfg.data_loader) - - with measure() as m: - dataset = data_loader.load( - data_path=spectrum_path, - predictions_path=predictions_path, - ) - from winnow.scripts.main import _filter_dataset - - dataset = _filter_dataset(dataset) - - return dataset, m["wall_time_s"] - - -@dataclass -class ScalingPoint: - """Timing measurements for one dataset-size fraction.""" - - fraction: float - n_spectra: int - stage_times: Dict[str, float] = field(default_factory=dict) - total_time: float = 0.0 - - -def run_at_size( - spec_path: Path, - pred_path: Path, - expected_n: int, - calibrator: ProbabilityCalibrator, - data_loader_name: str, - fraction: float, -) -> ScalingPoint: - """Run the full pipeline from disk at a given dataset size.""" - stage_times: Dict[str, float] = {} - - # Data loading (from subsampled files on disk) - dataset, load_time = load_dataset_timed( - str(spec_path), str(pred_path), data_loader_name - ) - n = len(dataset.metadata) - stage_times["Data loading"] = load_time - - # Feature computation (per feature) - feature_total = 0.0 - for name, feat in calibrator.feature_dict.items(): - with measure() as m: - feat.prepare(dataset=dataset) - feat.compute(dataset=dataset) - stage_times[f"Feature: {name}"] = m["wall_time_s"] - feature_total += m["wall_time_s"] - stage_times["Feature computation (total)"] = feature_total - - # MLP inference - with measure() as m: - calibrator.predict(dataset) - stage_times["MLP inference"] = m["wall_time_s"] - - # FDR / q-value - fdr = NonParametricFDRControl() - col = "calibrated_confidence" - with measure() as m: - fdr.fit(dataset=dataset.metadata[col]) - dataset.metadata = fdr.add_psm_pep(dataset.metadata, col) - dataset.metadata = fdr.add_psm_fdr(dataset.metadata, col) - dataset.metadata = fdr.add_psm_q_value(dataset.metadata, col) - cutoff = fdr.get_confidence_cutoff(threshold=0.05) - _ = dataset.metadata[dataset.metadata[col] >= cutoff] - stage_times["FDR / q-value"] = m["wall_time_s"] - - total = ( - load_time - + feature_total - + stage_times["MLP inference"] - + stage_times["FDR / q-value"] - ) - stage_times["End-to-end"] = total - - return ScalingPoint( - fraction=fraction, - n_spectra=n, - stage_times=stage_times, - total_time=total, - ) - - -def fit_exponent(sizes: List[int], times: List[float]) -> Tuple[float, float]: - """Fit t = c * n^alpha in log-log space; return (alpha, R²).""" - log_s = np.log(sizes) - log_t = np.log(times) - slope, intercept = np.polyfit(log_s, log_t, 1) - predicted = slope * log_s + intercept - ss_res = float(np.sum((log_t - predicted) ** 2)) - ss_tot = float(np.sum((log_t - np.mean(log_t)) ** 2)) - r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0 - return slope, r2 - - -def _save_fig(fig: plt.Figure, base_path: Path) -> None: - """Save figure as both PNG and PDF.""" - fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def plot_scaling(points: List[ScalingPoint], output_path: Path) -> None: - """Plot per-stage wall time vs. dataset size.""" - sizes = [p.n_spectra for p in points] - - stages_to_plot = [ - ("Data loading", "Data loading", "v", "-", _PALETTE[0]), - ("Feature computation", "Feature computation (total)", "o", "-", _PALETTE[2]), - ("MLP inference", "MLP inference", "^", "-", _PALETTE[1]), - ("FDR / q-value", "FDR / q-value", "D", "-", _PALETTE[3]), - ("End-to-end", "End-to-end", "s", "--", _PALETTE[5]), - ] - - fig, ax = plt.subplots(figsize=(6, 4)) - - for label, key, marker, linestyle, colour in stages_to_plot: - times = [p.stage_times[key] for p in points] - ax.plot( - sizes, - times, - marker=marker, - linestyle=linestyle, - label=label, - color=colour, - alpha=0.7, - ) - - max_size = max(sizes) - max_total = max(p.stage_times["End-to-end"] for p in points) - ref_sizes = np.linspace(0, max_size, 50) - ref_times = max_total * (ref_sizes / max_size) - ax.plot( - ref_sizes, - ref_times, - ls=":", - color=_PALETTE[6], - linewidth=1, - label="Linear reference", - ) - - ax.set_xlabel("Number of spectra") - ax.set_ylabel("Wall time (s)") - ax.set_ylim(top=300) - ax.set_title("Pipeline scaling\nexcluding Koina-dependent features") - ax.legend(loc="upper left", fontsize=9) - fig.tight_layout() - - base = output_path.with_suffix("") - _save_fig(fig, base) - print(f"Scaling plot saved to {base}.png and {base}.pdf") - - -def load_points_from_json(json_path: Path) -> List[ScalingPoint]: - """Load previously saved scaling measurements from JSON.""" - with open(json_path) as f: - data = json.load(f) - return [ - ScalingPoint( - fraction=p["fraction"], - n_spectra=p["n_spectra"], - stage_times=p["stage_times"], - total_time=p["total_time"], - ) - for p in data["points"] - ] - - -def parse_args() -> argparse.Namespace: - """Parse command-line arguments for the scaling benchmark.""" - parser = argparse.ArgumentParser( - description="Measure pipeline scaling at multiple dataset sizes.", - ) - parser.add_argument( - "--replot-json", - metavar="PATH", - help="Replot from a saved benchmark_scaling.json without rerunning benchmarks.", - ) - parser.add_argument( - "--spectrum-path", - help="Path to the spectrum data file.", - ) - parser.add_argument( - "--predictions-path", - help="Path to predictions file.", - ) - parser.add_argument( - "--model-path", - help="Path to a calibrator trained without Prosit features.", - ) - parser.add_argument( - "--data-loader", - default="instanovo", - help="Data loader to use (default: instanovo).", - ) - parser.add_argument( - "--fractions", - nargs="+", - type=float, - default=DEFAULT_FRACTIONS, - help="Dataset fractions to benchmark (default: 0.1 0.5 1.0).", - ) - parser.add_argument( - "--output-dir", - default="analysis", - help="Directory for output files.", - ) - args = parser.parse_args() - if args.replot_json is None: - for name in ("spectrum_path", "predictions_path", "model_path"): - if getattr(args, name) is None: - parser.error( - f"--{name.replace('_', '-')} is required unless --replot-json is set" - ) - return args - - -def main() -> None: - """Run scaling benchmarks or replot from a saved JSON file.""" - args = parse_args() - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - if args.replot_json: - replot_points = load_points_from_json(Path(args.replot_json)) - plot_scaling(replot_points, output_dir / "benchmark_scaling.png") - return - - # Load calibrator (without Prosit features) - calibrator = ProbabilityCalibrator.load( - pretrained_model_name_or_path=args.model_path - ) - to_remove = [ - name - for name, feat in calibrator.feature_dict.items() - if isinstance(feat, PROSIT_FEATURE_CLASSES) - ] - for name in to_remove: - calibrator.remove_feature(name) - - fractions = sorted(args.fractions) - tmpdir = Path(tempfile.mkdtemp(prefix="winnow_scaling_")) - points: List[ScalingPoint] = [] - - try: - # Pre-write all subsampled files - file_info: List[Tuple[float, Path, Path, int]] = [] - print("Preparing subsampled files ...") - for frac in fractions: - spec_p, pred_p, n = write_subsampled_files( - args.spectrum_path, args.predictions_path, frac, tmpdir - ) - file_info.append((frac, spec_p, pred_p, n)) - print(f" {frac:.0%}: {n:,} spectra -> {spec_p.name}, {pred_p.name}") - - for frac, spec_p, pred_p, expected_n in file_info: - print(f"\n>>> Running at {frac:.0%} ({expected_n:,} spectra) ...") - point = run_at_size( - spec_p, pred_p, expected_n, calibrator, args.data_loader, frac - ) - points.append(point) - print(f" {point.n_spectra:,} spectra -> {point.total_time:.2f} s total") - for stage, t in point.stage_times.items(): - if stage not in ("Feature computation (total)", "End-to-end"): - print(f" {stage}: {t:.3f} s") - - finally: - shutil.rmtree(tmpdir, ignore_errors=True) - - # Save JSON - json_path = output_dir / "benchmark_scaling.json" - json_data: Dict[str, Any] = { - "points": [ - { - "fraction": p.fraction, - "n_spectra": p.n_spectra, - "stage_times": p.stage_times, - "total_time": p.total_time, - } - for p in points - ], - } - with open(json_path, "w") as f: - json.dump(json_data, f, indent=2) - print(f"\nScaling data saved to {json_path}") - - # Plot - plot_path = output_dir / "benchmark_scaling.png" - plot_scaling(points, plot_path) - - -if __name__ == "__main__": - main() diff --git a/scripts/calibrator_generalisation_utils.py b/scripts/calibrator_generalisation_utils.py deleted file mode 100644 index 278cc5ef..00000000 --- a/scripts/calibrator_generalisation_utils.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Shared helpers for calibrator generalisation analysis.""" - -from __future__ import annotations - -import logging -import re -from pathlib import Path - -import polars as pl - -logger = logging.getLogger(__name__) - -HEPG2_SOURCE = "PXD019483" - -SPECIES_NAME_MAPPING: dict[str, str] = { - "gluc": "HeLa degradome", - "helaqc": "HeLa single shot", - "herceptin": "Herceptin", - "immuno": "Immunopeptidomics-1", - "celegans": "$\\it{C.\\;elegans}$", - "sbrodae": "$\\it{Scalindua\\;brodae}$", - HEPG2_SOURCE: "HepG2", - "snakevenoms": "Snake venomics", - "tplantibodies": "Therapeutic nanobodies", - "woundfluids": "Wound exudates", - "PXD014877": "$\\it{C.\\;elegans}$", -} - - -def extract_project_name(parquet_path: Path) -> str: - """Extract project name from ``dataset-helaqc-annotated-0000-0001.parquet``.""" - match = re.match(r"dataset-(.+?)-annotated", parquet_path.stem) - if match: - return match.group(1) - return parquet_path.stem - - -def build_experiment_source_mapping(biological_validation_dir: Path) -> dict[str, str]: - """Map every experiment in biological validation parquets to its source label.""" - mapping: dict[str, str] = {} - parquet_files = sorted(biological_validation_dir.glob("*.parquet")) - if not parquet_files: - raise FileNotFoundError( - f"No parquet files found in biological validation directory: " - f"{biological_validation_dir}" - ) - - for parquet_path in parquet_files: - project = extract_project_name(parquet_path) - experiments = ( - pl.scan_parquet(parquet_path) - .select("experiment_name") - .unique() - .collect()["experiment_name"] - .to_list() - ) - for experiment_name in experiments: - mapping[experiment_name] = project - - logger.info( - "Built experiment->source mapping for %d experiments across %d projects", - len(mapping), - len(parquet_files), - ) - return mapping - - -def annotate_train_source_labels( - train_parquet: Path, - train_predictions: Path, - biological_validation_dir: Path, -) -> None: - """Add a ``source`` column to the train parquet and predictions CSV. - - Experiments found in ``biological_validation_dir`` inherit that project name. - All other experiments are labelled as HepG2 (``PXD019483``). - """ - experiment_to_source = build_experiment_source_mapping(biological_validation_dir) - lookup = pl.DataFrame( - { - "experiment_name": list(experiment_to_source.keys()), - "source": list(experiment_to_source.values()), - } - ) - - spectra = pl.read_parquet(train_parquet) - if "source" not in spectra.columns: - spectra = spectra.join(lookup, on="experiment_name", how="left").with_columns( - pl.col("source").fill_null(HEPG2_SOURCE) - ) - spectra.write_parquet(train_parquet) - logger.info("Wrote source labels to %s", train_parquet) - else: - logger.info( - "Parquet already has source column, leaving %s unchanged", train_parquet - ) - - predictions = pl.read_csv(train_predictions) - if "source" not in predictions.columns: - source_by_spectrum = spectra.select("spectrum_id", "source") - predictions = predictions.join(source_by_spectrum, on="spectrum_id", how="left") - missing = predictions.filter(pl.col("source").is_null()) - if len(missing) > 0: - raise ValueError( - f"{len(missing)} prediction rows in {train_predictions} have no matching " - "spectrum_id in the train parquet" - ) - predictions.write_csv(train_predictions) - logger.info("Wrote source labels to %s", train_predictions) - else: - logger.info( - "Predictions CSV already has source column, leaving %s unchanged", - train_predictions, - ) diff --git a/scripts/evaluate_calibrator_generalisation.py b/scripts/evaluate_calibrator_generalisation.py deleted file mode 100644 index e443c530..00000000 --- a/scripts/evaluate_calibrator_generalisation.py +++ /dev/null @@ -1,415 +0,0 @@ -"""Evaluate calibrator generalisation by training on one source dataset and testing on all others. - -Uses the ``train_extra_small`` train parquet and predictions CSV, with a ``source`` -column derived from biological-validation experiment names (everything else is HepG2). -For each source, trains a fresh calibrator, evaluates it in-distribution (held-out -20 %) and out-of-distribution (every other source), then saves a combined results CSV. -""" - -import logging -import re -import sys -from pathlib import Path -from typing import Annotated, Dict, List, Optional - -import numpy as np -import pandas as pd -import yaml -from rich.logging import RichHandler -import typer - -from winnow.calibration.calibrator import ProbabilityCalibrator -from winnow.calibration.features import ( - BeamFeatures, - FragmentMatchFeatures, - MassErrorDaFeature, - RetentionTimeFeature, - TokenScoreFeatures, -) -from winnow.datasets.calibration_dataset import CalibrationDataset -from winnow.datasets.data_loaders import InstaNovoDatasetLoader - -_REPO_ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(_REPO_ROOT)) - -from scripts.calibrator_generalisation_utils import annotate_train_source_labels # noqa: E402 - -# --------------------------------------------------------------------------- -# Logging -# --------------------------------------------------------------------------- -logger = logging.getLogger("winnow.evaluate_generalization") -logger.setLevel(logging.INFO) -logger.propagate = False -logger.addHandler(RichHandler()) - -# --------------------------------------------------------------------------- -# Constants — loaded from the canonical Winnow YAML configs -# --------------------------------------------------------------------------- -SEED = 42 -TEST_SIZE = 0.2 - -_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "winnow" / "configs" - -with open(_CONFIGS_DIR / "residues.yaml") as _f: - RESIDUE_MASSES: dict[str, float] = yaml.safe_load(_f)["residue_masses"] - -with open(_CONFIGS_DIR / "data_loader" / "instanovo.yaml") as _f: - _instanovo_cfg = yaml.safe_load(_f) - RESIDUE_REMAPPING: dict[str, str] = _instanovo_cfg.get("residue_remapping", {}) - BEAM_COLUMNS: dict[str, str] | None = _instanovo_cfg.get("beam_columns") - -with open(_CONFIGS_DIR / "calibrator.yaml") as _f: - _calibrator_cfg = yaml.safe_load(_f) - _KOINA_CFG = _calibrator_cfg["koina"] - _KOINA_CONSTRAINTS = _KOINA_CFG["constraints"] - _KOINA_INPUT_CONSTANTS = _KOINA_CFG.get("input_constants") or { - "collision_energies": 27, - "fragmentation_types": "HCD", - } - _UNSUPPORTED_RESIDUES: list[str] = ( - _KOINA_CONSTRAINTS.get("unsupported_residues") or [] - ) - _MAX_PRECURSOR_CHARGE: int = _KOINA_CONSTRAINTS["max_precursor_charge"] - _MAX_PEPTIDE_LENGTH: int = _KOINA_CONSTRAINTS["max_peptide_length"] - _INTENSITY_MODEL: str = _KOINA_CFG["intensity_model"] - _IRT_MODEL: str = _KOINA_CFG["irt_model"] - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -_IRT_TRAIN_FRACTION_OVERRIDES: Dict[str, float] = { - "herceptin": 0.15, -} - -# Mirrors Makefile train-extra-small-mass-error-da / EXTRA_SMALL_* overrides. -_EXTRA_SMALL_FRAGMENT_EXCLUDE = [ - "spectral_angle", - "xcorr", - "complementary_ion_count", - "max_ion_gap", -] -_EXTRA_SMALL_BEAM_EXCLUDE = ["edit_distance"] - - -def initialise_calibrator( - *, - koina_server_url: Optional[str] = None, - koina_ssl: bool = True, - train_project: Optional[str] = None, -) -> ProbabilityCalibrator: - """Create a fresh calibrator matching train-extra-small-mass-error-da.""" - if koina_server_url is not None or not koina_ssl: - logger.warning( - "Ignoring koina_server_url/koina_ssl; Koina server overrides are no " - "longer supported (url=%s, ssl=%s).", - koina_server_url, - koina_ssl, - ) - - irt_train_fraction = _IRT_TRAIN_FRACTION_OVERRIDES.get(train_project or "", 0.1) - - calibrator = ProbabilityCalibrator( - hidden_dims=(50, 50), - dropout=0.3, - learning_rate=0.0001, - weight_decay=0.001, - max_epochs=1000, - batch_size=1024, - n_iter_no_change=10, - tol=0.0001, - seed=SEED, - val_early_stopping_max_psms=None, - val_subsample_seed=None, - ) - calibrator.add_feature(MassErrorDaFeature(residue_masses=RESIDUE_MASSES)) - calibrator.add_feature( - FragmentMatchFeatures( - mz_tolerance=20, - mz_tolerance_unit="ppm", - learn_from_missing=False, - intensity_model_name=_INTENSITY_MODEL, - max_precursor_charge=_MAX_PRECURSOR_CHARGE, - max_peptide_length=_MAX_PEPTIDE_LENGTH, - unsupported_residues=_UNSUPPORTED_RESIDUES, - model_input_constants=_KOINA_INPUT_CONSTANTS, - ) - ) - calibrator.add_feature( - RetentionTimeFeature( - train_fraction=irt_train_fraction, - min_train_points=3, - learn_from_missing=False, - irt_model_name=_IRT_MODEL, - max_peptide_length=_MAX_PEPTIDE_LENGTH, - unsupported_residues=_UNSUPPORTED_RESIDUES, - ) - ) - calibrator.add_feature(BeamFeatures()) - calibrator.add_feature(TokenScoreFeatures()) - # Former excluded_columns behaviour: train on a reduced feature subset. - training_columns = [ - col - for col in calibrator.columns - if col not in _EXTRA_SMALL_FRAGMENT_EXCLUDE - and col not in _EXTRA_SMALL_BEAM_EXCLUDE - ] - calibrator.set_training_feature_columns(training_columns) - return calibrator - - -def load_dataset(data_path: Path, predictions_path: Path) -> CalibrationDataset: - """Load the combined train_extra_small dataset.""" - logger.info("Loading dataset from %s and %s", data_path, predictions_path) - loader = InstaNovoDatasetLoader( - residue_masses=RESIDUE_MASSES, - residue_remapping=RESIDUE_REMAPPING, - beam_columns=BEAM_COLUMNS, - ) - return loader.load(data_path=data_path, predictions_path=predictions_path) - - -def subset_dataset(dataset: CalibrationDataset, idx: np.ndarray) -> CalibrationDataset: - """Return a row subset of *dataset* with aligned beam predictions.""" - meta = dataset.metadata.iloc[idx].reset_index(drop=True) - preds = ( - [dataset.predictions[i] for i in idx.tolist()] - if dataset.predictions is not None - else None - ) - return CalibrationDataset(metadata=meta, predictions=preds) - - -def split_dataset_by_source( - dataset: CalibrationDataset, -) -> Dict[str, CalibrationDataset]: - """Split a combined dataset into one CalibrationDataset per ``source`` label.""" - if "source" not in dataset.metadata.columns: - raise ValueError( - "Expected a 'source' column in the train parquet metadata. " - "Run annotate_train_source_labels() first." - ) - - datasets: Dict[str, CalibrationDataset] = {} - for source in sorted(dataset.metadata["source"].unique()): - idx = np.where(dataset.metadata["source"].values == source)[0] - datasets[source] = subset_dataset(dataset, idx) - return datasets - - -_MOD_RE = re.compile(r"\[UNIMOD:\d+\]") - - -def _peptide_key(tokens: object) -> str: - """Normalise a tokenised peptide to a modification-free, I/L-collapsed key. - - Matches the strategy in ``scripts/split_annotated_raw_parquets.py``: - strip UNIMOD modifications, normalise I→L. - """ - if not isinstance(tokens, list): - return "__MISSING__" - stripped = [_MOD_RE.sub("", tok).replace("I", "L") for tok in tokens] - return "".join(stripped) - - -def create_train_test_split( - dataset: CalibrationDataset, -) -> tuple[CalibrationDataset, CalibrationDataset]: - """Split a dataset 80/20 by peptide so no peptide appears in both folds.""" - meta = dataset.metadata - n = len(meta) - if n <= 1: - return dataset, dataset - - pep_keys = meta["sequence"].apply(_peptide_key) - unique_peptides = pep_keys.unique() - - rng = np.random.default_rng(SEED) - perm = rng.permutation(len(unique_peptides)) - n_train = int(len(unique_peptides) * (1 - TEST_SIZE)) - - train_peptides = set(unique_peptides[perm[:n_train]]) - train_mask = pep_keys.isin(train_peptides).values - - train_idx = np.where(train_mask)[0] - test_idx = np.where(~train_mask)[0] - - return subset_dataset(dataset, train_idx), subset_dataset(dataset, test_idx) - - -def evaluate_model( - model: ProbabilityCalibrator, - test_dataset: CalibrationDataset, - train_project: str, - test_project: str, - evaluation_type: str, -) -> pd.DataFrame: - """Run prediction and tag the results.""" - model.compute_features(test_dataset) - model.predict(test_dataset) - - results = test_dataset.metadata.copy() - results["trained_on_dataset"] = train_project - results["test_dataset"] = test_project - results["evaluation_type"] = evaluation_type - return results - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- -_DEFAULT_MODEL_OUTPUT_DIR = Path("models/generalisation") -_DEFAULT_RESULTS_OUTPUT_DIR = Path("results/generalisation") -_DEFAULT_TRAIN_PARQUET = Path("train_extra_small/train.parquet") -_DEFAULT_TRAIN_PREDS = Path("train_extra_small/train_preds.csv") -_DEFAULT_BIOLOGICAL_VALIDATION_DIR = Path( - "held_out_projects/biological_validation/annotated" -) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - - -@app.command() -def main( - train_parquet: Annotated[ - Path, typer.Option(help="Combined train_extra_small parquet file.") - ] = _DEFAULT_TRAIN_PARQUET, - train_predictions: Annotated[ - Path, typer.Option(help="Combined train_extra_small predictions CSV.") - ] = _DEFAULT_TRAIN_PREDS, - biological_validation_dir: Annotated[ - Path, - typer.Option( - help=( - "Directory of biological validation annotated parquets used to map " - "experiment_name values to source labels." - ) - ), - ] = _DEFAULT_BIOLOGICAL_VALIDATION_DIR, - model_output_dir: Annotated[ - Path, typer.Option(help="Directory to save trained models.") - ] = _DEFAULT_MODEL_OUTPUT_DIR, - results_output_dir: Annotated[ - Path, typer.Option(help="Directory to save evaluation results.") - ] = _DEFAULT_RESULTS_OUTPUT_DIR, - koina_server_url: Annotated[ - Optional[str], typer.Option(help="Koina server URL override.") - ] = None, - koina_ssl: Annotated[bool, typer.Option(help="Use SSL for Koina server.")] = True, -) -> None: - """Evaluate calibrator generalisation across train_extra_small source datasets.""" - model_output_dir.mkdir(parents=True, exist_ok=True) - results_output_dir.mkdir(parents=True, exist_ok=True) - - if not train_parquet.exists(): - logger.error("Train parquet not found: %s", train_parquet) - raise typer.Exit(1) - if not train_predictions.exists(): - logger.error("Train predictions CSV not found: %s", train_predictions) - raise typer.Exit(1) - if not biological_validation_dir.exists(): - logger.error( - "Biological validation directory not found: %s", biological_validation_dir - ) - raise typer.Exit(1) - - annotate_train_source_labels( - train_parquet, train_predictions, biological_validation_dir - ) - - full_dataset = load_dataset(train_parquet, train_predictions) - datasets = split_dataset_by_source(full_dataset) - logger.info("Found %d source datasets: %s", len(datasets), list(datasets.keys())) - for source, dataset in datasets.items(): - logger.info(" %s: %d samples", source, len(dataset.metadata)) - - # Train-on-each, evaluate-on-all - all_results: List[pd.DataFrame] = [] - for train_project in datasets: - logger.info("=== Training on %s ===", train_project) - - train_ds, in_dist_test_ds = create_train_test_split(datasets[train_project]) - logger.info( - " train: %d, in-dist test: %d", - len(train_ds.metadata), - len(in_dist_test_ds.metadata), - ) - - calibrator = initialise_calibrator( - koina_server_url=koina_server_url, - koina_ssl=koina_ssl, - train_project=train_project, - ) - calibrator.fit(train_ds) - - model_path = model_output_dir / f"trained_on_{train_project}" - ProbabilityCalibrator.save(calibrator, model_path) - - # In-distribution evaluation - logger.info( - " Evaluating in-distribution on %s (%d samples)", - train_project, - len(in_dist_test_ds.metadata), - ) - all_results.append( - evaluate_model( - calibrator, - in_dist_test_ds, - train_project, - train_project, - "in_distribution", - ) - ) - - # Out-of-distribution evaluation - for test_project in datasets: - if test_project == train_project: - continue - test_ds = datasets[test_project] - logger.info( - " Evaluating out-of-distribution on %s (%d samples)", - test_project, - len(test_ds.metadata), - ) - all_results.append( - evaluate_model( - calibrator, - test_ds, - train_project, - test_project, - "out_of_distribution", - ) - ) - - # Combine and save - combined = pd.concat(all_results, ignore_index=True) - - # Drop large array columns to save space - array_cols = [c for c in ["mz_array", "intensity_array"] if c in combined.columns] - if array_cols: - combined = combined.drop(columns=array_cols) - - results_path = results_output_dir / "calibrator_generalisation_results.csv" - combined.to_csv(results_path, index=False) - logger.info("Results saved to %s", results_path) - - # Summary - logger.info("Evaluation summary:") - summary = ( - combined.groupby(["trained_on_dataset", "test_dataset", "evaluation_type"]) - .size() - .reset_index(name="num_samples") - ) - for _, row in summary.iterrows(): - logger.info( - " Trained on %s, tested on %s (%s): %d samples", - row["trained_on_dataset"], - row["test_dataset"], - row["evaluation_type"], - row["num_samples"], - ) - - -if __name__ == "__main__": - app() diff --git a/scripts/fdr_tool_comparison_preprocess.py b/scripts/fdr_tool_comparison_preprocess.py deleted file mode 100644 index 2f936042..00000000 --- a/scripts/fdr_tool_comparison_preprocess.py +++ /dev/null @@ -1,809 +0,0 @@ -"""Shared preprocessing helpers for FDR tool comparisons. - -PSM comparison and the external peptide score-mixture share: -1. Method-specific load / NovoBoard mass-delta → ProForma conversion. -2. Pair-gated NovoBoard target-decoy filters (equal twin counts). -3. :func:`filter_prediction_table` / :func:`filter_novoboard_prediction_table`. - -Labelled correctness uses Novor token matching; proteome-hit proxies use -PTM-stripped I→L substring search against an I→L FASTA haystack. - -Peptide score-mixture only then: -4. :func:`max_score_per_peptide` (no re-filtering). -5. NovoBoard max-target → twin-decoy peptide TDC helpers. -""" - -from __future__ import annotations - -import logging -import re -from functools import lru_cache -from pathlib import Path -from typing import Iterable - -import numpy as np -import pandas as pd -import polars as pl -import yaml -from instanovo.utils.metrics import Metrics -from instanovo.utils.residues import ResidueSet - -from winnow.utils.proteome import ( - _batch_peptide_substring_hits, - processed_peptide_for_match, - residue_token_count, -) - -logger = logging.getLogger(__name__) - -_REPO_ROOT = Path(__file__).resolve().parent.parent -DEFAULT_RESIDUES_YAML = _REPO_ROOT / "winnow" / "configs" / "residues.yaml" - -MIN_PEPTIDE_LENGTH = 8 -# Labelled / reference sets use Novor agreement, so short peptides are valid. -# Keep a non-empty-key floor only. Unlabelled sets keep MIN_PEPTIDE_LENGTH -# because correctness is proteome substring membership. -LABELLED_MIN_PEPTIDE_LENGTH = 1 -_UNIMOD_RE = re.compile(r"\[UNIMOD:\d+\]") -_MOD_SQUARE = re.compile(r"\[.*?\]") -_MOD_PAREN = re.compile(r"\(.*?\)") -_NOVOBOARD_TO_PROFORMA = { - "C(+57.02)": "C[UNIMOD:4]", - "M(+15.99)": "M[UNIMOD:35]", - "N(+0.98)": "N[UNIMOD:7]", - "Q(+0.98)": "Q[UNIMOD:7]", - "S(+79.97)": "S[UNIMOD:21]", - "T(+79.97)": "T[UNIMOD:21]", - "Y(+79.97)": "Y[UNIMOD:21]", -} - - -def normalize_peptide_key(peptide: object) -> str: - """Normalise sequence-only peptide identity (strip PTMs, I→L).""" - if isinstance(peptide, list): - peptide = "".join(str(token) for token in peptide) - if pd.isna(peptide) or not isinstance(peptide, str): - return "" - if len(peptide) > 4 and peptide[1] == "." and peptide[-2] == ".": - peptide = peptide[2:-2] - seq = _MOD_SQUARE.sub("", peptide) - seq = _MOD_PAREN.sub("", seq) - seq = "".join(c for c in seq if c.isalpha()) - return seq.replace("I", "L") - - -def has_unsupported_unimod(peptide: object) -> bool: - """Return True when *peptide* still contains an unsupported ``[UNIMOD:n]`` token.""" - if pd.isna(peptide) or not isinstance(peptide, str): - return True - return bool(_UNIMOD_RE.search(peptide)) - - -def novoboard_to_proforma(peptide: object) -> object: - """Convert NovoBoard's supported mass-delta notation to ProForma.""" - if pd.isna(peptide) or not isinstance(peptide, str): - return peptide - converted = peptide - for novoboard_mod, proforma_mod in _NOVOBOARD_TO_PROFORMA.items(): - converted = converted.replace(novoboard_mod, proforma_mod) - return converted - - -def sequence_only_correct_prediction(sequence: object, prediction: object) -> bool: - """Full sequence equality after PTM stripping and I/L normalisation.""" - sequence_key = normalize_peptide_key(sequence) - prediction_key = normalize_peptide_key(prediction) - return bool(sequence_key) and sequence_key == prediction_key - - -def sequence_only_correctness_mask( - sequences: pd.Series, predictions: pd.Series -) -> np.ndarray: - """Vectorized strip-PTM I→L equality (not for labelled Novor eval).""" - return np.array( - [ - sequence_only_correct_prediction(sequence, prediction) - for sequence, prediction in zip(sequences, predictions) - ], - dtype=bool, - ) - - -def load_residue_masses(residues_yaml: Path | None = None) -> dict[str, float]: - """Load residue masses from Winnow's residues YAML.""" - path = residues_yaml if residues_yaml is not None else DEFAULT_RESIDUES_YAML - with path.open(encoding="utf-8") as handle: - return yaml.safe_load(handle)["residue_masses"] - - -@lru_cache(maxsize=4) -def _metrics_from_residue_masses_frozen( - residues_items: tuple[tuple[str, float], ...], -) -> Metrics: - residue_masses = dict(residues_items) - return Metrics( - residue_set=ResidueSet(residue_masses=residue_masses), - isotope_error_range=(0, 1), - ) - - -def metrics_from_residue_masses(residue_masses: dict[str, float]) -> Metrics: - """Build an InstaNovo ``Metrics`` instance for Novor matching.""" - items = tuple(sorted((str(k), float(v)) for k, v in residue_masses.items())) - return _metrics_from_residue_masses_frozen(items) - - -def novor_correct_prediction( - sequence: object, - prediction: object, - metrics: Metrics, -) -> bool: - """Winnow/InstaNovo Novor correctness: full residue-token match.""" - if isinstance(sequence, list): - gt = sequence - elif pd.isna(sequence) or not isinstance(sequence, str) or not sequence: - return False - else: - gt = metrics._split_peptide(sequence) - - if isinstance(prediction, list): - pred = prediction - elif pd.isna(prediction) or not isinstance(prediction, str) or not prediction: - return False - else: - pred = metrics._split_peptide(prediction) - - if not gt or not pred: - return False - num_matches = metrics._novor_match(gt, pred) - return bool(num_matches == len(gt) == len(pred)) - - -def novor_correctness_mask( - sequences: pd.Series | Iterable[object], - predictions: pd.Series | Iterable[object], - *, - residue_masses: dict[str, float] | None = None, - metrics: Metrics | None = None, -) -> np.ndarray: - """Vectorized Novor correctness (same rule as ``DatabaseGroundedFDRControl.fit``).""" - if metrics is None: - masses = residue_masses if residue_masses is not None else load_residue_masses() - metrics = metrics_from_residue_masses(masses) - return np.array( - [ - novor_correct_prediction(sequence, prediction, metrics) - for sequence, prediction in zip(sequences, predictions) - ], - dtype=bool, - ) - - -def dedupe_best_score_per_peptide( - df: pd.DataFrame, peptide_col: str, score_col: str -) -> pd.DataFrame: - """Keep the highest-scoring row per peptide key.""" - return ( - df.sort_values(score_col, ascending=False) - .groupby(peptide_col, as_index=False) - .first() - ) - - -def compute_q_values(fdr: np.ndarray) -> np.ndarray: - """Convert ranked FDR estimates to q-values using suffix minima.""" - values = np.asarray(fdr, dtype=float) - q_values = np.empty_like(values) - fdr_min = np.inf - for i in range(len(values) - 1, -1, -1): - fdr_min = min(fdr_min, values[i]) - q_values[i] = fdr_min - return q_values - - -def monotonize_q_by_confidence( - confidence: np.ndarray, q_value: np.ndarray -) -> np.ndarray: - """Enforce non-increasing q-values when confidence increases.""" - order = np.argsort(-np.asarray(confidence, dtype=float)) - q_sorted = np.asarray(q_value, dtype=float)[order] - q_mono = np.empty_like(q_sorted) - q_min = np.inf - for i in range(len(q_sorted) - 1, -1, -1): - if q_sorted[i] > q_min: - q_mono[i] = q_min - else: - q_mono[i] = q_sorted[i] - q_min = q_sorted[i] - out = np.empty_like(q_mono) - out[order] = q_mono - return out - - -def add_peptide_key( - df: pd.DataFrame, - peptide_col: str, - *, - key_col: str = "peptide_key", -) -> pd.DataFrame: - """Append normalised peptide keys.""" - work = df.copy() - work[key_col] = work[peptide_col].map(normalize_peptide_key) - return work - - -def filter_prediction_table( - df: pd.DataFrame, - peptide_col: str, - *, - min_length: int = MIN_PEPTIDE_LENGTH, - key_col: str = "peptide_key", - drop_unsupported_mods: bool = True, - log: bool = True, -) -> pd.DataFrame: - """Drop unsupported mods and peptides shorter than *min_length* (normalised).""" - work = add_peptide_key(df, peptide_col, key_col=key_col) - before = len(work) - if drop_unsupported_mods: - work = work[~work[peptide_col].map(has_unsupported_unimod)].copy() - work = work[work[key_col].str.len() >= min_length].copy() - dropped = before - len(work) - if log and dropped: - logger.info( - "Filtered %d/%d rows (unsupported mods and/or length < %d) on %s", - dropped, - before, - min_length, - peptide_col, - ) - return work.reset_index(drop=True) - - -def filter_novoboard_prediction_table( - df: pd.DataFrame, - *, - peptide_col: str = "Peptide", - min_length: int = MIN_PEPTIDE_LENGTH, - key_col: str = "_peptide_key", - log: bool = True, -) -> pd.DataFrame: - """Convert NovoBoard modifications to ProForma, then apply shared filters.""" - work = df.copy() - work[peptide_col] = work[peptide_col].map(novoboard_to_proforma) - return filter_prediction_table( - work, - peptide_col, - min_length=min_length, - key_col=key_col, - log=log, - ) - - -def filter_novoboard_target_decoy_pairs( - target: pd.DataFrame, - decoy: pd.DataFrame, - *, - peptide_col: str = "Peptide", - min_length: int = MIN_PEPTIDE_LENGTH, - key_col: str = "_peptide_key", - log: bool = True, -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Filter NovoBoard target/decoy as spectrum twins so pair counts stay equal. - - Both sides are converted to ProForma and passed through the shared - mod/length filter. Only ``_pair_key`` values present in **both** filtered - tables are kept, so dropping an unsupported-mod decoy also drops its - target (and vice versa). - """ - if "_pair_key" not in target.columns or "_pair_key" not in decoy.columns: - raise ValueError( - "filter_novoboard_target_decoy_pairs requires '_pair_key' on both tables" - ) - - target_f = filter_novoboard_prediction_table( - target, - peptide_col=peptide_col, - min_length=min_length, - key_col=key_col, - log=log, - ) - decoy_f = filter_novoboard_prediction_table( - decoy, - peptide_col=peptide_col, - min_length=min_length, - key_col=key_col, - log=log, - ) - - def _valid_pair_keys(series: pd.Series) -> set[str]: - keys = series.astype(str) - return {k for k in keys if k and k != "nan"} - - shared_keys = _valid_pair_keys(target_f["_pair_key"]) & _valid_pair_keys( - decoy_f["_pair_key"] - ) - target_out = target_f[target_f["_pair_key"].astype(str).isin(shared_keys)].copy() - decoy_out = decoy_f[decoy_f["_pair_key"].astype(str).isin(shared_keys)].copy() - n_target = target_out["_pair_key"].astype(str).nunique() - n_decoy = decoy_out["_pair_key"].astype(str).nunique() - if n_target != n_decoy: - raise AssertionError( - f"Pair filter left unequal twin counts: targets={n_target} decoys={n_decoy}" - ) - if log: - before_pairs = len( - _valid_pair_keys(target_f["_pair_key"]) - | _valid_pair_keys(decoy_f["_pair_key"]) - ) - logger.info( - "NovoBoard pair filter: %s → %s twin spectra " - "(target rows %s → %s, decoy rows %s → %s)", - before_pairs, - n_target, - len(target_f), - len(target_out), - len(decoy_f), - len(decoy_out), - ) - return target_out, decoy_out - - -def restrict_winnow_to_novoboard_spectra( - winnow: pd.DataFrame, novoboard: pd.DataFrame -) -> pd.DataFrame: - """Trim Winnow to NovoBoard twin-valid spectra under the subset invariant. - - After shared peptide filters and NovoBoard pair-gating, NovoBoard targets are - expected to be a subset of Winnow-filtered spectra (same InstaNovo - predictions; NovoBoard additionally drops pairs whose decoy fails). The - shared pool is therefore the NovoBoard spectrum set: only Winnow is trimmed. - - Raises: - AssertionError: If any NovoBoard spectrum is missing from Winnow. - """ - winnow_ids = set(winnow["spectrum_id"].astype(str)) - novoboard_ids = set(novoboard["spectrum_id"].astype(str)) - only_novoboard = novoboard_ids - winnow_ids - if only_novoboard: - examples = sorted(only_novoboard)[:5] - raise AssertionError( - "NovoBoard twin-valid spectra are not a subset of Winnow-filtered " - f"spectra ({len(only_novoboard)} missing); examples={examples}. " - "Expected identical InstaNovo predictions after ProForma remapping." - ) - winnow_shared = winnow[winnow["spectrum_id"].astype(str).isin(novoboard_ids)].copy() - logger.info( - "Shared spectrum pool: %d spectra (trimmed Winnow=%d; NovoBoard unchanged)", - len(novoboard_ids), - len(winnow) - len(winnow_shared), - ) - return winnow_shared - - -def assert_shared_prediction_keys( - winnow: pd.DataFrame, - novoboard: pd.DataFrame, - *, - winnow_peptide_col: str = "prediction", - novoboard_peptide_col: str = "Peptide", -) -> None: - """Require I/L-normalised prediction identity on the shared spectrum pool.""" - w_ids = winnow["spectrum_id"].astype(str) - nb_ids = novoboard["spectrum_id"].astype(str) - if w_ids.nunique() != len(winnow) or nb_ids.nunique() != len(novoboard): - raise AssertionError( - "Shared-pool tables must have one row per spectrum_id before " - f"prediction-key assert (winnow rows={len(winnow)} unique={w_ids.nunique()}, " - f"novoboard rows={len(novoboard)} unique={nb_ids.nunique()})" - ) - merged = ( - winnow[["spectrum_id", winnow_peptide_col]] - .assign(spectrum_id=w_ids) - .merge( - novoboard[["spectrum_id", novoboard_peptide_col]].assign( - spectrum_id=nb_ids - ), - on="spectrum_id", - how="inner", - validate="one_to_one", - ) - ) - if len(merged) != len(winnow) or len(merged) != len(novoboard): - raise AssertionError( - "Shared-pool spectrum_id join is not 1:1 " - f"(winnow={len(winnow)} novoboard={len(novoboard)} inner={len(merged)})" - ) - w_keys = merged[winnow_peptide_col].map(normalize_peptide_key) - nb_keys = merged[novoboard_peptide_col].map(normalize_peptide_key) - mismatch = w_keys != nb_keys - if bool(mismatch.any()): - bad = merged.loc[ - mismatch, ["spectrum_id", winnow_peptide_col, novoboard_peptide_col] - ] - examples = bad.head(5).to_dict(orient="records") - raise AssertionError( - "Winnow and NovoBoard predictions disagree after I/L-normalised " - f"peptide keys ({int(mismatch.sum())} spectra); examples={examples}" - ) - - -def label_series_by_spectrum_id(winnow: pd.DataFrame, label_col: str) -> pd.Series: - """Map ``spectrum_id`` → boolean label from a Winnow table.""" - if label_col not in winnow.columns: - raise KeyError(f"Missing label column {label_col!r}") - ids = winnow["spectrum_id"].astype(str) - if ids.duplicated().any(): - raise AssertionError( - f"Duplicate spectrum_id values when building {label_col} label map" - ) - return pd.Series( - winnow[label_col].astype(bool).to_numpy(), - index=ids, - name=label_col, - ) - - -def attach_labels_by_spectrum_id( - novoboard: pd.DataFrame, - label_by_id: pd.Series, - *, - label_col: str, -) -> pd.DataFrame: - """Attach a shared label column to NovoBoard rows by ``spectrum_id``.""" - out = novoboard.copy() - mapped = out["spectrum_id"].astype(str).map(label_by_id) - if mapped.isna().any(): - missing = out.loc[mapped.isna(), "spectrum_id"].astype(str).head(5).tolist() - raise AssertionError( - f"NovoBoard rows missing shared {label_col} labels; examples={missing}" - ) - out[label_col] = mapped.astype(bool) - return out - - -def _best_alc_per_pair_key(df: pd.DataFrame) -> pd.DataFrame: - """One highest-ALC row per ``_pair_key``.""" - work = df.dropna(subset=["ALC (%)", "_pair_key"]) - work = work[work["_pair_key"].astype(str) != "nan"] - return ( - work.sort_values("ALC (%)", ascending=False) - .groupby("_pair_key", as_index=False) - .first() - ) - - -def proteome_hit_mask( - peptides: pd.Series | Iterable[str], - haystack: str, - *, - min_length: int = MIN_PEPTIDE_LENGTH, -) -> np.ndarray: - """True when normalised peptide key (length ≥ *min_length*) hits the proteome.""" - keys = [normalize_peptide_key(p) for p in peptides] - eligible = [bool(k) and len(k) >= min_length for k in keys] - unique_keys = sorted({k for k, ok in zip(keys, eligible) if ok}) - hit_map: dict[str, bool] = {} - if unique_keys: - hits = _batch_peptide_substring_hits(unique_keys, haystack) - hit_map = dict(zip(unique_keys, hits)) - return np.array( - [bool(eligible[i] and hit_map.get(keys[i], False)) for i in range(len(keys))], - dtype=bool, - ) - - -def filter_and_annotate_preds( - preds: pl.DataFrame, - haystack: str, - metrics: Metrics, - min_residue_length: int, -) -> pl.DataFrame: - """Filter short peptides and annotate ``proteome_hit`` via ``winnow.utils.proteome``. - - Args: - preds: Polars frame with a ``prediction`` column. - haystack: I/L-normalised FASTA haystack from ``load_proteome_haystack``. - metrics: InstaNovo ``Metrics`` (uses ``metrics.residue_set`` for length). - min_residue_length: Drop PSMs with fewer than this many residue tokens. - """ - residue_set = metrics.residue_set - n_tok = preds["prediction"].map_elements( - lambda x: residue_token_count(x, residue_set), - return_dtype=pl.Int32, - ) - filtered = preds.with_columns(n_tok.alias("_n_residue_tokens")).filter( - pl.col("_n_residue_tokens") >= min_residue_length - ) - processed = filtered["prediction"].map_elements( - lambda x: processed_peptide_for_match(x) if isinstance(x, str) else "", - return_dtype=pl.Utf8, - ) - hits = _batch_peptide_substring_hits(processed.to_list(), haystack) - return filtered.drop("_n_residue_tokens").with_columns( - pl.Series("proteome_hit", hits, dtype=pl.Boolean) - ) - - -def max_score_per_peptide( - df: pd.DataFrame, - key_col: str, - score_col: str, -) -> pd.DataFrame: - """Keep the max-scoring row per peptide key (no filtering). - - Call after :func:`filter_prediction_table` or - :func:`filter_novoboard_prediction_table` so all methods share the same - filter → max-dedupe sequence in the peptide score-mixture benchmark. - """ - work = df.dropna(subset=[score_col, key_col]) - work = work[work[key_col].astype(str) != ""] - return dedupe_best_score_per_peptide(work, key_col, score_col).reset_index( - drop=True - ) - - -def confidence_to_log_prob(confidence: pd.Series | np.ndarray) -> np.ndarray: - """Map raw InstaNovo confidence in (0, 1] to Glissade-style log probabilities.""" - conf = np.asarray(confidence, dtype=float) - return np.log(np.clip(conf, 1e-300, 1.0)) - - -def _load_mgf_title_to_scan(mgf_path: Path) -> dict[str, str]: - """Parse TITLE→SCANS mapping from an MGF file.""" - mapping: dict[str, str] = {} - title: str | None = None - scan: str | None = None - with open(mgf_path, encoding="utf-8", errors="replace") as handle: - for line in handle: - value = line.strip() - if value.startswith("TITLE="): - title = value.removeprefix("TITLE=") - elif value.startswith("SCANS="): - scan = value.removeprefix("SCANS=") - elif value == "END IONS" and title is not None and scan is not None: - mapping[title] = scan - title = None - scan = None - return mapping - - -def attach_novoboard_pair_keys( - target: pd.DataFrame, - decoy: pd.DataFrame, - *, - novoboard_dir: Path, - split_prefix: str, -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Attach ``_pair_key`` using Scan identity or decoy-MGF TITLE→SCANS mapping.""" - target_out = target.copy() - decoy_out = decoy.copy() - if "Scan" not in target_out.columns or "Scan" not in decoy_out.columns: - raise ValueError("NovoBoard tables require a 'Scan' column for twin pairing") - - target_key = target_out["Scan"].astype(str) - decoy_key = decoy_out["Scan"].astype(str) - best_decoy_key = decoy_key - best_overlap = len(set(target_key) & set(decoy_key)) - - mgf_path = novoboard_dir.parent / f"{split_prefix}.mgf" - if mgf_path.is_file(): - title_to_scan = _load_mgf_title_to_scan(mgf_path) - if title_to_scan: - mapped = decoy_key.map(title_to_scan) - mapped_overlap = len(set(target_key) & set(mapped.dropna())) - if mapped_overlap > best_overlap: - best_decoy_key = mapped - best_overlap = mapped_overlap - - target_out["_pair_key"] = target_key - decoy_out["_pair_key"] = best_decoy_key - logger.info( - "NovoBoard %s pair-key overlap: target=%d decoy=%d overlap=%d", - split_prefix, - target_key.nunique(dropna=True), - pd.Series(best_decoy_key).nunique(dropna=True), - best_overlap, - ) - return target_out, decoy_out - - -def prepare_novoboard_decoy_by_pair( - decoy_df: pd.DataFrame, - *, - min_length: int = MIN_PEPTIDE_LENGTH, - already_filtered: bool = False, -) -> pd.DataFrame: - """Index the best decoy row per ``_pair_key``. - - Args: - decoy_df: Decoy table with ``_pair_key``. Prefer pair-gated output from - :func:`filter_novoboard_target_decoy_pairs`. - already_filtered: When True, skip ProForma/mod/length filtering (caller - already pair-filtered). - """ - if "_pair_key" not in decoy_df.columns: - raise ValueError("NovoBoard twin TDC requires '_pair_key' on decoy") - if already_filtered: - decoy = decoy_df.copy() - if "_peptide_key" not in decoy.columns: - if "peptide_key" in decoy.columns: - decoy["_peptide_key"] = decoy["peptide_key"] - else: - decoy = add_peptide_key(decoy, "Peptide", key_col="_peptide_key") - elif "_peptide_key" in decoy_df.columns and decoy_df["_peptide_key"].notna().all(): - decoy = decoy_df.copy() - else: - decoy = filter_novoboard_prediction_table( - decoy_df, min_length=min_length, key_col="_peptide_key" - ) - decoy = decoy.dropna(subset=["ALC (%)", "_pair_key"]) - decoy = decoy[ - (decoy["_pair_key"].astype(str) != "nan") & (decoy["_peptide_key"] != "") - ] - return _best_alc_per_pair_key(decoy).set_index("_pair_key", drop=False) - - -def novoboard_psm_tdc( - target_df: pd.DataFrame, - decoy_df: pd.DataFrame, - *, - min_length: int = MIN_PEPTIDE_LENGTH, -) -> pd.DataFrame: - """Recompute NovoBoard's pooled PSM TDC after pair-gated filtering. - - Target and decoy are filtered as spectrum twins so unsupported-mod drops - remove the pair. Competition uses one best-ALC row per twin on each side, - guaranteeing ``sum(is_target) == sum(~is_target)``. - """ - target, decoy = filter_novoboard_target_decoy_pairs( - target_df, decoy_df, min_length=min_length - ) - target = _best_alc_per_pair_key(target).assign(is_target=True) - decoy = _best_alc_per_pair_key(decoy).assign(is_target=False) - n_target = int(target["_pair_key"].nunique()) - n_decoy = int(decoy["_pair_key"].nunique()) - if n_target != n_decoy or len(target) != len(decoy): - raise AssertionError( - f"PSM TDC unbalanced after pair gate: " - f"target_rows={len(target)} decoy_rows={len(decoy)} " - f"target_pairs={n_target} decoy_pairs={n_decoy}" - ) - - combined = pd.concat([target, decoy], ignore_index=True, sort=False) - combined = combined.sort_values( - ["ALC (%)", "is_target"], ascending=[False, False] - ).reset_index(drop=True) - return _assign_cumulative_tdc_fdr(combined) - - -def _prepare_targets_for_twin_tdc( - target_df: pd.DataFrame, - decoy_df: pd.DataFrame, - *, - min_length: int, - decoy_by_pair: pd.DataFrame | None, -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Pair-gate or twin-filter targets and return ``(target, decoy_by_pair)``.""" - if "_pair_key" not in target_df.columns: - raise ValueError("NovoBoard twin TDC requires '_pair_key' on target") - if decoy_by_pair is None and "_pair_key" not in decoy_df.columns: - raise ValueError("NovoBoard twin TDC requires '_pair_key' on decoy") - - if decoy_by_pair is None: - target, decoy = filter_novoboard_target_decoy_pairs( - target_df, decoy_df, min_length=min_length, log=False - ) - decoy_by_pair = prepare_novoboard_decoy_by_pair( - decoy, min_length=min_length, already_filtered=True - ) - return target, decoy_by_pair - - target = target_df.copy() - if "_peptide_key" not in target.columns: - if "peptide_key" in target.columns: - target["_peptide_key"] = target["peptide_key"] - else: - target = filter_novoboard_prediction_table( - target, - min_length=min_length, - key_col="_peptide_key", - log=False, - ) - twin_keys = set(decoy_by_pair.index.astype(str)) - target = target[target["_pair_key"].astype(str).isin(twin_keys)] - return target, decoy_by_pair - - -def _assign_cumulative_tdc_fdr(combined: pd.DataFrame) -> pd.DataFrame: - """Add estimated FDR / q-value columns for a balanced target-decoy table.""" - out = combined.copy() - n_target = out["is_target"].astype(int).cumsum() - n_decoy = (~out["is_target"]).astype(int).cumsum() - out["estimated_fdr"] = np.divide( - n_decoy, - n_target, - out=np.ones(len(out), dtype=float), - where=n_target > 0, - ) - out["estimated_q_value"] = np.nan - target_mask = out["is_target"].to_numpy() - out.loc[target_mask, "estimated_q_value"] = compute_q_values( - out.loc[target_mask, "estimated_fdr"].to_numpy() - ) - return out - - -def novoboard_max_target_twin_decoy_tdc( - target_df: pd.DataFrame, - decoy_df: pd.DataFrame, - *, - min_length: int = MIN_PEPTIDE_LENGTH, - target_peptide_keys: set[str] | None = None, - decoy_by_pair: pd.DataFrame | None = None, - log_missing_twins: bool = True, -) -> pd.DataFrame: - """Peptide TDC: pair-gate, max ALC per target peptide, twin decoy by ``_pair_key``. - - Returns the combined ranked table with ``is_target``, ``estimated_fdr``, and - ``estimated_q_value`` (targets only). Targets without a twin-valid decoy are - dropped; the competition table is always 1:1. - - Args: - decoy_by_pair: Optional precomputed output of - :func:`prepare_novoboard_decoy_by_pair` from an already pair-gated - decoy table. When omitted, target/decoy are pair-filtered together. - """ - target, decoy_by_pair = _prepare_targets_for_twin_tdc( - target_df, - decoy_df, - min_length=min_length, - decoy_by_pair=decoy_by_pair, - ) - target = target.dropna(subset=["ALC (%)", "_pair_key"]) - target = target[ - (target["_pair_key"].astype(str) != "nan") & (target["_peptide_key"] != "") - ] - - if target_peptide_keys is not None: - target = target[target["_peptide_key"].isin(target_peptide_keys)] - - # Max-score only among twin-valid targets. - target_best = max_score_per_peptide(target, "_peptide_key", "ALC (%)") - pair_keys = target_best["_pair_key"].astype(str) - has_twin = pair_keys.isin(decoy_by_pair.index.astype(str)) - n_missing_twin = int((~has_twin).sum()) - if log_missing_twins and n_missing_twin: - logger.warning( - "NovoBoard twin-decoy TDC dropped %d/%d max-target peptides without twin", - n_missing_twin, - len(target_best), - ) - target_keep = target_best.loc[has_twin].copy() - if target_keep.empty: - return pd.DataFrame( - columns=[ - "spectrum_id", - "Peptide", - "ALC (%)", - "_peptide_key", - "_pair_key", - "is_target", - "estimated_fdr", - "estimated_q_value", - ] - ) - - decoy_keep = decoy_by_pair.loc[target_keep["_pair_key"].astype(str)].copy() - decoy_keep = decoy_keep.reset_index(drop=True) - target_keep = target_keep.assign(is_target=True).reset_index(drop=True) - decoy_keep = decoy_keep.assign(is_target=False) - if len(target_keep) != len(decoy_keep): - raise AssertionError( - f"Peptide TDC unbalanced: targets={len(target_keep)} decoys={len(decoy_keep)}" - ) - # Preserve 1:1 balance: one decoy row per retained target (do not dedupe decoys). - combined = pd.concat([target_keep, decoy_keep], ignore_index=True, sort=False) - combined = combined.sort_values( - ["ALC (%)", "is_target"], ascending=[False, False] - ).reset_index(drop=True) - return _assign_cumulative_tdc_fdr(combined) diff --git a/scripts/fdr_tool_comparison_summaries.py b/scripts/fdr_tool_comparison_summaries.py deleted file mode 100644 index b48d1dcc..00000000 --- a/scripts/fdr_tool_comparison_summaries.py +++ /dev/null @@ -1,533 +0,0 @@ -"""Summary tables for Winnow / NovoBoard / Glissade FDR tool comparisons. - -Produces two long-form CSVs: - -- ``*_acceptance.csv``: accepted counts and recovery at q-value thresholds. -- ``*_error_gain.csv``: observed FDP, excess over nominal FDR, optional mean - absolute q-value deviation vs a database-grounded reference (calibrated-score - DBG for Winnow; raw-score / ALC DBG for NovoBoard and Glissade), and relative - gain/loss of a primary method vs each comparator. -""" - -from __future__ import annotations - -import logging -from pathlib import Path -from typing import Iterable, Sequence, cast - -import numpy as np -import pandas as pd - -from scripts.fdr_tool_comparison_preprocess import compute_q_values - -logger = logging.getLogger(__name__) - -SUMMARY_THRESHOLDS: list[float] = [0.01, 0.05, 0.10] -# Match ``DatabaseGroundedFDRControl`` / PSM comparison default. -_DB_GROUNDED_DROP = 10 - -_KEY_COLS = ["dataset", "panel", "level", "method", "q_value_threshold"] - - -def _slug_comparator(name: str) -> str: - """Map a method label to a filesystem-/column-safe slug.""" - return ( - name.lower() - .replace(" ", "_") - .replace("(", "") - .replace(")", "") - .replace("-", "_") - ) - - -def acceptance_rows_from_q( - *, - dataset: str, - panel: str, - level: str, - method: str, - q_value: np.ndarray, - thresholds: Sequence[float] = SUMMARY_THRESHOLDS, - label_mask: np.ndarray | None = None, - recovery_denom: int | None = None, -) -> list[dict[str, object]]: - """Build acceptance/yield rows for one method at each q-value threshold. - - Args: - dataset: Dataset key (e.g. ``helaqc``). - panel: Evaluation panel (e.g. ``labelled_test``, ``unlabelled``, ``external``). - level: ``psm`` or ``peptide``. - method: Method display label. - q_value: Per-row estimated q-values. - thresholds: Nominal FDR thresholds. - label_mask: Optional boolean correctness / proteome-hit labels aligned with - ``q_value``. When provided, ``n_correct`` is filled. - recovery_denom: Denominator for recovery percentage. Defaults to the number - of True labels when ``label_mask`` is given. - - Returns: - One dict per threshold. - """ - q = np.asarray(q_value, dtype=float) - labels = None if label_mask is None else np.asarray(label_mask, dtype=bool) - if labels is not None and len(labels) != len(q): - raise ValueError( - f"label_mask length {len(labels)} does not match q_value length {len(q)}" - ) - if recovery_denom is None and labels is not None: - recovery_denom = int(labels.sum()) - - rows: list[dict[str, object]] = [] - for threshold in thresholds: - valid = ~np.isnan(q) - accepted = valid & (q <= threshold) - n_accepted = int(accepted.sum()) - n_correct: float | int = np.nan - recovery_pct: float = np.nan - if labels is not None: - n_correct = int((accepted & labels).sum()) - if recovery_denom and recovery_denom > 0: - recovery_pct = 100.0 * float(n_correct) / float(recovery_denom) - rows.append( - { - "dataset": dataset, - "panel": panel, - "level": level, - "method": method, - "q_value_threshold": float(threshold), - "n_accepted": n_accepted, - "n_correct": n_correct, - "recovery_pct": recovery_pct, - } - ) - return rows - - -def observed_fdp_at_thresholds( - q_value: np.ndarray, - label_mask: np.ndarray, - thresholds: Sequence[float] = SUMMARY_THRESHOLDS, -) -> list[float]: - """Return observed false-discovery proportion among accepted rows at each threshold. - - ``label_mask`` is True for correct (or proteome-hit) rows. Observed FDP is - ``1 - n_correct / n_accepted`` when any rows are accepted, else NaN. - """ - q = np.asarray(q_value, dtype=float) - labels = np.asarray(label_mask, dtype=bool) - if len(q) != len(labels): - raise ValueError( - f"label_mask length {len(labels)} does not match q_value length {len(q)}" - ) - out: list[float] = [] - for threshold in thresholds: - valid = ~np.isnan(q) - accepted = valid & (q <= threshold) - n_accepted = int(accepted.sum()) - if n_accepted == 0: - out.append(float("nan")) - continue - n_correct = int((accepted & labels).sum()) - out.append(1.0 - n_correct / n_accepted) - return out - - -def database_grounded_q_from_labels( - scores: np.ndarray, - labels: np.ndarray, - *, - drop: int = _DB_GROUNDED_DROP, -) -> np.ndarray: - """In-sample database-grounded q-values from ranked scores and boolean labels. - - Builds the empirical precision curve ``1 - cumsum(correct) / rank`` on scores - sorted descending (same construction as the proteome-hit shortcut in the PSM - comparison), drops the first *drop* ranks from the FDR map, assigns FDR by - score lookup, then converts to q-values. - - Args: - scores: Ranking scores (higher = more confident). - labels: Boolean correctness / hit labels aligned with *scores*. - drop: Leading ranks excluded from the FDR map (default 10). - - Returns: - q-value array aligned with *scores*. - """ - scores_a = np.asarray(scores, dtype=float) - labels_a = np.asarray(labels, dtype=bool) - n = len(scores_a) - if n == 0: - return np.asarray([], dtype=float) - if len(labels_a) != n: - raise ValueError( - f"labels length {len(labels_a)} does not match scores length {n}" - ) - - order = np.argsort(-scores_a, kind="mergesort") - precision = np.cumsum(labels_a[order].astype(float)) / np.arange(1, n + 1) - fdr_ranked = 1.0 - precision - drop_eff = min(drop, max(0, n - 1)) - fit_scores = scores_a[order][drop_eff:] - fit_fdr = fdr_ranked[drop_eff:] - n_fit = len(fit_scores) - - idx = np.searchsorted(-fit_scores, -scores_a, side="left") - fdr = np.empty(n, dtype=float) - below = (idx == n_fit) & (scores_a < fit_scores[-1]) - above = (idx == 0) & (scores_a > fit_scores[0]) - normal = ~(below | above) - fdr[below] = 1.0 - fdr[above] = float(fit_fdr[0]) - fdr[normal] = fit_fdr[np.clip(idx[normal], 0, n_fit - 1)] - - q_sorted = compute_q_values(fdr[order]) - q = np.empty(n, dtype=float) - q[order] = q_sorted - return q - - -def mean_abs_q_dev_vs_reference( - q_method: np.ndarray, - q_ref: np.ndarray, - thresholds: Sequence[float] = SUMMARY_THRESHOLDS, -) -> list[float]: - """Mean absolute q-value deviation vs a row-aligned reference at each threshold. - - For each threshold, restrict to rows accepted by either method - (``q_method <= t`` or ``q_ref <= t``) with finite q for both, then report - ``mean(|q_method - q_ref|)``. - """ - q_m = np.asarray(q_method, dtype=float) - q_r = np.asarray(q_ref, dtype=float) - if len(q_m) != len(q_r): - raise ValueError( - f"q_ref length {len(q_r)} does not match q_method length {len(q_m)}" - ) - out: list[float] = [] - for threshold in thresholds: - both_finite = ~np.isnan(q_m) & ~np.isnan(q_r) - either_accepted = both_finite & ((q_m <= threshold) | (q_r <= threshold)) - if not np.any(either_accepted): - out.append(float("nan")) - continue - out.append(float(np.mean(np.abs(q_m[either_accepted] - q_r[either_accepted])))) - return out - - -def error_rows_from_q( - *, - dataset: str, - panel: str, - level: str, - method: str, - q_value: np.ndarray, - thresholds: Sequence[float] = SUMMARY_THRESHOLDS, - label_mask: np.ndarray | None = None, - q_ref: np.ndarray | None = None, - observed_fdp: Sequence[float] | None = None, -) -> list[dict[str, object]]: - """Build error-metric rows for one method (without relative-gain columns). - - Args: - observed_fdp: Optional precomputed FDP values (e.g. from a mixture - benchmark). When omitted, FDP is derived from ``label_mask`` if given. - """ - if observed_fdp is not None and len(observed_fdp) != len(thresholds): - raise ValueError("observed_fdp length must match thresholds") - if observed_fdp is None and label_mask is not None: - fdp_values = observed_fdp_at_thresholds(q_value, label_mask, thresholds) - elif observed_fdp is not None: - fdp_values = [float(x) for x in observed_fdp] - else: - fdp_values = [float("nan")] * len(thresholds) - - if q_ref is not None: - q_dev = mean_abs_q_dev_vs_reference(q_value, q_ref, thresholds) - else: - q_dev = [float("nan")] * len(thresholds) - - rows: list[dict[str, object]] = [] - for threshold, fdp, dev in zip(thresholds, fdp_values, q_dev): - fdp_f = float(fdp) - excess = fdp_f - float(threshold) if np.isfinite(fdp_f) else float("nan") - rows.append( - { - "dataset": dataset, - "panel": panel, - "level": level, - "method": method, - "q_value_threshold": float(threshold), - "observed_fdp": fdp_f, - "fdp_excess": excess, - "mean_abs_q_dev_vs_db": float(dev), - } - ) - return rows - - -def _relative_gain_column(value_col: str, comparator: str) -> str: - """Return the plan-specified relative-gain column name for *value_col*.""" - slug = _slug_comparator(comparator) - if value_col == "n_accepted": - return f"accepted_pct_vs_{slug}" - if value_col == "recovery_pct": - return f"recovery_pct_vs_{slug}" - if value_col == "observed_fdp": - return f"fdp_delta_vs_{slug}" - if "fdp" in value_col: - return f"{value_col}_delta_vs_{slug}" - return f"{value_col}_pct_vs_{slug}" - - -def _relative_gain_value( - value_col: str, primary_val: object, comparator_val: object -) -> float: - """Compute primary-vs-comparator gain for one metric cell.""" - if pd.isna(primary_val) or pd.isna(comparator_val): - return float("nan") - primary = float(cast("float | int | str", primary_val)) - comparator = float(cast("float | int | str", comparator_val)) - if value_col == "observed_fdp" or "fdp" in value_col: - return primary - comparator - if comparator == 0: - return float("nan") - return 100.0 * (primary - comparator) / comparator - - -def _fill_primary_relative_gains( - work: pd.DataFrame, - group: pd.DataFrame, - *, - primary_method: str, - comparators: Sequence[str], - value_cols: Sequence[str], - group_cols: Sequence[str], - group_keys: tuple[object, ...], -) -> None: - """Write relative-gain columns onto the primary-method row for one group.""" - primary_rows = group[group["method"] == primary_method] - if primary_rows.empty: - return - primary = primary_rows.iloc[0] - mask = pd.Series(True, index=work.index) - for col, val in zip(group_cols, group_keys): - mask &= work[col] == val - primary_idx = work.index[mask & (work["method"] == primary_method)] - if len(primary_idx) == 0: - return - idx = primary_idx[0] - - for comparator in comparators: - comp_rows = group[group["method"] == comparator] - if comp_rows.empty: - continue - comp = comp_rows.iloc[0] - for col in value_cols: - out_col = _relative_gain_column(col, comparator) - gain = _relative_gain_value(col, primary[col], comp[col]) - if np.isfinite(gain): - work.at[idx, out_col] = gain - - -def add_relative_gain_columns( - df: pd.DataFrame, - *, - primary_method: str, - comparators: Iterable[str], - value_cols: Sequence[str], - group_cols: Sequence[str] = ("dataset", "panel", "level", "q_value_threshold"), -) -> pd.DataFrame: - """Attach primary-vs-comparator relative columns onto a long-form metrics table. - - For ``n_accepted`` / ``recovery_pct``, writes - ``100 * (primary - comparator) / comparator``. - For ``observed_fdp``, writes the signed difference ``primary - comparator``. - Relative columns are filled only on primary-method rows. - """ - if df.empty: - return df.copy() - - work = df.copy() - comparator_list = list(comparators) - for col in value_cols: - for comparator in comparator_list: - work[_relative_gain_column(col, comparator)] = np.nan - - group_list = list(group_cols) - for keys, group in work.groupby(group_list, dropna=False, sort=False): - if not isinstance(keys, tuple): - keys = (keys,) - _fill_primary_relative_gains( - work, - group, - primary_method=primary_method, - comparators=comparator_list, - value_cols=value_cols, - group_cols=group_list, - group_keys=keys, - ) - return work - - -def merge_acceptance_and_error( - acceptance: pd.DataFrame, - error: pd.DataFrame, - *, - key_cols: Sequence[str] | None = None, -) -> pd.DataFrame: - """Join acceptance counts onto error rows for relative-gain construction.""" - if acceptance.empty or error.empty: - return error.copy() - keys = list(key_cols) if key_cols is not None else list(_KEY_COLS) - keys = [c for c in keys if c in acceptance.columns and c in error.columns] - cols = [ - c - for c in ("n_accepted", "n_correct", "recovery_pct") - if c in acceptance.columns - ] - return error.merge( - acceptance[keys + cols], - on=keys, - how="left", - ) - - -def finalise_error_gain_table( - acceptance: pd.DataFrame, - error: pd.DataFrame, - *, - primary_method: str, - comparators: Sequence[str], - key_cols: Sequence[str] | None = None, - group_cols: Sequence[str] | None = None, -) -> pd.DataFrame: - """Merge counts into error rows and add primary-vs-comparator relative columns.""" - merged = merge_acceptance_and_error(acceptance, error, key_cols=key_cols) - value_cols = [ - c for c in ("n_accepted", "recovery_pct", "observed_fdp") if c in merged.columns - ] - gain_groups = ( - tuple(group_cols) - if group_cols is not None - else ("dataset", "panel", "level", "q_value_threshold") - ) - with_gain = add_relative_gain_columns( - merged, - primary_method=primary_method, - comparators=comparators, - value_cols=value_cols, - group_cols=gain_groups, - ) - # Keep error-table identity columns first; drop helper count cols that duplicate - # the acceptance table except when used only for gain calculation. - drop_helpers = [c for c in ("n_accepted", "n_correct") if c in with_gain.columns] - return with_gain.drop(columns=drop_helpers, errors="ignore") - - -def write_summary_tables( - acceptance_df: pd.DataFrame, - error_df: pd.DataFrame, - output_dir: Path, - stem: str, -) -> tuple[Path, Path]: - """Write acceptance and error/gain CSVs under *output_dir*.""" - output_dir.mkdir(parents=True, exist_ok=True) - acceptance_path = output_dir / f"{stem}_acceptance.csv" - error_path = output_dir / f"{stem}_error_gain.csv" - acceptance_df.to_csv(acceptance_path, index=False) - error_df.to_csv(error_path, index=False) - logger.info("Wrote %s", acceptance_path) - logger.info("Wrote %s", error_path) - return acceptance_path, error_path - - -def summarise_holdout_results( - raw: pd.DataFrame, - *, - thresholds: Sequence[float] = SUMMARY_THRESHOLDS, - primary_method: str = "Winnow", - comparators: Sequence[str] = ("NovoBoard", "Glissade"), - panel: str = "score_mixture", - level: str = "peptide", - group_extra: Sequence[str] = (), -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Aggregate mixture-benchmark iterations into the two summary tables. - - Args: - raw: Per-iteration rows from ``external_peptide_holdout_results.csv``. - thresholds: Nominal FDR thresholds to retain. - primary_method: Method used for relative gain/loss columns. - comparators: Comparator method labels. - panel: Panel name written into the summary tables. - level: Identification level written into the summary tables. - group_extra: Extra columns to group by (e.g. ``pi0_target``). - - Returns: - ``(acceptance_df, error_gain_df)``. - """ - if raw.empty: - return pd.DataFrame(), pd.DataFrame() - - filtered = raw[raw["q_value_threshold"].isin(thresholds)].copy() - if filtered.empty: - return pd.DataFrame(), pd.DataFrame() - - extra = [c for c in group_extra if c in filtered.columns] - group_cols = ["dataset", *extra, "method", "q_value_threshold"] - gain_group_cols = ("dataset", *extra, "panel", "level", "q_value_threshold") - agg_kwargs: dict[str, tuple[str, str]] = { - "n_accepted": ("accepted_peptides", "mean"), - "n_correct": ("true_correct_peptides", "mean"), - "recovery_pct": ("correct_discovery_pct", "mean"), - "observed_fdp": ("observed_fdp", "mean"), - "n_accepted_std": ("accepted_peptides", "std"), - "observed_fdp_std": ("observed_fdp", "std"), - "recovery_pct_std": ("correct_discovery_pct", "std"), - } - if "mean_abs_q_dev_vs_db" in filtered.columns: - agg_kwargs["mean_abs_q_dev_vs_db"] = ("mean_abs_q_dev_vs_db", "mean") - agg = ( - filtered.groupby(group_cols, as_index=False) - .agg(**agg_kwargs) - .sort_values(group_cols) - .reset_index(drop=True) - ) - agg["panel"] = panel - agg["level"] = level - agg["fdp_excess"] = agg["observed_fdp"] - agg["q_value_threshold"] - if "mean_abs_q_dev_vs_db" not in agg.columns: - agg["mean_abs_q_dev_vs_db"] = np.nan - - id_cols = ["dataset", *extra, "panel", "level", "method", "q_value_threshold"] - acceptance = agg[ - id_cols - + [ - "n_accepted", - "n_correct", - "recovery_pct", - "n_accepted_std", - "recovery_pct_std", - ] - ].copy() - - error = agg[ - id_cols - + [ - "observed_fdp", - "fdp_excess", - "mean_abs_q_dev_vs_db", - "observed_fdp_std", - ] - ].copy() - - error_gain = finalise_error_gain_table( - acceptance.drop( - columns=["n_accepted_std", "recovery_pct_std"], errors="ignore" - ), - error, - primary_method=primary_method, - comparators=comparators, - key_cols=id_cols, - group_cols=gain_group_cols, - ) - return acceptance, error_gain diff --git a/scripts/feature_subsets.py b/scripts/feature_subsets.py deleted file mode 100644 index 34712b4d..00000000 --- a/scripts/feature_subsets.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Feature column sets for subset calibrator training and evaluation.""" - -from __future__ import annotations - -from typing import TypedDict - - -class FeatureSubsetSpec(TypedDict): - """Metadata and column list for one feature-subset experiment.""" - - description: str - from_parquet: bool - columns: list[str] - - -# Full feature matrix columns (``confidence`` + features + ``correct`` label). -FULL_FEATURE_COLUMNS: list[str] = [ - "confidence", - "mass_error_ppm", - "ion_matches", - "ion_match_intensity", - "complementary_ion_count", - "max_ion_gap", - "spectral_angle", - "xcorr", - "irt_error", - "margin", - "median_margin", - "entropy", - "z-score", - "edit_distance", - "min_token_probability", - "std_token_probability", -] - -_NO_XCORR_SPECTRAL = {"spectral_angle", "xcorr"} -_NO_FRAGMENT_SIMILARITY = _NO_XCORR_SPECTRAL | { - "complementary_ion_count", - "max_ion_gap", - "edit_distance", -} - - -def _columns_excluding(*, drop: set[str]) -> list[str]: - return [c for c in FULL_FEATURE_COLUMNS if c not in drop] - - -FEATURE_SUBSETS: dict[str, FeatureSubsetSpec] = { - "no_xcorr_spectral": { - "description": "Exclude spectral_angle and xcorr only.", - "from_parquet": True, - "columns": _columns_excluding(drop=_NO_XCORR_SPECTRAL), - }, - "no_fragment_similarity": { - "description": ( - "Exclude spectral_angle, xcorr, complementary_ion_count, " - "max_ion_gap, and edit_distance." - ), - "from_parquet": True, - "columns": _columns_excluding(drop=_NO_FRAGMENT_SIMILARITY), - }, - "mass_error_da_no_similarity": { - "description": ( - "Exclude mass_error_ppm and fragment-similarity features; " - "use mass_error_da (Daltons) instead. Requires full winnow train " - "(recomputes features from raw spectra)." - ), - "from_parquet": False, - "columns": list( - _columns_excluding(drop=_NO_FRAGMENT_SIMILARITY | {"mass_error_ppm"}) - ) - + ["mass_error_da"], - }, -} diff --git a/scripts/plot_ablation_summary.py b/scripts/plot_ablation_summary.py deleted file mode 100644 index 862d1829..00000000 --- a/scripts/plot_ablation_summary.py +++ /dev/null @@ -1,641 +0,0 @@ -#!/usr/bin/env python3 -"""Bar charts of ablation calibration metrics from ``ablation_summary.csv``. - -Designed for publication main text: tail ECE at FDR operating points (and optionally -Brier) per feature-group config, with a reference line at the full ``All features`` model. -""" - -from __future__ import annotations - -import json -import logging -import sys -from pathlib import Path -from typing import Annotated, Literal - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import seaborn as sns -import typer - -_REPO_ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(_REPO_ROOT)) - -from scripts.plot_eval_results import ( # noqa: E402 - _display_name, - _fit_database_grounded_fdr, - _save_fig, - _style_ax, -) -from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 - -# Paul Tol qualitative palette (colour-blind safe) — canonical ablation colours. -_ABLATION_PALETTE = [ - "#4477AA", - "#EE6677", - "#228833", - "#CCBB44", - "#66CCEE", - "#AA3377", - "#EE7733", - "#0077BB", - "#33BBEE", - "#CC3311", -] - -# Ablation summary keys → ``plot_eval_results.DATASET_DISPLAY_NAMES`` keys. -_ABLATION_DATASET_KEYS: dict[str, str] = { - "Arabidopsis": "01747_C01_P018218_S00_I00_N03_R1", - "Astral": "astral", - "HCT116": "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46", -} - -logger = logging.getLogger(__name__) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -ABLATION_CONFIG_ORDER: list[str] = [ - "Confidence only", - "Confidence + mass error", - "Confidence + iRT error", - "Confidence + token-level", - "Confidence + beam search", - "Confidence + fragment matching", - "All features", -] - - -def assign_ablation_colors(config_names: list[str]) -> dict[str, str]: - """Assign a unique colour per ablation config (no palette wrap).""" - if len(config_names) > len(_ABLATION_PALETTE): - raise ValueError( - f"Need {len(config_names)} ablation colours but only " - f"{len(_ABLATION_PALETTE)} defined." - ) - return {name: _ABLATION_PALETTE[i] for i, name in enumerate(config_names)} - - -def ordered_ablation_configs(present: set[str]) -> list[str]: - """Canonical config order for ablation plots and colour assignment.""" - ordered = [c for c in ABLATION_CONFIG_ORDER if c in present] - extra = sorted(present - set(ordered)) - return ordered + extra - - -_CONFIG_SHORT_LABELS: dict[str, str] = { - "Confidence only": "Confidence", - "Confidence + mass error": "+ Mass", - "Confidence + iRT error": "+ iRT", - "Confidence + token-level": "+ Token", - "Confidence + beam search": "+ Beam", - "Confidence + fragment matching": "+ Fragment", - "All features": "All features", -} - -MetricName = Literal[ - "tail_ECE@5%FDR", - "tail_ECE@10%FDR", - "ECE", - "Brier", - "PR_AUC", - "fdr_bias@5%FDR", - "fdr_bias@10%FDR", - "q_dev@5%FDR", - "q_dev@10%FDR", -] - -FDR_TAIL_THRESHOLDS: tuple[float, ...] = (0.05, 0.10) -TAIL_ECE_COLUMN_BY_THRESHOLD: dict[float, str] = { - 0.05: "tail_ECE@5%FDR", - 0.10: "tail_ECE@10%FDR", -} -Q_DEV_COLUMN_BY_THRESHOLD: dict[float, str] = { - 0.05: "q_dev@5%FDR", - 0.10: "q_dev@10%FDR", -} -FDR_BIAS_COLUMN_BY_THRESHOLD: dict[float, str] = { - 0.05: "fdr_bias@5%FDR", - 0.10: "fdr_bias@10%FDR", -} - -_DEFAULT_SUMMARY = ( - Path.home() / "Documents/winnow/new_eval_sets_plots/ablations/ablation_summary.csv" -) - - -def load_ablation_summary(path: Path) -> pd.DataFrame: - """Load ``ablation_summary.csv`` or ``.json``.""" - if not path.is_file(): - raise FileNotFoundError(path) - if path.suffix == ".json": - with open(path) as f: - return pd.DataFrame(json.load(f)) - return pd.read_csv(path) - - -def compute_ece( - pred: np.ndarray, - labels: np.ndarray, - n_bins: int = 10, -) -> float: - """Expected calibration error.""" - bins = np.linspace(0.0, 1.0, n_bins + 1) - bin_indices = np.digitize(pred, bins) - 1 - bin_indices = np.clip(bin_indices, 0, n_bins - 1) - ece = 0.0 - for b in range(n_bins): - mask = bin_indices == b - if mask.sum() == 0: - continue - avg_conf = pred[mask].mean() - avg_acc = labels[mask].mean() - ece += mask.sum() / len(pred) * abs(avg_conf - avg_acc) - return float(ece) - - -def compute_tail_ece_at_fdr( - pred: np.ndarray, - labels: np.ndarray, - fdr_threshold: float, - *, - fdr_ctrl: NonParametricFDRControl | None = None, - n_bins: int = 10, -) -> float: - """ECE among PSMs accepted at a non-parametric FDR threshold.""" - if len(pred) == 0: - return float("nan") - - if fdr_ctrl is None: - fdr_ctrl = NonParametricFDRControl() - fdr_ctrl.fit(dataset=pd.Series(pred, name="score")) - - cutoff = fdr_ctrl.get_confidence_cutoff(threshold=fdr_threshold) - if np.isnan(cutoff): - return float("nan") - - mask = pred >= cutoff - if not mask.any(): - return float("nan") - - return compute_ece(pred[mask], labels[mask], n_bins=n_bins) - - -def compute_tail_ece_at_fdr_thresholds( - df: pd.DataFrame, - *, - confidence_col: str = "calibrated_confidence", - label_col: str = "correct", - fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, -) -> dict[float, float]: - """Tail ECE at each non-parametric FDR operating point.""" - work = df[[confidence_col, label_col]].dropna() - if work.empty: - return {threshold: float("nan") for threshold in fdr_thresholds} - - pred = work[confidence_col].to_numpy(dtype=float) - labels = work[label_col].to_numpy(dtype=float) - - fdr_ctrl = NonParametricFDRControl() - fdr_ctrl.fit(dataset=work[confidence_col]) - - return { - threshold: compute_tail_ece_at_fdr( - pred, - labels, - threshold, - fdr_ctrl=fdr_ctrl, - ) - for threshold in fdr_thresholds - } - - -def compute_fdr_bias_at_fdr_thresholds( - df: pd.DataFrame, - *, - confidence_col: str = "calibrated_confidence", - label_col: str = "correct", - fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, -) -> dict[float, float]: - """Signed FDR bias at each NP-FDR cutoff; equal to empirical sTECE.""" - work = df[[confidence_col, label_col]].dropna() - if work.empty: - return {threshold: float("nan") for threshold in fdr_thresholds} - - scores = work[confidence_col].to_numpy(dtype=float) - labels = work[label_col].to_numpy(dtype=float) - - fdr_ctrl = NonParametricFDRControl() - fdr_ctrl.fit(dataset=work[confidence_col]) - - results: dict[float, float] = {} - for threshold in fdr_thresholds: - cutoff = fdr_ctrl.get_confidence_cutoff(threshold=threshold) - if np.isnan(cutoff): - results[threshold] = float("nan") - continue - mask = scores >= cutoff - if not mask.any(): - results[threshold] = float("nan") - continue - - # E[1-S | S>=tau] - E[1-Y | S>=tau] = E[Y-S | S>=tau]. - results[threshold] = float(np.mean(labels[mask] - scores[mask])) - return results - - -def compute_pr_auc( - df: pd.DataFrame, - confidence_col: str = "calibrated_confidence", - label_col: str = "correct", -) -> float: - """Area under the ablation PR curve (matches ``run_feature_ablations`` plots).""" - work = df[[confidence_col, label_col]].dropna() - if work.empty: - return float("nan") - - sorted_data = work.sort_values(by=confidence_col, ascending=False) - cum_correct = np.cumsum(sorted_data[label_col].values) - precision = cum_correct / np.arange(1, len(sorted_data) + 1) - total_correct = cum_correct[-1] if len(cum_correct) else 0 - if total_correct <= 0 or len(precision) < 2: - return 0.0 - - recall = cum_correct / total_correct - from sklearn.metrics import auc - - return float(auc(recall, precision)) - - -def _vectorized_psm_fdr( - scores: np.ndarray, - ctrl: NonParametricFDRControl, -) -> np.ndarray: - """Map confidence scores to PSM FDR using a fitted controller.""" - conf = np.asarray(ctrl._confidence_scores, dtype=float) - fdr = np.asarray(ctrl._fdr_values, dtype=float) - scores = np.asarray(scores, dtype=float) - idx = np.searchsorted(-conf, -scores, side="left") - idx = np.clip(idx, 0, max(len(fdr) - 1, 0)) - if len(fdr) == 0: - return np.ones_like(scores) - - out = fdr[idx] - below = (idx == len(conf)) & (scores < conf[-1]) - above = (idx == 0) & (scores > conf[0]) - out[below] = 1.0 - out[above] = fdr[0] - return out - - -def _vectorized_psm_q_values( - scores: np.ndarray, - ctrl: NonParametricFDRControl, -) -> np.ndarray: - """Assign PSM q-values without per-row ``compute_fdr`` calls.""" - row_fdr = _vectorized_psm_fdr(scores, ctrl) - order = np.argsort(-scores) - sorted_fdr = row_fdr[order] - q_sorted = np.empty_like(sorted_fdr) - fdr_min = np.inf - for i in range(len(sorted_fdr) - 1, -1, -1): - current = sorted_fdr[i] - if current > fdr_min: - q_sorted[i] = fdr_min - else: - q_sorted[i] = current - fdr_min = current - q_values = np.empty_like(q_sorted) - q_values[order] = q_sorted - return q_values - - -def compute_q_value_deviations( - df: pd.DataFrame, - *, - confidence_col: str = "calibrated_confidence", - label_col: str = "correct", - fdr_thresholds: tuple[float, ...] = FDR_TAIL_THRESHOLDS, -) -> dict[float, float]: - """Mean absolute q-value deviation among NP-accepted PSMs at each FDR level.""" - work = df[[confidence_col, label_col]].dropna().copy() - if work.empty or label_col not in work.columns: - return {threshold: float("nan") for threshold in fdr_thresholds} - - np_fdr = NonParametricFDRControl() - np_fdr.fit(dataset=work[confidence_col]) - - dbg_ctrl = _fit_database_grounded_fdr( - work, - confidence_col=confidence_col, - correct_col=label_col, - drop=0 if len(work) <= 10 else 10, - ) - - scores = work[confidence_col].to_numpy(dtype=float) - est_q = _vectorized_psm_q_values(scores, np_fdr) - true_q = _vectorized_psm_q_values(scores, dbg_ctrl) - deviations = np.abs(est_q - true_q) - - results: dict[float, float] = {} - for threshold in fdr_thresholds: - mask = est_q <= threshold - if not mask.any(): - results[threshold] = float("nan") - else: - results[threshold] = float(np.mean(deviations[mask])) - return results - - -def metrics_from_eval_parquet(path: Path) -> dict[str, float | str]: - """Compute tail ECE, PR-AUC, and q-value metrics from one eval-results Parquet.""" - df = pd.read_parquet(path) - config_name = str(df["config_name"].iloc[0]) - dataset_name = str(df["dataset_name"].iloc[0]) - meta = df.drop(columns=["config_name", "dataset_name"], errors="ignore") - - tail_ece = compute_tail_ece_at_fdr_thresholds(meta) - fdr_bias = compute_fdr_bias_at_fdr_thresholds(meta) - pr_auc = compute_pr_auc(meta) - q_dev = compute_q_value_deviations(meta) - - return { - "config": config_name, - "dataset": dataset_name, - TAIL_ECE_COLUMN_BY_THRESHOLD[0.05]: round(tail_ece[0.05], 5), - TAIL_ECE_COLUMN_BY_THRESHOLD[0.10]: round(tail_ece[0.10], 5), - FDR_BIAS_COLUMN_BY_THRESHOLD[0.05]: round(fdr_bias[0.05], 5), - FDR_BIAS_COLUMN_BY_THRESHOLD[0.10]: round(fdr_bias[0.10], 5), - "PR_AUC": round(pr_auc, 5), - Q_DEV_COLUMN_BY_THRESHOLD[0.05]: round(q_dev[0.05], 5), - Q_DEV_COLUMN_BY_THRESHOLD[0.10]: round(q_dev[0.10], 5), - } - - -def enrich_summary_from_eval_results( - summary: pd.DataFrame, - eval_results_dir: Path, - *, - datasets: list[str] | None = None, -) -> pd.DataFrame: - """Add PR-AUC and q-value deviation columns using saved eval Parquets.""" - if not eval_results_dir.is_dir(): - raise FileNotFoundError(eval_results_dir) - - metric_rows: list[dict[str, float | str]] = [] - for path in sorted(eval_results_dir.glob("*.parquet")): - dataset_name = path.name.split("_", 1)[0] - if datasets is not None and dataset_name not in datasets: - continue - metric_rows.append(metrics_from_eval_parquet(path)) - - if not metric_rows: - raise FileNotFoundError( - f"No eval Parquets found under {eval_results_dir}" - + (f" for datasets {datasets!r}" if datasets else "") - ) - - metrics_df = pd.DataFrame(metric_rows) - merge_cols = ["config", "dataset"] - extra_cols = [ - *TAIL_ECE_COLUMN_BY_THRESHOLD.values(), - *FDR_BIAS_COLUMN_BY_THRESHOLD.values(), - "PR_AUC", - *Q_DEV_COLUMN_BY_THRESHOLD.values(), - "tail_ECE", - ] - summary = summary.drop(columns=extra_cols, errors="ignore") - return summary.merge(metrics_df, on=merge_cols, how="left") - - -def _ablation_dataset_display(dataset: str) -> str: - """Publication label via ``plot_eval_results._display_name``.""" - return _display_name(_ABLATION_DATASET_KEYS.get(dataset, dataset)) - - -def _wrap_title_before_dataset(title: str, *, max_line: int = 52) -> str: - """Break before ``on `` when the title would be too wide.""" - marker = " on " - if marker not in title or len(title) <= max_line: - return title - split = title.index(marker) - return f"{title[:split]}\n{title[split + 1 :]}" - - -def _metric_axis_label(metric: MetricName) -> str: - if metric == "tail_ECE@5%FDR": - return "Tail ECE at 5% FDR" - if metric == "tail_ECE@10%FDR": - return "Tail ECE at 10% FDR" - if metric == "ECE": - return "ECE" - if metric == "Brier": - return "Brier score" - if metric == "PR_AUC": - return "PR-AUC" - if metric == "fdr_bias@5%FDR": - return "FDR bias (= sTECE) at 5% FDR" - if metric == "fdr_bias@10%FDR": - return "FDR bias (= sTECE) at 10% FDR" - if metric == "q_dev@5%FDR": - return "Mean |q-value deviation| at 5% FDR" - return "Mean |q-value deviation| at 10% FDR" - - -def _metric_plot_title(metric: MetricName, dataset_display: str) -> str: - """Publication title: full sentence, ECE capitalised.""" - if metric == "tail_ECE@5%FDR": - title = ( - f"Tail expected calibration error among PSMs accepted at 5% FDR " - f"on {dataset_display}" - ) - elif metric == "tail_ECE@10%FDR": - title = ( - f"Tail expected calibration error among PSMs accepted at 10% FDR " - f"on {dataset_display}" - ) - elif metric == "ECE": - title = f"Expected calibration error (ECE) on {dataset_display}" - elif metric == "Brier": - title = f"Brier score on {dataset_display}." - elif metric == "PR_AUC": - title = f"Precision-recall AUC on {dataset_display}" - elif metric == "fdr_bias@5%FDR": - title = ( - f"FDR bias, equal to signed tail calibration error, " - f"at 5% FDR on {dataset_display}" - ) - elif metric == "fdr_bias@10%FDR": - title = ( - f"FDR bias, equal to signed tail calibration error, " - f"at 10% FDR on {dataset_display}" - ) - elif metric == "q_dev@5%FDR": - title = ( - f"Non-parametric q-value deviation from database-grounded q-values " - f"at 5% FDR on {dataset_display}" - ) - else: - title = ( - f"Non-parametric q-value deviation from database-grounded q-values " - f"at 10% FDR on {dataset_display}" - ) - return _wrap_title_before_dataset(title) - - -def plot_ablation_calibration_bars( - summary: pd.DataFrame, - dataset: str, - *, - metric: MetricName = "tail_ECE@5%FDR", - output_path: Path, - figsize: tuple[float, float] = (7.5, 4), -) -> pd.DataFrame: - """Bar chart of *metric* for one dataset; returns the plotted slice.""" - ds = summary.loc[summary["dataset"] == dataset].copy() - if ds.empty: - available = sorted(summary["dataset"].unique()) - raise ValueError(f"No rows for dataset {dataset!r}. Available: {available}") - - configs = ordered_ablation_configs(set(ds["config"])) - ds = ds.set_index("config").loc[configs].reset_index() - if metric not in ds.columns: - raise ValueError(f"Metric {metric!r} not in summary columns: {ds.columns}") - - values = ds[metric].to_numpy(dtype=float) - all_features_value = float(ds.loc[ds["config"] == "All features", metric].iloc[0]) - - colors = assign_ablation_colors(configs) - short_labels = [_CONFIG_SHORT_LABELS.get(c, c) for c in configs] - - fig, ax = plt.subplots(figsize=figsize) - x = np.arange(len(configs)) - bar_colors = [colors[c] for c in configs] - ax.bar(x, values, color=bar_colors, edgecolor="black", linewidth=0.6, zorder=2) - ax.axhline( - all_features_value, - color="#333333", - linestyle="--", - linewidth=1.2, - zorder=1, - label="All features", - ) - - display = _ablation_dataset_display(dataset) - ax.set_ylabel(_metric_axis_label(metric)) - ax.set_xlabel("Calibrator feature groups") - ax.set_title(_metric_plot_title(metric, display)) - ax.set_xticks(x) - ax.set_xticklabels(short_labels, rotation=35, ha="right") - ax.legend(loc="upper right") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_path) - logger.info("Wrote %s.png and %s.pdf", output_path, output_path) - return ds[["config", metric]] - - -@app.command() -def main( - summary: Annotated[ - Path, - typer.Option("--summary", help="ablation_summary.csv or .json"), - ] = _DEFAULT_SUMMARY, - dataset: Annotated[ - str, - typer.Option("--dataset", help="Dataset key in the summary table"), - ] = "Arabidopsis", - metric: Annotated[ - MetricName, - typer.Option("--metric", help="Calibration metric to plot"), - ] = "tail_ECE@5%FDR", - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for figure outputs"), - ] = _DEFAULT_SUMMARY.parent / "plots", - eval_results_dir: Annotated[ - Path | None, - typer.Option( - "--eval-results-dir", - help="Optional eval_results/ directory to enrich summary before plotting", - ), - ] = None, -) -> None: - """Plot ablation calibration bars for one dataset.""" - logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") - sns.set_theme(style="white", context="paper", font_scale=1.5) - - summary_df = load_ablation_summary(summary) - if eval_results_dir is not None: - summary_df = enrich_summary_from_eval_results( - summary_df, - eval_results_dir, - datasets=[dataset], - ) - output_dir.mkdir(parents=True, exist_ok=True) - slug = dataset.lower().replace(" ", "_") - metric_slug = metric.lower().replace("%", "pct").replace("@", "_at_") - out_base = output_dir / f"ablation_{metric_slug}_{slug}" - table = plot_ablation_calibration_bars( - summary_df, dataset, metric=metric, output_path=out_base - ) - print(table.to_string(index=False)) - - -@app.command("recompute-summary") -def recompute_summary( - eval_results_dir: Annotated[ - Path, - typer.Option("--eval-results-dir", help="Directory of eval_results Parquets"), - ], - summary: Annotated[ - Path | None, - typer.Option( - "--summary", - help="Existing ablation_summary.csv to merge with (optional)", - ), - ] = None, - datasets: Annotated[ - list[str] | None, - typer.Option( - "--datasets", - help="Restrict to these dataset keys (repeatable)", - ), - ] = None, - output: Annotated[ - Path, - typer.Option("--output", help="Output CSV path"), - ] = _DEFAULT_SUMMARY, -) -> None: - """Recompute PR-AUC and q-value deviation columns from eval Parquets.""" - logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") - - if summary is not None: - base = load_ablation_summary(summary) - if datasets is not None: - base = base.loc[base["dataset"].isin(datasets)].copy() - enriched = enrich_summary_from_eval_results( - base, - eval_results_dir, - datasets=datasets, - ) - else: - rows = [] - for path in sorted(eval_results_dir.glob("*.parquet")): - dataset_name = path.name.split("_", 1)[0] - if datasets is not None and dataset_name not in datasets: - continue - rows.append(metrics_from_eval_parquet(path)) - if not rows: - raise typer.BadParameter(f"No eval Parquets found under {eval_results_dir}") - enriched = pd.DataFrame(rows) - enriched = enriched.sort_values(["dataset", "config"]).reset_index(drop=True) - - output.parent.mkdir(parents=True, exist_ok=True) - enriched.to_csv(output, index=False) - logger.info("Wrote %s", output) - print(enriched.to_string(index=False)) - - -if __name__ == "__main__": - app() diff --git a/scripts/plot_acfm_minus_lcfm_fdr.py b/scripts/plot_acfm_minus_lcfm_fdr.py deleted file mode 100644 index 9dad0ebd..00000000 --- a/scripts/plot_acfm_minus_lcfm_fdr.py +++ /dev/null @@ -1,470 +0,0 @@ -"""Refit FDR on acfm predictions restricted to spectra not present in lcfm. - -For each external project, spectra are matched on ``spectrum_id``. The acfm -(unlabelled) set is filtered to ``spectrum_id`` values absent from the paired -lcfm (labelled) predictions, FDR is re-estimated on ``calibrated_confidence`` -for that subset only, and evaluation plots are written. -""" - -from __future__ import annotations - -import logging -import sys -from pathlib import Path -from typing import Annotated - -import pandas as pd -import typer -from rich.logging import RichHandler - -_REPO_ROOT = Path(__file__).resolve().parent.parent -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from scripts.plot_eval_results import ( # noqa: E402 - _compute_diagnostics, - _display_name, - _fit_database_grounded_fdr, - generate_all_plots, -) -from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 - -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -logger.propagate = False -if not logger.handlers: - logger.addHandler(RichHandler()) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - - -def _safe_basename(project: str) -> str: - """Flat basename for outputs; project keys may contain path separators.""" - return project.replace("/", "_") - - -# Per-run keys under new_eval_sets_results/{lcfm,acfm}/// -_PXD_RUN_PARENTS: tuple[str, ...] = ("PXD004452", "PXD006939", "PXD013868") - - -def _add_preds_candidate( - candidates: list[Path], - seen: set[Path], - root: Path, - *relative: str, - fname: str, -) -> None: - path = root.joinpath(*relative, fname) - if path not in seen: - seen.add(path) - candidates.append(path) - - -def _preds_csv_candidates(root: Path, project: str, *, role: str) -> list[Path]: - """Paths to try for ``preds_and_fdr_metrics.csv`` under *root*. - - Supports flat layouts (``{root}/{run}/``), S3-style nesting - (``{root}/PXD006939/{run}/``), and explicit ``PXD006939/run`` project keys. - """ - if role not in ("labelled", "unlabelled"): - raise ValueError(f"Unknown role {role!r}") - - role_suffix = "_labelled" if role == "labelled" else "_unlabelled" - fname = "preds_and_fdr_metrics.csv" - seen: set[Path] = set() - candidates: list[Path] = [] - - def add(*relative: str) -> None: - _add_preds_candidate(candidates, seen, root, *relative, fname=fname) - - add(project) - base = project.split("/")[-1] - add(f"{base}{role_suffix}") - if "/" in project: - add(*project.split("/")) - return candidates - - # S3 and download-new-eval-results: {root}/PXD*/{run}/; legacy flat {root}/{run}/. - pxd_seen: set[str] = set() - for pxd in _PXD_RUN_PARENTS: - add(pxd, project) - add(pxd, f"{project}{role_suffix}") - pxd_seen.add(pxd) - if root.is_dir(): - for child in sorted(root.iterdir()): - if child.is_dir() and child.name.startswith("PXD"): - if child.name not in pxd_seen: - add(child.name, project) - add(child.name, f"{project}{role_suffix}") - - return candidates - - -def _collect_alt_roots(root: Path, alt_roots: list[Path] | None) -> list[Path]: - """Build ordered list of roots to search, deduplicating by resolve().""" - roots_to_try: list[Path] = [root] - if not alt_roots: - return roots_to_try - for alt in alt_roots: - if alt.resolve() != root.resolve() and alt not in roots_to_try: - roots_to_try.append(alt) - return roots_to_try - - -def _resolve_preds_csv( - root: Path, - project: str, - *, - role: str, - alt_roots: list[Path] | None = None, -) -> Path: - """Resolve ``preds_and_fdr_metrics.csv`` for lcfm (labelled) or acfm (unlabelled).""" - tried: list[Path] = [] - for base in _collect_alt_roots(root, alt_roots): - for path in _preds_csv_candidates(base, project, role=role): - tried.append(path) - if not path.is_file(): - continue - if base.resolve() != root.resolve(): - logger.info( - "Using %s predictions at %s (not under %s)", - role, - path, - root, - ) - return path - - hint = "" - if root.is_dir(): - children = sorted(p.name for p in root.iterdir())[:12] - hint = f" Children of {root}: {children}" - raise FileNotFoundError( - f"Missing {role} predictions for {project!r} under {root} " - f"(tried: {', '.join(str(p) for p in tried)}){hint}" - ) - - -def _infer_predictions_root( - labelled_dir: Path | None, - unlabelled_dir: Path | None, -) -> Path: - """Infer a common parent when only per-tree dirs are passed.""" - if labelled_dir is not None and unlabelled_dir is not None: - if labelled_dir.parent == unlabelled_dir.parent: - return labelled_dir.parent - if labelled_dir is not None: - return labelled_dir.parent - if unlabelled_dir is not None: - return unlabelled_dir.parent - raise typer.BadParameter( - "Provide --predictions-root or at least one of --labelled-dir / --unlabelled-dir." - ) - - -def _resolve_explicit_tree_root( - predictions_root: Path, - explicit: Path, - *, - sub: str, -) -> tuple[Path, list[Path]]: - """Resolve an explicit tree path that may be relative to *predictions_root*.""" - alt_roots: list[Path] = [] - if len(explicit.parts) == 1: - under_predictions = predictions_root / explicit - if under_predictions.is_dir(): - root = under_predictions - if explicit.is_dir() and explicit.resolve() != root.resolve(): - alt_roots.append(explicit) - nested = predictions_root / sub - if nested.is_dir() and nested.resolve() != root.resolve(): - alt_roots.append(nested) - return root, alt_roots - return explicit, alt_roots - - -def _resolve_tree_root( - predictions_root: Path | None, - explicit: Path | None, - *, - role: str, -) -> tuple[Path, list[Path]]: - """Pick labelled (lcfm) or unlabelled (acfm) root; return alternates to try.""" - sub = "lcfm" if role == "labelled" else "acfm" - legacy = "labelled" if role == "labelled" else "unlabelled" - alt_roots: list[Path] = [] - - if explicit is not None: - if predictions_root is not None: - return _resolve_explicit_tree_root(predictions_root, explicit, sub=sub) - return explicit, alt_roots - - if predictions_root is None: - raise typer.BadParameter( - f"Missing --predictions-root and --{legacy}-dir for {role} tree." - ) - - for name in (sub, legacy): - candidate = predictions_root / name - if candidate.is_dir(): - return candidate, alt_roots - return predictions_root, alt_roots - - -def _lcfm_spectrum_ids( - labelled_root: Path, - project: str, - *, - labelled_alt_roots: list[Path] | None = None, -) -> set[str]: - labelled_path = _resolve_preds_csv( - labelled_root, - project, - role="labelled", - alt_roots=labelled_alt_roots, - ) - ids = pd.read_csv(labelled_path, usecols=["spectrum_id"])["spectrum_id"] - return set(ids.astype(str)) - - -def _load_acfm_unlabelled( - unlabelled_root: Path, - project: str, - *, - unlabelled_alt_roots: list[Path] | None = None, -) -> pd.DataFrame: - """Load acfm predict outputs with metadata merged (same as plot_eval_results).""" - preds_path = _resolve_preds_csv( - unlabelled_root, - project, - role="unlabelled", - alt_roots=unlabelled_alt_roots, - ) - folder = preds_path.parent - preds_df = pd.read_csv(preds_path) - meta_path = folder / "metadata.csv" - if meta_path.is_file(): - meta_df = pd.read_csv(meta_path) - overlap = [ - c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" - ] - if overlap: - meta_df = meta_df.drop(columns=overlap) - df = preds_df.merge(meta_df, on="spectrum_id", how="left") - else: - df = preds_df - - if "proteome_hit" not in df.columns: - raise ValueError( - f"Expected 'proteome_hit' column for unlabelled acfm in {preds_path}" - ) - df["correct"] = df["proteome_hit"].astype(float) - required = ["confidence", "calibrated_confidence", "correct"] - missing = [c for c in required if c not in df.columns] - if missing: - raise ValueError(f"Missing columns {missing} in {preds_path}") - return df - - -def filter_acfm_minus_lcfm(acfm_df: pd.DataFrame, lcfm_ids: set[str]) -> pd.DataFrame: - """Keep acfm rows whose ``spectrum_id`` is not in the lcfm set.""" - mask = ~acfm_df["spectrum_id"].astype(str).isin(lcfm_ids) - return acfm_df.loc[mask].copy() - - -def refit_fdr_on_confidence( - df: pd.DataFrame, - confidence_col: str = "calibrated_confidence", -) -> pd.DataFrame: - """Fit non-parametric FDR on *df* and attach PSM FDR / q-value / PEP columns.""" - out = df.copy() - for col in ("psm_fdr", "psm_q_value", "psm_pep"): - if col in out.columns: - out = out.drop(columns=[col]) - fdr_ctrl = NonParametricFDRControl() - fdr_ctrl.fit(dataset=out[confidence_col]) - out = fdr_ctrl.add_psm_fdr(out, confidence_col=confidence_col) - out = fdr_ctrl.add_psm_q_value(out, confidence_col=confidence_col) - out = fdr_ctrl.add_psm_pep(out, confidence_col=confidence_col) - return out - - -def process_project( - labelled_root: Path, - unlabelled_root: Path, - project: str, - output_dir: Path, - *, - labelled_alt_roots: list[Path] | None = None, - unlabelled_alt_roots: list[Path] | None = None, -) -> dict[str, int]: - """Filter acfm less lcfm, refit FDR, plot, and write tables for one project.""" - lcfm_ids = _lcfm_spectrum_ids( - labelled_root, project, labelled_alt_roots=labelled_alt_roots - ) - acfm_df = _load_acfm_unlabelled( - unlabelled_root, project, unlabelled_alt_roots=unlabelled_alt_roots - ) - subset_df = filter_acfm_minus_lcfm(acfm_df, lcfm_ids) - - counts = { - "n_lcfm_spectrum_ids": len(lcfm_ids), - "n_acfm": len(acfm_df), - "n_acfm_minus_lcfm": len(subset_df), - } - if counts["n_acfm_minus_lcfm"] == 0: - raise ValueError( - f"{project}: no acfm spectra remain after excluding lcfm spectrum_id values" - ) - - logger.info( - "%s: acfm=%s, lcfm ids=%s, acfm\\lcfm=%s", - project, - f"{counts['n_acfm']:,}", - f"{counts['n_lcfm_spectrum_ids']:,}", - f"{counts['n_acfm_minus_lcfm']:,}", - ) - - subset_df = refit_fdr_on_confidence(subset_df) - - project_dir = output_dir / project - project_dir.mkdir(parents=True, exist_ok=True) - safe = _safe_basename(project) - subset_df.to_csv(project_dir / "preds_and_fdr_metrics.csv", index=False) - - true_fdr_ctrl = _fit_database_grounded_fdr(subset_df) - db_fdr = true_fdr_ctrl.add_psm_fdr( - subset_df[["calibrated_confidence"]].copy(), - confidence_col="calibrated_confidence", - ) - subset_df["db_grounded_psm_fdr"] = db_fdr["psm_fdr"] - db_qval = true_fdr_ctrl.add_psm_q_value( - subset_df[["calibrated_confidence"]].copy(), - confidence_col="calibrated_confidence", - ) - subset_df["db_grounded_psm_q_value"] = db_qval["psm_q_value"] - - summary_cols = [ - c - for c in [ - "spectrum_id", - "prediction", - "confidence", - "calibrated_confidence", - "correct", - "psm_fdr", - "psm_q_value", - "db_grounded_psm_fdr", - "db_grounded_psm_q_value", - "proteome_hit", - ] - if c in subset_df.columns - ] - subset_df[summary_cols].to_csv(project_dir / f"{safe}_summary.csv", index=False) - - diag = _compute_diagnostics(subset_df, "unlabelled") - diag.to_csv(project_dir / f"{safe}_diagnostics.csv", index=False) - - pd.DataFrame([counts]).to_csv(project_dir / f"{safe}_counts.csv", index=False) - - generate_all_plots(subset_df, safe, "unlabelled", project_dir) - return counts - - -@app.command() -def main( - projects: Annotated[ - str, - typer.Option( - "--projects", - help="Space- or comma-separated project keys (e.g. 'PXD009935 PXD014877').", - ), - ], - output_dir: Annotated[ - Path, - typer.Option( - "--output-dir", - help="Directory for per-project plots and refitted prediction tables.", - ), - ], - predictions_root: Annotated[ - Path | None, - typer.Option( - "--predictions-root", - help=( - "Parent of labelled/unlabelled (or lcfm/acfm) trees. Optional when " - "both --labelled-dir and --unlabelled-dir are set (parent is inferred)." - ), - ), - ] = None, - labelled_dir: Annotated[ - Path | None, - typer.Option( - "--labelled-dir", - help=( - "Root with per-project lcfm folders ({project}/ or {project}_labelled/). " - "Use e.g. new_eval_sets_results/lcfm when mirroring S3. " - "If omitted, uses --predictions-root/lcfm when present." - ), - ), - ] = None, - unlabelled_dir: Annotated[ - Path | None, - typer.Option( - "--unlabelled-dir", - help=( - "Root with per-project acfm folders ({project}/ or {project}_unlabelled/). " - "Use e.g. new_eval_sets_results/acfm when mirroring S3. " - "If omitted, uses --predictions-root/acfm when present." - ), - ), - ] = None, -) -> None: - """Refit FDR on acfm less lcfm spectra and generate evaluation plots.""" - project_list = [p.strip() for p in projects.replace(",", " ").split() if p.strip()] - if not project_list: - raise typer.BadParameter("No projects specified.") - - preds_root = predictions_root - if preds_root is None: - preds_root = _infer_predictions_root(labelled_dir, unlabelled_dir) - - labelled_root, labelled_alt = _resolve_tree_root( - preds_root, labelled_dir, role="labelled" - ) - unlabelled_root, unlabelled_alt = _resolve_tree_root( - preds_root, unlabelled_dir, role="unlabelled" - ) - logger.info("Labelled (lcfm) root: %s", labelled_root.resolve()) - logger.info("Unlabelled (acfm) root: %s", unlabelled_root.resolve()) - - output_dir.mkdir(parents=True, exist_ok=True) - all_counts: list[dict[str, int | str]] = [] - - for project in project_list: - display = _display_name(project) - logger.info("Processing %s (%s)...", project, display) - try: - counts = process_project( - labelled_root, - unlabelled_root, - project, - output_dir, - labelled_alt_roots=labelled_alt, - unlabelled_alt_roots=unlabelled_alt, - ) - except FileNotFoundError as exc: - logger.warning("Skipping %s: %s", project, exc) - continue - all_counts.append({"project": project, **counts}) - - if all_counts: - pd.DataFrame(all_counts).to_csv(output_dir / "counts_summary.csv", index=False) - logger.info("Wrote counts summary to %s", output_dir / "counts_summary.csv") - else: - raise typer.Exit(code=1) - - logger.info("Done. Outputs in %s", output_dir) - - -if __name__ == "__main__": - app() diff --git a/scripts/plot_analysis.py b/scripts/plot_analysis.py deleted file mode 100644 index 74d49700..00000000 --- a/scripts/plot_analysis.py +++ /dev/null @@ -1,1122 +0,0 @@ -"""Generate analysis plots from Winnow predict outputs. - -Usage: - python scripts/plot_analysis.py \ - --predictions-dir results/instanovo_helaqc_predictions_test/ \ - --split test \ - --label-mode labelled \ - --fasta fasta/human.fasta \ - [--model-dir models/instanovo_helaqc] \ - [--output-dir results/instanovo_helaqc_predictions_test/plots/] -""" - -from __future__ import annotations - -import argparse -import sys -import warnings -from pathlib import Path - -import matplotlib.pyplot as plt -import numpy as np -import polars as pl -import seaborn as sns -import yaml -from matplotlib.patches import Patch -from scipy.stats import gaussian_kde -from sklearn.calibration import calibration_curve -from sklearn.decomposition import PCA -from sklearn.preprocessing import StandardScaler - -REPO_ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(REPO_ROOT)) - -from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 - filter_and_annotate_preds, -) -from winnow.utils.proteome import load_proteome_haystack # noqa: E402 -from winnow.calibration.calibrator import TrainingHistory # noqa: E402 -from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 - -# ── Style — Paul Tol "bright" palette (colour-blind safe) ──────────── -_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] -_CORRECT_COLOUR = _PALETTE[0] -_INCORRECT_COLOUR = _PALETTE[1] -_MAIN_LINE_COLOUR = _PALETTE[0] -_RAW_LINE_COLOUR = _PALETTE[5] -_IDEAL_LINE_COLOUR = _PALETTE[6] - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) -warnings.filterwarnings("ignore", module="winnow") - - -def _spine_fmt(ax: plt.Axes) -> None: - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -def _save(fig: plt.Figure, out_dir: Path, name: str) -> None: - base = out_dir / name - fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - print(f" saved {name}") - - -# ── Plot functions ──────────────────────────────────────────────────── - - -def _df_for_raw_confidence_plots(df: pl.DataFrame) -> pl.DataFrame: - """Drop PSMs with negative raw confidence (Casanovo mass-mismatch penalty scores).""" - if "confidence" not in df.columns: - return df - n_neg = int((df["confidence"] < 0).sum()) - if n_neg == 0: - return df - print( - f" excluding {n_neg} PSMs with negative raw confidence " - "from raw-confidence plots" - ) - return df.filter(pl.col("confidence") >= 0) - - -def plot_calibration_curves( - df: pl.DataFrame, - label_col: str, - title: str, - bins: int = 10, - df_raw: pl.DataFrame | None = None, -) -> plt.Figure: - """Plot reliability curves for calibrated and (optionally) raw confidence.""" - fig, ax = plt.subplots(figsize=(8, 6)) - - frac_pos, mean_pred = calibration_curve( - df[label_col].to_numpy(), - df["calibrated_confidence"].to_numpy(), - n_bins=bins, - strategy="uniform", - ) - ax.plot( - mean_pred, - frac_pos, - marker="o", - color=_MAIN_LINE_COLOUR, - label="Calibrated confidence", - linewidth=1.5, - markersize=6, - zorder=3, - ) - - if df_raw is not None and len(df_raw) > 0 and "confidence" in df_raw.columns: - frac_pos, mean_pred = calibration_curve( - df_raw[label_col].to_numpy(), - df_raw["confidence"].to_numpy(), - n_bins=bins, - strategy="uniform", - ) - ax.plot( - mean_pred, - frac_pos, - marker="D", - color=_RAW_LINE_COLOUR, - label="Raw confidence", - linewidth=1.5, - markersize=6, - zorder=3, - ) - - ax.plot( - [0, 1], - [0, 1], - "--", - color=_IDEAL_LINE_COLOUR, - label="Perfectly calibrated", - alpha=0.7, - zorder=2, - ) - ax.set_xlabel("Mean predicted probability") - ax.set_ylabel("Fraction of positives") - ax.set_title(title) - ax.legend(loc="lower right") - ax.set_xlim([0, 1.05]) - ax.set_ylim([0, 1.05]) - ax.grid(False) - _spine_fmt(ax) - return fig - - -def plot_pr_curves( - df: pl.DataFrame, - label_col: str, - title: str, - df_raw: pl.DataFrame | None = None, -) -> plt.Figure: - """Plot precision–recall curves for calibrated and (optionally) raw confidence.""" - fig, ax = plt.subplots(figsize=(8, 6)) - - sorted_cal = df.sort("calibrated_confidence", descending=True) - labels = sorted_cal[label_col].to_numpy() - cum = np.cumsum(labels) - precision = cum / np.arange(1, len(labels) + 1) - recall = cum / len(labels) - ax.plot( - recall, - precision, - color=_MAIN_LINE_COLOUR, - label="Calibrated confidence", - linewidth=1.5, - ) - - if df_raw is not None and len(df_raw) > 0 and "confidence" in df_raw.columns: - sorted_raw = df_raw.sort("confidence", descending=True) - labels = sorted_raw[label_col].to_numpy() - cum = np.cumsum(labels) - precision = cum / np.arange(1, len(labels) + 1) - recall = cum / len(labels) - ax.plot( - recall, - precision, - color=_RAW_LINE_COLOUR, - label="Raw confidence", - linewidth=1.5, - ) - - ax.set_xlabel("Recall") - ax.set_ylabel("Precision") - ax.set_title(title) - ax.set_xlim(0, 1.05) - ax.set_ylim(0, 1.05) - ax.legend(loc="lower left") - ax.grid(False) - _spine_fmt(ax) - return fig - - -def plot_confidence_histogram( - df: pl.DataFrame, - label_col: str, - conf_col: str, - col_label: str, - title: str, - bins: int = 50, -) -> plt.Figure: - """Plot confidence histograms with KDE overlays for correct vs incorrect PSMs.""" - fig, ax = plt.subplots(1, 1, figsize=(7, 5)) - pos = df.filter(pl.col(label_col)) - neg = df.filter(~pl.col(label_col)) - n_data = neg[conf_col].to_numpy() - p_data = pos[conf_col].to_numpy() - - ax.hist( - n_data, - bins=bins, - alpha=0.6, - label="Incorrect", - density=False, - edgecolor="black", - color=_INCORRECT_COLOUR, - ) - ax.hist( - p_data, - bins=bins, - alpha=0.6, - label="Correct", - density=False, - edgecolor="black", - color=_CORRECT_COLOUR, - ) - - x_min = min(n_data.min(), p_data.min()) - x_max = max(n_data.max(), p_data.max()) - x_grid = np.linspace(x_min, x_max, 300) - bin_width = (x_max - x_min) / bins if bins > 1 else 1.0 - - if len(n_data) > 1: - y_neg = gaussian_kde(n_data)(x_grid) * len(n_data) * bin_width - ax.plot(x_grid, y_neg, color=_INCORRECT_COLOUR, lw=1.5) - if len(p_data) > 1: - y_pos = gaussian_kde(p_data)(x_grid) * len(p_data) * bin_width - ax.plot(x_grid, y_pos, color=_CORRECT_COLOUR, lw=1.5) - - ax.set_xlabel(col_label) - ax.set_ylabel("Frequency") - ax.legend(loc="upper center") - ax.grid(False) - ax.set_title(title) - _spine_fmt(ax) - fig.tight_layout() - return fig - - -def _fit_db_fdr( - df: pl.DataFrame, - correct_col: str, - confidence_feature: str = "calibrated_confidence", - drop: int = 10, -) -> DatabaseGroundedFDRControl: - """Fit a DatabaseGroundedFDRControl from per-row correctness labels.""" - ctrl = DatabaseGroundedFDRControl( - confidence_feature=confidence_feature, - drop=drop, - ) - sorted_df = df.sort(confidence_feature, descending=True) - correct_vals = sorted_df[correct_col].to_numpy().astype(float) - confidence_vals = sorted_df[confidence_feature].to_numpy() - precision = np.cumsum(correct_vals) / np.arange(1, len(sorted_df) + 1) - ctrl._fdr_values = np.array(1 - precision[drop:]) - ctrl._confidence_scores = confidence_vals[drop:] - return ctrl - - -def plot_fdr_accuracy( - df: pl.DataFrame, - correct_col: str, - residue_masses: dict, - title: str, - metric: str = "fdr", - use_proteome_shortcut: bool = False, -) -> plt.Figure: - """Compare non-parametric vs database-grounded FDR or q-value vs confidence.""" - del residue_masses, use_proteome_shortcut # retained for call-site compatibility - fig, ax = plt.subplots(figsize=(8, 6)) - col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" - winnow_col = col_name - - ctrl = _fit_db_fdr(df, correct_col) - - if metric == "fdr": - db_pd = ctrl.add_psm_fdr(df.to_pandas(), "calibrated_confidence") - else: - df_pd = df.to_pandas() - if "psm_q_value" in df_pd.columns: - df_pd = df_pd.drop(columns=["psm_q_value"]) - db_pd = ctrl.add_psm_q_value(df_pd, "calibrated_confidence") - db_df = pl.from_pandas(db_pd).select(["spectrum_id", col_name]) - - merged = ( - df.select(["spectrum_id", "calibrated_confidence", winnow_col]) - .join(db_df, on="spectrum_id", how="inner", suffix="_db") - .sort("calibrated_confidence") - ) - - conf = merged["calibrated_confidence"].to_numpy() - ax.plot( - conf, - merged[winnow_col].to_numpy(), - color=_MAIN_LINE_COLOUR, - label="Non-parametric", - linewidth=1.5, - ) - ax.plot( - conf, - merged[f"{col_name}_db"].to_numpy(), - color=_RAW_LINE_COLOUR, - label="Database-grounded", - linewidth=1.5, - ) - - ax.set_xlabel("Calibrated confidence") - ylabel = "FDR" if metric == "fdr" else "Q-value" - ax.set_ylabel(ylabel) - ax.set_title(title) - ax.legend(loc="upper right") - ax.grid(False) - _spine_fmt(ax) - return fig - - -def plot_ranked_qvalue( - df: pl.DataFrame, - correct_col: str, - residue_masses: dict, - title: str, -) -> plt.Figure: - """Ranked predictions vs q-value (non-parametric & database-grounded).""" - ctrl = _fit_db_fdr(df, correct_col) - test_pd = df.to_pandas() - test_pd_no_q = test_pd.drop(columns=["psm_q_value"], errors="ignore") - db_q = ctrl.add_psm_q_value(test_pd_no_q, "calibrated_confidence") - - sorted_np = test_pd.sort_values( - "calibrated_confidence", ascending=False - ).reset_index(drop=True) - sorted_db = db_q.sort_values("calibrated_confidence", ascending=False).reset_index( - drop=True - ) - ranks = np.arange(1, len(sorted_np) + 1) - - fig, ax = plt.subplots(figsize=(8, 6)) - ax.plot( - ranks, - sorted_np["psm_q_value"].values, - color=_MAIN_LINE_COLOUR, - label="Non-parametric", - linewidth=1.5, - ) - ax.plot( - ranks, - sorted_db["psm_q_value"].values, - color=_RAW_LINE_COLOUR, - label="Database-grounded", - linewidth=1.5, - ) - ax.set_xlabel("Ranked predictions") - ax.set_ylabel("Q-value") - ax.set_title(title) - ax.legend(loc="upper left") - _spine_fmt(ax) - return fig - - -def plot_ranked_fdr_pep(df: pl.DataFrame, title: str) -> plt.Figure: - """Ranked predictions vs non-parametric FDR and PEP.""" - sorted_df = df.sort("calibrated_confidence", descending=True) - ranks = np.arange(1, len(sorted_df) + 1) - fig, ax = plt.subplots(figsize=(8, 6)) - ax.plot( - ranks, - sorted_df["psm_fdr"].to_numpy(), - color=_MAIN_LINE_COLOUR, - label="FDR", - linewidth=1.5, - ) - if "psm_pep" in sorted_df.columns: - ax.plot( - ranks, - sorted_df["psm_pep"].to_numpy(), - color=_PALETTE[3], - label="PEP", - linewidth=1.5, - ) - ax.set_xlabel("Ranked predictions") - ax.set_ylabel("Error rate") - ax.set_title(title) - ax.legend(loc="upper left") - _spine_fmt(ax) - return fig - - -def plot_ranked_fdr_raw_vs_cal( - df: pl.DataFrame, - correct_col: str, - residue_masses: dict, - title: str, - metric: str = "fdr", - df_raw: pl.DataFrame | None = None, -) -> plt.Figure: - """Ranked predictions vs FDR/q-value for non-parametric and database-grounded on raw+calibrated.""" - test_pd = df.to_pandas() - test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") - - np_cal = test_pd.sort_values("calibrated_confidence", ascending=False).reset_index( - drop=True - ) - - db_cal_ctrl = _fit_db_fdr( - df, correct_col, confidence_feature="calibrated_confidence" - ) - raw_df = df_raw if df_raw is not None else df - db_raw_ctrl = _fit_db_fdr(raw_df, correct_col, confidence_feature="confidence") - - col_name = "psm_fdr" if metric == "fdr" else "psm_q_value" - add_fn = "add_psm_fdr" if metric == "fdr" else "add_psm_q_value" - sort_col_cal = "calibrated_confidence" - sort_col_raw = "confidence" - - db_cal = getattr(db_cal_ctrl, add_fn)(test_pd_no_q.copy(), sort_col_cal) - db_cal = db_cal.sort_values(sort_col_cal, ascending=False).reset_index(drop=True) - - raw_pd_no_q = raw_df.to_pandas().drop( - columns=["psm_q_value", "psm_fdr"], errors="ignore" - ) - db_raw = getattr(db_raw_ctrl, add_fn)(raw_pd_no_q.copy(), sort_col_raw) - db_raw = db_raw.sort_values(sort_col_raw, ascending=False).reset_index(drop=True) - - ranks_cal = np.arange(1, len(np_cal) + 1) - ranks_raw = np.arange(1, len(db_raw) + 1) - - fig, ax = plt.subplots(figsize=(8, 6)) - ylabel = "FDR" if metric == "fdr" else "Q-value" - np_col = "psm_fdr" if metric == "fdr" else "psm_q_value" - ax.plot( - ranks_cal, - np_cal[np_col].values, - color=_MAIN_LINE_COLOUR, - label="Non-parametric (calibrated)", - linewidth=1.5, - ) - ax.plot( - ranks_cal, - db_cal[col_name].values, - color=_RAW_LINE_COLOUR, - label="Database-grounded (calibrated)", - linewidth=1.5, - ) - ax.plot( - ranks_raw, - db_raw[col_name].values, - color=_PALETTE[3], - label="Database-grounded (raw)", - linewidth=1.5, - ) - ax.set_xlabel("Ranked predictions") - ax.set_ylabel(ylabel) - ax.set_title(title) - ax.legend(loc="upper left") - _spine_fmt(ax) - return fig - - -def plot_bar_psms_fdr( - df: pl.DataFrame, - correct_col: str, - residue_masses: dict, - title: str, - df_raw: pl.DataFrame | None = None, -) -> plt.Figure: - """Bar plot of PSMs at q-value thresholds (calibrated vs raw, database-grounded).""" - test_pd = df.to_pandas() - test_pd_no_q = test_pd.drop(columns=["psm_q_value", "psm_fdr"], errors="ignore") - - db_cal_ctrl = _fit_db_fdr( - df, correct_col, confidence_feature="calibrated_confidence" - ) - raw_df = df_raw if df_raw is not None else df - db_raw_ctrl = _fit_db_fdr(raw_df, correct_col, confidence_feature="confidence") - - db_cal = db_cal_ctrl.add_psm_q_value(test_pd_no_q.copy(), "calibrated_confidence") - raw_pd_no_q = raw_df.to_pandas().drop( - columns=["psm_q_value", "psm_fdr"], errors="ignore" - ) - db_raw = db_raw_ctrl.add_psm_q_value(raw_pd_no_q.copy(), "confidence") - - thresholds = [0.001, 0.01, 0.05, 0.1] - counts_cal = [int((db_cal["psm_q_value"] <= t).sum()) for t in thresholds] - counts_raw = [int((db_raw["psm_q_value"] <= t).sum()) for t in thresholds] - - x = np.arange(len(thresholds)) - width, gap = 0.32, 0.04 - - fig, ax = plt.subplots(figsize=(8, 6)) - bars_cal = ax.bar( - x - width / 2 - gap / 2, - counts_cal, - width, - label="Calibrated confidence", - color=_MAIN_LINE_COLOUR, - edgecolor="black", - linewidth=1, - ) - bars_raw = ax.bar( - x + width / 2 + gap / 2, - counts_raw, - width, - label="Raw confidence", - color=_RAW_LINE_COLOUR, - edgecolor="black", - linewidth=1, - ) - ax.set_xlabel("FDR threshold") - ax.set_ylabel("Peptide-spectrum matches") - ax.set_title(title) - ax.set_xticks(x) - ax.set_xticklabels([str(t) for t in thresholds]) - ax.legend(loc="upper left") - - for bar_group in [bars_cal, bars_raw]: - for bar in bar_group: - h = bar.get_height() - ax.annotate( - f"{h:,}", - xy=(bar.get_x() + bar.get_width() / 2, h), - xytext=(0, 3), - textcoords="offset points", - ha="center", - va="bottom", - fontsize=10, - ) - _spine_fmt(ax) - return fig - - -def plot_raw_vs_cal_scatter( - df: pl.DataFrame, - label_col: str, - title: str, -) -> plt.Figure: - """Raw confidence vs calibrated confidence coloured by correctness.""" - fig, ax = plt.subplots(figsize=(8, 7)) - inc = df.filter(~pl.col(label_col)) - cor = df.filter(pl.col(label_col)) - ax.scatter( - inc["confidence"].to_numpy(), - inc["calibrated_confidence"].to_numpy(), - c=_INCORRECT_COLOUR, - label="Incorrect", - s=10, - alpha=0.3, - rasterized=True, - ) - ax.scatter( - cor["confidence"].to_numpy(), - cor["calibrated_confidence"].to_numpy(), - c=_CORRECT_COLOUR, - label="Correct", - s=10, - alpha=0.3, - rasterized=True, - ) - ax.plot( - [0, 1], - [0, 1], - color=_IDEAL_LINE_COLOUR, - linestyle="--", - linewidth=1, - label="Identity", - ) - ax.set_xlabel("Raw confidence") - ax.set_ylabel("Calibrated confidence") - ax.set_title(title) - ax.legend(loc="upper left") - _spine_fmt(ax) - return fig - - -def plot_pca_features( - df: pl.DataFrame, - label_col: str, - title: str, -) -> tuple[plt.Figure, PCA, list[str]]: - """PCA of calibrator features coloured by correctness.""" - feature_cols = [ - "confidence", - "mass_error_ppm", - "ion_matches", - "ion_match_intensity", - "complementary_ion_count", - "max_ion_gap", - "spectral_angle", - "xcorr", - "irt_error", - "margin", - "median_margin", - "entropy", - "z-score", - "edit_distance", - "min_token_probability", - "std_token_probability", - ] - available = [c for c in feature_cols if c in df.columns] - feat_df = df.select(available).to_pandas().dropna() - labels = df.filter(pl.all_horizontal([pl.col(c).is_not_null() for c in available]))[ - label_col - ].to_numpy() - - scaler = StandardScaler() - features_scaled = scaler.fit_transform(feat_df.values) - pca = PCA(n_components=2) - coords = pca.fit_transform(features_scaled) - - fig, ax = plt.subplots(figsize=(8, 7)) - mask_inc, mask_cor = ~labels, labels - ax.scatter( - coords[mask_inc, 0], - coords[mask_inc, 1], - c=_INCORRECT_COLOUR, - label="Incorrect", - s=10, - alpha=0.3, - rasterized=True, - ) - ax.scatter( - coords[mask_cor, 0], - coords[mask_cor, 1], - c=_CORRECT_COLOUR, - label="Correct", - s=10, - alpha=0.3, - rasterized=True, - ) - ax.set_xlabel(f"PC 1 ({pca.explained_variance_ratio_[0]:.1%} variance)") - ax.set_ylabel(f"PC 2 ({pca.explained_variance_ratio_[1]:.1%} variance)") - ax.set_title(title) - ax.legend(loc="upper left") - _spine_fmt(ax) - return fig, pca, available - - -def plot_pca_loadings( - pca: PCA, - feature_names: list[str], - title: str, -) -> plt.Figure: - """PCA loadings for PC1 and PC2, ordered by |PC1|.""" - pretty = { - "confidence": "Raw confidence", - "mass_error_ppm": "Log absolute mass error (ppm)", - "ion_matches": "Ion matches", - "ion_match_intensity": "Ion match intensity", - "complementary_ion_count": "Complementary ion count", - "max_ion_gap": "Maximum ion gap", - "spectral_angle": "Spectral angle", - "xcorr": "Cross-correlation", - "irt_error": "Retention time error", - "margin": "Margin", - "median_margin": "Median margin", - "entropy": "Entropy", - "z-score": "Z-score", - "edit_distance": "Edit distance", - "min_token_probability": "Minimum token probability", - "std_token_probability": "Token probability std. dev.", - } - pc1 = pca.components_[0] - pc2 = pca.components_[1] - names = [pretty.get(c, c) for c in feature_names] - order = np.argsort(np.abs(pc1))[::-1] - - y = np.arange(len(names)) - fig, ax = plt.subplots(figsize=(10, 7)) - ax.barh(y, pc1[order], color=_MAIN_LINE_COLOUR, alpha=0.6, edgecolor="black") - ax.barh(y, pc2[order], color=_RAW_LINE_COLOUR, alpha=0.4, edgecolor="black") - ax.set_yticks(y) - ax.set_yticklabels([names[i] for i in order]) - ax.invert_yaxis() - ax.set_xlabel("Loading value") - ax.set_title(title) - ax.axvline(0, color="black", linewidth=0.5) - ax.legend( - handles=[ - Patch(facecolor=_MAIN_LINE_COLOUR, alpha=0.6, label="PC 1 loading"), - Patch(facecolor=_RAW_LINE_COLOUR, alpha=0.4, label="PC 2 loading"), - ], - loc="lower right", - ) - _spine_fmt(ax) - return fig - - -def plot_scatter_feature_vs_conf( - df: pl.DataFrame, - label_col: str, - x_col: str, - y_col: str, - title: str, - x_label: str | None = None, - y_label: str | None = None, -) -> plt.Figure: - """Scatter of x_col vs y_col coloured by correctness.""" - fig, ax = plt.subplots(figsize=(8, 7)) - inc = df.filter(~pl.col(label_col)) - cor = df.filter(pl.col(label_col)) - ax.scatter( - inc[x_col].to_numpy(), - inc[y_col].to_numpy(), - c=_INCORRECT_COLOUR, - label="Incorrect", - s=10, - alpha=0.3, - rasterized=True, - ) - ax.scatter( - cor[x_col].to_numpy(), - cor[y_col].to_numpy(), - c=_CORRECT_COLOUR, - label="Correct", - s=10, - alpha=0.3, - rasterized=True, - ) - ax.set_xlabel(x_label or x_col.replace("_", " ").title()) - ax.set_ylabel(y_label or y_col.replace("_", " ").title()) - ax.set_title(title) - ax.legend(loc="upper left") - _spine_fmt(ax) - return fig - - -# ── Main logic ──────────────────────────────────────────────────────── - - -def _load_residue_masses() -> dict: - cfg_path = REPO_ROOT / "winnow" / "configs" / "residues.yaml" - with open(cfg_path) as f: - return yaml.safe_load(f)["residue_masses"] - - -def _load_data(predictions_dir: Path) -> pl.DataFrame: - preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") - meta_path = predictions_dir / "metadata.csv" - if meta_path.exists(): - meta = pl.read_csv(meta_path) - preds = preds.join(meta, on="spectrum_id", how="inner") - return preds - - -_SPLIT_DISPLAY_NAMES = { - "test": "test set", - "unlabelled": "unlabelled space", - "raw_less_train": "full search space", -} - - -def _split_display(split: str) -> str: - key = split.strip().replace("-", "_") - return _SPLIT_DISPLAY_NAMES.get(key, split) - - -def _dns_model_tag(dns_model: str | None) -> str: - return f" ({dns_model})" if dns_model else "" - - -def _split_title(split_label: str, dns_model: str | None = None) -> str: - return f"{split_label}{_dns_model_tag(dns_model)}" - - -def _eval_title( - prefix: str, - split_label: str, - eval_kind: str, - dns_model: str | None = None, -) -> str: - return f"{prefix} {_split_title(split_label, dns_model)}\nusing {eval_kind}" - - -def _save_training_history_plot(model_dir: Path, out_dir: Path) -> None: - hist_path = model_dir / "training_history.json" - if not hist_path.exists(): - return - print("Plotting training history") - th = TrainingHistory.load(str(hist_path)) - th.plot(output_path=out_dir / "training_history.png", show=False) - print(" saved training_history") - - -def _plot_calibration_and_pr( - df: pl.DataFrame, - df_raw_conf: pl.DataFrame, - split: str, - labelled: bool, - out_dir: Path, - dns_model: str | None = None, -) -> None: - split_label = _split_display(split) - print("Plotting calibration curves") - if labelled: - fig = plot_calibration_curves( - df, - "correct", - _eval_title( - "Calibration curves for", split_label, "database search", dns_model - ), - df_raw=df_raw_conf, - ) - _save(fig, out_dir, f"calibration_{split}_db_search") - - fig = plot_calibration_curves( - df, - "proteome_hit", - _eval_title( - "Calibration curves for", split_label, "proteome mapping", dns_model - ), - df_raw=df_raw_conf, - ) - _save(fig, out_dir, f"calibration_{split}_proteome") - - print("Plotting PR curves") - if labelled: - fig = plot_pr_curves( - df, - "correct", - _eval_title("PR curves for", split_label, "database search", dns_model), - df_raw=df_raw_conf, - ) - _save(fig, out_dir, f"pr_{split}_db_search") - - fig = plot_pr_curves( - df, - "proteome_hit", - _eval_title("PR curves for", split_label, "proteome mapping", dns_model), - df_raw=df_raw_conf, - ) - _save(fig, out_dir, f"pr_{split}_proteome") - - -def _plot_confidence_histograms( - df: pl.DataFrame, - df_raw_conf: pl.DataFrame, - split: str, - labelled: bool, - out_dir: Path, - dns_model: str | None = None, -) -> None: - split_label = _split_display(split) - print("Plotting confidence histograms") - for conf_col, conf_label, tag in [ - ("confidence", "Raw confidence", "raw"), - ("calibrated_confidence", "Calibrated confidence", "cal"), - ]: - hist_df = df_raw_conf if conf_col == "confidence" else df - if labelled: - fig = plot_confidence_histogram( - hist_df, - "correct", - conf_col, - conf_label, - _eval_title( - f"{conf_label} for", split_label, "database search", dns_model - ), - ) - _save(fig, out_dir, f"hist_{tag}_{split}_db_search") - - fig = plot_confidence_histogram( - hist_df, - "proteome_hit", - conf_col, - conf_label, - _eval_title( - f"{conf_label} for", split_label, "proteome mapping", dns_model - ), - ) - _save(fig, out_dir, f"hist_{tag}_{split}_proteome") - - -def _plot_fdr_accuracy_plots( - df: pl.DataFrame, - split: str, - labelled: bool, - residue_masses: dict, - out_dir: Path, - dns_model: str | None = None, -) -> None: - split_label = _split_display(split) - print("Plotting FDR accuracy") - use_shortcut = not labelled - for metric, tag in [("fdr", "fdr"), ("q_value", "qvalue")]: - metric_name = "FDR" if metric == "fdr" else "Q-value" - if labelled: - fig = plot_fdr_accuracy( - df, - "correct", - residue_masses, - _eval_title( - f"{metric_name} accuracy for", - split_label, - "database search", - dns_model, - ), - metric="fdr" if metric == "fdr" else "q_value", - use_proteome_shortcut=False, - ) - _save(fig, out_dir, f"{tag}_{split}_db_search") - - fig = plot_fdr_accuracy( - df, - "proteome_hit", - residue_masses, - _eval_title( - f"{metric_name} accuracy for", - split_label, - "proteome mapping", - dns_model, - ), - metric="fdr" if metric == "fdr" else "q_value", - use_proteome_shortcut=use_shortcut, - ) - _save(fig, out_dir, f"{tag}_{split}_proteome") - - -def _plot_labelled_diagnostics( - df: pl.DataFrame, - df_raw_conf: pl.DataFrame, - split: str, - residue_masses: dict, - out_dir: Path, - dns_model: str | None = None, -) -> None: - split_label = _split_display(split) - title_split = _split_title(split_label, dns_model) - label_col = "correct" - print("Plotting labelled-only diagnostics") - - fig = plot_ranked_qvalue( - df, - label_col, - residue_masses, - _eval_title( - "Ranked predictions vs q-value for", - split_label, - "database search", - dns_model, - ), - ) - _save(fig, out_dir, f"ranked_qvalue_{split}_db_search") - - fig = plot_ranked_fdr_pep( - df, f"Ranked predictions vs FDR and PEP for {title_split}" - ) - _save(fig, out_dir, f"ranked_fdr_pep_{split}_nonparametric") - - for metric, file_tag, metric_name in [ - ("fdr", "fdr", "FDR"), - ("q_value", "qvalue", "q-value"), - ]: - fig = plot_ranked_fdr_raw_vs_cal( - df, - label_col, - residue_masses, - _eval_title( - f"Ranked predictions vs {metric_name} for", - split_label, - "database search", - dns_model, - ), - metric=metric, - df_raw=df_raw_conf, - ) - _save(fig, out_dir, f"ranked_{file_tag}_raw_vs_cal_{split}_db_search") - - fig = plot_bar_psms_fdr( - df, - label_col, - residue_masses, - f"PSMs at database-grounded FDR thresholds for {title_split}", - df_raw=df_raw_conf, - ) - _save(fig, out_dir, f"bar_psms_fdr_thresholds_{split}_db_search") - - if len(df_raw_conf) > 0: - fig = plot_raw_vs_cal_scatter( - df_raw_conf, - label_col, - f"Raw vs calibrated confidence for {title_split}", - ) - _save(fig, out_dir, f"scatter_raw_vs_cal_confidence_{split}") - - if "margin" in df.columns and len(df_raw_conf) > 0: - fig = plot_scatter_feature_vs_conf( - df_raw_conf, - label_col, - "margin", - "confidence", - f"Raw confidence vs margin for {title_split}", - x_label="Margin", - y_label="Raw confidence", - ) - _save(fig, out_dir, f"scatter_raw_confidence_vs_margin_{split}") - - fig = plot_scatter_feature_vs_conf( - df, - label_col, - "margin", - "calibrated_confidence", - f"Calibrated confidence vs margin for {title_split}", - x_label="Margin", - y_label="Calibrated confidence", - ) - _save(fig, out_dir, f"scatter_cal_confidence_vs_margin_{split}") - - fig, pca_model, feat_names = plot_pca_features( - df, label_col, f"PCA of calibrator features for {title_split}" - ) - _save(fig, out_dir, f"pca_features_{split}") - - fig = plot_pca_loadings( - pca_model, feat_names, "PCA loadings for first two principal components" - ) - _save(fig, out_dir, f"pca_loadings_pc1_pc2_{split}") - - -def main() -> None: - """Load predictions, annotate proteome hits, and write analysis plots.""" - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument( - "--predictions-dir", - type=Path, - required=True, - help="Winnow predict output dir with preds_and_fdr_metrics.csv", - ) - parser.add_argument( - "--split", - type=str, - required=True, - help="Split id for filenames; titles use test set / unlabelled space / full search space", - ) - parser.add_argument( - "--label-mode", - choices=["labelled", "unlabelled"], - required=True, - help="labelled = has 'correct' column; unlabelled = proteome mapping only", - ) - parser.add_argument( - "--fasta", type=Path, required=True, help="FASTA file for proteome annotation" - ) - parser.add_argument( - "--output-dir", - type=Path, - default=None, - help="Output directory for plots (defaults to predictions-dir/plots/)", - ) - parser.add_argument( - "--model-dir", - type=Path, - default=None, - help="Model directory for training history plot", - ) - parser.add_argument( - "--dns-model", - type=str, - default=None, - help="Upstream DNS model name for plot titles (e.g. InstaNovo, Casanovo, $\\pi$-PrimeNovo)", - ) - args = parser.parse_args() - - out_dir = args.output_dir or (args.predictions_dir / "plots") - out_dir.mkdir(parents=True, exist_ok=True) - - split = args.split - dns_model = args.dns_model - labelled = args.label_mode == "labelled" - residue_masses = _load_residue_masses() - - print(f"Loading predictions from {args.predictions_dir}") - df = _load_data(args.predictions_dir) - - from instanovo.utils.metrics import Metrics - from instanovo.utils.residues import ResidueSet - - metrics = Metrics( - residue_set=ResidueSet(residue_masses=residue_masses), - isotope_error_range=(0, 1), - ) - - print(f"Annotating with proteome hits from {args.fasta}") - haystack = load_proteome_haystack(str(args.fasta)) - df = filter_and_annotate_preds(df, haystack, metrics, min_residue_length=7) - df_raw_conf = _df_for_raw_confidence_plots(df) - - if args.model_dir is not None: - _save_training_history_plot(args.model_dir, out_dir) - - _plot_calibration_and_pr( - df, df_raw_conf, split, labelled, out_dir, dns_model=dns_model - ) - _plot_confidence_histograms( - df, df_raw_conf, split, labelled, out_dir, dns_model=dns_model - ) - _plot_fdr_accuracy_plots( - df, split, labelled, residue_masses, out_dir, dns_model=dns_model - ) - - if labelled: - _plot_labelled_diagnostics( - df, df_raw_conf, split, residue_masses, out_dir, dns_model=dns_model - ) - - print(f"\nAll plots saved to {out_dir}") - - -if __name__ == "__main__": - main() diff --git a/scripts/plot_calibrator_generalisation_heatmap.py b/scripts/plot_calibrator_generalisation_heatmap.py deleted file mode 100644 index e82837da..00000000 --- a/scripts/plot_calibrator_generalisation_heatmap.py +++ /dev/null @@ -1,265 +0,0 @@ -"""Plot PR-AUC heatmaps for calibrator generalisation results. - -Reads the combined CSV produced by ``evaluate_calibrator_generalisation.py`` -and creates heatmaps comparing raw vs calibrated confidence PR-AUC values. -""" - -import logging -import sys -from pathlib import Path -from typing import Annotated - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import polars as pl -import seaborn as sns -from matplotlib.colors import LinearSegmentedColormap -from rich.logging import RichHandler -import typer - -_REPO_ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(_REPO_ROOT)) - -from scripts.calibrator_generalisation_utils import SPECIES_NAME_MAPPING # noqa: E402 - -# --------------------------------------------------------------------------- -# Logging -# --------------------------------------------------------------------------- -logger = logging.getLogger("winnow.plot_generalisation_heatmap") -logger.setLevel(logging.INFO) -logger.propagate = False -logger.addHandler(RichHandler()) - -# --------------------------------------------------------------------------- -# Style — Paul Tol "bright" palette + "sunset" diverging colourmap -# --------------------------------------------------------------------------- -_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] - -_SUNSET_COLORS = [ - "#364B9A", - "#4A7BB7", - "#6EA6CD", - "#98CAE1", - "#C2E4EF", - "#EAECCC", - "#FEDA8B", - "#FDB366", - "#F67E4B", - "#DD3D2D", - "#A50026", -] -_BAD_COLOUR = "#FFFFFF" - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=2) - - -def _diverging_cmap() -> LinearSegmentedColormap: - cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) - cmap.set_bad(color=_BAD_COLOUR) - return cmap - - -def _sequential_cmap() -> LinearSegmentedColormap: - cmap = LinearSegmentedColormap.from_list( - "tol_sunset_seq", _SUNSET_COLORS[5:], N=256 - ) - cmap.set_bad(color=_BAD_COLOUR) - return cmap - - -# --------------------------------------------------------------------------- -# PR-AUC computation -# --------------------------------------------------------------------------- -def compute_pr_auc( - input_dataset: pd.DataFrame, - confidence_column: str, - label_column: str, -) -> float: - """Compute Area Under Curve for precision-recall curve.""" - if len(input_dataset) == 0: - return 0.0 - - sorted_data = input_dataset[[confidence_column, label_column]].sort_values( - by=confidence_column, ascending=False - ) - - cum_correct = np.cumsum(sorted_data[label_column]) - precision = cum_correct / np.arange(1, len(sorted_data) + 1) - recall = ( - cum_correct / cum_correct.iloc[-1] - if cum_correct.iloc[-1] > 0 - else np.zeros_like(cum_correct) - ) - - if len(precision) < 2: - return 0.0 - - from sklearn.metrics import auc - - return auc(recall, precision) - - -# --------------------------------------------------------------------------- -# Heatmap creation -# --------------------------------------------------------------------------- -def _save_fig(fig: plt.Figure, base_path: Path) -> None: - """Save figure as both PNG and PDF.""" - fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def create_auc_heatmap( - auc_df: pd.DataFrame, - output_path: Path, - title: str = "Calibrator generalisation PR-AUC heatmap", -) -> None: - """Create and save a heatmap of PR-AUC values.""" - fig, ax = plt.subplots(figsize=(12, 10)) - - sns.heatmap( - auc_df, - annot=True, - fmt=".3f", - cmap=_sequential_cmap(), - cbar_kws={"label": "PR-AUC"}, - square=True, - linewidths=0.5, - ax=ax, - ) - - ax.set_title(title) - ax.set_xlabel("Test dataset") - ax.set_ylabel("Train dataset") - ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") - ax.set_yticklabels(ax.get_yticklabels(), rotation=0) - - base = str(output_path).removesuffix(".png") - _save_fig(fig, Path(base)) - logger.info("Heatmap saved to %s", output_path) - - -def create_comparison_heatmaps(results_path: Path, output_dir: Path) -> None: - """Create heatmaps comparing raw vs calibrated confidence PR-AUC values.""" - logger.info("Scanning results from %s", results_path) - results = pl.scan_csv(results_path) - - trained_datasets = sorted( - results.select(pl.col("trained_on_dataset")) - .unique() - .collect() - .to_series() - .to_list() - ) - test_datasets = sorted( - results.select(pl.col("test_dataset")).unique().collect().to_series().to_list() - ) - logger.info("Trained datasets: %s", trained_datasets) - logger.info("Test datasets: %s", test_datasets) - - trained_labels = [SPECIES_NAME_MAPPING.get(ds, ds) for ds in trained_datasets] - test_labels = [SPECIES_NAME_MAPPING.get(ds, ds) for ds in test_datasets] - - # Compute PR-AUC matrices for both confidence types - auc_matrices = {} - for conf_type in ["confidence", "calibrated_confidence"]: - auc_matrix = [] - for trained_dataset in trained_datasets: - auc_row = [] - for test_dataset in test_datasets: - logger.info( - "Computing PR-AUC (%s) for trained=%s, test=%s", - conf_type, - trained_dataset, - test_dataset, - ) - subset = ( - results.filter( - (pl.col("trained_on_dataset") == trained_dataset) - & (pl.col("test_dataset") == test_dataset) - ) - .collect() - .to_pandas() - ) - - if len(subset) > 0: - auc_row.append(compute_pr_auc(subset, conf_type, "correct")) - else: - auc_row.append(np.nan) - auc_matrix.append(auc_row) - - auc_matrices[conf_type] = pd.DataFrame( - auc_matrix, index=trained_labels, columns=test_labels - ) - - # Individual heatmaps - for conf_type, auc_df in auc_matrices.items(): - conf_name = conf_type.replace("_", " ") - output_path = ( - output_dir / f"calibrator_generalisation_{conf_type}_auc_heatmap.png" - ) - create_auc_heatmap( - auc_df, - output_path, - f"Calibrator generalisation {conf_name} PR-AUC", - ) - - # Difference heatmap (calibrated - raw) - diff_matrix = auc_matrices["calibrated_confidence"] - auc_matrices["confidence"] - - fig, ax = plt.subplots(figsize=(12, 10)) - sns.heatmap( - diff_matrix, - annot=True, - fmt=".3f", - cmap=_diverging_cmap(), - center=0, - cbar_kws={"label": r"PR-AUC difference $(\mathrm{calibrated} - \mathrm{raw})$"}, - square=True, - linewidths=0.5, - ax=ax, - ) - ax.set_title("Calibrator generalisation PR-AUC improvement") - ax.set_xlabel("Test dataset") - ax.set_ylabel("Train dataset") - ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right") - ax.set_yticklabels(ax.get_yticklabels(), rotation=0) - - diff_base = output_dir / "calibrator_generalisation_auc_difference_heatmap" - _save_fig(fig, diff_base) - logger.info("Difference heatmap saved to %s", diff_base) - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- -_DEFAULT_OUTPUT_DIR = Path("results/generalisation/plots") - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - - -@app.command() -def main( - results_path: Annotated[ - Path, typer.Option(help="Path to calibrator generalisation results CSV.") - ], - output_dir: Annotated[ - Path, typer.Option(help="Directory to save plots.") - ] = _DEFAULT_OUTPUT_DIR, -) -> None: - """Create PR-AUC heatmaps for calibrator generalisation results.""" - output_dir.mkdir(parents=True, exist_ok=True) - - if not results_path.exists(): - logger.error("Results file not found: %s", results_path) - raise typer.Exit(1) - - logger.info("Loading results from: %s", results_path) - logger.info("Saving plots to: %s", output_dir) - - create_comparison_heatmaps(results_path, output_dir) - - -if __name__ == "__main__": - app() diff --git a/scripts/plot_eval_results.py b/scripts/plot_eval_results.py deleted file mode 100644 index 32bccbfa..00000000 --- a/scripts/plot_eval_results.py +++ /dev/null @@ -1,984 +0,0 @@ -"""Generate publication-quality evaluation plots from ``winnow predict`` outputs. - -Supports both annotated (database-grounded) and raw (proteome-hit) evaluation -modes, producing six plots per project: precision-recall, FDR run, true vs -estimated FDR (full + zoomed), probability calibration, and before/after score -histograms. -""" - -from __future__ import annotations - -import logging -from pathlib import Path -from typing import Annotated -import warnings - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import seaborn as sns -import typer -from rich.logging import RichHandler - -from winnow.fdr.nonparametric import NonParametricFDRControl - -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -logger.propagate = False -if not logger.handlers: - logger.addHandler(RichHandler()) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -# Filter by the specific message or category -warnings.filterwarnings("ignore", message=".*range of fitted confidence scores.*") -warnings.filterwarnings("ignore", message=".*range of fitted FDR thresholds.*") - -# --------------------------------------------------------------------------- -# Dataset display names -# --------------------------------------------------------------------------- -DATASET_DISPLAY_NAMES: dict[str, str] = { - "gluc": "HeLa degradome", - "helaqc": "HeLa single shot", - "herceptin": "Herceptin", - "immuno": "Immunopeptidomics-1", - "celegans": "$\\it{C.\\;elegans}$", - "sbrodae": "$\\it{Scalindua\\;brodae}$", - "PXD019483": "HepG2", - "snakevenoms": "Snake venomics", - "tplantibodies": "Therapeutic nanobodies", - "woundfluids": "Wound exudates", - "PXD004732": "ProteomeTools-1", - "PXD014877": "$\\it{C.\\;elegans}$", - "PXD023064": "Immunopeptidomics-2", - "astral": "Astral $\\it{E.\\;coli}$", - "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", - "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", - "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", - "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", - "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", -} - -_FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") -_DISPLAY_NAME_LOOKUP = {k.lower(): v for k, v in DATASET_DISPLAY_NAMES.items()} - -# Paul Tol "bright" palette (colour-blind safe) -_PALETTE = [ - "#4477AA", - "#EE6677", - "#228833", - "#CCBB44", - "#66CCEE", - "#AA3377", - "#BBBBBB", -] -_CORRECT_COLOUR = _PALETTE[0] -_INCORRECT_COLOUR = _PALETTE[1] -_MAIN_LINE_COLOUR = _PALETTE[0] -_RAW_LINE_COLOUR = _PALETTE[5] -_IDEAL_LINE_COLOUR = _PALETTE[6] -_BAND_COLOUR = _PALETTE[0] - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -_DIAGNOSTIC_ALPHAS = (0.01, 0.05, 0.10) -_HOEFFDING_DELTA = 0.05 - - -def _normalize_project_key(key: str) -> str: - """Strip eval suffixes and nested path segments for display-name lookup.""" - key = key.strip() - if "/" in key: - key = key.rsplit("/", 1)[-1] - for suffix in _FOLDER_SUFFIXES: - if key.endswith(suffix): - return key[: -len(suffix)] - return key - - -def _display_name(key: str) -> str: - """Look up the publication-ready display name for a dataset key.""" - normalized = _normalize_project_key(key) - if normalized in DATASET_DISPLAY_NAMES: - return DATASET_DISPLAY_NAMES[normalized] - return _DISPLAY_NAME_LOOKUP.get(normalized.lower(), normalized) - - -def _ground_truth_qualifier(eval_type: str) -> str: - """Return the title-friendly ground truth qualifier for plot titles.""" - if eval_type in ("annotated", "labelled"): - return "using database search" - return "using proteome mapping" - - -def _save_fig(fig: plt.Figure, base_path: Path) -> None: - """Save figure as both PNG and PDF.""" - fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def _style_ax(ax: plt.Axes) -> None: - ax.grid(False) - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -# --------------------------------------------------------------------------- -# PR curve (non-standard cumulative definition) -# --------------------------------------------------------------------------- -def _compute_precision_recall( - df: pd.DataFrame, confidence_col: str = "calibrated_confidence" -) -> pd.DataFrame: - """Non-standard cumulative PR curve matching the codebase convention.""" - sorted_df = df.sort_values(confidence_col, ascending=False) - labels = sorted_df["correct"].values - cum_correct = np.cumsum(labels) - n = len(labels) - precision = cum_correct / np.arange(1, n + 1) - recall = cum_correct / n - return pd.DataFrame({"precision": precision, "recall": recall}) - - -def plot_precision_recall( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Plot precision-recall curve.""" - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - pr_cal = _compute_precision_recall(df, "calibrated_confidence") - pr_raw = _compute_precision_recall(df, "confidence") - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.plot( - pr_raw["recall"], - pr_raw["precision"], - color=_RAW_LINE_COLOUR, - lw=1.5, - label="Raw confidence", - ) - ax.plot( - pr_cal["recall"], - pr_cal["precision"], - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Calibrated confidence", - ) - ax.set_xlabel("Recall") - ax.set_ylabel("Precision") - ax.set_title(f"{display} precision-recall {qualifier}") - ax.set_xlim(0, 1) - ax.set_ylim(0, 1.02) - ax.legend(loc="lower right") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"pr_curve_{project}") - - -# --------------------------------------------------------------------------- -# FDR run plot -# --------------------------------------------------------------------------- -def plot_fdr_run( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Plot calibrated confidence vs estimated and true PSM FDR.""" - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - - df = df.sort_values("calibrated_confidence") - - true_fdr_ctrl = _fit_database_grounded_fdr(df) - true_fdr_df = true_fdr_ctrl.add_psm_fdr( - df.copy(), confidence_col="calibrated_confidence" - ) - true_fdr_df = true_fdr_df.sort_values("calibrated_confidence") - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.plot( - df["calibrated_confidence"].values, - df["psm_fdr"].values, - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Non-parametric", - ) - ax.plot( - true_fdr_df["calibrated_confidence"].values, - true_fdr_df["psm_fdr"].values, - color=_RAW_LINE_COLOUR, - lw=1.5, - label="Database-grounded", - ) - ax.set_xlabel("Calibrated confidence") - ax.set_ylabel("PSM FDR") - ax.set_title(f"{display} FDR run {qualifier}") - ax.legend(loc="upper right") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"fdr_run_{project}") - - -# --------------------------------------------------------------------------- -# Q-value run plot -# --------------------------------------------------------------------------- -def _fit_database_grounded_fdr( - df: pd.DataFrame, - confidence_col: str = "calibrated_confidence", - correct_col: str = "correct", - drop: int = 10, -) -> NonParametricFDRControl: - """Fit an FDR controller using ground-truth labels. - - Replicates the fitting logic of ``DatabaseGroundedFDRControl`` (computing - FDR as 1 − precision over sorted predictions, with the first *drop* entries - removed) without pulling in the instanovo dependency. - """ - sorted_desc = df.sort_values(confidence_col, ascending=False) - labels = sorted_desc[correct_col].values.astype(float) - precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) - confidence = sorted_desc[confidence_col].values - - ctrl = NonParametricFDRControl() - ctrl._fdr_values = (1.0 - precision)[drop:] - ctrl._confidence_scores = confidence[drop:] - return ctrl - - -def plot_q_value_run( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Plot calibrated confidence vs estimated and true PSM q-values.""" - if "psm_q_value" not in df.columns: - logger.warning( - "Skipping q-value run plot for %s: psm_q_value column missing", project - ) - return - - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - - sorted_df = df.sort_values("calibrated_confidence") - - true_fdr_ctrl = _fit_database_grounded_fdr(df) - qval_input = df[["calibrated_confidence"]].copy() - true_q_df = true_fdr_ctrl.add_psm_q_value( - qval_input, confidence_col="calibrated_confidence" - ) - true_q_df = true_q_df.sort_values("calibrated_confidence") - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.plot( - sorted_df["calibrated_confidence"].values, - sorted_df["psm_q_value"].values, - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Non-parametric", - ) - ax.plot( - true_q_df["calibrated_confidence"].values, - true_q_df["psm_q_value"].values, - color=_RAW_LINE_COLOUR, - lw=1.5, - label="Database-grounded", - ) - ax.set_xlabel("Calibrated confidence") - ax.set_ylabel("PSM q-value") - ax.set_title(f"{display} q-value run {qualifier}") - ax.legend(loc="upper right") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"qvalue_run_{project}") - - -# --------------------------------------------------------------------------- -# FDR / q-value run plots with Hoeffding confidence bands -# --------------------------------------------------------------------------- -def _hoeffding_band_arrays(n: int) -> np.ndarray: - """Compute pointwise Hoeffding half-widths for ranks 1..n (descending confidence).""" - ranks = np.arange(1, n + 1) - return np.sqrt(np.log(2.0 / _HOEFFDING_DELTA) / (2.0 * ranks)) - - -def plot_fdr_run_with_bands( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """FDR run plot with Hoeffding 95% confidence band on the non-parametric curve.""" - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - - df = df.sort_values("calibrated_confidence") - - true_fdr_ctrl = _fit_database_grounded_fdr(df) - true_fdr_df = true_fdr_ctrl.add_psm_fdr( - df.copy(), confidence_col="calibrated_confidence" - ) - true_fdr_df = true_fdr_df.sort_values("calibrated_confidence") - - fdr_vals = df["psm_fdr"].values - conf_vals = df["calibrated_confidence"].values - n = len(fdr_vals) - hw = _hoeffding_band_arrays(n)[::-1] - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.fill_between( - conf_vals, - np.clip(fdr_vals - hw, 0, None), - np.clip(fdr_vals + hw, None, 1), - color=_BAND_COLOUR, - alpha=0.2, - label="95% Hoeffding bound", - ) - ax.plot( - conf_vals, - fdr_vals, - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Non-parametric", - ) - ax.plot( - true_fdr_df["calibrated_confidence"].values, - true_fdr_df["psm_fdr"].values, - color=_RAW_LINE_COLOUR, - lw=1.5, - label="Database-grounded", - ) - ax.set_xlabel("Calibrated confidence") - ax.set_ylabel("PSM FDR") - ax.set_title(f"{display} FDR run with sampling error bounds {qualifier}") - ax.legend(loc="upper right") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"fdr_run_bands_{project}") - - -def plot_q_value_run_with_bands( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Q-value run plot with Hoeffding 95% confidence band on the non-parametric curve.""" - if "psm_q_value" not in df.columns: - logger.warning( - "Skipping banded q-value run plot for %s: psm_q_value column missing", - project, - ) - return - - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - - sorted_df = df.sort_values("calibrated_confidence") - - true_fdr_ctrl = _fit_database_grounded_fdr(df) - qval_input = df[["calibrated_confidence"]].copy() - true_q_df = true_fdr_ctrl.add_psm_q_value( - qval_input, confidence_col="calibrated_confidence" - ) - true_q_df = true_q_df.sort_values("calibrated_confidence") - - qvals = sorted_df["psm_q_value"].values - conf_vals = sorted_df["calibrated_confidence"].values - n = len(qvals) - hw = _hoeffding_band_arrays(n)[::-1] - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.fill_between( - conf_vals, - np.clip(qvals - hw, 0, None), - np.clip(qvals + hw, None, 1), - color=_BAND_COLOUR, - alpha=0.2, - label="95% Hoeffding bound", - ) - ax.plot( - conf_vals, - qvals, - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Non-parametric", - ) - ax.plot( - true_q_df["calibrated_confidence"].values, - true_q_df["psm_q_value"].values, - color=_RAW_LINE_COLOUR, - lw=1.5, - label="Database-grounded", - ) - ax.set_xlabel("Calibrated confidence") - ax.set_ylabel("PSM q-value") - ax.set_title(f"{display} q-value run with sampling error bounds {qualifier}") - ax.legend(loc="upper center") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"qvalue_run_bands_{project}") - - -# --------------------------------------------------------------------------- -# True FDR vs estimated FDR -# --------------------------------------------------------------------------- -def _compute_true_vs_estimated_fdr(df: pd.DataFrame) -> pd.DataFrame: - """Compute true and estimated FDR arrays, sorted by confidence descending.""" - sorted_df = df.sort_values("calibrated_confidence", ascending=False).reset_index( - drop=True - ) - - true_fdr_ctrl = _fit_database_grounded_fdr(sorted_df) - with_true_fdr = true_fdr_ctrl.add_psm_fdr( - sorted_df, confidence_col="calibrated_confidence" - ) - - return pd.DataFrame( - { - "estimated_fdr": sorted_df["psm_fdr"].values, - "true_fdr": with_true_fdr["psm_fdr"].values, - } - ) - - -def plot_true_vs_estimated_fdr( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, - *, - zoomed: bool = False, -) -> None: - """Plot true FDR vs estimated FDR.""" - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - fdr_data = _compute_true_vs_estimated_fdr(df) - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.plot( - fdr_data["estimated_fdr"], - fdr_data["true_fdr"], - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Observed", - ) - - # Only plot the ideal line up to the max extent of the observed line - max_x = float(fdr_data["estimated_fdr"].max()) - max_y = float(fdr_data["true_fdr"].max()) - lim = 0.1 if zoomed else 1.0 - ideal_end = min(lim, max(max_x, max_y)) - - ax.plot( - [0, ideal_end], - [0, ideal_end], - ls="--", - color=_IDEAL_LINE_COLOUR, - lw=1, - label="Perfectly calibrated", - ) - ax.set_xlabel("Non-parametric estimated FDR") - ax.set_ylabel("Database-grounded FDR") - zoom_suffix = " (0 to 0.1)" if zoomed else "" - ax.set_title(f"{display} true vs estimated FDR{zoom_suffix} {qualifier}") - if zoomed: - ax.set_xlim(0, 0.1) - ax.set_ylim(0, 0.1) - ax.legend(loc="upper left") - _style_ax(ax) - fig.tight_layout() - tag = "fdr_true_vs_est_zoom" if zoomed else "fdr_true_vs_est" - _save_fig(fig, output_dir / f"{tag}_{project}") - - -# --------------------------------------------------------------------------- -# True q-values vs estimated q-values -# --------------------------------------------------------------------------- -def _compute_true_vs_estimated_q_values(df: pd.DataFrame) -> pd.DataFrame: - """Compute true and estimated q-value arrays, sorted by confidence descending.""" - sorted_df = df.sort_values("calibrated_confidence", ascending=False).reset_index( - drop=True - ) - - true_q_val_ctrl = _fit_database_grounded_fdr(sorted_df) - qval_input = sorted_df[["calibrated_confidence"]].copy() - with_true_q_df = true_q_val_ctrl.add_psm_q_value( - qval_input, confidence_col="calibrated_confidence" - ) - return pd.DataFrame( - { - "estimated_q_value": sorted_df["psm_q_value"].values, - "true_q_value": with_true_q_df["psm_q_value"].values, - } - ) - - -def plot_true_vs_estimated_q_values( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, - *, - zoomed: bool = False, -) -> None: - """Plot true q-values vs estimated q-values.""" - if "psm_q_value" not in df.columns: - logger.warning( - "Skipping true vs estimated q-value plot for %s: psm_q_value column missing", - project, - ) - return - - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - q_value_data = _compute_true_vs_estimated_q_values(df) - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.plot( - q_value_data["estimated_q_value"], - q_value_data["true_q_value"], - color=_MAIN_LINE_COLOUR, - lw=1.5, - label="Observed", - ) - - # Only plot the ideal line up to the max extent of the observed line - max_x = float(q_value_data["estimated_q_value"].max()) - max_y = float(q_value_data["true_q_value"].max()) - lim = 0.1 if zoomed else 1.0 - ideal_end = min(lim, max(max_x, max_y)) - - ax.plot( - [0, ideal_end], - [0, ideal_end], - ls="--", - color=_IDEAL_LINE_COLOUR, - lw=1, - label="Perfectly calibrated", - ) - ax.set_xlabel("Non-parametric estimated q-values") - ax.set_ylabel("Database-grounded q-values") - zoom_suffix = " (0 to 0.1)" if zoomed else "" - ax.set_title(f"{display} true vs estimated q-values{zoom_suffix} {qualifier}") - if zoomed: - ax.set_xlim(0, 0.1) - ax.set_ylim(0, 0.1) - ax.legend(loc="upper left") - _style_ax(ax) - fig.tight_layout() - tag = "qvalue_true_vs_est_zoom" if zoomed else "qvalue_true_vs_est" - _save_fig(fig, output_dir / f"{tag}_{project}") - - -# --------------------------------------------------------------------------- -# Probability calibration (reliability diagram) -# --------------------------------------------------------------------------- -def _compute_calibration_curve( - df: pd.DataFrame, - pred_col: str, - label_col: str, - n_bins: int = 10, -) -> pd.DataFrame: - """Fixed-width bin calibration curve.""" - data = df[[pred_col, label_col]].dropna().copy() - data[pred_col] = data[pred_col].clip(0.0, 1.0) - bins = np.linspace(0.0, 1.0, n_bins + 1) - bin_cats = pd.cut(data[pred_col], bins=bins, include_lowest=True) - bin_cats.name = "bin" - grouped = ( - data.groupby(bin_cats, observed=True) - .agg( - pred_mean=(pred_col, "mean"), - empirical=(label_col, "mean"), - count=(label_col, "size"), - ) - .reset_index() - ) - grouped = grouped[grouped["count"] > 0] - grouped["bin_center"] = grouped["bin"].apply(lambda iv: (iv.left + iv.right) / 2) - return grouped[["pred_mean", "empirical", "count", "bin_center"]] - - -def _estimate_calibration_values( - df: pd.DataFrame, - pred_col: str, - label_col: str, - n_bins: int = 20, -) -> np.ndarray: - """Estimate c(s) for each PSM via binned calibration. - - Returns an array of the same length as *df* where each entry is the - empirical accuracy of the bin that PSM falls into. - """ - scores = df[pred_col].values.clip(0.0, 1.0) - bins = np.linspace(0.0, 1.0, n_bins + 1) - bin_idx = np.digitize(scores, bins) - 1 - bin_idx = np.clip(bin_idx, 0, n_bins - 1) - labels = df[label_col].values.astype(float) - bin_sums = np.bincount(bin_idx, weights=labels, minlength=n_bins) - bin_counts = np.bincount(bin_idx, minlength=n_bins).astype(float) - bin_counts[bin_counts == 0] = 1.0 - bin_means = bin_sums / bin_counts - return bin_means[bin_idx] - - -def _hoeffding_halfwidth(k: int, delta: float = _HOEFFDING_DELTA) -> float: - """Hoeffding 95% confidence half-width for a mean of *k* bounded [0,1] r.v.s.""" - if k <= 0: - return float("nan") - return float(np.sqrt(np.log(2.0 / delta) / (2.0 * k))) - - -def plot_calibration( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Plot probability calibration (reliability diagram).""" - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - cal_calibrated = _compute_calibration_curve(df, "calibrated_confidence", "correct") - cal_raw = _compute_calibration_curve(df, "confidence", "correct") - - fig, ax = plt.subplots(figsize=(6, 4)) - ax.plot( - cal_raw["pred_mean"], - cal_raw["empirical"], - marker="D", - color=_RAW_LINE_COLOUR, - label="Raw confidence", - ) - ax.plot( - cal_calibrated["pred_mean"], - cal_calibrated["empirical"], - marker="o", - color=_MAIN_LINE_COLOUR, - label="Calibrated confidence", - ) - ax.plot( - [0, 1], - [0, 1], - ls="--", - color=_IDEAL_LINE_COLOUR, - lw=1, - label="Perfectly calibrated", - ) - ax.set_xlabel("Mean predicted probability") - ax.set_ylabel("Empirical accuracy") - ax.set_title(f"{display} probability calibration {qualifier}") - ax.set_xlim(0, 1) - ax.set_ylim(0, 1) - ax.legend(loc="lower right") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"calibration_{project}") - - -# --------------------------------------------------------------------------- -# Before/after score histograms -# --------------------------------------------------------------------------- -def plot_score_histograms( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Plot before/after score histograms with correct/incorrect overlays.""" - display = _display_name(project) - qualifier = _ground_truth_qualifier(eval_type) - - correct_mask = df["correct"].astype(bool) - - fig, axes = plt.subplots(2, 1, figsize=(7, 6), sharex=False) - - # Before calibration - ax = axes[0] - bins_before = np.linspace(0, 1, 51) - ax.hist( - df.loc[correct_mask, "confidence"], - bins=bins_before, - alpha=0.5, - color=_CORRECT_COLOUR, - edgecolor="black", - label="Correct", - ) - ax.hist( - df.loc[~correct_mask, "confidence"], - bins=bins_before, - alpha=0.5, - color=_INCORRECT_COLOUR, - edgecolor="black", - label="Incorrect", - ) - ax.set_xlabel("Raw confidence") - ax.set_ylabel("Count") - ax.set_title("Before calibration") - ax.legend(loc="upper center") - _style_ax(ax) - - # After calibration - ax = axes[1] - bins_after = np.linspace(0, 1, 51) - ax.hist( - df.loc[correct_mask, "calibrated_confidence"], - bins=bins_after, - alpha=0.5, - color=_CORRECT_COLOUR, - edgecolor="black", - label="Correct", - ) - ax.hist( - df.loc[~correct_mask, "calibrated_confidence"], - bins=bins_after, - alpha=0.5, - color=_INCORRECT_COLOUR, - edgecolor="black", - label="Incorrect", - ) - ax.set_xlabel("Calibrated confidence") - ax.set_ylabel("Count") - ax.set_title("After calibration") - ax.legend(loc="upper center") - _style_ax(ax) - - fig.suptitle(f"{display} score distributions {qualifier}", fontsize=13) - fig.tight_layout() - _save_fig(fig, output_dir / f"score_histograms_{project}") - - -# --------------------------------------------------------------------------- -# Diagnostics CSV -# --------------------------------------------------------------------------- -def _is_labelled(eval_type: str) -> bool: - return eval_type in ("annotated", "labelled") - - -def _compute_diagnostics( - df: pd.DataFrame, - eval_type: str, - alphas: tuple[float, ...] = _DIAGNOSTIC_ALPHAS, -) -> pd.DataFrame: - """Compute FDR diagnostics at each target alpha. - - Label-dependent metrics (sTECE, TECE, realised FDR, etc.) are only - populated for annotated/labelled eval types. - """ - labelled = _is_labelled(eval_type) - - np_ctrl = NonParametricFDRControl() - np_ctrl.fit(dataset=df["calibrated_confidence"]) - - if labelled: - db_ctrl = _fit_database_grounded_fdr(df) - c_hat = _estimate_calibration_values(df, "calibrated_confidence", "correct") - scores = df["calibrated_confidence"].values.clip(0.0, 1.0) - - rows: list[dict] = [] - for alpha in alphas: - tau_hat = np_ctrl.get_confidence_cutoff(threshold=alpha) - if np.isnan(tau_hat): - rows.append({"alpha": alpha, "tau_hat": float("nan")}) - continue - - mask_hat = df["calibrated_confidence"].values >= tau_hat - k = int(mask_hat.sum()) - est_fdr = float(np_ctrl.compute_fdr(tau_hat)) - eps = _hoeffding_halfwidth(k) - - row: dict = { - "alpha": alpha, - "tau_hat": float(tau_hat), - "k_accepted": k, - "estimated_fdr": est_fdr, - "hoeffding_halfwidth": eps, - } - - if labelled: - residuals = c_hat[mask_hat] - scores[mask_hat] - row["stece"] = float(np.mean(residuals)) - row["tece"] = float(np.mean(np.abs(residuals))) - row["tece_2"] = float(np.sqrt(np.mean(residuals**2))) - - realised_fdr = float(db_ctrl.compute_fdr(tau_hat)) - row["realised_fdr"] = realised_fdr - row["fdr_bias"] = est_fdr - realised_fdr - - tau_star = db_ctrl.get_confidence_cutoff(threshold=alpha) - row["tau_star"] = float(tau_star) - if not np.isnan(tau_star): - k_star = int((df["calibrated_confidence"].values >= tau_star).sum()) - row["discovery_count_shift"] = k - k_star - else: - row["discovery_count_shift"] = float("nan") - - rows.append(row) - - return pd.DataFrame(rows) - - -# --------------------------------------------------------------------------- -# Orchestration -# --------------------------------------------------------------------------- -def _load_project_data( - predictions_root: Path, - project: str, - suffix: str, - eval_type: str, -) -> pd.DataFrame: - """Load and merge metadata.csv and preds_and_fdr_metrics.csv for a project.""" - # folder = predictions_root / f"{project}_{suffix}" - folder = predictions_root / f"{project}" - preds_path = folder / "preds_and_fdr_metrics.csv" - meta_path = folder / "metadata.csv" - if not preds_path.is_file(): - raise FileNotFoundError(f"Missing predictions file: {preds_path}") - - preds_df = pd.read_csv(preds_path) - - if meta_path.is_file(): - meta_df = pd.read_csv(meta_path) - # Drop columns already present in preds to avoid duplicates on merge - overlap = [ - c for c in meta_df.columns if c in preds_df.columns and c != "spectrum_id" - ] - if overlap: - meta_df = meta_df.drop(columns=overlap) - df = preds_df.merge(meta_df, on="spectrum_id", how="left") - else: - df = preds_df - - if eval_type in ("raw", "unlabelled"): - if "proteome_hit" not in df.columns: - raise ValueError( - f"Expected 'proteome_hit' column for eval-type={eval_type} in {preds_path}" - ) - df["correct"] = df["proteome_hit"].astype(float) - - required = ["confidence", "calibrated_confidence", "correct"] - missing = [c for c in required if c not in df.columns] - if missing: - raise ValueError(f"Missing columns {missing} in {preds_path}") - - return df - - -def generate_all_plots( - df: pd.DataFrame, - project: str, - eval_type: str, - output_dir: Path, -) -> None: - """Generate all plots for a single project.""" - output_dir.mkdir(parents=True, exist_ok=True) - - plot_precision_recall(df, project, eval_type, output_dir) - plot_fdr_run(df, project, eval_type, output_dir) - plot_fdr_run_with_bands(df, project, eval_type, output_dir) - plot_q_value_run(df, project, eval_type, output_dir) - plot_q_value_run_with_bands(df, project, eval_type, output_dir) - plot_true_vs_estimated_fdr(df, project, eval_type, output_dir, zoomed=False) - plot_true_vs_estimated_fdr(df, project, eval_type, output_dir, zoomed=True) - plot_true_vs_estimated_q_values(df, project, eval_type, output_dir, zoomed=False) - plot_true_vs_estimated_q_values(df, project, eval_type, output_dir, zoomed=True) - plot_calibration(df, project, eval_type, output_dir) - plot_score_histograms(df, project, eval_type, output_dir) - - -_EVAL_TYPE_SUFFIX: dict[str, str] = { - "annotated": "annotated", - "raw": "raw", - "labelled": "labelled", - "unlabelled": "unlabelled", -} - - -@app.command() -def main( - predictions_root: Annotated[ - Path, - typer.Option( - "--predictions-root", - help="Root directory containing per-project prediction folders.", - ), - ], - projects: Annotated[ - str, - typer.Option( - "--projects", - help="Space- or comma-separated project keys (e.g. 'helaqc,gluc' or 'helaqc gluc').", - ), - ], - eval_type: Annotated[ - str, - typer.Option( - "--eval-type", - help="Evaluation type: annotated, raw, labelled, or unlabelled.", - ), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory to save plots and summary CSVs."), - ], -) -> None: - """Generate evaluation plots from winnow predict outputs.""" - logging.basicConfig(level=logging.INFO, format="%(message)s", datefmt="%H:%M:%S") - - if eval_type not in _EVAL_TYPE_SUFFIX: - raise typer.BadParameter( - f"Unknown eval-type {eval_type!r}. Expected one of: {list(_EVAL_TYPE_SUFFIX)}" - ) - - project_list = [p.strip() for p in projects.replace(",", " ").split() if p.strip()] - if not project_list: - raise typer.BadParameter("No projects specified.") - - suffix = _EVAL_TYPE_SUFFIX[eval_type] - output_dir.mkdir(parents=True, exist_ok=True) - - for project in project_list: - display = _display_name(project) - logger.info("Processing %s (%s, eval-type=%s)...", project, display, eval_type) - - df = _load_project_data(predictions_root, project, suffix, eval_type) - logger.info(" Loaded %d rows", len(df)) - - true_fdr_ctrl = _fit_database_grounded_fdr(df) - db_fdr = true_fdr_ctrl.add_psm_fdr( - df[["calibrated_confidence"]].copy(), confidence_col="calibrated_confidence" - ) - df["db_grounded_psm_fdr"] = db_fdr["psm_fdr"] - db_qval = true_fdr_ctrl.add_psm_q_value( - df[["calibrated_confidence"]].copy(), confidence_col="calibrated_confidence" - ) - df["db_grounded_psm_q_value"] = db_qval["psm_q_value"] - - summary_cols = ["confidence", "calibrated_confidence", "correct"] - if "psm_fdr" in df.columns: - summary_cols.append("psm_fdr") - summary_cols.append("db_grounded_psm_fdr") - if "psm_q_value" in df.columns: - summary_cols.append("psm_q_value") - summary_cols.append("db_grounded_psm_q_value") - df[summary_cols].to_csv(output_dir / f"{project}_summary.csv", index=False) - - diag = _compute_diagnostics(df, eval_type) - diag.to_csv(output_dir / f"{project}_diagnostics.csv", index=False) - logger.info(" Diagnostics saved (%d alpha levels)", len(diag)) - - generate_all_plots(df, project, eval_type, output_dir) - logger.info(" Plots saved to %s", output_dir) - - logger.info("Done. All plots saved to %s", output_dir) - - -if __name__ == "__main__": - app() diff --git a/scripts/plot_fdr_method_comparison.py b/scripts/plot_fdr_method_comparison.py deleted file mode 100644 index 477738e7..00000000 --- a/scripts/plot_fdr_method_comparison.py +++ /dev/null @@ -1,1106 +0,0 @@ -#!/usr/bin/env python3 -"""Compare PSM-level FDR estimates from Winnow and NovoBoard. - -Plots and summaries cover labelled-test Novor correctness and unlabelled -reference-proteome membership at 1 %, 5 %, and 10 % FDR on a shared filtered -spectrum pool (NovoBoard mass-deltas converted to ProForma; unsupported -modifications dropped; NovoBoard target-decoy pairs gated). Unlabelled panels -also drop normalised peptides shorter than 8 residues (proteome-substring -proxy); labelled panels keep short peptides because correctness is Novor -agreement. The shared pool is the twin-valid NovoBoard set; Winnow is trimmed -to match under the invariant that NovoBoard ⊆ Winnow after identical InstaNovo -filters. - -A long-form ``fdr_method_comparison_curves.csv`` (per spectrum x method) is -written so plots and summary tables can be regenerated with ``--summarise-only``. - -External NovoBoard inputs (``--novoboard-root``) must follow -``{root}/{dataset}/novoboard/`` with target/decoy CSVs such as -``annotated_test.csv``, ``annotated_test_decoy_{rate}.csv``, -``raw_unlabelled.csv`` and ``raw_unlabelled_decoy_{rate}.csv``. Point -``--novoboard-root`` at the ``datasets`` directory of a NovoBoard checkout. -Local results were produced from the fork -``git@github.com:JemmaLDaniel/NovoBoard.git``, branch -``feat/adapt-to-instanovo`` at commit -``a9faab3ef1af06987599c2f01e6ba96072c80172``. -""" - -from __future__ import annotations - -import logging -import sys -from dataclasses import dataclass -from pathlib import Path -from typing import Annotated, Literal, Optional - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import polars as pl -import seaborn as sns -import typer - -_REPO_ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(_REPO_ROOT)) - -from winnow.utils.proteome import load_proteome_haystack # noqa: E402 -from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 - LABELLED_MIN_PEPTIDE_LENGTH, - MIN_PEPTIDE_LENGTH, - assert_shared_prediction_keys, - attach_labels_by_spectrum_id, - attach_novoboard_pair_keys, - compute_q_values, - filter_novoboard_target_decoy_pairs, - filter_prediction_table, - label_series_by_spectrum_id, - load_residue_masses, - novoboard_psm_tdc, - novor_correctness_mask, - proteome_hit_mask, - restrict_winnow_to_novoboard_spectra, -) -from scripts.fdr_tool_comparison_summaries import ( # noqa: E402 - SUMMARY_THRESHOLDS, - acceptance_rows_from_q, - error_rows_from_q, - finalise_error_gain_table, - write_summary_tables, -) -from scripts.plot_eval_results import ( # noqa: E402 - _MAIN_LINE_COLOUR, - _PALETTE, - _RAW_LINE_COLOUR, - _display_name, - _ground_truth_qualifier, - _save_fig, - _style_ax, -) -from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 -from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 - -PRIMARY_METHOD = "Winnow (non-parametric)" -DB_CAL_METHOD = "Database-grounded (calibrated confidence)" -DB_RAW_METHOD = "Database-grounded (raw confidence)" -NOVOBOARD_METHOD = "NovoBoard" -CURVES_CSV_NAME = "fdr_method_comparison_curves.csv" -_WINNOW_METHODS = (PRIMARY_METHOD, DB_CAL_METHOD, DB_RAW_METHOD) -_METHOD_ORDER = (*_WINNOW_METHODS, NOVOBOARD_METHOD) - -logger = logging.getLogger(__name__) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -FDR_THRESHOLDS = [0.01, 0.05, 0.10] -_DB_GROUNDED_DROP = 10 - -DEFAULT_WINNOW_RESULTS = _REPO_ROOT / "results" -DEFAULT_MODEL_ROOT = _REPO_ROOT / "models" -DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/fdr_method_comparison_psm" -DEFAULT_DATASETS = ["helaqc", "celegans"] -_METHOD_COLOURS = { - PRIMARY_METHOD: _MAIN_LINE_COLOUR, - DB_CAL_METHOD: _RAW_LINE_COLOUR, - DB_RAW_METHOD: _PALETTE[3], - NOVOBOARD_METHOD: _PALETTE[2], -} - -EvalType = Literal["labelled", "unlabelled"] - -_DATASET_META = { - "helaqc": { - "fasta": "fasta/human.fasta", - "novoboard_decoy": "0.50", - "winnow_suffix": "helaqc", - }, - "celegans": { - "fasta": "fasta/celegans.fasta", - "novoboard_decoy": "0.70", - "winnow_suffix": "celegans", - }, - "sbrodae": { - "fasta": "fasta/Sb_proteome.fasta", - "novoboard_decoy": "0.50", - "winnow_suffix": "sbrodae", - }, - "PXD019483": { - "fasta": "fasta/human.fasta", - "novoboard_decoy": "0.70", - "winnow_suffix": "pxd019483", - }, -} - - -@dataclass(frozen=True) -class DatasetConfig: - """Paths and metadata for one evaluation dataset.""" - - key: str - fasta: Path - winnow_unlabelled: Path - winnow_test: Path - novoboard_dir: Path - novoboard_decoy_rate: str - calibrator_train_metadata: Path - - -def build_dataset_configs( - winnow_results: Path = DEFAULT_WINNOW_RESULTS, - *, - novoboard_root: Path, - model_root: Path = DEFAULT_MODEL_ROOT, -) -> dict[str, DatasetConfig]: - """Build per-dataset path bundles from repo roots.""" - configs: dict[str, DatasetConfig] = {} - for key, meta in _DATASET_META.items(): - suffix = meta["winnow_suffix"] - configs[key] = DatasetConfig( - key=key, - fasta=_REPO_ROOT / meta["fasta"], - winnow_unlabelled=winnow_results - / f"instanovo_{suffix}_predictions_unlabelled", - winnow_test=winnow_results / f"instanovo_{suffix}_predictions_test", - novoboard_dir=novoboard_root / f"{key}/novoboard", - novoboard_decoy_rate=meta["novoboard_decoy"], - calibrator_train_metadata=model_root - / f"instanovo_{suffix}/metadata_train.parquet", - ) - return configs - - -@dataclass -class MethodCurve: - """One method's confidence and q-value arrays for curve plotting.""" - - label: str - color: str - confidence: np.ndarray - q_value: np.ndarray - - -@dataclass -class MethodCounts: - """PSM counts per q-value threshold for one method.""" - - label: str - color: str - counts: list[int] - - -@dataclass -class MethodRecovery: - """Correct labelled identifications recovered at q-value thresholds.""" - - label: str - color: str - q_value: np.ndarray - correct: np.ndarray - - -def _load_residue_masses() -> dict[str, float]: - return load_residue_masses() - - -def _fit_database_grounded_fdr( - df: pd.DataFrame, - correct_col: str, - confidence_col: str, - *, - drop: int = _DB_GROUNDED_DROP, -) -> DatabaseGroundedFDRControl: - """Fit ``DatabaseGroundedFDRControl`` from per-row correctness labels.""" - ctrl = DatabaseGroundedFDRControl( - confidence_feature=confidence_col, - drop=drop, - ) - sorted_df = df.sort_values(confidence_col, ascending=False) - labels = sorted_df[correct_col].astype(float).to_numpy() - conf = sorted_df[confidence_col].to_numpy() - precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) - ctrl._fdr_values = np.array(1.0 - precision)[drop:] - ctrl._confidence_scores = conf[drop:] - return ctrl - - -def load_winnow( - predictions_dir: Path, fasta: Path, eval_type: EvalType -) -> pd.DataFrame: - """Load Winnow preds + metadata; annotate proteome hits or use labelled ``correct``. - - Always drops unsupported ``[UNIMOD:n]`` tokens. Unlabelled panels also require - normalised peptide length ≥ :data:`MIN_PEPTIDE_LENGTH` (proteome-substring - proxy). Labelled panels only require a non-empty normalised key - (:data:`LABELLED_MIN_PEPTIDE_LENGTH`), because correctness is Novor agreement. - - Labelled ``correct`` keeps Winnow predict-time Novor labels when present; - otherwise recomputes Novor from ``sequence`` / ``prediction``. - """ - preds = pl.read_csv(predictions_dir / "preds_and_fdr_metrics.csv") - meta_path = predictions_dir / "metadata.csv" - if meta_path.exists(): - meta = pl.read_csv(meta_path, columns=["spectrum_id", "confidence"]) - preds = preds.join(meta, on="spectrum_id", how="inner") - - df = preds.to_pandas() - min_length = ( - LABELLED_MIN_PEPTIDE_LENGTH if eval_type == "labelled" else MIN_PEPTIDE_LENGTH - ) - df = filter_prediction_table( - df, "prediction", min_length=min_length, key_col="peptide_key" - ) - - if eval_type == "labelled": - if "correct" in df.columns: - pass - elif {"sequence", "prediction"}.issubset(df.columns): - df["correct"] = novor_correctness_mask(df["sequence"], df["prediction"]) - else: - raise ValueError( - f"Missing 'correct' (and sequence/prediction) in " - f"{predictions_dir}/preds_and_fdr_metrics.csv" - ) - return df - - haystack = load_proteome_haystack(fasta) - df["proteome_hit"] = proteome_hit_mask( - df["prediction"], haystack, min_length=MIN_PEPTIDE_LENGTH - ) - return df - - -def _effective_db_grounded_drop(n_rows: int, drop: int = _DB_GROUNDED_DROP) -> int: - """Cap drop so FDR fit retains at least one score when *n_rows* is small.""" - return min(drop, max(0, n_rows - 1)) - - -def _vectorized_fdr_from_control( - confidence: np.ndarray, ctrl: DatabaseGroundedFDRControl | NonParametricFDRControl -) -> np.ndarray: - """Vectorized equivalent of ``FDRControl.compute_fdr`` for an array of scores.""" - if ctrl._confidence_scores is None or ctrl._fdr_values is None: - raise AttributeError("FDR method not fitted, please call `fit()` first") - conf = np.asarray(confidence, dtype=float) - scores = np.asarray(ctrl._confidence_scores, dtype=float) - fdr_values = np.asarray(ctrl._fdr_values, dtype=float) - n = len(scores) - idx = np.searchsorted(-scores, -conf, side="left") - fdr = np.empty(len(conf), dtype=float) - below = (idx == n) & (conf < scores[-1]) - above = (idx == 0) & (conf > scores[0]) - normal = ~(below | above) - fdr[below] = 1.0 - fdr[above] = float(fdr_values[0]) - clipped = np.clip(idx[normal], 0, n - 1) - fdr[normal] = fdr_values[clipped] - return fdr - - -def _assign_q_values_fast( - df: pd.DataFrame, - confidence_col: str, - ctrl: DatabaseGroundedFDRControl | NonParametricFDRControl, - out_col: str, -) -> pd.DataFrame: - """Assign q-values without per-row ``compute_fdr`` applies (needed for large tables).""" - work = df.copy() - conf = work[confidence_col].to_numpy(dtype=float) - fdr = _vectorized_fdr_from_control(conf, ctrl) - order = np.argsort(-conf, kind="mergesort") - q_sorted = compute_q_values(fdr[order]) - q = np.empty_like(q_sorted) - q[order] = q_sorted - work[out_col] = q - return work - - -def _add_database_grounded_qvalues( - df: pd.DataFrame, - correct_col: str, - confidence_col: str, - out_col: str, - residue_masses: dict[str, float], - *, - fit_df: pd.DataFrame | None = None, - drop: int = _DB_GROUNDED_DROP, -) -> pd.DataFrame: - """Append database-grounded PSM q-values; fit on *fit_df* (defaults to *df*).""" - reference = fit_df if fit_df is not None else df - work = df.drop(columns=[out_col], errors="ignore").copy() - ctrl = _fit_database_grounded_fdr( - reference, - correct_col, - confidence_col, - drop=_effective_db_grounded_drop(len(reference), drop), - ) - return _assign_q_values_fast(work, confidence_col, ctrl, out_col) - - -def _prepare_winnow_psm_table( - df: pd.DataFrame, - correct_col: str, - residue_masses: dict[str, float], - *, - fit_df: pd.DataFrame | None = None, -) -> pd.DataFrame: - """Append Winnow PSM-level q-value columns while retaining labels.""" - reference = fit_df if fit_df is not None else df - db_cal = _add_database_grounded_qvalues( - df, - correct_col, - "calibrated_confidence", - "psm_q_value_db_cal", - residue_masses, - fit_df=reference, - ) - db_raw = _add_database_grounded_qvalues( - db_cal, - correct_col, - "confidence", - "psm_q_value_db_raw", - residue_masses, - fit_df=reference, - ) - return db_raw - - -def _curves_df_from_winnow_table( - table: pd.DataFrame, - *, - dataset: str, - panel: str, - label_col: str, -) -> pd.DataFrame: - """Long-form curve rows for the three Winnow PSM q-value methods.""" - if "spectrum_id" not in table.columns: - raise KeyError("Winnow curve export requires spectrum_id") - if label_col not in table.columns: - raise KeyError(f"Missing label column {label_col!r}") - label = table[label_col].astype(bool).to_numpy() - spectrum_id = table["spectrum_id"].astype(str) - specs = ( - (PRIMARY_METHOD, "calibrated_confidence", "psm_q_value"), - (DB_CAL_METHOD, "calibrated_confidence", "psm_q_value_db_cal"), - (DB_RAW_METHOD, "confidence", "psm_q_value_db_raw"), - ) - parts: list[pd.DataFrame] = [] - for method, score_col, q_col in specs: - if score_col not in table.columns or q_col not in table.columns: - raise KeyError(f"Missing {score_col!r} / {q_col!r} for {method}") - parts.append( - pd.DataFrame( - { - "dataset": dataset, - "panel": panel, - "method": method, - "spectrum_id": spectrum_id, - "score": table[score_col].to_numpy(dtype=float), - "q_value": table[q_col].to_numpy(dtype=float), - "label": label, - } - ) - ) - return pd.concat(parts, ignore_index=True) - - -def _curves_df_from_novoboard( - df: pd.DataFrame, - *, - dataset: str, - panel: str, - label_col: str, -) -> pd.DataFrame: - """Long-form curve rows for NovoBoard PSM TDC targets.""" - if "spectrum_id" not in df.columns: - raise KeyError("NovoBoard curve export requires spectrum_id") - if label_col not in df.columns: - raise KeyError(f"Missing label column {label_col!r}") - return pd.DataFrame( - { - "dataset": dataset, - "panel": panel, - "method": NOVOBOARD_METHOD, - "spectrum_id": df["spectrum_id"].astype(str), - "score": df["ALC (%)"].to_numpy(dtype=float), - "q_value": df["estimated_q_value"].to_numpy(dtype=float), - "label": df[label_col].astype(bool).to_numpy(), - } - ) - - -def load_novoboard_target_decoy( - novoboard_dir: Path, split: Literal["unlabelled", "test"], decoy_rate: str -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Load NovoBoard target/decoy tables and attach twin ``_pair_key`` values.""" - prefix = "raw_unlabelled" if split == "unlabelled" else "annotated_test" - target_path = novoboard_dir / f"{prefix}.csv" - decoy_path = novoboard_dir / f"{prefix}_decoy_{decoy_rate}.csv" - if not target_path.is_file(): - raise FileNotFoundError(target_path) - if not decoy_path.is_file(): - raise FileNotFoundError(decoy_path) - target = pd.read_csv(target_path) - decoy = pd.read_csv(decoy_path) - return attach_novoboard_pair_keys( - target, decoy, novoboard_dir=novoboard_dir, split_prefix=prefix - ) - - -def _restrict_winnow_to_novoboard_spectra( - winnow: pd.DataFrame, novoboard: pd.DataFrame -) -> pd.DataFrame: - """Trim Winnow to NovoBoard twin-valid spectra under the subset invariant.""" - return restrict_winnow_to_novoboard_spectra(winnow, novoboard) - - -def _assert_shared_prediction_keys( - winnow: pd.DataFrame, - novoboard: pd.DataFrame, - *, - winnow_peptide_col: str = "prediction", - novoboard_peptide_col: str = "Peptide", -) -> None: - """Require I/L-normalised prediction identity on the shared spectrum pool.""" - assert_shared_prediction_keys( - winnow, - novoboard, - winnow_peptide_col=winnow_peptide_col, - novoboard_peptide_col=novoboard_peptide_col, - ) - - -def _label_series_by_spectrum_id(winnow: pd.DataFrame, label_col: str) -> pd.Series: - """Map ``spectrum_id`` → boolean label from a Winnow table.""" - return label_series_by_spectrum_id(winnow, label_col) - - -def _attach_labels_by_spectrum_id( - novoboard: pd.DataFrame, - label_by_id: pd.Series, - *, - label_col: str, -) -> pd.DataFrame: - """Attach a shared label column to NovoBoard rows by ``spectrum_id``.""" - return attach_labels_by_spectrum_id(novoboard, label_by_id, label_col=label_col) - - -def _method_curves_from_panel(panel_df: pd.DataFrame) -> list[MethodCurve]: - """Rebuild plot curves from long-form curve rows for one panel.""" - curves: list[MethodCurve] = [] - for method in _METHOD_ORDER: - sub = panel_df.loc[panel_df["method"] == method] - if sub.empty: - continue - colour = _METHOD_COLOURS.get(str(method), _PALETTE[0]) - curves.append( - MethodCurve( - str(method), - colour, - sub["score"].to_numpy(dtype=float), - sub["q_value"].to_numpy(dtype=float), - ) - ) - return curves - - -def _recovery_series_from_panel(panel_df: pd.DataFrame) -> list[MethodRecovery]: - """Rebuild labelled recovery series from long-form curve rows.""" - series: list[MethodRecovery] = [] - for method in _METHOD_ORDER: - sub = panel_df.loc[panel_df["method"] == method] - if sub.empty: - continue - colour = _METHOD_COLOURS.get(str(method), _PALETTE[0]) - series.append( - MethodRecovery( - str(method), - colour, - sub["q_value"].to_numpy(dtype=float), - sub["label"].astype(bool).to_numpy(), - ) - ) - return series - - -def plot_dataset_from_curves( - curves: pd.DataFrame, dataset_key: str, output_dir: Path -) -> None: - """Write PSM comparison plots for one dataset from a curves table.""" - out = output_dir / dataset_key - out.mkdir(parents=True, exist_ok=True) - ds = curves.loc[curves["dataset"] == dataset_key] - if ds.empty: - raise ValueError(f"No curve rows for dataset {dataset_key!r}") - - panel_specs: tuple[tuple[str, EvalType, str], ...] = ( - ("unlabelled", "unlabelled", "unlabelled"), - ("labelled_test", "labelled", "test"), - ) - for panel, eval_type, stem in panel_specs: - panel_df = ds.loc[ds["panel"] == panel] - method_curves = _method_curves_from_panel(panel_df) - if not method_curves: - continue - plot_qvalue_by_rank( - method_curves, - dataset_key, - eval_type, - out / f"psm_qvalue_by_rank_{stem}_{dataset_key}", - ) - plot_threshold_barplot( - _bar_series_from_curves(method_curves), - dataset_key, - eval_type, - out / f"psm_counts_{stem}_{dataset_key}", - ) - - labelled = ds.loc[ds["panel"] == "labelled_test"] - recovery = _recovery_series_from_panel(labelled) - if recovery: - plot_recovery_curves( - recovery, - dataset_key, - out / f"psm_recovery_test_{dataset_key}", - ) - - -def write_curves_csv(curves: pd.DataFrame, output_dir: Path) -> Path: - """Write the long-form replot curves table.""" - output_dir.mkdir(parents=True, exist_ok=True) - path = output_dir / CURVES_CSV_NAME - # Preserve float64 q/score values so threshold edge cases survive round-trip. - curves.to_csv(path, index=False, float_format="%.17g") - logger.info("Wrote %s (%d rows)", path, len(curves)) - return path - - -def plot_qvalue_by_rank( - curves: list[MethodCurve], - dataset_key: str, - eval_type: EvalType, - output_path: Path, - *, - title_suffix: str = "", -) -> None: - """Plot q-value against native-score rank/accepted count.""" - display = _display_name(dataset_key) - qualifier = _ground_truth_qualifier( - "labelled" if eval_type == "labelled" else "unlabelled" - ) - title = f"{display} PSM q-value by rank {qualifier}{title_suffix}" - - fig, ax = plt.subplots(figsize=(8, 6)) - q_max = 0.0 - for curve in curves: - order = np.argsort(-np.asarray(curve.confidence, dtype=float)) - y = np.asarray(curve.q_value, dtype=float)[order] - rank = np.arange(1, len(y) + 1) - valid = ~np.isnan(y) - if not np.any(valid): - continue - q_max = max(q_max, float(np.nanmax(y[valid]))) - ax.plot( - rank[valid], - y[valid], - color=curve.color, - lw=1.5, - label=curve.label, - ) - - ax.set_xlabel("Accepted PSMs by native-score rank") - ax.set_ylabel("PSM q-value") - ax.set_title(title) - y_top = min(max(q_max * 1.15, 0.05), 1.0) - ax.set_ylim(0, y_top) - ax.legend(loc="upper left") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_path) - logger.info("Wrote %s", output_path) - - -def _count_at_thresholds( - q: np.ndarray, thresholds: list[float] = FDR_THRESHOLDS -) -> list[int]: - q = np.asarray(q, dtype=float) - valid = q[~np.isnan(q)] - return [int((valid <= t).sum()) for t in thresholds] - - -def _bar_series_from_curves(curves: list[MethodCurve]) -> list[MethodCounts]: - return [ - MethodCounts( - label=c.label, - color=c.color, - counts=_count_at_thresholds(c.q_value), - ) - for c in curves - ] - - -def _recovery_at_thresholds( - q: np.ndarray, - correct: np.ndarray, - thresholds: list[float] = FDR_THRESHOLDS, -) -> list[float]: - """Return correct-identification recovery percentage at each q-value threshold.""" - q = np.asarray(q, dtype=float) - correct = np.asarray(correct, dtype=bool) - denom = int(correct.sum()) - if denom == 0: - return [np.nan for _ in thresholds] - valid = ~np.isnan(q) - return [100.0 * int((valid & correct & (q <= t)).sum()) / denom for t in thresholds] - - -def plot_recovery_curves( - series: list[MethodRecovery], - dataset_key: str, - output_path: Path, -) -> None: - """Plot correct-identification recovery versus q-value threshold.""" - display = _display_name(dataset_key) - fig, ax = plt.subplots(figsize=(8, 6)) - for item in series: - y = _recovery_at_thresholds(item.q_value, item.correct) - ax.plot( - FDR_THRESHOLDS, - y, - marker="o", - lw=1.5, - label=item.label, - color=item.color, - ) - - ax.set_xlim(0, max(FDR_THRESHOLDS)) - ax.set_ylim(0, 100) - ax.set_xlabel("Estimated q-value threshold") - ax.set_ylabel("Correct PSM recovery\n(% of labelled correct PSMs)") - ax.set_title(f"{display} labelled PSM recovery by q-value threshold") - ax.legend(loc="upper left") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_path) - logger.info("Wrote %s", output_path) - - -def plot_threshold_barplot( - series: list[MethodCounts], - dataset_key: str, - eval_type: EvalType, - output_path: Path, -) -> None: - """Bar chart of identifications retained at each q-value threshold.""" - display = _display_name(dataset_key) - if eval_type == "labelled": - split_label = "labelled test set" - else: - split_label = "unlabelled set" - title = f"{display}: accepted PSMs on the {split_label} at q-value thresholds" - ylabel = "Peptide-spectrum matches" - - n_methods = len(series) - n_thresh = len(FDR_THRESHOLDS) - group_spacing = 0.825 - cluster_width = min(0.75, group_spacing * 0.92) - width = cluster_width / n_methods - x = np.arange(n_thresh) * group_spacing - - fig_w = max(10.0, 2.2 * n_thresh * group_spacing) - fig, ax = plt.subplots(figsize=(fig_w, 7)) - for i, item in enumerate(series): - offset = (i - (n_methods - 1) / 2) * width - bars = ax.bar( - x + offset, - item.counts, - width, - label=item.label, - color=item.color, - edgecolor="black", - linewidth=1, - ) - for bar in bars: - h = bar.get_height() - ax.annotate( - f"{int(h):,}", - xy=(bar.get_x() + bar.get_width() / 2, h), - xytext=(0, 3), - textcoords="offset points", - ha="center", - va="bottom", - fontsize=9, - ) - - max_count = max((c for item in series for c in item.counts), default=1) - y_headroom = (1.55 + 0.06 * n_methods) * (2 / 3) - ax.set_ylim(0, max_count * y_headroom) - - half_cluster = cluster_width / 2 - ax.set_xlim( - -half_cluster - 0.25, - (n_thresh - 1) * group_spacing + half_cluster + 0.25, - ) - - ax.set_xlabel("Q-value threshold") - ax.set_ylabel(ylabel) - ax.set_title(title) - ax.set_xticks(x) - ax.set_xticklabels([str(t) for t in FDR_THRESHOLDS]) - ax.legend(loc="upper left") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, output_path) - logger.info("Wrote %s", output_path) - - -def _append_method_summary_rows( - acceptance_rows: list[dict[str, object]], - error_rows: list[dict[str, object]], - *, - dataset: str, - panel: str, - level: str, - method: str, - q_value: np.ndarray, - label_mask: np.ndarray | None = None, - recovery_denom: int | None = None, - q_ref: np.ndarray | None = None, -) -> None: - """Append acceptance and error rows for one method.""" - acceptance_rows.extend( - acceptance_rows_from_q( - dataset=dataset, - panel=panel, - level=level, - method=method, - q_value=q_value, - thresholds=SUMMARY_THRESHOLDS, - label_mask=label_mask, - recovery_denom=recovery_denom, - ) - ) - error_rows.extend( - error_rows_from_q( - dataset=dataset, - panel=panel, - level=level, - method=method, - q_value=q_value, - thresholds=SUMMARY_THRESHOLDS, - label_mask=label_mask, - q_ref=q_ref, - ) - ) - - -def summary_rows_from_curves( - curves: pd.DataFrame, -) -> tuple[list[dict[str, object]], list[dict[str, object]]]: - """Build acceptance and error summary rows from a long-form curves table.""" - acceptance_rows: list[dict[str, object]] = [] - error_rows: list[dict[str, object]] = [] - required = {"dataset", "panel", "method", "spectrum_id", "q_value", "label"} - missing = required - set(curves.columns) - if missing: - raise ValueError(f"Curves table missing columns: {sorted(missing)}") - - for (dataset, panel), group in curves.groupby(["dataset", "panel"], sort=False): - cal = group.loc[group["method"] == DB_CAL_METHOD, ["spectrum_id", "q_value"]] - q_ref_cal_by_id = cal.drop_duplicates("spectrum_id").set_index("spectrum_id")[ - "q_value" - ] - raw = group.loc[group["method"] == DB_RAW_METHOD, ["spectrum_id", "q_value"]] - q_ref_raw_by_id = raw.drop_duplicates("spectrum_id").set_index("spectrum_id")[ - "q_value" - ] - for method, mdf in group.groupby("method", sort=False): - method_s = str(method) - q_value = mdf["q_value"].to_numpy(dtype=float) - label_mask = mdf["label"].astype(bool).to_numpy() - recovery_denom = int(label_mask.sum()) - q_ref: np.ndarray | None = None - # Winnow-family methods: deviation vs calibrated-confidence DBG. - # NovoBoard / Glissade: deviation vs raw-confidence DBG. - if method_s in _WINNOW_METHODS: - q_ref = ( - mdf["spectrum_id"] - .astype(str) - .map(q_ref_cal_by_id) - .to_numpy(dtype=float) - ) - elif method_s in (NOVOBOARD_METHOD, "Glissade"): - q_ref = ( - mdf["spectrum_id"] - .astype(str) - .map(q_ref_raw_by_id) - .to_numpy(dtype=float) - ) - _append_method_summary_rows( - acceptance_rows, - error_rows, - dataset=str(dataset), - panel=str(panel), - level="psm", - method=method_s, - q_value=q_value, - label_mask=label_mask, - recovery_denom=recovery_denom, - q_ref=q_ref, - ) - return acceptance_rows, error_rows - - -def _comparators_for_panel(panel: str, level: str) -> list[str]: - """Comparator method labels present in a given summary panel.""" - del panel, level - return ["NovoBoard"] - - -def _finalise_method_comparison_tables( - acceptance_rows: list[dict[str, object]], - error_rows: list[dict[str, object]], -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Build acceptance and error/gain DataFrames with per-panel relative columns.""" - acceptance = pd.DataFrame(acceptance_rows) - error = pd.DataFrame(error_rows) - if acceptance.empty: - return acceptance, error - - gain_parts: list[pd.DataFrame] = [] - group_cols = ["dataset", "panel", "level"] - for keys, acc_group in acceptance.groupby(group_cols, sort=False): - if not isinstance(keys, tuple): - keys = (keys,) - _, panel, level = keys - err_mask = True - for col, val in zip(group_cols, keys): - err_mask = err_mask & (error[col] == val) - err_group = error.loc[err_mask] - gain_parts.append( - finalise_error_gain_table( - acc_group, - err_group, - primary_method=PRIMARY_METHOD, - comparators=_comparators_for_panel(str(panel), str(level)), - ) - ) - error_gain = pd.concat(gain_parts, ignore_index=True) if gain_parts else error - return acceptance, error_gain - - -def process_dataset(cfg: DatasetConfig, output_dir: Path) -> pd.DataFrame: - """Generate comparison plots and return long-form curve rows for one dataset.""" - residue_masses = _load_residue_masses() - - winnow_unlabelled = load_winnow(cfg.winnow_unlabelled, cfg.fasta, "unlabelled") - winnow_test = load_winnow(cfg.winnow_test, cfg.fasta, "labelled") - nb_u_target, nb_u_decoy = load_novoboard_target_decoy( - cfg.novoboard_dir, "unlabelled", cfg.novoboard_decoy_rate - ) - nb_t_target, nb_t_decoy = load_novoboard_target_decoy( - cfg.novoboard_dir, "test", cfg.novoboard_decoy_rate - ) - nb_u_target, nb_u_decoy = filter_novoboard_target_decoy_pairs( - nb_u_target, nb_u_decoy, min_length=MIN_PEPTIDE_LENGTH - ) - nb_t_target, nb_t_decoy = filter_novoboard_target_decoy_pairs( - nb_t_target, nb_t_decoy, min_length=LABELLED_MIN_PEPTIDE_LENGTH - ) - # Pair-gated NovoBoard ⊆ Winnow after identical InstaNovo filters; only trim Winnow. - winnow_unlabelled = _restrict_winnow_to_novoboard_spectra( - winnow_unlabelled, nb_u_target - ) - winnow_test = _restrict_winnow_to_novoboard_spectra(winnow_test, nb_t_target) - _assert_shared_prediction_keys(winnow_test, nb_t_target) - _assert_shared_prediction_keys(winnow_unlabelled, nb_u_target) - - # Shared labels once on Winnow; NovoBoard reuses them by spectrum_id. - winnow_test = winnow_test.copy() - winnow_test["correct"] = novor_correctness_mask( - winnow_test["sequence"], - winnow_test["prediction"], - residue_masses=residue_masses, - ) - if "proteome_hit" not in winnow_unlabelled.columns: - raise KeyError("Expected proteome_hit on unlabelled Winnow table") - correct_by_id = _label_series_by_spectrum_id(winnow_test, "correct") - hit_by_id = _label_series_by_spectrum_id(winnow_unlabelled, "proteome_hit") - - novoboard_unlabelled = novoboard_psm_tdc( - nb_u_target, nb_u_decoy, min_length=MIN_PEPTIDE_LENGTH - ) - n_u_tgt = int(novoboard_unlabelled["is_target"].sum()) - n_u_dec = int((~novoboard_unlabelled["is_target"]).sum()) - if n_u_tgt != n_u_dec: - raise AssertionError( - f"Unlabelled PSM TDC unbalanced: targets={n_u_tgt} decoys={n_u_dec}" - ) - novoboard_unlabelled = novoboard_unlabelled[ - novoboard_unlabelled["is_target"] - ].copy() - novoboard_test = novoboard_psm_tdc( - nb_t_target, nb_t_decoy, min_length=LABELLED_MIN_PEPTIDE_LENGTH - ) - n_t_tgt = int(novoboard_test["is_target"].sum()) - n_t_dec = int((~novoboard_test["is_target"]).sum()) - if n_t_tgt != n_t_dec: - raise AssertionError( - f"Labelled PSM TDC unbalanced: targets={n_t_tgt} decoys={n_t_dec}" - ) - novoboard_test = novoboard_test[novoboard_test["is_target"]].copy() - novoboard_test = _attach_labels_by_spectrum_id( - novoboard_test, correct_by_id, label_col="correct" - ) - novoboard_unlabelled = _attach_labels_by_spectrum_id( - novoboard_unlabelled, hit_by_id, label_col="proteome_hit" - ) - - n_w_correct = int(winnow_test["correct"].sum()) - n_nb_correct = int(novoboard_test["correct"].sum()) - if n_w_correct != n_nb_correct: - raise AssertionError( - f"Shared labelled correct counts disagree: Winnow={n_w_correct} " - f"NovoBoard={n_nb_correct}" - ) - n_w_hit = int(winnow_unlabelled["proteome_hit"].sum()) - n_nb_hit = int(novoboard_unlabelled["proteome_hit"].sum()) - if n_w_hit != n_nb_hit: - raise AssertionError( - f"Shared proteome-hit counts disagree: Winnow={n_w_hit} NovoBoard={n_nb_hit}" - ) - - logger.info( - "%s shared PSM pools: unlabelled=%d labelled=%d " - "(NovoBoard twin-valid; Winnow trimmed; shared labels correct=%d hits=%d)", - cfg.key, - len(winnow_unlabelled), - len(winnow_test), - n_w_correct, - n_w_hit, - ) - - winnow_u_psm_table = _prepare_winnow_psm_table( - winnow_unlabelled, "proteome_hit", residue_masses - ) - winnow_t_psm_table = _prepare_winnow_psm_table( - winnow_test, "correct", residue_masses - ) - curves = pd.concat( - [ - _curves_df_from_winnow_table( - winnow_t_psm_table, - dataset=cfg.key, - panel="labelled_test", - label_col="correct", - ), - _curves_df_from_novoboard( - novoboard_test, - dataset=cfg.key, - panel="labelled_test", - label_col="correct", - ), - _curves_df_from_winnow_table( - winnow_u_psm_table, - dataset=cfg.key, - panel="unlabelled", - label_col="proteome_hit", - ), - _curves_df_from_novoboard( - novoboard_unlabelled, - dataset=cfg.key, - panel="unlabelled", - label_col="proteome_hit", - ), - ], - ignore_index=True, - ) - plot_dataset_from_curves(curves, cfg.key, output_dir) - return curves - - -@app.command() -def main( - novoboard_root: Annotated[ - Path, - typer.Option( - "--novoboard-root", - help=( - "Root of NovoBoard per-dataset tables: " - "{root}/{dataset}/novoboard/ with annotated_test*.csv and " - "raw_unlabelled*.csv target/decoy pairs (the datasets/ dir of " - "a NovoBoard checkout). Local runs used fork " - "JemmaLDaniel/NovoBoard, branch feat/adapt-to-instanovo " - "(commit a9faab3ef1af06987599c2f01e6ba96072c80172)." - ), - ), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for PNG/PDF outputs."), - ] = DEFAULT_OUTPUT_DIR, - datasets: Annotated[ - Optional[list[str]], - typer.Option("--datasets", help="Dataset keys to plot."), - ] = None, - winnow_results: Annotated[ - Path, - typer.Option("--winnow-results", help="Winnow results directory."), - ] = DEFAULT_WINNOW_RESULTS, - summarise_only: Annotated[ - Optional[Path], - typer.Option( - "--summarise-only", - help="Only write plots and summary CSVs from an existing curves CSV.", - ), - ] = None, -) -> None: - """Generate FDR method comparison plots and summary CSVs.""" - logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") - output_dir.mkdir(parents=True, exist_ok=True) - - if summarise_only is not None: - curves = pd.read_csv(summarise_only, float_precision="round_trip") - if "label" in curves.columns: - curves["label"] = curves["label"].astype(bool) - dataset_keys = ( - datasets - if datasets is not None - else sorted(curves["dataset"].astype(str).unique()) - ) - for key in dataset_keys: - logger.info("Replotting %s from curves CSV", key) - plot_dataset_from_curves(curves, str(key), output_dir) - curves_out = curves.loc[curves["dataset"].astype(str).isin(dataset_keys)] - acceptance_rows, error_rows = summary_rows_from_curves(curves_out) - acceptance, error_gain = _finalise_method_comparison_tables( - acceptance_rows, error_rows - ) - write_summary_tables( - acceptance, error_gain, output_dir, "fdr_method_comparison" - ) - if summarise_only.resolve() != (output_dir / CURVES_CSV_NAME).resolve(): - write_curves_csv(curves_out, output_dir) - return - - dataset_keys = datasets if datasets is not None else list(DEFAULT_DATASETS) - configs = build_dataset_configs(winnow_results, novoboard_root=novoboard_root) - - curve_parts: list[pd.DataFrame] = [] - for key in dataset_keys: - if key not in configs: - raise typer.BadParameter(f"Unknown dataset {key!r}") - logger.info("Processing %s", key) - curve_parts.append(process_dataset(configs[key], output_dir)) - - curves = pd.concat(curve_parts, ignore_index=True) - write_curves_csv(curves, output_dir) - acceptance_rows, error_rows = summary_rows_from_curves(curves) - acceptance, error_gain = _finalise_method_comparison_tables( - acceptance_rows, error_rows - ) - write_summary_tables(acceptance, error_gain, output_dir, "fdr_method_comparison") - - -if __name__ == "__main__": - app() diff --git a/scripts/plot_feature_investigation.py b/scripts/plot_feature_investigation.py deleted file mode 100644 index 2fd14664..00000000 --- a/scripts/plot_feature_investigation.py +++ /dev/null @@ -1,1323 +0,0 @@ -"""Generate feature investigation plots from calibrator training feature matrices. - -Produces KDE, scatter, violin, correlation, discriminative-power, pairplot, -mirror-spectrum, retention-time, token-stem, and beam-stem figures matching the -style of ``analysis/feature_investigation_new.ipynb``. - -Usage: - python scripts/plot_feature_investigation.py \ - --features-train models/instanovo_helaqc/features_train.parquet \ - [--features-val models/instanovo_helaqc/features_val.parquet] \ - [--metadata-train models/instanovo_helaqc/metadata_train.parquet] \ - [--metadata-val models/instanovo_helaqc/metadata_val.parquet] \ - [--predictions-csv held_out_projects/.../predictions.csv] \ - [--output-dir models/instanovo_helaqc/feature_investigation_plots] -""" - -from __future__ import annotations - -import argparse -import ast -import warnings -from pathlib import Path -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import polars as pl -import seaborn as sns -from matplotlib.colors import LinearSegmentedColormap -from matplotlib.patches import Patch -from scipy import stats as sp_stats -from scipy.stats import gaussian_kde -from sklearn.decomposition import PCA -from sklearn.metrics import roc_auc_score -from sklearn.preprocessing import StandardScaler - -# --------------------------------------------------------------------------- -# Style — Paul Tol "bright" palette (colour-blind safe) -# --------------------------------------------------------------------------- -_PALETTE = [ - "#4477AA", - "#EE6677", - "#228833", - "#CCBB44", - "#66CCEE", - "#AA3377", - "#BBBBBB", -] -_CORRECT_COLOUR = _PALETTE[0] -_INCORRECT_COLOUR = _PALETTE[1] -_NEUTRAL_COLOUR = _PALETTE[6] - -_HIGH_CONF_BEAM_COLOUR = _PALETTE[2] # green -_LOW_CONF_BEAM_COLOUR = _PALETTE[5] # purple -_MED_CONF_BEAM_COLOUR = _PALETTE[3] # yellow - -_OBS_COLOUR = _PALETTE[3] # yellow (observed spectrum) -_THEO_COLOUR = _PALETTE[5] # purple (predicted spectrum) - -HUE_LABEL_CORRECT = "Correct" -HUE_LABEL_INCORRECT = "Incorrect" -HUE_ORDER = [HUE_LABEL_CORRECT, HUE_LABEL_INCORRECT] -HUE_PALETTE = { - HUE_LABEL_CORRECT: _CORRECT_COLOUR, - HUE_LABEL_INCORRECT: _INCORRECT_COLOUR, -} - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -_SUNSET_COLORS = [ - "#364B9A", - "#4A7BB7", - "#6EA6CD", - "#98CAE1", - "#C2E4EF", - "#EAECCC", - "#FEDA8B", - "#FDB366", - "#F67E4B", - "#DD3D2D", - "#A50026", -] - - -def _diverging_cmap() -> LinearSegmentedColormap: - cmap = LinearSegmentedColormap.from_list("tol_sunset", _SUNSET_COLORS, N=256) - cmap.set_bad(color="#FFFFFF") - return cmap - - -# --------------------------------------------------------------------------- -# Feature definitions -# --------------------------------------------------------------------------- -FEATURE_COLUMNS = [ - "confidence", - "mass_error_ppm", - "ion_matches", - "ion_match_intensity", - "complementary_ion_count", - "max_ion_gap", - "spectral_angle", - "xcorr", - "irt_error", - "margin", - "median_margin", - "entropy", - "z-score", - "edit_distance", - "min_token_probability", - "std_token_probability", -] - -FRAGMENT_FEATURES = [ - "ion_matches", - "ion_match_intensity", - "complementary_ion_count", - "max_ion_gap", - "spectral_angle", - "xcorr", -] - -BEAM_FEATURES = ["margin", "median_margin", "entropy", "z-score", "edit_distance"] - -TOKEN_FEATURES = ["min_token_probability", "std_token_probability"] - -SKEWED_FEATURES = {"irt_error"} - -_NICE_LABELS: dict[str, str] = { - "ion_matches": "Ion match rate", - "ion_match_intensity": "Ion match intensity", - "complementary_ion_count": "Complementary ion count", - "max_ion_gap": "Max ion gap", - "spectral_angle": "Spectral angle", - "xcorr": "Cross-correlation (XCorr)", - "mass_error_ppm": "Precursor mass error", - "log_abs_mass_error_ppm": "Log-absolute mass error", - "log_abs_mass_error_da": "Log-absolute mass error", - "mass_error_da": "Precursor mass error", - "irt_error": "iRT prediction error", - "confidence": "Model confidence", - "margin": "Beam margin", - "median_margin": "Beam median margin", - "entropy": "Beam entropy", - "z-score": "Beam z-score", - "edit_distance": "Runner-up edit distance", - "min_token_probability": "Min. token probability", - "std_token_probability": "Std. token probability", - "predicted_irt": "Regressor-predicted iRT", - "irt": "Koina-predicted iRT", - "retention_time": "Retention time (s)", -} - - -def _nice_label(col: str) -> str: - return _NICE_LABELS.get(col, col.replace("_", " ").capitalize()) - - -# --------------------------------------------------------------------------- -# Plotting helpers -# --------------------------------------------------------------------------- -def _save_fig(fig: plt.Figure, name: str, output_dir: Path) -> None: - base = output_dir / name - fig.savefig(f"{base}.png", bbox_inches="tight", dpi=300) - fig.savefig(f"{base}.pdf", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def _style_ax(ax: plt.Axes) -> None: - ax.grid(False) - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -def _auto_ylim(df: pd.DataFrame, feature: str): - if feature in SKEWED_FEATURES: - q99 = df[feature].quantile(0.99) - q01 = df[feature].quantile(0.01) - margin = (q99 - q01) * 0.1 - return (q01 - margin, q99 + margin) - return None - - -def plot_feature_vs_confidence( - df: pd.DataFrame, - feature: str, - title: str | None = None, - ylim: tuple[float, float] | None = None, -) -> tuple[plt.Figure, plt.Axes]: - """Scatter plot of a feature against model confidence, coloured by class.""" - fig, ax = plt.subplots(figsize=(8, 6)) - for label in HUE_ORDER: - subset = df[df["hue"] == label] - ax.scatter( - subset["confidence"], - subset[feature], - c=HUE_PALETTE[label], - label=label, - alpha=0.3, - s=10, - rasterized=True, - ) - ax.set_xlabel(_nice_label("confidence")) - ax.set_ylabel(_nice_label(feature)) - if ylim is not None: - ax.set_ylim(ylim) - if title: - ax.set_title(title) - ax.legend() - _style_ax(ax) - fig.tight_layout() - return fig, ax - - -def _plot_peak_normalised_kde( - ax: plt.Axes, - subset: pd.Series, - colour: str, - label: str, - fill: bool, - clip: tuple[float, float] | None, -) -> None: - """Plot a peak-normalised KDE curve on *ax*.""" - kde = gaussian_kde(subset) - lo = subset.min() if clip is None else clip[0] - hi = subset.max() if clip is None else clip[1] - xs = np.linspace(lo, hi, 500) - ys = kde(xs) - ys /= ys.max() - ax.plot(xs, ys, color=colour, label=label, linewidth=1.5) - if fill: - ax.fill_between(xs, ys, alpha=0.3, color=colour) - - -def plot_kde_by_class( - df: pd.DataFrame, - feature: str, - title: str | None = None, - fill: bool = True, - clip: tuple[float, float] | None = None, - peak_normalise: bool = False, -) -> tuple[plt.Figure, plt.Axes]: - """KDE density plot of a feature split by correct/incorrect class.""" - fig, ax = plt.subplots(figsize=(8, 6)) - for label in HUE_ORDER: - subset = df[df["hue"] == label][feature].dropna() - if len(subset) < 2: - continue - if peak_normalise: - _plot_peak_normalised_kde(ax, subset, HUE_PALETTE[label], label, fill, clip) - else: - kw: dict = {} - if clip is not None: - kw["clip"] = clip - sns.kdeplot( - subset, - ax=ax, - color=HUE_PALETTE[label], - label=label, - fill=fill, - alpha=0.3, - linewidth=1.5, - **kw, - ) - ax.set_xlabel(_nice_label(feature)) - ax.set_ylabel("Peak-normalised density" if peak_normalise else "Density") - if title: - ax.set_title(title) - ax.legend(loc="upper center") - _style_ax(ax) - fig.tight_layout() - return fig, ax - - -def plot_mirror_spectrum( - obs_mz, - obs_int, - theo_mz, - theo_int, - annotations, - title: str, - ax: plt.Axes | None = None, -) -> tuple[plt.Figure, plt.Axes]: - """Mirror plot comparing observed vs predicted spectra.""" - own_fig = ax is None - if own_fig: - fig, ax = plt.subplots(figsize=(10, 4)) - else: - assert ax is not None - fig = ax.get_figure() - - obs_int_norm = np.array(obs_int) / max(obs_int) * 100 - theo_int_norm = np.array(theo_int) / max(theo_int) * 100 - - ax.vlines(obs_mz, 0, obs_int_norm, color=_OBS_COLOUR, linewidth=1.8) - ax.vlines(theo_mz, 0, -theo_int_norm, color=_THEO_COLOUR, linewidth=1.8) - - if annotations is not None: - for mz_val, intensity_val, ann in zip(theo_mz, theo_int_norm, annotations): - if intensity_val > 10: - label_text = ann.decode() if isinstance(ann, bytes) else str(ann) - ax.annotate( - label_text, - (mz_val, -intensity_val), - fontsize=8, - ha="center", - va="top", - rotation=90, - color=_THEO_COLOUR, - ) - - ax.axhline(0, color="black", linewidth=0.5) - ax.set_xlabel("m/z") - ax.set_ylabel("Relative intensity (%)") - ax.set_title(title) - - ax.text( - 0.99, - 0.95, - "Observed", - transform=ax.transAxes, - ha="right", - va="top", - fontsize=9, - color=_OBS_COLOUR, - fontweight="bold", - ) - ax.text( - 0.99, - 0.05, - "Predicted", - transform=ax.transAxes, - ha="right", - va="bottom", - fontsize=9, - color=_THEO_COLOUR, - fontweight="bold", - ) - - _style_ax(ax) - if own_fig: - fig.tight_layout() - return fig, ax - - -def compute_discriminative_stats(df: pd.DataFrame, features: list[str]) -> pd.DataFrame: - """Compute AUROC, KS statistic, and Cohen's d for each feature.""" - results = [] - labels = df["correct"].astype(int) - for feat in features: - vals = df[feat].dropna() - valid_mask = df[feat].notna() - valid_labels = labels[valid_mask] - valid_vals = vals - if len(valid_vals) < 10 or valid_labels.nunique() < 2: - results.append( - { - "feature": feat, - "auroc": np.nan, - "ks_stat": np.nan, - "cohens_d": np.nan, - } - ) - continue - try: - auroc = roc_auc_score(valid_labels, valid_vals) - auroc = max(auroc, 1 - auroc) - except ValueError: - auroc = np.nan - correct_vals = valid_vals[valid_labels == 1] - incorrect_vals = valid_vals[valid_labels == 0] - ks_stat, _ = sp_stats.ks_2samp(correct_vals, incorrect_vals) - pooled_std = np.sqrt( - ( - (len(correct_vals) - 1) * correct_vals.std() ** 2 - + (len(incorrect_vals) - 1) * incorrect_vals.std() ** 2 - ) - / (len(correct_vals) + len(incorrect_vals) - 2) - ) - cohens_d = ( - abs(correct_vals.mean() - incorrect_vals.mean()) / pooled_std - if pooled_std > 0 - else np.nan - ) - results.append( - {"feature": feat, "auroc": auroc, "ks_stat": ks_stat, "cohens_d": cohens_d} - ) - return ( - pd.DataFrame(results) - .sort_values("auroc", ascending=False) - .reset_index(drop=True) - ) - - -# --------------------------------------------------------------------------- -# Token / beam stem helpers -# --------------------------------------------------------------------------- -_STEM_FIGSIZE = (10, 4.5) - - -def _parse_token_probs(row): - """Extract token probabilities and residue labels from a row.""" - try: - token_probs = np.exp(np.array(ast.literal_eval(row["token_log_probs"]))) - except (ValueError, SyntaxError): - return np.array([]), [] - seq_str = row["prediction"] - residues: list[str] = [] - j = 0 - while j < len(seq_str): - if j + 1 < len(seq_str) and seq_str[j + 1] == "[": - end = seq_str.index("]", j + 1) + 1 - residues.append(seq_str[j:end]) - j = end - else: - residues.append(seq_str[j]) - j += 1 - n_tokens = min(len(token_probs), len(residues)) - return token_probs[:n_tokens], residues[:n_tokens] - - -def _plot_token_stem(row, beam_colour: str, title_suffix: str = ""): - """Stem plot of per-residue token probabilities for one PSM.""" - token_probs, residues = _parse_token_probs(row) - if len(token_probs) == 0: - return None - - n = len(token_probs) - fig, ax = plt.subplots(figsize=_STEM_FIGSIZE) - markerline, stemlines, baseline = ax.stem( - range(n), - token_probs, - linefmt="-", - markerfmt="o", - basefmt="k-", - ) - plt.setp(stemlines, color=beam_colour, linewidth=2.5) - plt.setp(markerline, color=beam_colour, markersize=7, zorder=5) - - ax.set_xticks(range(n)) - ax.set_xticklabels( - residues, - fontsize=11, - rotation=45, - ha="right", - rotation_mode="anchor", - ) - ax.set_xlim(-0.5, n - 0.5) - ax.set_ylim(-0.03, 1.05) - ax.set_ylabel("Token probability") - ax.set_xlabel("Residue") - - charge = int(row["precursor_charge"]) if "precursor_charge" in row.index else "?" - ax.set_title( - f"Token probabilities for {row['prediction']}, " - f"+{charge}, confidence={row['confidence']:.3f}{title_suffix}", - ) - _style_ax(ax) - fig.tight_layout() - return fig - - -def _infer_charge(row) -> str: - """Best-effort charge extraction from a beam CSV row.""" - for col in ("precursor_charge", "charge"): - if col in row.index and pd.notna(row[col]): - return str(int(row[col])) - return "?" - - -def _plot_beam_stem(row, beam_log_prob_cols, beam_seq_cols, colour: str): - """Stem plot of per-beam confidence for one spectrum.""" - probs: list[float] = [] - labels: list[str] = [] - for i, (lp_col, seq_col) in enumerate(zip(beam_log_prob_cols, beam_seq_cols)): - lp = row[lp_col] - seq = row[seq_col] - if pd.isna(lp) or np.isinf(lp): - continue - probs.append(np.exp(float(lp))) - label = str(seq) if pd.notna(seq) else f"beam {i}" - labels.append(label) - - if len(probs) < 2: - return None - - n = len(probs) - fig, ax = plt.subplots(figsize=_STEM_FIGSIZE) - markerline, stemlines, baseline = ax.stem( - range(n), - probs, - linefmt="-", - markerfmt="o", - basefmt="k-", - ) - plt.setp(stemlines, color=colour, linewidth=2.5) - plt.setp(markerline, color=colour, markersize=7, zorder=5) - - ax.set_xticks(range(n)) - ax.set_xlim(-0.5, n - 0.5) - ax.set_ylim(-max(probs) * 0.03, max(probs) * 1.15) - ax.set_ylabel("Beam confidence") - ax.set_xlabel("Beam prediction index") - ax.set_title(f"Beam confidence for {labels[0]}, +{_infer_charge(row)}") - _style_ax(ax) - fig.tight_layout() - return fig - - -# --------------------------------------------------------------------------- -# Section generators — each mirrors a notebook section -# --------------------------------------------------------------------------- -def plot_confidence(df: pd.DataFrame, output_dir: Path) -> None: - """Section 1: confidence distribution.""" - fig, ax = plot_kde_by_class(df, "confidence", title="Model confidence distribution") - _save_fig(fig, "01a_confidence_kde", output_dir) - - fig, ax = plt.subplots(figsize=(8, 6)) - for label in HUE_ORDER: - subset = df[df["hue"] == label] - ax.hist( - subset["confidence"], - bins=50, - alpha=0.5, - color=HUE_PALETTE[label], - edgecolor="black", - label=label, - density=True, - ) - ax.set_xlabel(_nice_label("confidence")) - ax.set_ylabel("Density") - ax.set_title("Model confidence histogram") - ax.legend(loc="upper center") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "01b_confidence_histogram", output_dir) - - -def _mass_error_log_column(df: pd.DataFrame) -> tuple[str, str]: - """Return (raw mass error column, log-absolute column) for plotting.""" - if "mass_error_da" in df.columns: - return "mass_error_da", "log_abs_mass_error_da" - if "mass_error_ppm" in df.columns: - return "mass_error_ppm", "log_abs_mass_error_ppm" - raise KeyError( - "Feature matrix must contain 'mass_error_da' or 'mass_error_ppm' for mass error plots" - ) - - -def plot_mass_error(df: pd.DataFrame, output_dir: Path) -> None: - """Section 2: mass error vs confidence (Da or ppm).""" - raw_col, log_col = _mass_error_log_column(df) - work = df.copy() - work[log_col] = np.log(work[raw_col].abs().clip(lower=1e-12)) - - fig, _ = plot_kde_by_class( - work, - raw_col, - title=f"{_nice_label(raw_col)} distribution", - ) - _save_fig(fig, f"02a_{raw_col}_kde", output_dir) - - fig, _ = plot_feature_vs_confidence( - work, - raw_col, - title=f"{_nice_label(raw_col)} vs model confidence", - ) - _save_fig(fig, f"02b_{raw_col}_vs_confidence", output_dir) - - fig, _ = plot_kde_by_class( - work, - log_col, - title="Log-absolute precursor mass error distribution", - ) - _save_fig(fig, "02c_mass_error_log_kde", output_dir) - - fig, _ = plot_feature_vs_confidence( - work, - log_col, - title="Log-absolute precursor mass error vs model confidence", - ) - _save_fig(fig, "02d_mass_error_log_vs_confidence", output_dir) - - if raw_col == "mass_error_da": - da_ylim = (-0.2, 0.2) - fig, _ = plot_kde_by_class( - work, - raw_col, - title=f"{_nice_label(raw_col)} distribution", - clip=da_ylim, - ) - _save_fig(fig, "02e_mass_error_da_kde_within_0.2da", output_dir) - - fig, _ = plot_feature_vs_confidence( - work, - raw_col, - title=f"{_nice_label(raw_col)} vs model confidence", - ylim=da_ylim, - ) - _save_fig(fig, "02f_mass_error_da_vs_confidence_within_0.2da", output_dir) - - -def plot_mirror_spectra(df_meta: pd.DataFrame, output_dir: Path) -> None: - """Section 3: mirror plots of observed vs predicted spectra.""" - required = { - "theoretical_mz", - "mz_array", - "intensity_array", - "theoretical_intensity", - } - if not required.issubset(df_meta.columns): - print(" Skipping mirror plots — missing spectrum columns in metadata.") - return - - valid_mirror = df_meta[ - df_meta["theoretical_mz"].apply(lambda x: x is not None and len(x) > 0) - & df_meta["mz_array"].apply(lambda x: x is not None and len(x) > 0) - ].copy() - - if len(valid_mirror) == 0: - print(" Skipping mirror plots — no rows with valid spectrum arrays.") - return - - has_annotations = "theoretical_annotation" in valid_mirror.columns - - def _add_mirror_margin(ax, y_frac=0.11): - ymin, ymax = ax.get_ylim() - y_pad = (ymax - ymin) * y_frac - ax.set_ylim(ymin - y_pad, ymax + y_pad) - - def _mirror_title(row) -> str: - pred = row.get("prediction", "?") - charge = ( - int(row["precursor_charge"]) if "precursor_charge" in row.index else "?" - ) - return f"Observed vs predicted spectrum for {pred}, +{charge}" - - correct_high = valid_mirror[valid_mirror["correct"]].nlargest(3, "confidence") - incorrect_low = valid_mirror[~valid_mirror["correct"]].nsmallest(3, "confidence") - - conf_middle_lo, conf_middle_hi = 0.45, 0.55 - middle_mask = valid_mirror["confidence"].between(conf_middle_lo, conf_middle_hi) - n_correct_mid = (valid_mirror["correct"] & middle_mask).sum() - n_incorrect_mid = (~valid_mirror["correct"] & middle_mask).sum() - correct_middle = valid_mirror[valid_mirror["correct"] & middle_mask].sample( - n=min(3, n_correct_mid), random_state=42 - ) - incorrect_middle = valid_mirror[~valid_mirror["correct"] & middle_mask].sample( - n=min(3, n_incorrect_mid), random_state=42 - ) - - groups = [ - (correct_high, "correct", "03a_mirror_high_conf"), - (incorrect_low, "incorrect", "03b_mirror_low_conf"), - (correct_middle, "correct", "03c_mirror_middle_conf_correct"), - (incorrect_middle, "incorrect", "03d_mirror_middle_conf_incorrect"), - ] - - for subset, _status, prefix in groups: - for i, (_, row) in enumerate(subset.iterrows()): - fig, ax = plt.subplots(figsize=(8, 5)) - annotations = row.get("theoretical_annotation") if has_annotations else None - plot_mirror_spectrum( - obs_mz=row["mz_array"], - obs_int=row["intensity_array"], - theo_mz=row["theoretical_mz"], - theo_int=row["theoretical_intensity"], - annotations=annotations, - title=_mirror_title(row), - ax=ax, - ) - _add_mirror_margin(ax) - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, f"{prefix}_{i}", output_dir) - - -def plot_fragment_features(df: pd.DataFrame, output_dir: Path) -> None: - """Section 4: fragment ion match features vs confidence.""" - available = [f for f in FRAGMENT_FEATURES if f in df.columns] - for feat in available: - fig, _ = plot_feature_vs_confidence( - df, - feat, - title=f"{_nice_label(feat)} vs model confidence", - ylim=_auto_ylim(df, feat), - ) - _save_fig(fig, f"04a_fragment_{feat}_vs_confidence", output_dir) - - fig, _ = plot_kde_by_class(df, feat, title=f"{_nice_label(feat)} distribution") - _save_fig(fig, f"04b_fragment_{feat}_kde", output_dir) - - -def plot_irt(df: pd.DataFrame, df_meta: pd.DataFrame | None, output_dir: Path) -> None: - """Section 7: iRT error plots + RT scatter when metadata is available.""" - if df_meta is not None: - has_rt = ( - "retention_time" in df_meta.columns and "predicted_irt" in df_meta.columns - ) - has_koina_irt = "irt" in df_meta.columns - - if has_rt and has_koina_irt: - fig, ax = plt.subplots(figsize=(8, 6)) - for label in HUE_ORDER: - subset = df_meta[df_meta["hue"] == label] - ax.scatter( - subset["retention_time"], - subset["irt"], - c=HUE_PALETTE[label], - label=label, - alpha=0.3, - s=10, - rasterized=True, - ) - ax.set_xlabel(_nice_label("retention_time")) - ax.set_ylabel(_nice_label("irt")) - ax.set_title("Retention time vs Koina-predicted iRT") - ax.legend(markerscale=3, frameon=True) - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "07a_rt_vs_koina_irt", output_dir) - - if has_koina_irt and has_rt: - fig, ax = plt.subplots(figsize=(8, 6)) - for label in HUE_ORDER: - subset = df_meta[df_meta["hue"] == label] - ax.scatter( - subset["predicted_irt"], - subset["irt"], - c=HUE_PALETTE[label], - label=label, - alpha=0.3, - s=10, - rasterized=True, - ) - ax.set_xlabel(_nice_label("predicted_irt")) - ax.set_ylabel(_nice_label("irt")) - ax.set_title("Koina-predicted iRT vs regressor-predicted iRT") - ax.legend(markerscale=3, frameon=True) - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "07b_predicted_vs_koina_irt", output_dir) - - if "irt_error" not in df.columns: - return - - irt_ylim = _auto_ylim(df, "irt_error") - fig, _ = plot_feature_vs_confidence( - df, "irt_error", title="iRT prediction error vs model confidence", ylim=irt_ylim - ) - _save_fig(fig, "07c_irt_error_vs_confidence", output_dir) - - fig, _ = plot_kde_by_class( - df, - "irt_error", - title="iRT prediction error distribution", - clip=(0, df["irt_error"].quantile(0.99)), - ) - _save_fig(fig, "07d_irt_error_kde", output_dir) - - -def plot_token_stems(df_meta: pd.DataFrame, output_dir: Path) -> None: - """Section 8: token-level probability stem plots from metadata.""" - if "token_log_probs" not in df_meta.columns: - print(" Skipping token stem plots — no token_log_probs column in metadata.") - return - if "prediction" not in df_meta.columns: - print(" Skipping token stem plots — no prediction column in metadata.") - return - - # High confidence - high_conf_pool = df_meta[df_meta["confidence"] >= 0.9] - high_samples = ( - high_conf_pool.sample(3, random_state=42) - if len(high_conf_pool) >= 3 - else high_conf_pool - ) - for i, (_, row) in enumerate(high_samples.iterrows()): - fig = _plot_token_stem(row, _HIGH_CONF_BEAM_COLOUR) - if fig is not None: - _save_fig(fig, f"08a_token_stem_high_conf_{i}", output_dir) - - # Medium confidence - med_conf_pool = df_meta[ - (df_meta["confidence"] >= 0.4) & (df_meta["confidence"] <= 0.7) - ] - med_samples = ( - med_conf_pool.sample(3, random_state=42) - if len(med_conf_pool) >= 3 - else med_conf_pool - ) - for i, (_, row) in enumerate(med_samples.iterrows()): - fig = _plot_token_stem(row, _MED_CONF_BEAM_COLOUR) - if fig is not None: - _save_fig(fig, f"08b_token_stem_med_conf_{i}", output_dir) - - # Low confidence - low_conf_pool = df_meta[df_meta["confidence"] <= 0.2] - low_samples = ( - low_conf_pool.sample(3, random_state=42) - if len(low_conf_pool) >= 3 - else low_conf_pool - ) - for i, (_, row) in enumerate(low_samples.iterrows()): - fig = _plot_token_stem(row, _LOW_CONF_BEAM_COLOUR) - if fig is not None: - _save_fig(fig, f"08c_token_stem_low_conf_{i}", output_dir) - - -def plot_beam_stems(predictions_csv: Path, output_dir: Path) -> None: - """Section 8b: beam confidence stem plots from the predictions CSV.""" - beam_csv = pd.read_csv(predictions_csv) - - beam_log_prob_cols = sorted( - [ - c - for c in beam_csv.columns - if c.startswith("predictions_log_probability_beam_") - ], - key=lambda c: int(c.rsplit("_", 1)[1]), - ) - beam_seq_cols = sorted( - [ - c - for c in beam_csv.columns - if c.startswith("predictions_beam_") - and "log_probability" not in c - and "token" not in c - ], - key=lambda c: int(c.rsplit("_", 1)[1]), - ) - - if not beam_log_prob_cols or not beam_seq_cols: - print(" Skipping beam stem plots — no beam columns in predictions CSV.") - return - - beam_csv["top_confidence"] = np.exp(beam_csv[beam_log_prob_cols[0]].astype(float)) - - # Filter rows where all beams are -inf or NaN - valid_beams = beam_csv.dropna(subset=beam_log_prob_cols, how="all").copy() - for col in beam_log_prob_cols: - valid_beams[col] = pd.to_numeric(valid_beams[col], errors="coerce") - valid_beams = valid_beams[ - valid_beams[beam_log_prob_cols].apply( - lambda row: not all(np.isinf(row) | row.isna()), axis=1 - ) - ] - valid_beams = valid_beams[ - valid_beams[beam_log_prob_cols].apply( - lambda row: any(np.exp(row.dropna()) > 1e-15), axis=1 - ) - ] - - if len(valid_beams) == 0: - print(" Skipping beam stem plots — no valid beam rows after filtering.") - return - - # High confidence beams - high_beam = valid_beams[valid_beams["top_confidence"] >= 0.9] - high_beam_samples = ( - high_beam.sample(3, random_state=42) if len(high_beam) >= 3 else high_beam - ) - for i, (_, row) in enumerate(high_beam_samples.iterrows()): - fig = _plot_beam_stem( - row, beam_log_prob_cols, beam_seq_cols, _HIGH_CONF_BEAM_COLOUR - ) - if fig is not None: - _save_fig(fig, f"08d_beam_conf_high_{i}", output_dir) - - # Low confidence beams - low_beam = valid_beams[valid_beams["top_confidence"] <= 0.2] - low_beam_samples = ( - low_beam.sample(3, random_state=42) if len(low_beam) >= 3 else low_beam - ) - for i, (_, row) in enumerate(low_beam_samples.iterrows()): - fig = _plot_beam_stem( - row, beam_log_prob_cols, beam_seq_cols, _LOW_CONF_BEAM_COLOUR - ) - if fig is not None: - _save_fig(fig, f"08e_beam_conf_low_{i}", output_dir) - - -def plot_beam_features(df: pd.DataFrame, output_dir: Path) -> None: - """Section 9: beam search features vs confidence.""" - available = [f for f in BEAM_FEATURES if f in df.columns] - for feat in available: - fig, _ = plot_feature_vs_confidence( - df, feat, title=f"{_nice_label(feat)} vs model confidence" - ) - _save_fig(fig, f"09a_beam_{feat}_scatter", output_dir) - - fig, _ = plot_kde_by_class(df, feat, title=f"{_nice_label(feat)} distribution") - _save_fig(fig, f"09b_beam_{feat}_kde", output_dir) - - -def plot_token_features(df: pd.DataFrame, output_dir: Path) -> None: - """Section 11: token-level features.""" - if "min_token_probability" not in df.columns: - return - - fig, _ = plot_kde_by_class( - df, "min_token_probability", title="Min. token probability distribution" - ) - _save_fig(fig, "11a_min_token_prob_kde", output_dir) - - fig, _ = plot_kde_by_class( - df, "std_token_probability", title="Std. token probability distribution" - ) - _save_fig(fig, "11b_std_token_prob_kde", output_dir) - - fig, _ = plot_feature_vs_confidence( - df, "min_token_probability", title="Min. token probability vs confidence" - ) - _save_fig(fig, "11c_min_token_prob_scatter", output_dir) - - fig, _ = plot_feature_vs_confidence( - df, "std_token_probability", title="Std. token probability vs confidence" - ) - _save_fig(fig, "11d_std_token_prob_scatter", output_dir) - - fig, ax = plt.subplots(figsize=(8, 6)) - for label in HUE_ORDER: - subset = df[df["hue"] == label] - ax.scatter( - subset["min_token_probability"], - subset["std_token_probability"], - c=HUE_PALETTE[label], - label=label, - alpha=0.3, - s=10, - rasterized=True, - ) - ax.set_xlabel(_nice_label("min_token_probability")) - ax.set_ylabel(_nice_label("std_token_probability")) - ax.set_title("Token-level feature space") - ax.legend(markerscale=3, frameon=True) - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "11e_token_feature_2d", output_dir) - - -def _plot_pca(df: pd.DataFrame, available: list[str], output_dir: Path) -> None: - """12f/12g — PCA scatter and loadings for calibrator features.""" - feat_df = df[available].dropna() - if len(feat_df) < 10: - return - - hue_pca = df.loc[feat_df.index, "hue"].values - scaler = StandardScaler() - x_scaled = scaler.fit_transform(feat_df.values) - pca = PCA(n_components=2) - z_pca = pca.fit_transform(x_scaled) - - fig, ax = plt.subplots(figsize=(8, 7)) - for label, colour in zip(reversed(HUE_ORDER), [_INCORRECT_COLOUR, _CORRECT_COLOUR]): - mask = hue_pca == label - ax.scatter( - z_pca[mask, 0], - z_pca[mask, 1], - c=colour, - label=label, - s=10, - alpha=0.3, - rasterized=True, - ) - ax.set_xlabel(f"PC 1 ({pca.explained_variance_ratio_[0]:.1%} variance)") - ax.set_ylabel(f"PC 2 ({pca.explained_variance_ratio_[1]:.1%} variance)") - ax.set_title("PCA of calibrator features") - ax.legend(loc="upper left") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "12f_pca_features", output_dir) - - pc1 = pca.components_[0] - pc2 = pca.components_[1] - names = [_nice_label(c) for c in available] - order = np.argsort(np.abs(pc1))[::-1] - - y = np.arange(len(names)) - fig, ax = plt.subplots(figsize=(10, 7)) - ax.barh( - y, - pc1[order], - color=_CORRECT_COLOUR, - alpha=0.6, - edgecolor="black", - linewidth=0.4, - ) - ax.barh( - y, pc2[order], color=_PALETTE[5], alpha=0.4, edgecolor="black", linewidth=0.4 - ) - ax.set_yticks(y) - ax.set_yticklabels([names[i] for i in order]) - ax.invert_yaxis() - ax.set_xlabel("Loading value") - ax.set_title("PCA loadings for first two principal components") - ax.axvline(0, color="black", linewidth=0.5) - ax.legend( - handles=[ - Patch(facecolor=_CORRECT_COLOUR, alpha=0.6, label="PC 1 loading"), - Patch(facecolor=_PALETTE[5], alpha=0.4, label="PC 2 loading"), - ], - loc="lower right", - ) - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "12g_pca_loadings", output_dir) - - -def plot_discriminative_analysis(df: pd.DataFrame, output_dir: Path) -> None: - """Section 12: discriminative stats, correlation, violins, pairplot.""" - available = [f for f in FEATURE_COLUMNS if f in df.columns] - disc_stats = compute_discriminative_stats(df, available) - - # 12a — AUROC bar chart - fig, ax = plt.subplots(figsize=(8, 7)) - colours = [ - _PALETTE[0] if v >= 0.7 else _NEUTRAL_COLOUR for v in disc_stats["auroc"] - ] - ax.barh(range(len(disc_stats)), disc_stats["auroc"], color=colours) - ax.set_yticks(range(len(disc_stats))) - ax.set_yticklabels([_nice_label(f) for f in disc_stats["feature"]], fontsize=9) - ax.set_xlabel("AUROC") - ax.set_title("Per-feature AUROC for separating correct vs incorrect") - ax.axvline(0.5, color="grey", linestyle="--", linewidth=0.8) - ax.invert_yaxis() - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, "12a_feature_auroc_ranking", output_dir) - - # 12b — correlation matrix - corr = df[available].corr() - fig, ax = plt.subplots(figsize=(14, 12)) - mask = np.triu(np.ones_like(corr, dtype=bool), k=1) - sns.heatmap( - corr, - mask=mask, - cmap=_diverging_cmap(), - center=0, - ax=ax, - xticklabels=[_nice_label(c) for c in available], - yticklabels=[_nice_label(c) for c in available], - annot=True, - fmt=".2f", - annot_kws={"size": 6}, - linewidths=0.5, - square=True, - vmin=-1, - vmax=1, - cbar_kws={"label": "Pearson r"}, - ) - ax.set_title("Feature correlation matrix") - ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right", fontsize=7) - ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=7) - fig.tight_layout() - _save_fig(fig, "12b_correlation_matrix", output_dir) - - # 12c — violin plots per feature group - feature_groups = { - "Fragment match": [f for f in FRAGMENT_FEATURES if f in df.columns], - "Beam search": [f for f in BEAM_FEATURES if f in df.columns], - "Token-level": [f for f in TOKEN_FEATURES if f in df.columns], - } - for _group_name, group_feats in feature_groups.items(): - for feat in group_feats: - fig, ax = plt.subplots(figsize=(6, 5)) - sns.violinplot( - data=df, - x="hue", - y=feat, - hue="hue", - ax=ax, - palette=HUE_PALETTE, - order=HUE_ORDER, - hue_order=HUE_ORDER, - inner="quartile", - cut=0, - linewidth=0.8, - legend=False, - ) - ax.set_xlabel("") - ax.set_ylabel(_nice_label(feat)) - ax.set_title(f"{_nice_label(feat)} by identification status") - _style_ax(ax) - fig.tight_layout() - _save_fig(fig, f"12c_violin_{feat}", output_dir) - - # 12e — pairplot of top-5 features - top5 = disc_stats.head(5)["feature"].tolist() - sample_size = min(2000, len(df)) - df_sample = df[top5 + ["hue"]].sample(n=sample_size, random_state=42) - - g = sns.pairplot( - df_sample, - vars=top5, - hue="hue", - palette=HUE_PALETTE, - hue_order=HUE_ORDER, - diag_kind="kde", - plot_kws={"alpha": 0.25, "s": 8, "rasterized": True}, - diag_kws={"fill": True, "alpha": 0.3}, - height=2.2, - ) - g.figure.suptitle("Pairplot of top-5 discriminative features", y=1.01, fontsize=13) - g._legend.set_title("Identification") - for ax_row in g.axes: - for ax_item in ax_row: - xl = ax_item.get_xlabel() - yl = ax_item.get_ylabel() - if xl: - ax_item.set_xlabel(_nice_label(xl), fontsize=7) - if yl: - ax_item.set_ylabel(_nice_label(yl), fontsize=7) - _style_ax(ax_item) - _save_fig(g.figure, "12e_pairplot_top5", output_dir) - - _plot_pca(df, available, output_dir) - - # Save discriminative stats as CSV for reference - disc_stats.to_csv(output_dir / "discriminative_stats.csv", index=False) - - -def print_summary(df: pd.DataFrame) -> None: - """Section 13: summary statistics printed to stdout.""" - available = [f for f in FEATURE_COLUMNS if f in df.columns] - disc_stats = compute_discriminative_stats(df, available) - - print("=" * 80) - print("DISCRIMINATIVE STATISTICS SUMMARY") - print("=" * 80) - n_correct = df["correct"].sum() - n_incorrect = (~df["correct"]).sum() - print( - f"\nDataset: {len(df):,} spectra | {n_correct:,} correct | {n_incorrect:,} incorrect" - ) - print(f"Class balance: {df['correct'].mean():.1%} correct\n") - - print("Per-feature discriminative power (sorted by AUROC):") - print("-" * 80) - print(disc_stats.to_string(index=False, float_format="%.3f")) - - print("\nTop-5 features by AUROC:") - for _, row in disc_stats.head(5).iterrows(): - print( - f" {_nice_label(row['feature']):40s} AUROC={row['auroc']:.3f} " - f"KS={row['ks_stat']:.3f} d={row['cohens_d']:.3f}" - ) - - print("\nBottom-5 features by AUROC:") - for _, row in disc_stats.tail(5).iterrows(): - print( - f" {_nice_label(row['feature']):40s} AUROC={row['auroc']:.3f} " - f"KS={row['ks_stat']:.3f} d={row['cohens_d']:.3f}" - ) - - print("\nConfidence statistics:") - print(f" Overall mean confidence: {df['confidence'].mean():.3f}") - print(f" Correct mean confidence: {df[df['correct']]['confidence'].mean():.3f}") - print(f" Incorrect mean confidence: {df[~df['correct']]['confidence'].mean():.3f}") - print( - f" Confidence AUROC: " - f"{roc_auc_score(df['correct'].astype(int), df['confidence']):.3f}" - ) - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- -def parse_args(argv: list[str] | None = None) -> argparse.Namespace: - """Parse command-line arguments for feature investigation plots.""" - parser = argparse.ArgumentParser( - description="Feature investigation plots from calibrator training matrices.", - ) - parser.add_argument( - "--features-train", - type=Path, - required=True, - help="Path to the training features parquet produced by `winnow compute-features`.", - ) - parser.add_argument( - "--features-val", - type=Path, - default=None, - help="Optional path to validation features parquet. When provided the train " - "and val splits are concatenated for richer plots.", - ) - parser.add_argument( - "--metadata-train", - type=Path, - default=None, - help="Optional path to full training metadata parquet (produced by compute-features " - "with metadata_output_path set). Enables mirror plots, RT scatter, and token stems.", - ) - parser.add_argument( - "--metadata-val", - type=Path, - default=None, - help="Optional path to full validation metadata parquet.", - ) - parser.add_argument( - "--predictions-csv", - type=Path, - default=None, - help="Optional path to InstaNovo-style predictions CSV with beam columns. " - "Enables beam confidence stem plots.", - ) - parser.add_argument( - "--output-dir", - type=Path, - default=None, - help="Directory for output plots. Defaults to a `feature_investigation_plots` " - "subdirectory next to --features-train.", - ) - return parser.parse_args(argv) - - -def _load_metadata(args: argparse.Namespace) -> pd.DataFrame | None: - """Load and concatenate metadata parquets when provided.""" - parts: list[pd.DataFrame] = [] - for path in (args.metadata_train, args.metadata_val): - if path is not None: - print(f"Loading metadata from {path}") - parts.append(pl.read_parquet(path).to_pandas()) - if not parts: - return None - df_meta = pd.concat(parts, ignore_index=True) - print(f"Combined metadata: {len(df_meta):,} rows, {len(df_meta.columns)} columns") - df_meta["hue"] = df_meta["correct"].map( - {True: HUE_LABEL_CORRECT, False: HUE_LABEL_INCORRECT} - ) - return df_meta - - -def main(argv: list[str] | None = None) -> None: - """Generate all feature investigation plots.""" - warnings.filterwarnings("ignore", category=FutureWarning) - args = parse_args(argv) - - # -- Load features (always required) -- - print(f"Loading training features from {args.features_train}") - df = pl.read_parquet(args.features_train).to_pandas() - - if args.features_val is not None: - print(f"Loading validation features from {args.features_val}") - df_val = pl.read_parquet(args.features_val).to_pandas() - df = pd.concat([df, df_val], ignore_index=True) - print(f"Combined dataset: {len(df):,} spectra") - - df["hue"] = df["correct"].map({True: HUE_LABEL_CORRECT, False: HUE_LABEL_INCORRECT}) - - # -- Load metadata (optional) -- - df_meta = _load_metadata(args) - - has_metadata = df_meta is not None - has_predictions_csv = args.predictions_csv is not None - - output_dir = args.output_dir - if output_dir is None: - output_dir = args.features_train.parent / "feature_investigation_plots" - output_dir.mkdir(parents=True, exist_ok=True) - print(f"Saving plots to {output_dir}/\n") - - print(f"Dataset shape: {df.shape}") - n_correct = df["correct"].sum() - n_incorrect = (~df["correct"]).sum() - print(f"Correct: {n_correct:,} | Incorrect: {n_incorrect:,} | Total: {len(df):,}") - print(f"Class balance: {df['correct'].mean():.1%} correct\n") - - n_steps = 7 + has_metadata * 2 + has_predictions_csv - step = 0 - - step += 1 - print(f"[{step}/{n_steps}] Confidence distribution...") - plot_confidence(df, output_dir) - - step += 1 - print(f"[{step}/{n_steps}] Mass error vs confidence...") - plot_mass_error(df, output_dir) - - if has_metadata: - step += 1 - print(f"[{step}/{n_steps}] Mirror spectrum plots...") - plot_mirror_spectra(df_meta, output_dir) - - step += 1 - print(f"[{step}/{n_steps}] Fragment ion match features...") - plot_fragment_features(df, output_dir) - - step += 1 - print(f"[{step}/{n_steps}] iRT error...") - plot_irt(df, df_meta, output_dir) - - if has_metadata: - step += 1 - print(f"[{step}/{n_steps}] Token-level stem plots...") - plot_token_stems(df_meta, output_dir) - - if has_predictions_csv: - step += 1 - print(f"[{step}/{n_steps}] Beam confidence stem plots...") - plot_beam_stems(args.predictions_csv, output_dir) - - step += 1 - print(f"[{step}/{n_steps}] Beam search features...") - plot_beam_features(df, output_dir) - - step += 1 - print(f"[{step}/{n_steps}] Token-level features...") - plot_token_features(df, output_dir) - - step += 1 - print( - f"[{step}/{n_steps}] Discriminative analysis (AUROC, correlation, violins, pairplot)..." - ) - plot_discriminative_analysis(df, output_dir) - - print() - print_summary(df) - - print(f"\nDone — plots saved to {output_dir}/") - - -if __name__ == "__main__": - main() diff --git a/scripts/run_external_peptide_holdout_benchmark.py b/scripts/run_external_peptide_holdout_benchmark.py deleted file mode 100644 index c84f91f8..00000000 --- a/scripts/run_external_peptide_holdout_benchmark.py +++ /dev/null @@ -1,1143 +0,0 @@ -#!/usr/bin/env python3 -"""Glissade-style external peptide score-mixture benchmark. - -Builds a **shared** matched pool S_m (labelled-test Novor-correct peptides) and -external pool S_e (unlabelled proteome-external peptides) after filter → -max-score-per-peptide (NovoBoard mass-deltas converted to ProForma; unsupported -mods dropped; NovoBoard target-decoy pairs gated so every retained key has a -twin). Unlabelled / external peptides require normalised length ≥ 8 (proteome -substring proxy); labelled matched peptides and Glissade's training-split -reference keep short peptides (Novor agreement). Novor and proteome-hit labels -are computed once on Winnow and reused for NovoBoard by ``spectrum_id``. -Method-specific scores are attached to the same peptide keys. Mixtures control -π₀ explicitly. - -Mixtures are drawn **without replacement** so every peptide key is unique. All -three methods therefore score the identical mixture and realise the same π₀; -NovoBoard's max-score-per-peptide step is a no-op, which is asserted per -mixture. - -Each tool draws its null/reference information from the same place, the -annotated *training* split of its own organism: Winnow through the pretrained -per-dataset calibrator, Glissade through the training-split matched score -distribution, NovoBoard through its training-tuned decoy masking rate. No tool -fits on the evaluation labels. - -NovoBoard peptide FDR uses max-target → twin-decoy TDC. Winnow uses max -calibrated confidence then nonparametric FDR (PSM-calibrator proxy). Glissade -uses native bootstrap FDR with NumPy seeded from the benchmark RNG. - -External tool checkouts for local results: - -- ``--novoboard-root``: ``{root}/{dataset}/novoboard/`` target/decoy CSVs (the - ``datasets`` dir of a NovoBoard checkout). Local runs used fork - ``git@github.com:JemmaLDaniel/NovoBoard.git``, branch - ``feat/adapt-to-instanovo`` at - ``a9faab3ef1af06987599c2f01e6ba96072c80172``. -- ``--glissade-repo``: clone root with importable ``glissade.glissade``. Local - runs used fork ``git@github.com:JemmaLDaniel/glissade.git``, branch - ``winnow-benchmark`` at ``6ee11b51b5f21ba8fdc1eb5821608352b082a533``. -""" - -from __future__ import annotations - -import importlib -import logging -import sys -from pathlib import Path -from typing import Annotated, Optional - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import polars as pl -import typer - -_REPO_ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(_REPO_ROOT)) - -from scripts.fdr_tool_comparison_preprocess import ( # noqa: E402 - LABELLED_MIN_PEPTIDE_LENGTH, - MIN_PEPTIDE_LENGTH, - assert_shared_prediction_keys, - attach_labels_by_spectrum_id, - confidence_to_log_prob, - filter_novoboard_target_decoy_pairs, - filter_prediction_table, - label_series_by_spectrum_id, - max_score_per_peptide, - novoboard_max_target_twin_decoy_tdc, - novor_correctness_mask, - prepare_novoboard_decoy_by_pair, - restrict_winnow_to_novoboard_spectra, -) -from scripts.fdr_tool_comparison_summaries import ( # noqa: E402 - SUMMARY_THRESHOLDS, - database_grounded_q_from_labels, - mean_abs_q_dev_vs_reference, - summarise_holdout_results, - write_summary_tables, -) -from scripts.plot_eval_results import _PALETTE, _display_name, _save_fig, _style_ax # noqa: E402 -from scripts.plot_fdr_method_comparison import ( # noqa: E402 - DEFAULT_MODEL_ROOT, - DEFAULT_WINNOW_RESULTS, - build_dataset_configs, - load_novoboard_target_decoy, - load_winnow, -) -from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 - -logger = logging.getLogger(__name__) -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark" -DEFAULT_DATASETS = ["helaqc", "celegans"] -DEFAULT_Q_THRESHOLDS = [round(float(x), 2) for x in np.linspace(0.0, 0.25, 26)] -DEFAULT_PI0_GRID = [0.5, 0.6, 0.7, 0.8, 0.9] -DEFAULT_N_ITERATIONS = 20 -# Matches the Glissade package default for run_bootstraps. -DEFAULT_N_BOOTSTRAPS = 10 -# Prefer the full matched pool; |S_c| is then capped so the highest π₀ remains -# drawable without replacement from S_e. -DEFAULT_HOLDOUT_FRAC = 1.0 -DEFAULT_SEED = 42 -METHODS = ("Winnow", "NovoBoard", "Glissade") - - -def max_correct_pool_for_pi0_grid(n_external_pool: int, pi0_grid: list[float]) -> int: - """Largest |S_c| such that every π₀ in ``pi0_grid`` fits without replacement. - - Requires ``round(π₀ / (1 - π₀) · |S_c|) ≤ |S_e|`` for each target π₀. - """ - if n_external_pool < 1: - return 0 - limit = n_external_pool - for pi0 in pi0_grid: - if not 0.0 < pi0 < 1.0: - continue - ratio = pi0 / (1.0 - pi0) - # Largest n with round(ratio * n) <= n_external_pool. - # For ratio = k integer (e.g. 0.9 → 9), this is floor(pool / k). - hi = int(n_external_pool / ratio) + 2 - n_ok = 0 - for n in range(1, hi + 1): - if int(round(ratio * n)) <= n_external_pool: - n_ok = n - else: - break - limit = min(limit, n_ok) - return max(0, limit) - - -def _load_glissade_functions(glissade_repo: Path): - repo = str(glissade_repo.resolve()) - if repo not in sys.path: - sys.path.insert(0, repo) - module = importlib.import_module("glissade.glissade") - return module.run_bootstraps, module.annotate_results, module.compute_fdr_transform - - -def _load_winnow_with_raw_confidence( - predictions_dir: Path, fasta: Path, eval_type: str -) -> pd.DataFrame: - """Load Winnow preds and ensure raw ``confidence`` is present.""" - df = load_winnow(predictions_dir, fasta, eval_type) # type: ignore[arg-type] - if "confidence" not in df.columns: - meta_path = predictions_dir / "metadata.csv" - if not meta_path.is_file(): - raise FileNotFoundError(meta_path) - meta = pd.read_csv(meta_path, usecols=["spectrum_id", "confidence"]) - df = df.merge(meta, on="spectrum_id", how="inner") - return df - - -def build_glissade_training_reference( - train_metadata: Path, - *, - min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, -) -> pd.DataFrame: - """Matched reference score distribution for Glissade, from the training split. - - Glissade anchors its null-fraction estimate on a database-matched score - distribution. Taking that anchor from the annotated training split puts it on - the same data the Winnow calibrator was trained on and the NovoBoard decoy - masking rate was tuned on, and keeps it disjoint from the evaluation spectra. - Short peptides are retained: the reference is labelled (Novor-correct). - - Args: - train_metadata: Calibrator training metadata parquet. - min_length: Minimum normalised peptide length (default: labelled floor). - - Returns: - One row per Novor-correct training peptide with ``raw_confidence`` and - ``score_glissade``. - """ - if not train_metadata.is_file(): - raise FileNotFoundError(train_metadata) - train = pl.read_parquet( - train_metadata, - columns=["spectrum_id", "prediction", "sequence", "confidence"], - ).to_pandas() - train = filter_prediction_table( - train, "prediction", min_length=min_length, key_col="peptide_key" - ) - train["correct"] = novor_correctness_mask(train["sequence"], train["prediction"]) - matched = max_score_per_peptide( - train.loc[train["correct"]], "peptide_key", "confidence" - ) - reference = matched[["peptide_key", "confidence"]].rename( - columns={"confidence": "raw_confidence"} - ) - reference["score_glissade"] = confidence_to_log_prob(reference["raw_confidence"]) - logger.info( - "Glissade training reference: %d matched peptides from %s", - len(reference), - train_metadata, - ) - return reference - - -def _namespace_pair_keys(df: pd.DataFrame, namespace: str) -> pd.DataFrame: - """Prefix ``_pair_key`` to avoid collisions across splits.""" - work = df.copy() - if "_pair_key" not in work.columns: - raise ValueError("Missing '_pair_key'") - work["_pair_key"] = namespace + ":" + work["_pair_key"].astype(str) - return work - - -def build_shared_score_tables( - *, - dataset: str, - winnow_results: Path, - novoboard_root: Path, - model_root: Path = DEFAULT_MODEL_ROOT, - unlabelled_min_length: int = MIN_PEPTIDE_LENGTH, - labelled_min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, -) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]: - """Return shared matched/external keys with per-method scores. - - All methods follow filter → :func:`max_score_per_peptide`. NovoBoard tables - are pair-gated (equal target/decoy twins) before max-dedupe so mixture keys - always have a twin. Labelled S_m and Glissade's training reference keep short - peptides (``labelled_min_length``); unlabelled S_e uses - ``unlabelled_min_length`` for the proteome-substring proxy. Labelled S_m uses - Novor correctness computed once on Winnow and reused for NovoBoard by - ``spectrum_id``; S_e uses proteome-hit the same way. Glissade scores are - ``log`` raw InstaNovo confidence on the shared keys. - - Returns: - matched_scores: one row per peptide_key in S_m with method scores. - external_scores: same for S_e. - nb_decoy_combined: namespaced twin-valid decoys for twin TDC. - glissade_reference: training-split matched scores for Glissade's anchor. - """ - cfg = build_dataset_configs( - winnow_results, novoboard_root=novoboard_root, model_root=model_root - )[dataset] - - winnow_test = _load_winnow_with_raw_confidence( - cfg.winnow_test, cfg.fasta, "labelled" - ) - winnow_unlab = _load_winnow_with_raw_confidence( - cfg.winnow_unlabelled, cfg.fasta, "unlabelled" - ) - # load_winnow already applied the labelled/unlabelled length floors; re-apply - # explicitly so callers can override without depending on that path. - winnow_test = filter_prediction_table( - winnow_test, - "prediction", - min_length=labelled_min_length, - key_col="peptide_key", - ) - w_unlab = filter_prediction_table( - winnow_unlab, - "prediction", - min_length=unlabelled_min_length, - key_col="peptide_key", - ) - if "proteome_hit" not in w_unlab.columns: - raise KeyError("Expected proteome_hit on unlabelled Winnow table") - - nb_test_target, nb_test_decoy = load_novoboard_target_decoy( - cfg.novoboard_dir, "test", cfg.novoboard_decoy_rate - ) - nb_unlab_target, nb_unlab_decoy = load_novoboard_target_decoy( - cfg.novoboard_dir, "unlabelled", cfg.novoboard_decoy_rate - ) - nb_test_target = _namespace_pair_keys(nb_test_target, "test") - nb_test_decoy = _namespace_pair_keys(nb_test_decoy, "test") - nb_unlab_target = _namespace_pair_keys(nb_unlab_target, "unlabelled") - nb_unlab_decoy = _namespace_pair_keys(nb_unlab_decoy, "unlabelled") - - nb_test_target, nb_test_decoy = filter_novoboard_target_decoy_pairs( - nb_test_target, - nb_test_decoy, - min_length=labelled_min_length, - key_col="peptide_key", - ) - nb_unlab_target, nb_unlab_decoy = filter_novoboard_target_decoy_pairs( - nb_unlab_target, - nb_unlab_decoy, - min_length=unlabelled_min_length, - key_col="peptide_key", - ) - nb_decoy_combined = pd.concat( - [nb_test_decoy, nb_unlab_decoy], ignore_index=True, sort=False - ) - - # Twin-valid NovoBoard ⊆ Winnow; shared labels once on Winnow. - winnow_test = restrict_winnow_to_novoboard_spectra(winnow_test, nb_test_target) - w_unlab = restrict_winnow_to_novoboard_spectra(w_unlab, nb_unlab_target) - assert_shared_prediction_keys(winnow_test, nb_test_target) - assert_shared_prediction_keys(w_unlab, nb_unlab_target) - - winnow_test = winnow_test.copy() - winnow_test["correct"] = novor_correctness_mask( - winnow_test["sequence"], winnow_test["prediction"] - ) - correct_by_id = label_series_by_spectrum_id(winnow_test, "correct") - hit_by_id = label_series_by_spectrum_id(w_unlab, "proteome_hit") - nb_test_target = attach_labels_by_spectrum_id( - nb_test_target, correct_by_id, label_col="correct" - ) - nb_unlab_f = attach_labels_by_spectrum_id( - nb_unlab_target, hit_by_id, label_col="proteome_hit" - ) - - w_matched = max_score_per_peptide( - winnow_test.loc[winnow_test["correct"]], - "peptide_key", - "calibrated_confidence", - ) - w_matched_raw = max_score_per_peptide( - winnow_test.loc[winnow_test["correct"]], - "peptide_key", - "confidence", - ) - nb_correct = nb_test_target.loc[nb_test_target["correct"]].copy() - nb_matched = max_score_per_peptide(nb_correct, "peptide_key", "ALC (%)") - - w_external = max_score_per_peptide( - w_unlab.loc[~w_unlab["proteome_hit"].astype(bool)], - "peptide_key", - "calibrated_confidence", - ) - w_external_raw = max_score_per_peptide( - w_unlab.loc[~w_unlab["proteome_hit"].astype(bool)], - "peptide_key", - "confidence", - ) - nb_external = max_score_per_peptide( - nb_unlab_f.loc[~nb_unlab_f["proteome_hit"].astype(bool)], - "peptide_key", - "ALC (%)", - ) - - # Shared membership = Winnow ∩ twin-valid NovoBoard keys. - sm_keys = ( - set(w_matched["peptide_key"]) - & set(nb_matched["peptide_key"]) - & set(w_matched_raw["peptide_key"]) - ) - se_keys = ( - set(w_external["peptide_key"]) - & set(nb_external["peptide_key"]) - & set(w_external_raw["peptide_key"]) - ) - if sm_keys != set(w_matched["peptide_key"]) or sm_keys != set( - nb_matched["peptide_key"] - ): - raise AssertionError( - f"{dataset} shared S_m keys disagree after shared Novor labels: " - f"winnow={len(w_matched)} novoboard={len(nb_matched)} " - f"intersection={len(sm_keys)}" - ) - if se_keys != set(w_external["peptide_key"]) or se_keys != set( - nb_external["peptide_key"] - ): - raise AssertionError( - f"{dataset} shared S_e keys disagree after shared proteome-hit labels: " - f"winnow={len(w_external)} novoboard={len(nb_external)} " - f"intersection={len(se_keys)}" - ) - if not sm_keys: - raise ValueError(f"Empty shared matched pool for {dataset}") - if not se_keys: - raise ValueError(f"Empty shared external pool for {dataset}") - - matched = pd.DataFrame({"peptide_key": sorted(sm_keys)}) - matched = matched.merge( - w_matched[["peptide_key", "calibrated_confidence"]].rename( - columns={"calibrated_confidence": "score_winnow"} - ), - on="peptide_key", - how="left", - ) - matched = matched.merge( - nb_matched[["peptide_key", "ALC (%)"]].rename( - columns={"ALC (%)": "score_novoboard"} - ), - on="peptide_key", - how="left", - ) - matched = matched.merge( - w_matched_raw[["peptide_key", "confidence"]].rename( - columns={"confidence": "raw_confidence"} - ), - on="peptide_key", - how="left", - ) - matched["score_glissade"] = confidence_to_log_prob(matched["raw_confidence"]) - - external = pd.DataFrame({"peptide_key": sorted(se_keys)}) - external = external.merge( - w_external[["peptide_key", "calibrated_confidence"]].rename( - columns={"calibrated_confidence": "score_winnow"} - ), - on="peptide_key", - how="left", - ) - external = external.merge( - nb_external[["peptide_key", "ALC (%)"]].rename( - columns={"ALC (%)": "score_novoboard"} - ), - on="peptide_key", - how="left", - ) - external = external.merge( - w_external_raw[["peptide_key", "confidence"]].rename( - columns={"confidence": "raw_confidence"} - ), - on="peptide_key", - how="left", - ) - external["score_glissade"] = confidence_to_log_prob(external["raw_confidence"]) - - # Attach NovoBoard pair metadata for twin TDC on mixture subsets. - nb_ext_meta = nb_unlab_f.loc[ - ~nb_unlab_f["proteome_hit"].astype(bool), - ["peptide_key", "Peptide", "ALC (%)", "spectrum_id", "_pair_key", "Scan"], - ].copy() - nb_ext_meta = ( - nb_ext_meta.sort_values("ALC (%)", ascending=False) - .groupby("peptide_key", as_index=False) - .first() - ) - external = external.merge( - nb_ext_meta.rename( - columns={ - "Peptide": "peptide_novoboard", - "spectrum_id": "spectrum_id_novoboard", - } - ), - on="peptide_key", - how="left", - ) - - nb_match_meta = ( - nb_correct[ - ["peptide_key", "Peptide", "ALC (%)", "spectrum_id", "_pair_key", "Scan"] - ] - .sort_values("ALC (%)", ascending=False) - .groupby("peptide_key", as_index=False) - .first() - .rename( - columns={ - "Peptide": "peptide_novoboard", - "spectrum_id": "spectrum_id_novoboard", - } - ) - ) - matched = matched.merge(nb_match_meta, on="peptide_key", how="left") - - twin_coverage_m = ( - float(matched["_pair_key"].notna().mean()) if len(matched) else 0.0 - ) - twin_coverage_e = ( - float(external["_pair_key"].notna().mean()) if len(external) else 0.0 - ) - if twin_coverage_m < 1.0 or twin_coverage_e < 1.0: - raise AssertionError( - f"{dataset} NovoBoard twin coverage incomplete: " - f"Sm={twin_coverage_m:.3f} Se={twin_coverage_e:.3f}" - ) - - glissade_reference = build_glissade_training_reference( - cfg.calibrator_train_metadata, min_length=labelled_min_length - ) - - logger.info( - "%s shared pools: matched=%d external=%d glissade_reference=%d " - "(shared Novor/proteome-hit labels; NB twin coverage 100%%)", - dataset, - len(matched), - len(external), - len(glissade_reference), - ) - return matched, external, nb_decoy_combined, glissade_reference - - -def _estimate_winnow_q_values(mixed: pd.DataFrame) -> pd.DataFrame: - work = mixed[["peptide_key", "score_winnow", "source"]].copy() - work = work.rename(columns={"score_winnow": "score"}) - work = work.dropna(subset=["score"]) - ctrl = NonParametricFDRControl() - ctrl.fit(work["score"]) - q_table = ctrl.add_psm_q_value(work.copy(), "score") - return pd.DataFrame( - { - "peptide_key": q_table["peptide_key"], - "score": q_table["score"], - "source": q_table["source"], - "q_value": q_table["psm_q_value"], - "method": "Winnow", - } - ) - - -def _estimate_novoboard_q_values( - mixed: pd.DataFrame, - decoy_by_pair: pd.DataFrame, -) -> pd.DataFrame: - target = pd.DataFrame( - { - "Peptide": mixed["peptide_novoboard"].fillna(mixed["peptide_key"]), - "ALC (%)": mixed["score_novoboard"], - "spectrum_id": mixed.get( - "spectrum_id_novoboard", - pd.Series(np.arange(len(mixed)), dtype=str), - ), - "_pair_key": mixed["_pair_key"], - "Scan": mixed.get("Scan", mixed["_pair_key"]), - "source": mixed["source"], - "peptide_key": mixed["peptide_key"], - } - ) - target = target.dropna(subset=["ALC (%)", "_pair_key"]) - # Restrict twin TDC to the mixture peptide keys only. - table = novoboard_max_target_twin_decoy_tdc( - target, - pd.DataFrame(), - target_peptide_keys=set(target["peptide_key"].astype(str)), - decoy_by_pair=decoy_by_pair, - log_missing_twins=False, - ) - targets = table[table["is_target"]].copy() - # Map back sources for FDP. - source_map = mixed.set_index("peptide_key")["source"].to_dict() - key_col = "_peptide_key" if "_peptide_key" in targets.columns else "peptide_key" - targets["peptide_key"] = targets[key_col] - targets["source"] = targets["peptide_key"].map(source_map) - return pd.DataFrame( - { - "peptide_key": targets["peptide_key"], - "score": targets["ALC (%)"], - "source": targets["source"], - "q_value": targets["estimated_q_value"], - "method": "NovoBoard", - } - ) - - -def _estimate_glissade_q_values( - mixed: pd.DataFrame, - matched_reference: pd.DataFrame, - *, - glissade_repo: Path, - n_bootstraps: int, - rng: np.random.Generator, -) -> pd.DataFrame: - run_bootstraps, annotate_results, compute_fdr_transform = _load_glissade_functions( - glissade_repo - ) - # Glissade FDR is defined on the mixture scores; the reference is the - # training-split matched distribution and is never scored itself. - mixed_scores = mixed["score_glissade"].astype(float).to_numpy() - matched_scores = matched_reference["score_glissade"].astype(float).to_numpy() - if len(mixed_scores) < 10 or len(matched_scores) < 10: - raise ValueError("Glissade FDR requires ≥10 reference and mixture scores") - - peptides = mixed["peptide_key"].astype(str).tolist() - np.random.seed(int(rng.integers(0, 2**32 - 1))) - fdrs, grid, _ = run_bootstraps( - matched_scores, - mixed_scores, - n_bootstraps=n_bootstraps, - ) - out_peptides, peptide_fdrs, scores = annotate_results( - peptides, mixed_scores, fdrs, grid - ) - peptide_fdrs = compute_fdr_transform(peptide_fdrs) - source_map = mixed.set_index("peptide_key")["source"].to_dict() - return pd.DataFrame( - { - "peptide_key": out_peptides, - "score": scores, - "source": [source_map.get(p) for p in out_peptides], - "q_value": peptide_fdrs, - "method": "Glissade", - } - ) - - -def _mixture_result_rows( - *, - dataset: str, - pi0: float, - true_pi0: float, - holdout_frac: float, - seed: int, - iteration: int, - method: str, - q_table: pd.DataFrame, - correct_keys: set[str], - n_external: int, - thresholds: list[float], - q_ref: np.ndarray | None = None, -) -> list[dict[str, object]]: - """Build per-threshold result rows for one method on one mixture.""" - work = q_table.dropna(subset=["q_value"]) - if q_ref is not None: - if len(q_ref) != len(q_table): - raise ValueError( - f"q_ref length {len(q_ref)} does not match q_table length {len(q_table)}" - ) - q_ref_aligned = pd.Series(np.asarray(q_ref, dtype=float), index=q_table.index) - q_ref_work = q_ref_aligned.loc[work.index].to_numpy(dtype=float) - q_devs = mean_abs_q_dev_vs_reference( - work["q_value"].to_numpy(dtype=float), q_ref_work, thresholds - ) - else: - q_devs = [float("nan")] * len(thresholds) - - rows: list[dict[str, object]] = [] - for threshold, q_dev in zip(thresholds, q_devs): - accepted = work[work["q_value"] <= threshold] - n_accepted = len(accepted) - n_true = int(accepted["peptide_key"].isin(correct_keys).sum()) - n_false = n_accepted - n_true - rows.append( - { - "dataset": dataset, - "pi0_target": float(pi0), - "true_pi0": float(true_pi0), - "holdout_frac": float(holdout_frac), - "seed": seed, - "iteration": iteration, - "method": method, - "q_value_threshold": float(threshold), - "mixed_external_peptides": n_external, - "correct_peptides": len(correct_keys), - "accepted_peptides": n_accepted, - "true_correct_peptides": n_true, - "false_external_peptides": n_false, - "observed_fdp": (n_false / n_accepted if n_accepted else np.nan), - "correct_discovery_pct": ( - 100.0 * n_true / len(correct_keys) if correct_keys else np.nan - ), - "mean_abs_q_dev_vs_db": float(q_dev), - } - ) - return rows - - -def _evaluate_mixture_methods( - *, - mixed: pd.DataFrame, - estimator_reference: pd.DataFrame, - decoy_by_pair: pd.DataFrame, - glissade_repo: Path, - n_bootstraps: int, - iter_rng: np.random.Generator, - dataset: str, - pi0: float, - true_pi0: float, - holdout_frac: float, - seed: int, - iteration: int, - correct_keys: set[str], - n_external: int, - thresholds: list[float], -) -> list[dict[str, object]]: - """Run Winnow / NovoBoard / Glissade FDR on one mixture and collect rows.""" - estimators: dict[str, object] = { - "Winnow": lambda m, _r: _estimate_winnow_q_values(m), - "NovoBoard": lambda m, _r: _estimate_novoboard_q_values(m, decoy_by_pair), - "Glissade": lambda m, r, rng=iter_rng: _estimate_glissade_q_values( - m, - r, - glissade_repo=glissade_repo, - n_bootstraps=n_bootstraps, - rng=rng, - ), - } - # Raw-confidence DBG on the mixture used for NovoBoard and Glissade q-deviation. Winnow keeps calibrated-score DBG unset here. - is_correct = mixed["peptide_key"].astype(str).isin(correct_keys).to_numpy() - q_db_raw = database_grounded_q_from_labels( - mixed["score_novoboard"].to_numpy(dtype=float), is_correct - ) - q_db_raw_by_key = dict(zip(mixed["peptide_key"].astype(str), q_db_raw, strict=True)) - - rows: list[dict[str, object]] = [] - mixture_keys = set(mixed["peptide_key"].astype(str)) - for method, estimator in estimators.items(): - try: - q_table = estimator(mixed, estimator_reference) # type: ignore[operator] - except Exception as exc: # noqa: BLE001 - boundary around external tool - logger.warning( - "%s FDR failed dataset=%s pi0=%.3g iter=%d: %s", - method, - dataset, - pi0, - iteration, - exc, - ) - continue - scored_keys = set(q_table["peptide_key"].astype(str)) - if scored_keys != mixture_keys: - raise AssertionError( - f"{method} scored a different mixture on dataset={dataset} " - f"pi0={pi0:.3g} iter={iteration}: mixture={len(mixture_keys)} " - f"scored={len(scored_keys)} " - f"missing={len(mixture_keys - scored_keys)} " - f"extra={len(scored_keys - mixture_keys)}" - ) - q_ref: np.ndarray | None = None - if method in ("NovoBoard", "Glissade"): - q_ref = ( - q_table["peptide_key"] - .astype(str) - .map(q_db_raw_by_key) - .to_numpy(dtype=float) - ) - rows.extend( - _mixture_result_rows( - dataset=dataset, - pi0=pi0, - true_pi0=true_pi0, - holdout_frac=holdout_frac, - seed=seed, - iteration=iteration, - method=method, - q_table=q_table, - correct_keys=correct_keys, - n_external=n_external, - thresholds=thresholds, - q_ref=q_ref, - ) - ) - return rows - - -def _sample_correct_component( - *, - dataset: str, - matched: pd.DataFrame, - external: pd.DataFrame, - pi0_grid: list[float], - holdout_frac: float, - rng: np.random.Generator, -) -> tuple[pd.DataFrame, set[str], float]: - """Sample S_c from matched, capped so the π₀ grid fits in S_e.""" - n_from_frac = max(1, int(round(len(matched) * holdout_frac))) - n_from_frac = min(n_from_frac, len(matched)) - n_cap = max_correct_pool_for_pi0_grid(len(external), pi0_grid) - if n_cap < 1: - raise ValueError( - f"{dataset}: external pool of {len(external)} cannot support any " - f"π₀ in {pi0_grid} without replacement" - ) - n_correct = min(n_from_frac, n_cap) - if n_correct < n_from_frac: - logger.info( - "%s capping |S_c| from %d to %d so π₀ grid %s fits in |S_e|=%d " - "without replacement", - dataset, - n_from_frac, - n_correct, - pi0_grid, - len(external), - ) - correct = matched.sample(n=n_correct, random_state=int(rng.integers(0, 2**32 - 1))) - correct = correct.copy() - correct["source"] = "correct" - correct_keys = set(correct["peptide_key"]) - effective_holdout_frac = n_correct / len(matched) if len(matched) else 0.0 - overlap = correct_keys & set(external["peptide_key"]) - if overlap: - raise AssertionError( - f"{dataset} S_c and S_e share {len(overlap)} peptide keys; " - "mixture sources would be ambiguous" - ) - return correct, correct_keys, effective_holdout_frac - - -def _external_draw_size( - *, - dataset: str, - pi0: float, - n_correct: int, - n_external_pool: int, -) -> int | None: - """Return |S_e'| for π₀, or None to skip; raise if the pool is too small.""" - if not 0.0 < pi0 < 1.0: - return None - n_external = int(round(pi0 / (1.0 - pi0) * n_correct)) - if n_external < 1: - logger.warning("Skipping pi0=%.3g: requested |S_e'|=%d", pi0, n_external) - return None - if n_external > n_external_pool: - raise AssertionError( - f"{dataset} pi0={pi0:.3g}: |S_e'|={n_external} exceeds pool " - f"{n_external_pool} after |S_c| cap {n_correct}" - ) - return n_external - - -def evaluate_controlled_mixtures( - *, - dataset: str, - matched: pd.DataFrame, - external: pd.DataFrame, - nb_decoy: pd.DataFrame, - glissade_reference: pd.DataFrame, - pi0_grid: list[float], - holdout_frac: float, - seed: int, - n_iterations: int, - thresholds: list[float], - glissade_repo: Path, - n_bootstraps: int, -) -> list[dict[str, object]]: - """Evaluate all methods on shared mixtures with controlled π₀. - - Uses as much of the matched pool as possible (up to holdout_frac), capped - so every π₀ in pi0_grid can draw its null component from S_e without - replacement. The null component is drawn without replacement, so mixture - peptide keys are unique and every method realises the same π₀ on the same - rows. - """ - rng = np.random.default_rng(seed) - correct, correct_keys, effective_holdout_frac = _sample_correct_component( - dataset=dataset, - matched=matched, - external=external, - pi0_grid=pi0_grid, - holdout_frac=holdout_frac, - rng=rng, - ) - decoy_by_pair = prepare_novoboard_decoy_by_pair(nb_decoy, already_filtered=True) - - rows: list[dict[str, object]] = [] - for pi0 in pi0_grid: - n_external = _external_draw_size( - dataset=dataset, - pi0=pi0, - n_correct=len(correct_keys), - n_external_pool=len(external), - ) - if n_external is None: - continue - - for iteration in range(n_iterations): - iter_rng = np.random.default_rng( - seed + 10_000 * iteration + int(1000 * pi0) - ) - ext_sample = external.sample( - n=n_external, - replace=False, - random_state=int(iter_rng.integers(0, 2**32 - 1)), - ).copy() - ext_sample["source"] = "external" - mixed = pd.concat([ext_sample, correct], ignore_index=True, sort=False) - if mixed["peptide_key"].duplicated().any(): - raise AssertionError( - f"{dataset} mixture has duplicate peptide keys at " - f"pi0={pi0:.3g} iter={iteration}" - ) - true_pi0 = len(ext_sample) / (len(ext_sample) + len(correct_keys)) - rows.extend( - _evaluate_mixture_methods( - mixed=mixed, - estimator_reference=glissade_reference, - decoy_by_pair=decoy_by_pair, - glissade_repo=glissade_repo, - n_bootstraps=n_bootstraps, - iter_rng=iter_rng, - dataset=dataset, - pi0=pi0, - true_pi0=true_pi0, - holdout_frac=effective_holdout_frac, - seed=seed, - iteration=iteration, - correct_keys=correct_keys, - n_external=len(ext_sample), - thresholds=thresholds, - ) - ) - return rows - - -def _plot_metric_by_pi0( - dataset_results: pd.DataFrame, - *, - dataset_label: str, - dataset_slug: str, - metric: str, - ylabel: str, - base_name: str, - percent_axis: bool, - output_dir: Path, -) -> None: - method_order = list(METHODS) - colors = {"Winnow": _PALETTE[0], "NovoBoard": _PALETTE[2], "Glissade": _PALETTE[4]} - pi0_values = sorted(dataset_results["pi0_target"].dropna().unique()) - n_pi0 = max(1, len(pi0_values)) - fig, axes = plt.subplots( - 1, n_pi0, figsize=(4.2 * n_pi0, 5.5), sharey=True, squeeze=False - ) - for ax, pi0 in zip(axes[0], pi0_values): - sub = dataset_results[dataset_results["pi0_target"] == pi0] - for method in method_order: - msub = sub[sub["method"] == method] - if msub.empty: - continue - summary = ( - msub.groupby("q_value_threshold", as_index=False)[metric] - .mean(numeric_only=True) - .sort_values("q_value_threshold") - ) - ax.plot( - summary["q_value_threshold"], - summary[metric], - lw=1.5, - label=method, - color=colors[method], - ) - if metric == "observed_fdp": - max_threshold = float(sub["q_value_threshold"].max()) - ax.plot( - [0.0, max_threshold], - [0.0, max_threshold], - color="#666666", - lw=1, - ls="--", - label="Nominal FDR", - ) - ax.set_ylim(bottom=0) - if percent_axis: - ax.set_ylim(0, 100) - ax.set_xlim(0, float(sub["q_value_threshold"].max())) - ax.set_xlabel("Estimated q-value threshold") - ax.set_title(f"π₀={pi0:g}") - _style_ax(ax) - axes[0][0].set_ylabel(ylabel) - axes[0][0].legend(loc="best", fontsize=9) - fig.suptitle(f"{dataset_label} external peptide score-mixture benchmark", y=1.02) - fig.tight_layout() - _save_fig(fig, output_dir / f"{base_name}_{dataset_slug}") - - -def plot_benchmark_results(results: pd.DataFrame, output_dir: Path) -> None: - """Save FDP and recovery plots faceted by π₀.""" - if results.empty: - return - output_dir.mkdir(parents=True, exist_ok=True) - specs = [ - ("observed_fdp", "Observed FDP", "external_peptide_score_mixture_fdp", False), - ( - "correct_discovery_pct", - "Correct peptide recovery\n(% of held-out correct peptides)", - "external_peptide_score_mixture_correct_discovery_pct", - True, - ), - ] - for dataset, dataset_results in results.groupby("dataset", sort=False): - for metric, ylabel, base_name, percent_axis in specs: - _plot_metric_by_pi0( - dataset_results, - dataset_label=_display_name(str(dataset)), - dataset_slug=str(dataset).replace("/", "_"), - metric=metric, - ylabel=ylabel, - base_name=base_name, - percent_axis=percent_axis, - output_dir=output_dir, - ) - - -def write_holdout_summary_tables( - results: pd.DataFrame, - output_dir: Path, - *, - thresholds: list[float] | None = None, -) -> tuple[Path, Path]: - """Aggregate raw mixture rows into acceptance and error/gain CSVs.""" - acceptance, error_gain = summarise_holdout_results( - results, - thresholds=thresholds if thresholds is not None else SUMMARY_THRESHOLDS, - group_extra=("pi0_target",), - ) - return write_summary_tables( - acceptance, error_gain, output_dir, "external_peptide_holdout" - ) - - -@app.command() -def main( - novoboard_root: Annotated[ - Path, - typer.Option( - "--novoboard-root", - help=( - "Root of NovoBoard per-dataset tables: " - "{root}/{dataset}/novoboard/ with annotated_test*.csv and " - "raw_unlabelled*.csv target/decoy pairs (the datasets/ dir of " - "a NovoBoard checkout). Local runs used fork " - "JemmaLDaniel/NovoBoard, branch feat/adapt-to-instanovo " - "(commit a9faab3ef1af06987599c2f01e6ba96072c80172)." - ), - ), - ], - glissade_repo: Annotated[ - Path, - typer.Option( - "--glissade-repo", - help=( - "Glissade clone root (must import as glissade.glissade). Local " - "runs used fork JemmaLDaniel/glissade, branch winnow-benchmark " - "(commit 6ee11b51b5f21ba8fdc1eb5821608352b082a533)." - ), - ), - ], - output_dir: Annotated[ - Path, - typer.Option("--output-dir", help="Directory for benchmark outputs."), - ] = DEFAULT_OUTPUT_DIR, - datasets: Annotated[ - Optional[list[str]], - typer.Option("--datasets", help="Dataset keys to benchmark."), - ] = None, - pi0_grid: Annotated[ - Optional[list[float]], - typer.Option("--pi0-grid", help="Target mixture null fractions."), - ] = None, - holdout_frac: Annotated[ - float, - typer.Option( - "--holdout-frac", - help=( - "Maximum fraction of shared matched peptides used as S_c " - "(default 1 = prefer the full pool). |S_c| is further capped so " - "every π₀ in --pi0-grid fits in S_e without replacement." - ), - ), - ] = DEFAULT_HOLDOUT_FRAC, - q_thresholds: Annotated[ - Optional[list[float]], - typer.Option( - "--q-thresholds", help="Estimated q-value thresholds to evaluate." - ), - ] = None, - seed: Annotated[int, typer.Option("--seed", help="Random seed.")] = DEFAULT_SEED, - n_iterations: Annotated[ - int, - typer.Option( - "--n-iterations", help="Number of external-score resampling iterations." - ), - ] = DEFAULT_N_ITERATIONS, - winnow_results: Annotated[ - Path, - typer.Option("--winnow-results", help="Winnow results directory."), - ] = DEFAULT_WINNOW_RESULTS, - model_root: Annotated[ - Path, - typer.Option( - "--model-root", - help="Per-dataset calibrator directories, used for Glissade's anchor.", - ), - ] = DEFAULT_MODEL_ROOT, - n_bootstraps: Annotated[ - int, - typer.Option( - "--n-bootstraps", help="Glissade bootstraps per mixture iteration." - ), - ] = DEFAULT_N_BOOTSTRAPS, - min_peptide_length: Annotated[ - int, - typer.Option( - "--min-peptide-length", - help=( - "Minimum normalised peptide length for unlabelled / external " - "pools. Labelled matched pools and Glissade's training reference " - "use the labelled floor (non-empty key only)." - ), - ), - ] = MIN_PEPTIDE_LENGTH, - plot: Annotated[bool, typer.Option(help="Create summary plots.")] = True, - summarise_only: Annotated[ - Optional[Path], - typer.Option( - "--summarise-only", - help="Only write summary CSVs/plots from an existing results CSV.", - ), - ] = None, -) -> None: - """Run the controlled-π₀ external peptide score-mixture benchmark.""" - logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") - output_dir.mkdir(parents=True, exist_ok=True) - - if summarise_only is not None: - results = pd.read_csv(summarise_only) - write_holdout_summary_tables(results, output_dir) - if plot: - plot_benchmark_results(results, output_dir / "plots") - return - - dataset_keys = datasets if datasets is not None else DEFAULT_DATASETS - pi0s = pi0_grid if pi0_grid is not None else list(DEFAULT_PI0_GRID) - thresholds = ( - q_thresholds if q_thresholds is not None else list(DEFAULT_Q_THRESHOLDS) - ) - - rows: list[dict[str, object]] = [] - for dataset in dataset_keys: - logger.info("Building shared score tables for %s", _display_name(dataset)) - matched, external, nb_decoy, glissade_reference = build_shared_score_tables( - dataset=dataset, - winnow_results=winnow_results, - novoboard_root=novoboard_root, - model_root=model_root, - unlabelled_min_length=min_peptide_length, - labelled_min_length=LABELLED_MIN_PEPTIDE_LENGTH, - ) - rows.extend( - evaluate_controlled_mixtures( - dataset=dataset, - matched=matched, - external=external, - nb_decoy=nb_decoy, - glissade_reference=glissade_reference, - pi0_grid=pi0s, - holdout_frac=holdout_frac, - seed=seed, - n_iterations=n_iterations, - thresholds=thresholds, - glissade_repo=glissade_repo, - n_bootstraps=n_bootstraps, - ) - ) - - results = pd.DataFrame(rows) - results_path = output_dir / "external_peptide_holdout_results.csv" - results.to_csv(results_path, index=False) - logger.info("Wrote %s (%d rows)", results_path, len(results)) - write_holdout_summary_tables(results, output_dir) - if plot: - plot_benchmark_results(results, output_dir / "plots") - - -if __name__ == "__main__": - app() diff --git a/scripts/run_feature_ablations.py b/scripts/run_feature_ablations.py deleted file mode 100644 index 17211d9b..00000000 --- a/scripts/run_feature_ablations.py +++ /dev/null @@ -1,1532 +0,0 @@ -"""Feature ablation study for Winnow calibrator. - -Trains MLP calibrators on subsets of pre-computed training feature matrices, -computes features from raw spectra for evaluation datasets, and produces -publication-quality plots of calibration, discrimination, and FDR behavior. -""" - -from __future__ import annotations - -import json -import logging -import sys -from collections import defaultdict -from dataclasses import dataclass, field -from pathlib import Path -from typing import Annotated, Iterable, Optional - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import polars as pl -import seaborn as sns -import torch -import typer -from rich.logging import RichHandler - -_REPO_ROOT = Path(__file__).resolve().parent.parent -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from scripts.feature_subsets import FEATURE_SUBSETS # noqa: E402 -from scripts.plot_ablation_summary import ( # noqa: E402 - FDR_BIAS_COLUMN_BY_THRESHOLD, - Q_DEV_COLUMN_BY_THRESHOLD, - TAIL_ECE_COLUMN_BY_THRESHOLD, - assign_ablation_colors, - compute_ece, - compute_fdr_bias_at_fdr_thresholds, - compute_pr_auc, - compute_q_value_deviations, - compute_tail_ece_at_fdr, - ordered_ablation_configs, -) - -from winnow.calibration.calibrator import ProbabilityCalibrator # noqa: E402 -from winnow.datasets.feature_dataset import FeatureDataset # noqa: E402 -from winnow.fdr.database_grounded import DatabaseGroundedFDRControl # noqa: E402 -from winnow.fdr.nonparametric import NonParametricFDRControl # noqa: E402 - -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -logger.propagate = False -if not logger.handlers: - logger.addHandler(RichHandler()) - -app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) - -# --------------------------------------------------------------------------- -# Plot theme — Paul Tol qualitative (colour-blind safe) -# --------------------------------------------------------------------------- -_PALETTE = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"] - -sns.set_theme(style="white", palette=_PALETTE, context="paper", font_scale=1.5) - -DATASET_DISPLAY_NAMES: dict[str, str] = { - "HCT116": "Human colon", - "gluc": "HeLa degradome", - "helaqc": "HeLa single shot", - "herceptin": "Herceptin", - "immuno": "Immunopeptidomics-1", - "celegans": "$\\it{C.\\;elegans}$", - "sbrodae": "$\\it{Scalindua\\;brodae}$", - "PXD019483": "HepG2", - "snakevenoms": "Snake venomics", - "tplantibodies": "Therapeutic nanobodies", - "woundfluids": "Wound exudates", - "PXD004732": "ProteomeTools-1", - "PXD014877": "$\\it{C.\\;elegans}$", - "PXD023064": "Immunopeptidomics-2", - "astral": "Astral $\\it{E.\\;coli}$", - "01747_C01_P018218_S00_I00_N03_R1": "$\\it{Arabidopsis\\;thaliana}$", - "Arabidopsis": "$\\it{Arabidopsis\\;thaliana}$", - "20150708_QE3_UPLC8_DBJ_QC_HELA_39frac_Chymotrypsin": "HeLa chymotrypsin", - "20151020_QE3_UPLC8_DBJ_SA_A549_Rep2_46": "Human lung", - "20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46": "Human colon", - "20170303_QEh1_LC2_FaMa_ChCh_SA_HLApI_JY_R1_exp2": "HLA Class I (JY cells)", - "20170609_QEh1_LC1_ChCh_FAMA_SA_HLAIIp_JY_all_R1": "HLA Class II (JY cells)", -} - -# --------------------------------------------------------------------------- -# Feature group definitions (reduced set: no xcorr, spectral_angle, gap/similarity, edit_distance) -# --------------------------------------------------------------------------- -_EXCLUDED_REDUCED = frozenset( - { - "xcorr", - "spectral_angle", - "complementary_ion_count", - "max_ion_gap", - "edit_distance", - } -) - -BEAM_COLUMNS = ["margin", "median_margin", "entropy", "z-score", "edit_distance"] -TOKEN_COLUMNS = ["min_token_probability", "std_token_probability"] -FRAGMENT_MATCH_COLUMNS = [ - "ion_matches", - "ion_match_intensity", - "complementary_ion_count", - "max_ion_gap", - "spectral_angle", - "xcorr", -] -RETENTION_TIME_COLUMNS = ["irt_error"] -MASS_ERROR_PPM = "mass_error_ppm" -MASS_ERROR_DA = "mass_error_da" - -REDUCED_BEAM_COLUMNS = [c for c in BEAM_COLUMNS if c not in _EXCLUDED_REDUCED] -REDUCED_FRAGMENT_COLUMNS = [ - c for c in FRAGMENT_MATCH_COLUMNS if c not in _EXCLUDED_REDUCED -] - -# Default training matrix columns (train_extra_small_matrix.parquet). -REDUCED_TRAIN_COLUMNS: list[str] = FEATURE_SUBSETS["no_fragment_similarity"]["columns"] - -# Hydra overrides aligned with Makefile ANALYSIS_REDUCED_FEATURE_OVERRIDES (mass_error_da model). -REDUCED_FEATURE_COMPUTE_OVERRIDES: list[str] = [ - "~calibrator.features.mass_error", - "+calibrator.features.mass_error_da._target_=winnow.calibration.calibration_features.MassErrorDaFeature", - "+calibrator.features.mass_error_da.residue_masses=${residue_masses}", -] - - -def _reference_model_columns(model_dir: Path | None) -> list[str] | None: - """Return ``feature_columns`` from a saved calibrator, if present.""" - if model_dir is None: - return None - config_path = model_dir / "config.json" - if not config_path.is_file(): - return None - with open(config_path) as f: - config = json.load(f) - cols = config.get("feature_columns") - return list(cols) if cols else None - - -def _resolve_mass_error_column( - df: pl.DataFrame, - reference_model_dir: Path | None, -) -> str: - """Pick mass-error column present in *df*, preferring the reference model.""" - ref_cols = _reference_model_columns(reference_model_dir) - if MASS_ERROR_DA in df.columns: - return MASS_ERROR_DA - if MASS_ERROR_PPM in df.columns: - if ref_cols and MASS_ERROR_DA in ref_cols: - logger.warning( - "Reference model uses %s but data has %s; using %s for ablations.", - MASS_ERROR_DA, - MASS_ERROR_PPM, - MASS_ERROR_PPM, - ) - return MASS_ERROR_PPM - raise ValueError( - f"No mass error column in data (tried {MASS_ERROR_DA}, {MASS_ERROR_PPM})" - ) - - -def _columns_available(df: pl.DataFrame, columns: list[str]) -> list[str]: - missing = [c for c in columns if c not in df.columns] - if missing: - raise ValueError(f"Missing columns: {missing}. Available: {df.columns}") - return columns - - -def resolve_all_feature_columns( - df: pl.DataFrame, - reference_model_dir: Path | None, -) -> list[str]: - """Full reduced feature set for the 'All features' ablation config.""" - ref_cols = _reference_model_columns(reference_model_dir) - if ref_cols: - cols = ["confidence"] - for col in ref_cols: - if ( - col == MASS_ERROR_DA - and col not in df.columns - and MASS_ERROR_PPM in df.columns - ): - cols.append(MASS_ERROR_PPM) - elif col in df.columns: - cols.append(col) - else: - cols = [c for c in REDUCED_TRAIN_COLUMNS if c in df.columns] - return _columns_available(df, cols) - - -def build_ablation_configs( - df: pl.DataFrame, - reference_model_dir: Path | None, -) -> dict[str, list[str]]: - """Build ablation configs using columns available in *df*.""" - mass_col = _resolve_mass_error_column(df, reference_model_dir) - all_features = resolve_all_feature_columns(df, reference_model_dir) - return { - "Confidence only": ["confidence"], - "Confidence + mass error": ["confidence", mass_col], - "Confidence + iRT error": ["confidence", *RETENTION_TIME_COLUMNS], - "Confidence + token-level": ["confidence", *TOKEN_COLUMNS], - "Confidence + beam search": ["confidence", *REDUCED_BEAM_COLUMNS], - "Confidence + fragment matching": ["confidence", *REDUCED_FRAGMENT_COLUMNS], - "All features": all_features, - } - - -ABLATION_CONFIGS: dict[str, list[str]] = {} - -ABLATION_COLORS: dict[str, str] = {} - - -def _dataset_display_name(key: str) -> str: - """Publication-ready dataset label for plot titles.""" - if key in EVAL_DATASETS: - return str(EVAL_DATASETS[key]["label"]) - return DATASET_DISPLAY_NAMES.get(key, key) - - -def _configure_ablation_colors(config_names: Iterable[str]) -> None: - global ABLATION_COLORS - ABLATION_COLORS = assign_ablation_colors( - ordered_ablation_configs(set(config_names)) - ) - - -# Default training hyperparameters (overridden by --hyperparams-from-model). -TRAIN_HYPERPARAMS = { - "hidden_dims": [128, 64], - "learning_rate": 0.0001, - "weight_decay": 0.0001, - "batch_size": 4096, - "max_epochs": 200, - "n_iter_no_change": 10, - "tol": 1e-4, -} - - -def train_hyperparams_from_model(model_dir: Path) -> dict[str, object]: - """Load MLP training hyperparameters from a saved calibrator ``config.json``.""" - config_path = model_dir / "config.json" - if not config_path.is_file(): - raise FileNotFoundError(f"No config.json at {model_dir}") - with open(config_path) as f: - config = json.load(f) - return { - "hidden_dims": tuple(config["hidden_dims"]), - "dropout": config["dropout"], - "learning_rate": config["learning_rate"], - "weight_decay": config["weight_decay"], - "batch_size": config["batch_size"], - "max_epochs": config["max_epochs"], - "n_iter_no_change": config["n_iter_no_change"], - "tol": config["tol"], - } - - -EVAL_DATASETS = { - "HCT116": { - "label": "Human colon", - "spectra": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.parquet", - "predictions": "new_eval_data/lcfm/PXD004452/20151020_QE3_UPLC8_DBJ_SA_HCT116_Rep2_46.csv", - "koina_mode": "columns", - }, - "Arabidopsis": { - "label": "Arabidopsis", - "spectra": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.parquet", - "predictions": "new_eval_data/lcfm/PXD013868/01747_C01_P018218_S00_I00_N03_R1.csv", - "koina_mode": "columns", - }, - "PXD023064": { - "label": "Immunopeptidomics-2", - "spectra": "held_out_projects/lcfm/PXD023064/", - "predictions": "held_out_projects/lcfm/PXD023064_predictions/PXD023064.csv", - "koina_mode": "columns", - }, -} - -# Residue masses for DatabaseGroundedFDRControl (loaded from config at runtime) -_RESIDUE_MASSES: dict[str, float] | None = None - - -def _get_residue_masses() -> dict[str, float]: - """Load residue masses from the winnow residues config.""" - global _RESIDUE_MASSES - if _RESIDUE_MASSES is None: - import yaml - - config_path = ( - Path(__file__).resolve().parent.parent - / "winnow" - / "configs" - / "residues.yaml" - ) - with open(config_path) as f: - cfg = yaml.safe_load(f) - _RESIDUE_MASSES = cfg["residue_masses"] - return _RESIDUE_MASSES - - -# --------------------------------------------------------------------------- -# Eval feature computation -# --------------------------------------------------------------------------- -def _feature_compute_overrides(reference_model_dir: Path | None) -> list[str]: - """Hydra overrides so eval features match a mass_error_da / reduced-feature model.""" - ref_cols = _reference_model_columns(reference_model_dir) - if ref_cols and MASS_ERROR_DA in ref_cols: - return list(REDUCED_FEATURE_COMPUTE_OVERRIDES) - return [] - - -def _compute_eval_features_for_dataset( - name: str, - spectra_path: str, - predictions_path: str, - cache_dir: Path, - koina_url: str, - koina_ssl: bool, - koina_mode: str = "columns", - feature_overrides: list[str] | None = None, -) -> Path: - """Compute the full feature matrix for an eval dataset and cache as Parquet. - - Args: - koina_mode: ``"columns"`` to read collision_energy / frag_type from - per-row metadata columns, or ``"constants"`` to use fixed values - (CE=27, HCD). - """ - cache_path = cache_dir / f"{name}.parquet" - if cache_path.exists(): - logger.info("Using cached eval features for %s at %s", name, cache_path) - return cache_path - - from hydra import compose, initialize_config_dir - from hydra.utils import instantiate - - from winnow.utils.config_path import get_primary_config_dir - - primary_config_dir = get_primary_config_dir(None) - - logger.info("Computing features for eval dataset %s ...", name) - - if koina_mode == "columns": - koina_overrides = [ - "+koina.input_columns.collision_energies=collision_energy", - "+koina.input_columns.fragmentation_types=frag_type", - "+calibrator.features.fragment_match_features.model_input_columns.collision_energies=collision_energy", - "+calibrator.features.fragment_match_features.model_input_columns.fragmentation_types=frag_type", - ] - else: - koina_overrides = [ - "+koina.input_constants.collision_energies=27", - "+koina.input_constants.fragmentation_types=HCD", - "+calibrator.features.fragment_match_features.model_input_constants.collision_energies=27", - "+calibrator.features.fragment_match_features.model_input_constants.fragmentation_types=HCD", - ] - - with initialize_config_dir( - config_dir=str(primary_config_dir), - version_base="1.3", - job_name=f"winnow_ablation_features_{name}", - ): - cfg = compose( - config_name="compute_features", - overrides=[ - f"dataset.spectrum_path_or_directory={spectra_path}", - f"dataset.predictions_path={predictions_path}", - f"koina.server_url={koina_url}", - f"koina.ssl={koina_ssl}", - *koina_overrides, - *(feature_overrides or []), - "labelled=true", - "filter_empty_predictions=true", - ], - ) - - data_loader = instantiate(cfg.data_loader) - calibrator = instantiate(cfg.calibrator) - - from winnow.scripts.main import ( - _compute_features_batched_metadata, - ) - - spectrum_path = Path(spectra_path) - preds_path = cfg.dataset.get("predictions_path", predictions_path) - - all_metadata = _compute_features_batched_metadata( - spectrum_path, - preds_path, - data_loader, - calibrator, - labelled=True, - ) - - combined_metadata = pd.concat(all_metadata, ignore_index=True) - logger.info( - " %s: %d spectra after feature computation", name, len(combined_metadata) - ) - - # Write the training matrix parquet with all feature columns + correct + extra cols for FDR - feature_columns = ["confidence"] + calibrator.columns - keep_cols = list(feature_columns) - if "correct" in combined_metadata.columns: - keep_cols.append("correct") - if "sequence" in combined_metadata.columns: - keep_cols.append("sequence") - if "prediction" in combined_metadata.columns: - keep_cols.append("prediction") - if "precursor_mz" in combined_metadata.columns: - keep_cols.append("precursor_mz") - if "precursor_charge" in combined_metadata.columns: - keep_cols.append("precursor_charge") - - # Deduplicate while preserving order - seen = set() - unique_cols = [] - for c in keep_cols: - if c not in seen and c in combined_metadata.columns: - seen.add(c) - unique_cols.append(c) - - training_df = pl.from_pandas(combined_metadata[unique_cols]) - cache_dir.mkdir(parents=True, exist_ok=True) - training_df.write_parquet(cache_path) - logger.info( - " Cached eval features to %s (%d rows, %d cols)", - cache_path, - len(training_df), - len(training_df.columns), - ) - return cache_path - - -def compute_all_eval_features( - output_dir: Path, - koina_url: str, - koina_ssl: bool, - astral_spectra: str | None, - astral_predictions: str | None, - skip_feature_compute: bool, - reference_model_dir: Path | None = None, -) -> dict[str, Path]: - """Compute (or locate cached) eval feature Parquets for all datasets.""" - cache_dir = output_dir / "eval_feature_cache" - result: dict[str, Path] = {} - feature_overrides = _feature_compute_overrides(reference_model_dir) - - for name, info in EVAL_DATASETS.items(): - if skip_feature_compute: - cache_path = cache_dir / f"{name}.parquet" - if not cache_path.exists(): - raise FileNotFoundError( - f"--skip-feature-compute set but cache not found: {cache_path}" - ) - result[name] = cache_path - else: - result[name] = _compute_eval_features_for_dataset( - name, - info["spectra"], - info["predictions"], - cache_dir, - koina_url, - koina_ssl, - koina_mode=info.get("koina_mode", "columns"), - feature_overrides=feature_overrides, - ) - - if astral_spectra and astral_predictions: - name = "Astral" - if skip_feature_compute: - cache_path = cache_dir / f"{name}.parquet" - if not cache_path.exists(): - raise FileNotFoundError( - f"--skip-feature-compute set but cache not found: {cache_path}" - ) - result[name] = cache_path - else: - result[name] = _compute_eval_features_for_dataset( - name, - astral_spectra, - astral_predictions, - cache_dir, - koina_url, - koina_ssl, - feature_overrides=feature_overrides, - ) - - return result - - -# --------------------------------------------------------------------------- -# Training -# --------------------------------------------------------------------------- -def _load_parquet_as_polars(path: str | Path) -> pl.DataFrame: - """Load a Parquet file or directory of Parquets into a single Polars DataFrame.""" - path = Path(path) - if path.is_dir(): - parquet_files = sorted(path.glob("*.parquet")) - if not parquet_files: - raise FileNotFoundError(f"No .parquet files in {path}") - return pl.concat([pl.read_parquet(f) for f in parquet_files]) - return pl.read_parquet(path) - - -def split_train_val_frames( - df: pl.DataFrame, - validation_fraction: float, - seed: int, -) -> tuple[pl.DataFrame, pl.DataFrame]: - """Random train/validation split with a fixed index permutation. - - Uses the same scheme as ``winnow.scripts.main._maybe_split_calibration_dataset``: - shuffle all row indices with *seed*, then assign the last ``validation_fraction`` - fraction to validation. The same split is reused for every ablation config. - """ - if "correct" not in df.columns: - raise ValueError("Training Parquet must contain a 'correct' column") - if not 0 < validation_fraction < 1: - raise ValueError( - f"validation_fraction must be in (0, 1), got {validation_fraction}" - ) - - n = len(df) - n_val = max(1, int(n * validation_fraction)) - rng = np.random.default_rng(seed) - perm = rng.permutation(n) - train_df = df[perm[: n - n_val].tolist()] - val_df = df[perm[n - n_val :].tolist()] - logger.info( - "Train/val split: %d train, %d val (fraction=%.2f, seed=%d)", - len(train_df), - len(val_df), - validation_fraction, - seed, - ) - return train_df, val_df - - -def _column_slice_to_feature_dataset( - df: pl.DataFrame, columns: list[str] -) -> FeatureDataset: - """Select columns from a Polars DataFrame and build a FeatureDataset.""" - if "correct" not in df.columns: - raise ValueError("Parquet must contain a 'correct' column") - missing = [c for c in columns if c not in df.columns] - if missing: - raise ValueError( - f"Missing columns in Parquet: {missing}. Available: {df.columns}" - ) - features = df.select(columns).to_numpy().astype(np.float32) - labels = df["correct"].to_numpy().astype(np.float32) - non_confidence = [c for c in columns if c != "confidence"] - return FeatureDataset(features=features, labels=labels, columns=non_confidence) - - -def _config_dir_name(config_name: str) -> str: - """Derive the on-disk directory name for an ablation config.""" - return config_name.lower().replace(" ", "_").replace("+", "and") - - -_LEGACY_DIR_NAMES: dict[str, list[str]] = { - "Confidence only": ["confidence_only"], - "Confidence + mass error": [ - "confidence_and_mass_error", - "confidence_and_mass_error_and_rt", - ], - "Confidence + iRT error": [ - "confidence_and_irt_error", - "confidence_and_mass_error_and_rt", - "confidence_and_fragment_matching", - "prosit", - ], - "Confidence + token-level": [ - "confidence_and_token_level", - "confidence_and_beam_search", - "beam_and_token", - ], - "Confidence + beam search": ["confidence_and_beam_search", "beam_and_token"], - "Confidence + fragment matching": [ - "confidence_and_fragment_matching", - "prosit", - ], - "All features": ["all_features", "full_model"], -} - - -def _resolve_model_dir(output_dir: Path, config_name: str) -> Path: - """Find the model directory for a config, falling back to legacy names.""" - candidates = _LEGACY_DIR_NAMES.get(config_name, [_config_dir_name(config_name)]) - for candidate in candidates: - model_dir = output_dir / "models" / candidate - if model_dir.exists(): - return model_dir - raise FileNotFoundError( - f"No saved model found for '{config_name}'. " - f"Searched: {[str(output_dir / 'models' / c) for c in candidates]}" - ) - - -def train_ablation_models( - train_df: pl.DataFrame, - val_df: pl.DataFrame, - output_dir: Path, - seed: int, - train_hyperparams: dict[str, object] | None = None, -) -> dict[str, ProbabilityCalibrator]: - """Train one calibrator per ablation config, return dict of fitted calibrators.""" - models: dict[str, ProbabilityCalibrator] = {} - hp = {**TRAIN_HYPERPARAMS, **(train_hyperparams or {})} - - for config_name, columns in ABLATION_CONFIGS.items(): - logger.info( - "Training ablation config: %s (%d features)", config_name, len(columns) - ) - - train_ds = _column_slice_to_feature_dataset(train_df, columns) - val_ds = _column_slice_to_feature_dataset(val_df, columns) - - calibrator = ProbabilityCalibrator( - seed=seed, - **hp, # type: ignore[arg-type] - ) - history = calibrator.fit_from_features(train_ds, val_ds) - - model_dir = output_dir / "models" / _config_dir_name(config_name) - ProbabilityCalibrator.save(calibrator, model_dir) - logger.info( - " Trained %s: %d epochs, best_epoch=%d", - config_name, - history.epochs_trained, - history.best_epoch, - ) - - models[config_name] = calibrator - - return models - - -def load_ablation_models( - output_dir: Path, -) -> dict[str, ProbabilityCalibrator]: - """Load pre-trained ablation calibrators from ``{output_dir}/models/``.""" - models: dict[str, ProbabilityCalibrator] = {} - - for config_name in ABLATION_CONFIGS: - model_dir = _resolve_model_dir(output_dir, config_name) - calibrator = ProbabilityCalibrator.load(model_dir) - logger.info(" Loaded %s from %s", config_name, model_dir) - models[config_name] = calibrator - - return models - - -# --------------------------------------------------------------------------- -# Evaluation helpers -# --------------------------------------------------------------------------- -def _predict_calibrated_scores( - calibrator: ProbabilityCalibrator, - features: np.ndarray, -) -> np.ndarray: - """Run forward pass through a fitted calibrator and return calibrated probabilities.""" - assert calibrator.network is not None - assert calibrator.feature_mean is not None - assert calibrator.feature_std is not None - - device = next(calibrator.network.parameters()).device - x = torch.as_tensor(features, dtype=torch.float32, device=device) - x = (x - calibrator.feature_mean) / calibrator.feature_std - - calibrator.network.eval() - with torch.no_grad(): - logits = calibrator.network(x) - probs = torch.sigmoid(logits).cpu().numpy().flatten() - - return probs - - -def compute_precision_recall_curve( - dataset: pd.DataFrame, - confidence_column: str, - label_column: str, - name: str, -) -> pd.DataFrame: - """Non-standard cumulative PR curve matching the casanovo notebook.""" - original = dataset[[confidence_column, label_column]] - original = original.sort_values(by=confidence_column, ascending=False) - cum_correct = np.cumsum(original[label_column].values) - precision = cum_correct / np.arange(1, len(original) + 1) - recall = cum_correct / len(original) - metrics = pd.DataFrame({"precision": precision, "recall": recall}).reset_index( - drop=True - ) - metrics["name"] = name - return metrics - - -def compute_calibration_curve( - df: pd.DataFrame, - pred_col: str, - label_col: str, - name: str, - n_bins: int = 10, -) -> pd.DataFrame: - """Fixed-width bin calibration curve matching the casanovo notebook.""" - data = df[[pred_col, label_col]].dropna().copy(deep=True) - data[pred_col] = data[pred_col].clip(0.0, 1.0) - bins = np.linspace(0.0, 1.0, n_bins + 1) - bin_cats = pd.cut(data[pred_col], bins=bins, include_lowest=True) - bin_cats.name = "bin" - grouped = ( - data.groupby(bin_cats, observed=True) - .agg( - pred_mean=(pred_col, "mean"), - empirical=(label_col, "mean"), - count=(label_col, "size"), - ) - .reset_index() - ) - grouped = grouped[grouped["count"] > 0] - grouped["bin_center"] = grouped["bin"].apply(lambda iv: (iv.left + iv.right) / 2) - grouped["name"] = name - return grouped[["pred_mean", "empirical", "count", "bin_center", "name"]] - - -def compute_brier_score(pred: np.ndarray, labels: np.ndarray) -> float: - """Brier score.""" - return float(np.mean((pred - labels) ** 2)) - - -def compute_ids_at_fdr( - calibrated_scores: np.ndarray, - labels: np.ndarray, - fdr_threshold: float, -) -> int: - """Count PSMs accepted at a given FDR threshold using NonParametricFDRControl.""" - fdr_ctrl = NonParametricFDRControl() - scores_series = pd.Series(calibrated_scores, name="score") - fdr_ctrl.fit(dataset=scores_series) - cutoff = fdr_ctrl.get_confidence_cutoff(threshold=fdr_threshold) - if np.isnan(cutoff): - return 0 - return int((calibrated_scores >= cutoff).sum()) - - -@dataclass -class EvalResult: - """Metrics and curves for a single ablation config evaluated on one dataset.""" - - config_name: str - dataset_name: str - ece: float - tail_ece_at_5pct: float - tail_ece_at_10pct: float - brier: float - ids_at_1pct: int - ids_at_5pct: int - ids_at_10pct: int - pr_auc: float - fdr_bias_at_5pct: float - fdr_bias_at_10pct: float - q_dev_at_5pct: float - q_dev_at_10pct: float - pr_curve: pd.DataFrame = field(repr=False) - calibration_curve: pd.DataFrame = field(repr=False) - calibrated_scores: np.ndarray = field(repr=False) - labels: np.ndarray = field(repr=False) - raw_confidence: np.ndarray = field(repr=False) - eval_df: pd.DataFrame = field(repr=False) - - -def evaluate_single( - config_name: str, - calibrator: ProbabilityCalibrator, - columns: list[str], - eval_df: pl.DataFrame, - dataset_name: str, -) -> EvalResult: - """Evaluate a single ablation config on a single eval dataset.""" - features = eval_df.select(columns).to_numpy().astype(np.float32) - labels = eval_df["correct"].to_numpy().astype(np.float32) - raw_confidence = eval_df["confidence"].to_numpy().astype(np.float64) - - calibrated = _predict_calibrated_scores(calibrator, features) - - # Build a pandas DataFrame for PR / calibration / FDR computations - meta = pd.DataFrame( - { - "confidence": raw_confidence, - "calibrated_confidence": calibrated, - "correct": labels, - } - ) - - # Carry over sequence and prediction for database-grounded FDR if available - if "sequence" in eval_df.columns: - meta["sequence"] = eval_df["sequence"].to_pandas() - if "prediction" in eval_df.columns: - meta["prediction"] = eval_df["prediction"].to_pandas() - - pr = compute_precision_recall_curve( - meta, "calibrated_confidence", "correct", config_name - ) - - cal = compute_calibration_curve( - meta, "calibrated_confidence", "correct", config_name - ) - - ece = compute_ece(calibrated, labels) - fdr_ctrl = NonParametricFDRControl() - fdr_ctrl.fit(dataset=pd.Series(calibrated, name="score")) - tail_ece_5 = compute_tail_ece_at_fdr(calibrated, labels, 0.05, fdr_ctrl=fdr_ctrl) - tail_ece_10 = compute_tail_ece_at_fdr(calibrated, labels, 0.10, fdr_ctrl=fdr_ctrl) - brier = compute_brier_score(calibrated, labels) - - ids_1 = compute_ids_at_fdr(calibrated, labels, 0.01) - ids_5 = compute_ids_at_fdr(calibrated, labels, 0.05) - ids_10 = compute_ids_at_fdr(calibrated, labels, 0.10) - - pr_auc = compute_pr_auc(meta) - fdr_bias = compute_fdr_bias_at_fdr_thresholds(meta) - q_dev = compute_q_value_deviations(meta) - - return EvalResult( - config_name=config_name, - dataset_name=dataset_name, - ece=ece, - tail_ece_at_5pct=tail_ece_5, - tail_ece_at_10pct=tail_ece_10, - brier=brier, - ids_at_1pct=ids_1, - ids_at_5pct=ids_5, - ids_at_10pct=ids_10, - pr_auc=pr_auc, - fdr_bias_at_5pct=fdr_bias[0.05], - fdr_bias_at_10pct=fdr_bias[0.10], - q_dev_at_5pct=q_dev[0.05], - q_dev_at_10pct=q_dev[0.10], - pr_curve=pr, - calibration_curve=cal, - calibrated_scores=calibrated, - labels=labels, - raw_confidence=raw_confidence, - eval_df=meta, - ) - - -# --------------------------------------------------------------------------- -# Plotting -# --------------------------------------------------------------------------- -def _style_axes(ax: plt.Axes) -> None: - """Apply standard axes formatting: no grid, black spines.""" - ax.set_axisbelow(True) - ax.grid(False) - for spine in ax.spines.values(): - spine.set_edgecolor("black") - spine.set_linewidth(0.8) - - -def _save_fig(fig: plt.Figure, base_path: Path, plot_format: str) -> None: - """Save figure in the requested format(s).""" - if plot_format in ("pdf", "both"): - fig.savefig(f"{base_path}.pdf", bbox_inches="tight", dpi=300) - if plot_format in ("png", "both"): - fig.savefig(f"{base_path}.png", bbox_inches="tight", dpi=300) - plt.close(fig) - - -def _lineplot( - ax: plt.Axes, - data: pd.DataFrame, - *, - x: str, - y: str, - label: str, - color: str, - linestyle: str = "-", - linewidth: float = 0.5, - marker: str | None = None, -) -> None: - """Line plot with consistent linewidth (seaborn, no auto legend).""" - kwargs: dict = { - "data": data, - "x": x, - "y": y, - "label": label, - "color": color, - "linestyle": linestyle, - "linewidth": linewidth, - "ax": ax, - "legend": False, - } - if marker is not None: - kwargs["marker"] = marker - sns.lineplot(**kwargs) - - -def _generate_plots_for_dataset( - ds_results: list[EvalResult], - ds_name: str, - plots_dir: Path, - plot_format: str, -) -> None: - """Generate all ablation figures for one dataset.""" - plot_precision_recall(ds_results, ds_name, plots_dir, plot_format) - plot_calibration(ds_results, ds_name, plots_dir, plot_format) - plot_fdr_vs_confidence(ds_results, ds_name, plots_dir, plot_format) - plot_fdr_accepted_psms(ds_results, ds_name, plots_dir, plot_format) - - -def plot_precision_recall( - results: list[EvalResult], - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """PR curve: one line per ablation config.""" - fig, ax = plt.subplots(figsize=(6, 4)) - - for r in results: - _lineplot( - ax, - r.pr_curve, - x="recall", - y="precision", - label=r.config_name, - color=ABLATION_COLORS[r.config_name], - ) - - display = _dataset_display_name(dataset_name) - ax.set( - xlabel="Recall", - ylabel="Precision", - title=f"{display} precision-recall by feature set", - ) - ax.legend(loc="lower left", fontsize=7) - _style_axes(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"pr_curve_{dataset_name}", plot_format) - - -def plot_calibration( - results: list[EvalResult], - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """Calibration diagram: reliability curves + diagonal.""" - fig, ax = plt.subplots(figsize=(6, 4)) - - for r in results: - _lineplot( - ax, - r.calibration_curve, - x="pred_mean", - y="empirical", - label=r.config_name, - color=ABLATION_COLORS[r.config_name], - marker="o", - ) - - display = _dataset_display_name(dataset_name) - ax.plot([0, 1], [0, 1], ls="--", color="gray", lw=0.5) - ax.set( - xlabel="Mean predicted probability", - ylabel="Empirical accuracy\n(database label)", - title=f"{display} probability calibration by feature set", - ) - ax.legend(loc="lower right", fontsize=7) - _style_axes(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"calibration_{dataset_name}", plot_format) - - -def plot_fdr_vs_confidence( - results: list[EvalResult], - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """PSM FDR vs calibrated confidence: non-parametric vs database-grounded per config.""" - n_configs = len(results) - fig, axes = plt.subplots(1, n_configs, figsize=(5 * n_configs, 4), squeeze=False) - - for i, r in enumerate(results): - ax = axes[0, i] - - np_fdr = NonParametricFDRControl() - np_fdr.fit(dataset=r.eval_df["calibrated_confidence"]) - winnow_metrics = np_fdr.add_psm_fdr( - r.eval_df.copy(), confidence_col="calibrated_confidence" - ) - - has_sequence = ( - "sequence" in r.eval_df.columns and "prediction" in r.eval_df.columns - ) - - if has_sequence: - dbg_fdr = DatabaseGroundedFDRControl( - confidence_feature="calibrated_confidence", - ) - try: - sorted_df = r.eval_df.sort_values( - "calibrated_confidence", ascending=False - ) - labels = sorted_df["correct"].astype(float).to_numpy() - conf = sorted_df["calibrated_confidence"].to_numpy() - drop = 10 - precision = np.cumsum(labels) / np.arange(1, len(labels) + 1) - dbg_fdr._fdr_values = np.array(1.0 - precision)[drop:] - dbg_fdr._confidence_scores = conf[drop:] - dbg_metrics = dbg_fdr.add_psm_fdr( - r.eval_df.copy(), confidence_col="calibrated_confidence" - ) - - sns.lineplot( - x=np.asarray(dbg_metrics["calibrated_confidence"], dtype=float), - y=np.asarray(dbg_metrics["psm_fdr"], dtype=float), - label="Database-grounded", - ax=ax, - color=_PALETTE[3], - linewidth=0.5, - legend=False, - ) - except Exception as e: - logger.warning( - "Database-grounded FDR failed for %s/%s: %s", - r.config_name, - dataset_name, - e, - ) - - sns.lineplot( - x=np.asarray(winnow_metrics["calibrated_confidence"], dtype=float), - y=np.asarray(winnow_metrics["psm_fdr"], dtype=float), - label="Winnow (non-parametric)", - ax=ax, - color=_PALETTE[0], - linewidth=0.5, - legend=False, - ) - - ax.set_xlabel("Calibrated confidence") - ax.set_ylabel("PSM FDR") - ax.set_title(r.config_name) - ax.legend(fontsize=7) - _style_axes(ax) - - display = _dataset_display_name(dataset_name) - fig.suptitle( - f"{display} PSM FDR vs calibrated confidence by feature set", fontsize=12 - ) - fig.tight_layout() - _save_fig(fig, output_dir / f"fdr_vs_confidence_{dataset_name}", plot_format) - - -def plot_fdr_accepted_psms( - results: list[EvalResult], - dataset_name: str, - output_dir: Path, - plot_format: str, -) -> None: - """Number of accepted PSMs vs q-value threshold.""" - fig, ax = plt.subplots(figsize=(6, 4)) - - thresholds = np.linspace(0.001, 0.10, 200) - - for r in results: - np_fdr = NonParametricFDRControl() - scores_series = pd.Series(r.calibrated_scores, name="score") - np_fdr.fit(dataset=scores_series) - - meta_with_q = np_fdr.add_psm_q_value( - pd.DataFrame({"calibrated_confidence": r.calibrated_scores}), - confidence_col="calibrated_confidence", - ) - - q_values = meta_with_q["psm_q_value"].values - counts = [] - for t in thresholds: - counts.append(int((q_values <= t).sum())) - - ax.plot( - thresholds, - counts, - label=r.config_name, - color=ABLATION_COLORS[r.config_name], - linewidth=0.5, - ) - - for fdr_line in [0.01, 0.05, 0.10]: - ax.axvline(fdr_line, ls="--", color="gray", lw=0.5, alpha=0.7) - - ax.relim() - ax.autoscale_view() - y_text = ax.get_ylim()[1] * 0.02 - for fdr_line in [0.01, 0.05, 0.10]: - ax.text( - fdr_line - 0.002, - y_text, - f"{fdr_line:.0%}", - ha="right", - va="bottom", - fontsize=7, - color="gray", - ) - - display = _dataset_display_name(dataset_name) - ax.set_xlabel("Non-parametric q-value threshold") - ax.set_ylabel("Accepted PSMs") - ax.set_title(f"{display} accepted PSMs at non-parametric q-value threshold") - ax.legend(loc="upper left", fontsize=7) - _style_axes(ax) - fig.tight_layout() - _save_fig(fig, output_dir / f"fdr_accepted_psms_{dataset_name}", plot_format) - - -# --------------------------------------------------------------------------- -# Saving eval results -# --------------------------------------------------------------------------- -def save_eval_results(all_results: list[EvalResult], output_dir: Path) -> None: - """Persist per-PSM eval DataFrames so plots can be reproduced without re-inference.""" - results_dir = output_dir / "eval_results" - results_dir.mkdir(parents=True, exist_ok=True) - - for r in all_results: - safe_config = r.config_name.lower().replace(" ", "_").replace("+", "and") - path = results_dir / f"{r.dataset_name}_{safe_config}.parquet" - df = r.eval_df.copy() - df["config_name"] = r.config_name - df["dataset_name"] = r.dataset_name - df.to_parquet(path, index=False) - - logger.info("Saved %d eval result Parquets to %s", len(all_results), results_dir) - - -def load_eval_results_for_plotting( - output_dir: Path, -) -> dict[str, list[EvalResult]]: - """Load saved eval Parquets and rebuild curve data for plotting.""" - results_dir = output_dir / "eval_results" - if not results_dir.is_dir(): - raise FileNotFoundError(f"No eval_results directory at {results_dir}") - - paths = sorted(results_dir.glob("*.parquet")) - if not paths: - raise FileNotFoundError(f"No eval result Parquets in {results_dir}") - - grouped: dict[str, list[EvalResult]] = defaultdict(list) - for path in paths: - df = pd.read_parquet(path) - config_name = str(df["config_name"].iloc[0]) - dataset_name = str(df["dataset_name"].iloc[0]) - meta = df.drop(columns=["config_name", "dataset_name"], errors="ignore") - calibrated = meta["calibrated_confidence"].to_numpy(dtype=np.float64) - labels = meta["correct"].to_numpy(dtype=np.float32) - pr = compute_precision_recall_curve( - meta, "calibrated_confidence", "correct", config_name - ) - cal = compute_calibration_curve( - meta, "calibrated_confidence", "correct", config_name - ) - grouped[dataset_name].append( - EvalResult( - config_name=config_name, - dataset_name=dataset_name, - ece=0.0, - tail_ece_at_5pct=float("nan"), - tail_ece_at_10pct=float("nan"), - brier=0.0, - ids_at_1pct=0, - ids_at_5pct=0, - ids_at_10pct=0, - pr_auc=0.0, - fdr_bias_at_5pct=float("nan"), - fdr_bias_at_10pct=float("nan"), - q_dev_at_5pct=float("nan"), - q_dev_at_10pct=float("nan"), - pr_curve=pr, - calibration_curve=cal, - calibrated_scores=calibrated, - labels=labels, - raw_confidence=meta["confidence"].to_numpy(dtype=np.float64), - eval_df=meta, - ) - ) - - for ds_name in grouped: - grouped[ds_name].sort(key=lambda r: r.config_name) - - return dict(grouped) - - -def _run_plots_only(output_dir: Path, plot_format: str) -> None: - """Regenerate plots from ``{output_dir}/eval_results`` without inference.""" - plots_dir = output_dir / "plots" - plots_dir.mkdir(parents=True, exist_ok=True) - - grouped = load_eval_results_for_plotting(output_dir) - config_names = [r.config_name for results in grouped.values() for r in results] - _configure_ablation_colors(config_names) - - for ds_name in sorted(grouped): - ds_results = grouped[ds_name] - logger.info("Generating plots for %s (%d configs)...", ds_name, len(ds_results)) - _generate_plots_for_dataset(ds_results, ds_name, plots_dir, plot_format) - - logger.info("Plots saved to %s", plots_dir) - - -# --------------------------------------------------------------------------- -# Summary -# --------------------------------------------------------------------------- -def build_summary_table(all_results: list[EvalResult]) -> pd.DataFrame: - """Aggregate all EvalResults into a single summary DataFrame.""" - rows = [] - for r in all_results: - rows.append( - { - "config": r.config_name, - "dataset": r.dataset_name, - "ECE": round(r.ece, 5), - TAIL_ECE_COLUMN_BY_THRESHOLD[0.05]: round(r.tail_ece_at_5pct, 5), - TAIL_ECE_COLUMN_BY_THRESHOLD[0.10]: round(r.tail_ece_at_10pct, 5), - "Brier": round(r.brier, 5), - "PR_AUC": round(r.pr_auc, 5), - FDR_BIAS_COLUMN_BY_THRESHOLD[0.05]: round(r.fdr_bias_at_5pct, 5), - FDR_BIAS_COLUMN_BY_THRESHOLD[0.10]: round(r.fdr_bias_at_10pct, 5), - Q_DEV_COLUMN_BY_THRESHOLD[0.05]: round(r.q_dev_at_5pct, 5), - Q_DEV_COLUMN_BY_THRESHOLD[0.10]: round(r.q_dev_at_10pct, 5), - "IDs@1%FDR": r.ids_at_1pct, - "IDs@5%FDR": r.ids_at_5pct, - "IDs@10%FDR": r.ids_at_10pct, - } - ) - return pd.DataFrame(rows) - - -_DEFAULT_OUTPUT_DIR = Path("analysis/hpo_ablation") - - -def _validate_training_inputs( - *, - skip_training: bool, - train_features: Path | None, - val_features: Path | None, - validation_fraction: float | None, -) -> None: - if skip_training: - return - if train_features is None: - raise typer.BadParameter( - "--train-features is required unless --skip-training is set." - ) - if val_features is None and validation_fraction is None: - raise typer.BadParameter( - "Provide --val-features or --validation-fraction when training." - ) - if val_features is not None and validation_fraction is not None: - logger.warning( - "Both --val-features and --validation-fraction set; using --val-features." - ) - - -def _configure_ablation_configs( - *, - skip_training: bool, - train_features: Path | None, - eval_dfs: dict[str, pl.DataFrame], - hyperparams_from_model: Path | None, -) -> None: - global ABLATION_CONFIGS - - if not skip_training: - assert train_features is not None - train_schema_df = _load_parquet_as_polars(train_features) - ABLATION_CONFIGS = build_ablation_configs( - train_schema_df, hyperparams_from_model - ) - else: - first_eval = next(iter(eval_dfs.values())) - ABLATION_CONFIGS = build_ablation_configs(first_eval, hyperparams_from_model) - _configure_ablation_colors(ABLATION_CONFIGS.keys()) - logger.info("Ablation configs: %s", list(ABLATION_CONFIGS.keys())) - - -def _train_or_load_ablation_models( - *, - skip_training: bool, - train_features: Path | None, - val_features: Path | None, - validation_fraction: float | None, - output_dir: Path, - seed: int, - hyperparams_from_model: Path | None, -) -> dict[str, ProbabilityCalibrator]: - if skip_training: - logger.info("Step 3: Loading pre-trained ablation models...") - return load_ablation_models(output_dir) - - logger.info("Step 3: Training ablation models...") - assert train_features is not None - full_train_df = _load_parquet_as_polars(train_features) - if val_features is not None: - train_df = full_train_df - val_df = _load_parquet_as_polars(val_features) - else: - assert validation_fraction is not None - train_df, val_df = split_train_val_frames( - full_train_df, validation_fraction, seed - ) - train_hp = None - if hyperparams_from_model is not None: - train_hp = train_hyperparams_from_model(hyperparams_from_model) - logger.info( - "Using training hyperparameters from %s: %s", - hyperparams_from_model, - train_hp, - ) - return train_ablation_models( - train_df, val_df, output_dir, seed, train_hyperparams=train_hp - ) - - -def _evaluate_ablations( - *, - eval_dfs: dict[str, pl.DataFrame], - models: dict[str, ProbabilityCalibrator], - plots_dir: Path, - plot_format: str, -) -> list[EvalResult]: - logger.info("Step 4: Evaluating ablation models...") - all_results: list[EvalResult] = [] - - for ds_name, ds_df in eval_dfs.items(): - ds_results: list[EvalResult] = [] - for config_name, columns in ABLATION_CONFIGS.items(): - result = evaluate_single( - config_name, models[config_name], columns, ds_df, ds_name - ) - ds_results.append(result) - all_results.append(result) - logger.info( - " %s / %s: ECE=%.4f, Brier=%.4f, IDs@1%%=%d, IDs@5%%=%d, IDs@10%%=%d", - ds_name, - config_name, - result.ece, - result.brier, - result.ids_at_1pct, - result.ids_at_5pct, - result.ids_at_10pct, - ) - - logger.info("Step 5: Generating plots for %s...", ds_name) - _generate_plots_for_dataset(ds_results, ds_name, plots_dir, plot_format) - - return all_results - - -def _write_ablation_summary(output_dir: Path, all_results: list[EvalResult]) -> None: - logger.info("Step 7: Writing summary...") - summary = build_summary_table(all_results) - summary.to_csv(output_dir / "ablation_summary.csv", index=False) - - summary_json = summary.to_dict(orient="records") - with open(output_dir / "ablation_summary.json", "w") as f: - json.dump(summary_json, f, indent=2) - - logger.info("Summary table:\n%s", summary.to_string(index=False)) - - -# --------------------------------------------------------------------------- -# Main CLI -# --------------------------------------------------------------------------- -@app.command() -def main( - train_features: Annotated[ - Optional[Path], - typer.Option( - help="Path to pre-computed training Parquet file or directory. " - "Required unless --skip-training is set.", - ), - ] = None, - val_features: Annotated[ - Optional[Path], - typer.Option( - help="Pre-computed validation Parquet. Omit if using --validation-fraction.", - ), - ] = None, - validation_fraction: Annotated[ - Optional[float], - typer.Option( - "--validation-fraction", - min=0.0, - max=1.0, - help=( - "Hold out this fraction of --train-features for validation " - "(same row split for every ablation model). Alternative to --val-features." - ), - ), - ] = None, - output_dir: Annotated[ - Path, - typer.Option(help="Directory for cached features, models, metrics, and plots."), - ] = _DEFAULT_OUTPUT_DIR, - astral_spectra: Annotated[ - Optional[str], - typer.Option(help="Optional: path to Astral spectra directory."), - ] = None, - astral_predictions: Annotated[ - Optional[str], - typer.Option(help="Optional: path to Astral predictions CSV."), - ] = None, - plot_format: Annotated[ - str, - typer.Option(help="Plot format: 'pdf', 'png', or 'both'."), - ] = "both", - seed: Annotated[ - int, - typer.Option(help="Random seed."), - ] = 42, - koina_url: Annotated[ - str, - typer.Option(help="Koina server URL for eval feature computation."), - ] = "koina.wilhelmlab.org:443", - koina_ssl: Annotated[ - bool, - typer.Option(help="Use SSL for Koina server."), - ] = True, - skip_feature_compute: Annotated[ - bool, - typer.Option( - "--skip-feature-compute", - help="Skip eval feature computation; assume cache exists.", - ), - ] = False, - skip_training: Annotated[ - bool, - typer.Option( - "--skip-training", - help="Load pre-trained ablation models from {output-dir}/models/ " - "instead of training from scratch.", - ), - ] = False, - hyperparams_from_model: Annotated[ - Optional[Path], - typer.Option( - help="Use training hyperparameters from this saved calibrator directory " - "(e.g. HPO best model). Reads config.json.", - ), - ] = None, - plots_only: Annotated[ - bool, - typer.Option( - "--plots-only", - help="Regenerate plots from {output-dir}/eval_results only " - "(no feature compute, training, or evaluation).", - ), - ] = False, -) -> None: - """Run feature ablation study for the Winnow calibrator.""" - output_dir.mkdir(parents=True, exist_ok=True) - - if plots_only: - _run_plots_only(output_dir, plot_format) - logger.info("Feature ablation plots complete.") - return - - _validate_training_inputs( - skip_training=skip_training, - train_features=train_features, - val_features=val_features, - validation_fraction=validation_fraction, - ) - - plots_dir = output_dir / "plots" - plots_dir.mkdir(parents=True, exist_ok=True) - - logger.info("Step 1: Computing eval features...") - eval_parquets = compute_all_eval_features( - output_dir, - koina_url, - koina_ssl, - astral_spectra, - astral_predictions, - skip_feature_compute, - reference_model_dir=hyperparams_from_model, - ) - - logger.info("Step 2: Loading Parquets...") - eval_dfs: dict[str, pl.DataFrame] = {} - for name, path in eval_parquets.items(): - eval_dfs[name] = _load_parquet_as_polars(path) - logger.info(" Loaded eval %s: %d rows", name, len(eval_dfs[name])) - - _configure_ablation_configs( - skip_training=skip_training, - train_features=train_features, - eval_dfs=eval_dfs, - hyperparams_from_model=hyperparams_from_model, - ) - models = _train_or_load_ablation_models( - skip_training=skip_training, - train_features=train_features, - val_features=val_features, - validation_fraction=validation_fraction, - output_dir=output_dir, - seed=seed, - hyperparams_from_model=hyperparams_from_model, - ) - all_results = _evaluate_ablations( - eval_dfs=eval_dfs, - models=models, - plots_dir=plots_dir, - plot_format=plot_format, - ) - - logger.info("Step 6: Saving eval results...") - save_eval_results(all_results, output_dir) - _write_ablation_summary(output_dir, all_results) - logger.info("Results saved to %s", output_dir) - logger.info("Feature ablation study complete.") - - -if __name__ == "__main__": - app() From 2b3f01af36e21e06614b598700aa90d35a7bcf8a Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:30:43 +0100 Subject: [PATCH 16/26] chore: add Make commands to clean paper reproduction outputs --- Makefile | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/Makefile b/Makefile index 2fb59b38..5f9033fd 100644 --- a/Makefile +++ b/Makefile @@ -164,3 +164,29 @@ predict-sample: ## Clean output directories (does not delete sample data) clean: rm -rf models/ results/ + +################################################################################# +## Paper analyses commands # +################################################################################# + +# Match Makefile.paper defaults +PAPER_DATA_DIR ?= paper_data +PAPER_RESULTS_DIR ?= paper_results +PAPER_PLOTS_DIR ?= paper_plots + +.PHONY: clean-paper-recompute clean-paper-plots clean-paper-data clean-paper + +## Remove paper-recompute-* outputs +clean-paper-recompute: + rm -rf $(PAPER_RESULTS_DIR)/ + +## Remove paper-plot-* figures +clean-paper-plots: + rm -rf $(PAPER_PLOTS_DIR)/ + +## Remove paper-setup downloads +clean-paper-data: + rm -rf $(PAPER_DATA_DIR)/ + +## Remove paper plot and recompute outputs (keeps paper_data/) +clean-paper-outputs: clean-paper-recompute clean-paper-plots From f27fdace38621221ae4c04bab9d77abd4b10d0aa Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:46:19 +0100 Subject: [PATCH 17/26] chore: remove unused method --- .../calibrator_generalisation_utils.py | 53 +------------------ 1 file changed, 1 insertion(+), 52 deletions(-) diff --git a/paper_scripts/calibrator_generalisation_utils.py b/paper_scripts/calibrator_generalisation_utils.py index 8847e3dd..475f73cb 100644 --- a/paper_scripts/calibrator_generalisation_utils.py +++ b/paper_scripts/calibrator_generalisation_utils.py @@ -10,8 +10,6 @@ logger = logging.getLogger(__name__) -HEPG2_SOURCE = "PXD019483" - SPECIES_NAME_MAPPING: dict[str, str] = { "gluc": "HeLa degradome", "helaqc": "HeLa single shot", @@ -19,7 +17,7 @@ "immuno": "Immunopeptidomics-1", "celegans": "$\\it{C.\\;elegans}$", "sbrodae": "$\\it{Scalindua\\;brodae}$", - HEPG2_SOURCE: "HepG2", + "PXD019483": "HepG2", "hepg2": "HepG2", "snakevenoms": "Snake venomics", "tplantibodies": "Therapeutic nanobodies", @@ -64,52 +62,3 @@ def build_experiment_source_mapping(biological_validation_dir: Path) -> dict[str len(parquet_files), ) return mapping - - -def annotate_train_source_labels( - train_parquet: Path, - train_predictions: Path, - biological_validation_dir: Path, -) -> None: - """Add a ``source`` column to the train parquet and predictions CSV. - - Experiments found in ``biological_validation_dir`` inherit that project name. - All other experiments are labelled as HepG2 (``PXD019483``). - """ - experiment_to_source = build_experiment_source_mapping(biological_validation_dir) - lookup = pl.DataFrame( - { - "experiment_name": list(experiment_to_source.keys()), - "source": list(experiment_to_source.values()), - } - ) - - spectra = pl.read_parquet(train_parquet) - if "source" not in spectra.columns: - spectra = spectra.join(lookup, on="experiment_name", how="left").with_columns( - pl.col("source").fill_null(HEPG2_SOURCE) - ) - spectra.write_parquet(train_parquet) - logger.info("Wrote source labels to %s", train_parquet) - else: - logger.info( - "Parquet already has source column, leaving %s unchanged", train_parquet - ) - - predictions = pl.read_csv(train_predictions) - if "source" not in predictions.columns: - source_by_spectrum = spectra.select("spectrum_id", "source") - predictions = predictions.join(source_by_spectrum, on="spectrum_id", how="left") - missing = predictions.filter(pl.col("source").is_null()) - if len(missing) > 0: - raise ValueError( - f"{len(missing)} prediction rows in {train_predictions} have no matching " - "spectrum_id in the train parquet" - ) - predictions.write_csv(train_predictions) - logger.info("Wrote source labels to %s", train_predictions) - else: - logger.info( - "Predictions CSV already has source column, leaving %s unchanged", - train_predictions, - ) From 44af530834f340d6151c526fca426150a542c31c Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:47:40 +0100 Subject: [PATCH 18/26] chore: fix .PHONY mismatch --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 5f9033fd..996982b3 100644 --- a/Makefile +++ b/Makefile @@ -174,7 +174,7 @@ PAPER_DATA_DIR ?= paper_data PAPER_RESULTS_DIR ?= paper_results PAPER_PLOTS_DIR ?= paper_plots -.PHONY: clean-paper-recompute clean-paper-plots clean-paper-data clean-paper +.PHONY: clean-paper-recompute clean-paper-plots clean-paper-data clean-paper-outputs ## Remove paper-recompute-* outputs clean-paper-recompute: From 4ab105743ee72959633bff449c7b3a6efb651c8f Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:50:22 +0100 Subject: [PATCH 19/26] docs: uncomment training_matrix_output_path in docs --- docs/configuration.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/configuration.md b/docs/configuration.md index 54325107..65a50d2f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -175,7 +175,7 @@ dataset: predictions_path: data/predictions.csv metadata_output_path: results/metadata.csv -# training_matrix_output_path: results/training_matrix.parquet +training_matrix_output_path: results/training_matrix.parquet labelled: true ``` From 0a119d96e7d9ab222b8fc1f1468ac36b33050937 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:36:02 +0100 Subject: [PATCH 20/26] fix: resolve FASTAs from the HF datasets snapshot --- paper_scripts/plot_fdr_method_comparison.py | 45 ++++++++++++++++++- .../run_external_peptide_holdout_benchmark.py | 18 +++++++- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/paper_scripts/plot_fdr_method_comparison.py b/paper_scripts/plot_fdr_method_comparison.py index ff6f443a..65ff64d6 100644 --- a/paper_scripts/plot_fdr_method_comparison.py +++ b/paper_scripts/plot_fdr_method_comparison.py @@ -102,8 +102,10 @@ DEFAULT_WINNOW_RESULTS = _REPO_ROOT / "results" DEFAULT_MODEL_ROOT = _REPO_ROOT / "models" +DEFAULT_FASTA_ROOT = _REPO_ROOT / "paper_data/winnow-ms-datasets" DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/fdr_method_comparison_psm" DEFAULT_DATASETS = ["helaqc", "celegans"] +_HF_DATASETS_ID = "InstaDeepAI/winnow-ms-datasets" _METHOD_COLOURS = { PRIMARY_METHOD: _MAIN_LINE_COLOUR, DB_CAL_METHOD: _RAW_LINE_COLOUR, @@ -150,11 +152,38 @@ class DatasetConfig: calibrator_train_metadata: Path +def _resolve_fasta(fasta_root: Path, relative_fasta: str) -> Path: + """Resolve a FASTA under the Hugging Face datasets snapshot. + + Args: + fasta_root: Root of the local ``winnow-ms-datasets`` snapshot + (``HF_DATASETS_DIR`` / ``paper_data/winnow-ms-datasets``). + relative_fasta: Path relative to that root (e.g. ``fasta/human.fasta``). + + Returns: + Absolute path to an existing FASTA file. + + Raises: + FileNotFoundError: If the file is missing. + """ + fasta = (fasta_root / relative_fasta).resolve() + if not fasta.is_file(): + raise FileNotFoundError( + f"Missing FASTA {fasta}. Proteomes ship in the Hugging Face " + f"dataset {_HF_DATASETS_ID} under fasta/. Run " + "`make -f Makefile.paper download-paper-datasets` " + "(or `hf download … --include 'fasta/**'` as in " + "paper_scripts/README.md)." + ) + return fasta + + def build_dataset_configs( winnow_results: Path = DEFAULT_WINNOW_RESULTS, *, novoboard_root: Path, model_root: Path = DEFAULT_MODEL_ROOT, + fasta_root: Path = DEFAULT_FASTA_ROOT, ) -> dict[str, DatasetConfig]: """Build per-dataset path bundles from repo roots.""" configs: dict[str, DatasetConfig] = {} @@ -162,7 +191,7 @@ def build_dataset_configs( suffix = meta["winnow_suffix"] configs[key] = DatasetConfig( key=key, - fasta=_REPO_ROOT / meta["fasta"], + fasta=_resolve_fasta(fasta_root, meta["fasta"]), winnow_unlabelled=winnow_results / f"instanovo_{suffix}_predictions_unlabelled", winnow_test=winnow_results / f"instanovo_{suffix}_predictions_test", @@ -1056,6 +1085,16 @@ def main( Path, typer.Option("--winnow-results", help="Winnow results directory."), ] = DEFAULT_WINNOW_RESULTS, + fasta_root: Annotated[ + Path, + typer.Option( + "--fasta-root", + help=( + "Root of the local Hugging Face winnow-ms-datasets snapshot " + "(FASTA paths are relative to this directory)." + ), + ), + ] = DEFAULT_FASTA_ROOT, summarise_only: Annotated[ Optional[Path], typer.Option( @@ -1099,7 +1138,9 @@ def main( ) dataset_keys = datasets if datasets is not None else list(DEFAULT_DATASETS) - configs = build_dataset_configs(winnow_results, novoboard_root=novoboard_root) + configs = build_dataset_configs( + winnow_results, novoboard_root=novoboard_root, fasta_root=fasta_root + ) curve_parts: list[pd.DataFrame] = [] for key in dataset_keys: diff --git a/paper_scripts/run_external_peptide_holdout_benchmark.py b/paper_scripts/run_external_peptide_holdout_benchmark.py index e791d80b..f71d45a8 100644 --- a/paper_scripts/run_external_peptide_holdout_benchmark.py +++ b/paper_scripts/run_external_peptide_holdout_benchmark.py @@ -85,6 +85,7 @@ ) from plot_eval_results import _PALETTE, _display_name, _save_fig, _style_ax # noqa: E402 from plot_fdr_method_comparison import ( # noqa: E402 + DEFAULT_FASTA_ROOT, DEFAULT_MODEL_ROOT, DEFAULT_WINNOW_RESULTS, build_dataset_configs, @@ -223,6 +224,7 @@ def build_shared_score_tables( winnow_results: Path, novoboard_root: Path, model_root: Path = DEFAULT_MODEL_ROOT, + fasta_root: Path = DEFAULT_FASTA_ROOT, unlabelled_min_length: int = MIN_PEPTIDE_LENGTH, labelled_min_length: int = LABELLED_MIN_PEPTIDE_LENGTH, ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]: @@ -244,7 +246,10 @@ def build_shared_score_tables( glissade_reference: training-split matched scores for Glissade's anchor. """ cfg = build_dataset_configs( - winnow_results, novoboard_root=novoboard_root, model_root=model_root + winnow_results, + novoboard_root=novoboard_root, + model_root=model_root, + fasta_root=fasta_root, )[dataset] winnow_test = _load_winnow_with_raw_confidence( @@ -1059,6 +1064,16 @@ def main( help="Per-dataset calibrator directories, used for Glissade's anchor.", ), ] = DEFAULT_MODEL_ROOT, + fasta_root: Annotated[ + Path, + typer.Option( + "--fasta-root", + help=( + "Root of the local Hugging Face winnow-ms-datasets snapshot " + "(FASTA paths are relative to this directory)." + ), + ), + ] = DEFAULT_FASTA_ROOT, n_bootstraps: Annotated[ int, typer.Option( @@ -1127,6 +1142,7 @@ def main( winnow_results=winnow_results, novoboard_root=novoboard_root, model_root=model_root, + fasta_root=fasta_root, unlabelled_min_length=min_peptide_length, labelled_min_length=LABELLED_MIN_PEPTIDE_LENGTH, ) From 813edd47a3f333da5e937e1535a4c74213121596 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:36:54 +0100 Subject: [PATCH 21/26] fix: fail without fallback when Figshare version is missing --- paper_scripts/download_figshare_article.py | 48 +++++++++++++++++----- 1 file changed, 37 insertions(+), 11 deletions(-) diff --git a/paper_scripts/download_figshare_article.py b/paper_scripts/download_figshare_article.py index 433e92d7..2200121a 100644 --- a/paper_scripts/download_figshare_article.py +++ b/paper_scripts/download_figshare_article.py @@ -74,18 +74,44 @@ def _http_get_json(url: str) -> dict[str, Any]: return data +def _latest_article_version(article_id: int) -> int | None: + """Latest published version from the unversioned article endpoint. + + Used only for error messages. Never used as a download source. + """ + try: + article = _http_get_json(_article_url(article_id, None)) + except RuntimeError: + return None + version = article.get("version") + if isinstance(version, int): + return version + if isinstance(version, str) and version.isdigit(): + return int(version) + return None + + def _fetch_article(article_id: int, version: int | None) -> dict[str, Any]: - """Load article metadata, preferring the versioned endpoint when set.""" - if version is not None: - url = _article_url(article_id, version) - try: - return _http_get_json(url) - except RuntimeError as exc: - logger.warning( - "Versioned endpoint failed (%s); falling back to current article metadata.", - exc, - ) - return _http_get_json(_article_url(article_id, None)) + """Load article metadata for an exact version pin when ``version`` is set. + + When ``version`` is set, only the versioned endpoint is used for files. + """ + if version is None: + return _http_get_json(_article_url(article_id, None)) + + url = _article_url(article_id, version) + try: + return _http_get_json(url) + except RuntimeError as exc: + latest = _latest_article_version(article_id) + if latest is not None: + latest_clause = f" Latest published version is {latest}." + else: + latest_clause = " Latest published version could not be determined." + raise RuntimeError( + f"Figshare article {article_id} version {version} is not available " + f"({exc}).{latest_clause}" + ) from exc def _relative_path(file_info: dict[str, Any], folder_structure: dict[str, Any]) -> str: From 8f2f58bdd4e07d904534a2f1b22d5f561698beb4 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:37:42 +0100 Subject: [PATCH 22/26] fix: fail-fast recompute umbrellas and update fasta paths --- Makefile.paper | 77 ++++++++++++++++++++++++-------------------------- 1 file changed, 37 insertions(+), 40 deletions(-) diff --git a/Makefile.paper b/Makefile.paper index 9b732b5f..b78e7ff6 100644 --- a/Makefile.paper +++ b/Makefile.paper @@ -132,10 +132,16 @@ HELAQC_DNS_TOOLS ?= instanovo casanovo primenovo HELAQC_DNS_LABEL_instanovo := InstaNovo HELAQC_DNS_MODEL_instanovo := $(HELAQC_MODEL) +HELAQC_DNS_LOADER_instanovo := instanovo +HELAQC_DNS_PREDS_DIR_instanovo := instanovo HELAQC_DNS_LABEL_casanovo := Casanovo HELAQC_DNS_MODEL_casanovo := $(CASANOVO_HELAQC_MODEL) +HELAQC_DNS_LOADER_casanovo := mztab +HELAQC_DNS_PREDS_DIR_casanovo := casanovo HELAQC_DNS_LABEL_primenovo := pi-PrimeNovo HELAQC_DNS_MODEL_primenovo := $(PRIMENOVO_HELAQC_MODEL) +HELAQC_DNS_LOADER_primenovo := primenovo +HELAQC_DNS_PREDS_DIR_primenovo := primenovo # FDR tool recompute datasets. Laptop keeps helaqc only (celegans holdout >5 min). FDR_RECOMPUTE_DATASETS_LAPTOP ?= helaqc @@ -326,29 +332,27 @@ $(foreach t,$(HELAQC_DNS),$(eval $(call HELAQC_PLOT_TOOL_RULE,$(t)))) paper-plot-helaqc-analysis: $(foreach t,$(HELAQC_DNS),paper-plot-helaqc-$(t)) -paper-recompute-helaqc: - @mkdir -p $(PAPER_RESULTS_DIR) - @for tool in $(HELAQC_DNS_TOOLS); do \ - for split in $(HELAQC_SPLITS); do \ - out=$(PAPER_RESULTS_DIR)/$${tool}_helaqc_predictions_$${split}; \ - echo "Predict $${tool} $${split} -> $${out}"; \ - case $${tool} in \ - instanovo) loader=instanovo; model=$(HELAQC_MODEL); preds=$(HELAQC_DATA)/instanovo/$${split}_preds.csv ;; \ - casanovo) loader=mztab; model=$(CASANOVO_HELAQC_MODEL); preds=$(HELAQC_DATA)/casanovo/$${split}_preds.csv ;; \ - primenovo) loader=primenovo; model=$(PRIMENOVO_HELAQC_MODEL); preds=$(HELAQC_DATA)/primenovo/$${split}_preds.csv ;; \ - *) echo "Unknown tool $${tool}"; exit 1 ;; \ - esac; \ - spectra=$(HELAQC_DATA)/$${split}.parquet; \ - $(WINNOW) predict \ - data_loader=$${loader} \ - dataset.spectrum_path_or_directory=$${spectra} \ - dataset.predictions_path=$${preds} \ - calibrator.pretrained_model_name_or_path=$${model} \ - fdr_control.fdr_threshold=$(PREDICT_FDR_THRESHOLD) \ - $(KOINA_FRAGMENT_MATCH_CONSTANTS) \ - output_folder=$${out}; \ - done; \ - done +# $(1) = tool (instanovo|casanovo|primenovo), $(2) = split (test|unlabelled|raw_less_train) +define HELAQC_RECOMPUTE_TOOL_SPLIT_RULE +.PHONY: paper-recompute-helaqc-$(1)-$(2) +paper-recompute-helaqc-$(1)-$(2): + @mkdir -p $$(PAPER_RESULTS_DIR) + @out=$$(PAPER_RESULTS_DIR)/$(1)_helaqc_predictions_$(2); \ + echo "Predict $(1) $(2) -> $$$$out"; \ + $$(WINNOW) predict \ + data_loader=$$(HELAQC_DNS_LOADER_$(1)) \ + dataset.spectrum_path_or_directory=$$(HELAQC_DATA)/$(2).parquet \ + dataset.predictions_path=$$(HELAQC_DATA)/$$(HELAQC_DNS_PREDS_DIR_$(1))/$(2)_preds.csv \ + calibrator.pretrained_model_name_or_path=$$(HELAQC_DNS_MODEL_$(1)) \ + fdr_control.fdr_threshold=$$(PREDICT_FDR_THRESHOLD) \ + $$(KOINA_FRAGMENT_MATCH_CONSTANTS) \ + output_folder=$$$$out +endef + +$(foreach t,$(HELAQC_DNS),$(foreach s,$(HELAQC_SPLITS),$(eval $(call HELAQC_RECOMPUTE_TOOL_SPLIT_RULE,$(t),$(s))))) + +paper-recompute-helaqc: \ + $(foreach t,$(HELAQC_DNS_TOOLS),$(foreach s,$(HELAQC_SPLITS),paper-recompute-helaqc-$(t)-$(s))) ################################################################################ # 3. FDR method comparison @@ -368,6 +372,7 @@ paper-recompute-fdr-method-comparison: $(PYTHON) $(PAPER_SCRIPTS)/plot_fdr_method_comparison.py \ --novoboard-root $(NOVOBOARD_ROOT) \ --winnow-results $(WINNOW_FDR_RESULTS) \ + --fasta-root $(HF_DATASETS_DIR) \ $(foreach d,$(FDR_RECOMPUTE_DATASETS),--datasets $(d)) \ --results-dir $(PAPER_RESULTS_DIR)/fdr_method_comparison \ --plots-dir $(PAPER_PLOTS_DIR)/fdr_method_comparison @@ -391,6 +396,7 @@ paper-recompute-external-peptide-holdout: --novoboard-root $(NOVOBOARD_ROOT) \ --winnow-results $(WINNOW_FDR_RESULTS) \ --model-root $(FDR_MODEL_ROOT) \ + --fasta-root $(HF_DATASETS_DIR) \ $(foreach d,$(FDR_RECOMPUTE_DATASETS),--datasets $(d)) \ --results-dir $(PAPER_RESULTS_DIR)/external_peptide_holdout \ --plots-dir $(PAPER_PLOTS_DIR)/external_peptide_holdout @@ -515,23 +521,14 @@ $(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_LABELLED_STEM_RULE,$(s)))) $(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_FULL_STEM_RULE,$(s)))) $(foreach s,$(GENERAL_STEMS),$(eval $(call GENERAL_BOTH_STEM_RULE,$(s)))) -paper-recompute-general-labelled: - @mkdir -p $(PAPER_RESULTS_DIR)/general_results/labelled - @for stem in $(GENERAL_RECOMPUTE_STEMS); do \ - $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-labelled-$${stem}; \ - done - -paper-recompute-general-full-small: - @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full - @for stem in $(GENERAL_FULL_SMALL_STEMS); do \ - $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ - done - -paper-recompute-general-full-large: - @mkdir -p $(PAPER_RESULTS_DIR)/general_results/full - @for stem in $(GENERAL_FULL_LARGE_STEMS); do \ - $(MAKE) -f Makefile.paper --no-print-directory paper-recompute-general-full-$${stem}; \ - done +paper-recompute-general-labelled: \ + $(foreach s,$(GENERAL_RECOMPUTE_STEMS),paper-recompute-general-labelled-$(s)) + +paper-recompute-general-full-small: \ + $(foreach s,$(GENERAL_FULL_SMALL_STEMS),paper-recompute-general-full-$(s)) + +paper-recompute-general-full-large: \ + $(foreach s,$(GENERAL_FULL_LARGE_STEMS),paper-recompute-general-full-$(s)) # Echo exact recipes for all nine stems × labelled/full (immuno2 always subsets). paper-recompute-general-print: From b1f1b14892fa191cc0b9a08ce230885ad4840ec9 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:38:27 +0100 Subject: [PATCH 23/26] chore: drop leftover experimental naming --- paper_scripts/analyze_fdr_overlap.py | 12 ++++++-- paper_scripts/analyze_upscored_fps.py | 4 +-- .../evaluate_calibrator_generalisation.py | 29 ++++++++++--------- paper_scripts/run_feature_ablations.py | 4 +-- 4 files changed, 29 insertions(+), 20 deletions(-) diff --git a/paper_scripts/analyze_fdr_overlap.py b/paper_scripts/analyze_fdr_overlap.py index 151cc3a1..80c81c22 100644 --- a/paper_scripts/analyze_fdr_overlap.py +++ b/paper_scripts/analyze_fdr_overlap.py @@ -136,7 +136,7 @@ def _display_name(key: str) -> str: def _project_key_from_folder(folder_name: str) -> str: - """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" + """Strip a known eval suffix to get the project key (e.g. ``PXD014877_unlabelled`` -> ``PXD014877``).""" for suffix in _FOLDER_SUFFIXES: if folder_name.endswith(suffix): return folder_name[: -len(suffix)] @@ -948,14 +948,20 @@ def main( Path, typer.Option( "--unlabelled-dir", - help="Root with per-project full-search folders (e.g. gluc_raw/, PXD014877_unlabelled/).", + help=( + "Root with per-project full-search folders " + "(e.g. paper_data/general_results/full/)." + ), ), ], labelled_dir: Annotated[ Path, typer.Option( "--labelled-dir", - help="Root with per-project database-search folders (e.g. gluc_annotated/, PXD014877_labelled/).", + help=( + "Root with per-project database-search folders " + "(e.g. paper_data/general_results/labelled/)." + ), ), ], results_dir: Annotated[ diff --git a/paper_scripts/analyze_upscored_fps.py b/paper_scripts/analyze_upscored_fps.py index 77faecce..fe8082c6 100644 --- a/paper_scripts/analyze_upscored_fps.py +++ b/paper_scripts/analyze_upscored_fps.py @@ -94,7 +94,7 @@ _FOLDER_SUFFIXES = ("_annotated", "_labelled", "_raw", "_unlabelled") -# new_eval_sets_results layout: lcfm/PXD004452//preds_and_fdr_metrics.csv +# Figshare general_results/labelled layout: PXD004452//preds_and_fdr_metrics.csv _PXD_ACCESSION_PREFIX = "PXD" FEATURE_COLUMNS_OF_INTEREST = [ @@ -212,7 +212,7 @@ def _strip_mods(seq: str) -> str: def _project_key_from_folder(folder_name: str) -> str: - """Strip a known eval suffix to get the project key (e.g. ``gluc_raw`` -> ``gluc``).""" + """Strip a known eval suffix to get the project key (e.g. ``PXD014877_labelled`` -> ``PXD014877``).""" for suffix in _FOLDER_SUFFIXES: if folder_name.endswith(suffix): return folder_name[: -len(suffix)] diff --git a/paper_scripts/evaluate_calibrator_generalisation.py b/paper_scripts/evaluate_calibrator_generalisation.py index 98259bc3..b3088b72 100644 --- a/paper_scripts/evaluate_calibrator_generalisation.py +++ b/paper_scripts/evaluate_calibrator_generalisation.py @@ -39,7 +39,7 @@ # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- -logger = logging.getLogger("winnow.evaluate_generalization") +logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) logger.propagate = False logger.addHandler(RichHandler()) @@ -84,21 +84,21 @@ "herceptin": 0.15, } -# Mirrors Makefile train-extra-small-mass-error-da / EXTRA_SMALL_* overrides. -_EXTRA_SMALL_FRAGMENT_EXCLUDE = [ +# Paper general-model feature set: reduced fragment/beam columns. +_REDUCED_FRAGMENT_EXCLUDE = [ "spectral_angle", "xcorr", "complementary_ion_count", "max_ion_gap", ] -_EXTRA_SMALL_BEAM_EXCLUDE = ["edit_distance"] +_REDUCED_BEAM_EXCLUDE = ["edit_distance"] def initialise_calibrator( *, train_project: Optional[str] = None, ) -> ProbabilityCalibrator: - """Create a fresh calibrator matching train-extra-small-mass-error-da.""" + """Create a fresh calibrator with the paper general-model feature set.""" irt_train_fraction = _IRT_TRAIN_FRACTION_OVERRIDES.get(train_project or "", 0.1) calibrator = ProbabilityCalibrator( @@ -139,19 +139,18 @@ def initialise_calibrator( ) calibrator.add_feature(BeamFeatures()) calibrator.add_feature(TokenScoreFeatures()) - # Former excluded_columns behaviour: train on a reduced feature subset. + # Train on a reduced feature subset (exclude some fragment/beam columns). training_columns = [ col for col in calibrator.columns - if col not in _EXTRA_SMALL_FRAGMENT_EXCLUDE - and col not in _EXTRA_SMALL_BEAM_EXCLUDE + if col not in _REDUCED_FRAGMENT_EXCLUDE and col not in _REDUCED_BEAM_EXCLUDE ] calibrator.set_training_feature_columns(training_columns) return calibrator def load_dataset(data_path: Path, predictions_path: Path) -> CalibrationDataset: - """Load the combined train_extra_small dataset.""" + """Load the HF general_model_training_set (or equivalent) dataset.""" logger.info("Loading dataset from %s and %s", data_path, predictions_path) loader = InstaNovoDatasetLoader( residue_masses=RESIDUE_MASSES, @@ -250,10 +249,14 @@ def evaluate_model( # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- -_DEFAULT_MODEL_OUTPUT_DIR = Path("models/generalisation") -_DEFAULT_RESULTS_OUTPUT_DIR = Path("results/generalisation") -_DEFAULT_TRAIN_PARQUET = Path("train_extra_small/train.parquet") -_DEFAULT_TRAIN_PREDS = Path("train_extra_small/train_preds.csv") +_DEFAULT_MODEL_OUTPUT_DIR = Path("paper_results/generalisation/models") +_DEFAULT_RESULTS_OUTPUT_DIR = Path("paper_results/generalisation") +_DEFAULT_TRAIN_PARQUET = Path( + "paper_data/winnow-ms-datasets/general_model_training_set/train.parquet" +) +_DEFAULT_TRAIN_PREDS = Path( + "paper_data/winnow-ms-datasets/general_model_training_set/train_preds.csv" +) app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) diff --git a/paper_scripts/run_feature_ablations.py b/paper_scripts/run_feature_ablations.py index 37bf606d..c56f1fc3 100644 --- a/paper_scripts/run_feature_ablations.py +++ b/paper_scripts/run_feature_ablations.py @@ -124,10 +124,10 @@ c for c in FRAGMENT_MATCH_COLUMNS if c not in _EXCLUDED_REDUCED ] -# Default training matrix columns (train_extra_small_matrix.parquet). +# Default training matrix columns. REDUCED_TRAIN_COLUMNS: list[str] = FEATURE_SUBSETS["no_fragment_similarity"]["columns"] -# Hydra overrides aligned with Makefile ANALYSIS_REDUCED_FEATURE_OVERRIDES (mass_error_da model). +# Hydra overrides for mass_error_da instead of mass_error_ppm. REDUCED_FEATURE_COMPUTE_OVERRIDES: list[str] = [ "~calibrator.features.mass_error", "+calibrator.features.mass_error_da._target_=winnow.calibration.calibration_features.MassErrorDaFeature", From 7b35e64c9f1a95382e601194d6b551473dec1e0b Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:49:03 +0100 Subject: [PATCH 24/26] fix: remove eager FASTA resolution --- paper_scripts/plot_fdr_method_comparison.py | 39 +++++++++++-------- .../run_external_peptide_holdout_benchmark.py | 1 + 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/paper_scripts/plot_fdr_method_comparison.py b/paper_scripts/plot_fdr_method_comparison.py index 65ff64d6..306c727f 100644 --- a/paper_scripts/plot_fdr_method_comparison.py +++ b/paper_scripts/plot_fdr_method_comparison.py @@ -126,16 +126,6 @@ "novoboard_decoy": "0.70", "winnow_suffix": "celegans", }, - "sbrodae": { - "fasta": "fasta/Sb_proteome.fasta", - "novoboard_decoy": "0.50", - "winnow_suffix": "sbrodae", - }, - "PXD019483": { - "fasta": "fasta/human.fasta", - "novoboard_decoy": "0.70", - "winnow_suffix": "pxd019483", - }, } @@ -184,10 +174,21 @@ def build_dataset_configs( novoboard_root: Path, model_root: Path = DEFAULT_MODEL_ROOT, fasta_root: Path = DEFAULT_FASTA_ROOT, + datasets: list[str] | None = None, ) -> dict[str, DatasetConfig]: - """Build per-dataset path bundles from repo roots.""" + """Build per-dataset path bundles from repo roots. + + FASTA paths are resolved only for the requested ``datasets`` keys (default: + all entries in ``_DATASET_META``). + """ + keys = list(datasets) if datasets is not None else list(_DATASET_META) configs: dict[str, DatasetConfig] = {} - for key, meta in _DATASET_META.items(): + for key in keys: + if key not in _DATASET_META: + raise ValueError( + f"Unknown dataset {key!r}. Known keys: {sorted(_DATASET_META)}" + ) + meta = _DATASET_META[key] suffix = meta["winnow_suffix"] configs[key] = DatasetConfig( key=key, @@ -1138,14 +1139,18 @@ def main( ) dataset_keys = datasets if datasets is not None else list(DEFAULT_DATASETS) - configs = build_dataset_configs( - winnow_results, novoboard_root=novoboard_root, fasta_root=fasta_root - ) + try: + configs = build_dataset_configs( + winnow_results, + novoboard_root=novoboard_root, + fasta_root=fasta_root, + datasets=dataset_keys, + ) + except ValueError as exc: + raise typer.BadParameter(str(exc)) from exc curve_parts: list[pd.DataFrame] = [] for key in dataset_keys: - if key not in configs: - raise typer.BadParameter(f"Unknown dataset {key!r}") logger.info("Processing %s", key) curve_parts.append(process_dataset(configs[key], plots_dir)) diff --git a/paper_scripts/run_external_peptide_holdout_benchmark.py b/paper_scripts/run_external_peptide_holdout_benchmark.py index f71d45a8..12f2a5c5 100644 --- a/paper_scripts/run_external_peptide_holdout_benchmark.py +++ b/paper_scripts/run_external_peptide_holdout_benchmark.py @@ -250,6 +250,7 @@ def build_shared_score_tables( novoboard_root=novoboard_root, model_root=model_root, fasta_root=fasta_root, + datasets=[dataset], )[dataset] winnow_test = _load_winnow_with_raw_confidence( From b2fe7009943083916936f8a7763b7a72a2023c53 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:09:27 +0100 Subject: [PATCH 25/26] chore: direct CLI defaults to paper reproduction folders --- paper_scripts/analyze_upscored_fps.py | 6 +++--- paper_scripts/benchmark_scaling.py | 4 ++-- paper_scripts/plot_ablation_summary.py | 4 ++-- .../plot_calibrator_generalisation_heatmap.py | 2 +- paper_scripts/plot_fdr_method_comparison.py | 13 +++++++------ .../run_external_peptide_holdout_benchmark.py | 7 ++++--- paper_scripts/run_feature_ablations.py | 2 +- 7 files changed, 20 insertions(+), 18 deletions(-) diff --git a/paper_scripts/analyze_upscored_fps.py b/paper_scripts/analyze_upscored_fps.py index fe8082c6..fd674c5d 100644 --- a/paper_scripts/analyze_upscored_fps.py +++ b/paper_scripts/analyze_upscored_fps.py @@ -596,9 +596,9 @@ def _levenshtein(s: str, t: str) -> int: # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- -_DEFAULT_PREDICTIONS_ROOT = Path("predictions/general_model") -_DEFAULT_RESULTS_DIR = Path("analysis/upscored_fps") -_DEFAULT_PLOTS_DIR = Path("analysis/upscored_fps/plots") +_DEFAULT_PREDICTIONS_ROOT = Path("paper_data/general_results/labelled") +_DEFAULT_RESULTS_DIR = Path("paper_results/upscored_fps") +_DEFAULT_PLOTS_DIR = Path("paper_plots/upscored_fps") @app.command() diff --git a/paper_scripts/benchmark_scaling.py b/paper_scripts/benchmark_scaling.py index 452874c0..0d3cec38 100644 --- a/paper_scripts/benchmark_scaling.py +++ b/paper_scripts/benchmark_scaling.py @@ -52,8 +52,8 @@ DEFAULT_FRACTIONS = [0.1, 0.5, 1.0] _SEED = 42 _DEFAULT_MODEL_OUTPUT_DIR = Path("paper_results/scaling/dummy_model") -_DEFAULT_RESULTS_DIR = Path("analysis") -_DEFAULT_PLOTS_DIR = Path("analysis") +_DEFAULT_RESULTS_DIR = Path("paper_results/scaling") +_DEFAULT_PLOTS_DIR = Path("paper_plots/scaling") @contextmanager diff --git a/paper_scripts/plot_ablation_summary.py b/paper_scripts/plot_ablation_summary.py index 7494ec1b..96cd7ece 100644 --- a/paper_scripts/plot_ablation_summary.py +++ b/paper_scripts/plot_ablation_summary.py @@ -124,8 +124,8 @@ def ordered_ablation_configs(present: set[str]) -> list[str]: 0.10: "fdr_bias@10%FDR", } -_DEFAULT_SUMMARY = Path("results/ablations/ablation_summary.csv") -_DEFAULT_OUTPUT_DIR = Path("results/ablations/plots") +_DEFAULT_SUMMARY = Path("paper_results/ablations/ablation_summary.csv") +_DEFAULT_OUTPUT_DIR = Path("paper_results/ablations/plots") def load_ablation_summary(path: Path) -> pd.DataFrame: diff --git a/paper_scripts/plot_calibrator_generalisation_heatmap.py b/paper_scripts/plot_calibrator_generalisation_heatmap.py index b8f292f0..174de702 100644 --- a/paper_scripts/plot_calibrator_generalisation_heatmap.py +++ b/paper_scripts/plot_calibrator_generalisation_heatmap.py @@ -238,7 +238,7 @@ def create_comparison_heatmaps(results_path: Path, output_dir: Path) -> None: # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- -_DEFAULT_OUTPUT_DIR = Path("results/generalisation/plots") +_DEFAULT_OUTPUT_DIR = Path("paper_plots/generalisation") app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) diff --git a/paper_scripts/plot_fdr_method_comparison.py b/paper_scripts/plot_fdr_method_comparison.py index 306c727f..32e19ee4 100644 --- a/paper_scripts/plot_fdr_method_comparison.py +++ b/paper_scripts/plot_fdr_method_comparison.py @@ -100,10 +100,11 @@ FDR_THRESHOLDS = [0.01, 0.05, 0.10] _DB_GROUNDED_DROP = 10 -DEFAULT_WINNOW_RESULTS = _REPO_ROOT / "results" -DEFAULT_MODEL_ROOT = _REPO_ROOT / "models" -DEFAULT_FASTA_ROOT = _REPO_ROOT / "paper_data/winnow-ms-datasets" -DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/fdr_method_comparison_psm" +DEFAULT_WINNOW_RESULTS = Path("paper_data/fdr_benchmark_inputs/winnow_results") +DEFAULT_MODEL_ROOT = Path("paper_data/fdr_benchmark_inputs/models") +DEFAULT_FASTA_ROOT = Path("paper_data/winnow-ms-datasets") +DEFAULT_RESULTS_DIR = Path("paper_results/fdr_method_comparison") +DEFAULT_PLOTS_DIR = Path("paper_plots/fdr_method_comparison") DEFAULT_DATASETS = ["helaqc", "celegans"] _HF_DATASETS_ID = "InstaDeepAI/winnow-ms-datasets" _METHOD_COLOURS = { @@ -1073,11 +1074,11 @@ def main( results_dir: Annotated[ Path, typer.Option("--results-dir", help="Directory for curves/summary CSVs."), - ] = DEFAULT_OUTPUT_DIR, + ] = DEFAULT_RESULTS_DIR, plots_dir: Annotated[ Path, typer.Option("--plots-dir", help="Directory for png/pdf figures."), - ] = DEFAULT_OUTPUT_DIR, + ] = DEFAULT_PLOTS_DIR, datasets: Annotated[ Optional[list[str]], typer.Option("--datasets", help="Dataset keys to plot."), diff --git a/paper_scripts/run_external_peptide_holdout_benchmark.py b/paper_scripts/run_external_peptide_holdout_benchmark.py index 12f2a5c5..032bbdc6 100644 --- a/paper_scripts/run_external_peptide_holdout_benchmark.py +++ b/paper_scripts/run_external_peptide_holdout_benchmark.py @@ -97,7 +97,8 @@ logger = logging.getLogger(__name__) app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False) -DEFAULT_OUTPUT_DIR = _REPO_ROOT / "results/external_peptide_holdout_benchmark" +DEFAULT_RESULTS_DIR = Path("paper_results/external_peptide_holdout") +DEFAULT_PLOTS_DIR = Path("paper_plots/external_peptide_holdout") DEFAULT_DATASETS = ["helaqc", "celegans"] DEFAULT_Q_THRESHOLDS = [round(float(x), 2) for x in np.linspace(0.0, 0.25, 26)] DEFAULT_PI0_GRID = [0.5, 0.6, 0.7, 0.8, 0.9] @@ -1014,7 +1015,7 @@ def main( results_dir: Annotated[ Path, typer.Option("--results-dir", help="Directory for results/summary CSVs."), - ] = DEFAULT_OUTPUT_DIR, + ] = DEFAULT_RESULTS_DIR, plots_dir: Annotated[ Optional[Path], typer.Option( @@ -1106,7 +1107,7 @@ def main( if plot: if plots_dir is None: - plots_dir = DEFAULT_OUTPUT_DIR / "plots" + plots_dir = DEFAULT_PLOTS_DIR results_dir.mkdir(parents=True, exist_ok=True) plots_dir.mkdir(parents=True, exist_ok=True) else: diff --git a/paper_scripts/run_feature_ablations.py b/paper_scripts/run_feature_ablations.py index c56f1fc3..a3d3fff3 100644 --- a/paper_scripts/run_feature_ablations.py +++ b/paper_scripts/run_feature_ablations.py @@ -1276,7 +1276,7 @@ def build_summary_table(all_results: list[EvalResult]) -> pd.DataFrame: return pd.DataFrame(rows) -_DEFAULT_OUTPUT_DIR = Path("analysis/hpo_ablation") +_DEFAULT_OUTPUT_DIR = Path("paper_results/ablations") def _validate_training_inputs( From ea26079ce873267885d04666913590f81b044a68 Mon Sep 17 00:00:00 2001 From: Jemma Daniel <134346753+JemmaLDaniel@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:23:36 +0100 Subject: [PATCH 26/26] fix: fail clearly on FDR benchmarking tool error --- .../run_external_peptide_holdout_benchmark.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/paper_scripts/run_external_peptide_holdout_benchmark.py b/paper_scripts/run_external_peptide_holdout_benchmark.py index 032bbdc6..6ff4250c 100644 --- a/paper_scripts/run_external_peptide_holdout_benchmark.py +++ b/paper_scripts/run_external_peptide_holdout_benchmark.py @@ -697,16 +697,11 @@ def _evaluate_mixture_methods( for method, estimator in estimators.items(): try: q_table = estimator(mixed, estimator_reference) # type: ignore[operator] - except Exception as exc: # noqa: BLE001 - boundary around external tool - logger.warning( - "%s FDR failed dataset=%s pi0=%.3g iter=%d: %s", - method, - dataset, - pi0, - iteration, - exc, - ) - continue + except Exception as exc: + raise RuntimeError( + f"{method} FDR failed for dataset={dataset} pi0={pi0:.3g} " + f"iter={iteration}: {exc}" + ) from exc scored_keys = set(q_table["peptide_key"].astype(str)) if scored_keys != mixture_keys: raise AssertionError(