{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "machine_shape": "hm", "gpuType": "L4", "include_colab_link": true }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "source": [ "# The AF2cycler\n", "This notebook contains the code to run the af2cycler and use for improvement of suboptimal designed protein backbones.\n", "Based on:\n", "\n", "**Alphafold2 refinement improves designability of large de novo proteins**\n", "\n", "Christopher Frank, Dominik Schiwietz, Lara Fuß, Sergey Ovchinnikov and Hendrik Dietz\n", "\n", "We recommend to run this Notebook with at leat a L4 or better a A100 GPU as the GPU memeory needed for ESMFold is quite significant\n", "\n" ], "metadata": { "id": "wsji9enfo8aM" } }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "yv2S1XouoxZF" }, "outputs": [], "source": [ "#@title setup\n", "%%time\n", "import os\n", "if not os.path.isdir(\"params\"):\n", " # get code\n", " os.system(\"pip -q install pyppeteer nest_asyncio\")\n", " os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git\")\n", " # for debugging\n", " os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n", " # download params\n", " os.system(\"mkdir params\")\n", " os.system(\"apt-get install aria2 -qq\")\n", " os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n", " os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n", "\n", "import warnings\n", "warnings.simplefilter(action='ignore', category=FutureWarning)\n", "\n", "import os\n", "from colabdesign import mk_afdesign_model, clear_mem\n", "from colabdesign.mpnn import mk_mpnn_model\n", "\n", "from IPython.display import HTML\n", "from google.colab import files\n", "import numpy as np\n", "\n", "import requests, time\n", "if not os.path.isfile(\"TMscore\"):\n", " os.system(\"wget -qnc https://zhanggroup.org/TM-score/TMscore.cpp\")\n", " os.system(\"g++ -static -O3 -ffast-math -lm -o TMscore TMscore.cpp\")\n", "def tmscore(x,y):\n", " # pass to TMscore\n", " output = os.popen(f'./TMscore {x} {y}')\n", " # parse outputs\n", " parse_float = lambda x: float(x.split(\"=\")[1].split()[0])\n", " o = {}\n", " for line in output:\n", " line = line.rstrip()\n", " if line.startswith(\"RMSD\"): o[\"rms\"] = parse_float(line)\n", " if line.startswith(\"TM-score\"): o[\"tms\"] = parse_float(line)\n", " if line.startswith(\"GDT-TS-score\"): o[\"gdt\"] = parse_float(line)\n", " return o\n", "\n", "import asyncio\n", "import nest_asyncio\n", "from pyppeteer import launch\n", "import base64\n", "\n", "# Apply nest_asyncio to enable nested event loops\n", "nest_asyncio.apply()\n", "\n", "async def fetch_blob_content(page, blob_url):\n", " blob_to_base64 = \"\"\"\n", " async (blobUrl) => {\n", " const blob = await fetch(blobUrl).then(r => r.blob());\n", " return new Promise((resolve) => {\n", " const reader = new FileReader();\n", " reader.onloadend = () => resolve(reader.result);\n", " reader.readAsDataURL(blob);\n", " });\n", " }\n", " \"\"\"\n", " base64_data = await page.evaluate(blob_to_base64, blob_url)\n", " _, encoded = base64_data.split(',', 1)\n", " return base64.b64decode(encoded)\n", "\n", "async def extract_pdb_file_download_link_and_content(url):\n", " browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])\n", " page = await browser.newPage()\n", " await page.goto(url, {'waitUntil': 'networkidle0'})\n", " elements = await page.querySelectorAll('a.btn.bg-purple')\n", " for element in elements:\n", " href = await page.evaluate('(element) => element.getAttribute(\"href\")', element)\n", " if 'blob:https://esmatlas.com/' in href:\n", " content = await fetch_blob_content(page, href)\n", " await browser.close()\n", " return href, content\n", " await browser.close()\n", " return \"No PDB file link found.\", None\n", "\n", "def esmfold_api(sequence):\n", " url = f'https://esmatlas.com/resources/fold/result?fasta_header=%3Eunnamed&sequence={sequence}'\n", " result = asyncio.get_event_loop().run_until_complete(extract_pdb_file_download_link_and_content(url))\n", " if result[1]:\n", " pdb_str = result[1].decode('utf-8')\n", " return pdb_str\n", " else:\n", " return \"Failed to retrieve PDB content.\"\n", "\n", "import jax\n", "import jax.numpy as jnp\n", "from colabdesign.af.alphafold.common import residue_constants\n", "\n", "if not os.path.exists('/content/in/'):\n", " os.mkdir('/content/in/')\n", "\n", "import py3Dmol\n", "\n", "def visualize_pdb_overlay(pdb1_path, pdb2_path):\n", " viewer = py3Dmol.view(width=800, height=600)\n", "\n", " with open(pdb1_path, 'r') as f:\n", " pdb1_data = f.read()\n", " viewer.addModel(pdb1_data, 'pdb')\n", " viewer.setStyle({'model': 0}, {'cartoon': {'color': 'grey'}})\n", "\n", " with open(pdb2_path, 'r') as f:\n", " pdb2_data = f.read()\n", " viewer.addModel(pdb2_data, 'pdb')\n", " viewer.setStyle({'model': 1}, {'cartoon': {'color' : 'red'}})\n", "\n", " viewer.zoomTo()\n", " viewer.show()\n", "\n", "\n" ] }, { "cell_type": "code", "source": [], "metadata": { "id": "bxL_tvmGqg-S" }, "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "source": [ "# Chroma Design\n", "\n", "On the start of this pipeline is the creation of draft backbones using Chroma or any other design method you want to use. We suggest you check out this two notebooks on Chroma to generate your desired proteins:\n", "**Chroma Quickstart**\n", "https://colab.research.google.com/github/generatebio/chroma/blob/main/notebooks/ChromaDemo.ipynb\n", "\n", "**Chroma API Tutorial**\n", "https://colab.research.google.com/github/generatebio/chroma/blob/main/notebooks/ChromaAPI.ipynb\n", "\n", "If you have your PDB files please upload them into the **in/** folder and proceed\n" ], "metadata": { "id": "WMpjGEyGqgSH" } }, { "cell_type": "code", "source": [ "#@title Monomer af2cycling\n", "#@markdown The af2cycler takes in the Chroma design and returns a new pdb file with improved structure\n", "iterations = 10 #@param {type:\"integer\"}\n", "#@markdown The af2cycled model is shown in red, while the chroma model is shown in grey\n", "\n", "\n", "\n", "import numpy as np\n", "import warnings\n", "\n", "warnings.simplefilter(action='ignore', category=FutureWarning)\n", "\n", "import os, re\n", "from colabdesign import mk_afdesign_model, clear_mem\n", "from colabdesign.mpnn import mk_mpnn_model\n", "\n", "import os, re\n", "from colabdesign import mk_afdesign_model, clear_mem\n", "from colabdesign.mpnn import mk_mpnn_model\n", "\n", "import numpy as np\n", "\n", "def sample_gumbel(shape, eps=1e-20):\n", " \"\"\"Sample from Gumbel(0, 1)\"\"\"\n", " U = np.random.uniform(size=shape)\n", " return -np.log(-np.log(U + eps) + eps)\n", "\n", "\n", "clear_mem()\n", "\n", "iters = iterations\n", "in_path = '/content/in/'\n", "out_path = 'out/'\n", "if not os.path.exists('/content/in/out/'):\n", " os.mkdir('/content/in/out/')\n", "\n", "\n", "starting_seq = \"\"\n", "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n", "\n", "\n", "file_list = os.listdir(in_path)\n", "\n", "clear_mem()\n", "mpnn_model = mk_mpnn_model()\n", "af_model = mk_afdesign_model(protocol=\"fixbb\",use_templates=True,use_initial_atom_pos=True,use_initial_guess=True)\n", "\n", "for file_name in file_list:\n", " if file_name[-1] == 'b':\n", "\n", " in_pdb = in_path + file_name\n", " out_pdb = in_path + out_path + 'Out_' + file_name\n", "\n", "\n", "\n", "\n", " af_model.prep_inputs(pdb_filename=in_pdb, chain='A')\n", "\n", " starting_seq = af_model._inputs['batch'][\"aatype\"]\n", "\n", " iterations = iters\n", "\n", " use_dropout = True\n", " num_recycles = 0\n", "\n", " mpnn_mode = \"conditional\"\n", " cmap_seqsep = 9\n", "\n", " cmap_num = 2\n", " #cmap_cutoff = 14\n", " L = sum(af_model._lengths)\n", " af_model.restart(mode=\"gumbel\")\n", " af_model._args[\"clear_prev\"] = False\n", " #af_model.set_opt(cmap_cutoff=cmap_cutoff)\n", " af_model.set_weights(helix=1e-8)\n", " # gather info about inputs\n", " if \"offset\" in af_model._inputs:\n", " offset = af_model._inputs\n", " else:\n", " idx = af_model._inputs[\"residue_index\"]\n", " offset = idx[:, None] - idx[None, :]\n", " # initialize sequence\n", " if len(starting_seq) > 1:\n", " af_model.set_seq(seq=starting_seq)\n", " # initialize coordinates\n", " af_model._inputs.pop(\"prev\", None)\n", " init = af_model._inputs[\"batch\"]['all_atom_positions'].copy()\n", "\n", " save_best = False\n", " for k in range(iterations):\n", "\n", " if k > (iterations - 10):\n", " use_dropout = False\n", " save_best = True\n", "\n", " # denoise\n", " aux = af_model.predict(return_aux=True, verbose=False,\n", " dropout=use_dropout,\n", " num_recycles=num_recycles)\n", "\n", " #af_model._inputs[\"prev\"] = aux[\"prev\"]\n", " #af_model._inputs[\"prev\"][\"prev_msa_first_row\"] *= 0\n", " #af_model._inputs[\"prev\"][\"prev_pos\"] *= 0\n", "\n", " cmap = aux[\"cmap\"] * (np.abs(offset) > cmap_seqsep)\n", " conf = np.sort(cmap)[:, -cmap_num:].mean(-1)\n", "\n", " plddt = aux[\"plddt\"]\n", " seq = aux[\"seq\"][\"hard\"][0].argmax(-1)\n", " xyz = aux[\"atom_positions\"].copy()\n", " # update inputs\n", " af_model._inputs[\"batch\"][\"aatype\"] = seq\n", " af_model._inputs[\"batch\"][\"all_atom_positions\"] = xyz\n", "\n", " if mpnn_mode != \"none\":\n", "\n", " mpnn_model.get_af_inputs(af_model)\n", " opt = {\"mask\": np.sqrt(conf)}\n", " if mpnn_mode == \"unconditional\":\n", " opt[\"ar_mask\"] = np.zeros((L, L))\n", " mpnn_out = mpnn_model.score(**opt)\n", " mpnn_logits = mpnn_out[\"logits\"][:, :20]\n", " aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"]\n", "\n", " c = conf[:, None]\n", "\n", " new_logits = (1 - c) * sample_gumbel(mpnn_logits.shape) + c * mpnn_logits\n", "\n", " af_model._params[\"seq\"] = 0.9 * af_model._params[\"seq\"] + 0.1 * new_logits\n", "\n", " # save results\n", " af_model._save_results(aux, save_best=save_best)\n", " af_model._k += 1\n", "\n", " af_model.save_pdb(out_pdb)\n", " visualize_pdb_overlay(in_pdb, out_pdb)\n", "\n" ], "metadata": { "cellView": "form", "id": "KeWmENFFqf1S" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "HTML(af_model.animate())" ], "metadata": { "id": "HtpnTpvZtHi6" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown #Redesign with solubleMPNN for ESMFold prediction\n", "#@markdown The standard manuscript settings were 8 sequences, 0.1 sampling temperature and the removal of cysteines\n", "import pickle\n", "import pandas as pd\n", "in_path = '/content/in/out/'\n", "out_path = 'out_sMPNN/'\n", "if not os.path.exists(out_path):\n", " os.mkdir(out_path)\n", "\n", "file_list = os.listdir(in_path)\n", "\n", "\n", "num_seqs = 8 #@param [\"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n", "mpnn_sampling_temp = 0.1 #@param [\"0.0001\", \"0.1\", \"0.15\", \"0.2\", \"0.25\", \"0.3\", \"0.5\", \"1.0\"] {type:\"raw\"}\n", "rm_aa = \"C\" #@param {type:\"string\"}\n", "use_solubleMPNN = True #@param {type:\"boolean\"}\n", "#@markdown - `mpnn_sampling_temp` - control diversity of sampled sequences. (higher = more diverse).\n", "#@markdown - `rm_aa='C'` - do not use [C]ysteines.\n", "#@markdown - `use_solubleMPNN` - use weights trained only on soluble proteins.\n", "#@markdown\n", "\n", "from colabdesign.shared.protein import alphabet_list as chain_list\n", "mpnn_model = mk_mpnn_model()\n", "\n", "\n", "for file in file_list:\n", " if file[-4:] == '.pdb':\n", "\n", " in_file1 = in_path + file\n", " mpnn_model.prep_inputs(pdb_filename=in_file1,\n", " chain='A',\n", " rm_aa=rm_aa,weights = \"soluble\")\n", " out = mpnn_model.sample(num=num_seqs//8,\n", " batch=8,\n", " temperature=mpnn_sampling_temp)\n", " for seq,score in zip(out[\"seq\"],out[\"score\"]):\n", " print(score,seq.split(\"/\")[0])\n", " df = pd.DataFrame(out[\"seq\"])\n", "\n", " # Define the output path for saving the sequences as a .pkl file\n", " output_pkl_file = out_path + file[:-4] + \"_sequences.pkl\"\n", "\n", " # Save the DataFrame to a .pkl file\n", " with open(output_pkl_file, 'wb') as f:\n", " pickle.dump(df, f)" ], "metadata": { "cellView": "form", "id": "KNcTvB5XzG-X" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown #Run ESMFold to test designability\n", "\n", "#@markdown **This cell is a little bit tricky. The problem is the compatibility between the JAX and PYTORCH frameworks between ColabDesign and ESMFold and GPU memory requirements**\n", "\n", "#@markdown Go to *Runtime* >> *Restart Session*\n", "\n", "#@markdown then Run this cell\n", "\n", "#@markdown Runtime 2-5 min\n", "\n", "#@markdown This cells runs ESMFold from huggingface and automatically calculates the RMSD to the designed backbone\n", "#@markdown NOTE: GPU memory can be a big problem here. If you get memory errors please restart the runtime and run this cell again. It should be self contained. Additionally, after finish the ESMFold prediction rerun the setup cell\n", "\n", "import os\n", "import pandas as pd\n", "from Bio.PDB import PDBParser, Superimposer\n", "import pickle\n", "import torch\n", "import numpy as np\n", "from transformers import AutoTokenizer, EsmForProteinFolding\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", "import py3Dmol\n", "\n", "def visualize_pdb_overlay(pdb1_path, pdb2_path):\n", " viewer = py3Dmol.view(width=800, height=600)\n", "\n", " with open(pdb1_path, 'r') as f:\n", " pdb1_data = f.read()\n", " viewer.addModel(pdb1_data, 'pdb')\n", " viewer.setStyle({'model': 0}, {'cartoon': {'color': 'grey'}})\n", "\n", " with open(pdb2_path, 'r') as f:\n", " pdb2_data = f.read()\n", " viewer.addModel(pdb2_data, 'pdb')\n", " viewer.setStyle({'model': 1}, {'cartoon': {'color' : 'red'}})\n", "\n", " viewer.zoomTo()\n", " viewer.show()\n", "\n", "\n", "\n", "\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esmfold_v1\")\n", "model = EsmForProteinFolding.from_pretrained(\"facebook/esmfold_v1\", low_cpu_mem_usage=True)\n", "\n", "device = 'cuda:0'\n", "model = model.cuda(device)\n", "model.esm = model.esm.half()\n", "model.trunk.set_chunk_size(64)\n", "torch.backends.cuda.matmul.allow_tf32 = True\n", "\n", "def convert_outputs_to_pdb(outputs):\n", " final_atom_positions = atom14_to_atom37(outputs[\"positions\"][-1], outputs)\n", " outputs = {k: v.to(\"cpu\").numpy() for k, v in outputs.items()}\n", " final_atom_positions = final_atom_positions.cpu().numpy()\n", " final_atom_mask = outputs[\"atom37_atom_exists\"]\n", " pdbs = []\n", " for i in range(outputs[\"aatype\"].shape[0]):\n", " aa = outputs[\"aatype\"][i]\n", " pred_pos = final_atom_positions[i]\n", " mask = final_atom_mask[i]\n", " resid = outputs[\"residue_index\"][i] + 1\n", " pred = OFProtein(\n", " aatype=aa,\n", " atom_positions=pred_pos,\n", " atom_mask=mask,\n", " residue_index=resid,\n", " b_factors=outputs[\"plddt\"][i],\n", " chain_index=outputs[\"chain_index\"][i] if \"chain_index\" in outputs else None,\n", " )\n", " pdbs.append(to_pdb(pred))\n", " return pdbs\n", "\n", "def calculate_ca_rmsd(pdb_file1, pdb_file2):\n", " parser = PDBParser(QUIET=True)\n", "\n", " structure1 = parser.get_structure(\"Protein1\", pdb_file1)\n", " structure2 = parser.get_structure(\"Protein2\", pdb_file2)\n", "\n", " ca_atoms1 = [atom for atom in structure1.get_atoms() if atom.get_name() == \"CA\"]\n", " ca_atoms2 = [atom for atom in structure2.get_atoms() if atom.get_name() == \"CA\"]\n", "\n", " super_imposer = Superimposer()\n", " super_imposer.set_atoms(ca_atoms1, ca_atoms2)\n", " super_imposer.apply(structure2.get_atoms())\n", " rmsd = super_imposer.rms\n", " return rmsd\n", "\n", "def process_sequences(seq_list, pdb_file,in_path_pdb):\n", " lowest_rmsd = float('inf')\n", " lowest_rmsd_data = None\n", " pdb_id = pdb_file[16:-4]\n", " out_ss_path = in_path_pdb + \"output/\"\n", "\n", " if not os.path.exists(out_ss_path):\n", " os.mkdir(out_ss_path)\n", "\n", " for test_protein in seq_list:\n", " data = {}\n", " tokenized_input = tokenizer([test_protein], return_tensors=\"pt\", add_special_tokens=False)['input_ids']\n", " tokenized_input = tokenized_input.cuda(device)\n", "\n", " with torch.no_grad():\n", " output = model(tokenized_input)\n", "\n", " data['out'] = output\n", " data[\"plddt\"] = torch.mean(output['plddt']).item()\n", " data['pae'] = torch.mean(output['predicted_aligned_error']).item()\n", "\n", " pdb_data = convert_outputs_to_pdb(output)\n", " tmp_pdb_file = os.path.join(out_ss_path, \"TMP.pdb\")\n", "\n", " with open(tmp_pdb_file, 'w') as file:\n", " for line in pdb_data:\n", " file.write(line)\n", "\n", " data['rmsd'] = calculate_ca_rmsd(tmp_pdb_file, pdb_file)\n", " print(f'Sequence: {test_protein}, plddt: {data[\"plddt\"]}, PAE: {data[\"pae\"]}, RMSD: {data[\"rmsd\"]}')\n", "\n", " if data['rmsd'] < lowest_rmsd:\n", " lowest_rmsd = data['rmsd']\n", " lowest_rmsd_data = data\n", "\n", " if lowest_rmsd_data is not None:\n", " print(f'Lowest RMSD: {lowest_rmsd}')\n", " best_pdb_data = convert_outputs_to_pdb(lowest_rmsd_data['out'])\n", " best_pdb_file = os.path.join(out_ss_path, f\"{pdb_id}_best_structure.pdb\")\n", "\n", " with open(best_pdb_file, 'w') as file:\n", " for line in best_pdb_data:\n", " file.write(line)\n", "\n", " original_dict = lowest_rmsd_data\n", " key_to_exclude = 'out'\n", " data_out = {k: v for k, v in original_dict.items() if k != key_to_exclude}\n", "\n", " with open(os.path.join(out_ss_path, f\"{pdb_id}_best_structure_data.pkl\"), 'wb') as f:\n", " pickle.dump(data_out, f)\n", "\n", " return lowest_rmsd, best_pdb_file, data_out\n", "\n", " return None, None, None\n", "\n", "\n", "in_path= '/content/out_sMPNN/'\n", "file_list = os.listdir(in_path)\n", "\n", "for file in file_list:\n", " if file[-1] =='l':\n", "\n", "\n", "\n", " output_pkl_file = in_path + file\n", "\n", " with open(output_pkl_file, 'rb') as f:\n", " seq = pickle.load(f)\n", " seq_list = []\n", " for i in np.asarray(seq):\n", " seq_list.append(i[0])\n", "\n", " pdb_file = \"/content/in/out/\" + file[:-14] + '.pdb'\n", " print(seq_list)\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", " lowest_rmsd, best_pdb_file, best_data = process_sequences(seq_list, pdb_file,in_path)\n", " if lowest_rmsd is not None:\n", " print(f\"Lowest RMSD: {lowest_rmsd}, Best PDB file: {best_pdb_file}\")\n", " else:\n", " print(\"No valid result found.\")\n", " #visualize_pdb_overlay(pdb_file, best_pdb_file)\n" ], "metadata": { "cellView": "form", "collapsed": true, "id": "Ctyijhvu0UVe" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown #Optional: Run AF2 based designability test\n", "\n", "\n", "#@markdown Rerun **setup cell** if you tested ESMFold prediction before!\n", "\n", "#@markdown This cell predicts the solubleMPNN generated seuqneces with AF2 with Initial Guess & All atom initialisation\n", "import re\n", "clear_mem()\n", "\n", "in_path = '/content/in/'\n", "out_path = 'out/'\n", "if not os.path.exists('/content/in/out/'):\n", " os.mkdir('/content/in/out/')\n", "\n", "\n", "starting_seq = \"\"\n", "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n", "\n", "\n", "file_list = os.listdir(in_path)\n", "\n", "clear_mem()\n", "\n", "af_model = mk_afdesign_model(protocol=\"fixbb\",use_initial_atom_pos=True,use_initial_guess=True)\n", "\n", "\n", "def process_sequences(seq_list, pdb_file,in_path_pdb):\n", " lowest_rmsd = float('inf')\n", " lowest_rmsd_data = None\n", " pdb_id = pdb_file[16:-4]\n", " out_ss_path = in_path_pdb + \"output_AF2/\"\n", "\n", " if not os.path.exists(out_ss_path):\n", " os.mkdir(out_ss_path)\n", " kk=0\n", " for test_protein in seq_list:\n", " data = {}\n", " af_model.prep_inputs(pdb_filename=pdb_file, chain='A')\n", " af_model.predict(seq=test_protein,num_recycles=3)\n", "\n", "\n", "\n", " data[\"plddt\"] = af_model.aux['losses']['plddt']\n", " data['pae'] = af_model.aux['losses']['pae']*31\n", " data['rmsd'] = af_model.aux['losses']['rmsd']\n", " af_model.save_pdb(f'{out_ss_path}{pdb_id}_{kk}.pdb')\n", " print(f'Sequence: {test_protein}, plddt: {data[\"plddt\"]}, PAE: {data[\"pae\"]}, RMSD: {data[\"rmsd\"]}')\n", "\n", " if data['rmsd'] < lowest_rmsd:\n", " lowest_rmsd = data['rmsd']\n", " lowest_rmsd_data = data\n", "\n", " if lowest_rmsd_data is not None:\n", " print(f'Lowest RMSD: {lowest_rmsd}')\n", "\n", " best_pdb_file = os.path.join(out_ss_path, f\"{pdb_id}_best_structure.pdb\")\n", "\n", " return lowest_rmsd, best_pdb_file\n", "\n", " return None, None, None\n", "\n", "\n", "in_path= '/content/out_sMPNN/'\n", "file_list = os.listdir(in_path)\n", "\n", "for file in file_list:\n", " if file[-1] =='l':\n", "\n", "\n", "\n", " output_pkl_file = in_path + file\n", "\n", " with open(output_pkl_file, 'rb') as f:\n", " seq = pickle.load(f)\n", " seq_list = []\n", " for i in np.asarray(seq):\n", " seq_list.append(i[0])\n", "\n", " pdb_file = \"/content/in/out/\" + file[:-14] + '.pdb'\n", " print(seq_list)\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", " lowest_rmsd, best_pdb_file = process_sequences(seq_list, pdb_file,in_path)\n", " if lowest_rmsd is not None:\n", " print(f\"Lowest RMSD: {lowest_rmsd}, Best PDB file: {best_pdb_file}\")\n", " else:\n", " print(\"No valid result found.\")\n", " #visualize_pdb_overlay(pdb_file, best_pdb_file)\n", "\n", "\n" ], "metadata": { "cellView": "form", "id": "s1EVl7Ee8ne-" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [], "metadata": { "id": "ZQljfnQl8oTc" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [], "metadata": { "id": "16XTY0u9YyxP" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [], "metadata": { "id": "u6mCraHrdqM3" }, "execution_count": null, "outputs": [] } ] }