{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.linalg import expm\n",
    "from scipy.integrate import quad\n",
    "from scipy.special import factorial\n",
    "from scipy.interpolate import PPoly\n",
    "from scipy.stats import sem\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib\n",
    "import seaborn as sns\n",
    "import msprime as msp\n",
    "import stdpopsim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "matplotlib.rcParams.update({\n",
    "    'font.family': 'Times New Roman',\n",
    "    'text.usetex': True,\n",
    "})\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "folder = '../../../exact_decoding_paper/figures/prior/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exact SMC Simulations\n",
    "### Simulation 1: Comparing renewal prior with Markov prior\n",
    "\n",
    "#### Define necessary functions for simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_expected_coal_times(times, eta):\n",
    "    r = eta[:-1]\n",
    "    a = times[:-1]\n",
    "    b = times[1:]\n",
    "    dt = b - a\n",
    "    # the exact formula is numerically unstable for small r, so use a taylor expansion.\n",
    "    r_small = (a + b) / 2. - dt ** 2 * r / 12\n",
    "    r_normal = b + 1. / r + dt / (np.expm1(-r * dt))\n",
    "    ary = np.where(r * dt < 1e-4, r_small, r_normal)\n",
    "    ret = np.concatenate([ary, times[None, -1] + 1 / eta[-1:]])\n",
    "    return ret\n",
    "\n",
    "def get_emissions(ts, w, L_w, mu, ect, w_end = None):\n",
    "    site_locations = [int(v.position)//w for v in ts.variants()]\n",
    "    loc, y = np.unique(site_locations, return_counts=True)\n",
    "    \n",
    "    if w_end is not None:\n",
    "        counts = np.zeros(L_w)\n",
    "        not_last = np.repeat(2 * mu * w, L_w - 1)\n",
    "        rate = np.r_[not_last, 2 * mu * w_end]\n",
    "        outer = np.outer(rate, ect)\n",
    "        \n",
    "    else:\n",
    "        counts = np.zeros(L_w)\n",
    "        rate = 2 * mu * w\n",
    "        outer = rate * ect[None,:]\n",
    "    \n",
    "    counts[loc] = y\n",
    "    \n",
    "    emissions = (outer ** counts[:,None]) * np.exp(-outer)/factorial(counts)[:,None]\n",
    "    return emissions\n",
    "\n",
    "def initial_distribution(times, eta):\n",
    "    x = np.r_[times, float('inf')]\n",
    "    y = PPoly(eta[None,:],x)\n",
    "    Y = y.antiderivative()\n",
    "    cdf_vals = 1 - np.exp(-np.r_[Y(times), float('inf')])\n",
    "    dist = cdf_vals[1:] - cdf_vals[:-1]\n",
    "    return dist\n",
    "\n",
    "def stationary_distribution(times, eta):\n",
    "    N0 = 1/eta[0]\n",
    "    vt = eta * N0\n",
    "    times_scaled = times / N0\n",
    "    x = np.r_[times_scaled, float('inf')]\n",
    "    y = PPoly(vt[None,:],x)\n",
    "    Y = y.antiderivative()\n",
    "    \n",
    "    def exp_int(tau):\n",
    "        return np.exp(-Y(tau))\n",
    "        \n",
    "    denominator = quad(exp_int, 0, float('inf'))[0]\n",
    "    def pi(t):\n",
    "        return (t * y(t) * exp_int(t))/denominator\n",
    "    \n",
    "    vals = np.r_[Y(times), float('inf')]\n",
    "    cdf_vals = np.r_[np.array([quad(pi, 0 ,t)[0] for t in times_scaled]), 1]\n",
    "    dist = cdf_vals[1:] - cdf_vals[:-1]\n",
    "    return dist\n",
    "\n",
    "def get_transitions_Markov(times, rho, eta):\n",
    "    x = np.r_[times, float('inf')]\n",
    "    y = PPoly(eta[None,:],x)\n",
    "    Y = y.antiderivative()\n",
    "    \n",
    "    A_rho = np.zeros((4,4))\n",
    "    A_rho[0] = [-1, 1, 0, 0]\n",
    "    A_rho[1] = [0, -0.5, 0.5, 0]\n",
    "    A_eta = np.zeros((4,4))\n",
    "    A_eta[1] = [1, -2, 0, 1]\n",
    "    A_eta[2] = [0, 4, -5, 1]\n",
    "    \n",
    "    def q(t, s):\n",
    "        if t > s:\n",
    "            e_j = expm(times[s]*rho*A_rho + Y(times[s])*A_eta)\n",
    "            return (e_j[0,1] + e_j[0,2]) * np.exp(-(Y(times[t-1])-Y(times[s]))) * (1-np.exp(-(Y(times[t])-Y(times[t-1]))))\n",
    "        \n",
    "        else:        \n",
    "            e_k1 = expm(times[t]*rho*A_rho + Y(times[t])*A_eta)\n",
    "            e_k2 = expm(times[t-1]*rho*A_rho + Y(times[t-1])*A_eta)\n",
    "\n",
    "            if t < s:\n",
    "                return e_k1[0,3] - e_k2[0,3]\n",
    "\n",
    "            else:\n",
    "                e_k3 = expm((times[t] - times[t-1])*rho*A_rho + (Y(times[t]) - Y(times[t-1]))*A_eta)\n",
    "\n",
    "                p = e_k1[0,0]\n",
    "                for i in range(3):\n",
    "                    p += e_k2[0,i] * e_k3[i,3]\n",
    "                return p\n",
    "        \n",
    "    transitions = np.empty((T,T))\n",
    "    for i in range(1,T+1):\n",
    "        for j in range(1,T+1):\n",
    "            if j < T:\n",
    "                transitions[i-1,j-1] = q(j,i)\n",
    "            else:\n",
    "                transitions[i-1,j-1] = np.max([1 - transitions[i-1,:j-1].sum(), 0])\n",
    "    return np.maximum(transitions, 0)\n",
    "\n",
    "def get_transitions_Renewal(times, ect, rho, eta):\n",
    "#     rows = stationary_distribution(times, eta)\n",
    "    rows = initial_distribution(times, eta)\n",
    "    cols = -np.expm1(-rho*ect)\n",
    "    diag = np.exp(-rho*ect)\n",
    "    transitions = np.outer(cols, rows) + np.diag(diag)\n",
    "    return transitions\n",
    "\n",
    "def forward(L_w, T, init_dist, emissions, transitions):\n",
    "    alpha = np.empty((L_w, T))\n",
    "    alpha[0] = init_dist * emissions[0]\n",
    "    c = np.empty(L_w)\n",
    "    c[0] = alpha[0].sum()\n",
    "    alpha[0] /= c[0]\n",
    "\n",
    "    for i in range(1, L_w):\n",
    "        alpha[i] = emissions[i] * (alpha[i-1][:,None] * transitions).sum(axis=0)\n",
    "        c[i] = alpha[i].sum()\n",
    "        alpha[i] /= c[i]\n",
    "    return alpha, c\n",
    "\n",
    "def backward(L_w, T, emissions, transitions, c):\n",
    "    beta = np.empty((L_w, T))\n",
    "    beta[-1] = np.ones(T)\n",
    "\n",
    "    for i in range(L_w-2, -1, -1):\n",
    "        beta[i] = (beta[i+1] * emissions[i+1,:] * transitions).sum(axis=1)\n",
    "        beta[i] /= c[i+1]\n",
    "    \n",
    "    return beta\n",
    "        \n",
    "def get_posterior(L_w, T, init_dist, emissions, transitions):\n",
    "    alpha, c = forward(L_w, T, init_dist, emissions, transitions)\n",
    "    beta = backward(L_w, T, emissions, transitions, c)\n",
    "    return np.maximum(alpha * beta,0)\n",
    "\n",
    "def get_viterbi(log_T, log_B, log_pi, L):\n",
    "    K = len(log_pi)\n",
    "    T1 = np.zeros([L, K])\n",
    "    T2 = np.zeros([L, K], dtype=int)\n",
    "    T1[0] = log_pi + log_B[0] \n",
    "    for i in range(1,L):\n",
    "        prob = T1[i-1][:,None] + log_T + log_B[i]\n",
    "        T1[i] = np.max(prob, axis=0)\n",
    "        T2[i] = np.argmax(prob, axis=0)\n",
    "    \n",
    "    Z = np.zeros(L, dtype=int)\n",
    "    Z[L-1] = np.argmax(T1[L-1])\n",
    "    \n",
    "    for i in range(L - 1, 0, -1):\n",
    "        Z[i - 1] = T2[i, Z[i]]\n",
    "    return Z\n",
    "\n",
    "def run_sim(num_sim, N0, mu, r, n, L, w, times, eta, rho, de=None):\n",
    "    L_w = L // w\n",
    "    T = times.size\n",
    "    \n",
    "    gammas_Markov = np.empty((num_sim, L_w, T))\n",
    "    gammas_Renewal = np.empty((num_sim, L_w, T))\n",
    "    viterbis_Markov = np.empty((num_sim, L_w), dtype=int)\n",
    "    viterbis_Renewal = np.empty((num_sim, L_w), dtype=int)\n",
    "    tss = []\n",
    "    \n",
    "    init_dist = initial_distribution(times, eta)\n",
    "    ect = get_expected_coal_times(times, eta)\n",
    "    transitions_Markov = get_transitions_Markov(times, rho, eta)\n",
    "    transitions_Renewal = get_transitions_Renewal(times, ect, rho, eta)\n",
    "    \n",
    "    for i in range(num_sim):\n",
    "        seed = i+1\n",
    "        ts = msp.simulate(Ne=N0, \n",
    "                          mutation_rate=mu, \n",
    "                          recombination_rate=r, \n",
    "                          sample_size=n, \n",
    "                          length=L,\n",
    "                          random_seed=seed,\n",
    "                          demographic_events=de)\n",
    "        tss.append(ts)\n",
    "\n",
    "        emissions = get_emissions(ts, w, L_w, mu, ect)\n",
    "\n",
    "\n",
    "        gammas_Markov[i] = get_posterior(L_w, T, init_dist, emissions, transitions_Markov)\n",
    "        gammas_Renewal[i] = get_posterior(L_w, T, init_dist, emissions, transitions_Renewal)\n",
    "        viterbis_Markov[i] = get_viterbi(np.log(transitions_Markov), np.log(emissions), np.log(init_dist), L_w)\n",
    "        viterbis_Renewal[i] = get_viterbi(np.log(transitions_Renewal), np.log(emissions), np.log(init_dist), L_w)\n",
    "        \n",
    "    return gammas_Markov, gammas_Renewal, viterbis_Markov, viterbis_Renewal, tss, ect\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Functions to plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_viterbi(viterbis_Markov, viterbis_renewal, tss, seed, ects, L, w, rs, name):\n",
    "    size = 25\n",
    "    L_w = L // w\n",
    "    scale_to_kb = 1000 / w\n",
    "    \n",
    "    fig, axs = plt.subplots(4, figsize=(20, 21), sharey=True)\n",
    "    fig.add_subplot(111, frameon=False)\n",
    "    plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)\n",
    "    \n",
    "    plot_labels = [r'Constant, Low $\\rho$', r'Expansion, Low $\\rho$', \n",
    "                   r'Constant, High $\\rho$', r'Expansion, High $\\rho$']\n",
    "    \n",
    "    for i, ax in enumerate(axs):\n",
    "        ts = tss[i][seed]\n",
    "        ect = ects[i]\n",
    "        \n",
    "        pmap = viterbis_Markov[i][seed]\n",
    "        dpmap, = np.diff(pmap).nonzero()\n",
    "        x = np.r_[0, dpmap, L_w - 1]\n",
    "        y = ect[pmap[x]]\n",
    "        ax.plot(x / scale_to_kb, y, drawstyle=\"steps-pre\", color=\"green\", label='Markov')\n",
    "\n",
    "        pmap = viterbis_Renewal[i][seed]\n",
    "        dpmap, = np.diff(pmap).nonzero()\n",
    "        x = np.r_[0, dpmap, L_w - 1]\n",
    "        y = ect[pmap[x]]\n",
    "        ax.plot(x / scale_to_kb, y, drawstyle=\"steps-pre\", color=\"blue\", label='Renewal')\n",
    "\n",
    "\n",
    "        truth = [(tree.interval[0] / w, tree.get_tmrca(0, 1)) for tree in ts.trees()]\n",
    "        x, y = zip(*truth)\n",
    "        x = np.r_[x, L_w]\n",
    "        y = np.r_[y, y[-1]]\n",
    "        ax.plot(x / scale_to_kb, y, drawstyle=\"steps-post\", color=\"red\", label='Truth')\n",
    "        \n",
    "        ax.set_yscale(\"log\")\n",
    "        ax.xaxis.set_tick_params(labelsize=size)\n",
    "        ax.yaxis.set_tick_params(labelsize=size)\n",
    "        ax.legend(loc = 'lower right', prop={'size': size})\n",
    "        \n",
    "        ax.text(0.01, .99, plot_labels[i], transform=ax.transAxes,\n",
    "                fontsize=size, va='top')\n",
    "        \n",
    "        if i < 2:\n",
    "            ax.set_xlim(0, L_w / scale_to_kb)\n",
    "        else:\n",
    "            ax.set_xlim(0, L_w / scale_to_kb * ratio)\n",
    "            \n",
    "        sns.despine()\n",
    "        \n",
    "        if i == 0:\n",
    "            ymin = np.amin(y)\n",
    "            ymax = np.amax(y)\n",
    "        \n",
    "        else:\n",
    "            ylow = np.amin(y)\n",
    "            if ylow < ymin:\n",
    "                ymin = ylow\n",
    "            \n",
    "            yhigh = np.amax(y)\n",
    "            if yhigh > ymax:\n",
    "                ymax = yhigh\n",
    "    \n",
    "    plt.setp(ax.get_xticklabels(), fontsize=size)   \n",
    "    plt.setp(ax.get_yticklabels(), fontsize=size)\n",
    "    plt.setp(ax, ylim=(0.75 * ymin, 1.25 * ymax))\n",
    "    \n",
    "    plt.xlabel('Position (kb)', fontsize=size, labelpad = size)\n",
    "    plt.ylabel('Coalescence time (Generations)', fontsize=size, labelpad = size)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(folder + name, bbox_inches='tight')\n",
    "    plt.show()\n",
    "    \n",
    "def plot_posterior(gammas_Markov, gammas_Renewal, tss, seed, ects, L, w, ratio, name):\n",
    "    size = 25\n",
    "    L_w = L // w\n",
    "    scale_to_kb = 1000 / w\n",
    "    \n",
    "    fig, axs = plt.subplots(8, figsize=(20, 28), sharey=True)\n",
    "    fig.add_subplot(111, frameon=False)\n",
    "    plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)\n",
    "    \n",
    "    plot_labels = [r'Constant, Low $\\rho$', r'Expansion, Low $\\rho$', \n",
    "                   r'Constant, High $\\rho$', r'Expansion, High $\\rho$']\n",
    "    \n",
    "    for i, ax in enumerate(axs):\n",
    "        if i % 2 == 0:\n",
    "            gamma = gammas_Markov[i//2][seed]\n",
    "            ax.text(0.01, .99, plot_labels[i//2], transform=ax.transAxes,\n",
    "                fontsize=size, va='top')\n",
    "            \n",
    "        else:\n",
    "            gamma = gammas_Renewal[i//2][seed]\n",
    "            \n",
    "        ts = tss[i//2][seed]\n",
    "        ect = ects[i//2]\n",
    "        \n",
    "        truth = [(tree.interval[0] / w, tree.get_tmrca(0, 1)) for tree in ts.trees()]\n",
    "        x, y = zip(*truth)\n",
    "        x = np.r_[x, L_w]\n",
    "        y = np.r_[y, y[-1]]\n",
    "        ax.plot(x / scale_to_kb, y, drawstyle=\"steps-post\", color=\"red\", label='Truth')\n",
    "        \n",
    "        X, Y = np.meshgrid(np.arange(L_w), ect)\n",
    "        ax.pcolormesh(X / scale_to_kb, Y, gamma.T, cmap=\"Greys\", rasterized=True)\n",
    "        \n",
    "        if i//2 < 2:\n",
    "            ax.set_xlim(0, L_w / scale_to_kb)\n",
    "        else:\n",
    "            ax.set_xlim(0, L_w / scale_to_kb * ratio)\n",
    "        \n",
    "        ax.set_yscale(\"log\")\n",
    "        ax.xaxis.set_tick_params(labelsize=size)\n",
    "        ax.yaxis.set_tick_params(labelsize=size)\n",
    "        \n",
    "        sns.despine()\n",
    "        \n",
    "        if i == 0:\n",
    "            ymin = np.amin(y)\n",
    "            ymax = np.amax(y)\n",
    "        \n",
    "        else:\n",
    "            ylow = np.amin(y)\n",
    "            if ylow < ymin:\n",
    "                ymin = ylow\n",
    "            \n",
    "            yhigh = np.amax(y)\n",
    "            if yhigh > ymax:\n",
    "                ymax = yhigh\n",
    "    \n",
    "    plt.setp(ax.get_xticklabels(), fontsize=size)\n",
    "    plt.setp(ax.get_yticklabels(), fontsize=size)\n",
    "    plt.setp(ax, ylim=(0.75 * ymin, 1.25 * ymax))\n",
    "    \n",
    "    plt.xlabel('Position (kb)', fontsize=size, labelpad = size)\n",
    "    plt.ylabel('Coalescence time (Generations)', fontsize=size, labelpad = size)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(folder + name, bbox_inches='tight')\n",
    "    plt.show()\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Functions to create tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_truth(tss, w, times):\n",
    "    T = times.size\n",
    "    num_sim = len(tss) \n",
    "    L = tss[0].sequence_length\n",
    "    L_w = int(L // w)\n",
    "    truth = np.empty((num_sim, L_w, T))\n",
    "    tmrca = np.empty((num_sim, L_w))\n",
    "    for i in range(len(tss)):\n",
    "        ts = tss[i]\n",
    "        cts = [(int(tree.interval[0] / w), \n",
    "                int(tree.interval[1] / w), \n",
    "                tree.get_tmrca(0,1)) for tree in ts.trees()]\n",
    "        x1, x2, y = zip(*cts)\n",
    "        for j in range(len(cts)):\n",
    "            left = x1[j]\n",
    "            right = x2[j]\n",
    "            t = y[j]\n",
    "            epoch = np.argmax(t < times) - 1\n",
    "            truth[i, left:right, epoch] = 1\n",
    "            tmrca[i, left:right] = t\n",
    "    return tmrca, truth\n",
    "\n",
    "def get_mean_se(gammas, dist, num_pos):\n",
    "    l1 = (gammas * dist).sum(axis=(1,2)) / num_pos\n",
    "    mean = np.average(l1)\n",
    "    se = sem(l1)\n",
    "    return mean, se\n",
    "\n",
    "def save_table(tab, name):\n",
    "    index_list = ['1','2','3','4']\n",
    "    tab['Scenario'] = index_list\n",
    "    tab = tab.set_index('Scenario').T\n",
    "    tab.to_latex(folder + name)\n",
    "    display(tab)\n",
    "    \n",
    "    \n",
    "def format_entry(mean, se, dec_places):\n",
    "    entry = f'{mean:.{dec_places}f} ({se:.{dec_places}f})'\n",
    "    return entry\n",
    "                \n",
    "mfunc = np.vectorize(format_entry)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create containers for all 4 scenarios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "gammas_Markov = [None] * 4\n",
    "gammas_Renewal = [None] * 4\n",
    "viterbis_Markov = [None] * 4\n",
    "viterbis_Renewal = [None] * 4\n",
    "tss = [None] * 4\n",
    "\n",
    "#These parameters are the same for each scenario\n",
    "num_sim = 25 #number of simulation runs\n",
    "N0 = 1e4 #base effective population size\n",
    "mu = 1.4e-8 #mutation rate\n",
    "L = 5_000_000 #sequence length\n",
    "n = 2 #sample size\n",
    "\n",
    "w = 100 #window size\n",
    "T = 32\n",
    "times = np.r_[0.0, np.geomspace(1e1, 1e6, T-1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 1\n",
    "* Constant population (N=1e4)\n",
    "* mutation rate = 1.4e-8\n",
    "* recombination rate = 1e-9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:186: RuntimeWarning: divide by zero encountered in log\n",
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:187: RuntimeWarning: divide by zero encountered in log\n"
     ]
    }
   ],
   "source": [
    "r = 1e-9\n",
    "eta = np.repeat(1/(2*N0), T)\n",
    "rho = 2*r*w\n",
    "\n",
    "ret = run_sim(num_sim, N0, mu, r, n, L, w, times, eta, rho)\n",
    "gammas_Markov[0], gammas_Renewal[0], viterbis_Markov[0], viterbis_Renewal[0], tss[0], ect_constant = ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 2\n",
    "* Population Expansion in the recent past (N=1e6 at time=0, N=5e-5 at time 1000, N=1e-4 at time 3000)\n",
    "* mutation rate = 1.4e-8\n",
    "* recombination rate = 1e-9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:186: RuntimeWarning: divide by zero encountered in log\n",
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:187: RuntimeWarning: divide by zero encountered in log\n"
     ]
    }
   ],
   "source": [
    "r = 1e-9\n",
    "eta = 1/np.r_[np.repeat(200 * N0, 13), np.repeat(N0, 3), np.repeat(2 * N0, T-16)]\n",
    "rho = 2*r*w\n",
    "\n",
    "#Add demographic event for variable population size\n",
    "de = [msp.PopulationParametersChange(time=0, initial_size=100 * N0),\n",
    "        msp.PopulationParametersChange(time=times[13], initial_size=0.5 * N0),\n",
    "        msp.PopulationParametersChange(time=times[16], initial_size=N0)]\n",
    "\n",
    "ret = run_sim(num_sim, N0, mu, r, n, L, w, times, eta, rho, de)\n",
    "gammas_Markov[1], gammas_Renewal[1], viterbis_Markov[1], viterbis_Renewal[1], tss[1], ect_variable = ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 3\n",
    "* Constant population (N=1e4)\n",
    "* mutation rate = 1.4e-8\n",
    "* recombination rate = 1.4e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:186: RuntimeWarning: divide by zero encountered in log\n",
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:187: RuntimeWarning: divide by zero encountered in log\n"
     ]
    }
   ],
   "source": [
    "r = 1.4e-8\n",
    "eta = np.repeat(1/(2*N0), T)\n",
    "rho = 2*r*w\n",
    "\n",
    "ret = run_sim(num_sim, N0, mu, r, n, L, w, times, eta, rho)\n",
    "gammas_Markov[2], gammas_Renewal[2], viterbis_Markov[2], viterbis_Renewal[2], tss[2], _ = ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 4\n",
    "* Population Expansion in the recent past (N=1e6 at time=0, N=5e-5 at time 1000, N=1e-4 at time 3000)\n",
    "* mutation rate = 1.4e-8\n",
    "* recombination rate = 1.4e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:186: RuntimeWarning: divide by zero encountered in log\n",
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:187: RuntimeWarning: divide by zero encountered in log\n"
     ]
    }
   ],
   "source": [
    "r = 1.4e-8\n",
    "eta = 1/np.r_[np.repeat(200 * N0, 13), np.repeat(N0, 3), np.repeat(2 * N0, T-16)]\n",
    "rho = 2*r*w\n",
    "\n",
    "#Add demographic event for variable population size\n",
    "de = [msp.PopulationParametersChange(time=0, initial_size=100 * N0),\n",
    "        msp.PopulationParametersChange(time=times[13], initial_size=0.5 * N0),\n",
    "        msp.PopulationParametersChange(time=times[16], initial_size=N0)]\n",
    "\n",
    "ret = run_sim(num_sim, N0, mu, r, n, L, w, times, eta, rho, de)\n",
    "gammas_Markov[3], gammas_Renewal[3], viterbis_Markov[3], viterbis_Renewal[3], tss[3], _ = ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot Viterbi and posterior for one run under each scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 103736 characters]\n",
      "text/plain": [
       "<Figure size 1440x1512 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_seed = 6\n",
    "ects= [ect_constant, ect_variable] * 2\n",
    "ratio = 1e-9/1.4e-8\n",
    "plot_viterbi(viterbis_Markov, viterbis_Renewal, tss, plot_seed, ects, L, w, ratio, 'viterbi1.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 304820 characters]\n",
      "text/plain": [
       "<Figure size 1440x2016 with 9 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_posterior(gammas_Markov, gammas_Renewal, tss, plot_seed, ects, L, w, ratio, 'posterior1.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['Markov', 'Renewal']\n",
    "df_abs = pd.DataFrame(columns=columns)\n",
    "df_log = pd.DataFrame(columns=columns)\n",
    "\n",
    "for i in range(4):\n",
    "    tmrca, truth = get_truth(tss[i], w, times)\n",
    "\n",
    "    if i%2==0:\n",
    "        dist_abs = np.abs(ect_constant - tmrca[:,:,None])\n",
    "        dist_log = np.abs(np.log10(ect_constant) - np.log10(tmrca[:,:,None]))\n",
    "    else:\n",
    "        dist_abs = np.abs(ect_variable - tmrca[:,:,None])\n",
    "        dist_log = np.abs(np.log10(ect_variable) - np.log10(tmrca[:,:,None]))\n",
    "\n",
    "    mean1, se1 = get_mean_se(gammas_Markov[i], dist_abs, L // w)\n",
    "    mean2, se2 = get_mean_se(gammas_Renewal[i], dist_abs, L // w)\n",
    "    mean3, se3 = get_mean_se(gammas_Markov[i], dist_log, L // w)\n",
    "    mean4, se4 = get_mean_se(gammas_Renewal[i], dist_log, L // w)\n",
    "\n",
    "    df_abs.loc[i] = [f'{mean1:.2f} ({se1:.2f})', f'{mean2:.2f} ({se2:.2f})']\n",
    "    df_log.loc[i] = [f'{mean3:.4f} ({se3:.4f})', f'{mean4:.4f} ({se4:.4f})']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Scenario</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <td>6075.50 (215.59)</td>\n",
       "      <td>5187.24 (187.27)</td>\n",
       "      <td>12037.25 (298.71)</td>\n",
       "      <td>12112.81 (106.44)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <td>6068.04 (209.11)</td>\n",
       "      <td>5187.49 (184.34)</td>\n",
       "      <td>11476.92 (282.36)</td>\n",
       "      <td>11571.28 (97.99)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Scenario                 1                 2                  3  \\\n",
       "Markov    6075.50 (215.59)  5187.24 (187.27)  12037.25 (298.71)   \n",
       "Renewal   6068.04 (209.11)  5187.49 (184.34)  11476.92 (282.36)   \n",
       "\n",
       "Scenario                  4  \n",
       "Markov    12112.81 (106.44)  \n",
       "Renewal    11571.28 (97.99)  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "save_table(df_abs, 'table1_1.tex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Scenario</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <td>0.1362 (0.0035)</td>\n",
       "      <td>0.1281 (0.0024)</td>\n",
       "      <td>0.3051 (0.0013)</td>\n",
       "      <td>0.3006 (0.0014)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <td>0.1437 (0.0026)</td>\n",
       "      <td>0.1359 (0.0025)</td>\n",
       "      <td>0.3451 (0.0046)</td>\n",
       "      <td>0.3422 (0.0014)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Scenario                1                2                3                4\n",
       "Markov    0.1362 (0.0035)  0.1281 (0.0024)  0.3051 (0.0013)  0.3006 (0.0014)\n",
       "Renewal   0.1437 (0.0026)  0.1359 (0.0025)  0.3451 (0.0046)  0.3422 (0.0014)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "save_table(df_log, 'table1_2.tex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = pd.MultiIndex.from_tuples([('Markov', 'Q1'), ('Renewal', 'Q1'), \n",
    "                                     ('Markov', 'Q2'), ('Renewal', 'Q2'),\n",
    "                                     ('Markov', 'Q3'), ('Renewal', 'Q3'),\n",
    "                                     ('Markov', 'Q4'), ('Renewal', 'Q4')])\n",
    "df_abs = pd.DataFrame(columns=columns)\n",
    "df_log = pd.DataFrame(columns=columns)\n",
    "\n",
    "count = np.empty((4, 4, num_sim))\n",
    "\n",
    "def get_row(l1, dec_places):\n",
    "    mean = np.average(l1, axis=0)\n",
    "    se = sem(l1, axis=0)\n",
    "    row = [f'{m:.{dec_places}f} ({s:.{dec_places}f})' for m, s in zip(mean, se)]\n",
    "    return row                       \n",
    "                                  \n",
    "\n",
    "for i in range(4):\n",
    "    df_abs.loc[i] = [None] * 8\n",
    "    df_log.loc[i] = [None] * 8\n",
    "    tmrca, truth = get_truth(tss[i], w, times)\n",
    "    l11 = np.empty((num_sim, 4))\n",
    "    l12 = np.empty((num_sim, 4))\n",
    "    l13 = np.empty((num_sim, 4))\n",
    "    l14 = np.empty((num_sim, 4))\n",
    "    \n",
    "    if i%2==0:\n",
    "        dist_abs = np.abs(ect_constant - tmrca[:,:,None])\n",
    "        dist_log = np.abs(np.log10(ect_constant) - np.log10(tmrca[:,:,None]))\n",
    "    else:\n",
    "        dist_abs = np.abs(ect_variable - tmrca[:,:,None])\n",
    "        dist_log = np.abs(np.log10(ect_variable) - np.log10(tmrca[:,:,None]))\n",
    "            \n",
    "    \n",
    "    for j in range(num_sim):\n",
    "        full_Markov_abs = gammas_Markov[i][j] * dist_abs[j]\n",
    "        full_Renewal_abs = gammas_Renewal[i][j] * dist_abs[j]\n",
    "        full_Markov_log = gammas_Markov[i][j] * dist_log[j]\n",
    "        full_Renewal_log = gammas_Renewal[i][j] * dist_log[j]\n",
    "\n",
    "        uniq_t = np.unique(tmrca)\n",
    "        breaks = np.quantile(uniq_t, [0, 0.25, 0.5, 0.75])\n",
    "        \n",
    "        for k in range(4):\n",
    "            if k != 3:\n",
    "                inds_k = (tmrca[j] >= breaks[k]) & (tmrca[j] < breaks[k+1])\n",
    "            else:\n",
    "                inds_k = (tmrca[j] >= breaks[3])\n",
    "                \n",
    "            num_pos = inds_k.sum()\n",
    "            count[i,k,j] = num_pos\n",
    "            l11[j,k] = full_Markov_abs[inds_k].sum() / num_pos\n",
    "            l12[j,k] = full_Renewal_abs[inds_k].sum() / num_pos\n",
    "            l13[j,k] = full_Markov_log[inds_k].sum() / num_pos\n",
    "            l14[j,k] = full_Renewal_log[inds_k].sum() / num_pos\n",
    "\n",
    "    df_abs.loc[i][::2] = get_row(l11, 2)\n",
    "    df_abs.loc[i][1::2] = get_row(l12, 2)\n",
    "    df_log.loc[i][::2] = get_row(l13, 4)\n",
    "    df_log.loc[i][1::2] = get_row(l14, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Scenario</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Q1</th>\n",
       "      <td>25878.20 (951.37)</td>\n",
       "      <td>28103.76 (1075.15)</td>\n",
       "      <td>27807.96 (580.26)</td>\n",
       "      <td>27659.64 (233.24)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Q2</th>\n",
       "      <td>12037.56 (646.81)</td>\n",
       "      <td>10589.40 (702.38)</td>\n",
       "      <td>10931.92 (331.94)</td>\n",
       "      <td>11149.32 (143.83)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Q3</th>\n",
       "      <td>7917.24 (461.61)</td>\n",
       "      <td>7231.24 (509.98)</td>\n",
       "      <td>7200.76 (178.59)</td>\n",
       "      <td>7119.12 (95.21)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Q4</th>\n",
       "      <td>4167.00 (360.71)</td>\n",
       "      <td>4075.60 (272.09)</td>\n",
       "      <td>4059.36 (124.89)</td>\n",
       "      <td>4071.92 (90.72)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Scenario                  1                   2                  3  \\\n",
       "Q1        25878.20 (951.37)  28103.76 (1075.15)  27807.96 (580.26)   \n",
       "Q2        12037.56 (646.81)   10589.40 (702.38)  10931.92 (331.94)   \n",
       "Q3         7917.24 (461.61)    7231.24 (509.98)   7200.76 (178.59)   \n",
       "Q4         4167.00 (360.71)    4075.60 (272.09)   4059.36 (124.89)   \n",
       "\n",
       "Scenario                  4  \n",
       "Q1        27659.64 (233.24)  \n",
       "Q2        11149.32 (143.83)  \n",
       "Q3          7119.12 (95.21)  \n",
       "Q4          4071.92 (90.72)  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mean = np.average(count, axis=2)\n",
    "se = sem(count, axis=2)\n",
    "entries = mfunc(mean, se, 2)\n",
    "df = pd.DataFrame(data=entries, columns=['Q1', 'Q2', 'Q3', 'Q4'])\n",
    "df['Scenario'] = ['1', '2', '3', '4']\n",
    "df = df.set_index('Scenario').T\n",
    "df.to_latex(folder + 'locicounts.tex')\n",
    "display(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Scenario</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q1</th>\n",
       "      <td>2840.79 (124.42)</td>\n",
       "      <td>2219.61 (101.38)</td>\n",
       "      <td>6820.01 (245.65)</td>\n",
       "      <td>6672.23 (88.84)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q1</th>\n",
       "      <td>2887.14 (129.53)</td>\n",
       "      <td>2273.91 (100.00)</td>\n",
       "      <td>5292.75 (180.86)</td>\n",
       "      <td>5256.89 (55.12)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q2</th>\n",
       "      <td>6524.06 (131.61)</td>\n",
       "      <td>6106.74 (122.89)</td>\n",
       "      <td>13346.56 (35.50)</td>\n",
       "      <td>13465.23 (57.68)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q2</th>\n",
       "      <td>6604.09 (117.57)</td>\n",
       "      <td>6123.72 (107.90)</td>\n",
       "      <td>11463.27 (38.05)</td>\n",
       "      <td>11665.01 (35.60)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q3</th>\n",
       "      <td>10543.82 (239.47)</td>\n",
       "      <td>9716.04 (219.67)</td>\n",
       "      <td>18831.10 (64.05)</td>\n",
       "      <td>18957.97 (67.16)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q3</th>\n",
       "      <td>10454.31 (249.99)</td>\n",
       "      <td>9528.27 (225.84)</td>\n",
       "      <td>19538.07 (72.65)</td>\n",
       "      <td>19682.76 (71.76)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q4</th>\n",
       "      <td>17242.02 (460.92)</td>\n",
       "      <td>16362.35 (471.86)</td>\n",
       "      <td>33016.17 (212.12)</td>\n",
       "      <td>33526.33 (226.10)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q4</th>\n",
       "      <td>16959.09 (505.93)</td>\n",
       "      <td>16341.50 (601.40)</td>\n",
       "      <td>40181.58 (263.39)</td>\n",
       "      <td>40141.12 (279.70)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Scenario                    1                  2                  3  \\\n",
       "Markov  Q1   2840.79 (124.42)   2219.61 (101.38)   6820.01 (245.65)   \n",
       "Renewal Q1   2887.14 (129.53)   2273.91 (100.00)   5292.75 (180.86)   \n",
       "Markov  Q2   6524.06 (131.61)   6106.74 (122.89)   13346.56 (35.50)   \n",
       "Renewal Q2   6604.09 (117.57)   6123.72 (107.90)   11463.27 (38.05)   \n",
       "Markov  Q3  10543.82 (239.47)   9716.04 (219.67)   18831.10 (64.05)   \n",
       "Renewal Q3  10454.31 (249.99)   9528.27 (225.84)   19538.07 (72.65)   \n",
       "Markov  Q4  17242.02 (460.92)  16362.35 (471.86)  33016.17 (212.12)   \n",
       "Renewal Q4  16959.09 (505.93)  16341.50 (601.40)  40181.58 (263.39)   \n",
       "\n",
       "Scenario                    4  \n",
       "Markov  Q1    6672.23 (88.84)  \n",
       "Renewal Q1    5256.89 (55.12)  \n",
       "Markov  Q2   13465.23 (57.68)  \n",
       "Renewal Q2   11665.01 (35.60)  \n",
       "Markov  Q3   18957.97 (67.16)  \n",
       "Renewal Q3   19682.76 (71.76)  \n",
       "Markov  Q4  33526.33 (226.10)  \n",
       "Renewal Q4  40141.12 (279.70)  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "save_table(df_abs, 'table1_3.tex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Scenario</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q1</th>\n",
       "      <td>0.1479 (0.0056)</td>\n",
       "      <td>0.1328 (0.0033)</td>\n",
       "      <td>0.3212 (0.0020)</td>\n",
       "      <td>0.3080 (0.0022)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q1</th>\n",
       "      <td>0.1563 (0.0045)</td>\n",
       "      <td>0.1413 (0.0034)</td>\n",
       "      <td>0.3299 (0.0057)</td>\n",
       "      <td>0.3119 (0.0018)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q2</th>\n",
       "      <td>0.1279 (0.0023)</td>\n",
       "      <td>0.1253 (0.0023)</td>\n",
       "      <td>0.2755 (0.0016)</td>\n",
       "      <td>0.2863 (0.0016)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q2</th>\n",
       "      <td>0.1373 (0.0024)</td>\n",
       "      <td>0.1336 (0.0024)</td>\n",
       "      <td>0.3277 (0.0019)</td>\n",
       "      <td>0.3480 (0.0022)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q3</th>\n",
       "      <td>0.1255 (0.0029)</td>\n",
       "      <td>0.1242 (0.0036)</td>\n",
       "      <td>0.2826 (0.0020)</td>\n",
       "      <td>0.2857 (0.0022)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q3</th>\n",
       "      <td>0.1324 (0.0033)</td>\n",
       "      <td>0.1293 (0.0039)</td>\n",
       "      <td>0.3769 (0.0024)</td>\n",
       "      <td>0.3857 (0.0024)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Markov</th>\n",
       "      <th>Q4</th>\n",
       "      <td>0.1197 (0.0036)</td>\n",
       "      <td>0.1216 (0.0032)</td>\n",
       "      <td>0.3168 (0.0026)</td>\n",
       "      <td>0.3178 (0.0021)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Renewal</th>\n",
       "      <th>Q4</th>\n",
       "      <td>0.1247 (0.0040)</td>\n",
       "      <td>0.1278 (0.0036)</td>\n",
       "      <td>0.4546 (0.0029)</td>\n",
       "      <td>0.4568 (0.0023)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Scenario                  1                2                3                4\n",
       "Markov  Q1  0.1479 (0.0056)  0.1328 (0.0033)  0.3212 (0.0020)  0.3080 (0.0022)\n",
       "Renewal Q1  0.1563 (0.0045)  0.1413 (0.0034)  0.3299 (0.0057)  0.3119 (0.0018)\n",
       "Markov  Q2  0.1279 (0.0023)  0.1253 (0.0023)  0.2755 (0.0016)  0.2863 (0.0016)\n",
       "Renewal Q2  0.1373 (0.0024)  0.1336 (0.0024)  0.3277 (0.0019)  0.3480 (0.0022)\n",
       "Markov  Q3  0.1255 (0.0029)  0.1242 (0.0036)  0.2826 (0.0020)  0.2857 (0.0022)\n",
       "Renewal Q3  0.1324 (0.0033)  0.1293 (0.0039)  0.3769 (0.0024)  0.3857 (0.0024)\n",
       "Markov  Q4  0.1197 (0.0036)  0.1216 (0.0032)  0.3168 (0.0026)  0.3178 (0.0021)\n",
       "Renewal Q4  0.1247 (0.0040)  0.1278 (0.0036)  0.4546 (0.0029)  0.4568 (0.0023)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "save_table(df_log, 'table1_4.tex')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulation 2: Effect of Demographic Prior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Redefine some functions for second simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_viterbi(viterbis, tss, ects, L_w, w, name):\n",
    "    size = 25\n",
    "    scale_to_kb = 1000 / w\n",
    "    width = 3600\n",
    "    \n",
    "    fig, axs = plt.subplots(3, figsize=(20, 13.5), sharex=True, sharey=True)\n",
    "    fig.add_subplot(111, frameon=False)\n",
    "    plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)\n",
    "    \n",
    "    labels = ['Africa', 'Zigzag', 'Constant']    \n",
    "    colors = ['purple', 'green', 'orange']\n",
    "    \n",
    "    for i, ax in enumerate(axs):\n",
    "        ts = tss[i]\n",
    "        for j in range(3):\n",
    "            ect = ects[j]\n",
    "        \n",
    "            pmap = viterbis[i,j]\n",
    "            dpmap, = np.diff(pmap).nonzero()\n",
    "            x = np.r_[0, dpmap, L_w-1]\n",
    "            y = ect[pmap[x]]\n",
    "            ax.plot(x / scale_to_kb, y, drawstyle=\"steps-pre\", color=colors[j], label=labels[j])\n",
    "\n",
    "        truth = [(tree.interval[0] / w, tree.get_tmrca(0, 1)) for tree in ts.trees()]\n",
    "        x, y = zip(*truth)\n",
    "        x = np.r_[x, L_w]\n",
    "        y = np.r_[y, y[-1]]\n",
    "        ax.plot(x / scale_to_kb, y, drawstyle=\"steps-post\", color=\"red\", label='Truth')\n",
    "        \n",
    "        ax.set_yscale(\"log\")\n",
    "        ax.xaxis.set_tick_params(labelsize=size)\n",
    "        ax.yaxis.set_tick_params(labelsize=size)\n",
    "        ax.legend(loc = 'lower right', prop={'size': size})\n",
    "        \n",
    "        ax.text(0.01, .99, labels[i], transform=ax.transAxes,\n",
    "                fontsize=size, va='top')\n",
    "            \n",
    "        sns.despine()\n",
    "        \n",
    "        if i == 0:\n",
    "            ymin = np.amin(y[:width])\n",
    "            ymax = np.amax(y[:width])\n",
    "        \n",
    "        else:\n",
    "            ylow = np.amin(y[:width])\n",
    "            if ylow < ymin:\n",
    "                ymin = ylow\n",
    "            \n",
    "            yhigh = np.amax(y[:width])\n",
    "            if yhigh > ymax:\n",
    "                ymax = yhigh\n",
    "    \n",
    "    plt.setp(ax.get_xticklabels(), fontsize=size)   \n",
    "    plt.setp(ax.get_yticklabels(), fontsize=size)\n",
    "    plt.setp(ax, xlim=(0, width/scale_to_kb))\n",
    "    plt.setp(ax, ylim=(0.75 * ymin, 1.25 * ymax))\n",
    "    \n",
    "    plt.xlabel('Position (kb)', fontsize=size, labelpad = size)\n",
    "    plt.ylabel('Coalescence time (Generations)', fontsize=size, labelpad = size)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(folder + name, bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_posterior(gammas, tss, ects, L_w, w, name):\n",
    "    size = 25\n",
    "    scale_to_kb = 1000 / w\n",
    "    width = 3600\n",
    "    \n",
    "    fig, axs = plt.subplots(9, figsize=(20, 27), sharex=True, sharey=True)\n",
    "    fig.add_subplot(111, frameon=False)\n",
    "    plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)\n",
    "    \n",
    "    labels = ['Africa', 'Zigzag', 'Constant'] \n",
    "    \n",
    "    for i, ax in enumerate(axs):\n",
    "        ind1 = i//3\n",
    "        ts = tss[ind1]\n",
    "        \n",
    "        ind2 = i % 3\n",
    "        if ind2 == 0:\n",
    "            ax.text(0.01, .99, labels[ind1], transform=ax.transAxes,\n",
    "                fontsize=size, va='top')\n",
    "\n",
    "        ect = ects[ind2]\n",
    "        gamma = gammas[ind1,ind2,:width]\n",
    "        \n",
    "        truth = [(tree.interval[0] / w, tree.get_tmrca(0, 1)) for tree in ts.trees()]\n",
    "        x, y = zip(*truth)\n",
    "        x = np.r_[x, L_w]\n",
    "        y = np.r_[y, y[-1]]\n",
    "        ax.plot(x / scale_to_kb, y, drawstyle=\"steps-post\", color=\"red\", label='Truth')\n",
    "        \n",
    "        X, Y = np.meshgrid(np.arange(width), ect)\n",
    "        im = ax.pcolormesh(X / scale_to_kb, Y, gamma.T, cmap=\"Greys\", rasterized=True)\n",
    "#         if unique_clim:\n",
    "#             im.set_clim(np.amin(gamma), np.amax(gamma))\n",
    "        \n",
    "        ax.set_yscale(\"log\")\n",
    "        ax.xaxis.set_tick_params(labelsize=size)\n",
    "        ax.yaxis.set_tick_params(labelsize=size)\n",
    "        \n",
    "        sns.despine()\n",
    "        \n",
    "        if i == 0:\n",
    "            ymin = np.amin(y)\n",
    "            ymax = np.amax(y)\n",
    "        \n",
    "        else:\n",
    "            ylow = np.amin(y)\n",
    "            if ylow < ymin:\n",
    "                ymin = ylow\n",
    "            \n",
    "            yhigh = np.amax(y)\n",
    "            if yhigh > ymax:\n",
    "                ymax = yhigh\n",
    "    \n",
    "    plt.setp(ax.get_xticklabels(), fontsize=size)\n",
    "    plt.setp(ax.get_yticklabels(), fontsize=size)\n",
    "    plt.setp(ax, ylim=(0.75 * ymin, 1.25 * ymax))\n",
    "    plt.setp(ax.set_xlim(0, width/scale_to_kb))\n",
    "    \n",
    "    plt.xlabel('Position (kb)', fontsize=size, labelpad = size)\n",
    "    plt.ylabel('Coalescence time (Generations)', fontsize=size, labelpad = size)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(folder + name, bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_truth(ts, w, times):\n",
    "    T = times.size\n",
    "    L = tss[0].sequence_length\n",
    "    L_w = int(np.ceil(L / w))\n",
    "    tmrca = np.empty(L_w)\n",
    "    \n",
    "    cts = [(round(tree.interval[0] / w), \n",
    "            round(tree.interval[1] / w), \n",
    "            tree.get_tmrca(0,1)) for tree in ts.trees()]\n",
    "    x1, x2, y = zip(*cts)\n",
    "    for j in range(len(cts)):\n",
    "        left = x1[j]\n",
    "        right = x2[j]\n",
    "        t = y[j]\n",
    "        tmrca[left:right] = t\n",
    "    return tmrca, truth\n",
    "\n",
    "def create_table(err, name, log=False):\n",
    "    columns = ['Africa', 'Zigzag', 'Constant']\n",
    "    mean = np.average(err, axis=2)\n",
    "    se = sem(err, axis=2)\n",
    "    if log:\n",
    "        dec_places = 4\n",
    "    else:\n",
    "        dec_places = 2\n",
    "    entries = mfunc(mean, se, dec_places)\n",
    "    df = pd.DataFrame(data=entries, index=columns, columns=columns)\n",
    "    df.to_latex(folder + name)\n",
    "    display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set simulation and model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sim = 25\n",
    "w = 100\n",
    "\n",
    "species = stdpopsim.get_species(\"HomSap\")\n",
    "chrom = 22\n",
    "contig = species.get_contig(f'chr{chrom}')\n",
    "r = contig.recombination_map.get_rates()[0]\n",
    "mu = contig.mutation_rate\n",
    "\n",
    "T = 64\n",
    "times = np.r_[0.0, np.geomspace(1e1, 1e6, T-1)]\n",
    "\n",
    "plot_seed = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "africa = species.get_demographic_model(\"Africa_1T12\")\n",
    "zigzag = species.get_demographic_model(\"Zigzag_1S14\")\n",
    "constant = stdpopsim.PiecewiseConstantSize(species.population_size)\n",
    "models = [africa, zigzag, constant]\n",
    "\n",
    "demo_africa = africa.get_demography_debugger()\n",
    "demo_zigzag = zigzag.get_demography_debugger()\n",
    "\n",
    "eta_africa = np.squeeze(1/(2*demo_africa.population_size_trajectory(steps=times)))\n",
    "eta_zigzag = np.squeeze(1/(2*demo_zigzag.population_size_trajectory(steps=times)))\n",
    "eta_constant = np.repeat(1/(2*species.population_size), T)\n",
    "etas = [eta_africa, eta_zigzag, eta_constant]\n",
    "\n",
    "L = species.genome.chromosomes[chrom-1].length\n",
    "L_w = int(np.ceil(L / w))\n",
    "\n",
    "w_end = None\n",
    "if L % w > 0:\n",
    "    w_end = L % w\n",
    "rho = 2*r*w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Population size trajectories "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31272 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(10,5))\n",
    "\n",
    "x = np.geomspace(1, 1e6, 1000)\n",
    "y = np.squeeze(2*demo_africa.population_size_trajectory(steps=x))\n",
    "plt.plot(x, y, color='red', label='Africa')\n",
    "\n",
    "y = np.squeeze(2*demo_africa.population_size_trajectory(steps=times))\n",
    "plt.plot(times, y, color=\"blue\", drawstyle='steps-post', label='Africa Piecewise')\n",
    "\n",
    "x = np.geomspace(1, 1e6, 1000)\n",
    "y = np.squeeze(2*demo_zigzag.population_size_trajectory(steps=x))\n",
    "plt.plot(x, y, color='orange', label='Zigzag')\n",
    "\n",
    "y = np.squeeze(2*demo_zigzag.population_size_trajectory(steps=times))\n",
    "plt.plot(times, y, color=\"green\", drawstyle='steps-post', label='Zigzag Piecewise')\n",
    "\n",
    "plt.plot([0,1e6], [2*species.population_size]*2, color='purple', label='Constant')\n",
    "\n",
    "size = 15\n",
    "plt.xticks(fontsize=size)\n",
    "plt.yticks(fontsize=size)\n",
    "\n",
    "plt.xlabel('Time (Generations)', fontsize=size)\n",
    "plt.ylabel('Effective Population Size', fontsize=size)\n",
    "sns.despine()\n",
    "\n",
    "plt.legend(fontsize=size)\n",
    "plt.yscale('log')\n",
    "plt.xscale('log')\n",
    "plt.savefig(folder + 'sizehistory.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/caleb/Desktop/Git/HMM-EM/venv/lib/python3.7/site-packages/ipykernel_launcher.py:40: RuntimeWarning: divide by zero encountered in log\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1\n",
      "0 2\n",
      "0 3\n",
      "0 4\n",
      "0 5\n",
      "0 6\n",
      "0 7\n",
      "0 8\n",
      "0 9\n",
      "0 10\n",
      "0 11\n",
      "0 12\n",
      "0 13\n",
      "0 14\n",
      "0 15\n",
      "0 16\n",
      "0 17\n",
      "0 18\n",
      "0 19\n",
      "0 20\n",
      "0 21\n",
      "0 22\n",
      "0 23\n",
      "0 24\n",
      "1 0\n",
      "1 1\n",
      "1 2\n",
      "1 3\n",
      "1 4\n",
      "1 5\n",
      "1 6\n",
      "1 7\n",
      "1 8\n",
      "1 9\n",
      "1 10\n",
      "1 11\n",
      "1 12\n",
      "1 13\n",
      "1 14\n",
      "1 15\n",
      "1 16\n",
      "1 17\n",
      "1 18\n",
      "1 19\n",
      "1 20\n",
      "1 21\n",
      "1 22\n",
      "1 23\n",
      "1 24\n",
      "2 0\n",
      "2 1\n",
      "2 2\n",
      "2 3\n",
      "2 4\n",
      "2 5\n",
      "2 6\n",
      "2 7\n",
      "2 8\n",
      "2 9\n",
      "2 10\n",
      "2 11\n",
      "2 12\n",
      "2 13\n",
      "2 14\n",
      "2 15\n",
      "2 16\n",
      "2 17\n",
      "2 18\n",
      "2 19\n",
      "2 20\n",
      "2 21\n",
      "2 22\n",
      "2 23\n",
      "2 24\n"
     ]
    }
   ],
   "source": [
    "ids = [None] * 3\n",
    "ects = [None] * 3\n",
    "transitions = [None] * 3\n",
    "\n",
    "gammas = np.empty((3,3,L_w,T))\n",
    "viterbis = np.empty((3,3,L_w), dtype=int)\n",
    "tss = []\n",
    "\n",
    "engine = stdpopsim.get_engine('msprime')\n",
    "\n",
    "err_A = np.empty((3, 3, num_sim))\n",
    "err_B = np.empty((3, 3, num_sim))\n",
    "\n",
    "for j in range(3):\n",
    "    ids[j] = initial_distribution(times, etas[j])\n",
    "    ects[j] = get_expected_coal_times(times, etas[j])\n",
    "    transitions[j] = get_transitions_Renewal(times, ects[j], rho, etas[j])\n",
    "\n",
    "for i in range(3):\n",
    "    model = models[i]\n",
    "    samples = model.get_samples(2)\n",
    "    for k in range(num_sim):\n",
    "        print(i, k)\n",
    "        ts = engine.simulate(model, contig, samples, seed=k+1)\n",
    "        if k+1 == plot_seed:\n",
    "            tss.append(ts)\n",
    "        \n",
    "        tmrca, truth = get_truth(ts, w, times)           \n",
    "\n",
    "        for j in range(3):\n",
    "            emissions = get_emissions(ts, w, L_w, mu, ects[j], w_end)\n",
    "            gamma = get_posterior(L_w, T, ids[j], emissions, transitions[j])\n",
    "            dist_abs = np.abs(ects[j] - tmrca[:,None])\n",
    "            dist_log = np.abs(np.log10(ects[j]) - np.log10(tmrca[:,None]))\n",
    "            err_A[i,j,k] = (gamma * dist_abs).sum()/L_w\n",
    "            err_B[i,j,k] = (gamma * dist_log).sum()/L_w\n",
    "\n",
    "            if k+1 == plot_seed:\n",
    "                gammas[i,j] = gamma\n",
    "                viterbis[i,j] = get_viterbi(np.log(transitions[j]), np.log(emissions), np.log(ids[j]), L_w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63500 characters]\n",
      "text/plain": [
       "<Figure size 1440x972 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_viterbi(viterbis, tss, ects, L_w, w, 'viterbi2.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 352260 characters]\n",
      "text/plain": [
       "<Figure size 1440x1944 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_posterior(gammas, tss, ects, L_w, w, 'posterior2.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Africa</th>\n",
       "      <th>Zigzag</th>\n",
       "      <th>Constant</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Africa</th>\n",
       "      <td>10144.40 (19.76)</td>\n",
       "      <td>10359.75 (20.84)</td>\n",
       "      <td>10663.13 (20.19)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Zigzag</th>\n",
       "      <td>5507.79 (132.03)</td>\n",
       "      <td>4962.75 (120.63)</td>\n",
       "      <td>5700.60 (137.21)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Constant</th>\n",
       "      <td>11584.24 (259.28)</td>\n",
       "      <td>11764.86 (263.18)</td>\n",
       "      <td>11898.49 (266.09)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                     Africa             Zigzag           Constant\n",
       "Africa     10144.40 (19.76)   10359.75 (20.84)   10663.13 (20.19)\n",
       "Zigzag     5507.79 (132.03)   4962.75 (120.63)   5700.60 (137.21)\n",
       "Constant  11584.24 (259.28)  11764.86 (263.18)  11898.49 (266.09)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "create_table(err_A, 'table2_1.tex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Africa</th>\n",
       "      <th>Zigzag</th>\n",
       "      <th>Constant</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Africa</th>\n",
       "      <td>0.3293 (0.0004)</td>\n",
       "      <td>0.3934 (0.0005)</td>\n",
       "      <td>0.3549 (0.0004)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Zigzag</th>\n",
       "      <td>0.3678 (0.0041)</td>\n",
       "      <td>0.3548 (0.0042)</td>\n",
       "      <td>0.3788 (0.0038)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Constant</th>\n",
       "      <td>0.3496 (0.0068)</td>\n",
       "      <td>0.3955 (0.0059)</td>\n",
       "      <td>0.3679 (0.0064)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   Africa           Zigzag         Constant\n",
       "Africa    0.3293 (0.0004)  0.3934 (0.0005)  0.3549 (0.0004)\n",
       "Zigzag    0.3678 (0.0041)  0.3548 (0.0042)  0.3788 (0.0038)\n",
       "Constant  0.3496 (0.0068)  0.3955 (0.0059)  0.3679 (0.0064)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "create_table(err_B, 'table2_2.tex', log=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tsmc",
   "language": "python",
   "name": "tsmc"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
