{"metadata":{"kaggle":{"accelerator":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":13482672,"sourceType":"datasetVersion","datasetId":8559888}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false},"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"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"1fe914b1d268858e","cell_type":"code","source":"# # Comprehensive Protein Visualization (Fixed)\n# ## ESM Embeddings, BioBERT Semantic Clustering & Physicochemical Properties\n# \n# **Pipeline:**\n# 1. **Data Loading**: Load Sequences, ESM Embeddings, and GO Annotations.\n# 2. **Physicochemical Encoding**: Calculate 14 features + 1 General PCA Score.\n# 3. **BioBERT Semantic Encoding**: Embed GO Terms -> Cluster Proteins.\n# 4. **Dimensionality Reduction**: UMAP on ESM embeddings.\n# 5. **Visualization**: Generate plots.\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:00.34543Z","iopub.execute_input":"2025-11-22T12:38:00.346978Z","iopub.status.idle":"2025-11-22T12:38:00.353878Z","shell.execute_reply.started":"2025-11-22T12:38:00.346899Z","shell.execute_reply":"2025-11-22T12:38:00.352426Z"}},"outputs":[],"execution_count":null},{"id":"af521458c53deaf7","cell_type":"code","source":"# 1. Install Dependencies\n!pip install -q obonet biopython umap-learn pandas matplotlib seaborn scikit-learn networkx transformers torch\n\nimport os\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport umap\nimport torch\nimport obonet\nimport networkx as nx\nfrom tqdm import tqdm\nfrom Bio import SeqIO\nfrom sklearn.decomposition import PCA\nfrom sklearn.cluster import KMeans\nfrom transformers import AutoTokenizer, AutoModel\n\n# Configuration\nCONFIG = {\n    'SEQ_PATH': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta',\n    'TERMS_PATH': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv',\n    'OBO_PATH': '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo',\n    'ESM_EMB_PATH': '/kaggle/input/cafa6-protein-embeddings-esm2/protein_embeddings.npy',\n    'ESM_IDS_PATH': '/kaggle/input/cafa6-protein-embeddings-esm2/protein_ids.csv',\n    'BIOBERT_MODEL': 'dmis-lab/biobert-v1.1',\n    'SAMPLE_SIZE': 5000,\n    'SEED': 999,\n    'DEVICE': 'cuda' if torch.cuda.is_available() else 'cpu'\n}\n\nnp.random.seed(CONFIG['SEED'])\ntorch.manual_seed(CONFIG['SEED'])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:00.355829Z","iopub.execute_input":"2025-11-22T12:38:00.356321Z","iopub.status.idle":"2025-11-22T12:38:05.485017Z","shell.execute_reply.started":"2025-11-22T12:38:00.356294Z","shell.execute_reply":"2025-11-22T12:38:05.483007Z"}},"outputs":[],"execution_count":null},{"id":"cb422c7df727017d","cell_type":"code","source":"# ## 1. Load Data & Align IDs\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:05.486627Z","iopub.execute_input":"2025-11-22T12:38:05.488071Z","iopub.status.idle":"2025-11-22T12:38:05.496713Z","shell.execute_reply.started":"2025-11-22T12:38:05.487993Z","shell.execute_reply":"2025-11-22T12:38:05.494332Z"}},"outputs":[],"execution_count":null},{"id":"abdad5e356f1784e","cell_type":"code","source":"print(\"Loading Data...\")\n\n# Load ESM Embeddings\nesm_emb_full = np.load(CONFIG['ESM_EMB_PATH'])\nesm_ids = pd.read_csv(CONFIG['ESM_IDS_PATH'])\nesm_ids['embedding_idx'] = range(len(esm_ids))\n\n# Load Sequences and clean IDs\nsequences = []\nfor record in SeqIO.parse(CONFIG['SEQ_PATH'], \"fasta\"):\n    # Fix ID format (TaxID|ProteinID -> ProteinID)\n    clean_id = record.id.split(\"|\")[1] if \"|\" in record.id else record.id\n    sequences.append({'protein_id': clean_id, 'sequence': str(record.seq)})\n\ndf_seq = pd.DataFrame(sequences)\n\n# Merge\ndf = esm_ids.merge(df_seq, on='protein_id', how='inner')\n\n# Load Terms to ensure we only keep annotated proteins for clustering\ndf_terms = pd.read_csv(CONFIG['TERMS_PATH'], sep=\"\\t\", names=[\"protein_id\", \"go_id\", \"aspect\"])\nannotated_proteins = set(df_terms['protein_id'].unique())\ndf = df[df['protein_id'].isin(annotated_proteins)]\n\n# Sample\nif len(df) > CONFIG['SAMPLE_SIZE']:\n    df = df.sample(n=CONFIG['SAMPLE_SIZE'], random_state=CONFIG['SEED']).reset_index(drop=True)\n\n# Get Arrays\nX_esm = esm_emb_full[df['embedding_idx'].values]\n\nprint(f\"✅ Analyzed Dataset: {len(df)} proteins\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:05.498845Z","iopub.execute_input":"2025-11-22T12:38:05.50051Z","iopub.status.idle":"2025-11-22T12:38:09.135917Z","shell.execute_reply.started":"2025-11-22T12:38:05.500454Z","shell.execute_reply":"2025-11-22T12:38:09.134639Z"}},"outputs":[],"execution_count":null},{"id":"5b127c64e135d558","cell_type":"code","source":"# ## 2. Physicochemical Properties\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:09.138691Z","iopub.execute_input":"2025-11-22T12:38:09.139013Z","iopub.status.idle":"2025-11-22T12:38:09.143708Z","shell.execute_reply.started":"2025-11-22T12:38:09.138991Z","shell.execute_reply":"2025-11-22T12:38:09.142573Z"}},"outputs":[],"execution_count":null},{"id":"4b8b2f44bef6db40","cell_type":"code","source":"Physicochemical_Properties = {\n    \"A\": [ 0.2369,  0.0616,  0.5125,  1.0000,  0.5043,  0.1085,  0.1652,  0.0000,  0.0000,  0.0000,  0.2490,  0.3951,  0.4043,  0.4819],\n    \"C\": [ 0.0000,  0.5517,  0.4961,  0.2996,  0.4866,  0.3566,  0.3304,  0.0000,  0.0000,  0.0000,  0.2048,  0.0741,  0.2879,  0.2229],\n    \"D\": [ 0.7539,  0.5065,  0.1405,  0.5054,  0.0000,  0.4496,  0.3739,  0.0000,  0.0000,  1.0000,  0.8675,  1.0000,  0.0000,  0.8042],\n    \"E\": [ 0.8507,  0.0197,  0.7938,  0.6060,  0.3918,  0.5581,  0.5304,  0.0000,  0.0000,  0.6667,  0.8112,  0.9136,  0.0563,  0.8283],\n    \"F\": [ 0.1062,  0.2591,  0.8465,  0.4681,  0.5952,  0.6977,  0.7565,  1.0000,  0.0000,  0.0000,  0.0763,  0.0370,  0.3392,  0.0000],\n    \"G\": [ 0.3021,  0.8810,  0.7751,  0.8580,  0.8643,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  1.0000,  0.5062,  0.4005,  0.6687],\n    \"H\": [ 0.5290,  0.3071,  0.3929,  0.1769,  0.5154,  0.6202,  0.6087,  0.0000,  0.3333,  0.3333,  0.1124,  0.6790,  0.6033,  0.5482],\n    \"I\": [ 0.0328,  0.2710,  0.8771,  0.6817,  0.6610,  0.4341,  0.6609,  0.0000,  0.0000,  0.0000,  0.6707,  0.0370,  0.4068,  0.0964],\n    \"K\": [ 1.0000,  0.2671,  0.6737,  0.5005,  0.7965,  0.5504,  0.7565,  0.0000,  0.6667,  0.0000,  0.6867,  0.7901,  0.8723,  0.9006],\n    \"L\": [ 0.1021,  0.1490,  0.4143,  0.9178,  0.3795,  0.4341,  0.6609,  0.0000,  0.0000,  0.0000,  0.2811,  0.0000,  0.4018,  0.1265],\n    \"M\": [ 0.2142,  0.0000,  0.8883,  0.3037,  0.7255,  0.5736,  0.6609,  0.0000,  0.0000,  0.0000,  0.0000,  0.0988,  0.3717,  0.2470],\n    \"N\": [ 0.7209,  0.6524,  0.7712,  0.5297,  0.6801,  0.4419,  0.4174,  0.0000,  0.3333,  0.3333,  0.6747,  0.8272,  0.3304,  0.9247],\n    \"P\": [ 0.4827,  1.0000,  0.3986,  0.6893,  0.3014,  0.3101,  0.3652,  0.0000,  0.0000,  0.0000,  0.8594,  0.3827,  0.4418,  0.6054],\n    \"Q\": [ 0.7164,  0.3731,  0.2234,  0.4394,  0.2263,  0.5504,  0.5739,  0.0000,  0.3333,  0.3333,  0.7952,  0.6914,  0.3605,  0.9006],\n    \"R\": [ 0.9077,  0.4075,  0.7970,  0.6944,  1.0000,  0.7674,  0.8696,  0.0000,  1.0000,  0.0000,  0.9398,  0.6914,  1.0000,  1.0000],\n    \"S\": [ 0.3513,  0.8108,  0.0000,  0.7566,  0.0969,  0.2326,  0.2174,  0.0000,  0.3333,  0.3333,  0.8514,  0.5309,  0.3642,  0.6325],\n    \"T\": [ 0.4130,  0.5132,  0.8875,  0.8210,  0.7420,  0.3411,  0.3913,  0.0000,  0.3333,  0.3333,  0.5984,  0.4568,  0.3542,  0.5723],\n    \"V\": [ 0.0019,  0.3454,  0.5366,  0.9113,  0.3225,  0.3256,  0.4957,  0.0000,  0.0000,  0.0000,  0.3655,  0.1235,  0.3992,  0.2259],\n    \"W\": [ 0.2357,  0.4252,  0.6914,  0.0000,  0.4981,  1.0000,  1.0000,  1.0000,  0.3333,  0.0000,  0.0402,  0.0617,  0.3905,  0.0392],\n    \"Y\": [ 0.5050,  0.6530,  1.0000,  0.3488,  0.7744,  0.8217,  0.8087,  1.0000,  0.3333,  0.3333,  0.5020,  0.1605,  0.3617,  0.2500],\n    \"B\": [ 0.7376,  0.5795,  0.4558,  0.5176,  0.3400,  0.4457,  0.3957,  0.0000,  0.1667,  0.6667,  0.7711,  0.9136,  0.1652,  0.8645],\n    \"Z\": [ 0.7836,  0.1964,  0.5086,  0.5227,  0.3090,  0.5543,  0.5522,  0.0000,  0.1667,  0.5000,  0.8032,  0.8025,  0.2090,  0.8645],\n    \"X\": [ 0.4231,  0.4227,  0.6058,  0.5754,  0.5281,  0.4651,  0.4957,  0.0000,  0.1667,  0.1667,  0.4217,  0.3827,  0.4043,  0.5181],\n}\n\ndef sequence_to_physio(seq):\n    vecs = [Physicochemical_Properties.get(aa, Physicochemical_Properties[\"X\"]) for aa in seq]\n    if not vecs: return Physicochemical_Properties[\"X\"]\n    return np.mean(vecs, axis=0)\n\nprint(\"Calculating Physicochemical Features...\")\nphysio_vecs = np.array([sequence_to_physio(s) for s in df['sequence']])\n\nfor i in range(14):\n    df[f'Physio_Dim_{i+1}'] = physio_vecs[:, i]\n\npca = PCA(n_components=1)\ndf['General_Physio_Score'] = pca.fit_transform(physio_vecs)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:09.145042Z","iopub.execute_input":"2025-11-22T12:38:09.145374Z","iopub.status.idle":"2025-11-22T12:38:09.225316Z","shell.execute_reply.started":"2025-11-22T12:38:09.145347Z","shell.execute_reply":"2025-11-22T12:38:09.224073Z"}},"outputs":[],"execution_count":null},{"id":"f070a8bafef018c9","cell_type":"code","source":"# ## 3. BioBERT Semantic Encoding & Clustering\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:09.226518Z","iopub.execute_input":"2025-11-22T12:38:09.226898Z","iopub.status.idle":"2025-11-22T12:38:09.231714Z","shell.execute_reply.started":"2025-11-22T12:38:09.226864Z","shell.execute_reply":"2025-11-22T12:38:09.230679Z"}},"outputs":[],"execution_count":null},{"id":"e353bf3653fdd499","cell_type":"code","source":"print(\"Parsing OBO & Loading BioBERT...\")\ngraph = obonet.read_obo(CONFIG['OBO_PATH'])\n\ndef get_def(data):\n    raw = data.get('def', '')\n    if '\"' in raw:\n        return raw.split('\"')[1]\n    return raw\n\ngo_data = {\n    id_: f\"{data.get('name', '')}. {get_def(data)}\" \n    for id_, data in graph.nodes(data=True)\n}\n\ntokenizer = AutoTokenizer.from_pretrained(CONFIG['BIOBERT_MODEL'])\nmodel = AutoModel.from_pretrained(CONFIG['BIOBERT_MODEL']).to(CONFIG['DEVICE'])\nmodel.eval()\n\ndef embed_texts(texts):\n    inputs = tokenizer(texts, padding=True, truncation=True, max_length=64, return_tensors=\"pt\").to(CONFIG['DEVICE'])\n    with torch.no_grad():\n        outputs = model(**inputs)\n        embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()\n    return embeddings\n\ncurrent_terms = df_terms[df_terms['protein_id'].isin(df['protein_id'])]['go_id'].unique()\nprint(f\"Embedding {len(current_terms)} unique GO terms...\")\n\nterm_embedding_map = {}\nbatch_size = 32\nunique_terms_list = [t for t in current_terms if t in go_data]\n\nfor i in tqdm(range(0, len(unique_terms_list), batch_size)):\n    batch_terms = unique_terms_list[i:i+batch_size]\n    batch_texts = [go_data[t] for t in batch_terms]\n    batch_embs = embed_texts(batch_texts)\n    for t, emb in zip(batch_terms, batch_embs):\n        term_embedding_map[t] = emb\n\n# Aggregate for Proteins\nprint(\"Creating Protein Semantic Vectors...\")\nprotein_semantic_vectors = []\nprot_groups = df_terms[df_terms['protein_id'].isin(df['protein_id'])].groupby('protein_id')['go_id'].apply(list)\n\nfor idx, pid in enumerate(df['protein_id']):\n    if pid in prot_groups:\n        terms = prot_groups[pid]\n        embs = [term_embedding_map[t] for t in terms if t in term_embedding_map]\n        if embs:\n            protein_semantic_vectors.append(np.mean(embs, axis=0))\n        else:\n            protein_semantic_vectors.append(np.zeros(768))\n    else:\n        protein_semantic_vectors.append(np.zeros(768))\n\nX_semantic = np.array(protein_semantic_vectors)\n\n# --- FIXED: Clustering Interpretation ---\nprint(\"Clustering via BioBERT Semantics...\")\n\ndef get_cluster_name(labels, n_clusters):\n    \"\"\"Interprets clusters by finding the dominant GO term\"\"\"\n    names = {}\n    for i in range(n_clusters):\n        cid = str(i)\n        # Compare strings to strings!\n        mask = (labels == cid)\n        \n        if not mask.any(): \n            names[cid] = f\"Cluster {cid} (Empty)\"\n            continue\n            \n        cluster_pids = df['protein_id'][mask]\n        cluster_terms = df_terms[df_terms['protein_id'].isin(cluster_pids)]['go_id']\n        \n        if cluster_terms.empty:\n            names[cid] = f\"Cluster {cid} (No Terms)\"\n        else:\n            top_term_id = cluster_terms.value_counts().index[0]\n            term_name = graph.nodes[top_term_id].get('name', top_term_id)\n            names[cid] = f\"{cid}: {term_name}\"\n            \n    return names\n\n# K=5\nkmeans_5 = KMeans(n_clusters=5, random_state=CONFIG['SEED'])\ndf['cluster_k5_id'] = kmeans_5.fit_predict(X_semantic).astype(str)\nnames_k5 = get_cluster_name(df['cluster_k5_id'], 5)\ndf['BioBERT Cluster (K=5)'] = df['cluster_k5_id'].map(names_k5)\n\n# K=20\nkmeans_20 = KMeans(n_clusters=20, random_state=CONFIG['SEED'])\ndf['cluster_k20_id'] = kmeans_20.fit_predict(X_semantic).astype(str)\nnames_k20 = get_cluster_name(df['cluster_k20_id'], 20)\ndf['BioBERT Cluster (K=20)'] = df['cluster_k20_id'].map(names_k20)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:09.233031Z","iopub.execute_input":"2025-11-22T12:38:09.233423Z","iopub.status.idle":"2025-11-22T12:38:36.928262Z","shell.execute_reply.started":"2025-11-22T12:38:09.233393Z","shell.execute_reply":"2025-11-22T12:38:36.927118Z"}},"outputs":[],"execution_count":null},{"id":"1fe673e6b36baed2","cell_type":"code","source":"# ## 4. UMAP Projections\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:36.929984Z","iopub.execute_input":"2025-11-22T12:38:36.930325Z","iopub.status.idle":"2025-11-22T12:38:36.935469Z","shell.execute_reply.started":"2025-11-22T12:38:36.930302Z","shell.execute_reply":"2025-11-22T12:38:36.934072Z"}},"outputs":[],"execution_count":null},{"id":"a8c8ba6293b61dfb","cell_type":"code","source":"print(\"Running UMAP on ESM (Protein Space)...\")\numap_esm = umap.UMAP(n_neighbors=20, min_dist=0.1, n_components=2, random_state=CONFIG['SEED'])\nemb_2d = umap_esm.fit_transform(X_esm)\ndf['umap_x'] = emb_2d[:, 0]\ndf['umap_y'] = emb_2d[:, 1]\n\nprint(\"Running UMAP on BioBERT (Term Space)...\")\nterm_ids = list(term_embedding_map.keys())\nterm_matrix = np.array([term_embedding_map[t] for t in term_ids])\numap_bert = umap.UMAP(n_neighbors=15, min_dist=0.1, n_components=2, random_state=CONFIG['SEED'])\nterm_2d = umap_bert.fit_transform(term_matrix)\n\ndf_terms_viz = pd.DataFrame({\n    'go_id': term_ids,\n    'umap_x': term_2d[:, 0],\n    'umap_y': term_2d[:, 1],\n    'name': [graph.nodes[t].get('name', 'Unknown') for t in term_ids],\n    'namespace': [graph.nodes[t].get('namespace', 'Unknown') for t in term_ids]\n})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:36.936508Z","iopub.execute_input":"2025-11-22T12:38:36.936954Z","iopub.status.idle":"2025-11-22T12:38:37.476862Z","shell.execute_reply.started":"2025-11-22T12:38:36.93692Z","shell.execute_reply":"2025-11-22T12:38:37.475382Z"}},"outputs":[],"execution_count":null},{"id":"afcafa88ddb2b61e","cell_type":"code","source":"# ## 5. Visualizations\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:37.477907Z","iopub.execute_input":"2025-11-22T12:38:37.478222Z","iopub.status.idle":"2025-11-22T12:38:37.482918Z","shell.execute_reply.started":"2025-11-22T12:38:37.478199Z","shell.execute_reply":"2025-11-22T12:38:37.481919Z"}},"outputs":[],"execution_count":null},{"id":"ffe66c241af7e517","cell_type":"code","source":"def plot_umap(data, x, y, hue, title, is_categorical=True, palette='tab10', cmap='viridis'):\n    plt.figure(figsize=(16, 10))\n    \n    if is_categorical:\n        ax = sns.scatterplot(\n            data=data, x=x, y=y, hue=hue, \n            palette=palette, s=15, alpha=0.7, linewidth=0\n        )\n        # Check if legend exists before moving it\n        if ax.get_legend() is not None:\n            sns.move_legend(ax, \"upper left\", bbox_to_anchor=(1, 1))\n    else:\n        sc = plt.scatter(\n            data[x], data[y], c=data[hue], \n            cmap=cmap, s=15, alpha=0.7\n        )\n        cbar = plt.colorbar(sc)\n        cbar.set_label(hue, rotation=270, labelpad=20)\n        \n    plt.title(title, fontsize=16, weight='bold')\n    plt.axis('off')\n    plt.tight_layout()\n    plt.show()\n\n# 1. UMAP of GO TERMS K=5 (BioBERT Clusters)\nplot_umap(df, 'umap_x', 'umap_y', 'BioBERT Cluster (K=5)', \n          \"Protein Space (ESM) colored by BioBERT Semantic Cluster (K=5)\", palette='Set1')\n\n# 2. UMAP of GO TERMS K=20 (BioBERT Clusters)\nplot_umap(df, 'umap_x', 'umap_y', 'BioBERT Cluster (K=20)', \n          \"Protein Space (ESM) colored by BioBERT Semantic Cluster (K=20)\", palette='tab20')\n\n# 3. UMAP of General Physicochemical Properties\nplot_umap(df, 'umap_x', 'umap_y', 'General_Physio_Score', \n          \"Protein Space (ESM) colored by General Physio Score (PCA)\", is_categorical=False, cmap='magma')\n\n# 4. UMAPs of each Physicochemical Encoding column\nfor i in range(1, 15):\n    col = f'Physio_Dim_{i}'\n    plot_umap(df, 'umap_x', 'umap_y', col, \n              f\"Protein Space (ESM) colored by Physio Property Dim {i}\", is_categorical=False, cmap='viridis')\n\n# 5. UMAP of BioBERT encoded GO Terms (The Term Space)\nplot_umap(df_terms_viz, 'umap_x', 'umap_y', 'namespace', \n          \"GO Term Space (BioBERT) colored by Namespace\", palette='Dark2')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-22T12:38:37.484091Z","iopub.execute_input":"2025-11-22T12:38:37.484524Z","iopub.status.idle":"2025-11-22T12:38:43.812391Z","shell.execute_reply.started":"2025-11-22T12:38:37.484491Z","shell.execute_reply":"2025-11-22T12:38:43.811254Z"}},"outputs":[],"execution_count":null},{"id":"9a9e2f755cc83ee6","cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}