output_pdbs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Output PDBs with escape scores as B factors\n",
    "This Python Jupyter notebook outputs PDBs with the escape scores as B factors.\n",
    "\n",
    "Though we will want more elaborate series of commands to codify our visualization of these RBD structures colored by escape, the series of commands below, when executed in a `PyMol` session with one of these PDBs open, will color the RBD surface according to escape scores.\n",
    "\n",
    "For example, to normalize each structure colored by the max mut effect, we might want to have a white to red scale from 0 to 1:\n",
    "\n",
    "     create RBD, chain E\n",
    "     hide all; show cartoon, chain A; color gray20, chain A\n",
    "     show surface, RBD; spectrum b, white red, RBD, minimum=0, maximum=1\n",
    "     \n",
    "For something like total escape, maybe we want each structure normalized to the maximum total escape in that structure, in which case we can just leave the maximum argument empty:\n",
    "\n",
    "     create RBD, chain E\n",
    "     hide all; show cartoon, chain A; color gray20, chain A\n",
    "     show surface, RBD; spectrum b, white red, RBD, minimum=0\n",
    "     \n",
    "We write PDBs with B factors indicating the total site escape and maximum mutation escape at each site, and the same with these values normalized to the maximum for the full structure (the latter are easier to process in `Chimera`).\n",
    "\n",
    "First, import Python modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import copy\n",
    "import os\n",
    "import warnings\n",
    "\n",
    "import Bio.PDB\n",
    "\n",
    "import dms_variants.pdb_utils\n",
    "\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read the configuration file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('config.yaml') as f:\n",
    "    config = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read configuration for outputting PDBs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading PDB output configuration from data/output_pdbs_config.yaml\n"
     ]
    }
   ],
   "source": [
    "print(f\"Reading PDB output configuration from {config['output_pdbs_config']}\")\n",
    "with open(config['output_pdbs_config']) as f:\n",
    "    output_pdbs_config = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make output directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(config['pdb_outputs_dir'], exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read escape fractions and compute **total** and **maximum** escape at each site, and also the total and maximum escape at each site normalized to be between 0 and 1 for each selection:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading escape fractions from results/escape_scores/escape_fracs.csv\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>selection</th>\n",
       "      <th>site</th>\n",
       "      <th>total_escape</th>\n",
       "      <th>max_escape</th>\n",
       "      <th>max_total_escape</th>\n",
       "      <th>max_max_escape</th>\n",
       "      <th>norm_total_escape</th>\n",
       "      <th>norm_max_escape</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>CAB-A17_139</td>\n",
       "      <td>331</td>\n",
       "      <td>0.016274</td>\n",
       "      <td>0.001090</td>\n",
       "      <td>0.577178</td>\n",
       "      <td>0.5368</td>\n",
       "      <td>0.028196</td>\n",
       "      <td>0.002031</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>CAB-A17_139</td>\n",
       "      <td>332</td>\n",
       "      <td>0.020121</td>\n",
       "      <td>0.002068</td>\n",
       "      <td>0.577178</td>\n",
       "      <td>0.5368</td>\n",
       "      <td>0.034861</td>\n",
       "      <td>0.003852</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>CAB-A17_139</td>\n",
       "      <td>333</td>\n",
       "      <td>0.017871</td>\n",
       "      <td>0.001585</td>\n",
       "      <td>0.577178</td>\n",
       "      <td>0.5368</td>\n",
       "      <td>0.030964</td>\n",
       "      <td>0.002953</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>CAB-A17_139</td>\n",
       "      <td>334</td>\n",
       "      <td>0.018437</td>\n",
       "      <td>0.001169</td>\n",
       "      <td>0.577178</td>\n",
       "      <td>0.5368</td>\n",
       "      <td>0.031944</td>\n",
       "      <td>0.002178</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>CAB-A17_139</td>\n",
       "      <td>335</td>\n",
       "      <td>0.018819</td>\n",
       "      <td>0.001354</td>\n",
       "      <td>0.577178</td>\n",
       "      <td>0.5368</td>\n",
       "      <td>0.032605</td>\n",
       "      <td>0.002522</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(f\"Reading escape fractions from {config['escape_fracs']}\")\n",
    "\n",
    "escape_fracs = (\n",
    "    pd.read_csv(config['escape_fracs'])\n",
    "    .query('library == \"average\"')\n",
    "    .assign(site=lambda x: x['label_site'])\n",
    "    .groupby(['selection', 'site'])\n",
    "    .aggregate(total_escape=pd.NamedAgg(config['mut_metric'], 'sum'),\n",
    "               max_escape=pd.NamedAgg(config['mut_metric'], 'max')\n",
    "               )\n",
    "    .reset_index()\n",
    "    .assign(max_total_escape=lambda x: x.groupby('selection')['total_escape'].transform('max'),\n",
    "            max_max_escape=lambda x: x.groupby('selection')['max_escape'].transform('max'),\n",
    "            norm_total_escape=lambda x: x['total_escape'] / x['max_total_escape'],\n",
    "            norm_max_escape=lambda x: x['max_escape'] / x['max_max_escape'])\n",
    "    )\n",
    "\n",
    "display(HTML(escape_fracs.head().to_html(index=False)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now map the escape metrics to the B-factors.\n",
    "For sites where no mutations have escape scores:\n",
    " - In the RBD chain(s) fill the B-factor for non-normalized scores to -1 to enable collapsing to zero or callout as a a separate class, depending how we choose to color sites for different visualizations. For normalized scores, fill to 0.\n",
    " - In other chains, always fill missing B factors to 0.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Making PDB mappings for 6m0j to data/pdbs/6M0J.pdb\n",
      "Making mappings for 3 conditions.\n",
      "Mapping to the following chains: E\n",
      "  Writing B-factor re-assigned PDBs for CAB-A17_139 to:\n",
      "    results/pdb_outputs/CAB-A17_139_6m0j_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-A17_139_6m0j_max_escape.pdb\n",
      "    results/pdb_outputs/CAB-A17_139_6m0j_norm_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-A17_139_6m0j_norm_max_escape.pdb\n",
      "  Writing B-factor re-assigned PDBs for CAB-A49_222 to:\n",
      "    results/pdb_outputs/CAB-A49_222_6m0j_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-A49_222_6m0j_max_escape.pdb\n",
      "    results/pdb_outputs/CAB-A49_222_6m0j_norm_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-A49_222_6m0j_norm_max_escape.pdb\n",
      "  Writing B-factor re-assigned PDBs for CAB-C19_184 to:\n",
      "    results/pdb_outputs/CAB-C19_184_6m0j_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-C19_184_6m0j_max_escape.pdb\n",
      "    results/pdb_outputs/CAB-C19_184_6m0j_norm_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-C19_184_6m0j_norm_max_escape.pdb\n",
      "\n",
      "Making PDB mappings for BA1_A17 to data/pdbs/BA1_A17.pdb\n",
      "Making mappings for 1 conditions.\n",
      "Mapping to the following chains: A\n",
      "  Writing B-factor re-assigned PDBs for CAB-A17_139 to:\n",
      "    results/pdb_outputs/CAB-A17_139_BA1_A17_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-A17_139_BA1_A17_max_escape.pdb\n",
      "    results/pdb_outputs/CAB-A17_139_BA1_A17_norm_total_escape.pdb\n",
      "    results/pdb_outputs/CAB-A17_139_BA1_A17_norm_max_escape.pdb\n"
     ]
    }
   ],
   "source": [
    "for name, specs in output_pdbs_config.items():\n",
    "    print(f\"\\nMaking PDB mappings for {name} to {specs['pdbfile']}\")\n",
    "    assert os.path.isfile(specs['pdbfile'])\n",
    "    \n",
    "    # get escape fracs just for conditions of interest\n",
    "    if isinstance(specs['conditions'], str) and specs['conditions'].upper() == 'ALL':\n",
    "        conditions = escape_fracs['selection'].unique().tolist()\n",
    "    else:\n",
    "        assert isinstance(specs['conditions'], list)\n",
    "        conditions = specs['conditions']\n",
    "    print(f\"Making mappings for {len(conditions)} conditions.\")\n",
    "    df = escape_fracs.query('selection in @conditions')\n",
    "    \n",
    "    # get chains\n",
    "    assert isinstance(specs['chains'], list)\n",
    "    print('Mapping to the following chains: ' + ', '.join(specs['chains']))\n",
    "    df = pd.concat([df.assign(chain=chain) for chain in specs['chains']], ignore_index=True)\n",
    "    \n",
    "    # make mappings for each condition and metric\n",
    "    for condition, df in df.groupby('selection'):\n",
    "        print(f\"  Writing B-factor re-assigned PDBs for {condition} to:\")\n",
    "    \n",
    "        for metric in ['total_escape', 'max_escape', 'norm_total_escape', 'norm_max_escape']:\n",
    "        \n",
    "            # what do we assign to missing sites?\n",
    "            missing_metric = collections.defaultdict(lambda: 0)  # non-RBD chains always fill to zero\n",
    "            for chain in specs['chains']:\n",
    "                if 'norm' in metric:\n",
    "                    missing_metric[chain] = 0  # missing sites in RBD are 0 for normalized metric PDBs\n",
    "                else:\n",
    "                    missing_metric[chain] = -1  # missing sites in RBD are -1 for non-normalized metric PDBs\n",
    "        \n",
    "            fname = os.path.join(config['pdb_outputs_dir'], f\"{condition}_{name}_{metric}.pdb\")\n",
    "            print(f\"    {fname}\")\n",
    "            \n",
    "            dms_variants.pdb_utils.reassign_b_factor(input_pdbfile=specs['pdbfile'],\n",
    "                                                     output_pdbfile=fname,\n",
    "                                                     df=df,\n",
    "                                                     metric_col=metric,\n",
    "                                                     missing_metric=missing_metric)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}