ProteinMPNN_ddG_binding.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/PeptoneLtd/proteinmpnn_ddg/blob/main/ProteinMPNN_ddG_binding.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#ProteinMPNN-ddG for binding affinity\n",
        "\n",
        "Scores all possible point mutations of a protein to improving binding at an interface\n"
      ],
      "metadata": {
        "id": "IVmFMidn965N"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install ProteinMPNN-ddG (and colabdesign)\n",
        "import os\n",
        "try:\n",
        "  import proteinmpnn_ddg\n",
        "except:\n",
        "  os.system(\"pip install -q proteinmpnn_ddg[cuda12]@git+https://github.com/PeptoneLtd/proteinmpnn_ddg.git@paper\")\n",
        "\n",
        "from proteinmpnn_ddg import predict_logits_for_all_point_mutations_of_single_pdb\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "pd.options.display.float_format = '{:,.2f}'.format\n",
        "\n",
        "from google.colab import files\n",
        "from google.colab import data_table\n",
        "data_table.disable_dataframe_formatter()\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\""
      ],
      "metadata": {
        "id": "tnwjNa-IWMSA",
        "cellView": "form"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import warnings, os, re\n",
        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
        "\n",
        "os.system(\"mkdir -p output\")\n",
        "\n",
        "# USER OPTIONS\n",
        "#@markdown # ProteinMPNN options\n",
        "model_name = \"v_48_020\" #@param [\"v_48_002\", \"v_48_010\", \"v_48_020\", \"v_48_030\"]\n",
        "#@markdown (v_48_020 recommended)\n",
        "\n",
        "#@markdown # Input Options\n",
        "pdb='7WPH' #@param {type:\"string\"}\n",
        "#@markdown (leave `pdb` as  blank to get an upload prompt)\n",
        "binder_chains = \"H,L\" #@param {type:\"string\"}\n",
        "receptor_chains = \"A\" #@param {type:\"string\"}\n",
        "#@markdown (You can specify several chains, separating by commas e.g. \"A,C\")\n",
        "\n",
        "#@markdown Only the chains specified will be loaded from the PDB file for prediction\n",
        "# chains_to_predict = \"\" #@param {type:\"string\"}\n",
        "# #@markdown (Leave `chains_to_predict` empty to predict all chains)\n",
        "\n",
        "nrepeats = 1\n",
        "seed = 42\n",
        "\n",
        "# cleaning user options\n",
        "binder_chains, receptor_chains = (re.sub(\"[^A-Za-z]+\",\",\", chains).split(',') for chains in (\n",
        "   binder_chains, receptor_chains)\n",
        ")\n",
        "chains = list(binder_chains) + list(receptor_chains)\n",
        "\n",
        "pdb_path = get_pdb(pdb)"
      ],
      "metadata": {
        "id": "GjdIxO4j-Hnn",
        "cellView": "form"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run ProteinMPNN-binding-ddG\n",
        "#@markdown We compute the logit difference with and without the binding partner\n",
        "#@markdown then take the difference to identify mutations that should improve\n",
        "#@markdown binding to the receptor chains.\n",
        "\n",
        "#@markdown We would suggest taking the mutations with the\n",
        "#@markdown highest binding ddG after filtering for unbound ddG > 0\n",
        "\n",
        "#@markdown (Positive values are good mutations, which strengthen binding/stability)\n",
        "\n",
        "# # @markdown ```df[df.unbound_ddg>0].sort_values('binding_ddg', ascending=False)```\n",
        "data = {}\n",
        "for source, chains_to_predict, context_chains in [\n",
        "    ('unbound', binder_chains, []),\n",
        "    ('bound', binder_chains, receptor_chains)\n",
        "]:\n",
        "  dfs = []\n",
        "  for chain in chains_to_predict:\n",
        "    df = predict_logits_for_all_point_mutations_of_single_pdb(\n",
        "        model_name,\n",
        "        chains_to_predict+context_chains,\n",
        "        pdb_path,\n",
        "        nrepeat=nrepeats,\n",
        "        seed=seed,\n",
        "        chain_to_predict=chain,\n",
        "        pad_inputs=False,\n",
        "        apply_ddG_correction=True)\n",
        "    df['chain'] = chain\n",
        "    dfs.append(df)\n",
        "  df = pd.concat(dfs).set_index(['chain','pre','pos', 'post']).rename(\n",
        "      columns={'logit_difference_ddg': f'{source}_ddg'})[f'{source}_ddg']\n",
        "  data[source] = df\n",
        "\n",
        "df = pd.concat(data.values(), axis=1)\n",
        "df['binding_ddg'] = df['bound_ddg'] - df['unbound_ddg']\n",
        "df.to_csv('predictions.csv')\n",
        "\n",
        "display(df[df.unbound_ddg>0].sort_values(\n",
        "    'binding_ddg', ascending=False\n",
        "    ).head(\n",
        "      10\n",
        "      ).style.hide().format(\n",
        "          precision=2, decimal=\".\"\n",
        "          )\n",
        "      )"
      ],
      "metadata": {
        "cellView": "form",
        "id": "JuRjhXhr9wte",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 363
        },
        "outputId": "e2d197fb-3b8a-4c3e-8312-b0eee92889f7"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<pandas.io.formats.style.Styler at 0x7a0bdc111c30>"
            ],
            "text/html": [
              "<style type=\"text/css\">\n",
              "</style>\n",
              "<table id=\"T_33d0f\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr>\n",
              "      <th id=\"T_33d0f_level0_col0\" class=\"col_heading level0 col0\" >unbound_ddg</th>\n",
              "      <th id=\"T_33d0f_level0_col1\" class=\"col_heading level0 col1\" >bound_ddg</th>\n",
              "      <th id=\"T_33d0f_level0_col2\" class=\"col_heading level0 col2\" >binding_ddg</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row0_col0\" class=\"data row0 col0\" >0.49</td>\n",
              "      <td id=\"T_33d0f_row0_col1\" class=\"data row0 col1\" >1.37</td>\n",
              "      <td id=\"T_33d0f_row0_col2\" class=\"data row0 col2\" >0.88</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row1_col0\" class=\"data row1 col0\" >0.52</td>\n",
              "      <td id=\"T_33d0f_row1_col1\" class=\"data row1 col1\" >1.12</td>\n",
              "      <td id=\"T_33d0f_row1_col2\" class=\"data row1 col2\" >0.60</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row2_col0\" class=\"data row2 col0\" >0.49</td>\n",
              "      <td id=\"T_33d0f_row2_col1\" class=\"data row2 col1\" >1.09</td>\n",
              "      <td id=\"T_33d0f_row2_col2\" class=\"data row2 col2\" >0.60</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row3_col0\" class=\"data row3 col0\" >1.20</td>\n",
              "      <td id=\"T_33d0f_row3_col1\" class=\"data row3 col1\" >1.77</td>\n",
              "      <td id=\"T_33d0f_row3_col2\" class=\"data row3 col2\" >0.56</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row4_col0\" class=\"data row4 col0\" >0.74</td>\n",
              "      <td id=\"T_33d0f_row4_col1\" class=\"data row4 col1\" >1.29</td>\n",
              "      <td id=\"T_33d0f_row4_col2\" class=\"data row4 col2\" >0.55</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row5_col0\" class=\"data row5 col0\" >0.69</td>\n",
              "      <td id=\"T_33d0f_row5_col1\" class=\"data row5 col1\" >1.17</td>\n",
              "      <td id=\"T_33d0f_row5_col2\" class=\"data row5 col2\" >0.48</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row6_col0\" class=\"data row6 col0\" >1.35</td>\n",
              "      <td id=\"T_33d0f_row6_col1\" class=\"data row6 col1\" >1.81</td>\n",
              "      <td id=\"T_33d0f_row6_col2\" class=\"data row6 col2\" >0.46</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row7_col0\" class=\"data row7 col0\" >1.06</td>\n",
              "      <td id=\"T_33d0f_row7_col1\" class=\"data row7 col1\" >1.47</td>\n",
              "      <td id=\"T_33d0f_row7_col2\" class=\"data row7 col2\" >0.42</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row8_col0\" class=\"data row8 col0\" >0.54</td>\n",
              "      <td id=\"T_33d0f_row8_col1\" class=\"data row8 col1\" >0.95</td>\n",
              "      <td id=\"T_33d0f_row8_col2\" class=\"data row8 col2\" >0.41</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_33d0f_row9_col0\" class=\"data row9 col0\" >0.19</td>\n",
              "      <td id=\"T_33d0f_row9_col1\" class=\"data row9 col1\" >0.56</td>\n",
              "      <td id=\"T_33d0f_row9_col2\" class=\"data row9 col2\" >0.37</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title download predictions (optional)\n",
        "from google.colab import files\n",
        "files.download(f'predictions.csv')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ZOtuzwwUAgHj",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "outputId": "b24bce95-842c-45df-e9fc-bd331f713d16"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "    async function download(id, filename, size) {\n",
              "      if (!google.colab.kernel.accessAllowed) {\n",
              "        return;\n",
              "      }\n",
              "      const div = document.createElement('div');\n",
              "      const label = document.createElement('label');\n",
              "      label.textContent = `Downloading \"${filename}\": `;\n",
              "      div.appendChild(label);\n",
              "      const progress = document.createElement('progress');\n",
              "      progress.max = size;\n",
              "      div.appendChild(progress);\n",
              "      document.body.appendChild(div);\n",
              "\n",
              "      const buffers = [];\n",
              "      let downloaded = 0;\n",
              "\n",
              "      const channel = await google.colab.kernel.comms.open(id);\n",
              "      // Send a message to notify the kernel that we're ready.\n",
              "      channel.send({})\n",
              "\n",
              "      for await (const message of channel.messages) {\n",
              "        // Send a message to notify the kernel that we're ready.\n",
              "        channel.send({})\n",
              "        if (message.buffers) {\n",
              "          for (const buffer of message.buffers) {\n",
              "            buffers.push(buffer);\n",
              "            downloaded += buffer.byteLength;\n",
              "            progress.value = downloaded;\n",
              "          }\n",
              "        }\n",
              "      }\n",
              "      const blob = new Blob(buffers, {type: 'application/binary'});\n",
              "      const a = document.createElement('a');\n",
              "      a.href = window.URL.createObjectURL(blob);\n",
              "      a.download = filename;\n",
              "      div.appendChild(a);\n",
              "      a.click();\n",
              "      div.remove();\n",
              "    }\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "download(\"download_d4ca1f0a-728a-4800-9683-5ab394c01607\", \"predictions.csv\", 358877)"
            ]
          },
          "metadata": {}
        }
      ]
    }
  ]
}