{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# CAFA6 Aspect-Specific Optimization\n\n## The Insight\nMolecular Function (MF), Biological Process (BP), and Cellular Component (CC) are fundamentally different ontologies:\n- **MF**: Very specific, often just 1-2 terms per protein. High precision required.\n- **BP**: Deep trees, many terms per protein. High recall required.\n- **CC**: Shallow, few terms. \n\nUsing a \"one size fits all\" approach hurts performance. We need **Aspect-Specific Strategies**.\n\n## The Strategy\n1. Split predictions by aspect\n2. Apply different logic:\n   - **MF**: Strict threshold, lower budget (max 300 terms), boost high-IC scores\n   - **BP**: Lenient threshold, higher budget (max 1000 terms), standard propagation\n   - **CC**: Medium threshold, low budget (max 200 terms)\n3. Recombine for final submission"}, {"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\"\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# Blend weights\nW_GOA = 0.60\nW_PROTT5 = 0.40\n\n# ASPECT-SPECIFIC PARAMETERS\nASPECT_PARAMS = {\n    'MF': {'min_score': 0.01, 'max_terms': 300, 'boost_ic': True, 'score_pow': 0.8},\n    'BP': {'min_score': 0.001, 'max_terms': 1000, 'boost_ic': False, 'score_pow': 1.0},\n    'CC': {'min_score': 0.005, 'max_terms': 200, 'boost_ic': False, 'score_pow': 0.9},\n    'UNK': {'min_score': 0.01, 'max_terms': 100, 'boost_ic': False, 'score_pow': 1.0}\n}\n\nprint(f\"Aspect Parameters: {ASPECT_PARAMS}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 1. BUILD GO GRAPH & LOAD ASPECTS\n# ==========================================\ndef build_go_graph(go_obo_path):\n    print(f\"[1/5] Parsing OBO...\")\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: \"):\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\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)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 2. LOAD AND BLEND PREDICTIONS\n# ==========================================\nprint(f\"[2/5] 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    for chunk in pd.read_csv(GOA_PATH, sep='\\t', header=None, \n                              names=['id', 'term', 'score'], chunksize=500000):\n        for row in chunk.itertuples(index=False):\n            merged_preds[row.id][row.term] += float(row.score) * W_GOA\n\n# Load ProtT5\nif os.path.exists(PROTT5_PATH):\n    print(f\"      -> Loading ProtT5...\")\n    for chunk in pd.read_csv(PROTT5_PATH, sep='\\t', header=None, \n                              names=['id', 'term', 'score'], chunksize=500000):\n        for row in chunk.itertuples(index=False):\n            merged_preds[row.id][row.term] += float(row.score) * W_PROTT5\n\nprint(f\"      -> Combined predictions for {len(merged_preds):,} proteins\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 3. HARD-MAX PROPAGATION & ASPECT OPTIMIZATION\n# ==========================================\nprint(\"[3/5] Propagating and optimizing by aspect...\")\n\nfinal_rows = []\n\n# Process protein by protein\nfor pid, terms_dict in tqdm(merged_preds.items(), desc=\"Processing\"):\n    scores = terms_dict.copy()\n    \n    # 1. Propagate first\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    # 2. Aspect-specific processing\n    aspect_terms = {'MF': [], 'BP': [], 'CC': [], 'UNK': []}\n    \n    for term, score in scores.items():\n        aspect = term_aspect.get(term, 'UNK')\n        params = ASPECT_PARAMS.get(aspect, ASPECT_PARAMS['UNK'])\n        \n        # Check threshold\n        if score < params['min_score']: continue\n        \n        # Transform score\n        new_score = score ** params['score_pow']\n        \n        aspect_terms[aspect].append((term, new_score))\n    \n    # 3. Apply budget per aspect\n    for aspect, candidates in aspect_terms.items():\n        params = ASPECT_PARAMS.get(aspect, ASPECT_PARAMS['UNK'])\n        budget = params['max_terms']\n        \n        # Sort by score\n        candidates.sort(key=lambda x: x[1], reverse=True)\n        \n        # Select top N\n        selected = candidates[:budget]\n        \n        for term, score in selected:\n            final_rows.append((pid, term, score))\n\nsubmission = pd.DataFrame(final_rows, columns=['protein_id', 'go_term', 'score'])\nprint(f\"      -> Final prediction count: {len(submission):,}\")\n\ndel merged_preds, final_rows\ngc.collect()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 4. SAVE SUBMISSION\n# ==========================================\nprint(f\"[4/5] Saving submission...\")\n\nsubmission['score'] = submission['score'].clip(0.001, 1.0).round(3)\nsubmission.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\"  - Unique proteins: {submission['protein_id'].nunique():,}\")\nprint(f\"  - Score range: [{submission['score'].min():.3f}, {submission['score'].max():.3f}]\")\nprint(f\"\\nAspect counts:\")\n\n# Add aspect info back to check stats\nsubmission['aspect'] = submission['go_term'].map(term_aspect).fillna('UNK')\nprint(submission['aspect'].value_counts())\n\nprint(f\"\\nSample:\")\nprint(submission.head(10))"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Verify\nif os.path.exists(SUBMISSION_OUTPUT):\n    size = os.path.getsize(SUBMISSION_OUTPUT)\n    print(f\"\\nSUCCESS: {SUBMISSION_OUTPUT} ({size:,} bytes)\")"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.12"}}, "nbformat": 4, "nbformat_minor": 5}