{"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}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This approach leverages the KNN algorithms:\n- Step 1: Search for K nearest neighbours\n- Step 2: Use the labels from these neighbours","metadata":{}},{"cell_type":"code","source":"%pip install biopython","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:37:59.489246Z","iopub.execute_input":"2026-01-02T17:37:59.489608Z","iopub.status.idle":"2026-01-02T17:38:07.406731Z","shell.execute_reply.started":"2026-01-02T17:37:59.489579Z","shell.execute_reply":"2026-01-02T17:38:07.405463Z"}},"outputs":[],"execution_count":null},{"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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:07.409076Z","iopub.execute_input":"2026-01-02T17:38:07.409472Z","iopub.status.idle":"2026-01-02T17:38:07.813214Z","shell.execute_reply.started":"2026-01-02T17:38:07.409433Z","shell.execute_reply":"2026-01-02T17:38:07.811892Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nfrom Bio import SeqIO\nimport unicodedata\nimport re\nfrom joblib import Parallel, delayed\nfrom time import time\nfrom collections import defaultdict\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.neighbors import NearestNeighbors\nimport gc","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:07.814294Z","iopub.execute_input":"2026-01-02T17:38:07.814937Z","iopub.status.idle":"2026-01-02T17:38:09.192071Z","shell.execute_reply.started":"2026-01-02T17:38:07.814905Z","shell.execute_reply":"2026-01-02T17:38:09.190823Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Data Preparation","metadata":{}},{"cell_type":"code","source":"# load the train sequences\ntrain_sequence_file = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\ntotal = sum(1 for line in open(train_sequence_file, \"r\", encoding=\"utf-8\") if line.startswith(\">\"))\nprint(f\"Total number of train records: {total}\")\n\ntrain_sequence_df = []\nfor record in tqdm(SeqIO.parse(train_sequence_file, \"fasta\"), total=total, desc=\"Parsing FASTA\", unit=\"seq\"):\n    r = [\n        record.id, \n        record.id.split('|')[1],\n        record.description, \n        str(record.seq)\n    ]\n    train_sequence_df.append(r)\n    \n# convert to pandas data frame\ntrain_sequence_df = pd.DataFrame(train_sequence_df, columns=['full_prot_id', 'prot_id', 'description', 'sequence'])\ntrain_sequence_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:09.194803Z","iopub.execute_input":"2026-01-02T17:38:09.195259Z","iopub.status.idle":"2026-01-02T17:38:11.050396Z","shell.execute_reply.started":"2026-01-02T17:38:09.195215Z","shell.execute_reply":"2026-01-02T17:38:11.049394Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load train-terms\ntrain_term_df = pd.read_csv(\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n    sep=\"\\t\"\n)\ntrain_term_df = train_term_df.rename(columns={'EntryID': 'prot_id'})\ntrain_term_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:11.051537Z","iopub.execute_input":"2026-01-02T17:38:11.051815Z","iopub.status.idle":"2026-01-02T17:38:11.439984Z","shell.execute_reply.started":"2026-01-02T17:38:11.051788Z","shell.execute_reply":"2026-01-02T17:38:11.438993Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# merge to get labels\ntrain_sequence_df = train_sequence_df.merge(\n    train_term_df,\n    on='prot_id'\n)\ntrain_sequence_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:11.441113Z","iopub.execute_input":"2026-01-02T17:38:11.441411Z","iopub.status.idle":"2026-01-02T17:38:11.715699Z","shell.execute_reply.started":"2026-01-02T17:38:11.441385Z","shell.execute_reply":"2026-01-02T17:38:11.714415Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get the train series\n# note: two proteins can have the same sequence\ntrain = train_sequence_df.drop_duplicates(subset=['prot_id'])\ntrain.index = train.prot_id\ntrain = train['sequence']\nprint(f\"No. proteins: {train.index.nunique()}\")\nprint(f\"No. sequences: {train.drop_duplicates().shape[0]}\")\ntrain","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:11.716814Z","iopub.execute_input":"2026-01-02T17:38:11.717175Z","iopub.status.idle":"2026-01-02T17:38:11.827464Z","shell.execute_reply.started":"2026-01-02T17:38:11.717138Z","shell.execute_reply":"2026-01-02T17:38:11.826221Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load test data\ntest_superset_file = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\n# Count sequences quickly (without parsing)\ntotal = sum(1 for line in open(test_superset_file, \"r\", encoding=\"utf-8\") if line.startswith(\">\"))\nprint(f\"Total number of records: {total}\")\n\n# Load all sequences from a FASTA file\ntest_df = []\nfor record in tqdm(SeqIO.parse(test_superset_file, \"fasta\"), total=total, desc=\"Parsing FASTA\", unit=\"seq\"):\n    r = [record.id, record.description, str(record.seq)]\n    test_df.append(r)\n\ntest_df = pd.DataFrame(test_df, columns=['prot_id', 'description', 'sequence'])\nprint(f\"No. proteins: {test_df.prot_id.nunique()}\")\nprint(f\"No. sequences: {test_df.sequence.nunique()}\")\ntest_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:11.828507Z","iopub.execute_input":"2026-01-02T17:38:11.828789Z","iopub.status.idle":"2026-01-02T17:38:15.800730Z","shell.execute_reply.started":"2026-01-02T17:38:11.828763Z","shell.execute_reply":"2026-01-02T17:38:15.799661Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get test series\ntest = test_df.copy()\ntest.index = test.prot_id\ntest = test.sequence\ntest","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:15.801831Z","iopub.execute_input":"2026-01-02T17:38:15.802150Z","iopub.status.idle":"2026-01-02T17:38:15.859865Z","shell.execute_reply.started":"2026-01-02T17:38:15.802109Z","shell.execute_reply":"2026-01-02T17:38:15.858749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # only for testing\n# test = test.iloc[:1000]\n# test","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:47:55.120037Z","iopub.execute_input":"2026-01-02T17:47:55.120437Z","iopub.status.idle":"2026-01-02T17:47:55.125333Z","shell.execute_reply.started":"2026-01-02T17:47:55.120397Z","shell.execute_reply":"2026-01-02T17:47:55.124113Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get the vocab from train\ntrain_vocab = set().union(*train.apply(lambda x: set(x)).tolist())\nprint(f\"train vocab size: {len(train_vocab)}\")\ntrain_vocab","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:15.872247Z","iopub.execute_input":"2026-01-02T17:38:15.872644Z","iopub.status.idle":"2026-01-02T17:38:16.685788Z","shell.execute_reply.started":"2026-01-02T17:38:15.872604Z","shell.execute_reply":"2026-01-02T17:38:16.684934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get the vocab from test\ntest_vocab = set().union(*test.apply(lambda x: set(x)).tolist())\nprint(f\"test vocab size: {len(test_vocab)}\")\ntest_vocab","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:16.686935Z","iopub.execute_input":"2026-01-02T17:38:16.687903Z","iopub.status.idle":"2026-01-02T17:38:16.704884Z","shell.execute_reply.started":"2026-01-02T17:38:16.687864Z","shell.execute_reply":"2026-01-02T17:38:16.703363Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Test Amino acid not in Train: {test_vocab.difference(train_vocab)}\")\nprint(f\"Train Amino acid not in Test: {train_vocab.difference(test_vocab)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:16.706967Z","iopub.execute_input":"2026-01-02T17:38:16.707353Z","iopub.status.idle":"2026-01-02T17:38:16.727395Z","shell.execute_reply.started":"2026-01-02T17:38:16.707296Z","shell.execute_reply":"2026-01-02T17:38:16.726328Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. KNN","metadata":{}},{"cell_type":"code","source":"# BoW embeddings\nt_start = time()\n\nvectorizer = TfidfVectorizer(\n    analyzer='char',\n    ngram_range=(1,2)\n)\nX_train = vectorizer.fit_transform(train)\nX_test = vectorizer.transform(test)\nprint(f\"X_train shape: {X_train.shape}\")\nprint(f\"X_test shape: {X_test.shape}\")\n\n\nprint(f\"Finished encoding in {time() - t_start} seconds\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:16.728774Z","iopub.execute_input":"2026-01-02T17:38:16.729094Z","iopub.status.idle":"2026-01-02T17:38:42.337937Z","shell.execute_reply.started":"2026-01-02T17:38:16.729049Z","shell.execute_reply":"2026-01-02T17:38:42.336926Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# KNN\nprint(f\"fitting KNN...\")\nt_start = time()\nnn = NearestNeighbors(\n    n_neighbors=10,\n    metric='cosine',       # change to 'euclidean' if you prefer\n    algorithm='brute',     # brute is reliable for sparse + cosine\n    n_jobs=-1              # use all cores if available\n)\n\n# Fit on the training set\nnn.fit(X_train)\nprint(f\"KNN fitted in {time() - t_start}\")\n\nprint(f\"querrying...\")\nt_start = time()\n\n# Chunked inference\nn_test = X_test.shape[0]\nk = 10\nbatch_size = 2000  # tune based on your RAM/CPU\nall_indices = np.empty((n_test, k), dtype=np.int32)\nall_distances = np.empty((n_test, k), dtype=np.float32)\n\nfor start in tqdm(range(0, n_test, batch_size)):\n    end = min(start + batch_size, n_test)\n    X_batch = X_test[start:end]\n    distances, indices = nn.kneighbors(X_batch, n_neighbors=k, return_distance=True)\n    all_indices[start:end] = indices\n    all_distances[start:end] = distances\n\n# Optionally convert to similarities\nall_similarities = 1.0 - all_distances\n\nprint(f\"finished searching in {time() - t_start} seconds\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:38:42.339093Z","iopub.execute_input":"2026-01-02T17:38:42.339508Z","iopub.status.idle":"2026-01-02T17:39:05.510418Z","shell.execute_reply.started":"2026-01-02T17:38:42.339460Z","shell.execute_reply":"2026-01-02T17:39:05.509538Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# convert result to long-form data frame\ntest_ids = np.repeat(np.arange(n_test, dtype=np.int32), k)\ntrain_ids = all_indices.ravel().astype(np.int32)\nsimilarity_scores = all_similarities.ravel().astype(np.float32)\n\nknn_df = pd.DataFrame({\n    \"test_id\": test_ids,\n    \"train_id\": train_ids,\n    \"similarity_score\": similarity_scores\n})\n\ntest_ind2id = dict(zip(range(test.shape[0]), test.index))\ntrain_ind2id = dict(zip(range(train.shape[0]), train.index))\n\nknn_df['test_prot'] = knn_df.test_id.map(test_ind2id)\nknn_df['train_prot'] = knn_df.train_id.map(train_ind2id)\nknn_df = knn_df.merge(\n    train_sequence_df[['prot_id', 'term']],\n    left_on='train_prot',\n    right_on='prot_id',\n    how='left'\n)\nknn_df = knn_df[['test_prot', 'term', 'similarity_score']]\nknn_df = knn_df\\\n    .sort_values(\n        by=['test_prot', 'term', 'similarity_score'], \n        ascending=[True, True, False])\\\n    .drop_duplicates(\n        subset=['test_prot', 'term'],\n        keep='first'\n    ).reset_index(drop=True)\n\n# use average similarity score as the probability\n# knn_df = knn_df.groupby(['test_prot', 'term'], as_index=False)['similarity_score'].mean()\n\nknn_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:39:05.511732Z","iopub.execute_input":"2026-01-02T17:39:05.512070Z","iopub.status.idle":"2026-01-02T17:39:05.711093Z","shell.execute_reply.started":"2026-01-02T17:39:05.512026Z","shell.execute_reply":"2026-01-02T17:39:05.710029Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# remove negative relationships\n# ==========================================\n# 1. OPTIMIZED OBO PARSER (Pre-compute Descendants)\n# ==========================================\ndef parse_obo_and_get_descendants(go_obo_path):\n    \"\"\"\n    Parses OBO and returns a dictionary mapping every GO term \n    to a set of ALL its descendants (transitive closure).\n    \"\"\"\n    print(f\"[1/5] Parsing OBO Ontology from {go_obo_path}...\")\n    children = defaultdict(set)\n    \n    if not os.path.exists(go_obo_path):\n        print(\"Warning: OBO file not found.\")\n        return {}\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                # Parent logic: if \"id\" is_a \"pid\", then \"id\" is a child of \"pid\"\n                pid = line.split()[1].strip()\n                if cur_id:\n                    children[pid].add(cur_id)\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                        children[pid].add(cur_id)\n\n    # Pre-compute transitive closure (all descendants) for every node\n    # Using a cache to avoid re-calculating common paths\n    print(f\"[1/5] Calculating transitive closure for {len(children)} parent nodes...\")\n    descendants_cache = {}\n\n    def get_all_descendants(term):\n        if term in descendants_cache:\n            return descendants_cache[term]\n        \n        direct_children = children.get(term, set())\n        all_desc = set(direct_children)\n        \n        for child in direct_children:\n            all_desc |= get_all_descendants(child)\n            \n        descendants_cache[term] = all_desc\n        return all_desc\n\n    # Force computation for all keys\n    # We iterate list(children.keys()) to avoid runtime dict change errors\n    for term in list(children.keys()):\n        get_all_descendants(term)\n        \n    print(f\"[1/5] OBO Parsed. Knowledge base ready.\")\n    return descendants_cache\n    \n# ==========================================\n# 2. LOAD & MERGE NEGATIVE ANNOTATIONS\n# ==========================================\ndef get_negative_pairs(file_paths, descendants_map):\n    \"\"\"\n    Loads multiple GOA files, filters for 'NOT' qualifier, \n    propagates negatives to descendants, and returns a set of unique keys.\n    \"\"\"\n    unique_negatives = set()\n    \n    for path in file_paths:\n        if not os.path.exists(path):\n            continue\n            \n        print(f\"[2/5] Processing annotations: {os.path.basename(path)}\")\n        # Determine separator based on file extension\n        sep = '\\t' if path.endswith('.tsv') else ','\n        \n        # specific loading for efficiency\n        df = pd.read_csv(path, sep=sep, dtype=str, usecols=['protein_id', 'go_term', 'qualifier'])\n        \n        # Filter for Negative qualifiers (NOT)\n        neg_df = df[df['qualifier'].str.contains('NOT', na=False)]\n        print(f\"      Found {len(neg_df)} explicit negative annotations.\")\n        \n        # Group by Protein to speed up propagation\n        # protein -> [term1, term2, ...]\n        prot_to_terms = neg_df.groupby('protein_id')['go_term'].apply(list).to_dict()\n        \n        count = 0\n        for protein, terms in tqdm(prot_to_terms.items(), desc=\"Propagating Negatives\"):\n            for term in terms:\n                # 1. The term itself is negative\n                unique_negatives.add(f\"{protein}_{term}\")\n                \n                # 2. All descendants of the term are negative (Negative Propagation Rule)\n                # If protein is NOT involved in 'Metabolism', it is NOT involved in 'Gluconeogenesis'\n                if term in descendants_map:\n                    for desc in descendants_map[term]:\n                        unique_negatives.add(f\"{protein}_{desc}\")\n        \n        del df, neg_df, prot_to_terms\n        gc.collect()\n\n    return unique_negatives\n\n\n# --- Configuration ---\nOBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nGOA_FILES = [\n    '/kaggle/input/protein-go-annotations/goa_uniprot_ver227.csv',\n    '/kaggle/input/protein-go-annotations/goa_uniprot_ver228.tsv'\n]\n\n# Step 1: Get Ontology Structure\ndescendants_map = parse_obo_and_get_descendants(OBO_PATH)\n\n# Step 2: Build Global Negative Dictionary\nnegative_keys = get_negative_pairs(GOA_FILES, descendants_map)\nprint(f\"[+] Total unique Negative Protein-GO pairs (Propagated): {len(negative_keys)}\")\n\n# Free up memory (we only need the set `negative_keys` now)\ndel descendants_map\ngc.collect()\n\nknn_df['key'] = knn_df['test_prot'] + \"_\" + knn_df['term']\nknn_df = knn_df[~knn_df.key.isin(negative_keys)]\nknn_df = knn_df.drop(columns=['key'])\nknn_df.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:47:26.691124Z","iopub.execute_input":"2026-01-02T17:47:26.691603Z","iopub.status.idle":"2026-01-02T17:47:33.201661Z","shell.execute_reply.started":"2026-01-02T17:47:26.691573Z","shell.execute_reply":"2026-01-02T17:47:33.200557Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# save to file\nknn_df.to_csv('submission.tsv', index=False, sep='\\t', header=False)\nknn_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T17:47:33.203354Z","iopub.execute_input":"2026-01-02T17:47:33.203684Z","iopub.status.idle":"2026-01-02T17:47:33.269013Z","shell.execute_reply.started":"2026-01-02T17:47:33.203656Z","shell.execute_reply":"2026-01-02T17:47:33.267729Z"}},"outputs":[],"execution_count":null}]}