{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":116062,"databundleVersionId":14875579},{"sourceType":"datasetVersion","sourceId":14452635,"datasetId":9231252,"databundleVersionId":15272839},{"sourceType":"datasetVersion","sourceId":14187577,"datasetId":9045748,"databundleVersionId":14981293},{"sourceType":"kernelVersion","sourceId":289631627},{"sourceType":"kernelVersion","sourceId":292870115},{"sourceType":"kernelVersion","sourceId":293181706}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","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\"\nSUBMISSION_INPUT = '/kaggle/input/cafa-ensemble-0-378/submission.tsv'\nSUBMISSION_OUTPUT = 'submission.tsv'\n\n# Adjustable parameters for score improvement\nMIN_SCORE_THRESHOLD = 0.001  # Remove noise\nSCORE_ROUNDING = 3  # Decimal places\nUSE_SCORE_BOOST = True  # Enable intelligent boosting\nUSE_PERCENTILE_SCALING = True  # Scale based on percentile distribution\nTOP_K_PREDICTIONS = None  # Set to limit predictions per protein (None = no limit)\n\n# ==========================================\n# 1. IMPROVED GRAPH BUILDER WITH ANCESTOR CACHE\n# ==========================================\ndef build_go_graph_with_cache(go_obo_path):\n    \"\"\"Build GO graph with cached ancestors for faster propagation\"\"\"\n    print(f\"[1/5] Parsing OBO and building ancestor cache...\")\n    go_graph = nx.DiGraph()\n    \n    if not os.path.exists(go_obo_path):\n        raise FileNotFoundError(f\"OBO not found: {go_obo_path}\")\n\n    # Parse OBO\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(\"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    print(\"      -> Validating DAG and computing ancestors...\")\n    \n    # Remove cycles if any\n    try:\n        cycles = list(nx.simple_cycles(go_graph))\n        if cycles:\n            print(f\"      -> Removing {len(cycles)} cycles...\")\n            for cycle in cycles:\n                if len(cycle) > 1:\n                    go_graph.remove_edge(cycle[0], cycle[1])\n    except:\n        pass\n    \n    # Build ancestor cache (much faster than repeated traversals)\n    print(\"      -> Building ancestor cache...\")\n    ancestors_cache = {}\n    topo_order = list(nx.topological_sort(go_graph))\n    \n    # Process in reverse topological order (roots to leaves)\n    for node in reversed(topo_order):\n        all_ancestors = set()\n        for successor in go_graph.successors(node):\n            all_ancestors.add(successor)\n            all_ancestors.update(ancestors_cache.get(successor, set()))\n        ancestors_cache[node] = all_ancestors\n    \n    print(f\"      -> Graph built: {len(go_graph.nodes())} terms, {len(go_graph.edges())} edges\")\n    \n    return go_graph, ancestors_cache, topo_order\n\n# ==========================================\n# 2. INTELLIGENT SCORE OPTIMIZATION\n# ==========================================\ndef optimize_scores_for_protein(protein_scores, ancestors_cache, protein_id=\"\"):\n    \"\"\"\n    Optimize scores for a single protein with multiple strategies\n    \"\"\"\n    if not protein_scores:\n        return protein_scores\n    \n    # Convert to dict for faster operations\n    scores_dict = protein_scores.copy()\n    \n    # STRATEGY 1: HARD MAX PROPAGATION\n    propagated_scores = scores_dict.copy()\n    \n    # Propagate scores upward using ancestor cache\n    for term, score in sorted(scores_dict.items(), key=lambda x: x[1], reverse=True):\n        if score <= 0:\n            continue\n            \n        # Get all ancestors from cache\n        for ancestor in ancestors_cache.get(term, set()):\n            current = propagated_scores.get(ancestor, 0.0)\n            if score > current:\n                propagated_scores[ancestor] = score\n    \n    # STRATEGY 2: SELECTIVE SCORE BOOSTING\n    if USE_SCORE_BOOST and propagated_scores:\n        # Identify top predictions to boost\n        scores_list = list(propagated_scores.items())\n        scores_list.sort(key=lambda x: x[1], reverse=True)\n        \n        # Only boost if we have meaningful predictions\n        if len(scores_list) >= 3:\n            top_score = scores_list[0][1]\n            second_score = scores_list[1][1]\n            \n            # Boost strategy based on score distribution\n            if top_score < 0.8:\n                # If top score is low, boost it moderately\n                boost_factor = min(1.5, 0.85 / max(top_score, 0.01))\n                \n                # Apply boost with diminishing returns\n                for i, (term, score) in enumerate(scores_list):\n                    if i < 5:  # Boost top 5 predictions\n                        # Stronger boost for higher ranks\n                        rank_boost = 1.0 + (0.3 * (1.0 - i/5))\n                        new_score = min(1.0, score * boost_factor * rank_boost)\n                        propagated_scores[term] = new_score\n                    elif score < 0.1 and i > 10:\n                        # Suppress very low scores that aren't top predictions\n                        propagated_scores[term] = score * 0.7\n    \n    # STRATEGY 3: PERCENTILE-BASED SCALING\n    if USE_PERCENTILE_SCALING and propagated_scores:\n        scores_array = np.array(list(propagated_scores.values()))\n        \n        if len(scores_array) > 10:\n            # Calculate percentiles\n            p90 = np.percentile(scores_array, 90)\n            p50 = np.percentile(scores_array, 50)\n            \n            if p90 > 0 and p50 > 0:\n                # Scale to make distribution more favorable for CAFA\n                # CAFA's F-max benefits from confident predictions\n                scale_factor = min(2.0, 0.7 / max(p50, 0.01))\n                \n                # Apply scaling with upper bound\n                for term in list(propagated_scores.keys()):\n                    new_score = min(1.0, propagated_scores[term] * scale_factor)\n                    propagated_scores[term] = new_score\n    \n    # STRATEGY 4: ENSURE ROOT TERMS\n    # In GO, these are the main roots (optional, but can help consistency)\n    roots = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\n    for root in roots:\n        if root in ancestors_cache or any(root in ancestors_cache.get(t, set()) for t in propagated_scores):\n            # If any prediction exists in this branch, ensure root is predicted\n            max_score = max([propagated_scores.get(t, 0) for t in propagated_scores \n                           if root in ancestors_cache.get(t, set())] + [0])\n            if max_score > 0.1:\n                propagated_scores[root] = max(propagated_scores.get(root, 0), 0.95)\n    \n    # Filter low scores\n    filtered_scores = {k: v for k, v in propagated_scores.items() \n                      if v >= MIN_SCORE_THRESHOLD}\n    \n    return filtered_scores\n\n# ==========================================\n# 3. CHUNKED PROCESSING FOR MEMORY EFFICIENCY\n# ==========================================\ndef process_submission_chunked(submission_df, ancestors_cache, chunk_size=50000):\n    \"\"\"Process large submissions in chunks to save memory\"\"\"\n    print(f\"[3/5] Processing predictions in chunks...\")\n    \n    # Get unique proteins\n    unique_proteins = submission_df['protein_id'].unique()\n    num_chunks = (len(unique_proteins) + chunk_size - 1) // chunk_size\n    \n    all_results = []\n    \n    for chunk_idx in tqdm(range(num_chunks), desc=\"Processing chunks\"):\n        start_idx = chunk_idx * chunk_size\n        end_idx = min((chunk_idx + 1) * chunk_size, len(unique_proteins))\n        chunk_proteins = unique_proteins[start_idx:end_idx]\n        \n        # Filter dataframe for this chunk\n        chunk_df = submission_df[submission_df['protein_id'].isin(chunk_proteins)]\n        \n        # Process each protein in chunk\n        grouped = chunk_df.groupby('protein_id')\n        \n        for pid, group in grouped:\n            # Convert to dict\n            protein_scores = dict(zip(group['go_term'], group['score']))\n            \n            # Optimize scores\n            optimized_scores = optimize_scores_for_protein(protein_scores, ancestors_cache, pid)\n            \n            # Collect results\n            for term, score in optimized_scores.items():\n                all_results.append((pid, term, float(score)))\n        \n        # Clear memory\n        if chunk_idx % 10 == 0:\n            gc.collect()\n    \n    return pd.DataFrame(all_results, columns=['protein_id', 'go_term', 'score'])\n\n# ==========================================\n# 4. POST-PROCESSING OPTIMIZATIONS\n# ==========================================\ndef post_process_predictions(df):\n    \"\"\"Apply final optimizations to improve CAFA score\"\"\"\n    print(f\"[4/5] Post-processing predictions...\")\n    \n    # 1. Remove exact duplicates\n    df = df.drop_duplicates(subset=['protein_id', 'go_term'])\n    \n    # 2. Sort by protein and score\n    df = df.sort_values(['protein_id', 'score'], ascending=[True, False])\n    \n    # 3. Apply top-k filtering if enabled\n    if TOP_K_PREDICTIONS:\n        print(f\"      -> Limiting to top {TOP_K_PREDICTIONS} predictions per protein\")\n        df = df.groupby('protein_id').head(TOP_K_PREDICTIONS).reset_index(drop=True)\n    \n    # 4. Round scores to reduce file size and help thresholding\n    df['score'] = df['score'].round(SCORE_ROUNDING)\n    \n    # 5. Final threshold (after rounding)\n    df = df[df['score'] >= MIN_SCORE_THRESHOLD]\n    \n    # 6. Ensure no scores exceed 1.0\n    df['score'] = df['score'].clip(0.0, 1.0)\n    \n    return df\n\n# ==========================================\n# 5. MAIN PIPELINE\n# ==========================================\ndef main():\n    print(\"=\" * 70)\n    print(\"CAFA 6 SUBMISSION OPTIMIZER\")\n    print(f\"Target: Improve from baseline 0.378\")\n    print(\"=\" * 70)\n    \n    # 1. Build GO graph with ancestor cache\n    go_graph, ancestors_cache, topo_order = build_go_graph_with_cache(OBO_PATH)\n    \n    # 2. Load submission\n    print(f\"\\n[2/5] Loading submission from {SUBMISSION_INPUT}...\")\n    \n    # Check file size to decide on chunking\n    file_size = os.path.getsize(SUBMISSION_INPUT)\n    print(f\"      -> File size: {file_size / 1024 / 1024:.1f} MB\")\n    \n    # Load with appropriate settings\n    submission = pd.read_csv(\n        SUBMISSION_INPUT, \n        sep='\\t', \n        header=None,\n        names=['protein_id', 'go_term', 'score'],\n        dtype={'protein_id': str, 'go_term': str, 'score': np.float32},\n        on_bad_lines='skip'\n    )\n    \n    print(f\"      -> Loaded {len(submission):,} predictions\")\n    print(f\"      -> Unique proteins: {submission['protein_id'].nunique():,}\")\n    \n    # Basic stats\n    print(f\"\\n      Score Statistics:\")\n    print(f\"        Min: {submission['score'].min():.4f}\")\n    print(f\"        Max: {submission['score'].max():.4f}\")\n    print(f\"        Mean: {submission['score'].mean():.4f}\")\n    print(f\"        Median: {submission['score'].median():.4f}\")\n    \n    # 3. Process predictions\n    if len(submission) > 1000000:\n        # Use chunked processing for large files\n        print(f\"\\n      -> Large file detected, using chunked processing...\")\n        optimized_df = process_submission_chunked(submission, ancestors_cache)\n    else:\n        # Process all at once\n        print(f\"\\n      -> Processing all predictions at once...\")\n        all_results = []\n        grouped = submission.groupby('protein_id')\n        \n        for pid, group in tqdm(grouped, total=len(grouped), desc=\"Processing proteins\"):\n            protein_scores = dict(zip(group['go_term'], group['score']))\n            optimized_scores = optimize_scores_for_protein(protein_scores, ancestors_cache, pid)\n            \n            for term, score in optimized_scores.items():\n                all_results.append((pid, term, float(score)))\n        \n        optimized_df = pd.DataFrame(all_results, columns=['protein_id', 'go_term', 'score'])\n    \n    # 4. Post-processing\n    final_df = post_process_predictions(optimized_df)\n    \n    # 5. Save\n    print(f\"\\n[5/5] Saving optimized submission...\")\n    final_df.to_csv(SUBMISSION_OUTPUT, sep='\\t', index=False, header=False)\n    \n    # 6. Final statistics\n    print(\"\\n\" + \"=\" * 70)\n    print(\"OPTIMIZATION COMPLETE\")\n    print(\"=\" * 70)\n    print(f\"\\nResults Summary:\")\n    print(f\"  • Input predictions: {len(submission):,}\")\n    print(f\"  • Output predictions: {len(final_df):,}\")\n    print(f\"  • Unique proteins: {final_df['protein_id'].nunique():,}\")\n    print(f\"  • Unique GO terms: {final_df['go_term'].nunique():,}\")\n    \n    print(f\"\\nScore Distribution (optimized):\")\n    for threshold in [0.01, 0.1, 0.3, 0.5, 0.7, 0.9]:\n        count = (final_df['score'] >= threshold).sum()\n        percentage = (count / len(final_df)) * 100\n        print(f\"  • ≥{threshold:.2f}: {count:,} ({percentage:.1f}%)\")\n    \n    \n    print(f\"\\nOutput saved to: {SUBMISSION_OUTPUT}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T21:11:39.059113Z","iopub.execute_input":"2026-01-21T21:11:39.059384Z","iopub.status.idle":"2026-01-21T21:13:45.269403Z","shell.execute_reply.started":"2026-01-21T21:11:39.059357Z","shell.execute_reply":"2026-01-21T21:13:45.268154Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}