notebooks/06_global_analysis_of_protein_folding_mutation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import itertools\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import proteinsolver\n",
    "import psutil\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "import torch\n",
    "from kmbio import PDB\n",
    "from scipy import stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEBUG = \"CI\" not in os.environ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    %load_ext autoreload\n",
    "    %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "try:\n",
    "    inline_rc\n",
    "except NameError:\n",
    "    inline_rc = mpl.rcParams.copy()\n",
    "    \n",
    "mpl.rcParams.update({\"font.size\": 12})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNIQUE_ID = \"191f05de\"  # No attention\n",
    "# UNIQUE_ID = \"0007604c\"  # 5-layer graph-conv with attention, batch_size=1\n",
    "# UNIQUE_ID = \"91fc9ab9\"  # 4-layer graph-conv with attention, batch_size=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "BEST_STATE_FILES = {\n",
    "    #\n",
    "    \"191f05de\": \"protein_train/191f05de/e53-s1952148-d93703104.state\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('06_global_analysis_of_protein_folding')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NOTEBOOK_NAME = \"06_global_analysis_of_protein_folding\"\n",
    "NOTEBOOK_PATH = Path(NOTEBOOK_NAME)\n",
    "NOTEBOOK_PATH.mkdir(exist_ok=True)\n",
    "NOTEBOOK_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab1/database_data/datapkg_input_dir')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INPUT_PATH = Path(os.getenv(\"DATAPKG_INPUT_DIR\"))\n",
    "INPUT_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab1/database_data/datapkg_output_dir')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATAPKG_DATA_DIR = Path(f\"~/datapkg_data_dir\").expanduser().resolve()\n",
    "DATAPKG_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/kimlab1/database_data/datapkg_output_dir'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proteinsolver.settings.data_url = DATAPKG_DATA_DIR.as_posix()\n",
    "proteinsolver.settings.data_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "toc-hr-collapsed": false
   },
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aan0693_designed-PDB-files\t   hhpred\n",
      "aan0693_designed-PDB-files.gz.zip  hhpred2\n",
      "aan0693_SI_datasets\t\t   mutation_structures_for_rosetta\n",
      "aan0693_SI_datasets.tar.gz.zip\t   swissmodel\n"
     ]
    }
   ],
   "source": [
    "!ls {INPUT_PATH}/global_analysis_of_protein_folding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aan0693_designed-PDB-files.gz  nmr  other  other2  rd1\trd2  rd3  rd4\n"
     ]
    }
   ],
   "source": [
    "!ls {INPUT_PATH}/global_analysis_of_protein_folding/aan0693_designed-PDB-files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "counts_and_ec50s.tar.gz\t\t  protein_and_dna_sequences.tar.gz\n",
      "design_scripts.tgz\t\t  stability_scores\n",
      "design_structural_metrics\t  stability_scores.tar.gz\n",
      "design_structural_metrics.tar.gz  unfolded_state_model_params\n",
      "fig1_thermodynamic_data.csv\n"
     ]
    }
   ],
   "source": [
    "!ls {INPUT_PATH}/global_analysis_of_protein_folding/aan0693_SI_datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### aan0693_SI_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rd1_stability_scores  rd3_stability_scores  ssm2_stability_scores\n",
      "rd2_stability_scores  rd4_stability_scores\n"
     ]
    }
   ],
   "source": [
    "!ls {INPUT_PATH}/global_analysis_of_protein_folding/aan0693_SI_datasets/stability_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_controls(df):\n",
    "    df = df[\n",
    "        (~df[\"name\"].str.endswith(\"_hp\")) & (~df[\"name\"].str.endswith(\"_random\")) & (~df[\"name\"].str.endswith(\"_buryD\"))\n",
    "    ]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_stability_scores(key):\n",
    "    stability_scores = pd.read_csv(\n",
    "        INPUT_PATH\n",
    "        / \"global_analysis_of_protein_folding\"\n",
    "        / \"aan0693_SI_datasets\"\n",
    "        / \"stability_scores\"\n",
    "        / f\"{key}_stability_scores\",\n",
    "        sep=\"\\t\",\n",
    "    )\n",
    "    stability_scores = remove_controls(stability_scores)\n",
    "\n",
    "    for energy_function in [\"talaris2013\", \"betanov15\"]:\n",
    "        rosetta_energies_file = (\n",
    "            INPUT_PATH\n",
    "            / \"global_analysis_of_protein_folding\"\n",
    "            / \"aan0693_SI_datasets\"\n",
    "            / \"design_structural_metrics\"\n",
    "            / f\"{key}_relax_scored_{'filtered_' if energy_function == 'betanov15' else ''}{energy_function}.sc\"\n",
    "        )\n",
    "        if not rosetta_energies_file.is_file():\n",
    "            print(f\"Not loading Rosetta energies for {energy_function}!\")\n",
    "            continue\n",
    "\n",
    "        before_ = len(stability_scores)\n",
    "        relax_scored_filtered = pd.read_csv(\n",
    "            rosetta_energies_file, sep=\"\\t\" if energy_function == \"betanov15\" else \" +\", engine=\"python\"\n",
    "        ).rename(columns={\"description\": \"name\", \"total_score\": f\"{energy_function}_score\"})\n",
    "        stability_scores = stability_scores.merge(\n",
    "            relax_scored_filtered[[\"name\", f\"{energy_function}_score\"]], on=\"name\", how=\"outer\"\n",
    "        )\n",
    "#         assert len(stability_scores) == before_, (len(stability_scores), before_)\n",
    "\n",
    "    stability_scores[\"library_name\"] = key\n",
    "    return stability_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### stability_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stability_scores = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for key in [\"rd1\", \"rd2\", \"rd3\", \"rd4\", \"ssm2\"]:\n",
    "#     stability_scores[key] = load_stability_scores(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stability_scores[\"fig1\"] = pd.read_csv(\n",
    "#     INPUT_PATH / \"global_analysis_of_protein_folding\" / \"aan0693_SI_datasets\" / \"fig1_thermodynamic_data.csv\"\n",
    "# ).assign(library_name=\"fig1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores = torch.load(NOTEBOOK_PATH.joinpath(\"stability_scores.torch\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run protein_train/{UNIQUE_ID}/model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1\n",
    "num_features = 20\n",
    "adj_input_size = 2\n",
    "hidden_size = 128\n",
    "frac_present = 0.5\n",
    "frac_present_valid = frac_present\n",
    "info_size= 1024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'protein_train/191f05de/e53-s1952148-d93703104.state'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state_file = BEST_STATE_FILES[UNIQUE_ID]\n",
    "state_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net(\n",
    "    x_input_size=num_features + 1, adj_input_size=adj_input_size, hidden_size=hidden_size, output_size=num_features\n",
    ")\n",
    "net.load_state_dict(torch.load(state_file, map_location=device))\n",
    "net.eval()\n",
    "net = net.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mutation probabilities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "structure_file = INPUT_PATH.joinpath(\n",
    "    \"global_analysis_of_protein_folding\", \"aan0693_designed-PDB-files\", \"rd1\", \"EEHEE_rd1_0004.pdb\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "structure = PDB.load(structure_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdata = proteinsolver.utils.extract_seq_and_adj(structure, 'A')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 4.2337e+00,  1.4978e+00,  7.5491e+00,  9.7959e-02, -1.7546e+00,\n",
       "         2.8181e+00,  5.4245e+00,  1.3950e+00,  2.0784e+00,  3.1730e+00,\n",
       "         1.3765e+00,  6.6040e+00, -9.9698e-04,  4.2144e+00,  1.3128e+00,\n",
       "         6.9296e-02,  5.1866e+00,  6.3043e+00,  1.9462e+00,  1.1357e+00,\n",
       "         4.4304e+00,  6.7671e+00,  2.5967e+00,  1.5917e+00,  7.2876e-01,\n",
       "         5.6794e+00,  4.6740e+00,  2.3055e-01,  2.4577e-01,  7.3126e+00,\n",
       "         2.6852e+00,  5.3051e-01,  1.5237e+00,  4.0610e+00,  4.3447e-02,\n",
       "         3.3018e+00,  4.3885e+00,  3.5008e+00,  4.2183e+00,  2.1397e+00,\n",
       "        -2.2222e-01])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = proteinsolver.datasets.protein.row_to_data(pdata)\n",
    "data = proteinsolver.datasets.protein.transform_edge_attr(data)\n",
    "data.to(device)\n",
    "\n",
    "proteinsolver.utils.scan_with_mask(\n",
    "    net, data.x, data.edge_index, data.edge_attr, 20\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### mutation_proba_ssm2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mutation_proba_ssm2(net, filename, sequence):\n",
    "    if not filename.endswith(\".pdb\"):\n",
    "        filename = filename + \".pdb\"\n",
    "\n",
    "    try:\n",
    "        library_name = filename.split(\"_\")[1]\n",
    "        if library_name not in [\"rd1\", \"rd2\", \"rd3\", \"rd4\"]:\n",
    "            library_name = \"other\"\n",
    "    except IndexError:\n",
    "        library_name = \"other\"\n",
    "\n",
    "    filepath = (\n",
    "        INPUT_PATH / \"global_analysis_of_protein_folding\" / \"aan0693_designed-PDB-files\" / library_name / filename\n",
    "    )\n",
    "    if not filepath.is_file():\n",
    "        return np.nan\n",
    "\n",
    "    structure = PDB.load(filepath)\n",
    "    chain = \"X\" if library_name == \"other\" else \"A\"\n",
    "    pdata = proteinsolver.utils.extract_seq_and_adj(structure, chain)\n",
    "\n",
    "    if filename in [\"EHEE_rd2_0005.pdb\", \"EHEE_rd3_0015.pdb\"]:\n",
    "        # Skip because the number of amino acids in the PDB does not match the number of amino acids in the mutated sequence\n",
    "        return np.nan\n",
    "\n",
    "    if sum([(a == b) for a, b in zip(sequence, pdata.sequence)]) < (len(sequence) * 0.8):\n",
    "        print(f\"Warning, sequence does not match for protien {filename}\")\n",
    "        return np.nan\n",
    "\n",
    "    pdata = pdata._replace(sequence=sequence)\n",
    "\n",
    "    data = proteinsolver.datasets.protein.row_to_data(pdata)\n",
    "    data = proteinsolver.datasets.protein.transform_edge_attr(data)\n",
    "    data = data.to(device)\n",
    "\n",
    "    log_prob = proteinsolver.utils.scan_with_mask(\n",
    "        net, data.x, data.edge_index, data.edge_attr, 20\n",
    "    )\n",
    "\n",
    "    return log_prob.sum().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"ssm2\"][\"network_score\"] = [\n",
    "    get_mutation_proba_ssm2(net, filename, sequence)\n",
    "    for filename, sequence in stability_scores[\"ssm2\"][[\"my_wt\", \"sequence\"]].values\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### mutation_proba_fig1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mutation_proba_fig1(net, filename, sequence):\n",
    "    if not filename.endswith(\".pdb\"):\n",
    "        filename = filename + \".pdb\"\n",
    "\n",
    "    try:\n",
    "        library_name = filename.split(\"_\")[1]\n",
    "        if library_name not in [\"rd1\", \"rd2\", \"rd3\", \"rd4\"]:\n",
    "            library_name = \"other2\"\n",
    "    except IndexError:\n",
    "        library_name = \"other2\"\n",
    "\n",
    "    filepath = (\n",
    "        INPUT_PATH / \"global_analysis_of_protein_folding\" / \"aan0693_designed-PDB-files\" / library_name / filename\n",
    "    )\n",
    "    if not filepath.is_file():\n",
    "        return np.nan\n",
    "\n",
    "    structure = PDB.load(filepath)\n",
    "    chain = \"X\" if library_name == \"other2\" else \"A\"\n",
    "    pdata = proteinsolver.utils.extract_seq_and_adj(structure, chain)\n",
    "    if filename == \"villin.pdb\":\n",
    "        pdata = pdata._replace(sequence=pdata.sequence + \"SSSGSSGS\")\n",
    "    elif filename == \"BBL.pdb\":\n",
    "        pdata = pdata._replace(sequence=pdata.sequence + \"N\")\n",
    "\n",
    "    if len(pdata.sequence) != len(sequence):\n",
    "        if filename.startswith(\"villin\") or filename.startswith(\"BBL\"):\n",
    "            # Checked and those two are OK to truncate\n",
    "            pass\n",
    "        else:\n",
    "            print(\n",
    "                f\"Warning, sequence lengths do not match for protien {filename} ({len(pdata.sequence)}, {len(sequence)})\"\n",
    "            )\n",
    "            return np.nan\n",
    "    if sum([a == b for a, b in zip(sequence, pdata.sequence)]) < (len(sequence) * 0.6):\n",
    "        print(f\"Warning, sequences do not match for protien {filename}\")\n",
    "        return np.nan\n",
    "    pdata = pdata._replace(sequence=sequence)\n",
    "\n",
    "    data = proteinsolver.datasets.protein.row_to_data(pdata)\n",
    "    data = proteinsolver.datasets.protein.transform_edge_attr(data)\n",
    "    data = data.to(device)\n",
    "\n",
    "    node_outputs = proteinsolver.utils.scan_with_mask(\n",
    "        net, data.x, data.edge_index, data.edge_attr, 20\n",
    "    )\n",
    "\n",
    "    return node_outputs.mean().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"fig1\"][\"my_wt\"] = stability_scores[\"fig1\"][\"name\"].apply(\n",
    "    lambda name: {\"BBL\": \"BBL_H142W.pdb\"}.get(name.split(\"_\")[0], name.split(\"_\")[0] + \"_wt.pdb\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"fig1\"][\"network_score\"] = [\n",
    "    get_mutation_proba_fig1(net, filename, sequence)\n",
    "    for filename, sequence in stability_scores[\"fig1\"][[\"my_wt\", \"sequence\"]].values\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(stability_scores, NOTEBOOK_PATH.joinpath(\"stability_scores_for_mutations.raw.scan.torch\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add deltas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_diff(df, colname):\n",
    "    if f\"{colname}_wt\" not in df:\n",
    "        before_ = len(df)\n",
    "        ext = \".pdb\" if not df.iloc[0][\"name\"].endswith(\".pdb\") else \"\"\n",
    "        row_is_wt = (df[\"name\"] + ext) == df[\"my_wt\"]\n",
    "        assert set(df[row_is_wt][\"my_wt\"]) == set(df[\"my_wt\"])\n",
    "        df2 = df[row_is_wt][[\"my_wt\", colname]].rename(columns={colname: f\"{colname}_wt\"})\n",
    "        df = df.merge(df2, on=\"my_wt\")\n",
    "        assert before_ == len(df), (before_, len(df))\n",
    "    df[f\"{colname}_change\"] = df[f\"{colname}\"] - df[f\"{colname}_wt\"]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['name', 'sequence', 'my_wt', 'pos', 'mut', 'wt_aa', 'ec50_t',\n",
       "       'delta_ec50_t', 'ec50_95ci_lbound_t', 'ec50_95ci_ubound_t',\n",
       "       'ec50_95ci_t', 'ec50_pred_t', 'delta_pred_vs_wt_t', 'ec50_rise_t',\n",
       "       'stabilityscore_t', 'ec50_c', 'delta_ec50_c', 'ec50_95ci_lbound_c',\n",
       "       'ec50_95ci_ubound_c', 'ec50_95ci_c', 'ec50_pred_c',\n",
       "       'delta_pred_vs_wt_c', 'ec50_rise_c', 'stabilityscore_c',\n",
       "       'ec50_rise_c_adj', 'stabilityscore_c_adj', 'consensus_ec50_rise',\n",
       "       'consensus_stability_score', 'library_name', 'network_score'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stability_scores[\"ssm2\"].columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"ssm2\"][\"consensus_stability_score2\"] = (\n",
    "    stability_scores[\"ssm2\"][\"stabilityscore_t\"] + stability_scores[\"ssm2\"][\"stabilityscore_c\"]\n",
    ") / 2\n",
    "\n",
    "for column in [\n",
    "    \"network_score\",\n",
    "    \"stabilityscore_t\",\n",
    "    \"stabilityscore_c\",\n",
    "    \"consensus_stability_score\",\n",
    "    \"consensus_stability_score2\",\n",
    "]:\n",
    "    stability_scores[\"ssm2\"] = add_diff(stability_scores[\"ssm2\"], column)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores_wpreds = stability_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(stability_scores, NOTEBOOK_PATH.joinpath(\"stability_scores_form_mutations.logprob.incremental.torch\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## For Carles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_for_rosetta = stability_scores_wpreds[\"ssm2\"][\n",
    "    stability_scores_wpreds[\"ssm2\"][\"network_score_change\"].notnull()\n",
    "][[\"my_wt\", \"pos\", \"wt_aa\", \"mut\", \"sequence\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(mutations_for_rosetta[\"my_wt\"].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutations_for_rosetta.to_csv(\n",
    "    \"/home/kimlab1/database_data/datapkg_input_dir/global_analysis_of_protein_folding/mutation_structures_for_rosetta/gapf_mutations_for_rosetta.csv\",\n",
    "    index=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"fig1\"] = stability_scores[\"fig1\"].rename(\n",
    "    columns={\"deltaGunf thermal\": \"dg_thermal\", \"deltaGunf chemical\": \"dg_chemical\"}\n",
    ")\n",
    "\n",
    "stability_scores[\"fig1\"][\"consensus_stability_score\"] = (\n",
    "    stability_scores[\"fig1\"][\"stabilityscore_t\"] + stability_scores[\"fig1\"][\"stabilityscore_c\"]\n",
    ") / 2\n",
    "\n",
    "stability_scores[\"fig1\"][\"consensus_dg\"] = (\n",
    "    stability_scores[\"fig1\"][\"dg_thermal\"] + stability_scores[\"fig1\"][\"dg_chemical\"]\n",
    ") / 2\n",
    "\n",
    "for column in [\n",
    "    \"network_score\",\n",
    "    \"stabilityscore_t\",\n",
    "    \"stabilityscore_c\",\n",
    "    \"consensus_stability_score\",\n",
    "    \"dg_thermal\",\n",
    "    \"dg_chemical\",\n",
    "    \"consensus_dg\",\n",
    "    \"Tm\",\n",
    "]:\n",
    "    stability_scores[\"fig1\"] = add_diff(stability_scores[\"fig1\"], column)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate correlations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_conf_interval(r, num):\n",
    "    import math\n",
    "\n",
    "    stderr = 1.0 / math.sqrt(num - 3)\n",
    "    delta = 1.96 * stderr\n",
    "    lower = math.tanh(math.atanh(r) - delta)\n",
    "    upper = math.tanh(math.atanh(r) + delta)\n",
    "    return r - lower, upper - r "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping EHEE_0882.pdb!\n",
      "Skipping EHEE_rd2_0005.pdb!\n",
      "Skipping EHEE_rd3_0015.pdb!\n",
      "Skipping HHH_0142.pdb!\n"
     ]
    }
   ],
   "source": [
    "x_col = \"network_score_change\"\n",
    "y_col = \"consensus_stability_score_change\"\n",
    "library = \"ssm2\"\n",
    "\n",
    "correlations = []\n",
    "\n",
    "for my_wt, gp in stability_scores[library].groupby(\"my_wt\"):\n",
    "    x = gp[[x_col, y_col]].dropna()\n",
    "    if x.empty:\n",
    "        print(f\"Skipping {my_wt}!\")\n",
    "        continue\n",
    "    corr = stats.spearmanr(x[x_col], x[y_col])\n",
    "    corr_conf = get_conf_interval(corr[0], len(x))\n",
    "    correlations.append(\n",
    "        {\"my_wt\": my_wt, \"corr\": corr[0], \"pvalue\": corr[1], \"corr_lower_bound\": corr_conf[0], \"corr_upper_bound\": corr_conf[1]}\n",
    "    )\n",
    "correlations_df = pd.DataFrame(correlations, columns=correlations[0].keys())\n",
    "#     plt.title(f\"Spearman R: {corr[0]:.4f} ({corr[1]:.3f})\")\n",
    "#     plt.plot(x[x_col], x[y_col], \"r.\")\n",
    "#     plt.suptitle(my_wt)\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>my_wt</th>\n",
       "      <th>corr</th>\n",
       "      <th>pvalue</th>\n",
       "      <th>corr_lower_bound</th>\n",
       "      <th>corr_upper_bound</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>EEHEE_rd3_0037.pdb</td>\n",
       "      <td>0.394147</td>\n",
       "      <td>3.306064e-30</td>\n",
       "      <td>0.061183</td>\n",
       "      <td>0.057878</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>EEHEE_rd3_1498.pdb</td>\n",
       "      <td>0.393333</td>\n",
       "      <td>4.440973e-30</td>\n",
       "      <td>0.061226</td>\n",
       "      <td>0.057925</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>EEHEE_rd3_1702.pdb</td>\n",
       "      <td>0.149497</td>\n",
       "      <td>2.932858e-05</td>\n",
       "      <td>0.069584</td>\n",
       "      <td>0.068134</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>EEHEE_rd3_1716.pdb</td>\n",
       "      <td>0.434199</td>\n",
       "      <td>5.645419e-37</td>\n",
       "      <td>0.058951</td>\n",
       "      <td>0.055452</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>HEEH_rd2_0779.pdb</td>\n",
       "      <td>0.480017</td>\n",
       "      <td>6.493175e-46</td>\n",
       "      <td>0.056094</td>\n",
       "      <td>0.052426</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>HEEH_rd3_0223.pdb</td>\n",
       "      <td>0.393417</td>\n",
       "      <td>4.307661e-30</td>\n",
       "      <td>0.061221</td>\n",
       "      <td>0.057920</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>HEEH_rd3_0726.pdb</td>\n",
       "      <td>0.256082</td>\n",
       "      <td>4.534633e-13</td>\n",
       "      <td>0.067015</td>\n",
       "      <td>0.064641</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>HEEH_rd3_0872.pdb</td>\n",
       "      <td>0.237012</td>\n",
       "      <td>2.342070e-11</td>\n",
       "      <td>0.067597</td>\n",
       "      <td>0.065378</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>HHH_rd2_0134.pdb</td>\n",
       "      <td>0.295338</td>\n",
       "      <td>4.596301e-17</td>\n",
       "      <td>0.065648</td>\n",
       "      <td>0.062973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>HHH_rd3_0138.pdb</td>\n",
       "      <td>0.429652</td>\n",
       "      <td>3.688339e-36</td>\n",
       "      <td>0.059216</td>\n",
       "      <td>0.055738</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Pin1</td>\n",
       "      <td>0.434542</td>\n",
       "      <td>9.606567e-34</td>\n",
       "      <td>0.061974</td>\n",
       "      <td>0.058115</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>hYAP65</td>\n",
       "      <td>0.442019</td>\n",
       "      <td>5.715900e-41</td>\n",
       "      <td>0.056488</td>\n",
       "      <td>0.053187</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>villin</td>\n",
       "      <td>0.533106</td>\n",
       "      <td>1.275038e-47</td>\n",
       "      <td>0.058296</td>\n",
       "      <td>0.053639</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 my_wt      corr        pvalue  corr_lower_bound  \\\n",
       "0   EEHEE_rd3_0037.pdb  0.394147  3.306064e-30          0.061183   \n",
       "1   EEHEE_rd3_1498.pdb  0.393333  4.440973e-30          0.061226   \n",
       "2   EEHEE_rd3_1702.pdb  0.149497  2.932858e-05          0.069584   \n",
       "3   EEHEE_rd3_1716.pdb  0.434199  5.645419e-37          0.058951   \n",
       "4    HEEH_rd2_0779.pdb  0.480017  6.493175e-46          0.056094   \n",
       "5    HEEH_rd3_0223.pdb  0.393417  4.307661e-30          0.061221   \n",
       "6    HEEH_rd3_0726.pdb  0.256082  4.534633e-13          0.067015   \n",
       "7    HEEH_rd3_0872.pdb  0.237012  2.342070e-11          0.067597   \n",
       "8     HHH_rd2_0134.pdb  0.295338  4.596301e-17          0.065648   \n",
       "9     HHH_rd3_0138.pdb  0.429652  3.688339e-36          0.059216   \n",
       "10                Pin1  0.434542  9.606567e-34          0.061974   \n",
       "11              hYAP65  0.442019  5.715900e-41          0.056488   \n",
       "12              villin  0.533106  1.275038e-47          0.058296   \n",
       "\n",
       "    corr_upper_bound  \n",
       "0           0.057878  \n",
       "1           0.057925  \n",
       "2           0.068134  \n",
       "3           0.055452  \n",
       "4           0.052426  \n",
       "5           0.057920  \n",
       "6           0.064641  \n",
       "7           0.065378  \n",
       "8           0.062973  \n",
       "9           0.055738  \n",
       "10          0.058115  \n",
       "11          0.053187  \n",
       "12          0.053639  "
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "correlations_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'category_names' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-48-77db8755ad43>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m category_colors = {\n\u001b[0;32m----> 2\u001b[0;31m     \u001b[0;34m**\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mcmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mcategory_names\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstartswith\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"mae\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstartswith\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"rmsd\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;34m**\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mcmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mcategory_names\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstartswith\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"mae\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstartswith\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"rmsd\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m }\n",
      "\u001b[0;31mNameError\u001b[0m: name 'category_names' is not defined"
     ]
    }
   ],
   "source": [
    "category_colors = {\n",
    "    **{c: cmap(0) for c in category_names if c.startswith(\"mae\") or c.startswith(\"rmsd\")},\n",
    "    **{c: cmap(1) for c in category_names if not c.startswith(\"mae\") and not c.startswith(\"rmsd\")},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_spider_plot(title, row, color, ax):\n",
    "    other_color = \"k\"\n",
    "\n",
    "    yticks = np.array([0.2, 0.4, 0.6, 0.8, 1.0])\n",
    "\n",
    "    # Get a list of categories\n",
    "    categories = row.index.values.tolist()\n",
    "\n",
    "    # Set the angles of each category\n",
    "    angles = np.linspace(0, 2 * np.pi, len(categories) + 1)\n",
    "\n",
    "    # Put the first category on top and going clockwise\n",
    "    ax.set_theta_offset(np.pi / 2)\n",
    "    ax.set_theta_direction(-1)\n",
    "\n",
    "    labels = []\n",
    "    for c in categories:\n",
    "        cs = c.rsplit(\"_\", 1)\n",
    "        if len(cs) == 1:\n",
    "            labels.append(c)\n",
    "        else:\n",
    "            labels.append(f\"{cs[0]}\\n_{cs[1]}\")\n",
    "\n",
    "    # Configure X axis\n",
    "    ax.set_xticks(angles[:-1])\n",
    "    ax.set_xticklabels(labels, size=10)\n",
    "    for c, xtick in zip(labels, ax.get_xticklabels()):\n",
    "        xtick.set_color(other_color)\n",
    "\n",
    "    # Configure Y axis\n",
    "    ax.set_rlabel_position(0)\n",
    "    ax.set_yticks(yticks)\n",
    "    ax.set_yticklabels([f\"{v:.1f}\" for v in yticks], color=other_color)\n",
    "    ax.set_ylim(0, 0.7)\n",
    "\n",
    "    # Configure grid\n",
    "    ax.grid(color=\"k\", linestyle=\":\", linewidth=0.5)\n",
    "\n",
    "    # Plot values\n",
    "    ax.plot(angles, np.r_[df[\"corr\"], df[\"corr\"][0]], color=color, linewidth=2, linestyle=\"solid\")\n",
    "    ax.errorbar(\n",
    "        angles,\n",
    "        np.r_[df[\"corr\"], df[\"corr\"][0]],\n",
    "        color=color,\n",
    "#         ecolor=other_color,\n",
    "#         linewidth=2,\n",
    "#         linestyle=\"solid\",\n",
    "        yerr=np.c_[df[[\"corr_upper_bound\", \"corr_lower_bound\"]].values.T, df[[\"corr_upper_bound\", \"corr_lower_bound\"]].values.T[:, :1]],\n",
    "    )\n",
    "    ax.fill(angles, np.r_[df[\"corr\"], df[\"corr\"][0]], color=color, alpha=0.4)\n",
    "\n",
    "    # Add a title\n",
    "    if title:\n",
    "        ax.set_title(title, size=12, pad=20)\n",
    "\n",
    "\n",
    "#     add_yticks(ax, angles[len(angles) // 2], yticks, yticks, color=category_colors[categories[3]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = correlations_df.assign(my_wt=correlations_df[\"my_wt\"].str.replace(\".pdb\", \"\")).set_index(\"my_wt\")\n",
    "df = pd.concat([df.iloc[-1:], df.iloc[:-1]])\n",
    "\n",
    "cmap = plt.cm.get_cmap(\"Set1\", 8)\n",
    "\n",
    "fg, ax = plt.subplots(figsize=(4.5, 4.5), subplot_kw=dict(polar=True))\n",
    "make_spider_plot(\"\", df, cmap(0), ax=ax)\n",
    "fg.tight_layout()\n",
    "fg.savefig(NOTEBOOK_PATH.joinpath(\"correlations_with_mutations.svg\"))\n",
    "fg.savefig(NOTEBOOK_PATH.joinpath(\"correlations_with_mutations.pdf\"))\n",
    "fg.savefig(NOTEBOOK_PATH.joinpath(\"correlations_with_mutations.svg\"), dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "correlations_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "correlations_df.sort_values(\"my_wt\", ascending=False)[\"my_wt\"].str.replace(\".pdb\", \"\")).set_index(\"my_wt\")[\"corr\"].in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_col = \"consensus_stability_score_change\"\n",
    "y_col = \"consensus_dg_change\"\n",
    "\n",
    "x = stability_scores[\"fig1\"][[x_col, y_col]].dropna()\n",
    "\n",
    "corr = stats.spearmanr(x[x_col], x[y_col])\n",
    "\n",
    "plt.title(f\"Spearman R: {corr[0]:.4f} ({corr[1]:.3f})\")\n",
    "plt.plot(x[x_col], x[y_col], \"r.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"consensus_stability_score_wt\" not in stability_scores[\"ssm2\"]:\n",
    "    before_ = len(stability_scores[\"ssm2\"])\n",
    "    row_is_wt = stability_scores[\"ssm2\"][\"name\"] == stability_scores[\"ssm2\"][\"my_wt\"]\n",
    "    df = stability_scores[\"ssm2\"][row_is_wt][[\"my_wt\", \"consensus_stability_score\"]].rename(\n",
    "        columns={\"consensus_stability_score\": \"consensus_stability_score_wt\"}\n",
    "    )\n",
    "    stability_scores[\"ssm2\"] = stability_scores[\"ssm2\"].merge(df, on=\"my_wt\")\n",
    "    assert before_ == len(stability_scores[\"ssm2\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"ssm2\"][\"consensus_stability_score_change\"] = (\n",
    "    stability_scores[\"ssm2\"][\"consensus_stability_score\"] - stability_scores[\"ssm2\"][\"consensus_stability_score_wt\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = stability_scores[\"ssm2\"][[\"sum_log_prob_change\", \"consensus_stability_score_change\"]].dropna()\n",
    "\n",
    "corr = stats.spearmanr(\n",
    "    x[\"sum_log_prob_change\"],\n",
    "    x[\"consensus_stability_score_change\"],\n",
    ")\n",
    "\n",
    "fg, ax = plt.subplots()\n",
    "\n",
    "ax.plot(\n",
    "    x[\"sum_log_prob_change\"],\n",
    "    x[\"consensus_stability_score_change\"],\n",
    "    'r.',\n",
    "    alpha=0.5\n",
    ")\n",
    "ax.set_title(f\"Spearman R: {corr[0]:.4f} ({corr[1]:.3f})\")\n",
    "fg.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def myplot(x, y, color, ax):\n",
    "    x, y = pd.DataFrame({'x': x, 'y': y}).copy().dropna().values.T\n",
    "    \n",
    "    corr = stats.spearmanr(x, y)\n",
    "    ax.plot(x, y, color=color)\n",
    "    ax.set_title(f\"N: {len(x):,};   Spearman R: {corr[0]:.3f} ({corr[1]:.2f})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for library_name in list(stability_scores)[:3]:\n",
    "    df = stability_scores[library_name][\n",
    "        (stability_scores[library_name][\"sum_log_prob\"].isnull())\n",
    "        & ((stability_scores[library_name][\"talaris2013_score\"].notnull()))\n",
    "    ]\n",
    "    print(library_name, len(df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for library_name in list(stability_scores)[:3]:\n",
    "    df = stability_scores[library_name][\n",
    "        (stability_scores[library_name][\"sum_log_prob\"].notnull())\n",
    "        & ((stability_scores[library_name][\"talaris2013_score\"].isnull()))\n",
    "    ]\n",
    "    print(library_name, len(df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for library_name in list(stability_scores)[:3]:\n",
    "    df = stability_scores[library_name][\n",
    "        ((stability_scores[library_name][\"talaris2013_score\"].isnull()))\n",
    "    ]\n",
    "    print(library_name, len(df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stability_scores[\"rd1\"].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "\n",
    "for library in [\"rd1\", \"rd2\", \"rd3\", \"rd4\"]:\n",
    "    df = stability_scores[library].dropna(subset=[\"sum_log_prob\", \"talaris2013_score\", \"betanov15_score\"])\n",
    "    df[\"domain\"] = df[\"name\"].str.split(\"_\").str[0]\n",
    "    for domain, gp in df.groupby(\"domain\"):\n",
    "        for score in [\"sum_log_prob\", \"talaris2013_score\", \"betanov15_score\"]:\n",
    "            corr_t = stats.spearmanr(gp[score], gp[\"stabilityscore_t\"])\n",
    "            corr_c = stats.spearmanr(gp[score], gp[\"stabilityscore_c\"])\n",
    "            results.append((library, domain, score, corr_t[0], corr_t[1], corr_c[0], corr_c[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_col = \"sum_log_prob_normed\"\n",
    "df = stability_scores[library_name].dropna(subset=[\"sum_log_prob\", \"talaris2013_score\", \"betanov15_score\"])\n",
    "\n",
    "fg, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "myplot(df[x_col], df[\"stabilityscore_t\"], ax=axs[0])\n",
    "myplot(df[x_col], df[\"stabilityscore_c\"], ax=axs[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_col = \"talaris2013_score_normed\"\n",
    "df = stability_scores[library_name].dropna(subset=[\"sum_log_prob\", \"talaris2013_score\", \"betanov15_score\"])\n",
    "\n",
    "fg, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "myplot(df[x_col], df[\"stabilityscore_t\"], ax=axs[0])\n",
    "myplot(df[x_col], df[\"stabilityscore_c\"], ax=axs[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_col = \"betanov15_score_normed\"\n",
    "df = stability_scores[library_name].dropna(subset=[\"sum_log_prob\", \"talaris2013_score\", \"betanov15_score\"])\n",
    "\n",
    "fg, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "myplot(df[x_col], df[\"stabilityscore_t\"], ax=axs[0])\n",
    "myplot(df[x_col], df[\"stabilityscore_c\"], ax=axs[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_col = \"talaris2013_score_normed\"\n",
    "df = stability_scores[\"rd1\"]\n",
    "\n",
    "fg, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "myplot(stability_scores[\"rd1\"][x_col], stability_scores[\"rd1\"][\"stabilityscore_t\"], ax=axs[0])\n",
    "myplot(stability_scores[\"rd1\"][x_col], stability_scores[\"rd1\"][\"stabilityscore_c\"], ax=axs[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = stability_scores[\"rd1\"]\n",
    "\n",
    "fg, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "axs[0].plot(-df[\"talaris2013_score\"] / df[\"sequence_lenght\"], df[\"stabilityscore_t\"], \"r.\")\n",
    "axs[1].plot(-df[\"talaris2013_score\"] / df[\"sequence_lenght\"], df[\"stabilityscore_c\"], \"r.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = stability_scores[\"rd1\"]\n",
    "\n",
    "fg, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "axs[0].plot(-df[\"betanov15_score\"] / df[\"sequence_lenght\"], df[\"stabilityscore_t\"], \"r.\")\n",
    "axs[1].plot(-df[\"betanov15_score\"] / df[\"sequence_lenght\"], df[\"stabilityscore_c\"], \"r.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_structure_proba(net, \"rd1\", \"EEHEE_rd1_0001.pdb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rd3_stability_scores[\"sequence_lenght\"] = [\n",
    "    get_sequence_length(net, rd3_structure_path / name) for name in rd3_stability_scores[\"name\"]\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(rd3_stability_scores))\n",
    "df = rd3_stability_scores.dropna(how='any').copy()\n",
    "print(len(df))\n",
    "\n",
    "stats.spearmanr(df[\"sum_log_proba\"] / df['sequence_lenght'], df[\"stabilityscore_t\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(df[\"sum_log_proba\"] / df['sequence_lenght'], df[\"stabilityscore_t\"], 'r.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['type'] = df['name'].str.split(\"_\").str[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for type, gp in df.groupby(\"type\"):\n",
    "    print(type)\n",
    "    c = \"ec50_rise_t\"\n",
    "    print(stats.spearmanr(gp[\"sum_log_proba\"] / gp['sequence_lenght'], gp[c]))\n",
    "    plt.plot(gp[\"sum_log_proba\"] / gp['sequence_lenght'], gp[c], 'r.')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}