{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "VD9K5H2cnFxL" }, "source": [ "#AF_pseudo_diffusion + proteinMPNN\n", "Hacking AlphaFold to be a diffusion model (for backbone generation) via recycling mechanism. At each step add logits from proteinMPNN.\n", "\n", "\n", "**WARNING**: This notebook is experimental, designed as a control. Not intended for practical use at this stage." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "YCRzPGdTZfEe" }, "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 git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\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, re\n", "from colabdesign import mk_afdesign_model, clear_mem\n", "from colabdesign.mpnn import mk_mpnn_model\n", "from colabdesign.af.alphafold.common import residue_constants\n", "from colabdesign.shared.protein import _np_get_cb\n", "\n", "from IPython.display import HTML\n", "from google.colab import files\n", "import numpy as np\n", "import jax.numpy as jnp\n", "import jax\n", "from scipy.special import softmax, log_softmax\n", "\n", "import tqdm.notebook\n", "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\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", "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)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "wxYMB4A9Zmrf" }, "outputs": [], "source": [ "#@title initialize the model\n", "length = 100 #@param {type:\"integer\"}\n", "#@markdown Provide a starting point (optional)\n", "starting_seq = \"\" #@param {type:\"string\"}\n", "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n", "#@markdown - if `starting_seq` provided the `length` option will be overwritten by length of starting sequence.\n", "\n", "if len(starting_seq) > 0:\n", " length = len(starting_seq)\n", "\n", "clear_mem()\n", "af_model = mk_afdesign_model(protocol=\"hallucination\")\n", "af_model.prep_inputs(length=length)\n", "mpnn_model = mk_mpnn_model()\n", "\n", "af_model_test = mk_afdesign_model(protocol=\"fixbb\", best_metric=\"rmsd\")\n", "mpnn_model_test = mk_mpnn_model()\n", "\n", "print(\"lengths\",af_model._lengths)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "3Dt8i00UbxtW", "cellView": "form" }, "outputs": [], "source": [ "#@title run protocol\n", "#@markdown Optimization options\n", "iterations = 100 #@param [\"50\", \"100\", \"200\"] {type:\"raw\"}\n", "\n", "#AlphaFold options\n", "use_dropout = True\n", "num_recycles = 0\n", "\n", "#@markdown proteinMPNN options (set to `none` to disable)\n", "mpnn_mode = \"conditional\" #@param [\"none\",\"conditional\", \"unconditional\"]\n", "\n", "#@markdown proteinMPNN contact map masking options\n", "cmap_seqsep = 6 #@param {type:\"raw\"}\n", "cmap_num = 1 #@param {type:\"integer\"}\n", "cmap_cutoff = 8 #@param {type:\"raw\"}\n", "\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", "\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", "\n", "# initialize sequence\n", "if len(starting_seq) > 1:\n", " af_model.set_seq(seq=starting_seq)\n", "\n", "# initialize coordinates\n", "af_model._inputs.pop(\"prev\",None)\n", "af_model._inputs[\"batch\"] = {\"aatype\":np.zeros(L).astype(int),\n", " \"all_atom_mask\":np.zeros((L,37)),\n", " \"all_atom_positions\":np.zeros((L,37,3))}\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", "\n", " # denoise\n", " aux = af_model.predict(return_aux=True, verbose=False,\n", " dropout=use_dropout,\n", " num_recycles=num_recycles)\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", " # per position confidence\n", " cmap = aux[\"cmap\"] * (np.abs(offset) > cmap_seqsep)\n", " conf = np.sort(cmap)[:,-cmap_num:].mean(-1)\n", "\n", " # gather features\n", " plddt = aux[\"plddt\"]\n", " seq = aux[\"seq\"][\"hard\"][0].argmax(-1)\n", " xyz = aux[\"atom_positions\"].copy()\n", " \n", " # update inputs \n", " af_model._inputs[\"batch\"][\"aatype\"] = seq\n", " af_model._inputs[\"batch\"][\"all_atom_positions\"] = xyz\n", "\n", " # add logits from proteinmpnn at each stage\n", " if mpnn_mode != \"none\": \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", " new_logits = (1 - c) * sample_gumbel(mpnn_logits.shape) + c * mpnn_logits\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(\"init.pdb\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "YeYD4KF8MUA_" }, "outputs": [], "source": [ "af_model.plot_pdb()\n", "af_model.get_seqs()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6cJhizcYcUxm" }, "outputs": [], "source": [ "HTML(af_model.animate(dpi=100))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "3H_B0AoYIiWH" }, "outputs": [], "source": [ "#@title sample new sequences using proteinMPNN and rescore with alphafold (w/o template)\n", "#@markdown #### MPNN Options\n", "num_seqs = 16 #@param [\"8\", \"16\", \"32\", \"64\", \"128\", \"256\", \"512\", \"1024\"] {type:\"raw\"}\n", "sampling_temp = 0.1 \n", "#@markdown #### AlphaFold Options\n", "alphafold_model = \"model_4_ptm\" #@param [\"model_1_ptm\", \"model_2_ptm\", \"model_3_ptm\", \"model_4_ptm\", \"model_5_ptm\"]\n", "num_recycles = 3 #@param [\"0\", \"1\", \"2\", \"3\"] {type:\"raw\"}\n", "import pandas as pd\n", "\n", "# zero out template inputs\n", "af_model_test.prep_inputs(\"init.pdb\")\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", "af_terms = [\"plddt\",\"ptm\",\"pae\",\"rmsd\",\"dgram_cce\"]\n", "for k in af_terms: out[k] = []\n", "os.system(\"mkdir -p output/all_pdb\")\n", "\n", "with tqdm.notebook.tqdm(total=out[\"S\"].shape[0], bar_format=TQDM_BAR_FORMAT) as pbar:\n", " with open(\"design.fasta\",\"w\") as fasta:\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 t in af_terms:\n", " out[t].append(af_model_test.aux[\"log\"][t])\n", " out[\"pae\"][-1] = out[\"pae\"][-1] * 31\n", " af_model_test._save_results(save_best=True, verbose=False)\n", " af_model_test.save_current_pdb(f\"output/all_pdb/n{n}.pdb\")\n", " af_model_test._k += 1\n", "\n", " line = f'>mpnn:{out[\"score\"][n]:.3f}_plddt:{out[\"plddt\"][n]:.3f}_ptm:{out[\"ptm\"][n]:.3f}_pae:{out[\"pae\"][n]:.3f}\\n{out[\"seq\"][n]}'\n", " fasta.write(line+\"\\n\")\n", " pbar.update(1)\n", "\n", "af_model_test.save_pdb(\"final.pdb\")\n", "\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", "\n", "df = pd.DataFrame(data, columns=labels)\n", "df.to_csv('output/mpnn_results.csv')\n", "df.round(3).sort_values(\"rmsd\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "v3dPQeEQcAcT" }, "outputs": [], "source": [ "af_model_test.plot_pdb()\n", "af_model_test.get_seqs()" ] } ], "metadata": { "accelerator": "GPU", "colab": { "provenance": [], "include_colab_link": true }, "gpuClass": "standard", "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }