ProteinMPNN_ddG.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.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\n",
        "\n",
        "Scores all possible point mutations of a protein, to identify those which improve stability and or expression\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 numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\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",
        "\n",
        "#@markdown # Input Options\n",
        "pdb='6MRR' #@param {type:\"string\"}\n",
        "#@markdown (leave `pdb` as  blank to get an upload prompt)\n",
        "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",
        "chains = re.sub(\"[^A-Za-z]+\",\",\", chains)\n",
        "chains = chains.split(',')\n",
        "\n",
        "pdb_path = get_pdb(pdb)"
      ],
      "metadata": {
        "id": "GjdIxO4j-Hnn",
        "cellView": "form"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run ProteinMPNN-ddG\n",
        "\n",
        "#@markdown (Positive values are good mutations, which strengthen stability and expression)\n",
        "\n",
        "dfs = []\n",
        "for chain in chains:\n",
        "  df = predict_logits_for_all_point_mutations_of_single_pdb(\n",
        "      model_name,\n",
        "      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)\n",
        "df = df.rename(columns={'logit_difference_ddg': 'proteinmpnn_ddg_score'})\n",
        "df = df[['chain','pre','pos','post','proteinmpnn_ddg_score']]\n",
        "df.to_csv('predictions.csv')\n",
        "\n",
        "display(df.sort_values(\n",
        "    'proteinmpnn_ddg_score', ascending=False\n",
        "    ).head(\n",
        "      10\n",
        "      ).style.hide().format(\n",
        "          precision=1, decimal=\".\"\n",
        "          )\n",
        "      )"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 363
        },
        "cellView": "form",
        "id": "JuRjhXhr9wte",
        "outputId": "a236265b-68f4-40f8-d365-59f8a77a3773"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<pandas.io.formats.style.Styler at 0x7f7f1bcafbe0>"
            ],
            "text/html": [
              "<style type=\"text/css\">\n",
              "</style>\n",
              "<table id=\"T_bd9b9\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr>\n",
              "      <th id=\"T_bd9b9_level0_col0\" class=\"col_heading level0 col0\" >chain</th>\n",
              "      <th id=\"T_bd9b9_level0_col1\" class=\"col_heading level0 col1\" >pre</th>\n",
              "      <th id=\"T_bd9b9_level0_col2\" class=\"col_heading level0 col2\" >pos</th>\n",
              "      <th id=\"T_bd9b9_level0_col3\" class=\"col_heading level0 col3\" >post</th>\n",
              "      <th id=\"T_bd9b9_level0_col4\" class=\"col_heading level0 col4\" >proteinmpnn_ddg_score</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row0_col0\" class=\"data row0 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row0_col1\" class=\"data row0 col1\" >R</td>\n",
              "      <td id=\"T_bd9b9_row0_col2\" class=\"data row0 col2\" >10</td>\n",
              "      <td id=\"T_bd9b9_row0_col3\" class=\"data row0 col3\" >V</td>\n",
              "      <td id=\"T_bd9b9_row0_col4\" class=\"data row0 col4\" >4.1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row1_col0\" class=\"data row1 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row1_col1\" class=\"data row1 col1\" >V</td>\n",
              "      <td id=\"T_bd9b9_row1_col2\" class=\"data row1 col2\" >36</td>\n",
              "      <td id=\"T_bd9b9_row1_col3\" class=\"data row1 col3\" >I</td>\n",
              "      <td id=\"T_bd9b9_row1_col4\" class=\"data row1 col4\" >3.4</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row2_col0\" class=\"data row2 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row2_col1\" class=\"data row2 col1\" >H</td>\n",
              "      <td id=\"T_bd9b9_row2_col2\" class=\"data row2 col2\" >9</td>\n",
              "      <td id=\"T_bd9b9_row2_col3\" class=\"data row2 col3\" >Y</td>\n",
              "      <td id=\"T_bd9b9_row2_col4\" class=\"data row2 col4\" >3.4</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row3_col0\" class=\"data row3 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row3_col1\" class=\"data row3 col1\" >S</td>\n",
              "      <td id=\"T_bd9b9_row3_col2\" class=\"data row3 col2\" >3</td>\n",
              "      <td id=\"T_bd9b9_row3_col3\" class=\"data row3 col3\" >D</td>\n",
              "      <td id=\"T_bd9b9_row3_col4\" class=\"data row3 col4\" >3.3</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row4_col0\" class=\"data row4 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row4_col1\" class=\"data row4 col1\" >S</td>\n",
              "      <td id=\"T_bd9b9_row4_col2\" class=\"data row4 col2\" >3</td>\n",
              "      <td id=\"T_bd9b9_row4_col3\" class=\"data row4 col3\" >N</td>\n",
              "      <td id=\"T_bd9b9_row4_col4\" class=\"data row4 col4\" >3.1</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row5_col0\" class=\"data row5 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row5_col1\" class=\"data row5 col1\" >G</td>\n",
              "      <td id=\"T_bd9b9_row5_col2\" class=\"data row5 col2\" >41</td>\n",
              "      <td id=\"T_bd9b9_row5_col3\" class=\"data row5 col3\" >C</td>\n",
              "      <td id=\"T_bd9b9_row5_col4\" class=\"data row5 col4\" >2.5</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row6_col0\" class=\"data row6 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row6_col1\" class=\"data row6 col1\" >W</td>\n",
              "      <td id=\"T_bd9b9_row6_col2\" class=\"data row6 col2\" >2</td>\n",
              "      <td id=\"T_bd9b9_row6_col3\" class=\"data row6 col3\" >M</td>\n",
              "      <td id=\"T_bd9b9_row6_col4\" class=\"data row6 col4\" >2.4</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row7_col0\" class=\"data row7 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row7_col1\" class=\"data row7 col1\" >H</td>\n",
              "      <td id=\"T_bd9b9_row7_col2\" class=\"data row7 col2\" >9</td>\n",
              "      <td id=\"T_bd9b9_row7_col3\" class=\"data row7 col3\" >W</td>\n",
              "      <td id=\"T_bd9b9_row7_col4\" class=\"data row7 col4\" >2.4</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row8_col0\" class=\"data row8 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row8_col1\" class=\"data row8 col1\" >H</td>\n",
              "      <td id=\"T_bd9b9_row8_col2\" class=\"data row8 col2\" >9</td>\n",
              "      <td id=\"T_bd9b9_row8_col3\" class=\"data row8 col3\" >F</td>\n",
              "      <td id=\"T_bd9b9_row8_col4\" class=\"data row8 col4\" >2.3</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td id=\"T_bd9b9_row9_col0\" class=\"data row9 col0\" >A</td>\n",
              "      <td id=\"T_bd9b9_row9_col1\" class=\"data row9 col1\" >V</td>\n",
              "      <td id=\"T_bd9b9_row9_col2\" class=\"data row9 col2\" >38</td>\n",
              "      <td id=\"T_bd9b9_row9_col3\" class=\"data row9 col3\" >M</td>\n",
              "      <td id=\"T_bd9b9_row9_col4\" class=\"data row9 col4\" >2.2</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": "e3cc2b1e-e7c9-4334-c05f-1aee854de814"
      },
      "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_f589180b-bf70-44f8-88a2-6a77089d2f85\", \"predictions.csv\", 32042)"
            ]
          },
          "metadata": {}
        }
      ]
    }
  ]
}