{ "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). 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.\n", "\n", "---\n", "**NEW**\n", "For latest version of notebook (reconfigured to work in distogram space) go here:\n", "[af_pseudo_diffusion_dgram.ipynb](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_pseudo_diffusion_dgram.ipynb)" ] }, { "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 IPython.display import HTML\n", "from google.colab import files\n", "import numpy as np\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", "#@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", "use_starting_pdb = False #@param {type:\"boolean\"}\n", "pdb = \"\" #@param {type:\"string\"}\n", "#@markdown - specify PDB or Uniprot code or leave pdb blank for upload prompt\n", "chains = \"A\" #@param {type:\"string\"}\n", "#@markdown - (example `A` or `A,B` for complexes)\n", "fix_pos = \"\" #@param {type:\"string\"}\n", "#@markdown - specify which positions to keep fixed in the sequence (example: `1,2-10`)\n", "#@markdown - you can also specify chain specific constraints (example: `A1-10,B1-20`)\n", "#@markdown - you can also specify to fix entire chain(s) (example: `A`)\n", "\n", "if len(starting_seq) > 0:\n", " length = len(starting_seq)\n", "\n", "clear_mem()\n", "if use_starting_pdb:\n", " af_model = mk_afdesign_model(protocol=\"fixbb\", use_templates=True)\n", " af_model.prep_inputs(get_pdb(pdb), chains, fix_pos=fix_pos) \n", "else:\n", " af_model = mk_afdesign_model(protocol=\"hallucination\", use_templates=True)\n", " af_model.prep_inputs(length=length)\n", "mpnn_model = mk_mpnn_model()\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\"] {type:\"raw\"}\n", "use_xyz_noise = True #@param {type:\"boolean\"}\n", "use_seq_noise = True #@param {type:\"boolean\"}\n", "use_dropout = True #@param {type:\"boolean\"}\n", "use_plddt = True #@param {type:\"boolean\"}\n", "store_denoised = True #@param {type:\"boolean\"}\n", "#@markdown - this does not change results, but toggle between storing noised vs denoised coordinates for animation.\n", "\n", "#@markdown AlphaFold options\n", "sample_models = False #@param {type:\"boolean\"}\n", "rm_template_seq = True #@param {type:\"boolean\"}\n", "\n", "#@markdown proteinMPNN options (disable to keep sequence the same)\n", "use_mpnn = True #@param {type:\"boolean\"}\n", "mpnn_mode = \"conditional\" #@param [\"conditional\", \"unconditional\"]\n", "\n", "af_model.restart(mode=\"gumbel\")\n", "\n", "if len(starting_seq) > 1:\n", " af_model.set_seq(seq=starting_seq)\n", "elif use_starting_pdb:\n", " af_model.set_seq(mode=\"wildtype\")\n", "\n", "af_model._inputs[\"rm_template_seq\"] = rm_template_seq\n", "L = sum(af_model._lengths)\n", "\n", "af_model._inputs[\"bias\"] = np.zeros((L,20))\n", "\n", "if not use_starting_pdb:\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", "for k in range(iterations):\n", " # add noise\n", " if use_seq_noise:\n", " af_model._inputs[\"bias\"] = 0.1 * sample_gumbel((L,20))\n", "\n", " if use_xyz_noise:\n", " n = np.random.normal(size=(L,37,3)) * (1-k/iterations)\n", " af_model._inputs[\"batch\"][\"all_atom_positions\"] += n\n", " if not store_denoised and k > 0:\n", " af_model._tmp[\"traj\"][\"xyz\"][-1] += n[:,1]\n", "\n", " # denoise\n", " aux = af_model.predict(return_aux=True, verbose=False,\n", " sample_models=sample_models, dropout=use_dropout)\n", " plddt = af_model.aux[\"plddt\"]\n", " # update inputs\n", " af_model._inputs[\"batch\"][\"aatype\"] = af_model.aux[\"seq\"][\"hard\"].argmax(-1)[0]\n", " af_model._inputs[\"batch\"][\"all_atom_mask\"][:,:4] = 1\n", " if use_plddt:\n", " af_model._inputs[\"batch\"][\"all_atom_mask\"][:,(1,3)] = np.sqrt(plddt)[:,None]\n", " af_model._inputs[\"batch\"][\"all_atom_positions\"] = af_model.aux[\"atom_positions\"].copy()\n", "\n", " # add logits from proteinmpnn at each stage\n", " if use_mpnn:\n", " mpnn_model.get_af_inputs(af_model)\n", " opt = {} if mpnn_mode == \"conditional\" else {\"ar_mask\":np.zeros((L,L))}\n", " mpnn_out = mpnn_model.score(**opt)\n", " aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"]\n", " mpnn_logits = mpnn_out[\"logits\"][:,:20]\n", " m = (k/iterations)\n", " af_model._params[\"seq\"] = (1-m) * af_model._params[\"seq\"] + m * mpnn_logits\n", "\n", " # save results\n", " af_model._save_results(aux)\n", " af_model._k += 1" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "v3dPQeEQcAcT" }, "outputs": [], "source": [ "af_model.plot_pdb()" ] }, { "cell_type": "code", "source": [ "af_model.save_pdb(\"tmp.pdb\")\n", "af_model.get_seqs()" ], "metadata": { "id": "6Ck-ME5_99qE" }, "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": [], "metadata": { "id": "aVS6VC7wKyhD" }, "execution_count": null, "outputs": [] } ], "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 }