{ "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 distogram. 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\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": { "id": "wxYMB4A9Zmrf", "cellView": "form" }, "outputs": [], "source": [ "#@title initialize the model\n", "length = 100 #@param {type:\"integer\"}\n", "#symmetry = \"C\" #@param [\"C\"]\n", "#copies = 1 #@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", "#@markdown Experimental options\n", "use_multimer = False \n", "mode = \"dgram_retrain\" #@param [\"dgram\",\"dgram_retrain\"]\n", "##@markdown - `xyz` - use structure output as template input\n", "#@markdown - `dgram` - use distogram output as template input\n", "#@markdown - `dgram_retrain` - replace distogram head from AlphaFold with one retrained to map output bins to template bins.\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", " use_templates=True,\n", " debug=True, \n", " use_multimer=use_multimer)\n", "af_model.prep_inputs(length=length)\n", "mpnn_model = mk_mpnn_model()\n", "print(\"lengths\",af_model._lengths)\n", "\n", "if \"dgram\" in mode:\n", " if \"retrain\" in mode and not use_multimer:\n", " # update distogram head to return all 39 bins\n", " af_model._cfg.model.heads.distogram.first_break = 3.25\n", " af_model._cfg.model.heads.distogram.last_break = 50.75\n", " af_model._cfg.model.heads.distogram.num_bins = 39\n", " af_model._model = af_model._get_model(af_model._cfg)\n", " from colabdesign.af.weights import __file__ as af_path\n", " template_dgram_head = np.load(os.path.join(os.path.dirname(af_path),'template_dgram_head.npy'))\n", " for k in range(len(af_model._model_params)):\n", " params = {\"weights\":jnp.array(template_dgram_head[k]),\"bias\":jnp.zeros(39)}\n", " af_model._model_params[k][\"alphafold/alphafold_iteration/distogram_head/half_logits\"] = params\n", " else:\n", " dgram_map = np.eye(39)[np.repeat(np.append(0,np.arange(15)),4)]\n", " dgram_map[-1,:] = 0 \n", "\n", "def get_dgram(positions, num_bins=39, min_bin=3.25, max_bin=50.75):\n", " atom_idx = residue_constants.atom_order\n", " atoms = {k:positions[...,atom_idx[k],:] for k in [\"N\",\"CA\",\"C\"]}\n", " cb = _np_get_cb(**atoms, use_jax=False)\n", " dist2 = np.square(cb[None,:] - cb[:,None]).sum(-1,keepdims=True)\n", " lower_breaks = np.linspace(min_bin, max_bin, num_bins)\n", " lower_breaks = np.square(lower_breaks)\n", " upper_breaks = np.concatenate([lower_breaks[1:],np.array([1e8], dtype=jnp.float32)], axis=-1)\n", " return ((dist2 > lower_breaks) * (dist2 < upper_breaks)).astype(float)" ] }, { "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\"] {type:\"raw\"}\n", "use_dgram_noise = True #@param {type:\"boolean\"}\n", "use_seq_noise = True #@param {type:\"boolean\"}\n", "use_dropout = True #@param {type:\"boolean\"}\n", "seqsep_mask = 6 #@param {type:\"integer\"}\n", "\n", "#@markdown AlphaFold options\n", "sample_models = True #@param {type:\"boolean\"}\n", "num_recycles = 0 #@param [\"0\", \"1\", \"2\", \"3\"] {type:\"raw\"}\n", "\n", "#@markdown proteinMPNN options (set to `none` to disable)\n", "mpnn_mode = \"conditional\" #@param [\"none\",\"sample\", \"conditional\", \"unconditional\"]\n", "\n", "L = sum(af_model._lengths)\n", "af_model.restart(mode=\"gumbel\")\n", "af_model._inputs[\"rm_template_seq\"] = True\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", "af_model._inputs[\"bias\"] = np.zeros((L,20))\n", "\n", "# initialize coordinates/dgram\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", " \"dgram\":np.zeros((L,L,39))}\n", "\n", "for k in range(iterations):\n", "\n", " # disable stochastic part for the last 10 steps\n", " if k > (iterations - 10):\n", " use_dropout = False\n", " sample_models = False\n", " use_seq_noise = False\n", " seqsep_mask = 0.0\n", "\n", " # noise\n", " if k > 0:\n", " dgram_xyz = get_dgram(xyz)\n", " dgram_prob = softmax(dgram_logits,-1)\n", "\n", " if use_seq_noise:\n", " af_model._inputs[\"bias\"] = 0.1 * sample_gumbel((L,20))\n", "\n", " if mode == \"xyz\":\n", " dgram = dgram_xyz\n", " if mode == \"dgram\":\n", " dgram = dgram_prob @ dgram_map\n", " dgram[...,14:] = dgram_xyz[...,14:] * dgram_prob[...,-1:]\n", " if mode == \"dgram_retrain\":\n", " dgram = dgram_prob\n", " \n", " if use_dgram_noise:\n", " noise = sample_gumbel(dgram.shape) * (1 - k/iterations)\n", " dgram = softmax(np.log(dgram + 1e-8) + noise, -1)\n", "\n", " # add mask to avoid local contacts being fixed (otherwise there is a bias toward helix)\n", " mask = np.abs(offset) > seqsep_mask\n", " af_model._inputs[\"batch\"][\"dgram\"] = dgram * mask[:,:,None]\n", "\n", " # denoise\n", " aux = af_model.predict(return_aux=True, verbose=False,\n", " sample_models=sample_models,\n", " dropout=use_dropout, num_recycles=num_recycles)\n", " plddt = aux[\"plddt\"]\n", " seq = aux[\"seq\"][\"hard\"][0].argmax(-1)\n", " xyz = aux[\"atom_positions\"].copy()\n", " dgram_logits = aux[\"debug\"][\"outputs\"][\"distogram\"][\"logits\"] \n", " \n", " # update inputs \n", " af_model._inputs[\"batch\"][\"aatype\"] = seq\n", " af_model._inputs[\"batch\"][\"all_atom_mask\"][:,:4] = np.sqrt(plddt)[:,None]\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", " if mpnn_mode == \"sample\":\n", " mpnn_out = mpnn_model.sample(temp = 1-k/iterations)\n", " mpnn_logits = mpnn_out[\"logits\"][0,:,:20]\n", " aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"][0]\n", " else:\n", " opt = {} if mpnn_mode == \"conditional\" else {\"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", " beta = np.square(k/iterations) * plddt[:,None]\n", " af_model._params[\"seq\"] = (1-beta) * af_model._params[\"seq\"] + beta * mpnn_logits\n", "\n", " # save results\n", " af_model._save_results(aux)\n", " af_model._k += 1" ] }, { "cell_type": "code", "source": [ "af_model.save_pdb(\"0.pdb\")\n", "af_model.plot_pdb()\n", "af_model.get_seqs()" ], "metadata": { "id": "YeYD4KF8MUA_" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6cJhizcYcUxm" }, "outputs": [], "source": [ "HTML(af_model.animate(dpi=100))" ] }, { "cell_type": "code", "source": [ "#@title sample new sequences using proteinMPNN and rescore with alphafold (w/o template)\n", "#@markdown #### Design Options\n", "num_seqs = 32 #@param [\"32\", \"64\", \"128\", \"256\", \"512\", \"1024\"] {type:\"raw\"}\n", "sampling_temp = 0.1 \n", "import pandas as pd\n", "from google.colab import data_table\n", "data_table.enable_dataframe_formatter()\n", "\n", "# zero out template inputs\n", "out = mpnn_model.sample(num=num_seqs//32, batch=32,\n", " temperature=sampling_temp)\n", "af_terms = [\"plddt\",\"ptm\",\"pae\"]\n", "for k in af_terms: out[k] = []\n", "os.system(\"mkdir -p output/all_pdb\")\n", "\n", "af_model._inputs[\"batch\"][\"dgram\"] = np.zeros((L,L,39))\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.predict(seq=seq,\n", " num_recycles=1,\n", " num_models=1,\n", " verbose=False)\n", "\n", " for t in af_terms:\n", " out[t].append(af_model.aux[\"log\"][t])\n", " out[\"pae\"][-1] = out[\"pae\"][-1] * 31\n", " af_model._save_results(save_best=True, verbose=False)\n", " af_model.save_current_pdb(f\"output/all_pdb/n{n}.pdb\")\n", " af_model._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", "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", "data_table.DataTable(df.round(3).sort_values(\"pae\"))" ], "metadata": { "cellView": "form", "id": "3H_B0AoYIiWH" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "v3dPQeEQcAcT" }, "outputs": [], "source": [ "af_model.save_pdb(\"1.pdb\")\n", "af_model.plot_pdb()\n", "af_model.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 }