{"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":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":292870115,"sourceType":"kernelVersion"}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\nimport os\nimport gc\nimport pandas as pd\nimport numpy as np\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm\n\n# ==========================================\n# 1. OPTIMIZED OBO PARSER\n# ==========================================\ndef parse_obo_parents(go_obo_path):\n    print(f\"[1/5] Parsing OBO Ontology...\")\n    term_parents = defaultdict(set)\n    roots = set(['GO:0005575', 'GO:0015404', 'GO:0025045'])\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            elif line.startswith(\"is_a: \"):\n                pid = line.split()[1].strip()\n                if cur_id: term_parents[cur_id].add(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: term_parents[cur_id].add(pid)\n    return term_parents, roots\n\ndef get_ancestors_map(term_parents):\n    print(\"[1/5] Building Ancestor Map...\")\n    ancestors = {}\n    def get_all_ancestors(term):\n        if term in ancestors: return ancestors[term]\n        parents = term_parents.get(term, set())\n        all_anc = set(parents)\n        for p in parents:\n            all_anc |= get_all_ancestors(p)\n        ancestors[term] = all_anc\n        return all_anc\n    \n    for term in tqdm(list(term_parents.keys())):\n        get_all_ancestors(term)\n    return ancestors\n\n# ==========================================\n# 2. LOGIC: PROPAGATION + NORMALIZATION\n# ==========================================\ndef process_predictions(df, ancestors_map, roots):\n    print(\"[3/5] Processing Predictions (Propagate + Normalize)...\")\n    \n    # 1. Convert to Dict Structure for fast access\n    # { protein_id: { go_term: score } }\n    data_map = defaultdict(dict)\n    values = df.values # protein, term, score\n    \n    for pid, term, score in tqdm(values, desc=\"Grouping\"):\n        data_map[pid][term] = float(score)\n    \n    new_rows = []\n    \n    # 2. Iterate Per Protein\n    for pid, terms_dict in tqdm(data_map.items(), desc=\"Optimizing\"):\n        final_scores = terms_dict.copy()\n        \n        # --- A. POSITIVE PROPAGATION ---\n        # Ensure Parent >= Child\n        # Sort terms to process leaves first (optional but helps)\n        current_terms = list(terms_dict.keys())\n        for term in current_terms:\n            s = terms_dict[term]\n            if term in ancestors_map:\n                for anc in ancestors_map[term]:\n                    final_scores[anc] = max(final_scores.get(anc, 0.0), s)\n        \n        # --- B. FORCE ROOTS ---\n        # If we have any prediction, the roots must be 1.0\n        if len(final_scores) > 0:\n            for r in roots:\n                final_scores[r] = 1.0\n        \n        # --- C. RANK NORMALIZATION (The Boost) ---\n        # Find the max score for this protein (excluding roots usually 1.0)\n        # We want to boost the best \"non-root\" prediction to a high confidence\n        # to ensure it survives thresholding.\n        \n        max_val = 0.0\n        for t, s in final_scores.items():\n            if t not in roots:\n                max_val = max(max_val, s)\n        \n        # If the best prediction is weak (e.g., 0.3), scale everything up\n        # Target: Make the max score at least 0.95\n        if max_val > 0 and max_val < 0.65:\n            scale_factor = 0.95 / max_val\n            for t in final_scores:\n                if t not in roots: # Don't scale roots > 1.0\n                    final_scores[t] = min(1.0, final_scores[t] * scale_factor)\n\n        # Collect\n        for go_term, score in final_scores.items():\n            # Optimization: Drop extremely low scores to reduce file size\n            if score >= 0.001:\n                new_rows.append((pid, go_term, score))\n                \n    return pd.DataFrame(new_rows, columns=['protein_id', 'go_term', 'score'])\n\n# ==========================================\n# 3. MAIN PIPELINE\n# ==========================================\n\n# Paths\nOBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nSUBMISSION_INPUT = '/kaggle/input/cafa6-protein-function-enhanced-nb-v2/submission.tsv'\nSUBMISSION_OUTPUT = 'submission.tsv'\n\n# 1. Load Ontology\nterm_parents, roots = parse_obo_parents(OBO_PATH)\nancestors_map = get_ancestors_map(term_parents)\n\n# 2. Load Submission\nprint(f\"[2/5] Loading submission...\")\n# If OOM occurs, add nrows=10_000_000 or use chunks\nsubmission = pd.read_csv(SUBMISSION_INPUT, sep='\\t', header=None, names=['protein_id', 'go_term', 'score', 'key'])\nsubmission = submission[['protein_id', 'go_term', 'score']]\n\n# 3. Process\nfinal_df = process_predictions(submission, ancestors_map, roots)\n\n# 4. Save\nprint(f\"[4/5] Saving {len(final_df)} rows...\")\n# Sorting ensures deterministic file for hashing/checks\nfinal_df.sort_values(['protein_id', 'score'], ascending=[True, False], inplace=True)\nfinal_df.to_csv(SUBMISSION_OUTPUT, sep='\\t', index=False, header=False)\n\nprint(f\"[✅] Done. Saved to {SUBMISSION_OUTPUT}\")\nprint(final_df.head())","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-29T10:01:24.898052Z","iopub.execute_input":"2026-01-29T10:01:24.898297Z","iopub.status.idle":"2026-01-29T10:08:08.096372Z","shell.execute_reply.started":"2026-01-29T10:01:24.898268Z","shell.execute_reply":"2026-01-29T10:08:08.095444Z"}},"outputs":[],"execution_count":null}]}