{"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":14187577,"sourceType":"datasetVersion","datasetId":9045748},{"sourceId":287175447,"sourceType":"kernelVersion"},{"sourceId":292870115,"sourceType":"kernelVersion"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# CAFA6 SOTA SOLUTION #17JAN (single block)\nimport os, gc\nimport pandas as pd\nimport numpy as np\nimport networkx as nx\nfrom tqdm.auto import tqdm\n\nprint(\"Starting...\")\n\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\"\nMIN_SCORE_THRESHOLD = 0.001\n\ndef build_go_graph(go_obo_path: str):\n    go_graph = nx.DiGraph()\n    if not os.path.exists(go_obo_path):\n        raise FileNotFoundError(f\"OBO not found: {go_obo_path}\")\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    if not nx.is_directed_acyclic_graph(go_graph):\n        for cycle in list(nx.simple_cycles(go_graph)):\n            go_graph.remove_edge(cycle[0], cycle[1])\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    return topo_order, term_parents_map\n\ndef propagate_hard_max(df: pd.DataFrame, topo_order, term_parents_map):\n    groups = df.groupby(\"protein_id\")\n    final_rows = []\n    for pid, group in tqdm(groups, desc=\"Propagating\"):\n        scores = dict(zip(group[\"go_term\"], group[\"score\"]))\n        for child in topo_order:\n            if child in scores:\n                child_score = scores[child]\n                parents = term_parents_map.get(child, ())\n                for parent in parents:\n                    if child_score > scores.get(parent, 0.0):\n                        scores[parent] = child_score\n        for term, score in scores.items():\n            if score >= MIN_SCORE_THRESHOLD:\n                final_rows.append((pid, term, score))\n    return pd.DataFrame(final_rows, columns=[\"protein_id\", \"go_term\", \"score\"])\n\ntopo_order, term_parents_map = build_go_graph(OBO_PATH)\n\nsubmission = pd.read_csv(\n    SUBMISSION_INPUT,\n    sep=\"\\t\",\n    header=None,\n    names=[\"protein_id\", \"go_term\", \"score\"],\n    dtype={\"score\": np.float32},\n    on_bad_lines=\"skip\",\n)\n\nsubmission_improved = propagate_hard_max(submission, topo_order, term_parents_map)\nsubmission_improved[\"score\"] = submission_improved[\"score\"].round(6)\nsubmission_improved.to_csv(SUBMISSION_OUTPUT, sep=\"\\t\", index=False, header=False)\n\ngc.collect()\nprint(\"Completed. Submit now!!\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}