{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14187577,"sourceType":"datasetVersion","datasetId":9045748},{"sourceId":14452635,"sourceType":"datasetVersion","datasetId":9231252},{"sourceId":287175447,"sourceType":"kernelVersion"},{"sourceId":291608916,"sourceType":"kernelVersion"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nfrom collections import defaultdict, deque\nfrom tqdm.auto import tqdm\nimport pickle\nfrom pathlib import Path\n\n# ==================== CONFIGURATION ====================\nCOMPETITION_DATA = '/kaggle/input/cafa-6-protein-function-prediction'\nPREDICTION_DATA = '/kaggle/input/cafa6-goa-predictions'\nGOA = '/kaggle/input/cafa-6-goa-prott5-ensemble-0-370-f2fcb6/submission.tsv'\n\nCONFIG = {\n    # Ensemble weights\n    \"WEIGHT_GOA\": 0.68,\n    \"WEIGHT_PROTT5\": 0.32,\n    \n    # ProtT5 penalties\n    \"PROTT5_ONLY_PENALTY_BASE\": 0.78,\n    \"PROTT5_ONLY_PENALTY_SUPPORTED\": 0.90,\n    \"PROTT5_ONLY_GAMMA\": 1.15,\n    \"AGREE_BONUS\": 0.03,\n    \n    # Propagation\n    \"UP_DECAY\": 0.985,\n    \"UP_MIN_SCORE\": 0.015,\n    \"CAP_CHILD_BY_PARENT\": True,\n    \"POWER\": 0.92,\n    \"DO_PER_PROTEIN_MAX_NORM\": False,\n    \n    # Output\n    \"TOP_K\": 250,\n    \"SCORE_THRESHOLD\": 0.0005,\n}\n\nROOTS = {\"GO:0003674\", \"GO:0008150\", \"GO:0005575\"}\n\n# ==================== CACHED ONTOLOGY LOADER ====================\nclass OntologyCache:\n    \"\"\"Cache ontology to avoid re-parsing on multiple runs\"\"\"\n    _cache_file = \"ontology_cache.pkl\"\n    \n    @classmethod\n    def load_or_create(cls, obo_path):\n        if Path(cls._cache_file).exists():\n            print(\"Loading cached ontology...\")\n            with open(cls._cache_file, 'rb') as f:\n                return pickle.load(f)\n        \n        print(\"Parsing Ontology (caching for future runs)...\")\n        term_parents = defaultdict(set)\n        \n        with open(obo_path, 'r') as f:\n            lines = f.readlines()\n        \n        current_id = None\n        for line in lines:\n            line = line.strip()\n            if line.startswith(\"id: \"):\n                current_id = line[4:].strip()\n            elif line.startswith(\"is_a: \") and current_id:\n                parent = line.split()[1].strip()\n                term_parents[current_id].add(parent)\n            elif line.startswith(\"relationship: part_of \") and current_id:\n                parent = line.split()[2].strip()\n                term_parents[current_id].add(parent)\n        \n        # Build ancestors with iterative BFS (faster than recursion)\n        ancestors_map = {}\n        for term in term_parents:\n            if term in ancestors_map:\n                continue\n            stack = [term]\n            ancestors = set()\n            while stack:\n                t = stack.pop()\n                if t in ancestors_map:\n                    ancestors.update(ancestors_map[t])\n                    continue\n                for p in term_parents.get(t, []):\n                    ancestors.add(p)\n                    stack.append(p)\n            ancestors_map[term] = ancestors\n        \n        # Save cache\n        with open(cls._cache_file, 'wb') as f:\n            pickle.dump((dict(term_parents), ancestors_map), f)\n        \n        return dict(term_parents), ancestors_map\n\n# ==================== OPTIMIZED DATA LOADER ====================\ndef load_predictions_optimized(filepath, use_numpy=True):\n    \"\"\"Optimized loader using numpy for numerical operations\"\"\"\n    data = {}\n    with open(filepath, 'r') as f:\n        # Pre-allocate batches for faster processing\n        batch_size = 100000\n        batch = []\n        \n        for line in tqdm(f, desc=f\"Loading {os.path.basename(filepath)}\"):\n            parts = line.rstrip(\"\\n\").split(\"\\t\")\n            if len(parts) >= 3:\n                batch.append((parts[0], parts[1], float(parts[2])))\n            \n            if len(batch) >= batch_size:\n                for pid, go, score in batch:\n                    if pid not in data:\n                        data[pid] = {}\n                    data[pid][go] = max(data[pid].get(go, 0.0), score)\n                batch = []\n        \n        # Process remaining\n        for pid, go, score in batch:\n            if pid not in data:\n                data[pid] = {}\n            data[pid][go] = max(data[pid].get(go, 0.0), score)\n    \n    return data\n\n# ==================== VECTORIZED ENSEMBLING ====================\nclass ProteinEnsembler:\n    \"\"\"Optimized ensembler with vectorized operations\"\"\"\n    \n    def __init__(self, config, term_parents, ancestors_map):\n        self.config = config\n        self.term_parents = term_parents\n        self.ancestors_map = ancestors_map\n    \n    def ensemble_protein(self, pid, goa, pt):\n        \"\"\"Ensemble scores for a single protein\"\"\"\n        if not goa and not pt:\n            return {}\n        \n        merged = {}\n        goa_term_set = set(goa.keys())\n        \n        # Vectorize GOA scores (already present)\n        for term, s in goa.items():\n            merged[term] = s\n        \n        # Process ProtT5 scores with numpy operations\n        for term, s_pt in pt.items():\n            s_goa = goa.get(term, 0.0)\n            \n            if s_goa > 0:\n                # Agreement case\n                s = (self.config[\"WEIGHT_GOA\"] * s_goa + \n                     self.config[\"WEIGHT_PROTT5\"] * s_pt)\n                s = min(1.0, s + self.config[\"AGREE_BONUS\"] * min(s_goa, s_pt))\n            else:\n                # ProtT5-only case\n                s_shaped = s_pt ** self.config[\"PROTT5_ONLY_GAMMA\"]\n                \n                # Check support using precomputed ancestors\n                if any(a in goa_term_set for a in self.ancestors_map.get(term, set())):\n                    pen = self.config[\"PROTT5_ONLY_PENALTY_SUPPORTED\"]\n                else:\n                    pen = self.config[\"PROTT5_ONLY_PENALTY_BASE\"]\n                \n                s = s_shaped * pen\n            \n            # Keep max score\n            if s > merged.get(term, 0.0):\n                merged[term] = s\n        \n        return merged\n    \n    def propagate_scores(self, scores):\n        \"\"\"Propagate and constrain scores\"\"\"\n        if not scores:\n            return {r: 1.0 for r in ROOTS}\n        \n        updated = dict(scores)\n        \n        # Fast upward propagation\n        if updated:\n            # Use priority queue for propagation (highest scores first)\n            high_score_terms = [(s, t) for t, s in updated.items() \n                               if s >= self.config[\"UP_MIN_SCORE\"] and t not in ROOTS]\n            high_score_terms.sort(reverse=True)\n            \n            for s_term, term in high_score_terms:\n                for parent in self.term_parents.get(term, []):\n                    s_parent = s_term * self.config[\"UP_DECAY\"]\n                    if s_parent > updated.get(parent, 0.0):\n                        updated[parent] = s_parent\n        \n        # Downward constraint (True Path Rule)\n        if self.config[\"CAP_CHILD_BY_PARENT\"] and updated:\n            # Single efficient pass\n            for term, score in list(updated.items()):\n                if term in ROOTS:\n                    continue\n                \n                parents = self.term_parents.get(term)\n                if not parents:\n                    continue\n                \n                parent_scores = [updated.get(p, 0.0) for p in parents if p in updated]\n                if parent_scores:\n                    best_parent = max(parent_scores)\n                    if score > best_parent:\n                        updated[term] = 0.5 * score + 0.5 * best_parent\n        \n        # Apply power scaling\n        if self.config[\"POWER\"] != 1.0:\n            for term in list(updated.keys()):\n                if term not in ROOTS:\n                    updated[term] = updated[term] ** self.config[\"POWER\"]\n        \n        # Add roots\n        for r in ROOTS:\n            updated[r] = 1.0\n        \n        return updated\n\n# ==================== QUICK PREDICTION FUNCTION ====================\ndef make_quick_predictions(proteins_to_predict=None, limit=None):\n    \"\"\"Optimized main function with quick prediction capability\"\"\"\n    \n    # Load ontology\n    obo_path = f\"{COMPETITION_DATA}/Train/go-basic.obo\"\n    term_parents, ancestors_map = OntologyCache.load_or_create(obo_path)\n    \n    # Load predictions (only for requested proteins if specified)\n    print(\"Loading predictions...\")\n    goa_preds = load_predictions_optimized(GOA)\n    prott5_preds = load_predictions_optimized(f\"{PREDICTION_DATA}/prott5_interpro_predictions.tsv\")\n    \n    # Initialize ensembler\n    ensembler = ProteinEnsembler(CONFIG, term_parents, ancestors_map)\n    \n    # Determine which proteins to process\n    all_proteins = set(goa_preds.keys()) | set(prott5_preds.keys())\n    if proteins_to_predict:\n        proteins_to_process = [p for p in proteins_to_predict if p in all_proteins]\n    else:\n        proteins_to_process = list(all_proteins)\n    \n    if limit:\n        proteins_to_process = proteins_to_process[:limit]\n    \n    print(f\"Processing {len(proteins_to_process)} proteins...\")\n    \n    # Process proteins\n    final_rows = []\n    for pid in tqdm(proteins_to_process, desc=\"Ensembling & Propagating\"):\n        goa = goa_preds.get(pid, {})\n        pt = prott5_preds.get(pid, {})\n        \n        # Ensemble\n        ensemble_scores = ensembler.ensemble_protein(pid, goa, pt)\n        \n        # Propagate\n        propagated = ensembler.propagate_scores(ensemble_scores)\n        \n        # Filter and select top K\n        items = [(t, s) for t, s in propagated.items() \n                if s >= CONFIG[\"SCORE_THRESHOLD\"]]\n        items.sort(key=lambda x: x[1], reverse=True)\n        \n        if len(items) > CONFIG[\"TOP_K\"]:\n            items = items[:CONFIG[\"TOP_K\"]]\n        \n        # Format output\n        for term, score in items:\n            final_rows.append(f\"{pid}\\t{term}\\t{score:.5f}\")\n    \n    # Write output\n    OUTPUT_FILE = \"submission.tsv\"\n    with open(OUTPUT_FILE, \"w\") as f:\n        f.write(\"\\n\".join(final_rows))\n    \n    print(f\"File saved: {OUTPUT_FILE}\")\n    print(f\"Total Predictions: {len(final_rows):,}\")\n    \n    return final_rows\n\n# ==================== UTILITY FOR BATCH PROCESSING ====================\ndef batch_predict(protein_batches, batch_size=1000):\n    \"\"\"Process proteins in batches for memory efficiency\"\"\"\n    all_predictions = []\n    \n    for i in range(0, len(protein_batches), batch_size):\n        batch = protein_batches[i:i + batch_size]\n        print(f\"Processing batch {i//batch_size + 1}/{(len(protein_batches)-1)//batch_size + 1}\")\n        predictions = make_quick_predictions(batch)\n        all_predictions.extend(predictions)\n    \n    return all_predictions\n\n# ==================== MAIN EXECUTION ====================\nif __name__ == \"__main__\":\n    # Quick single protein prediction\n    # predictions = make_quick_predictions([\"protein_id_1\", \"protein_id_2\"])\n    \n    # Full prediction (same as original)\n    make_quick_predictions()\n    \n    # For batch processing large datasets:\n    # all_proteins = [...]  # Your list of proteins\n    # batch_predict(all_proteins, batch_size=5000)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import os\n# import gc\n# import pandas as pd\n# import numpy as np\n# from collections import defaultdict\n# from tqdm.auto import tqdm\n# from multiprocessing import Pool, cpu_count\n# import warnings\n# warnings.filterwarnings('ignore')\n\n# # ==========================================\n# # 1. ULTRA-FAST OBO PARSER (Vectorized)\n# # ==========================================\n# def parse_obo_parents(go_obo_path):\n#     \"\"\"Optimized OBO parser using list operations\"\"\"\n#     print(f\"[1/5] Parsing OBO Ontology...\")\n#     term_parents = {}\n#     roots = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\n    \n#     with open(go_obo_path, \"r\") as f:\n#         lines = f.readlines()\n    \n#     cur_id = None\n#     parent_list = []\n    \n#     for line in lines:\n#         line = line.strip()\n#         if line == \"[Term]\":\n#             if cur_id and parent_list:\n#                 term_parents[cur_id] = set(parent_list)\n#             cur_id = None\n#             parent_list = []\n#         elif line.startswith(\"id: \"):\n#             cur_id = line[4:].strip()\n#         elif line.startswith(\"is_a: \"):\n#             parent_list.append(line.split()[1])\n#         elif line.startswith(\"relationship: part_of \"):\n#             parts = line.split()\n#             if len(parts) >= 3:\n#                 parent_list.append(parts[2])\n    \n#     if cur_id and parent_list:\n#         term_parents[cur_id] = set(parent_list)\n    \n#     return term_parents, roots\n\n# def get_ancestors_map_fast(term_parents):\n#     \"\"\"Iterative BFS approach - much faster than recursion\"\"\"\n#     print(\"[1/5] Building Ancestor Map (Fast BFS)...\")\n#     ancestors = {}\n    \n#     # Process in batches\n#     all_terms = list(term_parents.keys())\n    \n#     for term in tqdm(all_terms, desc=\"Computing ancestors\"):\n#         if term in ancestors:\n#             continue\n            \n#         visited = set()\n#         queue = list(term_parents.get(term, set()))\n        \n#         while queue:\n#             current = queue.pop(0)\n#             if current in visited:\n#                 continue\n#             visited.add(current)\n#             queue.extend(term_parents.get(current, set()))\n        \n#         ancestors[term] = visited\n    \n#     return ancestors\n\n# # ==========================================\n# # 2. VECTORIZED PROCESSING (NumPy Based)\n# # ==========================================\n# def process_predictions_vectorized(df, ancestors_map, roots):\n#     \"\"\"Ultra-fast vectorized processing using NumPy and dict operations\"\"\"\n#     print(\"[3/5] Processing Predictions (Vectorized Approach)...\")\n    \n#     # Convert to numpy for speed\n#     proteins = df['protein_id'].values\n#     terms = df['go_term'].values\n#     scores = df['score'].values.astype(np.float32)\n    \n#     # Group by protein using dict\n#     protein_data = defaultdict(lambda: defaultdict(float))\n    \n#     print(\"   → Grouping data...\")\n#     for i in tqdm(range(len(proteins)), desc=\"Grouping\", disable=True):\n#         p, t, s = proteins[i], terms[i], scores[i]\n#         protein_data[p][t] = max(protein_data[p][t], s)\n    \n#     # Pre-allocate result lists\n#     result_proteins = []\n#     result_terms = []\n#     result_scores = []\n    \n#     print(\"   → Propagating & normalizing...\")\n#     for pid, terms_dict in tqdm(protein_data.items(), desc=\"Processing\"):\n#         final_scores = {}\n        \n#         # Copy original scores\n#         for t, s in terms_dict.items():\n#             final_scores[t] = s\n        \n#         # PROPAGATION: Parent >= Child\n#         for term, score in terms_dict.items():\n#             if term in ancestors_map:\n#                 for anc in ancestors_map[term]:\n#                     final_scores[anc] = max(final_scores.get(anc, 0.0), score)\n        \n#         # FORCE ROOTS to 1.0\n#         if final_scores:\n#             for r in roots:\n#                 final_scores[r] = 1.0\n        \n#         # NORMALIZATION: Boost max to 0.95\n#         non_root_scores = [s for t, s in final_scores.items() if t not in roots]\n#         if non_root_scores:\n#             max_val = max(non_root_scores)\n#             if 0 < max_val < 0.95:\n#                 scale = 0.95 / max_val\n#                 for t in final_scores:\n#                     if t not in roots:\n#                         final_scores[t] = min(1.0, final_scores[t] * scale)\n        \n#         # Collect results (filter low scores)\n#         for go_term, score in final_scores.items():\n#             if score >= 0.001:\n#                 result_proteins.append(pid)\n#                 result_terms.append(go_term)\n#                 result_scores.append(score)\n    \n#     # Create DataFrame efficiently\n#     print(\"   → Creating output DataFrame...\")\n#     return pd.DataFrame({\n#         'protein_id': result_proteins,\n#         'go_term': result_terms,\n#         'score': result_scores\n#     })\n\n# # ==========================================\n# # 3. CHUNKED PROCESSING FOR LARGE FILES\n# # ==========================================\n# def process_in_chunks(submission_path, ancestors_map, roots, chunk_size=2_000_000):\n#     \"\"\"Process large files in chunks to avoid memory issues\"\"\"\n#     print(f\"[2/5] Loading submission in chunks...\")\n    \n#     chunks_processed = []\n#     chunk_num = 0\n    \n#     for chunk in pd.read_csv(submission_path, sep='\\t', header=None, \n#                               names=['protein_id', 'go_term', 'score', 'key'],\n#                               usecols=['protein_id', 'go_term', 'score'],\n#                               chunksize=chunk_size,\n#                               dtype={'protein_id': str, 'go_term': str, 'score': np.float32}):\n        \n#         chunk_num += 1\n#         print(f\"   Processing chunk {chunk_num} ({len(chunk):,} rows)...\")\n        \n#         processed = process_predictions_vectorized(chunk, ancestors_map, roots)\n#         chunks_processed.append(processed)\n        \n#         del chunk\n#         gc.collect()\n    \n#     print(\"   → Combining all chunks...\")\n#     return pd.concat(chunks_processed, ignore_index=True)\n\n# # ==========================================\n# # 4. MAIN PIPELINE (OPTIMIZED)\n# # ==========================================\n# def main():\n#     # Paths\n#     OBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n#     SUBMISSION_INPUT = '/kaggle/input/goa-negative-propagation/submission.tsv'\n#     SUBMISSION_OUTPUT = 'submission.tsv'\n    \n#     print(\"=\"*60)\n#     print(\"  CAFA-6 OPTIMIZED SUBMISSION GENERATOR\")\n#     print(\"=\"*60)\n    \n#     # Step 1: Parse Ontology (Fast)\n#     term_parents, roots = parse_obo_parents(OBO_PATH)\n#     ancestors_map = get_ancestors_map_fast(term_parents)\n#     print(f\"   ✓ Loaded {len(term_parents):,} GO terms\")\n#     print(f\"   ✓ Computed {len(ancestors_map):,} ancestor relationships\")\n    \n#     gc.collect()\n    \n#     # Step 2 & 3: Process Predictions\n#     try:\n#         # Try loading full file first\n#         print(f\"[2/5] Attempting full file load...\")\n#         submission = pd.read_csv(SUBMISSION_INPUT, sep='\\t', header=None,\n#                                 names=['protein_id', 'go_term', 'score', 'key'],\n#                                 usecols=['protein_id', 'go_term', 'score'],\n#                                 dtype={'protein_id': str, 'go_term': str, 'score': np.float32})\n        \n#         print(f\"   ✓ Loaded {len(submission):,} predictions\")\n#         final_df = process_predictions_vectorized(submission, ancestors_map, roots)\n        \n#         del submission\n#         gc.collect()\n        \n#     except MemoryError:\n#         print(\"   ⚠ Memory limit - switching to chunked processing...\")\n#         final_df = process_in_chunks(SUBMISSION_INPUT, ancestors_map, roots)\n    \n#     # Step 4: Remove duplicates (keep highest score)\n#     print(f\"[4/5] Deduplicating {len(final_df):,} rows...\")\n#     final_df = final_df.groupby(['protein_id', 'go_term'], as_index=False)['score'].max()\n    \n#     # Step 5: Sort & Save\n#     print(f\"[5/5] Sorting and saving {len(final_df):,} rows...\")\n#     final_df.sort_values(['protein_id', 'score'], ascending=[True, False], inplace=True)\n#     final_df.to_csv(SUBMISSION_OUTPUT, sep='\\t', index=False, header=False)\n    \n#     print(\"=\"*60)\n#     print(f\"✅ COMPLETE! Saved to: {SUBMISSION_OUTPUT}\")\n#     print(f\"   Total 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#     print(\"=\"*60)\n#     print(\"\\nFirst 10 rows:\")\n#     print(final_df.head(10))\n    \n#     return final_df\n\n# # ==========================================\n# # RUN\n# # ==========================================\n# if __name__ == \"__main__\":\n#     final_df = main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}