{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 1;\n",
       "                var nbb_unformatted_code = \"%load_ext nb_black\\n%matplotlib inline\";\n",
       "                var nbb_formatted_code = \"%load_ext nb_black\\n%matplotlib inline\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%load_ext nb_black\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 2;\n",
       "                var nbb_unformatted_code = \"## setup code\\nimport matplotlib.pyplot as plt\\n\\nplt.style.use(\\\"seaborn\\\")\\nplt.rc(\\\"figure\\\", figsize=(15, 3))\\nimport numpy as np\\nimport msprime as msp\\n\\nnp.random.seed(1)\";\n",
       "                var nbb_formatted_code = \"## setup code\\nimport matplotlib.pyplot as plt\\n\\nplt.style.use(\\\"seaborn\\\")\\nplt.rc(\\\"figure\\\", figsize=(15, 3))\\nimport numpy as np\\nimport msprime as msp\\n\\nnp.random.seed(1)\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## setup code\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.style.use(\"seaborn\")\n",
    "plt.rc(\"figure\", figsize=(15, 3))\n",
    "import numpy as np\n",
    "import msprime as msp\n",
    "\n",
    "np.random.seed(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quick start guide\n",
    "This packages implements methods described in [the paper](https://doi.org/10.1101/2020.09.21.307355):\n",
    "\n",
    "    Exact decoding of the sequentially Markov coalescent\n",
    "    Caleb Ki, Jonathan Terhorst\n",
    "    bioRxiv 2020.09.21.307355; doi:10.1101/2020.09.21.307355\n",
    "    \n",
    "The methods are implemented in a Python package called `xsmc`. The package has two main uses: drawing exact samples from the posterior distribution, and fast estimation of the maximum *a posteriori* (MAP) path. Convenience methods are also included for plotting."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will demonstrate how the method works on simulated data using `msprime`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 3;\n",
       "                var nbb_unformatted_code = \"def simulate(recombination_rate, length, sample_size=2):\\n    return msp.simulate(\\n        sample_size=sample_size,\\n        mutation_rate=1e-8,\\n        recombination_rate=recombination_rate,\\n        Ne=1e4,\\n        length=length,\\n        random_seed=np.random.randint(1, int(4e9)),\\n        demographic_events=[\\n            msp.PopulationParametersChange(time=0.0, initial_size=1e5),\\n            msp.PopulationParametersChange(time=1000.0, initial_size=1e4),\\n        ],\\n    )\";\n",
       "                var nbb_formatted_code = \"def simulate(recombination_rate, length, sample_size=2):\\n    return msp.simulate(\\n        sample_size=sample_size,\\n        mutation_rate=1e-8,\\n        recombination_rate=recombination_rate,\\n        Ne=1e4,\\n        length=length,\\n        random_seed=np.random.randint(1, int(4e9)),\\n        demographic_events=[\\n            msp.PopulationParametersChange(time=0.0, initial_size=1e5),\\n            msp.PopulationParametersChange(time=1000.0, initial_size=1e4),\\n        ],\\n    )\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def simulate(recombination_rate, length, sample_size=2):\n",
    "    return msp.simulate(\n",
    "        sample_size=sample_size,\n",
    "        mutation_rate=1e-8,\n",
    "        recombination_rate=recombination_rate,\n",
    "        Ne=1e4,\n",
    "        length=length,\n",
    "        random_seed=np.random.randint(1, int(4e9)),\n",
    "        demographic_events=[\n",
    "            msp.PopulationParametersChange(time=0.0, initial_size=1e5),\n",
    "            msp.PopulationParametersChange(time=1000.0, initial_size=1e4),\n",
    "        ],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All methods are accessed by creating an instance of the `xsmc.XSMC` class. The class has three required arguments:\n",
    "\n",
    "- `ts`, a `tskit.TreeSequence` instance containing the data.\n",
    "- `focal`, the node id of the focal haplotype in `ts`.\n",
    "- `panel`, a list of node ids in `ts` which are potential targets for `focal` to copy from.\n",
    "\n",
    "First, we\n",
    "focus on the simplest type of analysis, where we have two chromosomes sampled from a singled diploid individual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 4;\n",
       "                var nbb_unformatted_code = \"from xsmc import XSMC, Segmentation\\n\\ndata = simulate(recombination_rate=1e-9, length=1e6, sample_size=2)\\nxs = XSMC(data, focal=0, panel=[1], rho_over_theta=0.1)\";\n",
       "                var nbb_formatted_code = \"from xsmc import XSMC, Segmentation\\n\\ndata = simulate(recombination_rate=1e-9, length=1e6, sample_size=2)\\nxs = XSMC(data, focal=0, panel=[1], rho_over_theta=0.1)\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from xsmc import XSMC, Segmentation\n",
    "\n",
    "data = simulate(recombination_rate=1e-9, length=1e6, sample_size=2)\n",
    "xs = XSMC(data, focal=0, panel=[1], rho_over_theta=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exact sampling\n",
    "To sample paths from the posterior, use the `XSMC.sample()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 5;\n",
       "                var nbb_unformatted_code = \"paths = xs.sample(k=100, seed=1)\";\n",
       "                var nbb_formatted_code = \"paths = xs.sample(k=100, seed=1)\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "paths = xs.sample(k=100, seed=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returned object is a list of length `k` of `Segmentation` objects. These objects can be easily plotted to visualize the posterior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81020 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 6;\n",
       "                var nbb_unformatted_code = \"for p in paths:\\n    p.draw(alpha=0.1)\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_formatted_code = \"for p in paths:\\n    p.draw(alpha=0.1)\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for p in paths:\n",
    "    p.draw(alpha=0.1)\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compare the posterior with the truth, we can use built-in convenience methods to turn data into a `Segmentation` object and plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6204 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 7;\n",
       "                var nbb_unformatted_code = \"truth = Segmentation.from_ts(data, focal=0, panel=[1])\\ntruth.draw()\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_formatted_code = \"truth = Segmentation.from_ts(data, focal=0, panel=[1])\\ntruth.draw()\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "truth = Segmentation.from_ts(data, focal=0, panel=[1])\n",
    "truth.draw()\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rescaling the output\n",
    "You may notice that the $y$ scale of the above plots is different. By default, `Segmentation`s returned by `XSMC` are in units of coalescent time. This matches the scaling of the Li & Durbin's PSMC.\n",
    "\n",
    "To convert these times to generations, you must know the biological mutation rate $\\mu$, and then estimate $N_0$ by the formula $N_0 = \\theta / (4 \\mu)$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 8;\n",
       "                var nbb_unformatted_code = \"N0 = xs.theta / (4 * 1e-8)\";\n",
       "                var nbb_formatted_code = \"N0 = xs.theta / (4 * 1e-8)\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "N0 = xs.theta / (4 * 1e-8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, you can convert time to generations by multiplying by $2N_0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 60716 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 9;\n",
       "                var nbb_unformatted_code = \"for p in paths:\\n    p.rescale(2 * N0).draw(\\n        color=\\\"tab:blue\\\", alpha=0.1\\n    )  # divide by the mutation rate used in the simulation, convert to haploid\\ntruth.draw(color=\\\"tab:red\\\")\\nplt.yscale(\\\"log\\\")\\n# plt.ylim(0, 1e4)\";\n",
       "                var nbb_formatted_code = \"for p in paths:\\n    p.rescale(2 * N0).draw(\\n        color=\\\"tab:blue\\\", alpha=0.1\\n    )  # divide by the mutation rate used in the simulation, convert to haploid\\ntruth.draw(color=\\\"tab:red\\\")\\nplt.yscale(\\\"log\\\")\\n# plt.ylim(0, 1e4)\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for p in paths:\n",
    "    p.rescale(2 * N0).draw(\n",
    "        color=\"tab:blue\", alpha=0.1\n",
    "    )  # divide by the mutation rate used in the simulation, convert to haploid\n",
    "truth.draw(color=\"tab:red\")\n",
    "plt.yscale(\"log\")\n",
    "# plt.ylim(0, 1e4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing functionals of the posterior\n",
    "The `Segmentation.to_pp()` method can be used to get access to the raw stepwise function it represents (returned as a **p**iecewise **p**olynomial, whence the name). By calling this method on paths returned from `XSMC.sample_paths()`, this you can compute functionals of the posterior distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32556 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 10;\n",
       "                var nbb_unformatted_code = \"x = np.linspace(0, data.get_sequence_length(), 1000)\\ny = np.array([p.to_pp()(x) for p in paths])\\nq = (\\n    np.quantile(y, [0.25, 0.5, 0.75], axis=0) * 2 * N0\\n)  # faster than calling .rescale() on each path\\nplt.plot(x, q[1])\\nplt.fill_between(x, q[0], q[2], alpha=0.3)\\ntruth.draw()\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_formatted_code = \"x = np.linspace(0, data.get_sequence_length(), 1000)\\ny = np.array([p.to_pp()(x) for p in paths])\\nq = (\\n    np.quantile(y, [0.25, 0.5, 0.75], axis=0) * 2 * N0\\n)  # faster than calling .rescale() on each path\\nplt.plot(x, q[1])\\nplt.fill_between(x, q[0], q[2], alpha=0.3)\\ntruth.draw()\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = np.linspace(0, data.get_sequence_length(), 1000)\n",
    "y = np.array([p.to_pp()(x) for p in paths])\n",
    "q = (\n",
    "    np.quantile(y, [0.25, 0.5, 0.75], axis=0) * 2 * N0\n",
    ")  # faster than calling .rescale() on each path\n",
    "plt.plot(x, q[1])\n",
    "plt.fill_between(x, q[0], q[2], alpha=0.3)\n",
    "truth.draw()\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiple samples\n",
    "We can analyze multiple samples using a \"trunk approximation\" described in the paper. In this case, the returned segmentation approximates the earliest time to coalescence with any member of the panel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 11;\n",
       "                var nbb_unformatted_code = \"data = simulate(recombination_rate=1e-8, length=1e6, sample_size=20)\\nfocal = 0\\npanel = list(range(1, 20))\\nxs = XSMC(data, focal=focal, panel=panel, rho_over_theta=1.0)\\npaths = xs.sample(k=100, seed=1)\";\n",
       "                var nbb_formatted_code = \"data = simulate(recombination_rate=1e-8, length=1e6, sample_size=20)\\nfocal = 0\\npanel = list(range(1, 20))\\nxs = XSMC(data, focal=focal, panel=panel, rho_over_theta=1.0)\\npaths = xs.sample(k=100, seed=1)\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = simulate(recombination_rate=1e-8, length=1e6, sample_size=20)\n",
    "focal = 0\n",
    "panel = list(range(1, 20))\n",
    "xs = XSMC(data, focal=focal, panel=panel, rho_over_theta=1.0)\n",
    "paths = xs.sample(k=100, seed=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56040 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 12;\n",
       "                var nbb_unformatted_code = \"N0 = xs.theta / (4 * 1e-8)\\nfor p in paths:\\n    p.rescale(2 * N0).draw(color=\\\"tab:blue\\\", alpha=0.05)\\ntruth = Segmentation.from_ts(data, focal=focal, panel=panel)\\ntruth.draw(color=\\\"tab:red\\\")\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_formatted_code = \"N0 = xs.theta / (4 * 1e-8)\\nfor p in paths:\\n    p.rescale(2 * N0).draw(color=\\\"tab:blue\\\", alpha=0.05)\\ntruth = Segmentation.from_ts(data, focal=focal, panel=panel)\\ntruth.draw(color=\\\"tab:red\\\")\\nplt.yscale(\\\"log\\\")\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "N0 = xs.theta / (4 * 1e-8)\n",
    "for p in paths:\n",
    "    p.rescale(2 * N0).draw(color=\"tab:blue\", alpha=0.05)\n",
    "truth = Segmentation.from_ts(data, focal=focal, panel=panel)\n",
    "truth.draw(color=\"tab:red\")\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Viterbi Decoding\n",
    "The other function of the `XSMC` class is to compute the Viterbi/most probable path under the posterior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f0493c47190>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 17588 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 13;\n",
       "                var nbb_unformatted_code = \"data = simulate(recombination_rate=1e-9, length=10e6, sample_size=10)\\nxs = XSMC(data, 0, [1, 2], rho_over_theta=1e-9 / 1.4e-8)\\nN0 = xs.theta / (4 * 1e-8)\\nmap_path = xs.viterbi()\\nmap_path.rescale(2 * N0).draw(label=\\\"viterbi\\\")\\nSegmentation.from_ts(data, 0, [1, 2]).draw(label=\\\"truth\\\")\\nplt.yscale(\\\"log\\\")\\nplt.legend()\";\n",
       "                var nbb_formatted_code = \"data = simulate(recombination_rate=1e-9, length=10e6, sample_size=10)\\nxs = XSMC(data, 0, [1, 2], rho_over_theta=1e-9 / 1.4e-8)\\nN0 = xs.theta / (4 * 1e-8)\\nmap_path = xs.viterbi()\\nmap_path.rescale(2 * N0).draw(label=\\\"viterbi\\\")\\nSegmentation.from_ts(data, 0, [1, 2]).draw(label=\\\"truth\\\")\\nplt.yscale(\\\"log\\\")\\nplt.legend()\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = simulate(recombination_rate=1e-9, length=10e6, sample_size=10)\n",
    "xs = XSMC(data, 0, [1, 2], rho_over_theta=1e-9 / 1.4e-8)\n",
    "N0 = xs.theta / (4 * 1e-8)\n",
    "map_path = xs.viterbi()\n",
    "map_path.rescale(2 * N0).draw(label=\"viterbi\")\n",
    "Segmentation.from_ts(data, 0, [1, 2]).draw(label=\"truth\")\n",
    "plt.yscale(\"log\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that each path also includes a haplotype component which is not plotted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Segment(hap=2, interval=array([     0, 476205]), height=0.08770837623933274, mutations=21)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 14;\n",
       "                var nbb_unformatted_code = \"map_path.segments[0]\";\n",
       "                var nbb_formatted_code = \"map_path.segments[0]\";\n",
       "                var nbb_cells = Jupyter.notebook.get_cells();\n",
       "                for (var i = 0; i < nbb_cells.length; ++i) {\n",
       "                    if (nbb_cells[i].input_prompt_number == nbb_cell_id) {\n",
       "                        if (nbb_cells[i].get_text() == nbb_unformatted_code) {\n",
       "                             nbb_cells[i].set_text(nbb_formatted_code);\n",
       "                        }\n",
       "                        break;\n",
       "                    }\n",
       "                }\n",
       "            }, 500);\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "map_path.segments[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This says that the first segment in the MAP path had height 0.09 (in coalecent units) and coalesced with panel haplotype 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
