fixing the use_esm_1b notebook, updating to use alphafolddb v3

Author: sokrypton

af/examples/use_esm_1b_bias.ipynb

@@ -10,6 +10,59 @@
         "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/use_esm_1b_bias.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Gather inputs"
+      ],
+      "metadata": {
+        "id": "M9xamJDvx1pd"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import requests, re\n",
+        "from google.colab import files\n",
+        "\n",
+        "def get_uniprot_seq(uid):\n",
+        "  url = f'https://rest.uniprot.org/uniprotkb/stream?compressed=false&format=fasta&query={uid}'\n",
+        "  sequence = \"\".join(re.split(r'\\n(?=>)', requests.get(url).text)[0].split(\"\\n\")[1:])\n",
+        "  return sequence\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\",alphafold_model=False):\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",
+        "  else:\n",
+        "    if alphafold_model:\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",
+        "    else:\n",
+        "      os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "      return f\"{pdb_code}.pdb\""
+      ],
+      "metadata": {
+        "id": "lev6Nc_-vx2D"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "UNIPROT = \"P0A6A8\"\n",
+        "SEQUENCE = get_uniprot_seq(UNIPROT)\n",
+        "PDB_FILENAME = get_pdb(UNIPROT, alphafold_model=True)"
+      ],
+      "metadata": {
+        "id": "VkPdQj-Fv1LQ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -93,7 +146,7 @@
       "outputs": [],
       "source": [
         "# get bias\n",
-        "seq = \"MSTIEERVKKIIGEQLGVKQEEVTNNASFVEDLGADSLDTVELVMALEEEFDTEIPDEEAEKITTVQAAIDYINGHQA\"\n",
+        "seq = SEQUENCE\n",
         "bias = get_bias_from_esm(seq)\n",
         "np.savetxt(\"bias.txt\",bias)"
       ]
@@ -107,7 +160,7 @@
       "outputs": [],
       "source": [
         "import matplotlib.pyplot as plt\n",
-        "plt.imshow(bias.T)"
+        "plt.imshow(bias.T,cmap=\"bwr_r\",vmin=-10,vmax=10)"
       ]
     },
     {
@@ -157,34 +210,21 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "Vt7G_nbNeSQ3"
+        "id": "Vt7G_nbNeSQ3",
+        "cellView": "form"
       },
       "outputs": [],
       "source": [
-        "#import libraries\n",
+        "#@title import libraries\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.af.alphafold.common import residue_constants\n",
         "from IPython.display import HTML\n",
-        "from google.colab import files\n",
         "import numpy as np\n",
-        "from alphafold.common import residue_constants\n",
-        "\n",
-        "def get_pdb(pdb_code=\"\",alphafold_model=False):\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",
-        "  else:\n",
-        "    if alphafold_model:\n",
-        "      os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v2.pdb\")\n",
-        "      return f\"AF-{pdb_code}-F1-model_v2.pdb\"      \n",
-        "    else:\n",
-        "      os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
-        "      return f\"{pdb_code}.pdb\""
+        "\n"
       ]
     },
     {
@@ -197,11 +237,8 @@
       "source": [
         "clear_mem()\n",
         "model = mk_afdesign_model(protocol=\"fixbb\",\n",
-        "                        use_templates=False) # set True to constrain structure\n",
-        "\n",
-        "pdb_filename = get_pdb(\"P0A6A8\",alphafold_model=True)\n",
-        "model.prep_inputs(pdb_filename, chain=\"A\")\n",
-        "\n",
+        "                          use_templates=False) # set True to constrain structure\n",
+        "model.prep_inputs(PDB_FILENAME, chain=\"A\")\n",
         "print(\"length\",  model._len)"
       ]
     },
@@ -227,8 +264,7 @@
       "outputs": [],
       "source": [
         "model.restart()\n",
-        "model.opt[\"bias\"] = bias\n",
-        "\n",
+        "model.set_seq(bias=bias)\n",
         "model.design_3stage(50,50,10)"
       ]
     },
@@ -268,9 +304,7 @@
     },
     {
       "cell_type": "code",
-      "source": [
-        ""
-      ],
+      "source": [],
       "metadata": {
         "id": "0Xd7zQ-PqTly"
       },
@@ -298,4 +332,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}