{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","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":14571222,"sourceType":"datasetVersion","datasetId":9307683,"isSourceIdPinned":true}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pandas as pd\nimport networkx as nx\nfrom tqdm.auto import tqdm\n\nobo_file_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nsubmission_files = [\n    \"/kaggle/input/cafa-dataset/submission_fold1.tsv\",\n    \"/kaggle/input/cafa-dataset/submission_fold2.tsv\",\n    \"/kaggle/input/cafa-dataset/submission_fold3.tsv\",\n]\nfinal_output_file = \"submission.tsv\"\n\ndef create_ontology_graph(obo_path):\n    if not os.path.exists(obo_path):\n        raise FileNotFoundError(f\"File not found: {obo_path}\")\n\n    ontology_graph = nx.DiGraph()\n    current_id = None\n\n    with open(obo_path, \"r\") as file:\n        for line in file:\n            line = line.strip()\n            if line == \"[Term]\":\n                current_id = None\n            elif line.startswith(\"id: \"):\n                current_id = line.split(\"id: \", 1)[1].strip()\n                ontology_graph.add_node(current_id)\n            elif line.startswith(\"is_a: \"):\n                parent_id = line.split()[1].strip()\n                if current_id:\n                    ontology_graph.add_edge(current_id, parent_id)\n            elif line.startswith(\"relationship: \"):\n                parts = line.split()\n                if len(parts) >= 3 and parts[1] in {\"part_of\", \"regulates\", \"positively_regulates\", \"negatively_regulates\"}:\n                    parent_id = parts[2].strip()\n                    if current_id:\n                        ontology_graph.add_edge(current_id, parent_id)\n                        \n    if not nx.is_directed_acyclic_graph(ontology_graph):\n        for cycle in nx.simple_cycles(ontology_graph):\n            ontology_graph.remove_edge(cycle[0], cycle[1])\n\n    topological_order = list(nx.topological_sort(ontology_graph))\n    child_to_parents = {node: list(ontology_graph.successors(node)) for node in ontology_graph.nodes()}\n    return topological_order, child_to_parents, set(ontology_graph.nodes())\n\ndef propagate_scores(predictions_df, topological_order, child_to_parents, valid_go_terms):\n    mask = predictions_df['go_term'].isin(valid_go_terms)\n    filtered_df = predictions_df[mask].copy()\n\n    grouped_by_protein = filtered_df.groupby('protein_id')\n    results = []\n\n    for protein_id, protein_group in tqdm(grouped_by_protein, total=grouped_by_protein.ngroups):\n        term_scores = dict(zip(protein_group['go_term'], protein_group['score']))\n\n        relevant_terms = set(term_scores.keys()) & set(topological_order)\n        for child_term in topological_order:\n            if child_term not in relevant_terms:\n                continue\n            child_score = term_scores[child_term]\n            for parent_term in child_to_parents.get(child_term, []):\n                current_parent_score = term_scores.get(parent_term, 0.0)\n                if child_score > current_parent_score:\n                    term_scores[parent_term] = child_score\n                    relevant_terms.add(parent_term)\n                    \n        for term, score in term_scores.items():\n            clamped_score = max(0.0, min(1.0, score))\n            if clamped_score >= 0.0013:\n                results.append((protein_id, term, round(clamped_score, 3)))\n\n    return pd.DataFrame(results, columns=['protein_id', 'go_term', 'score'])\n\nall_predictions = []\n\nfor file_path in submission_files:\n    fold_predictions = pd.read_csv(\n        file_path, sep='\\t', header=None,\n        names=['protein_id', 'go_term', 'score']\n    )\n    all_predictions.append(fold_predictions)\n\ncombined_predictions = pd.concat(all_predictions, ignore_index=True)\naveraged_predictions = combined_predictions.groupby(['protein_id', 'go_term'])['score'].mean().reset_index()\n\nontology_order, parent_mapping, valid_go_set = create_ontology_graph(obo_file_path)\n\nfinal_predictions = propagate_scores(averaged_predictions, ontology_order, parent_mapping, valid_go_set)\nfinal_predictions.drop_duplicates(subset=['protein_id', 'go_term'], inplace=True)\n\nfinal_predictions.to_csv(final_output_file, sep='\\t', index=False, header=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}