{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# CAFA6 GOA + ProtT5 IC-Aware Pruning\n\n**Objective**: Improve precision by pruning low-IC (generic) terms.\n\n**Key Insight**: F-max is weighted by Information Accretion (IA). \n- Low-IA terms add to the denominator but contribute little to the numerator\n- Over-predicting generic terms KILLS precision\n- We should predict FEWER terms with HIGHER confidence\n\n**Approach**:\n1. Blend GOA + ProtT5 (60/40)\n2. Apply hard-max propagation\n3. PRUNE low-IC terms (IC < threshold)\n4. Keep only top N terms per protein, ranked by Score * IC"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nimport gc\nimport pandas as pd\nimport numpy as np\nimport networkx as nx\nfrom tqdm.auto import tqdm\nfrom collections import defaultdict\n\n# ==========================================\n# CONFIGURATION\n# ==========================================\nOBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nIA_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\"\nGOA_PATH = '/kaggle/input/cafa6-goa-predictions/goa_submission.tsv'\nPROTT5_PATH = '/kaggle/input/cafa6-goa-predictions/prott5_interpro_predictions.tsv'\nSUBMISSION_OUTPUT = 'submission.tsv'\n\n# Weights for blending\nW_GOA = 0.60\nW_PROTT5 = 0.40\n\n# IC-aware pruning parameters\nMIN_IC_THRESHOLD = 1.0  # Prune terms with IC < 1.0 (very generic)\nMAX_TERMS_PER_PROTEIN = 1200  # Limit predictions per protein\nMIN_SCORE_THRESHOLD = 0.01  # Higher threshold for cleaner predictions\n\nprint(f\"OBO Path: {OBO_PATH}\")\nprint(f\"IA Path: {IA_PATH}\")\nprint(f\"GOA Path: {GOA_PATH}\")\nprint(f\"ProtT5 Path: {PROTT5_PATH}\")\nprint(f\"\\nWeights: GOA={W_GOA}, ProtT5={W_PROTT5}\")\nprint(f\"IC-aware: min_ic={MIN_IC_THRESHOLD}, max_terms={MAX_TERMS_PER_PROTEIN}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 1. BUILD GO GRAPH & LOAD IA\n# ==========================================\ndef build_go_graph(go_obo_path):\n    print(f\"[1/6] Parsing OBO into DAG...\")\n    go_graph = nx.DiGraph()\n    term_aspect = {}\n    \n    with open(go_obo_path, \"r\") as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == \"[Term]\":\n                cur_id = None\n            elif line.startswith(\"id: GO:\"):\n                cur_id = line.split(\"id: \")[1].strip()\n                go_graph.add_node(cur_id)\n            elif line.startswith(\"namespace: \"):\n                ns = line.split(\"namespace: \")[1].strip()\n                if cur_id:\n                    if ns == 'molecular_function': aspect = 'MF'\n                    elif ns == 'biological_process': aspect = 'BP'\n                    elif ns == 'cellular_component': aspect = 'CC'\n                    else: aspect = 'UNK'\n                    term_aspect[cur_id] = aspect\n            elif line.startswith(\"is_a: \"):\n                pid = line.split()[1].strip()\n                if cur_id:\n                    go_graph.add_edge(cur_id, pid)\n            elif line.startswith(\"relationship: part_of \"):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id:\n                        go_graph.add_edge(cur_id, pid)\n\n    if not nx.is_directed_acyclic_graph(go_graph):\n        cycles = list(nx.simple_cycles(go_graph))\n        for cycle in cycles:\n            go_graph.remove_edge(cycle[0], cycle[1])\n\n    topo_order = list(nx.topological_sort(go_graph))\n    term_parents_map = {n: list(go_graph.successors(n)) for n in go_graph.nodes()}\n    \n    print(f\"      -> Built graph with {len(go_graph)} nodes\")\n    return topo_order, term_parents_map, term_aspect\n\ntopo_order, term_parents_map, term_aspect = build_go_graph(OBO_PATH)\n\n# Load IA scores\nprint(f\"      -> Loading IA scores from {IA_PATH}...\")\nia_df = pd.read_csv(IA_PATH, sep='\\t', header=None, names=['term', 'ia'])\nterm_ia = dict(zip(ia_df['term'], ia_df['ia']))\nprint(f\"      -> Loaded IA for {len(term_ia):,} terms\")\nprint(f\"      -> IA range: [{ia_df['ia'].min():.4f}, {ia_df['ia'].max():.4f}]\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 2. LOAD AND BLEND PREDICTIONS\n# ==========================================\nprint(f\"[2/6] Loading and blending predictions...\")\n\nmerged_preds = defaultdict(lambda: defaultdict(float))\n\n# Load GOA\nif os.path.exists(GOA_PATH):\n    print(f\"      -> Loading GOA...\")\n    goa_count = 0\n    for chunk in pd.read_csv(GOA_PATH, sep='\\t', header=None, \n                              names=['id', 'term', 'score'], \n                              chunksize=500000):\n        for row in chunk.itertuples(index=False):\n            merged_preds[row.id][row.term] += float(row.score) * W_GOA\n            goa_count += 1\n    print(f\"      -> Loaded {goa_count:,} GOA predictions\")\n\n# Load ProtT5\nif os.path.exists(PROTT5_PATH):\n    print(f\"      -> Loading ProtT5...\")\n    prott5_count = 0\n    for chunk in pd.read_csv(PROTT5_PATH, sep='\\t', header=None, \n                              names=['id', 'term', 'score'], \n                              chunksize=500000):\n        for row in chunk.itertuples(index=False):\n            merged_preds[row.id][row.term] += float(row.score) * W_PROTT5\n            prott5_count += 1\n    print(f\"      -> Loaded {prott5_count:,} ProtT5 predictions\")\n\nprint(f\"      -> Combined predictions for {len(merged_preds):,} proteins\")\ngc.collect()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 3. CONVERT TO DATAFRAME\n# ==========================================\nprint(f\"[3/6] Converting to DataFrame...\")\n\nrows = []\nfor protein_id, terms_dict in tqdm(merged_preds.items(), desc=\"Converting\"):\n    for term, score in terms_dict.items():\n        if score >= 0.001:  # Minimal threshold first\n            rows.append((protein_id, term, score))\n\nsubmission = pd.DataFrame(rows, columns=['protein_id', 'go_term', 'score'])\nsubmission['score'] = submission['score'].astype(np.float32)\n\nprint(f\"      -> Created DataFrame with {len(submission):,} predictions\")\n\ndel merged_preds, rows\ngc.collect()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 4. HARD-MAX PROPAGATION\n# ==========================================\ndef propagate_hard_max(df, topo_order, term_parents_map):\n    print(\"[4/6] Propagating Scores (Hard Max)...\")\n    \n    groups = df.groupby('protein_id')\n    final_rows = []\n    \n    for pid, group in tqdm(groups, desc=\"Propagating\"):\n        scores = dict(zip(group['go_term'], group['score']))\n        \n        for child in topo_order:\n            if child in scores:\n                child_score = scores[child]\n                if child in term_parents_map:\n                    for parent in term_parents_map[child]:\n                        current_parent_score = scores.get(parent, 0.0)\n                        if child_score > current_parent_score:\n                            scores[parent] = child_score\n        \n        for term, score in scores.items():\n            if score >= 0.001:\n                final_rows.append((pid, term, score))\n\n    return pd.DataFrame(final_rows, columns=['protein_id', 'go_term', 'score'])\n\nsubmission_propagated = propagate_hard_max(submission, topo_order, term_parents_map)\nprint(f\"      -> After propagation: {len(submission_propagated):,} predictions\")\n\ndel submission\ngc.collect()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 5. IC-AWARE PRUNING\n# ==========================================\nprint(f\"[5/6] Applying IC-aware pruning...\")\n\n# Add IC and aspect columns\nsubmission_propagated['ic'] = submission_propagated['go_term'].map(term_ia).fillna(0.0)\nsubmission_propagated['aspect'] = submission_propagated['go_term'].map(term_aspect).fillna('UNK')\n\nprint(f\"      -> Before pruning: {len(submission_propagated):,} predictions\")\n\n# Filter by minimum IC\nsubmission_filtered = submission_propagated[\n    (submission_propagated['ic'] >= MIN_IC_THRESHOLD) | \n    (submission_propagated['score'] >= 0.5)  # Keep high-confidence even if low-IC\n].copy()\n\nprint(f\"      -> After IC filter: {len(submission_filtered):,} predictions\")\n\n# Apply minimum score threshold\nsubmission_filtered = submission_filtered[submission_filtered['score'] >= MIN_SCORE_THRESHOLD]\nprint(f\"      -> After score filter: {len(submission_filtered):,} predictions\")\n\n# Rank by Score * IC (prioritize specific, high-confidence terms)\nsubmission_filtered['rank_score'] = submission_filtered['score'] * (submission_filtered['ic'] + 0.1)\n\n# Select top N per protein\ndef select_top_per_protein(group):\n    return group.nlargest(MAX_TERMS_PER_PROTEIN, 'rank_score')\n\nsubmission_final = submission_filtered.groupby('protein_id', group_keys=False).apply(select_top_per_protein)\n\nprint(f\"      -> After top-N selection: {len(submission_final):,} predictions\")\nprint(f\"      -> Unique proteins: {submission_final['protein_id'].nunique():,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 6. SAVE SUBMISSION\n# ==========================================\nprint(f\"[6/6] Saving submission ({len(submission_final):,} rows)...\")\n\n# Prepare final output\noutput_df = submission_final[['protein_id', 'go_term', 'score']].copy()\noutput_df['score'] = output_df['score'].clip(0.001, 1.0).round(3)\n\n# Save without header\noutput_df.to_csv(SUBMISSION_OUTPUT, sep='\\t', index=False, header=False)\n\nprint(\"------------------------------------------------\")\nprint(f\"[Done] Saved to {SUBMISSION_OUTPUT}\")\nprint(\"------------------------------------------------\")\nprint(f\"Final stats:\")\nprint(f\"  - Total predictions: {len(output_df):,}\")\nprint(f\"  - Unique proteins: {output_df['protein_id'].nunique():,}\")\nprint(f\"  - Score range: [{output_df['score'].min():.3f}, {output_df['score'].max():.3f}]\")\nprint(f\"\\nParameters used:\")\nprint(f\"  - Blend: GOA={W_GOA}, ProtT5={W_PROTT5}\")\nprint(f\"  - Min IC: {MIN_IC_THRESHOLD}\")\nprint(f\"  - Min Score: {MIN_SCORE_THRESHOLD}\")\nprint(f\"  - Max terms/protein: {MAX_TERMS_PER_PROTEIN}\")\nprint(\"\\nSample:\")\nprint(output_df.head(10))"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Verify file\nif os.path.exists(SUBMISSION_OUTPUT):\n    size = os.path.getsize(SUBMISSION_OUTPUT)\n    print(f\"SUCCESS: {SUBMISSION_OUTPUT} created ({size:,} bytes)\")\nelse:\n    print(f\"ERROR: {SUBMISSION_OUTPUT} not found!\")"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.12"}}, "nbformat": 4, "nbformat_minor": 5}