{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --quiet --upgrade --force-reinstall \"fair-esm[esmfold]\" biopython==1.81 omegaconf einops ml-collections","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-01T02:12:23.861533Z","iopub.execute_input":"2025-11-01T02:12:23.861958Z","iopub.status.idle":"2025-11-01T02:12:29.657319Z","shell.execute_reply.started":"2025-11-01T02:12:23.861902Z","shell.execute_reply":"2025-11-01T02:12:29.656337Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working\n!git clone https://github.com/aqlaboratory/openfold.git\n!pip install --quiet \"biopython==1.81\" omegaconf einops ml-collections\n!pip install --quiet -r openfold/requirements.txt\n!pip install --quiet -e openfold","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, re\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom Bio import SeqIO\nimport torch\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (kích hoạt 3D proj)\n\n# ==== ĐƯỜNG DẪN SỬA TÙY BẠN ====\nROOT = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\n\nP = {\n    \"train_fasta\": ROOT/\"Train\"/\"train_sequences.fasta\",\n    \"test_fasta\":  ROOT/\"Test\"/\"testsuperset.fasta\",\n}\n\nOUTDIR = \"/kaggle/working/pred_struct\"\n\n# ==== GPU ====\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"DEVICE:\", DEVICE)\nprint(\"GPU count:\", torch.cuda.device_count())\n\nMIXED_PRECISION = True  # dùng fp16 autocast cho nhanh/đỡ VRAM\nMAX_LEN = 1022          # ESMFold thường giới hạn ~1022 aa; quá dài thì bỏ qua/handle riêng\nBATCH_SIZE = 2          # batch inference nhiều protein 1 lúc (tối đa tùy VRAM T4x2)\nSKIP_IF_EXISTS = True   # không chạy lại protein đã fold rồi","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def norm_uniprot_id(header_id: str) -> str:\n    \"\"\"\n    Convert 'sp|P9WHI7|RECN_MYCT ...' -> 'P9WHI7'\n    \"\"\"\n    rid = str(header_id).split()[0]\n    if \"|\" in rid:\n        parts = rid.split(\"|\")\n        if len(parts) >= 2 and parts[1]:\n            return parts[1]\n    return rid\n\ndef load_fasta_list(fasta_path: Path):\n    \"\"\"\n    Trả về list dict:\n      {\n        \"target\": accession_id,\n        \"seq\":    amino acid sequence (str),\n        \"len\":    sequence length\n      }\n    Không lẫn taxon gì ở đây, vì DeepFRI chỉ cần seq → cấu trúc.\n    \"\"\"\n    rows = []\n    for rec in SeqIO.parse(str(fasta_path), \"fasta\"):\n        acc = norm_uniprot_id(rec.id)\n        seq = str(rec.seq)\n        rows.append({\n            \"target\": acc,\n            \"seq\": seq,\n            \"len\": len(seq),\n        })\n    return rows\n\ntrain_list = load_fasta_list(P[\"train_fasta\"])\ntest_list  = load_fasta_list(P[\"test_fasta\"])\n\nprint(\"Train proteins:\", len(train_list))\nprint(\"Test  proteins:\", len(test_list))\nprint(\"Sample train entry:\", train_list[0])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import esm\n\n# Load pretrained ESMFold\nesmfold_model = esm.pretrained.esmfold_v1()\nesmfold_model = esmfold_model.eval()  # inference mode\n\n# Đưa lên GPU\nesmfold_model = esmfold_model.to(DEVICE)\n\n# Dùng DataParallel nếu có >1 GPU\nif torch.cuda.device_count() > 1:\n    print(\"Using DataParallel across\", torch.cuda.device_count(), \"GPUs\")\n    esmfold_model = nn.DataParallel(esmfold_model)\n\n# Hàm tiện: model nhận list[str] (các chuỗi aa), trả batch output\ndef fold_batch_sequences(seq_list):\n    \"\"\"\n    Input:\n      seq_list: list of aa strings length B\n    Output:\n      coords_ca_list: list of np.array [L,3] (C-alpha coords)\n    \"\"\"\n    with torch.no_grad():\n        if MIXED_PRECISION and DEVICE == \"cuda\":\n            autocast_ctx = torch.cuda.amp.autocast(dtype=torch.float16)\n        else:\n            # dummy context manager\n            class Dummy: \n                def __enter__(self): pass\n                def __exit__(self,a,b,c): pass\n            autocast_ctx = Dummy()\n\n        with autocast_ctx:\n            # ESMFold forward\n            # esmfold_model expects a list of sequences\n            output = esmfold_model(seq_list)\n\n            # output[\"positions\"]: shape [B, L, n_atoms, 3]\n            # L có thể khác nhau giữa các seq trong batch => ESMFold pad nội bộ.\n            # Với DataParallel, output có thể là dict of tensors hoặc list of dict,\n            # tuỳ version fair-esm. Chúng ta handle 2 case phổ biến:\n\n            if isinstance(output, dict):\n                pos = output[\"positions\"]  # [B, L, n_atoms, 3] tensor\n                coords_all = pos.detach().float().cpu().numpy()\n            elif isinstance(output, list):\n                # một số bản DataParallel trả list dict per device, ta gộp tay\n                # đơn giản nhất: concat theo batch dim\n                coords_list = []\n                for out_i in output:\n                    coords_list.append(out_i[\"positions\"].detach().float().cpu().numpy())\n                coords_all = np.concatenate(coords_list, axis=0)\n            else:\n                raise RuntimeError(\"Unexpected ESMFold output type\")\n\n    # Lấy C-alpha (CA) = index 1 trong n_atoms\n    # coords_all shape [B, L, n_atoms, 3] -> [B, L, 3]\n    coords_ca_list = []\n    for b in range(coords_all.shape[0]):\n        coords_ca = coords_all[b, :, 1, :]  # lấy atom index=1 => CA\n        coords_ca_list.append(coords_ca)\n    return coords_ca_list","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def save_coords_npz(prot_id, seq, coords, out_dir: Path):\n    out_path = out_dir / f\"{prot_id}.npz\"\n    np.savez_compressed(out_path, coords=coords.astype(np.float32), seq=np.array(seq))\n    return out_path\n\ndef generate_structures_for_list(protein_list, out_dir: Path, split_name=\"train\"):\n    \"\"\"\n    protein_list: list of dict {\"target\",\"seq\",\"len\"}\n    \"\"\"\n    miss_too_long = 0\n    already = 0\n    done = 0\n\n    # tạo danh sách các protein mà ta sẽ thực sự fold\n    jobs = []\n    for p in protein_list:\n        prot_id = p[\"target\"]\n        seq     = p[\"seq\"]\n        L       = p[\"len\"]\n\n        out_path = out_dir / f\"{prot_id}.npz\"\n        if SKIP_IF_EXISTS and out_path.exists():\n            already += 1\n            continue\n\n        if L > MAX_LEN:\n            # skip protein quá dài\n            miss_too_long += 1\n            continue\n\n        jobs.append((prot_id, seq))\n\n    print(f\"[{split_name}] total={len(protein_list)}, to_run={len(jobs)}, already={already}, too_long={miss_too_long}\")\n\n    # chạy theo batch\n    for i in tqdm(range(0, len(jobs), BATCH_SIZE), desc=f\"Folding {split_name}\"):\n        batch_jobs = jobs[i:i+BATCH_SIZE]\n        batch_ids  = [j[0] for j in batch_jobs]\n        batch_seqs = [j[1] for j in batch_jobs]\n\n        coords_batch = fold_batch_sequences(batch_seqs)  # list np[L,3]\n        for pid, seq, coords in zip(batch_ids, batch_seqs, coords_batch):\n            save_coords_npz(pid, seq, coords, out_dir)\n            done += 1\n\n    print(f\"[{split_name}] generated new {done} structures\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"generate_structures_for_list(train_list, OUTDIR, split_name=\"train\")\ngenerate_structures_for_list(test_list,  OUTDIR, split_name=\"test\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_protein_coords(coords, title=\"Protein 3D\"):\n    xs = coords[:,0]\n    ys = coords[:,1]\n    zs = coords[:,2]\n    fig = plt.figure(figsize=(5,5))\n    ax = fig.add_subplot(111, projection='3d')\n    ax.scatter(xs, ys, zs, s=8)\n    ax.plot(xs, ys, zs, linewidth=1, alpha=0.4)\n    ax.set_xlabel(\"x\")\n    ax.set_ylabel(\"y\")\n    ax.set_zlabel(\"z\")\n    ax.set_title(title)\n    plt.show()\n\n# ví dụ: chọn 1 protein đã sinh\nexample_pid = train_list[0][\"target\"]\nexample_path = OUTDIR / f\"{example_pid}.npz\"\nif example_path.exists():\n    data = np.load(example_path)\n    coords_ex = data[\"coords\"]  # (L,3)\n    plot_protein_coords(coords_ex, title=f\"{example_pid} (len={coords_ex.shape[0]})\")\nelse:\n    print(\"Chưa có file ví dụ để plot. Chạy Cell 6 trước đã.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}