{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "9c6f23e0", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[✓] All required Python packages already installed\n", "[*] fpocket not found — installing from source...\n", "/mnt/disk3/Arun_work/esmfold/fpocket\n", "cd src/qhull/ && make\n", "make[1]: Entering directory '/mnt/disk3/Arun_work/esmfold/fpocket/src/qhull'\n", "mkdir -p bin lib\n", "make[1]: Leaving directory '/mnt/disk3/Arun_work/esmfold/fpocket/src/qhull'\n", "/mnt/disk3/Arun_work/esmfold\n", "[✓] fpocket is working correctly\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/arunraj/miniconda3/envs/phosbind_env/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "[*] Using device: cpu\n", "[*] Loading ESMFold model (will download only first time)...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n", "Loading weights: 100%|██████████| 4533/4533 [00:00<00:00, 9524.03it/s]\n", "\u001b[1mEsmForProteinFolding LOAD REPORT\u001b[0m from: facebook/esmfold_v1\n", "Key | Status | \n", "-----------------------------------+------------+-\n", "esm.embeddings.position_ids | UNEXPECTED | \n", "esm.contact_head.regression.bias | MISSING | \n", "esm.contact_head.regression.weight | MISSING | \n", "\n", "Notes:\n", "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n", "- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "[✓] ESMFold ready\n" ] } ], "source": [ "# =========================================\n", "# 📦 CHECK & INSTALL PACKAGES AND FUNCTIONS\n", "# =========================================\n", "\n", "import importlib\n", "\n", "def is_installed(pkg):\n", " return importlib.util.find_spec(pkg) is not None\n", "\n", "required_packages = {\n", " \"transformers\": \"transformers\",\n", " \"accelerate\": \"accelerate\",\n", " \"Bio\": \"biopython\",\n", " \"py3Dmol\": \"py3Dmol\"\n", "}\n", "\n", "missing = [pip_name for mod, pip_name in required_packages.items() if not is_installed(mod)]\n", "\n", "if missing:\n", " print(f\"[*] Installing missing packages: {missing}\")\n", " !pip install -q {' '.join(missing)}\n", "else:\n", " print(\"[✓] All required Python packages already installed\")\n", "\n", "# ================================\n", "# 🔧 CHECK & INSTALL FPOCKET\n", "# ================================\n", "\n", "import os\n", "import shutil\n", "from pathlib import Path\n", "import subprocess\n", "\n", "fpocket_path = shutil.which(\"fpocket\")\n", "\n", "if fpocket_path:\n", " print(f\"[✓] fpocket already available at: {fpocket_path}\")\n", "else:\n", " print(\"[*] fpocket not found — installing from source...\")\n", "\n", " if not Path(\"fpocket\").exists():\n", " !git clone --depth 1 https://github.com/Discngine/fpocket.git\n", "\n", " %cd fpocket\n", " !make\n", "\n", " # Add to PATH\n", " os.environ[\"PATH\"] += \":\" + os.getcwd() + \"/bin\"\n", "\n", " %cd ..\n", "\n", " fpocket_path = shutil.which(\"fpocket\")\n", "\n", "# Verify fpocket works\n", "if fpocket_path:\n", " result = subprocess.run(\n", " [\"fpocket\", \"-h\"],\n", " stdout=subprocess.PIPE,\n", " stderr=subprocess.PIPE,\n", " text=True\n", " )\n", " if result.returncode == 0:\n", " print(\"[✓] fpocket is working correctly\")\n", " else:\n", " print(\"[!] fpocket found but may not run correctly\")\n", "else:\n", " print(\"[✗] fpocket installation failed\")\n", "\n", "# ================================\n", "# 🧠 LOAD ESMFOLD MODEL\n", "# ================================\n", "\n", "import torch\n", "from transformers import AutoTokenizer, EsmForProteinFolding\n", "from accelerate.test_utils.testing import get_backend\n", "\n", "DEVICE, _, _ = get_backend()\n", "print(f\"[*] Using device: {DEVICE}\")\n", "\n", "print(\"[*] Loading ESMFold model (will download only first time)...\")\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esmfold_v1\")\n", "model = EsmForProteinFolding.from_pretrained(\n", " \"facebook/esmfold_v1\",\n", " low_cpu_mem_usage=True\n", ").to(DEVICE)\n", "\n", "# Memory optimization\n", "if DEVICE == \"cuda\":\n", " model.esm = model.esm.half()\n", " torch.backends.cuda.matmul.allow_tf32 = True\n", "\n", "model.trunk.set_chunk_size(64)\n", "\n", "print(\"[✓] ESMFold ready\")" ] }, { "cell_type": "code", "execution_count": 2, "id": "8462f364", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[+] Job name : 9SV1_1\n", "[+] Sequence length: 100 aa\n", "[+] Output folder : 9SV1_1_metadata/9SV1_1\n" ] } ], "source": [ "# ================================\n", "# 🔧 USER INPUT\n", "# ================================\n", "\n", "from pathlib import Path\n", "import re\n", "\n", "USER_SEQUENCE = \"MSKVCIIAWVYGRVQGVGFRYTTQYEAKRLGLTGYAKNLDDGSVEVVACGEEGQVEKLMQWLKSGGPRSARVERVLSEPHHPSGELTDFRIRLEHHHHHH\"\n", "JOB_NAME = \"9SV1_1\"\n", "\n", "\n", "BASE_OUTPUT_DIR = f\"./{JOB_NAME}_metadata\"\n", "USER_SEQUENCE = re.sub(r\"\\s+\", \"\", USER_SEQUENCE).upper()\n", "if not USER_SEQUENCE:\n", " raise ValueError(\"USER_SEQUENCE is empty.\")\n", "\n", "valid_aa = set(\"ACDEFGHIKLMNPQRSTVWY\")\n", "invalid_chars = sorted(set(USER_SEQUENCE) - valid_aa)\n", "\n", "if invalid_chars:\n", " raise ValueError(f\"Invalid amino acids found in sequence: {invalid_chars}\")\n", "\n", "JOB_NAME = JOB_NAME.strip()\n", "if not JOB_NAME:\n", " raise ValueError(\"JOB_NAME cannot be empty.\")\n", "SAFE_JOB_NAME = re.sub(r\"[^A-Za-z0-9_.-]\", \"_\", JOB_NAME)\n", "\n", "# Create output directory\n", "JOB_DIR = Path(BASE_OUTPUT_DIR) / SAFE_JOB_NAME\n", "JOB_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "# Summary\n", "print(f\"[+] Job name : {SAFE_JOB_NAME}\")\n", "print(f\"[+] Sequence length: {len(USER_SEQUENCE)} aa\")\n", "print(f\"[+] Output folder : {JOB_DIR}\")" ] }, { "cell_type": "code", "execution_count": 3, "id": "b809f34f", "metadata": {}, "outputs": [], "source": [ "# ================================\n", "# 🧰 HELPER FUNCTIONS\n", "# ================================\n", "\n", "import re\n", "import json\n", "import subprocess\n", "from pathlib import Path\n", "\n", "from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein\n", "from transformers.models.esm.openfold_utils.feats import atom14_to_atom37\n", "\n", "\n", "def convert_outputs_to_pdb(outputs):\n", " \"\"\"\n", " Convert ESMFold model outputs to PDB strings.\n", " Returns a list of PDB strings, one per input sequence.\n", " \"\"\"\n", " final_atom_positions = atom14_to_atom37(outputs[\"positions\"][-1], outputs)\n", "\n", " outputs_np = {\n", " k: v.detach().cpu().numpy() if torch.is_tensor(v) else v\n", " for k, v in outputs.items()\n", " }\n", "\n", " final_atom_positions = final_atom_positions.detach().cpu().numpy()\n", " final_atom_mask = outputs_np[\"atom37_atom_exists\"]\n", "\n", " pdbs = []\n", " for i in range(outputs_np[\"aatype\"].shape[0]):\n", " pred = OFProtein(\n", " aatype=outputs_np[\"aatype\"][i],\n", " atom_positions=final_atom_positions[i],\n", " atom_mask=final_atom_mask[i],\n", " residue_index=outputs_np[\"residue_index\"][i] + 1,\n", " b_factors=outputs_np[\"plddt\"][i],\n", " chain_index=outputs_np[\"chain_index\"][i] if \"chain_index\" in outputs_np else None,\n", " )\n", " pdbs.append(to_pdb(pred))\n", "\n", " return pdbs\n", "\n", "\n", "def save_pdb_files(pdb_strings, out_dir, prefix):\n", " \"\"\"\n", " Save PDB strings to files.\n", " Returns a list of saved file paths.\n", " \"\"\"\n", " out_dir = Path(out_dir)\n", " out_dir.mkdir(parents=True, exist_ok=True)\n", "\n", " saved_paths = []\n", " for i, pdb_str in enumerate(pdb_strings):\n", " pdb_path = out_dir / f\"{prefix}_{i}.pdb\"\n", " pdb_path.write_text(pdb_str)\n", " saved_paths.append(pdb_path)\n", " print(f\"[+] Saved PDB: {pdb_path}\")\n", "\n", " return saved_paths\n", "\n", "\n", "def run_fpocket(pdb_path):\n", " \"\"\"\n", " Run fpocket on a PDB file.\n", " Returns the fpocket output directory.\n", " \"\"\"\n", " pdb_path = Path(pdb_path).resolve()\n", "\n", " result = subprocess.run(\n", " [\"fpocket\", \"-f\", str(pdb_path)],\n", " stdout=subprocess.PIPE,\n", " stderr=subprocess.PIPE,\n", " text=True\n", " )\n", "\n", " if result.returncode != 0:\n", " print(result.stdout)\n", " print(result.stderr)\n", " raise RuntimeError(f\"fpocket failed for {pdb_path.name}\")\n", "\n", " out_dir = pdb_path.parent / f\"{pdb_path.stem}_out\"\n", " print(f\"[+] fpocket completed: {out_dir}\")\n", " return out_dir\n", "\n", "\n", "def parse_pocket_summary(out_dir):\n", " \"\"\"\n", " Parse fpocket *_info.txt summary file.\n", " Returns a list of pocket dictionaries.\n", " \"\"\"\n", " out_dir = Path(out_dir)\n", " info_files = list(out_dir.glob(\"*_info.txt\"))\n", "\n", " if not info_files:\n", " print(\"[!] No fpocket summary file found.\")\n", " return []\n", "\n", " info_file = info_files[0]\n", " pockets = []\n", " current = None\n", "\n", " with open(info_file, \"r\") as f:\n", " for line in f:\n", " line = line.rstrip()\n", "\n", " pocket_match = re.match(r\"Pocket\\s+(\\d+)\\s*:\", line)\n", " if pocket_match:\n", " if current is not None:\n", " pockets.append(current)\n", " current = {\"id\": int(pocket_match.group(1))}\n", " continue\n", "\n", " kv_match = re.match(r\"\\s+(.+?)\\s*:\\s*(.+)\", line)\n", " if kv_match and current is not None:\n", " key = kv_match.group(1).strip()\n", " value = kv_match.group(2).strip()\n", " current[key] = value\n", "\n", " if current is not None:\n", " pockets.append(current)\n", "\n", " return pockets\n", "\n", "\n", "def find_pocket_pdb_files(out_dir):\n", " \"\"\"\n", " Find pocket atom PDB files produced by fpocket.\n", " Returns a dict: pocket_id -> file path\n", " \"\"\"\n", " out_dir = Path(out_dir)\n", " pockets_dir = out_dir / \"pockets\"\n", "\n", " pocket_files = {}\n", " if pockets_dir.exists():\n", " for f in sorted(pockets_dir.glob(\"pocket*_atm.pdb\")):\n", " match = re.search(r\"pocket(\\d+)_atm\\.pdb\", f.name)\n", " if match:\n", " pocket_id = int(match.group(1))\n", " pocket_files[pocket_id] = f\n", "\n", " return pocket_files\n", "\n", "\n", "def save_json(data, out_file):\n", " \"\"\"\n", " Save a Python object as JSON.\n", " \"\"\"\n", " out_file = Path(out_file)\n", " with open(out_file, \"w\") as f:\n", " json.dump(data, f, indent=2)\n", " print(f\"[+] Saved JSON: {out_file}\")" ] }, { "cell_type": "code", "execution_count": 4, "id": "857ed980", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[*] Running ESMFold for job: 9SV1_1\n", "[*] Sequence length: 100 aa\n", "[✓] ESMFold prediction complete\n", "[+] Mean pLDDT: 0.76\n", "[+] Min pLDDT: 0.24\n", "[+] Max pLDDT: 0.96\n", "[!] Warning: overall structure confidence is low; pocket predictions may be less reliable. Consider using a pdb file from a high-quality source.\n" ] } ], "source": [ "# ================================\n", "# 🔬 RUN ESMFOLD\n", "# ================================\n", "\n", "# Tokenize input sequence\n", "tokenized_input = tokenizer(\n", " [USER_SEQUENCE],\n", " return_tensors=\"pt\",\n", " add_special_tokens=False\n", ")[\"input_ids\"].to(DEVICE)\n", "\n", "print(f\"[*] Running ESMFold for job: {SAFE_JOB_NAME}\")\n", "print(f\"[*] Sequence length: {len(USER_SEQUENCE)} aa\")\n", "\n", "# Run structure prediction\n", "with torch.no_grad():\n", " output = model(tokenized_input)\n", "\n", "# Basic confidence summary\n", "mean_plddt = output[\"plddt\"].mean().item()\n", "min_plddt = output[\"plddt\"].min().item()\n", "max_plddt = output[\"plddt\"].max().item()\n", "\n", "print(\"[✓] ESMFold prediction complete\")\n", "print(f\"[+] Mean pLDDT: {mean_plddt:.2f}\")\n", "print(f\"[+] Min pLDDT: {min_plddt:.2f}\")\n", "print(f\"[+] Max pLDDT: {max_plddt:.2f}\")\n", "\n", "if mean_plddt < 0.90:\n", " print(\"[!] Warning: overall structure confidence is low; pocket predictions may be less reliable. Consider using a pdb file from a high-quality source.\")" ] }, { "cell_type": "code", "execution_count": 5, "id": "d30d3d70", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[+] Saved PDB: 9SV1_1_metadata/9SV1_1/9SV1_1_0.pdb\n", "[✓] Saved 1 PDB file(s)\n" ] } ], "source": [ "# ================================\n", "# 💾 CONVERT OUTPUT TO PDB AND SAVE\n", "# ================================\n", "\n", "# Convert ESMFold outputs to PDB strings\n", "pdb_strings = convert_outputs_to_pdb(output)\n", "\n", "# Save PDB files\n", "pdb_paths = save_pdb_files(\n", " pdb_strings=pdb_strings,\n", " out_dir=JOB_DIR,\n", " prefix=SAFE_JOB_NAME\n", ")\n", "\n", "print(f\"[✓] Saved {len(pdb_paths)} PDB file(s)\")" ] }, { "cell_type": "code", "execution_count": null, "id": "5318125a", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "======================================================================\n", "Processing: 9SV1_1_0.pdb\n", "======================================================================\n", "\n", "Pocket 1\n", "Binding-site residue indices: [4, 5, 6, 52, 55, 56, 59, 79]\n", "Binding-site residue names: ['VAL', 'CYS', 'ILE', 'GLU', 'VAL', 'GLU', 'MET', 'PRO']\n", "\n", "Pocket 2\n", "Binding-site residue indices: [18, 21, 22, 25, 26, 61, 66, 67]\n", "Binding-site residue names: ['GLY', 'TYR', 'THR', 'TYR', 'GLU', 'TRP', 'GLY', 'PRO']\n", "\n", "Pocket 3\n", "Binding-site residue indices: [3, 4, 5, 31, 33, 49, 50, 51, 82, 84, 85, 86]\n", "Binding-site residue names: ['LYS', 'VAL', 'CYS', 'GLY', 'THR', 'CYS', 'GLY', 'GLU', 'PRO', 'GLY', 'GLU', 'LEU']\n", "\n", "Pocket 4\n", "Binding-site residue indices: [1, 2, 51, 52, 53, 54]\n", "Binding-site residue names: ['MET', 'SER', 'GLU', 'GLU', 'GLY', 'GLN']\n", "\n", "Pocket 5\n", "Binding-site residue indices: [3, 4, 5, 79, 80, 81, 82]\n", "Binding-site residue names: ['LYS', 'VAL', 'CYS', 'PRO', 'HIS', 'HIS', 'PRO']\n", "\n", "Pocket 6\n", "Binding-site residue indices: [34, 35, 85, 86, 87, 88, 90, 91, 92]\n", "Binding-site residue names: ['GLY', 'TYR', 'GLU', 'LEU', 'THR', 'ASP', 'ARG', 'ILE', 'ARG']\n", "\n", "Pocket 7\n", "Binding-site residue indices: [35, 47, 80, 82, 83, 84, 85, 86, 92]\n", "Binding-site residue names: ['TYR', 'VAL', 'HIS', 'PRO', 'SER', 'GLY', 'GLU', 'LEU', 'ARG']\n", "\n", "Pocket 8\n", "Binding-site residue indices: [20, 38, 92, 93, 94, 96]\n", "Binding-site residue names: ['ARG', 'ASN', 'ARG', 'LEU', 'GLU', 'HIS']\n", "\n", "Pocket 9\n", "Binding-site residue indices: [14, 15, 16, 17, 18, 19, 20, 21, 38]\n", "Binding-site residue names: ['VAL', 'GLN', 'GLY', 'VAL', 'GLY', 'PHE', 'ARG', 'TYR', 'ASN']\n" ] } ], "source": [ "# ================================\n", "# 🧪 FPOCKET + 5 Å BINDING-SITE RESIDUES\n", "# ================================\n", "\n", "from pathlib import Path\n", "import re\n", "import math\n", "import subprocess\n", "\n", "DIST_CUTOFF = 5.0\n", "\n", "AA3_TO_AA1 = {\n", " \"ALA\":\"A\",\"ARG\":\"R\",\"ASN\":\"N\",\"ASP\":\"D\",\"CYS\":\"C\",\n", " \"GLU\":\"E\",\"GLN\":\"Q\",\"GLY\":\"G\",\"HIS\":\"H\",\"ILE\":\"I\",\n", " \"LEU\":\"L\",\"LYS\":\"K\",\"MET\":\"M\",\"PHE\":\"F\",\"PRO\":\"P\",\n", " \"SER\":\"S\",\"THR\":\"T\",\"TRP\":\"W\",\"TYR\":\"Y\",\"VAL\":\"V\"\n", "}\n", "\n", "# --------------------------------\n", "# Run fpocket\n", "# --------------------------------\n", "def run_fpocket(pdb_path):\n", " pdb_path = Path(pdb_path).resolve()\n", "\n", " result = subprocess.run(\n", " [\"fpocket\", \"-f\", str(pdb_path)],\n", " stdout=subprocess.PIPE,\n", " stderr=subprocess.PIPE,\n", " text=True\n", " )\n", "\n", " if result.returncode != 0:\n", " print(result.stdout)\n", " print(result.stderr)\n", " raise RuntimeError(f\"fpocket failed for {pdb_path.name}\")\n", "\n", " return pdb_path.parent / f\"{pdb_path.stem}_out\"\n", "\n", "\n", "# --------------------------------\n", "# Parse fpocket summary\n", "# --------------------------------\n", "def parse_fpocket_info(info_file):\n", " pockets = {}\n", " current = None\n", "\n", " with open(info_file, \"r\") as f:\n", " for line in f:\n", " m = re.match(r\"Pocket\\s+(\\d+)\\s*:\", line)\n", " if m:\n", " current = int(m.group(1))\n", " pockets[current] = {}\n", " continue\n", "\n", " if current is None:\n", " continue\n", "\n", " fields = [\n", " (\"score\", r\"Score\\s*:\\s*(.+)\"),\n", " (\"druggability\", r\"Druggability Score\\s*:\\s*(.+)\"),\n", " (\"volume\", r\"Volume\\s*:\\s*(.+)\"),\n", " (\"n_spheres\", r\"Number of alpha spheres\\s*:\\s*(.+)\")\n", " ]\n", "\n", " for key, pattern in fields:\n", " hit = re.search(pattern, line)\n", " if hit:\n", " val = hit.group(1).strip()\n", " try:\n", " pockets[current][key] = int(val) if key == \"n_spheres\" else float(val)\n", " except ValueError:\n", " pockets[current][key] = val\n", "\n", " return pockets\n", "\n", "\n", "# --------------------------------\n", "# Find pocket files\n", "# --------------------------------\n", "def find_pocket_files(out_dir):\n", " pocket_dir = Path(out_dir) / \"pockets\"\n", " files = {}\n", "\n", " if not pocket_dir.exists():\n", " return files\n", "\n", " for p in pocket_dir.glob(\"pocket*_atm.pdb\"):\n", " m = re.search(r\"pocket(\\d+)_atm\\.pdb\", p.name)\n", " if m:\n", " pid = int(m.group(1))\n", " files[pid] = {\n", " \"atm\": p,\n", " \"vert\": pocket_dir / f\"pocket{pid}_vert.pqr\"\n", " }\n", "\n", " return files\n", "\n", "\n", "# --------------------------------\n", "# Read protein atoms\n", "# --------------------------------\n", "def parse_protein_atoms(pdb_file):\n", " atoms = []\n", "\n", " with open(pdb_file, \"r\") as f:\n", " for line in f:\n", " if not line.startswith(\"ATOM\"):\n", " continue\n", "\n", " try:\n", " atoms.append({\n", " \"resname\": line[17:20].strip(),\n", " \"chain\": line[21].strip() or \"-\",\n", " \"resid\": int(line[22:26].strip()),\n", " \"coord\": (\n", " float(line[30:38]),\n", " float(line[38:46]),\n", " float(line[46:54])\n", " )\n", " })\n", " except ValueError:\n", " continue\n", "\n", " return atoms\n", "\n", "\n", "# --------------------------------\n", "# Read alpha sphere coordinates\n", "# --------------------------------\n", "def get_alpha_sphere_coords(vert_file):\n", " coords = []\n", "\n", " if not Path(vert_file).exists():\n", " return coords\n", "\n", " with open(vert_file, \"r\") as f:\n", " for line in f:\n", " if not line.startswith((\"ATOM\", \"HETATM\")):\n", " continue\n", " try:\n", " coords.append((\n", " float(line[30:38]),\n", " float(line[38:46]),\n", " float(line[46:54])\n", " ))\n", " except ValueError:\n", " continue\n", "\n", " return coords\n", "\n", "\n", "# --------------------------------\n", "# Distance\n", "# --------------------------------\n", "def distance(a, b):\n", " return math.sqrt(\n", " (a[0] - b[0])**2 +\n", " (a[1] - b[1])**2 +\n", " (a[2] - b[2])**2\n", " )\n", "\n", "\n", "# --------------------------------\n", "# Residues within cutoff of alpha spheres\n", "# --------------------------------\n", "def get_binding_site_residues(protein_pdb, vert_file, cutoff=5.0):\n", " protein_atoms = parse_protein_atoms(protein_pdb)\n", " alpha_coords = get_alpha_sphere_coords(vert_file)\n", "\n", " selected = {}\n", "\n", " if not alpha_coords:\n", " return []\n", "\n", " for atom in protein_atoms:\n", " min_dist = min(distance(atom[\"coord\"], ac) for ac in alpha_coords)\n", "\n", " if min_dist <= cutoff:\n", " key = (atom[\"chain\"], atom[\"resid\"], atom[\"resname\"])\n", "\n", " if key not in selected or min_dist < selected[key][\"min_dist\"]:\n", " selected[key] = {\n", " \"index\": atom[\"resid\"],\n", " \"resid1\": atom[\"resid\"],\n", " \"chain\": atom[\"chain\"],\n", " \"resname\": atom[\"resname\"],\n", " \"min_dist\": min_dist\n", " }\n", "\n", " residues = sorted(selected.values(), key=lambda x: x[\"index\"])\n", " return residues\n", "\n", "\n", "# --------------------------------\n", "# Rank pockets\n", "# --------------------------------\n", "def rank_pockets(pocket_info):\n", " def safe_float(x, default=float(\"-inf\")):\n", " try:\n", " return float(x)\n", " except (TypeError, ValueError):\n", " return default\n", "\n", " ranked = sorted(\n", " pocket_info.items(),\n", " key=lambda kv: (\n", " safe_float(kv[1].get(\"score\")),\n", " safe_float(kv[1].get(\"druggability\")),\n", " safe_float(kv[1].get(\"n_spheres\"))\n", " ),\n", " reverse=True\n", " )\n", " return ranked\n", "\n", "\n", "# --------------------------------\n", "# Main\n", "# --------------------------------\n", "all_binding_sites = {}\n", "\n", "for pdb_path in pdb_paths:\n", " pdb_path = Path(pdb_path)\n", "\n", " print(f\"\\n{'='*70}\")\n", " print(f\"Processing: {pdb_path.name}\")\n", " print(f\"{'='*70}\")\n", "\n", " out_dir = run_fpocket(pdb_path)\n", " info_file = out_dir / f\"{pdb_path.stem}_info.txt\"\n", "\n", " if not info_file.exists():\n", " print(\"[!] fpocket info file not found\")\n", " continue\n", "\n", " pocket_info = parse_fpocket_info(info_file)\n", " pocket_files = find_pocket_files(out_dir)\n", "\n", " if not pocket_info or not pocket_files:\n", " print(\"[!] No pockets found\")\n", " continue\n", "\n", " ranked = rank_pockets(pocket_info)\n", " all_binding_sites[str(pdb_path)] = []\n", "\n", " for display_idx, (pid, info) in enumerate(ranked, start=1):\n", " if pid not in pocket_files:\n", " continue\n", "\n", " residues = get_binding_site_residues(\n", " protein_pdb=pdb_path,\n", " vert_file=pocket_files[pid][\"vert\"],\n", " cutoff=DIST_CUTOFF\n", " )\n", "\n", " residue_indices = [r[\"index\"] for r in residues]\n", " residue_names = [r[\"resname\"] for r in residues]\n", "\n", " all_binding_sites[str(pdb_path)].append({\n", " \"ranked_pocket_number\": display_idx,\n", " \"fpocket_id\": pid,\n", " \"score\": info.get(\"score\"),\n", " \"druggability\": info.get(\"druggability\"),\n", " \"volume\": info.get(\"volume\"),\n", " \"residue_indices\": residue_indices,\n", " \"residue_names\": residue_names\n", " })\n", "\n", " print(f\"\\nPocket {display_idx}\")\n", " #print(f\"fpocket ID: {pid}\")\n", " #print(f\"Score: {info.get('score', 'N/A')}\")\n", " #print(f\"Druggability Score: {info.get('druggability', 'N/A')}\")\n", " #print(f\"Volume: {info.get('volume', 'N/A')}\")\n", " print(f\"Binding-site residue indices: {residue_indices}\")\n", " print(f\"Binding-site residue names: {residue_names}\")" ] }, { "cell_type": "code", "execution_count": 10, "id": "e38f74ad", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "9SV1_1_0_pocket_1\n", " indices: [2, 3, 4, 78, 79, 80, 81]\n", " residues: ['LYS', 'VAL', 'CYS', 'PRO', 'HIS', 'HIS', 'PRO']\n", "9SV1_1_0_pocket_2\n", " indices: [19, 37, 91, 92, 93, 95]\n", " residues: ['ARG', 'ASN', 'ARG', 'LEU', 'GLU', 'HIS']\n", "9SV1_1_0_pocket_3\n", " indices: [13, 14, 15, 16, 17, 18, 19, 20, 37]\n", " residues: ['VAL', 'GLN', 'GLY', 'VAL', 'GLY', 'PHE', 'ARG', 'TYR', 'ASN']\n", "9SV1_1_0_pocket_4\n", " indices: [34, 46, 79, 81, 82, 83, 84, 85, 91]\n", " residues: ['TYR', 'VAL', 'HIS', 'PRO', 'SER', 'GLY', 'GLU', 'LEU', 'ARG']\n", "9SV1_1_0_pocket_5\n", " indices: [33, 34, 84, 85, 86, 87, 89, 90, 91]\n", " residues: ['GLY', 'TYR', 'GLU', 'LEU', 'THR', 'ASP', 'ARG', 'ILE', 'ARG']\n", "9SV1_1_0_pocket_6\n", " indices: [17, 20, 21, 24, 25, 60, 65, 66]\n", " residues: ['GLY', 'TYR', 'THR', 'TYR', 'GLU', 'TRP', 'GLY', 'PRO']\n", "9SV1_1_0_pocket_7\n", " indices: [2, 3, 4, 30, 32, 48, 49, 50, 81, 83, 84, 85]\n", " residues: ['LYS', 'VAL', 'CYS', 'GLY', 'THR', 'CYS', 'GLY', 'GLU', 'PRO', 'GLY', 'GLU', 'LEU']\n", "9SV1_1_0_pocket_8\n", " indices: [0, 1, 50, 51, 52, 53]\n", " residues: ['MET', 'SER', 'GLU', 'GLU', 'GLY', 'GLN']\n", "9SV1_1_0_pocket_9\n", " indices: [3, 4, 5, 51, 54, 55, 58, 78]\n", " residues: ['VAL', 'CYS', 'ILE', 'GLU', 'VAL', 'GLU', 'MET', 'PRO']\n" ] } ], "source": [ "# ================================\n", "# 🧠 EACH POCKET EMBEDDING AS SEPARATE VARIABLE\n", "# ================================\n", "\n", "import re\n", "import math\n", "import subprocess\n", "import importlib\n", "from pathlib import Path\n", "import torch\n", "\n", "DIST_CUTOFF = 5.0\n", "TOP_K_POCKETS = 10 # number of pockets to keep\n", "\n", "POSITIVE_RES = {\"ARG\", \"LYS\", \"HIS\"}\n", "POLAR_RES = {\"SER\", \"THR\", \"ASN\", \"GLN\", \"CYS\", \"TYR\"}\n", "NEGATIVE_RES = {\"ASP\", \"GLU\"}\n", "HYDROPHOBIC_RES = {\"ALA\", \"VAL\", \"ILE\", \"LEU\", \"MET\", \"PHE\", \"TRP\", \"PRO\"}\n", "\n", "W_POSITIVE = 3.0\n", "W_POLAR = 1.5\n", "W_NEGATIVE = -3.0\n", "W_HYDROPHOBIC = -1.0\n", "\n", "# --------------------------------\n", "# Install/load esm\n", "# --------------------------------\n", "if importlib.util.find_spec(\"esm\") is None:\n", " print(\"[*] Installing esm package...\")\n", " get_ipython().system(\"pip install -q fair-esm\")\n", "\n", "import esm\n", "\n", "# --------------------------------\n", "# Load ESM2\n", "# --------------------------------\n", "esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n", "batch_converter = alphabet.get_batch_converter()\n", "\n", "esm_model.eval()\n", "if torch.cuda.is_available():\n", " esm_model = esm_model.cuda()\n", "\n", "# --------------------------------\n", "# Your embedding functions\n", "# --------------------------------\n", "def get_set_of_embeddings_at_binding_site_esm(model, alphabet, batch_converter, seq_data, indices_for_averaging):\n", " batch_labels, batch_strs, batch_tokens = batch_converter([(seq_data[0], seq_data[1])])\n", "\n", " if torch.cuda.is_available():\n", " batch_tokens = batch_tokens.cuda()\n", "\n", " with torch.no_grad():\n", " results = model(batch_tokens, repr_layers=[33], return_contacts=False)\n", "\n", " token_representations = results[\"representations\"][33]\n", " return [token_representations[0][i] for i in indices_for_averaging]\n", "\n", "\n", "def preprocess_sequence(sequence, important_residues, esm_model, alphabet, batch_converter):\n", " set_of_embeddings = get_set_of_embeddings_at_binding_site_esm(\n", " esm_model, alphabet, batch_converter, (\"protein\", sequence), important_residues\n", " )\n", " avg_embedding = torch.mean(torch.stack(set_of_embeddings), dim=0)\n", " return avg_embedding.detach().cpu().numpy()\n", "\n", "# --------------------------------\n", "# fpocket helpers\n", "# --------------------------------\n", "def run_fpocket(pdb_path):\n", " pdb_path = Path(pdb_path).resolve()\n", " result = subprocess.run(\n", " [\"fpocket\", \"-f\", str(pdb_path)],\n", " stdout=subprocess.PIPE,\n", " stderr=subprocess.PIPE,\n", " text=True\n", " )\n", " if result.returncode != 0:\n", " print(result.stdout)\n", " print(result.stderr)\n", " raise RuntimeError(f\"fpocket failed for {pdb_path.name}\")\n", " return pdb_path.parent / f\"{pdb_path.stem}_out\"\n", "\n", "\n", "def parse_fpocket_info(info_file):\n", " pockets = {}\n", " current = None\n", "\n", " with open(info_file, \"r\") as f:\n", " for line in f:\n", " m = re.match(r\"Pocket\\s+(\\d+)\\s*:\", line)\n", " if m:\n", " current = int(m.group(1))\n", " pockets[current] = {}\n", " continue\n", "\n", " if current is None:\n", " continue\n", "\n", " fields = [\n", " (\"score\", r\"Score\\s*:\\s*(.+)\"),\n", " (\"volume\", r\"Volume\\s*:\\s*(.+)\"),\n", " (\"n_spheres\", r\"Number of alpha spheres\\s*:\\s*(.+)\")\n", " ]\n", "\n", " for key, pattern in fields:\n", " hit = re.search(pattern, line)\n", " if hit:\n", " val = hit.group(1).strip()\n", " try:\n", " pockets[current][key] = int(val) if key == \"n_spheres\" else float(val)\n", " except ValueError:\n", " pockets[current][key] = val\n", " return pockets\n", "\n", "\n", "def find_pocket_files(out_dir):\n", " pocket_dir = Path(out_dir) / \"pockets\"\n", " files = {}\n", " if not pocket_dir.exists():\n", " return files\n", "\n", " for p in pocket_dir.glob(\"pocket*_atm.pdb\"):\n", " m = re.search(r\"pocket(\\d+)_atm\\.pdb\", p.name)\n", " if m:\n", " pid = int(m.group(1))\n", " files[pid] = {\n", " \"atm\": p,\n", " \"vert\": pocket_dir / f\"pocket{pid}_vert.pqr\"\n", " }\n", " return files\n", "\n", "\n", "def parse_protein_atoms(pdb_file):\n", " atoms = []\n", " with open(pdb_file, \"r\") as f:\n", " for line in f:\n", " if not line.startswith(\"ATOM\"):\n", " continue\n", " try:\n", " atoms.append({\n", " \"resname\": line[17:20].strip(),\n", " \"chain\": line[21].strip() or \"-\",\n", " \"resid\": int(line[22:26].strip()),\n", " \"coord\": (\n", " float(line[30:38]),\n", " float(line[38:46]),\n", " float(line[46:54])\n", " )\n", " })\n", " except ValueError:\n", " continue\n", " return atoms\n", "\n", "\n", "def get_alpha_sphere_coords(vert_file):\n", " coords = []\n", " vert_file = Path(vert_file)\n", " if not vert_file.exists():\n", " return coords\n", "\n", " with open(vert_file, \"r\") as f:\n", " for line in f:\n", " if not line.startswith((\"ATOM\", \"HETATM\")):\n", " continue\n", " try:\n", " coords.append((\n", " float(line[30:38]),\n", " float(line[38:46]),\n", " float(line[46:54])\n", " ))\n", " except ValueError:\n", " continue\n", " return coords\n", "\n", "\n", "def distance(a, b):\n", " return math.sqrt((a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2)\n", "\n", "\n", "def get_binding_site_residues(protein_pdb, vert_file, cutoff=5.0):\n", " protein_atoms = parse_protein_atoms(protein_pdb)\n", " alpha_coords = get_alpha_sphere_coords(vert_file)\n", "\n", " selected = {}\n", " if not alpha_coords:\n", " return []\n", "\n", " for atom in protein_atoms:\n", " min_dist = min(distance(atom[\"coord\"], ac) for ac in alpha_coords)\n", "\n", " if min_dist <= cutoff:\n", " key = (atom[\"chain\"], atom[\"resid\"], atom[\"resname\"])\n", " if key not in selected or min_dist < selected[key][\"min_dist\"]:\n", " seq_index = atom[\"resid\"] - 1\n", " selected[key] = {\n", " \"index\": seq_index,\n", " \"resname\": atom[\"resname\"],\n", " \"min_dist\": min_dist\n", " }\n", "\n", " return sorted(selected.values(), key=lambda x: x[\"index\"])\n", "\n", "\n", "def pocket_residue_counts(residues):\n", " counts = {\"positive\": 0, \"polar\": 0, \"negative\": 0, \"hydrophobic\": 0, \"other\": 0}\n", " for r in residues:\n", " res = r[\"resname\"]\n", " if res in POSITIVE_RES:\n", " counts[\"positive\"] += 1\n", " elif res in POLAR_RES:\n", " counts[\"polar\"] += 1\n", " elif res in NEGATIVE_RES:\n", " counts[\"negative\"] += 1\n", " elif res in HYDROPHOBIC_RES:\n", " counts[\"hydrophobic\"] += 1\n", " else:\n", " counts[\"other\"] += 1\n", " return counts\n", "\n", "\n", "def anion_preference_score(residues):\n", " counts = pocket_residue_counts(residues)\n", " score = (\n", " W_POSITIVE * counts[\"positive\"] +\n", " W_POLAR * counts[\"polar\"] +\n", " W_NEGATIVE * counts[\"negative\"] +\n", " W_HYDROPHOBIC * counts[\"hydrophobic\"]\n", " )\n", " return score, counts\n", "\n", "\n", "def select_top_anion_pockets(pdb_path, pocket_info, pocket_files, cutoff=5.0, top_k=3):\n", " candidates = []\n", "\n", " for pid, info in pocket_info.items():\n", " if pid not in pocket_files:\n", " continue\n", "\n", " residues = get_binding_site_residues(\n", " protein_pdb=pdb_path,\n", " vert_file=pocket_files[pid][\"vert\"],\n", " cutoff=cutoff\n", " )\n", "\n", " pref_score, counts = anion_preference_score(residues)\n", "\n", " candidates.append({\n", " \"pid\": pid,\n", " \"fpocket_score\": info.get(\"score\", None),\n", " \"volume\": info.get(\"volume\", None),\n", " \"n_spheres\": info.get(\"n_spheres\", None),\n", " \"anion_pref_score\": pref_score,\n", " \"counts\": counts,\n", " \"residues\": residues\n", " })\n", "\n", " def safe_float(x, default=float(\"-inf\")):\n", " try:\n", " return float(x)\n", " except (TypeError, ValueError):\n", " return default\n", "\n", " candidates = sorted(\n", " candidates,\n", " key=lambda x: (\n", " x[\"anion_pref_score\"],\n", " safe_float(x[\"fpocket_score\"]),\n", " safe_float(x[\"n_spheres\"])\n", " ),\n", " reverse=True\n", " )\n", "\n", " return candidates[:top_k]\n", "\n", "# --------------------------------\n", "# MAIN\n", "# --------------------------------\n", "pocket_embeddings = {}\n", "\n", "for pdb_path in pdb_paths:\n", " pdb_path = Path(pdb_path)\n", " out_dir = run_fpocket(pdb_path)\n", " info_file = out_dir / f\"{pdb_path.stem}_info.txt\"\n", "\n", " pocket_info = parse_fpocket_info(info_file)\n", " pocket_files = find_pocket_files(out_dir)\n", "\n", " selected_pockets = select_top_anion_pockets(\n", " pdb_path=pdb_path,\n", " pocket_info=pocket_info,\n", " pocket_files=pocket_files,\n", " cutoff=DIST_CUTOFF,\n", " top_k=TOP_K_POCKETS\n", " )\n", "\n", " for rank_i, pocket in enumerate(selected_pockets, start=1):\n", " pocket_indices = [r[\"index\"] for r in pocket[\"residues\"]]\n", " pocket_resnames = [r[\"resname\"] for r in pocket[\"residues\"]]\n", "\n", " if not pocket_indices:\n", " continue\n", "\n", " emb = preprocess_sequence(\n", " USER_SEQUENCE,\n", " pocket_indices,\n", " esm_model,\n", " alphabet,\n", " batch_converter\n", " )\n", "\n", " var_name = f\"{pdb_path.stem}_pocket_{rank_i}\"\n", " pocket_embeddings[var_name] = {\n", " \"embedding\": emb,\n", " \"indices\": pocket_indices,\n", " \"resnames\": pocket_resnames,\n", " \"fpocket_id\": pocket[\"pid\"],\n", " \"fpocket_score\": pocket[\"fpocket_score\"],\n", " \"anion_preference_score\": pocket[\"anion_pref_score\"]\n", " }\n", "\n", " print(f\"{var_name}\")\n", " print(f\" indices: {pocket_indices}\")\n", " print(f\" residues: {pocket_resnames}\")" ] }, { "cell_type": "code", "execution_count": 11, "id": "12eca5ed", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_1\n", "======================================================================\n", "fpocket ID: 5\n", "Residue indices: [2, 3, 4, 78, 79, 80, 81]\n", "Residue names: ['LYS', 'VAL', 'CYS', 'PRO', 'HIS', 'HIS', 'PRO']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 0\n", "Predicted Class : phosphate\n", "Confidence : 0.7295\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.7295\n", " Class 1 (sulfate): 0.2048\n", " Class 2 (chloride): 0.0654\n", " Class 3 (nitrate): 0.0001\n", " Class 4 (carbonate): 0.0003\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_2\n", "======================================================================\n", "fpocket ID: 7\n", "Residue indices: [19, 37, 91, 92, 93, 95]\n", "Residue names: ['ARG', 'ASN', 'ARG', 'LEU', 'GLU', 'HIS']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 2\n", "Predicted Class : chloride\n", "Confidence : 0.7145\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.0014\n", " Class 1 (sulfate): 0.2841\n", " Class 2 (chloride): 0.7145\n", " Class 3 (nitrate): 0.0001\n", " Class 4 (carbonate): 0.0000\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_3\n", "======================================================================\n", "fpocket ID: 9\n", "Residue indices: [13, 14, 15, 16, 17, 18, 19, 20, 37]\n", "Residue names: ['VAL', 'GLN', 'GLY', 'VAL', 'GLY', 'PHE', 'ARG', 'TYR', 'ASN']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 0\n", "Predicted Class : phosphate\n", "Confidence : 0.7601\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.7601\n", " Class 1 (sulfate): 0.2398\n", " Class 2 (chloride): 0.0001\n", " Class 3 (nitrate): 0.0000\n", " Class 4 (carbonate): 0.0000\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_4\n", "======================================================================\n", "fpocket ID: 4\n", "Residue indices: [34, 46, 79, 81, 82, 83, 84, 85, 91]\n", "Residue names: ['TYR', 'VAL', 'HIS', 'PRO', 'SER', 'GLY', 'GLU', 'LEU', 'ARG']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 1\n", "Predicted Class : sulfate\n", "Confidence : 0.9441\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.0438\n", " Class 1 (sulfate): 0.9441\n", " Class 2 (chloride): 0.0114\n", " Class 3 (nitrate): 0.0005\n", " Class 4 (carbonate): 0.0002\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_5\n", "======================================================================\n", "fpocket ID: 2\n", "Residue indices: [33, 34, 84, 85, 86, 87, 89, 90, 91]\n", "Residue names: ['GLY', 'TYR', 'GLU', 'LEU', 'THR', 'ASP', 'ARG', 'ILE', 'ARG']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 1\n", "Predicted Class : sulfate\n", "Confidence : 0.8455\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.0022\n", " Class 1 (sulfate): 0.8455\n", " Class 2 (chloride): 0.1511\n", " Class 3 (nitrate): 0.0007\n", " Class 4 (carbonate): 0.0004\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_6\n", "======================================================================\n", "fpocket ID: 3\n", "Residue indices: [17, 20, 21, 24, 25, 60, 65, 66]\n", "Residue names: ['GLY', 'TYR', 'THR', 'TYR', 'GLU', 'TRP', 'GLY', 'PRO']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 1\n", "Predicted Class : sulfate\n", "Confidence : 0.8267\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.1668\n", " Class 1 (sulfate): 0.8267\n", " Class 2 (chloride): 0.0063\n", " Class 3 (nitrate): 0.0000\n", " Class 4 (carbonate): 0.0002\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_7\n", "======================================================================\n", "fpocket ID: 1\n", "Residue indices: [2, 3, 4, 30, 32, 48, 49, 50, 81, 83, 84, 85]\n", "Residue names: ['LYS', 'VAL', 'CYS', 'GLY', 'THR', 'CYS', 'GLY', 'GLU', 'PRO', 'GLY', 'GLU', 'LEU']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 1\n", "Predicted Class : sulfate\n", "Confidence : 0.8070\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.1622\n", " Class 1 (sulfate): 0.8070\n", " Class 2 (chloride): 0.0221\n", " Class 3 (nitrate): 0.0075\n", " Class 4 (carbonate): 0.0012\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_8\n", "======================================================================\n", "fpocket ID: 6\n", "Residue indices: [0, 1, 50, 51, 52, 53]\n", "Residue names: ['MET', 'SER', 'GLU', 'GLU', 'GLY', 'GLN']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 0\n", "Predicted Class : phosphate\n", "Confidence : 0.5607\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.5607\n", " Class 1 (sulfate): 0.4375\n", " Class 2 (chloride): 0.0012\n", " Class 3 (nitrate): 0.0001\n", " Class 4 (carbonate): 0.0005\n", "\n", "======================================================================\n", "Pocket: 9SV1_1_0_pocket_9\n", "======================================================================\n", "fpocket ID: 8\n", "Residue indices: [3, 4, 5, 51, 54, 55, 58, 78]\n", "Residue names: ['VAL', 'CYS', 'ILE', 'GLU', 'VAL', 'GLU', 'MET', 'PRO']\n", "\n", "==============================\n", "PREDICTION RESULT (MLP)\n", "==============================\n", "Predicted Class ID : 2\n", "Predicted Class : chloride\n", "Confidence : 0.7859\n", "\n", "Full Probability Distribution:\n", " Class 0 (phosphate): 0.0026\n", " Class 1 (sulfate): 0.0487\n", " Class 2 (chloride): 0.7859\n", " Class 3 (nitrate): 0.1625\n", " Class 4 (carbonate): 0.0003\n" ] } ], "source": [ "# ================================\n", "# 🔮 PREDICT FOR EACH POCKET EMBEDDING\n", "# ================================\n", "\n", "import os\n", "import json\n", "import numpy as np\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "\n", "# ---- same MLP definition (must match training) ----\n", "class MLP(nn.Module):\n", " def __init__(self, in_dim, h1, h2, h3, out_dim, p_drop=0.3, use_bn=True):\n", " super().__init__()\n", " self.fc1 = nn.Linear(in_dim, h1)\n", " self.bn1 = nn.BatchNorm1d(h1) if use_bn else nn.Identity()\n", " self.fc2 = nn.Linear(h1, h2)\n", " self.bn2 = nn.BatchNorm1d(h2) if use_bn else nn.Identity()\n", " self.fc3 = nn.Linear(h2, h3)\n", " self.bn3 = nn.BatchNorm1d(h3) if use_bn else nn.Identity()\n", " self.out = nn.Linear(h3, out_dim)\n", " self.drop = nn.Dropout(p_drop)\n", "\n", " def forward(self, x):\n", " x = self.fc1(x); x = self.bn1(x); x = F.relu(x)\n", " x = self.fc2(x); x = self.bn2(x); x = F.relu(x)\n", " x = self.fc3(x); x = self.bn3(x); x = F.relu(x)\n", " return self.out(x)\n", "\n", "def load_mlp(save_dir, device=None):\n", " device = device or torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "\n", " with open(os.path.join(save_dir, \"metadata.json\"), \"r\") as f:\n", " meta = json.load(f)\n", "\n", " in_dim = meta[\"feature_dim\"]\n", " out_dim = meta[\"n_classes\"]\n", "\n", " hidden_sizes = meta[\"architecture\"][\"hidden_sizes\"]\n", " h1, h2, h3 = hidden_sizes\n", "\n", " best_params = meta[\"best_params\"]\n", "\n", " model = MLP(\n", " in_dim=in_dim,\n", " h1=h1,\n", " h2=h2,\n", " h3=h3,\n", " out_dim=out_dim,\n", " p_drop=best_params[\"dropout\"],\n", " use_bn=best_params[\"use_bn\"],\n", " ).to(device)\n", "\n", " sd = torch.load(os.path.join(save_dir, \"mlp_state_dict.pt\"), map_location=device)\n", " model.load_state_dict(sd)\n", " model.eval()\n", " return model, meta, device\n", "\n", "def predict_one(vec_1d, model, meta, device, topk=3):\n", " x = torch.tensor(np.asarray(vec_1d, dtype=np.float32), device=device).view(1, -1)\n", "\n", " if x.shape[1] != meta[\"feature_dim\"]:\n", " raise ValueError(f\"Expected {meta['feature_dim']} features, got {x.shape[1]}\")\n", "\n", " with torch.no_grad():\n", " logits = model(x)\n", " probs = torch.softmax(logits, dim=1).cpu().numpy()[0]\n", "\n", " pred_class = int(np.argmax(probs))\n", " pred_name = meta[\"class_id_to_name\"].get(str(pred_class), str(pred_class))\n", " conf = float(probs[pred_class])\n", "\n", " print(\"\\n==============================\")\n", " print(\"PREDICTION RESULT (MLP)\")\n", " print(\"==============================\")\n", " print(f\"Predicted Class ID : {pred_class}\")\n", " print(f\"Predicted Class : {pred_name}\")\n", " print(f\"Confidence : {conf:.4f}\")\n", "\n", " print(\"\\nFull Probability Distribution:\")\n", " for i, p in enumerate(probs):\n", " cname = meta[\"class_id_to_name\"].get(str(i), str(i))\n", " print(f\" Class {i} ({cname}): {p:.4f}\")\n", "\n", " return pred_class, pred_name, conf, probs\n", "\n", "\n", "# Load trained model\n", "model, meta, device = load_mlp(\n", " \"/mnt/disk3/Arun_work/phospher_work/Results/results_mlp_BW_hparam_sweep/best_model_20260320-103259\"\n", ")\n", "\n", "all_predictions = {}\n", "\n", "for name, data in pocket_embeddings.items():\n", " emb = data[\"embedding\"]\n", "\n", " print(f\"\\n{'='*70}\")\n", " print(f\"Pocket: {name}\")\n", " print(f\"{'='*70}\")\n", " print(f\"fpocket ID: {data['fpocket_id']}\")\n", " print(f\"Residue indices: {data['indices']}\")\n", " print(f\"Residue names: {data['resnames']}\")\n", "\n", " pred_class, pred_name, conf, probs = predict_one(\n", " emb,\n", " model,\n", " meta,\n", " device,\n", " topk=3\n", " )\n", "\n", " all_predictions[name] = {\n", " \"fpocket_id\": data[\"fpocket_id\"],\n", " \"indices\": data[\"indices\"],\n", " \"resnames\": data[\"resnames\"],\n", " \"predicted_class_id\": pred_class,\n", " \"predicted_class_name\": pred_name,\n", " \"confidence\": conf,\n", " \"probabilities\": probs.tolist()\n", " }" ] } ], "metadata": { "kernelspec": { "display_name": "phosbind_env", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.12" } }, "nbformat": 4, "nbformat_minor": 5 }