{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "OA2k3sAYuiXe" }, "source": [ "\n", "\n", "\n", "#Protein Design using Relaxed Sequence Optimization\n", "\n", "\n", "**Scalable protein design using optimization in a relaxed sequence space**\n", "\n", "\n", "\n", "\n", "Christopher Frank, Ali Khoshouei, Lara Fuß, Lara Weber Dominik Schiewitz,Zhixuan Zhao, Motoyuki Hattori, Yosta de Stigter, Shihao Feng, Sergey Ovchinnikov and Hendrik Dietz\n", "\n", "\n", "This notebook contains code to run relaxed sequence optimisation for de novo protein design as described in the manuscript. There are additional options to modify the pipeline according to ones needs\n", "\n", "We recommend using at least an L4 GPU to run this notebook, as the free T4 GPU struggles with larger proteins\n", "\n", "Alternativly a local installation of ColabDesign is strongly recommendet, especially for the design of larger proteins.\n", "\n", "For questions feel free to reach out to the authors\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "-AXy0s_4cKaK" }, "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" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "sZnYfCbfEvol", "cellView": "form" }, "outputs": [], "source": [ "#@title # Unconditional Generation (Custom)\n", "#@markdown For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n", "LENGTH = 100 #@param {type:\"integer\"}\n", "#@markdown With copies you can specify the number of identical sequences design, resulting in homo oligomers. Copies = 1 is the standard, resulting in a monomer\n", "\n", "COPIES = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\"] {type:\"raw\"}\n", "MODE = \"manuscript\"\n", "\n", "#@markdown Select the losses you want to use. For unconditional generation as reported in the manuscript use all the losses. To increase the diversity of designes remove confidence losses and/or increase the weight of the helix loss.\n", "\n", "use_rg_loss = True #@param {type:\"boolean\"}\n", "#@markdown A too strong rg loss can lead to problems and clashes. Use 0.1 for backbones smaller then 600 AA and 0.01 for larger proteins (0.001 for 1000 AA).\n", "rg_weight = 0.1 #@param {type:\"raw\"}\n", "use_helix_loss = True #@param {type:\"boolean\"}\n", "use_con_loss = True #@param {type:\"boolean\"}\n", "use_confidence_loss = True #@param {type:\"boolean\"}\n", "#@markdown How many halluicnation iteration you want to perform. The standard in the manuscript is 100.\n", "\n", "iters = 50 #@param [\"100\", \"50\", \"30\"] {type:\"raw\"}\n", "\n", "\n", "#@markdown Select if you want to use the 'standard\" ProteinMPNN weights or the soluble ones. The soluble ones usually result in higher in silico as well as experimental sucess, but will increase the negative net charge of the protein which sould potentially interfer with certain protein design problems. The manuscript settings are soluble MPNN\n", "\n", "use_solubleMPNN = True #@param {type:\"boolean\"}\n", "#@markdown Select this to use an experimental ProteinMPNN loss, also backpropagating through ProteinMPNN. This was not used in the manuscript\n", "\n", "use_mpnn_loss = False #@param {type:\"boolean\"}\n", "#@markdown\n", "\n", "def add_rg_loss(self, weight=0.1):\n", " '''add radius of gyration loss'''\n", " def loss_fn(inputs, outputs):\n", " xyz = outputs[\"structure_module\"]\n", " ca = xyz[\"final_atom_positions\"][:,residue_constants.atom_order[\"CA\"]]\n", " if self.protocol == \"binder\":\n", " ca = ca[-self._binder_len:]\n", " #This uses a scaled version of the rg loss, only looking at every 5th residue\n", " if MODE == \"manuscript\":\n", " ca = ca[::5]\n", " rg = jnp.sqrt(jnp.square(ca - ca.mean(0)).sum(-1).mean() + 1e-8)\n", "\n", " if MODE == \"original\":\n", " rg_th = 2.38 * ca.shape[0] ** 0.365\n", " rg = jax.nn.elu(rg - rg_th)\n", " return {\"rg\":rg}\n", " self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n", " self.opt[\"weights\"][\"rg\"] = weight\n", "\n", "def add_mpnn_loss(self, mpnn=0.1, mpnn_seq=0.0):\n", " '''\n", " add mpnn loss\n", " mpnn = maximize confidence of proteinmpnn\n", " mpnn_seq = push designed sequence to match proteinmpnn logits\n", " '''\n", "\n", " self._mpnn = mk_mpnn_model(weights = \"soluble\" if use_solubleMPNN else \"original\")\n", " def loss_fn(inputs, outputs, aux, key):\n", "\n", " # get structure\n", " atom_idx = tuple(residue_constants.atom_order[k] for k in [\"N\",\"CA\",\"C\",\"O\"])\n", " I = {\"S\": inputs[\"aatype\"],\n", " \"residue_idx\": inputs[\"residue_index\"],\n", " \"chain_idx\": inputs[\"asym_id\"],\n", " \"X\": outputs[\"structure_module\"][\"final_atom_positions\"][:,atom_idx],\n", " \"mask\": outputs[\"structure_module\"][\"final_atom_mask\"][:,1],\n", " \"lengths\": self._lengths,\n", " \"key\": key}\n", "\n", " if \"offset\" in inputs:\n", " I[\"offset\"] = inputs[\"offset\"]\n", "\n", " # set autoregressive mask\n", " L = sum(self._lengths)\n", " if self.protocol == \"binder\":\n", " I[\"ar_mask\"] = 1 - np.eye(L)\n", " I[\"ar_mask\"][-self._len:,-self._len:] = 0\n", " else:\n", " I[\"ar_mask\"] = np.zeros((L,L))\n", "\n", " # get logits\n", " logits = self._mpnn._score(**I)[\"logits\"][:,:20]\n", " if self.protocol == \"binder\":\n", " logits = logits[-self._len:]\n", " else:\n", " logits = logits[:self._len]\n", " aux[\"mpnn_logits\"] = logits\n", "\n", " # compute loss\n", " log_q = jax.nn.log_softmax(logits)\n", " p = inputs[\"seq\"][\"hard\"]\n", " q = jax.nn.softmax(logits)\n", " losses = {}\n", " losses[\"mpnn\"] = -log_q.max(-1).mean()\n", " losses[\"mpnn_seq\"] = -(p * jax.lax.stop_gradient(log_q)).sum(-1).mean()\n", " return losses\n", "\n", " self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n", " self.opt[\"weights\"][\"mpnn\"] = mpnn\n", " self.opt[\"weights\"][\"mpnn_seq\"] = mpnn_seq\n", "\n", "clear_mem()\n", "af_model = mk_afdesign_model(protocol=\"hallucination\")\n", "af_model.prep_inputs(length=LENGTH, copies=COPIES)\n", "\n", "# add extra losses\n", "\n", "if use_mpnn_loss: add_mpnn_loss(af_model)\n", "\n", "print(\"length\",af_model._lengths)\n", "print(\"weights\",af_model.opt[\"weights\"])" ] }, { "cell_type": "code", "source": [ "#This cell runs the design loop. Run this in a for loop for design of multiple proteins\n", "\n", "af_model.restart()\n", "af_model.set_seq(mode=[\"gumbel\",\"soft\"])\n", "if use_rg_loss: add_rg_loss(af_model,rg_weight)\n", "if use_helix_loss : af_model.set_weights(helix=-0.2)\n", "if use_con_loss : af_model.set_weights(con=1.0)\n", "if use_confidence_loss : af_model.set_weights(plddt=0.5, pae=0.5)\n", "print(\"weights\",af_model.opt[\"weights\"])\n", "af_model.design_logits(iters-10)\n", "af_model.design_logits(10, save_best=True)" ], "metadata": { "id": "f76xqCkw0vj9" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "A1GxeLZdTTya" }, "outputs": [], "source": [ "#This cell plots and saves the results as a pdb file\n", "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n", "af_model.plot_pdb()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "L2E9Tn2Acchj" }, "outputs": [], "source": [ "HTML(af_model.animate())" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "YSKWYu0_GlUH" }, "outputs": [], "source": [ "af_model.get_seqs()" ] }, { "cell_type": "code", "source": [ "import pandas as pd\n", "#@title # Designability test\n", "#@markdown Test the designability of the backbone, taking in the backbone, generating sequences with solubleMPNN and predicting the sequence with AF2 in single sequence mode.\n", "#@markdown Use Initial Guess (IG) and All Atom Initialisation (AA) for larger proteins\n", "\n", "AA = False #@param {type:\"boolean\"}\n", "IG = False #@param {type:\"boolean\"}\n", "#@markdown NOTE: we remove cysteines from all designed proteins. Additionally for large proteins we also exclude methions to reduce the number of internal start codons\n", "\n", "def designability_test(af_model_test, mpnn_model_test,\n", " num_seqs=8, sampling_temp=0.1, num_recycles=3,\n", " model_num=4, best_metric=\"rmsd\",\n", " in_pdb=\"init.pdb\", out_pdb=\"final.pdb\",\n", " verbose=False):\n", " alphafold_model = f\"model_{model_num}_ptm\"\n", "\n", " af_model_test.prep_inputs(in_pdb)\n", " af_model_test.restart(rm_aa=\"C,M\")\n", " af_model_test._args[\"best_metric\"] = best_metric\n", " L = sum(af_model_test._lengths)\n", " mpnn_model_test.get_af_inputs(af_model_test)\n", " out = mpnn_model_test.sample(num=num_seqs // 8, batch=8,\n", " temperature=sampling_temp)\n", "\n", " af_terms = [\"plddt\", \"ptm\", \"pae\", \"rmsd\", \"dgram_cce\"]\n", " for k in af_terms: out[k] = []\n", "\n", " for n in range(num_seqs):\n", " seq = out[\"seq\"][n]\n", " af_model_test.predict(seq=seq,\n", " num_recycles=num_recycles,\n", " num_models=1,\n", " verbose=False,\n", " models=alphafold_model)\n", "\n", " for k in af_terms: out[k].append(af_model_test.aux[\"log\"][k])\n", " out[\"pae\"][-1] = out[\"pae\"][-1] * 31\n", " af_model_test._save_results(save_best=True, verbose=verbose)\n", " af_model_test._k += 1\n", "\n", " af_model_test.save_pdb(out_pdb)\n", " labels = [\"score\"] + af_terms + [\"seq\"]\n", " data = [[out[k][n] for k in labels] for n in range(num_seqs)]\n", " labels[0] = \"mpnn\"\n", " df = pd.DataFrame(data, columns=labels)\n", " return df\n", "\n", "af_model_test = mk_afdesign_model(protocol=\"fixbb\",best_metric=\"rmsd\",use_initial_guess=IG,use_initial_atom_pos=AA,use_templates=False)\n", "mpnn_model_test = mk_mpnn_model(weights=\"soluble\")\n", "\n", "\n", "lowest_rmsd = float('inf')\n", "lowest_rmsd_data = None\n", "\n", "in_pdb = f\"{af_model.protocol}.pdb\"\n", "out_pdb = f\"{af_model.protocol}_out.pdb\"\n", "\n", "\n", "out = designability_test(af_model_test, mpnn_model_test,\n", " num_seqs=8, sampling_temp=0.1, num_recycles=3,\n", " model_num=4, best_metric=\"rmsd\",\n", " in_pdb=in_pdb, out_pdb=out_pdb,\n", " verbose=True)\n", "\n" ], "metadata": { "cellView": "form", "id": "Y_-8_Ya4Q8oo" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "from colabdesign import mk_afdesign_model, clear_mem\n", "from colabdesign.af.alphafold.common import residue_constants\n", "import jax\n", "import jax.numpy as jnp\n", "#@title # OPTIONAL Unconditional Generation (Manuscript Code)\n", "\n", "#@markdown This code generates a sample of 10 unconditional proteins for lengths between 100 and 800 AA exactly as in the manuscript. For larger proteins CUDA_UNIFIED_MEMORY is needed. This can be done by localy running the code on a CUDA capeable GPU with sufficient memory (A100 80GB e.g.) and running the code with the environment variables XLA_PYTHON_CLIENT_MEM_FRACTION=100.0 TF_FORCE_UNIFIED_MEMORY=1\n", "def rg_loss(inputs, outputs):\n", " positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n", " ca = positions[::5,residue_constants.atom_order[\"CA\"]]\n", " center = ca.mean(0)\n", " rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n", " rg_th = 2.38 * ca.shape[0] ** 0.365\n", " rg = jax.nn.elu(rg - rg_th)\n", " return {\"rg\":rg}\n", "\n", "\n", "\n", "\n", "for length in [100,200,300,400,500,600,700,800]:\n", " model = mk_afdesign_model(protocol=\"hallucination\",loss_callback=rg_loss)\n", " model.prep_inputs(length=length)\n", " print(\"weights\",model.opt[\"weights\"])\n", " print('Starting up and compiling JAX model....')\n", "\n", " for i in range(10):\n", " model.restart(mode=[\"gumbel\", \"soft\"],rm_aa=\"C\")\n", " model.opt[\"weights\"][\"rg\"] = 0.1\n", " if length > 600:\n", " model.opt[\"weights\"][\"rg\"] = 0.01\n", " #model.opt[\"weights\"]['helix'] = -0.1\n", " model.opt[\"weights\"]['plddt'] = 1.0\n", " model.opt[\"weights\"]['pae'] = 1.0\n", " model.opt[\"weights\"]['helix'] = -0.1\n", " print(\"weights\", model.opt[\"weights\"])\n", " model.design_logits(100)\n", "\n", " #change the output path for local execution\n", " model.save_pdb(f\"Hallo_{i}.pdb\")" ], "metadata": { "cellView": "form", "id": "gPby7oP25b5i" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown #Redesign with ProteinMPNN 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", "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 = False #@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", "mpnn_model.prep_inputs(pdb_filename=f\"{af_model.protocol}.pdb\",\n", " chain=\",\".join(chain_list[:COPIES]),\n", " homooligmer=COPIES>1,\n", " rm_aa=rm_aa,\n", " weights = \"soluble\" if use_solubleMPNN else\"original\")\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 = \"redesigned_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": { "id": "m2qAYsDsCfqJ", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown #Run ESMFold to test designability\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", "output_pkl_file = \"redesigned_sequences.pkl\"\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 = \"hallucination.pdb\"\n", "print(seq_list)\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):\n", " lowest_rmsd = float('inf')\n", " lowest_rmsd_data = None\n", " out_ss_path = \"./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, \"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, \"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", "\n", "lowest_rmsd, best_pdb_file, best_data = process_sequences(seq_list, pdb_file)\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" ], "metadata": { "id": "Ey29NmNAFtK0", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title # Heterodimer Design Prep\n", "#@markdown Design a set of heterodimeric proteins with two chains making a complex. The settings are excatly the ones used in the manuscript to design the heterodimer binders.\n", "LENGTH1 = 100 #@param {type:\"integer\"}\n", "LENGTH2 = 100 #@param {type:\"integer\"}\n", "\n", "#@markdown ProteinMPNN Settings\n", "use_solubleMPNN = True #@param {type:\"boolean\"}\n", "#@markdown\n", "\n", "\n", "from colabdesign.af.alphafold.common import residue_constants\n", "import jax\n", "import jax.numpy as jnp\n", "\n", "def hd_loss(inputs, outputs):\n", " positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n", " ca1 = positions[:LENGTH1, residue_constants.atom_order[\"CA\"]]\n", " center1 = ca1.mean(0)\n", " rg1 = jnp.sqrt(jnp.square(ca1 - center1).sum(-1).mean() + 1e-8)\n", " rg_th = 2.38 * ca1.shape[0] ** 0.365\n", " rg1 = jax.nn.elu(rg1 - rg_th)\n", "\n", "\n", " ca2 = positions[LENGTH2:, residue_constants.atom_order[\"CA\"]]\n", " center2 = ca2.mean(0)\n", " rg2 = jnp.sqrt(jnp.square(ca2 - center2).sum(-1).mean() + 1e-8)\n", " rg_th = 2.38 * ca2.shape[0] ** 0.365\n", " rg2 = jax.nn.elu(rg2 - rg_th)\n", "\n", "\n", "\n", " return {\"hd\":rg1+rg2}\n", "\n", "total_length = LENGTH1 + LENGTH2\n", "clear_mem()\n", "af_model = mk_afdesign_model(protocol=\"hallucination\", loss_callback=hd_loss)\n", "af_model.prep_inputs(length=total_length)\n", "af_model._inputs['residue_index'][LENGTH1:] = np.arange(LENGTH2) + 50 + LENGTH1\n", "# add extra losses\n", "af_model.restart(mode=[\"gumbel\", \"soft\"])\n", "af_model.opt[\"weights\"][\"hd\"] = 0.1\n", "af_model.opt[\"weights\"]['plddt'] = 1.0\n", "af_model.opt[\"weights\"]['pae'] = 1.0\n", "af_model.opt[\"weights\"]['helix'] = -0.5\n", "print(\"weights\", af_model.opt[\"weights\"])\n", "print('Starting up and compiling JAX model....')\n" ], "metadata": { "id": "88HCs0ID6CGc", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title # Run Design\n", "af_model.design_logits(100)\n", "af_model.save_pdb(\"Heterodimer.pdb\")" ], "metadata": { "id": "5KVtj6LM7RQ4" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "af_model.save_pdb(\"Heterodimer.pdb\")\n", "af_model.plot_pdb()" ], "metadata": { "id": "qfOEMo4qpWfm" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title # Design Sequence using Homooligomer Filter\n", "#@markdown We first test if the two protomers are predictd to fold into a high confidence protein on their own, removing proteins that are not likely to be expressed on their own. Then we predict the heterodimer using the AF multimer model. Generally the AF multimer model has a hard time predicting de novo designed proteins. This is why we use templates and remove any interchain information. Finally we predict each individual protomer with a copy of itself, testing for homooligomerisation.\n", "\n", "\n", "file_path =\"Heterodimer.pdb\"\n", "\n", "folder_path = \"/content/\"\n", "\n", "######## make A - B chain file\n", "\n", "from Bio.PDB import PDBParser, PDBIO, Chain\n", "\n", "# Set the input and output PDB file names\n", "input_pdb_file = file_path\n", "if not os.path.exists(os.path.join(folder_path, 'AB')):\n", " os.mkdir(os.path.join(folder_path, 'AB'))\n", "output_pdb_file = os.path.join(folder_path, 'AB',\"Heterodimer.pdb\")\n", "\n", "# Create a PDB parser and read the input PDB file\n", "parser = PDBParser()\n", "structure = parser.get_structure(\"input_structure\", input_pdb_file)\n", "\n", "# Find the initial chain id\n", "initial_chain_id = None\n", "for chain in structure[0]:\n", " initial_chain_id = chain.get_id()\n", " break\n", "\n", "# Create new chains A and B\n", "chain_A = Chain.Chain(\"A\")\n", "chain_B = Chain.Chain(\"B\")\n", "\n", "# Iterate over the residues in the original chain\n", "for residue in structure[0][initial_chain_id]:\n", " res_id = residue.get_id()[1]\n", "\n", " # Add residues 1-200 to chain A\n", " if 1 <= res_id <= 100:\n", " chain_A.add(residue.copy())\n", "\n", " # Add residues 201-400 to chain B\n", " elif 151 <= res_id <= 450:\n", " chain_B.add(residue.copy())\n", "\n", "# Remove the existing chain\n", "for model in structure:\n", " model.detach_child(initial_chain_id)\n", "\n", "# Add the new chains to the model\n", "structure[0].add(chain_A)\n", "structure[0].add(chain_B)\n", "\n", "# Save the modified structure to a new PDB file\n", "io = PDBIO()\n", "io.set_structure(structure)\n", "io.save(output_pdb_file)\n", "\n", "\n", "clear_mem()\n", "he_model = mk_afdesign_model(protocol=\"fixbb\", use_templates=True, use_multimer=True)\n", "\n", "\n", "ho_model = mk_afdesign_model(protocol=\"hallucination\")\n", "ho_model.prep_inputs(length=LENGTH1, copies=2)\n", "\n", "ho_model.set_weights(i_pae=1.0)\n", "s_model = mk_afdesign_model(protocol=\"hallucination\")\n", "s_model.prep_inputs(length=LENGTH2)\n", "mpnn_model = mk_mpnn_model(weights=\"soluble\")\n", "\n", "\n", "mpnn_model.prep_inputs(pdb_filename=output_pdb_file, chain='A,B',rm_aa=\"C\")\n", "samples = mpnn_model.sample_parallel(8)\n", "\n", "he_model.prep_inputs(pdb_filename=output_pdb_file, chain='A,B',rm_template_ic=True)\n", "he_model._inputs['residue_index'][LENGTH1:] = np.arange(LENGTH2) + 50 + LENGTH1\n", "\n", "k = 0\n", "for seq in samples['seq']:\n", " print('Predicting Protomer 1...')\n", " s_model.predict(seq=seq[:LENGTH1], num_recycles=3)\n", " plddt1 = s_model.aux['losses']['plddt']\n", " print('Predicting Protomer 2...')\n", " s_model.predict(seq=seq[LENGTH1+1:], num_recycles=3)\n", " plddt2 = s_model.aux['losses']['plddt']\n", " k = k + 1\n", " if plddt1 < 0.20 and plddt2 < 0.20:\n", " print('Passed Protomer Check! Predicting Heterodimer...')\n", " he_model.predict(seq=''.join([seq[:LENGTH1], seq[LENGTH1+1:]]), num_recycles=3)\n", "\n", " if he_model.aux['losses']['plddt'] < 0.15 and he_model.aux['losses']['rmsd'] < 2.0:\n", " print('Passed Heterodimer Check! Predicting Homodimer 1...')\n", " ho_model.predict(seq=seq[:LENGTH1],num_recycles=3)\n", " print('Predicting Homodimer 2...')\n", " ipae1 = ho_model.aux['losses']['i_pae']\n", " ho_model.predict(seq=seq[LENGTH1+1:],num_recycles=3)\n", " ipae2 = ho_model.aux['losses']['i_pae']\n", " if ipae1 > 0.8 and ipae2 > 0.8:\n", " print('Passed Homodimer check!')\n", " he_model.save_pdb(f'Heterodimer_seq_{k}.pdb')\n", "\n" ], "metadata": { "id": "mII01sjwhAqX", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "def get_pdb(pdb_code=\"\"):\n", " if pdb_code is None or pdb_code == \"\":\n", " upload_dict = files.upload()\n", " pdb_string = upload_dict[list(upload_dict.keys())[0]]\n", " with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n", " return \"tmp.pdb\"\n", " elif os.path.isfile(pdb_code):\n", " return pdb_code\n", " elif len(pdb_code) == 4:\n", " os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n", " return f\"{pdb_code}.pdb\"\n", " else:\n", " os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n", " return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n", "\n", "def add_rg_loss(self, weight=0.1):\n", " '''add radius of gyration loss'''\n", " def loss_fn(inputs, outputs):\n", " xyz = outputs[\"structure_module\"]\n", " ca = xyz[\"final_atom_positions\"][:,residue_constants.atom_order[\"CA\"]]\n", "\n", " ca = ca[-self._binder_len:]\n", "\n", " rg = jnp.sqrt(jnp.square(ca - ca.mean(0)).sum(-1).mean() + 1e-8)\n", " rg_th = 2.38 * ca.shape[0] ** 0.365\n", " rg = jax.nn.elu(rg - rg_th)\n", " return {\"rg\":rg}\n", " self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n", " self.opt[\"weights\"][\"rg\"] = weight\n", "\n", "\n", "\n", "#@title # Binder Design\n", "#@markdown For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n", "LENGTH = 100 #@param {type:\"integer\"}\n", "binder_pdb = '5NGV' #@param {type:\"string\"}\n", "binder_chain ='A' #@param {type:\"string\"}\n", "hotspot ='' #@param {type:\"string\"}\n", "if hotspot == \"\": hotspot = None\n", "#@markdown ProteinMPNN Settings\n", "use_solubleMPNN = True #@param {type:\"boolean\"}\n", "#@markdown\n", "\n", "clear_mem()\n", "af_model = mk_afdesign_model(protocol=\"binder\")\n", "add_rg_loss(af_model)\n", "af_model.prep_inputs(pdb_filename=get_pdb(binder_pdb), chain=binder_chain,hotspot=hotspot, binder_len=LENGTH)\n", "\n", "\n", "af_model.restart(mode=[\"gumbel\", \"soft\"])\n", "\n", "af_model.opt[\"weights\"][\"rg\"] = 0.5\n", "\n", "af_model.opt[\"weights\"]['helix'] = -0.2\n", "af_model.opt[\"weights\"]['plddt'] = 0.1\n", "af_model.opt[\"weights\"]['pae'] = 0.1\n", "af_model.opt[\"weights\"]['i_pae'] = 0.1\n", "af_model.opt[\"weights\"]['i_con'] = 2.0\n", "\n", "print(\"weights\", af_model.opt[\"weights\"])\n", "print('Starting up and compiling JAX model....')\n" ], "metadata": { "id": "H4WqduyJ785H", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "af_model.design_logits(100)\n", "af_model.save_pdb(\"Binder.pdb\")" ], "metadata": { "id": "bS83MeVd99FN" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "af_model.plot_pdb()" ], "metadata": { "id": "UQMYKLCNZK2L" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "\n", "#@title # Binder Sequence Design with AF Multimer filtering\n", "#@markdown Use this to generate sequences for the binder candidate generated in the previous step\n", "\n", "#@markdown First we use the AF2 PTM model to predict the binder without receptor, acting as a fast pre filter. Then we use the AF Multimer model to predict the Receptor Binder complex. Again we use a template for the binder to help AF Multimer predicting the de novo designed protein\n", "\n", "binder_model = mk_afdesign_model(protocol=\"binder\",use_multimer=True,use_initial_guess=True)\n", "hall_model = mk_afdesign_model(protocol=\"fixbb\")\n", "\n", "\n", "binder_model.set_weights(i_pae=1.0)\n", "\n", "\n", "mpnn_model = mk_mpnn_model(weights=\"soluble\")\n", "mpnn_model.prep_inputs(pdb_filename=\"Binder.pdb\", chain='A,B', fix_pos='A',rm_aa=\"C\")\n", "\n", "samples = mpnn_model.sample_parallel(8,temperature=0.01)\n", "hall_model.prep_inputs(pdb_filename=\"Binder.pdb\", chain='B')\n", "binder_model.prep_inputs(pdb_filename=\"Binder.pdb\", chain='A', binder_chain='B',use_binder_template=True,rm_template_ic=True)\n", "k=0\n", "for seq in samples['seq']:\n", " print(\"Predicting binder only\")\n", " hall_model.predict(seq=seq[-LENGTH:], num_recycles=3)\n", " if hall_model.aux['losses']['rmsd'] < 2.0 :\n", " print(\"Passed! Predicting binder with receptor using AF Multimer\")\n", " binder_model.predict(seq=seq[-LENGTH:], num_recycles=3)\n", " plddt1 = binder_model.aux['losses']['plddt']\n", " i_pae = binder_model.aux['losses']['i_pae']\n", " if plddt1 < 0.15 and i_pae < 0.4:\n", " print(f\"Passed! Final I_PAE is {i_pae*31}\")\n", " binder_model.save_pdb(f'Binder_seq_{k}.pdb')\n", " binder_model.plot_pdb()\n", "\n", " k = k + 1" ], "metadata": { "id": "QfrRhvwTAyK-", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title # Site scaffolding example\n", "#@markdown This cell provides the code to perform the site scaffolding in bulk.\n", "#@markdown Just go to the commented section with names, contigs and length to insert the desired PDB identifier, contigs and final size and start designing.\n", "#@markdown Num_designs controls how many backbones one designes per PDB file\n", "\n", "num_designs = 1 #@param {type:\"integer\"}\n", "\n", "\n", "def get_pdb(pdb_code=\"\"):\n", " if pdb_code is None or pdb_code == \"\":\n", " upload_dict = files.upload()\n", " pdb_string = upload_dict[list(upload_dict.keys())[0]]\n", " with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n", " return \"tmp.pdb\"\n", " elif os.path.isfile(pdb_code):\n", " return pdb_code\n", " elif len(pdb_code) == 4:\n", " os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n", " return f\"{pdb_code}.pdb\"\n", " else:\n", " os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n", " return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n", "\n", "\n", "\n", "from colabdesign import mk_afdesign_model, clear_mem\n", "import contextlib\n", "\n", "from colabdesign.af.alphafold.common import residue_constants\n", "import jax\n", "import jax.numpy as jnp\n", "import pickle\n", "from colabdesign.mpnn import mk_mpnn_model\n", "\n", "import re\n", "import os\n", "\n", "\n", "#Add the names of the PDB files for the scaffolding problem here\n", "names = [\n", " \"1PRW\"\n", "]\n", "print(len(names))\n", "#Add the design contigs here\n", "\n", "inputs = [\n", " \"5-20,A16-35,10-25,A52-71,5-20\"\n", "]\n", "#Add the total length here. We only use the maximum length specified\n", "lengths = [\n", " \"60-105\"\n", "]\n", "\n", "\n", "def rg_loss(inputs, outputs):\n", " positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n", " ca = positions[::5, residue_constants.atom_order[\"CA\"]]\n", " center = ca.mean(0)\n", " rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n", " rg_th = 2.38 * ca.shape[0] ** 0.365\n", " rg = jax.nn.elu(rg - rg_th)\n", " return {\"rg\": rg}\n", "\n", "\n", "clear_mem()\n", "\n", "for _name, _input, _length in zip(\n", " names, inputs, lengths\n", "):\n", " print(f\"Starting on {_name}\")\n", " _input = _input.replace(\" \", \"\")\n", " __name = _name.split(\"_\")[0]\n", " model = mk_afdesign_model(\n", " protocol=\"partial\"\n", " )\n", " wire_loop_repr = [\"l\" if re.search(\"[A-Z]\", x) else \"w\" for x in _input.split(\",\")]\n", "\n", " _lengths = []\n", " for _id, rep in zip(wire_loop_repr, _input.split(\",\")):\n", " if \"-\" in rep: # loop or range\n", " if _id == \"l\": # loop\n", " rep = rep[1:]\n", " _len = int(rep.split(\"-\")[1]) - int(rep.split(\"-\")[0]) + 1\n", " else: # range\n", " _len = int(rep.split(\"-\")[1])\n", " else:\n", " if _id == \"l\":\n", " rep = rep[1:]\n", " _len = int(1)\n", " _lengths.append(_len)\n", "\n", " overall_length = sum(_lengths)\n", " print(overall_length)\n", " old_pos = list(filter(lambda x: re.search(\"[A-Z]\", x), _input.split(\",\")))\n", " order = list(range(len(old_pos)))\n", " old_pos = \",\".join(old_pos)\n", " wires = list(filter(lambda x: not re.search(\"[A-Z]\", x), _input.split(\",\")))\n", " wires = [\n", " int(wire) if \"-\" not in wire else int(wire.split(\"-\")[1]) for wire in wires\n", " ]\n", " offset = wires[0] if not wire_loop_repr[0] == \"l\" else 0\n", " if wire_loop_repr[0] == \"w\":\n", " wires = wires[1:]\n", " if wire_loop_repr[-1] == \"w\":\n", " wires = wires[:-1]\n", "\n", " chain = re.findall(\"[A-Z]\", _input)\n", " chain = list(set(chain))\n", " assert len(chain) == 1\n", " chain = chain[0]\n", " if \"-\" in _length:\n", " _length = _length.split(\"-\")[1]\n", " _length = int(_length)\n", " if _length < overall_length:\n", " _length = overall_length\n", "\n", " debug = False\n", " if debug:\n", " print(\"chain \" + str(chain))\n", " print(\"old_pos \" + str(old_pos))\n", " print(\"wires \" + str(wires))\n", " print(\"offset \" + str(offset))\n", " print(\"_length \" + str(_length))\n", " print(\"order \" + str(order))\n", " print(_name)\n", " pdb_file = get_pdb(_name)\n", " model.prep_inputs(\n", " pdb_file,\n", " chain=chain,\n", " pos=old_pos,\n", " length=_length,\n", " fix_seq=True,\n", " )\n", "\n", "\n", " model.rewire(\n", " order=order, # set order of segments\n", " loops=wires, # change loop length inbetween segments\n", " offset=offset,\n", " ) # essentially loop length at the N term\n", "\n", " print(\" Starting up and compiling JAX model....\")\n", "\n", " for i in range(num_designs):\n", " print(f\" Iteration {i} of 100\")\n", " model.restart(mode=[\"gumbel\", \"soft\"], rm_aa=\"C\")\n", " model.opt[\"weights\"][\"rg\"] = 0.1\n", " model.opt[\"weights\"][\"dgram_cce\"] = 2.0\n", " model.opt[\"weights\"][\"plddt\"] = 0.1\n", " model.opt[\"weights\"][\"pae\"] = 0.1\n", " model.opt[\"weights\"][\"rmsd\"] = 1.0\n", " model.opt[\"weights\"]['sc_rmsd'] = 1.0\n", " # model.opt[\"weights\"]['fape'] = 1.0\n", "\n", " model.design_logits(190)\n", " model.design_logits(10, save_best=True)\n", " outfile = f\"out_sc/{_name}_resesigned/{_name}_redesigned_{i}.pdb\"\n", " os.makedirs(os.path.dirname(outfile), exist_ok=True)\n", " model.save_pdb(outfile)\n", " mpnn_model = mk_mpnn_model()\n", "\n", " p = (\n", " []\n", " ) # [homo if not n in _interfaceFixturesIndexSecChain else hetero for n, (homo, hetero) in enumerate(zip(list(ho2), list(he[-len(ho2):])))]\n", " for k in model.opt[\"pos\"]:\n", " p.append(str(k + 1)) # Might be wrong\n", " p.append(\",\")\n", " posf = \"\".join(p[:-1])\n", "\n", " repredictionModel = mk_afdesign_model(\n", " protocol=\"fixbb\", use_templates=False\n", " )\n", " os.makedirs(os.path.dirname('out_sc_Redesigned/'), exist_ok=True)\n", "\n", " for j in range(num_designs):\n", " print(f\" Reprediction Iteration {j} of 100\")\n", " repredictionModel.prep_inputs(\n", " f\"out_sc/{_name}_resesigned/{_name}_redesigned_{j}.pdb\"\n", " )\n", "\n", " mpnn_model.prep_inputs(\n", " pdb_filename=f\"out_sc/{_name}_resesigned/{_name}_redesigned_{j}.pdb\",\n", " chain=\"A\",\n", " fix_pos=posf,\n", " rm_aa=\"C\",\n", " )\n", " out = mpnn_model.sample(num=1, batch=8, temperature=0.1)\n", "\n", " for n, i in enumerate(out[\"seq\"]):\n", "\n", " repredictionModel.predict(seq=i, num_recycles=3)\n", " if (\n", " repredictionModel.aux[\"log\"][\"rmsd\"] < 2.0\n", " and repredictionModel.aux[\"log\"][\"plddt\"] > 0.85\n", " ):\n", " filename = f'out_sc_Redesigned/{_name}_resesigned/{_name}_redesigned-{j}_num-{n}_rmsd-{int(repredictionModel.aux[\"log\"][\"rmsd\"]*100)}.pdb'\n", " os.makedirs(os.path.dirname(filename), exist_ok=True)\n", " repredictionModel.save_pdb(filename)\n", "\n", "for _name, _input, _length in zip(\n", " names, inputs, lengths\n", "):\n", " print(f\"Starting on {_name}\")\n", " clear_mem()\n", "\n", " _input = _input.replace(\" \", \"\")\n", " __name = _name.split(\"_\")[0]\n", "\n", " test_model = mk_afdesign_model(protocol='fixbb')\n", " model = mk_afdesign_model(\n", " protocol=\"partial\", use_templates=False\n", " ) # set True to constrain positions using template input\n", " # define positions we want to constrain (input PDB numbering)\n", "\n", " wire_loop_repr = [\"l\" if re.search(\"[A-Z]\", x) else \"w\" for x in _input.split(\",\")]\n", "\n", " _lengths = []\n", " for _id, rep in zip(wire_loop_repr, _input.split(\",\")):\n", " if \"-\" in rep: # loop or range\n", " if _id == \"l\": # loop\n", " rep = rep[1:]\n", " _len = int(rep.split(\"-\")[1]) - int(rep.split(\"-\")[0]) + 1\n", " else: # range\n", " _len = int(rep.split(\"-\")[1])\n", " else:\n", " if _id == \"l\":\n", " rep = rep[1:]\n", " _len = 1\n", " _lengths.append(_len)\n", "\n", " overall_length = sum(_lengths)\n", "\n", " old_pos = list(filter(lambda x: re.search(\"[A-Z]\", x), _input.split(\",\")))\n", " order = list(range(len(old_pos)))\n", " old_pos = \",\".join(old_pos)\n", " wires = list(filter(lambda x: not re.search(\"[A-Z]\", x), _input.split(\",\")))\n", " wires = [\n", " int(wire) if \"-\" not in wire else int(wire.split(\"-\")[1]) for wire in wires\n", " ]\n", " offset = wires[0] if not wire_loop_repr[0] == \"l\" else 0\n", " if wire_loop_repr[0] == \"w\":\n", " wires = wires[1:]\n", " if wire_loop_repr[-1] == \"w\":\n", " wires = wires[:-1]\n", "\n", " chain = re.findall(\"[A-Z]\", _input)\n", " chain = list(set(chain))\n", " assert len(chain) == 1\n", " chain = chain[0]\n", " if \"-\" in _length:\n", " _length = _length.split(\"-\")[1]\n", " _length = int(_length)\n", " if _length < overall_length:\n", " _length = overall_length\n", "\n", "\n", " print(_name)\n", " pdb_file = get_pdb(_name)\n", "\n", "\n", "\n", " model.prep_inputs(\n", " pdb_file,\n", " chain=chain,\n", " pos=old_pos, # define positions to contrain\n", " length=_length, # define if the desired length is different from input PDB\n", " fix_seq=True,\n", " ) # set True to constrain the sequence\n", "\n", " # set positions (if different from PDB)\n", " # reorder the segments,\n", " model.rewire(\n", " order=order, # set order of segments\n", " loops=wires, # change loop length inbetween segments\n", " offset=offset,\n", " ) # essentially loop length at the N term\n", "\n", " in_files = os.listdir(f'out_sc_Redesigned/{_name}_resesigned/')\n", " if not os.path.exists(f'out_sc_Redesigned/{_name}_resesigned/out/'):\n", " os.mkdir(f'out_sc_Redesigned/{_name}_resesigned/out/')\n", " for ii in in_files:\n", " if ii[-1] == 'b':\n", " test_model.prep_inputs(pdb_filename=f'out_sc_Redesigned/{_name}_resesigned/{ii}')\n", " seq = test_model._inputs['batch'][\"aatype\"]\n", " #print(seq)\n", " model.predict(seq=seq, num_recycles=3)\n", " if model.aux[\"losses\"][\"rmsd\"] < 1.0:\n", " model.save_pdb(f'out_sc_Redesigned/{_name}_resesigned/out/{ii}')\n", " with open(f'out_sc_Redesigned/{_name}_resesigned/out/{ii[:-4]}_data.pkl', 'wb') as f:\n", " pickle.dump(model.aux[\"losses\"][\"rmsd\"], f)\n", "\n" ], "metadata": { "id": "8sLixVvGGJ47", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [], "metadata": { "id": "QHFiJeHjB4sM" }, "execution_count": null, "outputs": [] } ], "metadata": { "accelerator": "GPU", "colab": { "provenance": [], "machine_shape": "hm", "gpuType": "L4", "include_colab_link": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }