analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook contains code and analysis for the simulations of affinity maturation presented in\n",
    "\n",
    "### Role of framework mutations and antibody flexibility in the evolution of broadly neutralizing antibodies  \n",
    "\n",
    "Victor Ovchinnikov$^{1,*}$, Joy E. Louveau$^{2,*}$, John P. Barton$^{3-6,*,†}$, Martin Karplus$^{1,7,\\#}$, and Arup K. Chakraborty$^{3-6,8,9,\\#}$  \n",
    "\n",
    "$^1$ Harvard University, Department of Chemistry and Chemical Biology (Cambridge, MA, United States)  \n",
    "$^2$ Massachusetts Institute of Technology, Harvard-MIT Division of Health Sciences and Technology (Cambridge, MA, United States)  \n",
    "$^3$ Massachusetts Institute of Technology, Department of Chemical Engineering (Cambridge, MA, United States)  \n",
    "$^4$ Massachusetts Institute of Technology, Department of Physics (Cambridge, MA, United States)  \n",
    "$^5$ Massachusetts Institute of Technology, Institute for Medical Engineering and Science (Cambridge, MA, United States)  \n",
    "$^6$ Ragon Institute of MGH, MIT and Harvard (Cambridge, MA, United States)  \n",
    "$^7$ Université de Strasbourg, Laboratoire de Chimie Biophysique, ISIS (Strasbourg, France)  \n",
    "$^8$ Massachusetts Institute of Technology, Department of Biological Engineering (Cambridge, MA, United States)  \n",
    "$^9$ Massachusetts Institute of Technology, Department of Chemistry (Cambridge, MA, United States)  \n",
    "$^*$ these authors contributed equally  \n",
    "$^\\#$ correspondence to arupc@mit.edu or marci@tammy.harvard.edu  \n",
    "$^†$ present address: University of California, Riverside, Department of Physics and Astronomy (Riverside, CA, United States)\n",
    "\n",
    "# Contents\n",
    "\n",
    "In addition to this notebook, this repository contains\n",
    "\n",
    "`am.py` - Python code used for affinity maturation simulations  \n",
    "`mplot.py` - a Python library for generating plots through matplotlib, see the GitHub repository [here](https://github.com/johnbarton/mplot)  \n",
    "`/data` - a folder containing output from the affinity maturation simulations and other analysis in this notebook  \n",
    "`/figures` - a folder containings figures generated in this notebook\n",
    "\n",
    "This notebook contains the following sections\n",
    "1. Loading libraries/functions\n",
    "2. Data analysis and plotting figures\n",
    "3. Recording of parameters for `am.py` used in alternative simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load necessary libraries and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This notebook was prepared using:\n",
      "python version 3.6.4 (default, Jan  6 2018, 11:51:59) \n",
      "[GCC 4.2.1 Compatible Apple LLVM 9.0.0 (clang-900.0.39.2)]\n",
      "numpy version 1.14.0\n",
      "pandas version 0.22.0\n",
      "matplotlib version 2.1.2\n"
     ]
    }
   ],
   "source": [
    "# Full library list and version numbers\n",
    "\n",
    "print('This notebook was prepared using:')\n",
    "\n",
    "import sys\n",
    "print('python version %s' % sys.version)\n",
    "\n",
    "import numpy as np\n",
    "print('numpy version %s' % np.__version__)\n",
    "\n",
    "import pandas as pd\n",
    "print('pandas version %s' % pd.__version__)\n",
    "\n",
    "import matplotlib\n",
    "print('matplotlib version %s' % matplotlib.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure settings and libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import deepcopy\n",
    "from importlib import reload\n",
    "\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.pyplot as plot\n",
    "import matplotlib.gridspec as gridspec\n",
    "import matplotlib.image as mpimg\n",
    "\n",
    "import mplot as mp\n",
    "\n",
    "# Plot conventions\n",
    "\n",
    "def cm2inch(x): return float(x)/2.54\n",
    "single_column   = cm2inch(8.5)\n",
    "one_five_column = cm2inch(11.4)\n",
    "double_column   = cm2inch(17.4)\n",
    "\n",
    "goldr        = (1.0 + np.sqrt(5)) / 2.0\n",
    "bkcolor      = '#252525'\n",
    "altcolor     = '#3182bd'\n",
    "fontfamily   = 'Arial'\n",
    "sizeline     = 0.6\n",
    "sizesublabel = 9\n",
    "sizelabel    = 8\n",
    "sizetick     = 7\n",
    "ticklength   = 3\n",
    "tickpad      = 3\n",
    "axwidth      = 0.4\n",
    "\n",
    "figprops = {\n",
    "    'transparent' : True,\n",
    "    'bbox_inches' : 'tight'\n",
    "    }\n",
    "\n",
    "def_barprops = {\n",
    "    'lw'          : sizeline/2,\n",
    "    'width'       : 0.25,\n",
    "    'edgecolor'   : bkcolor,\n",
    "    'align'       : 'center', #other option: edge\n",
    "    'orientation' : 'vertical'\n",
    "    }\n",
    "\n",
    "def_histprops = {\n",
    "    'histtype'    : 'bar',\n",
    "    'lw'          : sizeline/2,\n",
    "    'rwidth'      : 0.8,\n",
    "    'ls'          : 'solid',\n",
    "    'edgecolor'   : bkcolor,\n",
    "    'alpha'       : 0.5\n",
    "    }\n",
    "\n",
    "def_lineprops = {\n",
    "    'lw' : sizeline, \n",
    "    'ls' : '-'\n",
    "    }\n",
    "\n",
    "def_labelprops = {\n",
    "    'family' : fontfamily,\n",
    "    'size'   : sizelabel,\n",
    "    'color'  : bkcolor\n",
    "    }\n",
    "\n",
    "def_sublabelprops = {\n",
    "    'family' : fontfamily,\n",
    "    'size'   : sizesublabel,\n",
    "    'weight' : 'bold',\n",
    "    'ha'     : 'center',\n",
    "    'va'     : 'center',\n",
    "    'color'  : 'k'\n",
    "    }\n",
    "\n",
    "def_ticklabelprops = {\n",
    "    'family' : fontfamily,\n",
    "    'size'   : sizetick,\n",
    "    'color'  : bkcolor\n",
    "    }\n",
    "\n",
    "def_tickprops = {\n",
    "    'length'    : ticklength,\n",
    "    'width'     : axwidth/2,\n",
    "    'pad'       : tickpad,\n",
    "    'axis'      : 'both',\n",
    "    'direction' : 'out',\n",
    "    'colors'    : bkcolor,\n",
    "    'bottom'    : True,\n",
    "    'left'      : True,\n",
    "    'top'       : False,\n",
    "    'right'     : False\n",
    "    }\n",
    "\n",
    "def_minortickprops = {\n",
    "    'length'    : ticklength-1.25,\n",
    "    'width'     : axwidth/2,\n",
    "    'axis'      : 'both',\n",
    "    'direction' : 'out',\n",
    "    'which'     : 'minor',\n",
    "    'color'     : bkcolor\n",
    "    }\n",
    "\n",
    "def_axprops = {\n",
    "    'linewidth' : axwidth,\n",
    "    'linestyle' : '-',\n",
    "    'color'     : bkcolor\n",
    "    }\n",
    "\n",
    "params = {'text.usetex': False, 'mathtext.fontset': 'stixsans', 'mathtext.default' : 'regular'}\n",
    "plot.rcParams.update(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Important functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getdE(o=3.0, mu=1.9, sigma=0.5):\n",
    "    \"\"\" Return a random (change in) binding energy following a shifted log-normal distribution. \"\"\"\n",
    "    return o - np.exp(np.random.normal(mu, sigma))\n",
    "\n",
    "def getE(Ec, Ev, Q, E0=4., lam=0.9):\n",
    "    \"\"\" Return binding energy given raw binding energies with conserved and variable regions, and flexibility. \"\"\"\n",
    "    return (Q * ((lam * Ec) + ((1. - lam) * Ev))) + ((1. - Q) * E0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Analysis and figures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure 4-figure supplement 1. Population trajectories of successful and unsusccessful GCs, single antigen test case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data here -- single antigen case\n",
    "\n",
    "df    = pd.read_csv(        'data/1ag-output-total.csv.gz', memory_map=True)\n",
    "dfLC  = pd.read_csv('data/1ag-output-largest-clone.csv.gz', memory_map=True)\n",
    "dfMem = pd.read_csv(       'data/1ag-output-memory.csv.gz', memory_map=True)\n",
    "\n",
    "nTrials = 10000\n",
    "finalE  = []\n",
    "nMut    = []\n",
    "\n",
    "for t in range(nTrials):\n",
    "    dfT = dfMem[dfMem['trial']==t]\n",
    "    if len(dfT)>0:\n",
    "        idx = np.argmax(np.array(dfT['number']))\n",
    "        nMut.append(dfT.iloc[idx].mutations)\n",
    "        tempEc = dfT.iloc[idx].Ec\n",
    "        tempEv = dfT.iloc[idx].Ev0\n",
    "        tempQ  = dfT.iloc[idx].Q\n",
    "        finalE.append(getE(tempEc, tempEv, tempQ, E0=4., lam=0.9))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting GC trajectories...\n",
      "Plotting supplementary info...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Plot the figure\n",
    "\n",
    "nTrials = 10000\n",
    "plotCut = 1000  # number of actual lines to plot\n",
    "nSurv   = 0\n",
    "hSurv   = []\n",
    "hDie    = []\n",
    "\n",
    "w     = one_five_column\n",
    "goldh = (1.0*w)/goldr\n",
    "fig   = plot.figure(figsize=(w, goldh))\n",
    "gs1   = gridspec.GridSpec(3, 1, width_ratios=[1.00], height_ratios=[1.0,1.0,1.0])\n",
    "gs1.update(left=0.02, right=0.58, hspace=0.2)\n",
    "ax1   = plot.subplot(gs1[:2, 0]) # trajectories\n",
    "ax2   = plot.subplot(gs1[ 2, 0]) # end cycle histogram\n",
    "gs2   = gridspec.GridSpec(2, 1, width_ratios=[1.00], height_ratios=[1.0,1.0])\n",
    "gs2.update(left=0.72, right=0.98, bottom=0.3, hspace=0.8)\n",
    "ax3   = plot.subplot(gs2[ 0, 0]) # kde - mutations\n",
    "ax4   = plot.subplot(gs2[ 1, 0]) # kde - final energy\n",
    "\n",
    "\n",
    "# PART 1: Trajectories\n",
    "\n",
    "print('Plotting GC trajectories...')\n",
    "\n",
    "legendprops = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "lineprops   = { 'lw' : sizeline, 'ls' : '-', 'alpha' : 0.2 }\n",
    "cSurv       = '#FF7F00'\n",
    "cDie        = bkcolor\n",
    "\n",
    "pprops = { 'xlim'      : [ 0, 250],\n",
    "           'xticks'    : [],\n",
    "           'ylim'      : [ 0, 1536],\n",
    "           'yticks'    : [ 0,  250, 500, 750, 1000, 1250, 1500],\n",
    "           'ylabel'    : 'Number of B cells',\n",
    "           'plotprops' : lineprops,\n",
    "           'hide'      : ['bottom'],\n",
    "           'sublabel'  : 'A',\n",
    "           'sublabelcoords' : [-0.18, 1.00] }\n",
    "\n",
    "for t in range(nTrials):\n",
    "    dfT  = df[df['trial']==t]\n",
    "    x, y = np.array(dfT['cycle']), np.array(dfT['number recycled'])\n",
    "    if y[-1]>0: \n",
    "        nSurv += 1\n",
    "        hSurv.append(x[-1])\n",
    "        if t<plotCut: mp.line(ax=ax1, x=[x], y=[y], colors=[cSurv], **pprops)\n",
    "    else:\n",
    "        hDie.append(x[-1])\n",
    "        if t<plotCut: mp.line(ax=ax1, x=[x], y=[y], colors=[cDie],  **pprops)\n",
    "\n",
    "\n",
    "offPlot    = [[-10,-10],[-10,-10]]\n",
    "survString = '%d%%' % np.round(float(nSurv)/float(nTrials) * 100) + ' succeed'\n",
    "dieString  = '%d%%' % np.round(100. - (float(nSurv)/float(nTrials) * 100)) + ' die'\n",
    "ax1.text(185, 250, survString, ha='left', va='center', color=cSurv, family=fontfamily, size=sizelabel)\n",
    "ax1.text(185, 110,  dieString, ha='left', va='center', color= cDie, family=fontfamily, size=sizelabel)\n",
    "mp.plot(type='line', ax=ax1, x=offPlot, y=offPlot, colors=[cSurv,cDie], **pprops)\n",
    "\n",
    "\n",
    "# PART 2: Histograms\n",
    "\n",
    "print('Plotting supplementary info...')\n",
    "\n",
    "histprops                = { 'histtype' : 'bar', 'lw' : sizeline/2, 'rwidth' : 0.8, 'ls' : 'solid', 'alpha' : 0.7,\n",
    "                             'align' : 'left' }\n",
    "pprops['hide']           = []\n",
    "pprops['xlim']           = [ 0, 250]\n",
    "pprops['xticks']         = [ 0,  50, 100, 150, 200, 250]\n",
    "pprops['xaxy']           = 1\n",
    "pprops['xlabel']         = 'Cycles'\n",
    "pprops['ylim']           = [0, 0.15]\n",
    "pprops['yticks']         = [0, 0.05, 0.10, 0.15]\n",
    "pprops['yticklabels']    = ['0', '5', '10', '15']\n",
    "pprops['ylabel']         = 'Frequency (%)'\n",
    "pprops['plotprops']      = histprops\n",
    "pprops['sublabel']       = 'B'\n",
    "pprops['sublabelcoords'] = [-0.18, 1.00]\n",
    "\n",
    "mp.plot(type='hist', ax=ax2, x=[hSurv, hDie], colors=[cSurv, cDie], **pprops)\n",
    "\n",
    "\n",
    "# PART 3: KDEs\n",
    "cKDE = '#3182bd'\n",
    "\n",
    "pprops['plotprops']      = histprops\n",
    "pprops['bins']           = np.arange(0,15+1,1)\n",
    "pprops['xlim']           = [ 0, 15]\n",
    "pprops['xticks']         = [ 0, 5, 10, 15]\n",
    "pprops['xlabel']         = 'Number of mutations'\n",
    "pprops['ylim']           = [0, 0.30]\n",
    "pprops['yticks']         = [0, 0.10, 0.20, 0.30]\n",
    "pprops['yticklabels']    = ['0', '10', '20', '30']\n",
    "pprops['ylabel']         = 'Frequency (%)'\n",
    "pprops['sublabel']       = 'C'\n",
    "pprops['sublabelcoords'] = [-0.40, 1.00]\n",
    "\n",
    "mp.plot(type='hist', ax=ax3, x=[nMut], colors=[cKDE], **pprops)\n",
    "\n",
    "del pprops['yticklabels']\n",
    "pprops['plotprops']      = def_lineprops\n",
    "pprops['xlim']           = [ 0, 10]\n",
    "pprops['xticks']         = [ 0, 2.5, 5, 7.5, 10]\n",
    "pprops['xticklabels']    = [r'$0$', r'$-2.5$', r'$-5$', r'$-7.5$', r'$-10$']\n",
    "pprops['xlabel']         = 'Binding free energy'\n",
    "pprops['ylim']           = [0, 0.3]\n",
    "pprops['yticks']         = [0, 0.1, 0.2, 0.3]\n",
    "pprops['ylabel']         = 'Density'\n",
    "pprops['sublabel']       = 'D'\n",
    "pprops['sublabelcoords'] = [-0.40, 1.00]\n",
    "\n",
    "mp.plot(type='kde', ax=ax4, x=[finalE], colors=[cKDE], **pprops)\n",
    "\n",
    "\n",
    "# Save data\n",
    "plot.savefig('figures/figure-4-figure-supplement-1.pdf', **figprops)\n",
    "plot.close(fig)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure 6. Accumulation of mutations in the FR and CDR regions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data here\n",
    "\n",
    "dfHi = pd.read_csv('data/hiEc-output-largest-clone.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-output-largest-clone.csv.gz', memory_map=True)\n",
    "        \n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "genCut = [25, 150, 400]\n",
    "\n",
    "xHi, yHi = [], []\n",
    "xLo, yLo = [], []\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        xHi.append(np.array(nextdfHi['CDR_mutations']))\n",
    "        yHi.append(np.array(nextdfHi['FR_mutations']))\n",
    "        xLo.append(np.array(nextdfLo['CDR_mutations']))\n",
    "        yLo.append(np.array(nextdfLo['FR_mutations']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting number of mutations...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Plot the figure\n",
    "\n",
    "w     = double_column\n",
    "goldh = (2.0/3.)*w\n",
    "fig   = plot.figure(figsize=(w, goldh))\n",
    "\n",
    "HiBox = [0.20, 0.98, 0.60, 0.90]\n",
    "LoBox = [0.20, 0.98, 0.16, 0.46]\n",
    "\n",
    "gsS   = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsS.update(left=HiBox[0], right=HiBox[1], bottom=0.93, top=1.0)\n",
    "axS   = plot.subplot(gsS[ 0, 0])\n",
    "\n",
    "gsHi  = gridspec.GridSpec(1, 3, width_ratios=[1.0, 1.0, 1.0])\n",
    "gsHi.update(left=HiBox[0], right=HiBox[1], bottom=HiBox[2], top=HiBox[3], hspace=0.5, wspace=0.5)\n",
    "gsLo  = gridspec.GridSpec(1, 3, width_ratios=[1.0, 1.0, 1.0])\n",
    "gsLo.update(left=LoBox[0], right=LoBox[1], bottom=LoBox[2], top=LoBox[3], hspace=0.5, wspace=0.5)\n",
    "\n",
    "axHi  = [plot.subplot(gsHi[i]) for i in range(len(genCut))]\n",
    "axLo  = [plot.subplot(gsLo[i]) for i in range(len(genCut))]\n",
    "\n",
    "# Overhead line\n",
    "\n",
    "lineprops = { 'lw' : axwidth/2., 'ls' : '-', 'alpha' : 1.0 }\n",
    "pprops    = { 'xlim' : [0, 1], 'xticks' : [], 'ylim' : [0, 1], 'yticks' : [], \n",
    "              'hide' : ['top','bottom','left','right'], 'plotprops' : lineprops }\n",
    "tprops    = dict(color=bkcolor, family=fontfamily, size=sizelabel)\n",
    "axS.text(3./6.-0.005, 1.0, 'Time (generations)', ha='center', va='top',    **tprops)\n",
    "axS.text(1./6.-0.045, 0.3,              r'$25$', ha='center', va='center', **tprops)\n",
    "axS.text(3./6.-0.005, 0.3,             r'$150$', ha='center', va='center', **tprops)\n",
    "axS.text(5./6.+0.055, 0.3,             r'$400$', ha='center', va='center', **tprops)\n",
    "ay = 0.05\n",
    "axS.arrow(x=0, y=ay, dx=1, dy=0, head_width=0.25, head_length=0.025, overhang=0.5, fc=bkcolor,\n",
    "             ec=None, lw=sizeline/2., clip_on=False, zorder=100)\n",
    "mp.plot(type='line', ax=axS, x=[[0,1]], y=[[ay,ay]], colors=[bkcolor], **pprops)\n",
    "\n",
    "# Number of CDR/FWR mutations\n",
    "\n",
    "print('Plotting number of mutations...')\n",
    "\n",
    "levels = [0.02] + list(0.02 * np.logspace(1., 5., num=5, base=1.5))\n",
    "\n",
    "legendprops  = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                 'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "contourprops = { 'cmap' : 'Greens', 'colors' : bkcolor, 'extend' : 'max', 'linewidths' : sizeline/2.,\n",
    "                 'vmin' : 0.001, 'vmax' : 0.15 , 'levels' : levels }\n",
    "\n",
    "dx     = 0.1\n",
    "dy     = dx\n",
    "cx     = np.arange(0., 12.+dx, dx)\n",
    "cy     = np.arange(0., 12.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.9\n",
    "sy = sx\n",
    "\n",
    "pprops = { 'xlim'      : [ 0., 12.],\n",
    "           'xticks'    : [ 0., 4., 8., 12.],\n",
    "           'ylim'      : [ 0., 12.],\n",
    "           'yticks'    : [ 0., 4., 8., 12.],\n",
    "           'xlabel'    : 'CDR mutations',\n",
    "           'ylabel'    : 'Framework mutations',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "for i in range(len(xHi)):\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xHi[i])):\n",
    "        zz += np.exp(-((xHi[i][k]-xx)**2/sx**2) - ((yHi[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xHi[i]))\n",
    "    mp.plot(type='contour', ax=axHi[i], x=[xx], y=[yy], z=[zz], **pprops)\n",
    "\n",
    "for i in range(len(xLo)):\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xLo[i])):\n",
    "        zz += np.exp(-((xLo[i][k]-xx)**2/sx**2) - ((yLo[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xLo[i]))\n",
    "    mp.plot(type='contour', ax=axLo[i], x=[xx], y=[yy], z=[zz], **pprops)\n",
    "\n",
    "# Background\n",
    "\n",
    "axS.text(HiBox[0]-0.06, HiBox[3]-0.01, 'A', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS.text(LoBox[0]-0.06, LoBox[3]-0.01, 'B', transform=fig.transFigure, **def_sublabelprops)\n",
    "\n",
    "cBG   = '#F5F5F5'\n",
    "dx    = 0.08\n",
    "dy    = 0.09\n",
    "bgHi  = axS.axis()\n",
    "recHi = matplotlib.patches.Rectangle(xy=(HiBox[0]-dx,HiBox[2]-dy), width=0.8+dx, height=HiBox[3]-HiBox[2]+(1.2*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recHi = axS.add_patch(recHi)\n",
    "\n",
    "bgLo  = axS.axis()\n",
    "recLo = matplotlib.patches.Rectangle(xy=(LoBox[0]-dx,LoBox[2]-dy), width=0.8+dx, height=LoBox[3]-LoBox[2]+(1.2*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recLo = axS.add_patch(recLo)\n",
    "    \n",
    "# Save data\n",
    "plot.savefig('figures/figure-6.pdf', dpi = 1000, \n",
    "              facecolor = fig.get_facecolor(), edgecolor=None, **figprops)\n",
    "plot.close(fig)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analysis for Figure 4 and Figure 4-figure supplement 2. Breadth comparisons at different binding energy cutoffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create alternate \"test panel\" of 100 sequences with different Ev\n",
    "\n",
    "nTrials   = 10000\n",
    "nPanel    = 100\n",
    "testPanel = np.array([getdE(o=0, mu=3.0, sigma=1.0) for i in range(nPanel)])\n",
    "\n",
    "# Read through largest clone histories and output binding energies with the test panel\n",
    "\n",
    "fIn  = ['data/hiEC-output-largest-clone.csv.gz',\n",
    "        'data/loEC-output-largest-clone.csv.gz',\n",
    "        'data/hiEC-l0.8-output-largest-clone.csv.gz',\n",
    "        'data/loEC-l0.8-output-largest-clone.csv.gz']\n",
    "fOut = ['data/hiEc-panel-breadth.csv.gz',\n",
    "        'data/loEc-panel-breadth.csv.gz',\n",
    "        'data/hiEc-l0.8-panel-breadth.csv.gz', \n",
    "        'data/loEc-l0.8-panel-breadth.csv.gz']\n",
    "\n",
    "for k in range(len(fIn)):\n",
    "    if os.path.isfile(fOut[k]):\n",
    "        continue\n",
    "    else:\n",
    "        f  = open(fOut[k], 'w')\n",
    "        f.write('trial,generation,Q,Ec%s\\n' % (''.join([',Ev%d' % (i+1) for i in range(nPanel)])))\n",
    "\n",
    "        df = pd.read_csv(fIn[k], memory_map=True)\n",
    "\n",
    "        for t in range(nTrials):\n",
    "            trialDf = df[df.trial==t]\n",
    "\n",
    "            if len(trialDf)==0: continue\n",
    "\n",
    "            # Set starting binding energies with test panel\n",
    "\n",
    "            panelEc = trialDf.iloc[0].Ec\n",
    "            panelEv = deepcopy(testPanel)\n",
    "\n",
    "            f.write('%d,%d,%lf,%lf' % (trialDf.iloc[0].trial,trialDf.iloc[0].generation,trialDf.iloc[0].Q,panelEc))\n",
    "            for i in range(len(testPanel)): f.write(',%lf' % panelEv[i])\n",
    "            f.write('\\n')\n",
    "\n",
    "            # Iterate over successive mutations and update binding energies\n",
    "\n",
    "            for i in range(1,len(trialDf)):\n",
    "\n",
    "                if trialDf.iloc[i].Q==trialDf.iloc[i-1].Q:\n",
    "                    dEc = trialDf.iloc[i].Ec-trialDf.iloc[i-1].Ec\n",
    "                    if dEc==0:\n",
    "                        for j in range(len(testPanel)): panelEv[j] += getdE()\n",
    "                    else: panelEc += dEc\n",
    "\n",
    "                f.write('%d,%d,%lf,%lf' % (trialDf.iloc[i].trial,trialDf.iloc[i].generation,trialDf.iloc[i].Q,panelEc))\n",
    "                for j in range(len(testPanel)): f.write(',%lf' % panelEv[j])\n",
    "                f.write('\\n')\n",
    "                f.flush()\n",
    "\n",
    "        f.close() # files can be gzipped below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compress panel files\n",
    "\n",
    "fin  = ['data/hiEc-panel-breadth.csv',\n",
    "        'data/loEc-panel-breadth.csv',\n",
    "        'data/hiEc-l0.8-panel-breadth.csv', \n",
    "        'data/loEc-l0.8-panel-breadth.csv']\n",
    "fout = ['data/hiEc-panel-breadth',\n",
    "        'data/loEc-panel-breadth',\n",
    "        'data/hiEc-l0.8-panel-breadth', \n",
    "        'data/loEc-l0.8-panel-breadth']\n",
    "\n",
    "for k in range(len(fin)):\n",
    "    df = pd.read_csv(fin[k], memory_map=True)\n",
    "    df.to_csv(fout[k]+'.csv.gz', compression='gzip')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure 4. Trajectories of binding energy/flexibility during affinity maturation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data: Q vs Ec plot\n",
    "\n",
    "dfHi = pd.read_csv('data/hiEc-output-largest-clone.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-output-largest-clone.csv.gz', memory_map=True)\n",
    "\n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "genCut = [25, 150, 400]\n",
    "\n",
    "xHi, yHi = [], []\n",
    "xLo, yLo = [], []\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        xHi.append(np.array(nextdfHi['Ec']))\n",
    "        yHi.append(np.array(nextdfHi['Q']))\n",
    "        xLo.append(np.array(nextdfLo['Ec']))\n",
    "        yLo.append(np.array(nextdfLo['Q']))\n",
    "        \n",
    "# Load data: breadth plot\n",
    "\n",
    "ECut = [3.0, 6.0, 9.0]\n",
    "\n",
    "dfHi = pd.read_csv('data/hiEc-panel-breadth.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-panel-breadth.csv.gz', memory_map=True)\n",
    "\n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "nPanel     = 100\n",
    "bHi        = []\n",
    "bLo        = []\n",
    "bHiTrueMed = []\n",
    "bLoTrueMed = []\n",
    "EvLabel    = ['Ev%d' % (i+1) for i in range(nPanel)]\n",
    "denom      = float(nPanel)\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        tempQ  = np.array([ q * np.ones(nPanel) for  q in nextdfHi.Q])\n",
    "        tempEc = np.array([ec * np.ones(nPanel) for ec in nextdfHi.Ec])\n",
    "        EHi    = np.array(getE(tempEc, np.array(nextdfHi[EvLabel]), tempQ, E0=4., lam=0.9))\n",
    "        \n",
    "        tempQ  = np.array([ q * np.ones(nPanel) for  q in nextdfLo.Q])\n",
    "        tempEc = np.array([ec * np.ones(nPanel) for ec in nextdfLo.Ec])\n",
    "        ELo    = np.array(getE(tempEc, np.array(nextdfLo[EvLabel]), tempQ, E0=4., lam=0.9))\n",
    "        \n",
    "        bHi.append([[np.sum(E>c)/denom for E in EHi] for c in ECut])\n",
    "        bLo.append([[np.sum(E>c)/denom for E in ELo] for c in ECut])\n",
    "        \n",
    "        bHiTrueMed.append([np.median(E) for E in EHi])\n",
    "        bLoTrueMed.append([np.median(E) for E in ELo])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting densities...\n",
      "Plotting breadth...\n",
      "Plotting Q vs breadth...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Plot the figure\n",
    "\n",
    "nTrials = 10000\n",
    "reload(mp)\n",
    "\n",
    "w     = double_column\n",
    "goldh = (3.3/3.0)*w\n",
    "fig   = plot.figure(figsize=(w, goldh))\n",
    "\n",
    "HiBox = [0.20, 0.98, 0.61, 0.91]\n",
    "LoBox = [0.20, 0.98, 0.23, 0.53]\n",
    "LLBox = [0.10, 0.30, 0.00, 0.16]\n",
    "\n",
    "gsS   = gridspec.GridSpec(1, 2, width_ratios=[1.0, 2.5], height_ratios=[1.0])\n",
    "gsS.update(left=0, right=HiBox[1], bottom=0.93, top=1.0)\n",
    "gsHiS = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsHiS.update(left=0.02, right=0.12, bottom=HiBox[2]+((HiBox[3]-HiBox[2])*0.5/1.3), top=HiBox[3])\n",
    "gsLoS = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsLoS.update(left=0.02, right=0.12, bottom=LoBox[2]+((LoBox[3]-LoBox[2])*0.5/1.3), top=LoBox[3])\n",
    "axS   = [plot.subplot(gsS[ 0, 1]), plot.subplot(gsHiS[ 0, 0]), plot.subplot(gsLoS[ 0, 0])]\n",
    "axKey = plot.subplot(gsS[0, 0])\n",
    "\n",
    "gsHi  = gridspec.GridSpec(2, 3, height_ratios=[1.0, 0.55], width_ratios=[1.0, 1.0, 1.0])\n",
    "gsHi.update(left=HiBox[0], right=HiBox[1], bottom=HiBox[2], top=HiBox[3], hspace=0.2, wspace=0.2)\n",
    "gsLo  = gridspec.GridSpec(2, 3, height_ratios=[1.0, 0.55], width_ratios=[1.0, 1.0, 1.0])\n",
    "gsLo.update(left=LoBox[0], right=LoBox[1], bottom=LoBox[2], top=LoBox[3], hspace=0.2, wspace=0.2)\n",
    "axHi  = [[plot.subplot(gsHi[0, i]), plot.subplot(gsHi[1, i])] for i in range(len(genCut))]\n",
    "axLo  = [[plot.subplot(gsLo[0, i]), plot.subplot(gsLo[1, i])] for i in range(len(genCut))]\n",
    "\n",
    "gsLL = gridspec.GridSpec(1, 1, height_ratios=[1.0], width_ratios=[1.0])\n",
    "gsLL.update(left=LLBox[0], right=LLBox[1], bottom=LLBox[2], top=LLBox[3], hspace=0.4, wspace=0.1)\n",
    "axLL = [plot.subplot(gsLL[0, 0])]\n",
    "\n",
    "# Overhead line\n",
    "\n",
    "lineprops = { 'lw' : axwidth/2., 'ls' : '-', 'alpha' : 1.0 }\n",
    "pprops    = { 'xlim' : [0, 1], 'xticks' : [], 'ylim' : [0, 1], 'yticks' : [], \n",
    "              'hide' : ['top','bottom','left','right'], 'plotprops' : lineprops }\n",
    "tprops    = dict(color=bkcolor, family=fontfamily, size=sizelabel)\n",
    "axS[0].text(2.4/6.-0.005, 0.6, 'Time (generations)', ha='center', va='top',    **tprops)\n",
    "axS[0].text(           0, 0.2,              r'$25$', ha='center', va='center', **tprops)\n",
    "axS[0].text(2.4/6.-0.005, 0.2,             r'$150$', ha='center', va='center', **tprops)\n",
    "axS[0].text(4.9/6.+0.015, 0.2,             r'$400$', ha='center', va='center', **tprops)\n",
    "ay = 0.05\n",
    "axS[0].arrow(x=-0.30, y=ay, dx=1+0.30, dy=0, head_width=0.25, head_length=0.025, overhang=0.5, fc=bkcolor,\n",
    "             ec=None, lw=sizeline/2., clip_on=False, zorder=100)\n",
    "mp.plot(type='line', ax=axS[0], x=[[0,1]], y=[[ay,ay]], colors=[bkcolor], **pprops)\n",
    "\n",
    "# PART 1: densities\n",
    "\n",
    "print('Plotting densities...')\n",
    "\n",
    "legendprops  = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                 'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "contourprops = { 'cmap' : 'Oranges', 'colors' : bkcolor, 'vmin' : 0.001, 'vmax' : 0.10, 'extend' : 'max',\n",
    "                 'linewidths' : sizeline/2., 'levels' : [0.001, 0.005, 0.015, 0.035, 0.075, 0.15]}\n",
    "\n",
    "dx     = 0.1\n",
    "dy     = dx * (1. / 16.)\n",
    "cx     = np.arange(-2., 14.+dx, dx)\n",
    "cy     = np.arange( 0.,  1.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.5\n",
    "sy = sx * (1. / 16.)\n",
    "    \n",
    "pprops = { 'xlim'      : [-2., 14.],\n",
    "           'xticks'    : [0.,  4.,  8., 12.],\n",
    "           'xticklabels' : [r'$0$', r'$-4$', r'$-8$', r'$-12$'],\n",
    "           'xminorticks' : [-2., 2., 6., 10., 14.],\n",
    "           'ylim'      : [ 0., 1.],\n",
    "           'yticks'    : [ 0., 0.25, 0.50, 0.75, 1.00],\n",
    "           'xlabel'    : r'$E_c$',\n",
    "           'ylabel'    : r'$Q$',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "r      = 1.\n",
    "cprops = dict(width=r, height=r/14., ec='k', fc='w', alpha=1.0, zorder=100)\n",
    "\n",
    "for i in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if i>0:\n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks']   = []\n",
    "        tempprops['hide']     = ['left']\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xHi[i])):\n",
    "        zz += np.exp(-((xHi[i][k]-xx)**2/sx**2) - ((yHi[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xHi[i]))\n",
    "    mp.plot(type='contour', ax=axHi[i][0], x=[xx], y=[yy], z=[zz], **tempprops)\n",
    "    circHi = matplotlib.patches.Ellipse(xy=(4.,0.8), **cprops)\n",
    "    axHi[i][0].add_patch(circHi)\n",
    "    if i==0:\n",
    "        axHi[i][0].annotate('Initial\\nparameters', xy=(3.7, 0.78), xytext=(1.75, 0.50), ha='center', va='center',\n",
    "                            arrowprops=dict(arrowstyle='simple', fc='k', mutation_scale=10., shrinkA=3., shrinkB=3.), \n",
    "                            **def_labelprops)\n",
    "    if i==len(genCut)-1:\n",
    "        axHi[i][0].text(2., 0.90, 'CH103-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "\n",
    "for i in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if i>0:\n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks']   = []\n",
    "        tempprops['hide']     = ['left']\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xLo[i])):\n",
    "        zz += np.exp(-((xLo[i][k]-xx)**2/sx**2) - ((yLo[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xLo[i]))\n",
    "    mp.plot(type='contour', ax=axLo[i][0], x=[xx], y=[yy], z=[zz], **tempprops)\n",
    "    circLo = matplotlib.patches.Ellipse(xy=(0.,0.8), **cprops)\n",
    "    axLo[i][0].add_patch(circLo)\n",
    "    if i==len(genCut)-1:\n",
    "        axLo[i][0].text(-0.5, 0.90, 'PGT121-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "        axLo[i][0].text(-2., 0.35, '3BNC60-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "\n",
    "# PART 2: breadth\n",
    "\n",
    "print('Plotting breadth...')\n",
    "\n",
    "legendprops = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "lineprops   = { 'lw' : sizeline, 'ls' : '-', 'alpha' : 1.0 }\n",
    "cHist       = '#3182bd'\n",
    "\n",
    "histprops = { 'histtype' : 'bar', 'lw' : sizeline/2, 'rwidth' : 0.8, 'ls' : 'solid', 'alpha' : 0.7, 'ec' : bkcolor }\n",
    "\n",
    "pprops = { 'xlim'        : [-2., 10.],\n",
    "           'xticks'      : [0.,  4., 8.],\n",
    "           'xticklabels' : [r'$0$', r'$-4$', r'$-8$'],\n",
    "           'xminorticks' : [-2,  2., 6., 10.],\n",
    "           'ylim'        : [ 0., 0.50],\n",
    "           'yticks'      : [ 0., 0.25, 0.50],\n",
    "           'yticklabels' : [  0,   25,   50],\n",
    "           'xlabel'      : 'Median binding free energy',\n",
    "           'ylabel'      : 'Frequency (%)',\n",
    "           'plotprops'   : histprops,\n",
    "           'bins'        : np.arange(-2.0, 10.5, 0.5),\n",
    "           'combine'     : True }\n",
    "\n",
    "for g in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if g>0: \n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks'] = []\n",
    "        tempprops['hide']   = ['left']\n",
    "    mp.plot(type='hist', ax=axHi[g][1], x=[bHiTrueMed[g]], colors=[cHist], **tempprops)\n",
    "    mp.plot(type='hist', ax=axLo[g][1], x=[bLoTrueMed[g]], colors=[cHist], **tempprops)\n",
    "    \n",
    "# PART 3: Q vs breadth\n",
    "\n",
    "print('Plotting Q vs breadth...')\n",
    "\n",
    "dx     = 0.05\n",
    "dy     = dx * (1. / 7.)\n",
    "cx     = np.arange(1., 8.+dx, dx)\n",
    "cy     = np.arange(0., 1.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.4\n",
    "sy = sx * (1. / 7.)\n",
    "\n",
    "zz = np.zeros(xx.shape)\n",
    "for k in range(len(bLoTrueMed[-1])):\n",
    "    zz += np.exp(-((bLoTrueMed[-1][k]-xx)**2/sx**2) - ((yLo[-1][k]-yy)**2/sy**2))\n",
    "zz /= float(len(bLoTrueMed[-1]))\n",
    "\n",
    "contourprops['cmap'] = 'Blues'\n",
    "scatterprops = { 'lw' : 0, 's' : 8., 'marker' : 'o' }\n",
    "\n",
    "pprops = { 'xlim'      : [ 1., 8.],\n",
    "           'xticks'    : [ 2., 4., 6., 8.],\n",
    "           'xticklabels' : [r'$-2$', r'$-4$', r'$-6$', r'$-8$'],\n",
    "           'xaxstart'  : 1,\n",
    "           'ylim'      : [ 0., 1.],\n",
    "           'yticks'    : [ 0., 0.25, 0.50, 0.75, 1.00],\n",
    "           'xlabel'    : 'Median binding free energy',\n",
    "           'ylabel'    : r'$Q$',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "mp.plot(type='contour', ax=axLL[0], x=[xx], y=[yy], z=[zz], **pprops)\n",
    "axLL[0].text(6.5, 0.80, 'PGT121-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "axLL[0].text(5.5, 0.25, '3BNC60-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "    \n",
    "# Schematic\n",
    "\n",
    "hiSchem  = mpimg.imread('figures/hiEc.png')\n",
    "loSchem  = mpimg.imread('figures/loEc.png')\n",
    "keySchem = mpimg.imread('figures/legend2.png')\n",
    "axS[1].imshow(hiSchem)\n",
    "axS[2].imshow(loSchem)\n",
    "axKey.imshow(keySchem, zorder=-200)\n",
    "for ax in [axS[1], axS[2], axKey]:\n",
    "    for a in ['left', 'right', 'bottom', 'top']: \n",
    "        ax.spines[a].set_visible(False)\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "axS[2].text( 0.02, LoBox[3]-0.01, 'B', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS[1].text( 0.02, HiBox[3]-0.01, 'A', transform=fig.transFigure, **def_sublabelprops)\n",
    "axLL[0].text(0.02, LoBox[2]-0.08, 'C', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS[1].text(0.07, (HiBox[2]+HiBox[3])/2.-0.03, 'Strong initial ' + r'$E_c$', ha='center', va='center',\n",
    "            transform=fig.transFigure, **def_labelprops)\n",
    "axS[2].text(0.07, (LoBox[2]+LoBox[3])/2.-0.03, 'Weak initial ' + r'$E_c$', ha='center', va='center', \n",
    "            transform=fig.transFigure, **def_labelprops)\n",
    "tprops['size'] = sizelabel\n",
    "axKey.text(0.080, 0.965,        'viral\\nspike', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "axKey.text(0.175, 0.965,   'conserved\\nregion', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "axKey.text(0.290, 0.965, 'antibody\\nfootprint', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "\n",
    "# Background\n",
    "\n",
    "cBG   = '#F5F5F5'\n",
    "dy    = 0.07\n",
    "bgHi  = axS[1].axis()\n",
    "recHi = matplotlib.patches.Rectangle(xy=(0,HiBox[2]-0.8*dy), width=1., height=HiBox[3]-HiBox[2]+(1.*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recHi = axS[1].add_patch(recHi)\n",
    "\n",
    "bgLo  = axS[2].axis()\n",
    "recLo = matplotlib.patches.Rectangle(xy=(0,LoBox[2]-0.8*dy), width=1., height=LoBox[3]-LoBox[2]+(1.*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recLo = axS[2].add_patch(recLo)\n",
    "    \n",
    "# Save data\n",
    "plot.savefig('figures/figure-4.pdf', dpi = 1000, \n",
    "              facecolor = fig.get_facecolor(), edgecolor=None, **figprops)\n",
    "plot.close(fig)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure 4-figure supplement 2. Trajectories of binding energy/flexibility during affinity maturation, smaller overlap ($\\lambda = 0.8$)\n",
    "\n",
    "Note: This figure is the same as Figure 4, just using different input files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data: Q vs Ec plot\n",
    "\n",
    "dfHi = pd.read_csv('data/hiEc-l0.8-output-largest-clone.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-l0.8-output-largest-clone.csv.gz', memory_map=True)\n",
    "\n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "genCut = [25, 150, 400]\n",
    "\n",
    "xHi, yHi = [], []\n",
    "xLo, yLo = [], []\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        xHi.append(np.array(nextdfHi['Ec']))\n",
    "        yHi.append(np.array(nextdfHi['Q']))\n",
    "        xLo.append(np.array(nextdfLo['Ec']))\n",
    "        yLo.append(np.array(nextdfLo['Q']))\n",
    "        \n",
    "# Load data: breadth plot\n",
    "\n",
    "ECut = [3.0, 6.0, 9.0]\n",
    "\n",
    "dfHi = pd.read_csv('data/hiEc-l0.8-panel-breadth.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-l0.8-panel-breadth.csv.gz', memory_map=True)\n",
    "\n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "nPanel     = 100\n",
    "bHi        = []\n",
    "bLo        = []\n",
    "bHiTrueMed = []\n",
    "bLoTrueMed = []\n",
    "EvLabel    = ['Ev%d' % (i+1) for i in range(nPanel)]\n",
    "denom      = float(nPanel)\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        tempQ  = np.array([ q * np.ones(nPanel) for  q in nextdfHi.Q])\n",
    "        tempEc = np.array([ec * np.ones(nPanel) for ec in nextdfHi.Ec])\n",
    "        EHi    = np.array(getE(tempEc, np.array(nextdfHi[EvLabel]), tempQ, E0=4., lam=0.8))\n",
    "         \n",
    "        tempQ  = np.array([ q * np.ones(nPanel) for  q in nextdfLo.Q])\n",
    "        tempEc = np.array([ec * np.ones(nPanel) for ec in nextdfLo.Ec])\n",
    "        ELo    = np.array(getE(tempEc, np.array(nextdfLo[EvLabel]), tempQ, E0=4., lam=0.8))\n",
    "        \n",
    "        bHi.append([[np.sum(E>c)/denom for E in EHi] for c in ECut])\n",
    "        bLo.append([[np.sum(E>c)/denom for E in ELo] for c in ECut])\n",
    "        \n",
    "        bHiTrueMed.append([np.median(E) for E in EHi])\n",
    "        bLoTrueMed.append([np.median(E) for E in ELo])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting densities...\n",
      "Plotting breadth...\n",
      "Plotting Q vs breadth...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Plot the figure\n",
    "\n",
    "nTrials = 10000\n",
    "reload(mp)\n",
    "\n",
    "w     = double_column\n",
    "goldh = (3.3/3.0)*w\n",
    "fig   = plot.figure(figsize=(w, goldh))\n",
    "\n",
    "HiBox = [0.20, 0.98, 0.61, 0.91]\n",
    "LoBox = [0.20, 0.98, 0.23, 0.53]\n",
    "LLBox = [0.10, 0.30, 0.00, 0.16]\n",
    "\n",
    "gsS   = gridspec.GridSpec(1, 2, width_ratios=[1.0, 2.5], height_ratios=[1.0])\n",
    "gsS.update(left=0, right=HiBox[1], bottom=0.93, top=1.0)\n",
    "gsHiS = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsHiS.update(left=0.02, right=0.12, bottom=HiBox[2]+((HiBox[3]-HiBox[2])*0.5/1.3), top=HiBox[3])\n",
    "gsLoS = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsLoS.update(left=0.02, right=0.12, bottom=LoBox[2]+((LoBox[3]-LoBox[2])*0.5/1.3), top=LoBox[3])\n",
    "axS   = [plot.subplot(gsS[ 0, 1]), plot.subplot(gsHiS[ 0, 0]), plot.subplot(gsLoS[ 0, 0])]\n",
    "axKey = plot.subplot(gsS[0, 0])\n",
    "\n",
    "gsHi  = gridspec.GridSpec(2, 3, height_ratios=[1.0, 0.55], width_ratios=[1.0, 1.0, 1.0])\n",
    "gsHi.update(left=HiBox[0], right=HiBox[1], bottom=HiBox[2], top=HiBox[3], hspace=0.2, wspace=0.2)\n",
    "gsLo  = gridspec.GridSpec(2, 3, height_ratios=[1.0, 0.55], width_ratios=[1.0, 1.0, 1.0])\n",
    "gsLo.update(left=LoBox[0], right=LoBox[1], bottom=LoBox[2], top=LoBox[3], hspace=0.2, wspace=0.2)\n",
    "axHi  = [[plot.subplot(gsHi[0, i]), plot.subplot(gsHi[1, i])] for i in range(len(genCut))]\n",
    "axLo  = [[plot.subplot(gsLo[0, i]), plot.subplot(gsLo[1, i])] for i in range(len(genCut))]\n",
    "\n",
    "gsLL = gridspec.GridSpec(1, 1, height_ratios=[1.0], width_ratios=[1.0])\n",
    "gsLL.update(left=LLBox[0], right=LLBox[1], bottom=LLBox[2], top=LLBox[3], hspace=0.4, wspace=0.1)\n",
    "axLL = [plot.subplot(gsLL[0, 0])]\n",
    "\n",
    "# Overhead line\n",
    "\n",
    "lineprops = { 'lw' : axwidth/2., 'ls' : '-', 'alpha' : 1.0 }\n",
    "pprops    = { 'xlim' : [0, 1], 'xticks' : [], 'ylim' : [0, 1], 'yticks' : [], \n",
    "              'hide' : ['top','bottom','left','right'], 'plotprops' : lineprops }\n",
    "tprops    = dict(color=bkcolor, family=fontfamily, size=sizelabel)\n",
    "axS[0].text(2.4/6.-0.005, 0.6, 'Time (generations)', ha='center', va='top',    **tprops)\n",
    "axS[0].text(           0, 0.2,              r'$25$', ha='center', va='center', **tprops)\n",
    "axS[0].text(2.4/6.-0.005, 0.2,             r'$150$', ha='center', va='center', **tprops)\n",
    "axS[0].text(4.9/6.+0.015, 0.2,             r'$400$', ha='center', va='center', **tprops)\n",
    "ay = 0.05\n",
    "axS[0].arrow(x=-0.30, y=ay, dx=1+0.30, dy=0, head_width=0.25, head_length=0.025, overhang=0.5, fc=bkcolor,\n",
    "             ec=None, lw=sizeline/2., clip_on=False, zorder=100)\n",
    "mp.plot(type='line', ax=axS[0], x=[[0,1]], y=[[ay,ay]], colors=[bkcolor], **pprops)\n",
    "\n",
    "# PART 1: densities\n",
    "\n",
    "print('Plotting densities...')\n",
    "\n",
    "legendprops  = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                 'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "contourprops = { 'cmap' : 'Oranges', 'colors' : bkcolor, 'vmin' : 0.001, 'vmax' : 0.10, 'extend' : 'max',\n",
    "                 'linewidths' : sizeline/2., 'levels' : [0.001, 0.005, 0.015, 0.035, 0.075, 0.15]}\n",
    "\n",
    "dx     = 0.1\n",
    "dy     = dx * (1. / 16.)\n",
    "cx     = np.arange(-2., 14.+dx, dx)\n",
    "cy     = np.arange( 0.,  1.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.5\n",
    "sy = sx * (1. / 16.)\n",
    "    \n",
    "pprops = { 'xlim'      : [-2., 14.],\n",
    "           'xticks'    : [0.,  4.,  8., 12.],\n",
    "           'xticklabels' : [r'$0$', r'$-4$', r'$-8$', r'$-12$'],\n",
    "           'xminorticks' : [-2., 2., 6., 10., 14.],\n",
    "           'ylim'      : [ 0., 1.],\n",
    "           'yticks'    : [ 0., 0.25, 0.50, 0.75, 1.00],\n",
    "           'xlabel'    : r'$E_c$',\n",
    "           'ylabel'    : r'$Q$',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "r      = 1.\n",
    "cprops = dict(width=r, height=r/14., ec='k', fc='w', alpha=1.0, zorder=100)\n",
    "\n",
    "for i in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if i>0:\n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks']   = []\n",
    "        tempprops['hide']     = ['left']\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xHi[i])):\n",
    "        zz += np.exp(-((xHi[i][k]-xx)**2/sx**2) - ((yHi[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xHi[i]))\n",
    "    mp.plot(type='contour', ax=axHi[i][0], x=[xx], y=[yy], z=[zz], **tempprops)\n",
    "    circHi = matplotlib.patches.Ellipse(xy=(4.,0.8), **cprops)\n",
    "    axHi[i][0].add_patch(circHi)\n",
    "    if i==0:\n",
    "        axHi[i][0].annotate('Initial\\nparameters', xy=(3.7, 0.78), xytext=(1.75, 0.50), ha='center', va='center',\n",
    "                            arrowprops=dict(arrowstyle='simple', fc='k', mutation_scale=10., shrinkA=3., shrinkB=3.), \n",
    "                            **def_labelprops)\n",
    "    if i==len(genCut)-1:\n",
    "        axHi[i][0].text(2., 0.90, 'CH103-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "\n",
    "for i in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if i>0:\n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks']   = []\n",
    "        tempprops['hide']     = ['left']\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xLo[i])):\n",
    "        zz += np.exp(-((xLo[i][k]-xx)**2/sx**2) - ((yLo[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xLo[i]))\n",
    "    mp.plot(type='contour', ax=axLo[i][0], x=[xx], y=[yy], z=[zz], **tempprops)\n",
    "    circLo = matplotlib.patches.Ellipse(xy=(0.,0.8), **cprops)\n",
    "    axLo[i][0].add_patch(circLo)\n",
    "    if i==len(genCut)-1:\n",
    "        axLo[i][0].text(-0.5, 0.90, 'PGT121-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "        axLo[i][0].text(-2., 0.35, '3BNC60-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "\n",
    "# PART 2: breadth\n",
    "\n",
    "print('Plotting breadth...')\n",
    "\n",
    "legendprops = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "lineprops   = { 'lw' : sizeline, 'ls' : '-', 'alpha' : 1.0 }\n",
    "cHist       = '#3182bd'\n",
    "\n",
    "histprops = { 'histtype' : 'bar', 'lw' : sizeline/2, 'rwidth' : 0.8, 'ls' : 'solid', 'alpha' : 0.7, 'ec' : bkcolor }\n",
    "\n",
    "pprops = { 'xlim'        : [-4., 8.],\n",
    "           'xticks'      : [-4, 0.,  4., 8.],\n",
    "           'xticklabels' : [r'$4$', r'$0$', r'$-4$', r'$-8$'],\n",
    "           'xminorticks' : [-2,  2., 6.],\n",
    "           'ylim'        : [ 0., 0.50],\n",
    "           'yticks'      : [ 0., 0.25, 0.50],\n",
    "           'yticklabels' : [  0,   25,   50],\n",
    "           'xlabel'      : 'Median binding free energy',\n",
    "           'ylabel'      : 'Frequency (%)',\n",
    "           'plotprops'   : histprops,\n",
    "           'bins'        : np.arange(-4.0, 8.5, 0.5),\n",
    "           'combine'     : True }\n",
    "\n",
    "for g in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if g>0: \n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks'] = []\n",
    "        tempprops['hide']   = ['left']\n",
    "    mp.plot(type='hist', ax=axHi[g][1], x=[bHiTrueMed[g]], colors=[cHist], **tempprops)\n",
    "    mp.plot(type='hist', ax=axLo[g][1], x=[bLoTrueMed[g]], colors=[cHist], **tempprops)\n",
    "    \n",
    "# PART 3: Q vs breadth\n",
    "\n",
    "print('Plotting Q vs breadth...')\n",
    "\n",
    "dx     = 0.05\n",
    "dy     = dx * (1. / 7.)\n",
    "cx     = np.arange(1., 8.+dx, dx)\n",
    "cy     = np.arange(0., 1.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.4\n",
    "sy = sx * (1. / 7.)\n",
    "\n",
    "zz = np.zeros(xx.shape)\n",
    "for k in range(len(bLoTrueMed[-1])):\n",
    "    zz += np.exp(-((bLoTrueMed[-1][k]-xx)**2/sx**2) - ((yLo[-1][k]-yy)**2/sy**2))\n",
    "zz /= float(len(bLoTrueMed[-1]))\n",
    "\n",
    "contourprops['cmap'] = 'Blues'\n",
    "scatterprops = { 'lw' : 0, 's' : 8., 'marker' : 'o' }\n",
    "\n",
    "pprops = { 'xlim'      : [ 1., 8.],\n",
    "           'xticks'    : [ 2., 4., 6., 8.],\n",
    "           'xticklabels' : [r'$-2$', r'$-4$', r'$-6$', r'$-8$'],\n",
    "           'xaxstart'  : 1,\n",
    "           'ylim'      : [ 0., 1.],\n",
    "           'yticks'    : [ 0., 0.25, 0.50, 0.75, 1.00],\n",
    "           'xlabel'    : 'Median binding free energy',\n",
    "           'ylabel'    : r'$Q$',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "mp.plot(type='contour', ax=axLL[0], x=[xx], y=[yy], z=[zz], **pprops)\n",
    "axLL[0].text(4.5, 0.80, 'PGT121-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "axLL[0].text(4.5, 0.25, '3BNC60-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "    \n",
    "# Schematic\n",
    "\n",
    "hiSchem  = mpimg.imread('figures/hiEc.png')\n",
    "loSchem  = mpimg.imread('figures/loEc.png')\n",
    "keySchem = mpimg.imread('figures/legend2.png')\n",
    "axS[1].imshow(hiSchem)\n",
    "axS[2].imshow(loSchem)\n",
    "axKey.imshow(keySchem, zorder=-200)\n",
    "for ax in [axS[1], axS[2], axKey]:\n",
    "    for a in ['left', 'right', 'bottom', 'top']: \n",
    "        ax.spines[a].set_visible(False)\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "axS[2].text( 0.02, LoBox[3]-0.01, 'B', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS[1].text( 0.02, HiBox[3]-0.01, 'A', transform=fig.transFigure, **def_sublabelprops)\n",
    "axLL[0].text(0.02, LoBox[2]-0.08, 'C', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS[1].text(0.07, (HiBox[2]+HiBox[3])/2.-0.03, 'Strong initial ' + r'$E_c$', ha='center', va='center',\n",
    "            transform=fig.transFigure, **def_labelprops)\n",
    "axS[2].text(0.07, (LoBox[2]+LoBox[3])/2.-0.03, 'Weak initial ' + r'$E_c$', ha='center', va='center', \n",
    "            transform=fig.transFigure, **def_labelprops)\n",
    "tprops['size'] = sizelabel\n",
    "axKey.text(0.080, 0.965,        'viral\\nspike', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "axKey.text(0.175, 0.965,   'conserved\\nregion', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "axKey.text(0.290, 0.965, 'antibody\\nfootprint', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "\n",
    "# Background\n",
    "\n",
    "cBG   = '#F5F5F5'\n",
    "dy    = 0.07\n",
    "bgHi  = axS[1].axis()\n",
    "recHi = matplotlib.patches.Rectangle(xy=(0,HiBox[2]-0.8*dy), width=1., height=HiBox[3]-HiBox[2]+(1.*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recHi = axS[1].add_patch(recHi)\n",
    "\n",
    "bgLo  = axS[2].axis()\n",
    "recLo = matplotlib.patches.Rectangle(xy=(0,LoBox[2]-0.8*dy), width=1., height=LoBox[3]-LoBox[2]+(1.*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recLo = axS[2].add_patch(recLo)\n",
    "    \n",
    "# Save data\n",
    "plot.savefig('figures/figure-4-figure-supplement-2.pdf', dpi = 1000, \n",
    "              facecolor = fig.get_facecolor(), edgecolor=None, **figprops)\n",
    "plot.close(fig)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure 4-figure supplement 3. Trajectories of binding energy/flexibility during affinity maturation with high/low $E_0$\n",
    "\n",
    "Note: This figure is the same as Figures 4 and S5, just using different input files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create alternate \"test panel\" of 100 sequences with different Ev\n",
    "\n",
    "nTrials   = 10000\n",
    "nPanel    = 100\n",
    "testPanel = np.array([getdE(o=0, mu=3.0, sigma=1.0) for i in range(nPanel)])\n",
    "\n",
    "# Read through largest clone histories and output binding energies with the test panel\n",
    "\n",
    "fIn  = ['data/loEc-hiE0-output-largest-clone.csv.gz',\n",
    "        'data/loEc-loE0-output-largest-clone.csv.gz',\n",
    "        'data/loEc-simE0-output-largest-clone.csv.gz']\n",
    "fOut = ['data/loEc-hiE0-panel-breadth.csv.gz',\n",
    "        'data/loEc-loE0-panel-breadth.csv.gz',\n",
    "        'data/loEc-simE0-panel-breadth.csv.gz']\n",
    "\n",
    "for k in range(len(fIn)):\n",
    "    if os.path.isfile(fOut[k]):\n",
    "        continue\n",
    "    else:\n",
    "        f  = open(fOut[k], 'w')\n",
    "        f.write('trial,generation,Q,Ec%s\\n' % (''.join([',Ev%d' % (i+1) for i in range(nPanel)])))\n",
    "\n",
    "        df = pd.read_csv(fIn[k], memory_map=True)\n",
    "\n",
    "        for t in range(nTrials):\n",
    "            trialDf = df[df.trial==t]\n",
    "\n",
    "            if len(trialDf)==0: continue\n",
    "\n",
    "            # Set starting binding energies with test panel\n",
    "\n",
    "            panelEc = trialDf.iloc[0].Ec\n",
    "            panelEv = deepcopy(testPanel)\n",
    "\n",
    "            f.write('%d,%d,%lf,%lf' % (trialDf.iloc[0].trial,trialDf.iloc[0].generation,trialDf.iloc[0].Q,panelEc))\n",
    "            for i in range(len(testPanel)): f.write(',%lf' % panelEv[i])\n",
    "            f.write('\\n')\n",
    "\n",
    "            # Iterate over successive mutations and update binding energies\n",
    "\n",
    "            for i in range(1,len(trialDf)):\n",
    "\n",
    "                if trialDf.iloc[i].Q==trialDf.iloc[i-1].Q:\n",
    "                    dEc = trialDf.iloc[i].Ec-trialDf.iloc[i-1].Ec\n",
    "                    if dEc==0:\n",
    "                        for j in range(len(testPanel)): panelEv[j] += getdE()\n",
    "                    else: panelEc += dEc\n",
    "\n",
    "                f.write('%d,%d,%lf,%lf' % (trialDf.iloc[i].trial,trialDf.iloc[i].generation,trialDf.iloc[i].Q,panelEc))\n",
    "                for j in range(len(testPanel)): f.write(',%lf' % panelEv[j])\n",
    "                f.write('\\n')\n",
    "                f.flush()\n",
    "\n",
    "        f.close() # files can be gzipped below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compress panel files\n",
    "\n",
    "fin  = ['data/loEc-hiE0-output-largest-clone.csv',\n",
    "        'data/loEc-hiE0-output-memory.csv',\n",
    "        'data/loEc-hiE0-output-total.csv',\n",
    "        'data/loEc-hiE0-panel-breadth.csv',\n",
    "        'data/loEc-loE0-output-largest-clone.csv',\n",
    "        'data/loEc-loE0-output-memory.csv',\n",
    "        'data/loEc-loE0-output-total.csv',\n",
    "        'data/loEc-loE0-panel-breadth.csv',\n",
    "        'data/loEc-simE0-output-largest-clone.csv',\n",
    "        'data/loEc-simE0-output-memory.csv',\n",
    "        'data/loEc-simE0-output-total.csv',\n",
    "        'data/loEc-simE0-panel-breadth.csv']\n",
    "fout = ['data/loEc-hiE0-output-largest-clone',\n",
    "        'data/loEc-hiE0-output-memory',\n",
    "        'data/loEc-hiE0-output-total',\n",
    "        'data/loEc-hiE0-panel-breadth',\n",
    "        'data/loEc-loE0-output-largest-clone',\n",
    "        'data/loEc-loE0-output-memory',\n",
    "        'data/loEc-loE0-output-total',\n",
    "        'data/loEc-loE0-panel-breadth',\n",
    "        'data/loEc-simE0-output-largest-clone',\n",
    "        'data/loEc-simE0-output-memory',\n",
    "        'data/loEc-simE0-output-total',\n",
    "        'data/loEc-simE0-panel-breadth']\n",
    "\n",
    "for k in range(len(fin)):\n",
    "    df = pd.read_csv(fin[k], memory_map=True)\n",
    "    df.to_csv(fout[k]+'.csv.gz', compression='gzip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data: Q vs Ec plot\n",
    "\n",
    "dfHi = pd.read_csv('data/loEc-hiE0-output-largest-clone.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-loE0-output-largest-clone.csv.gz', memory_map=True)\n",
    "\n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "genCut = [25, 150, 400]\n",
    "\n",
    "xHi, yHi = [], []\n",
    "xLo, yLo = [], []\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        xHi.append(np.array(nextdfHi['Ec']))\n",
    "        yHi.append(np.array(nextdfHi['Q']))\n",
    "        xLo.append(np.array(nextdfLo['Ec']))\n",
    "        yLo.append(np.array(nextdfLo['Q']))\n",
    "        \n",
    "# Load data: breadth plot\n",
    "\n",
    "ECut = [3.0, 6.0, 9.0]\n",
    "\n",
    "dfHi = pd.read_csv('data/loEc-hiE0-panel-breadth.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-loE0-panel-breadth.csv.gz', memory_map=True)\n",
    "\n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "nPanel     = 100\n",
    "bHi        = []\n",
    "bLo        = []\n",
    "bHiTrueMed = []\n",
    "bLoTrueMed = []\n",
    "EvLabel    = ['Ev%d' % (i+1) for i in range(nPanel)]\n",
    "denom      = float(nPanel)\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        tempQ  = np.array([ q * np.ones(nPanel) for  q in nextdfHi.Q])\n",
    "        tempEc = np.array([ec * np.ones(nPanel) for ec in nextdfHi.Ec])\n",
    "        EHi    = np.array(getE(tempEc, np.array(nextdfHi[EvLabel]), tempQ, E0=6., lam=0.9))\n",
    "        \n",
    "        tempQ  = np.array([ q * np.ones(nPanel) for  q in nextdfLo.Q])\n",
    "        tempEc = np.array([ec * np.ones(nPanel) for ec in nextdfLo.Ec])\n",
    "        ELo    = np.array(getE(tempEc, np.array(nextdfLo[EvLabel]), tempQ, E0=2., lam=0.9))\n",
    "        \n",
    "        bHi.append([[np.sum(E>c)/denom for E in EHi] for c in ECut])\n",
    "        bLo.append([[np.sum(E>c)/denom for E in ELo] for c in ECut])\n",
    "        \n",
    "        bHiTrueMed.append([np.median(E) for E in EHi])\n",
    "        bLoTrueMed.append([np.median(E) for E in ELo])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting densities...\n",
      "Plotting breadth...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Plot the figure\n",
    "\n",
    "nTrials = 10000\n",
    "reload(mp)\n",
    "\n",
    "w     = double_column\n",
    "goldh = (3.3/3.0)*w\n",
    "fig   = plot.figure(figsize=(w, goldh))\n",
    "\n",
    "HiBox = [0.20, 0.98, 0.61, 0.91]\n",
    "LoBox = [0.20, 0.98, 0.23, 0.53]\n",
    "\n",
    "gsS   = gridspec.GridSpec(1, 2, width_ratios=[1.0, 2.5], height_ratios=[1.0])\n",
    "gsS.update(left=0, right=HiBox[1], bottom=0.93, top=1.0)\n",
    "gsHiS = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsHiS.update(left=0.02, right=0.12, bottom=HiBox[2]+((HiBox[3]-HiBox[2])*0.5/1.3), top=HiBox[3])\n",
    "gsLoS = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsLoS.update(left=0.02, right=0.12, bottom=LoBox[2]+((LoBox[3]-LoBox[2])*0.5/1.3), top=LoBox[3])\n",
    "axS   = [plot.subplot(gsS[ 0, 1]), plot.subplot(gsHiS[ 0, 0]), plot.subplot(gsLoS[ 0, 0])]\n",
    "axKey = plot.subplot(gsS[0, 0])\n",
    "\n",
    "gsHi  = gridspec.GridSpec(2, 3, height_ratios=[1.0, 0.55], width_ratios=[1.0, 1.0, 1.0])\n",
    "gsHi.update(left=HiBox[0], right=HiBox[1], bottom=HiBox[2], top=HiBox[3], hspace=0.2, wspace=0.2)\n",
    "gsLo  = gridspec.GridSpec(2, 3, height_ratios=[1.0, 0.55], width_ratios=[1.0, 1.0, 1.0])\n",
    "gsLo.update(left=LoBox[0], right=LoBox[1], bottom=LoBox[2], top=LoBox[3], hspace=0.2, wspace=0.2)\n",
    "axHi  = [[plot.subplot(gsHi[0, i]), plot.subplot(gsHi[1, i])] for i in range(len(genCut))]\n",
    "axLo  = [[plot.subplot(gsLo[0, i]), plot.subplot(gsLo[1, i])] for i in range(len(genCut))]\n",
    "\n",
    "# Overhead line\n",
    "\n",
    "lineprops = { 'lw' : axwidth/2., 'ls' : '-', 'alpha' : 1.0 }\n",
    "pprops    = { 'xlim' : [0, 1], 'xticks' : [], 'ylim' : [0, 1], 'yticks' : [], \n",
    "              'hide' : ['top','bottom','left','right'], 'plotprops' : lineprops }\n",
    "tprops    = dict(color=bkcolor, family=fontfamily, size=sizelabel)\n",
    "axS[0].text(2.4/6.-0.005, 0.6, 'Time (generations)', ha='center', va='top',    **tprops)\n",
    "axS[0].text(           0, 0.2,              r'$25$', ha='center', va='center', **tprops)\n",
    "axS[0].text(2.4/6.-0.005, 0.2,             r'$150$', ha='center', va='center', **tprops)\n",
    "axS[0].text(4.9/6.+0.015, 0.2,             r'$400$', ha='center', va='center', **tprops)\n",
    "ay = 0.05\n",
    "axS[0].arrow(x=-0.30, y=ay, dx=1+0.30, dy=0, head_width=0.25, head_length=0.025, overhang=0.5, fc=bkcolor,\n",
    "             ec=None, lw=sizeline/2., clip_on=False, zorder=100)\n",
    "mp.plot(type='line', ax=axS[0], x=[[0,1]], y=[[ay,ay]], colors=[bkcolor], **pprops)\n",
    "\n",
    "# PART 1: densities\n",
    "\n",
    "print('Plotting densities...')\n",
    "\n",
    "legendprops  = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                 'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "contourprops = { 'cmap' : 'Oranges', 'colors' : bkcolor, 'vmin' : 0.001, 'vmax' : 0.10, 'extend' : 'max',\n",
    "                 'linewidths' : sizeline/2., 'levels' : [0.001, 0.005, 0.015, 0.035, 0.075, 0.15]}\n",
    "\n",
    "dx     = 0.1\n",
    "dy     = dx * (1. / 16.)\n",
    "cx     = np.arange(-2., 14.+dx, dx)\n",
    "cy     = np.arange( 0.,  1.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.5\n",
    "sy = sx * (1. / 16.)\n",
    "    \n",
    "pprops = { 'xlim'      : [-2., 14.],\n",
    "           'xticks'    : [0.,  4.,  8., 12.],\n",
    "           'xticklabels' : [r'$0$', r'$-4$', r'$-8$', r'$-12$'],\n",
    "           'xminorticks' : [-2., 2., 6., 10., 14.],\n",
    "           'ylim'      : [ 0., 1.],\n",
    "           'yticks'    : [ 0., 0.25, 0.50, 0.75, 1.00],\n",
    "           'xlabel'    : r'$E_c$',\n",
    "           'ylabel'    : r'$Q$',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "r      = 1.\n",
    "cprops = dict(width=r, height=r/14., ec='k', fc='w', alpha=1.0, zorder=100)\n",
    "\n",
    "for i in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if i>0:\n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks']   = []\n",
    "        tempprops['hide']     = ['left']\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xHi[i])):\n",
    "        zz += np.exp(-((xHi[i][k]-xx)**2/sx**2) - ((yHi[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xHi[i]))\n",
    "    mp.plot(type='contour', ax=axHi[i][0], x=[xx], y=[yy], z=[zz], **tempprops)\n",
    "    circHi = matplotlib.patches.Ellipse(xy=(0.,0.8), **cprops)\n",
    "    axHi[i][0].add_patch(circHi)\n",
    "    if i==len(genCut)-1:\n",
    "        axHi[i][0].text(-2., 0.25, '3BNC60-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "\n",
    "for i in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if i>0:\n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks']   = []\n",
    "        tempprops['hide']     = ['left']\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xLo[i])):\n",
    "        zz += np.exp(-((xLo[i][k]-xx)**2/sx**2) - ((yLo[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xLo[i]))\n",
    "    mp.plot(type='contour', ax=axLo[i][0], x=[xx], y=[yy], z=[zz], **tempprops)\n",
    "    circLo = matplotlib.patches.Ellipse(xy=(0.,0.8), **cprops)\n",
    "    axLo[i][0].add_patch(circLo)\n",
    "    if i==len(genCut)-1:\n",
    "        axLo[i][0].text(-1.0, 0.90, 'CH103-like', size=sizelabel, color='k', ha='left', va='center',\n",
    "                        bbox={'facecolor':'w', 'alpha':0.5, 'pad':2, 'lw':0})\n",
    "        \n",
    "\n",
    "# PART 2: breadth\n",
    "\n",
    "print('Plotting breadth...')\n",
    "\n",
    "legendprops = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "lineprops   = { 'lw' : sizeline, 'ls' : '-', 'alpha' : 1.0 }\n",
    "cHist       = '#3182bd'\n",
    "\n",
    "histprops = { 'histtype' : 'bar', 'lw' : sizeline/2, 'rwidth' : 0.8, 'ls' : 'solid', 'alpha' : 0.7, 'ec' : bkcolor }\n",
    "\n",
    "pprops = { 'xlim'        : [-2., 10.],\n",
    "           'xticks'      : [0.,  4., 8.],\n",
    "           'xticklabels' : [r'$0$', r'$-4$', r'$-8$'],\n",
    "           'xminorticks' : [-2,  2., 6., 10.],\n",
    "           'ylim'        : [ 0., 0.50],\n",
    "           'yticks'      : [ 0., 0.25, 0.50],\n",
    "           'yticklabels' : [  0,   25,   50],\n",
    "           'xlabel'      : 'Median binding free energy',\n",
    "           'ylabel'      : 'Frequency (%)',\n",
    "           'plotprops'   : histprops,\n",
    "           'bins'        : np.arange(-2.0, 10.5, 0.5),\n",
    "           'combine'     : True }\n",
    "\n",
    "for g in range(len(genCut)):\n",
    "    tempprops = pprops.copy()\n",
    "    if g>0: \n",
    "        del tempprops['ylabel']\n",
    "        tempprops['yticks'] = []\n",
    "        tempprops['hide']   = ['left']\n",
    "    mp.plot(type='hist', ax=axHi[g][1], x=[bHiTrueMed[g]], colors=[cHist], **tempprops)\n",
    "    mp.plot(type='hist', ax=axLo[g][1], x=[bLoTrueMed[g]], colors=[cHist], **tempprops)\n",
    "    \n",
    "# Schematic\n",
    "\n",
    "hiSchem  = mpimg.imread('figures/loEc.png')\n",
    "loSchem  = mpimg.imread('figures/loEc.png')\n",
    "keySchem = mpimg.imread('figures/legend2.png')\n",
    "axS[1].imshow(hiSchem)\n",
    "axS[2].imshow(loSchem)\n",
    "axKey.imshow(keySchem, zorder=-200)\n",
    "for ax in [axS[1], axS[2], axKey]:\n",
    "    for a in ['left', 'right', 'bottom', 'top']: \n",
    "        ax.spines[a].set_visible(False)\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "axS[2].text( 0.02, LoBox[3]-0.01, 'B', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS[1].text( 0.02, HiBox[3]-0.01, 'A', transform=fig.transFigure, **def_sublabelprops)\n",
    "axLL[0].text(0.02, LoBox[2]-0.08, 'C', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS[1].text(0.07, (HiBox[2]+HiBox[3])/2.-0.04, 'Weak initial ' + r'$E_c$' + '\\nStrong ' + r'$E_0$', \n",
    "            ha='center', va='center', transform=fig.transFigure, **def_labelprops)\n",
    "axS[2].text(0.07, (LoBox[2]+LoBox[3])/2.-0.04, 'Weak initial ' + r'$E_c$' + '\\nWeak ' + r'$E_0$', \n",
    "            ha='center', va='center', transform=fig.transFigure, **def_labelprops)\n",
    "tprops['size'] = sizelabel\n",
    "axKey.text(0.080, 0.965,        'viral\\nspike', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "axKey.text(0.175, 0.965,   'conserved\\nregion', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "axKey.text(0.290, 0.965, 'antibody\\nfootprint', transform=fig.transFigure, ha='center', va='center', **tprops)\n",
    "\n",
    "# Background\n",
    "\n",
    "cBG   = '#F5F5F5'\n",
    "dy    = 0.07\n",
    "bgHi  = axS[1].axis()\n",
    "recHi = matplotlib.patches.Rectangle(xy=(0,HiBox[2]-0.8*dy), width=1., height=HiBox[3]-HiBox[2]+(1.*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recHi = axS[1].add_patch(recHi)\n",
    "\n",
    "bgLo  = axS[2].axis()\n",
    "recLo = matplotlib.patches.Rectangle(xy=(0,LoBox[2]-0.8*dy), width=1., height=LoBox[3]-LoBox[2]+(1.*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recLo = axS[2].add_patch(recLo)\n",
    "    \n",
    "# Save data\n",
    "plot.savefig('figures/figure-4-figure-supplement-3.pdf', dpi = 1000, \n",
    "              facecolor = fig.get_facecolor(), edgecolor=None, **figprops)\n",
    "plot.close(fig)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Figure 6-figure supplement 1. Accumulation of mutations in the FR and CDR regions with high/low $E_0$\n",
    "\n",
    "Note: This figure is the same as Figure 6, just using different input files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data here\n",
    "\n",
    "dfHi = pd.read_csv('data/loEc-hiE0-output-largest-clone.csv.gz', memory_map=True)\n",
    "dfLo = pd.read_csv('data/loEc-loE0-output-largest-clone.csv.gz', memory_map=True)\n",
    "        \n",
    "dfHi.set_index(['trial'], inplace=True)\n",
    "dfLo.set_index(['trial'], inplace=True)\n",
    "\n",
    "genCut = [25, 150, 400]\n",
    "\n",
    "xHi, yHi = [], []\n",
    "xLo, yLo = [], []\n",
    "\n",
    "lastdfHi = dfHi[dfHi['generation']==0]\n",
    "lastdfLo = dfLo[dfLo['generation']==0]\n",
    "\n",
    "for t in range(1,501):\n",
    "    nextdfHi = dfHi[dfHi['generation']==t].combine_first(lastdfHi)\n",
    "    lastdfHi = nextdfHi\n",
    "    nextdfLo = dfLo[dfLo['generation']==t].combine_first(lastdfLo)\n",
    "    lastdfLo = nextdfLo\n",
    "    \n",
    "    if t in genCut:\n",
    "        xHi.append(np.array(nextdfHi['CDR_mutations']))\n",
    "        yHi.append(np.array(nextdfHi['FR_mutations']))\n",
    "        xLo.append(np.array(nextdfLo['CDR_mutations']))\n",
    "        yLo.append(np.array(nextdfLo['FR_mutations']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting number of mutations...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Plot the figure\n",
    "\n",
    "w     = double_column\n",
    "goldh = (2.0/3.)*w\n",
    "fig   = plot.figure(figsize=(w, goldh))\n",
    "\n",
    "HiBox = [0.20, 0.98, 0.60, 0.90]\n",
    "LoBox = [0.20, 0.98, 0.16, 0.46]\n",
    "\n",
    "gsS   = gridspec.GridSpec(1, 1, width_ratios=[1.0], height_ratios=[1.0])\n",
    "gsS.update(left=HiBox[0], right=HiBox[1], bottom=0.93, top=1.0)\n",
    "axS   = plot.subplot(gsS[ 0, 0])\n",
    "\n",
    "gsHi  = gridspec.GridSpec(1, 3, width_ratios=[1.0, 1.0, 1.0])\n",
    "gsHi.update(left=HiBox[0], right=HiBox[1], bottom=HiBox[2], top=HiBox[3], hspace=0.5, wspace=0.5)\n",
    "gsLo  = gridspec.GridSpec(1, 3, width_ratios=[1.0, 1.0, 1.0])\n",
    "gsLo.update(left=LoBox[0], right=LoBox[1], bottom=LoBox[2], top=LoBox[3], hspace=0.5, wspace=0.5)\n",
    "\n",
    "axHi  = [plot.subplot(gsHi[i]) for i in range(len(genCut))]\n",
    "axLo  = [plot.subplot(gsLo[i]) for i in range(len(genCut))]\n",
    "\n",
    "# Overhead line\n",
    "\n",
    "lineprops = { 'lw' : axwidth/2., 'ls' : '-', 'alpha' : 1.0 }\n",
    "pprops    = { 'xlim' : [0, 1], 'xticks' : [], 'ylim' : [0, 1], 'yticks' : [], \n",
    "              'hide' : ['top','bottom','left','right'], 'plotprops' : lineprops }\n",
    "tprops    = dict(color=bkcolor, family=fontfamily, size=sizelabel)\n",
    "axS.text(3./6.-0.005, 1.0, 'Time (generations)', ha='center', va='top',    **tprops)\n",
    "axS.text(1./6.-0.045, 0.3,              r'$25$', ha='center', va='center', **tprops)\n",
    "axS.text(3./6.-0.005, 0.3,             r'$150$', ha='center', va='center', **tprops)\n",
    "axS.text(5./6.+0.055, 0.3,             r'$400$', ha='center', va='center', **tprops)\n",
    "ay = 0.05\n",
    "axS.arrow(x=0, y=ay, dx=1, dy=0, head_width=0.25, head_length=0.025, overhang=0.5, fc=bkcolor,\n",
    "             ec=None, lw=sizeline/2., clip_on=False, zorder=100)\n",
    "mp.plot(type='line', ax=axS, x=[[0,1]], y=[[ay,ay]], colors=[bkcolor], **pprops)\n",
    "\n",
    "# Number of CDR/FWR mutations\n",
    "\n",
    "print('Plotting number of mutations...')\n",
    "\n",
    "levels = [0.02] + list(0.02 * np.logspace(1., 5., num=5, base=1.5))\n",
    "\n",
    "legendprops  = { 'loc' : 4, 'frameon' : False, 'scatterpoints' : 1, 'handletextpad' : 0.1, \n",
    "                 'prop' : {'size' : sizelabel}, 'ncol' : 1 }\n",
    "contourprops = { 'cmap' : 'Greens', 'colors' : bkcolor, 'extend' : 'max', 'linewidths' : sizeline/2.,\n",
    "                 'vmin' : 0.001, 'vmax' : 0.15 , 'levels' : levels }\n",
    "\n",
    "dx     = 0.1\n",
    "dy     = dx\n",
    "cx     = np.arange(0., 12.+dx, dx)\n",
    "cy     = np.arange(0., 12.+dy, dy)\n",
    "xx, yy = np.meshgrid(cx, cy)\n",
    "\n",
    "sx = 0.9\n",
    "sy = sx\n",
    "\n",
    "pprops = { 'xlim'      : [ 0., 12.],\n",
    "           'xticks'    : [ 0., 4., 8., 12.],\n",
    "           'ylim'      : [ 0., 12.],\n",
    "           'yticks'    : [ 0., 4., 8., 12.],\n",
    "           'xlabel'    : 'CDR mutations',\n",
    "           'ylabel'    : 'Framework mutations',\n",
    "           'plotprops' : contourprops,\n",
    "           'zorder'    : 10 }\n",
    "\n",
    "for i in range(len(xHi)):\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xHi[i])):\n",
    "        zz += np.exp(-((xHi[i][k]-xx)**2/sx**2) - ((yHi[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xHi[i]))\n",
    "    mp.plot(type='contour', ax=axHi[i], x=[xx], y=[yy], z=[zz], **pprops)\n",
    "\n",
    "for i in range(len(xLo)):\n",
    "    zz = np.zeros(xx.shape)\n",
    "    for k in range(len(xLo[i])):\n",
    "        zz += np.exp(-((xLo[i][k]-xx)**2/sx**2) - ((yLo[i][k]-yy)**2/sy**2))\n",
    "    zz /= float(len(xLo[i]))\n",
    "    mp.plot(type='contour', ax=axLo[i], x=[xx], y=[yy], z=[zz], **pprops)\n",
    "\n",
    "# Background\n",
    "\n",
    "axS.text(HiBox[0]-0.06, HiBox[3]-0.01, 'A', transform=fig.transFigure, **def_sublabelprops)\n",
    "axS.text(LoBox[0]-0.06, LoBox[3]-0.01, 'B', transform=fig.transFigure, **def_sublabelprops)\n",
    "\n",
    "cBG   = '#F5F5F5'\n",
    "dx    = 0.08\n",
    "dy    = 0.09\n",
    "bgHi  = axS.axis()\n",
    "recHi = matplotlib.patches.Rectangle(xy=(HiBox[0]-dx,HiBox[2]-dy), width=0.8+dx, height=HiBox[3]-HiBox[2]+(1.2*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recHi = axS.add_patch(recHi)\n",
    "\n",
    "bgLo  = axS.axis()\n",
    "recLo = matplotlib.patches.Rectangle(xy=(LoBox[0]-dx,LoBox[2]-dy), width=0.8+dx, height=LoBox[3]-LoBox[2]+(1.2*dy), \n",
    "                                     transform=fig.transFigure, ec=None, fc=cBG, clip_on=False, zorder=-100)\n",
    "recLo = axS.add_patch(recLo)\n",
    "    \n",
    "# Save data\n",
    "plot.savefig('figures/figure-6-figure-supplement-1.pdf', dpi = 1000, \n",
    "              facecolor = fig.get_facecolor(), edgecolor=None, **figprops)\n",
    "plot.close(fig)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Parameters for affinity maturation simulations\n",
    "\n",
    "`am.py` uses global parameters set at the start of the function for simulations. The parameters that we used in the main cases considered in the paper,\n",
    "\n",
    "1. single antigen,\n",
    "2. high/low $E_c$ (standard), and \n",
    "3. high/low $E_c$ (low overlap, $\\lambda=0.8$)\n",
    "4. high/low $E_0$, low $E_c$\n",
    "\n",
    "are given below. Any parameters not listed below are the same for all data sets and match with those currently in `am.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters for single antigen case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_var  = 0.10        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 0.        # starting value for Ec\n",
    "nb_Ag    = 1         # number of antigens\n",
    "conc     = 1.11      # antigen concentration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters for standard case with multiple antigens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# High Ec\n",
    "\n",
    "p_var  = 0.10        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 4.        # starting value for Ec\n",
    "nb_Ag    = 10        # number of antigens\n",
    "conc     = 0.98      # antigen concentration\n",
    "\n",
    "# Low Ec\n",
    "\n",
    "p_var  = 0.10        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 0.        # starting value for Ec\n",
    "nb_Ag    = 10        # number of antigens\n",
    "conc     = 1.18      # antigen concentration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters for low overlap ($\\lambda=0.8$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# High Ec\n",
    "\n",
    "p_var  = 0.20        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 4.        # starting value for Ec\n",
    "nb_Ag    = 5         # number of antigens\n",
    "conc     = 1.08      # antigen concentration\n",
    "\n",
    "# Low Ec\n",
    "\n",
    "p_var  = 0.20        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 0.        # starting value for Ec\n",
    "nb_Ag    = 5         # number of antigens\n",
    "conc     = 1.21      # antigen concentration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters for varying $E_0$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# High E0\n",
    "\n",
    "E0     = 6.          # mean binding energy for mixing with flexibility\n",
    "\n",
    "p_var  = 0.10        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 4.        # starting value for Ec\n",
    "nb_Ag    = 10        # number of antigens\n",
    "conc     = 1.11      # antigen concentration\n",
    "\n",
    "# Low E0\n",
    "\n",
    "E0     = 2.          # mean binding energy for mixing with flexibility\n",
    "\n",
    "p_var  = 0.10        # probability that a CDR mutation affects the variable region\n",
    "p_cons = 1.0 - p_var # probability that a CDR mutation affects the constant region\n",
    "\n",
    "Ec_start = 0.        # starting value for Ec\n",
    "nb_Ag    = 10        # number of antigens\n",
    "conc     = 1.23      # antigen concentration"
   ]
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}