{
 "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",
       "                var nbb_formatted_code = \"%load_ext nb_black\";\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1000 Genomes analysis\n",
    "The tree sequence files for 1kg are available [here](https://zenodo.org/record/3051855). The other metadata files are part of the public 1kg release."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 2;\n",
       "                var nbb_unformatted_code = \"import os\\nimport tskit\\nimport numpy as np\\nimport xsmc\\nimport matplotlib.pyplot as plt\\nfrom concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor\\nfrom xsmc.supporting.kde_ne import kde_ne\\nfrom xsmc.supporting.plotting import *\\nimport logging\\nimport os\\n\\n\\nlogging.getLogger(\\\"xsmc\\\").setLevel(logging.INFO)\\nPAPER_ROOT = os.path.expanduser(os.environ.get(\\\"PAPER_ROOT\\\", \\\".\\\"))\";\n",
       "                var nbb_formatted_code = \"import os\\nimport tskit\\nimport numpy as np\\nimport xsmc\\nimport matplotlib.pyplot as plt\\nfrom concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor\\nfrom xsmc.supporting.kde_ne import kde_ne\\nfrom xsmc.supporting.plotting import *\\nimport logging\\nimport os\\n\\n\\nlogging.getLogger(\\\"xsmc\\\").setLevel(logging.INFO)\\nPAPER_ROOT = os.path.expanduser(os.environ.get(\\\"PAPER_ROOT\\\", \\\".\\\"))\";\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": [
    "import os\n",
    "import tskit\n",
    "import numpy as np\n",
    "import xsmc\n",
    "import matplotlib.pyplot as plt\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor\n",
    "from xsmc.supporting.kde_ne import kde_ne\n",
    "from xsmc.supporting.plotting import *\n",
    "import logging\n",
    "import os\n",
    "\n",
    "\n",
    "logging.getLogger(\"xsmc\").setLevel(logging.INFO)\n",
    "PAPER_ROOT = os.path.expanduser(os.environ.get(\"PAPER_ROOT\", \".\"))"
   ]
  },
  {
   "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 = \"np.random.seed(1)\\n\\n\\ndef seed():\\n    return np.random.randint(1, np.iinfo(np.int32).max)\";\n",
       "                var nbb_formatted_code = \"np.random.seed(1)\\n\\n\\ndef seed():\\n    return np.random.randint(1, np.iinfo(np.int32).max)\";\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": [
    "np.random.seed(1)\n",
    "\n",
    "\n",
    "def seed():\n",
    "    return np.random.randint(1, np.iinfo(np.int32).max)"
   ]
  },
  {
   "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 = \"full_chroms = [tskit.load(f\\\"/scratch/1kg/trees/1kg_chr{i}.trees\\\") for i in range(1, 23)]\\n# chr1 = chr1.delete_intervals([[0, 16e6]], simplify=False).trim()\";\n",
       "                var nbb_formatted_code = \"full_chroms = [tskit.load(f\\\"/scratch/1kg/trees/1kg_chr{i}.trees\\\") for i in range(1, 23)]\\n# chr1 = chr1.delete_intervals([[0, 16e6]], simplify=False).trim()\";\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": [
    "full_chroms = [tskit.load(f\"/scratch/1kg/trees/1kg_chr{i}.trees\") for i in range(1, 23)]\n",
    "# chr1 = chr1.delete_intervals([[0, 16e6]], simplify=False).trim()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accessibility masking\n",
    "The 1kg tree sequences feature some trees that span huge intervals. These represent centromeres, telomeres, or other inaccessible regions. To our method they appear as huge tracts of IBD, leading to downward bias in the size history estimates. To properly correct for this we should use the 1kg accessibility masks, but here I adopt the simpler approach of heuristically chopping out trees that span large segments (>1Mb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: [Interval(left=121347335.0, right=142643153.0)],\n",
       " 2: [Interval(left=90490322.0, right=91634056.0),\n",
       "  Interval(left=92102205.0, right=95333904.0)],\n",
       " 3: [Interval(left=90291128.0, right=93527692.0)],\n",
       " 4: [Interval(left=49629442.0, right=52700457.0)],\n",
       " 5: [Interval(left=46118241.0, right=49571959.0)],\n",
       " 6: [Interval(left=58665913.0, right=61967504.0)],\n",
       " 7: [Interval(left=57669630.0, right=62465883.0)],\n",
       " 8: [Interval(left=43529404.0, right=47458633.0)],\n",
       " 9: [Interval(left=47309856.0, right=65510886.0)],\n",
       " 10: [Interval(left=39043557.0, right=42746452.0)],\n",
       " 11: [Interval(left=50381334.0, right=51436554.0),\n",
       "  Interval(left=51510699.0, right=55254669.0)],\n",
       " 12: [Interval(left=34340941.0, right=38610698.0)],\n",
       " 13: [Interval(left=0.0, right=19037838.0)],\n",
       " 14: [Interval(left=0.0, right=19083007.0)],\n",
       " 15: [Interval(left=0.0, right=20253797.0)],\n",
       " 16: [Interval(left=35146682.0, right=46554124.0)],\n",
       " 17: [Interval(left=22155426.0, right=25392232.0)],\n",
       " 18: [Interval(left=15357425.0, right=18535477.0)],\n",
       " 19: [Interval(left=24310556.0, right=28366182.0)],\n",
       " 20: [Interval(left=26193540.0, right=29862565.0)],\n",
       " 21: [Interval(left=0.0, right=9491942.0),\n",
       "  Interval(left=11010799.0, right=14383028.0)],\n",
       " 22: [Interval(left=0.0, right=16103536.0)]}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 5;\n",
       "                var nbb_unformatted_code = \"long_spans = {}\\nfor i, chrom in enumerate(full_chroms, 1):\\n    long_spans[i] = [t.interval for t in chrom.trees() if np.diff(t.interval) > 1e6]\\nlong_spans\";\n",
       "                var nbb_formatted_code = \"long_spans = {}\\nfor i, chrom in enumerate(full_chroms, 1):\\n    long_spans[i] = [t.interval for t in chrom.trees() if np.diff(t.interval) > 1e6]\\nlong_spans\";\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": [
    "long_spans = {}\n",
    "for i, chrom in enumerate(full_chroms, 1):\n",
    "    long_spans[i] = [t.interval for t in chrom.trees() if np.diff(t.interval) > 1e6]\n",
    "long_spans"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we simply chop these intervals out. This technically allows IBD tracts to span these gaps during inference, but the effect should be minimal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 6;\n",
       "                var nbb_unformatted_code = \"chroms = [\\n    chrom.delete_intervals(long_spans[i], simplify=False).trim()\\n    for i, chrom in enumerate(full_chroms, 1)\\n]\";\n",
       "                var nbb_formatted_code = \"chroms = [\\n    chrom.delete_intervals(long_spans[i], simplify=False).trim()\\n    for i, chrom in enumerate(full_chroms, 1)\\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": [
    "chroms = [\n",
    "    chrom.delete_intervals(long_spans[i], simplify=False).trim()\n",
    "    for i, chrom in enumerate(full_chroms, 1)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 7;\n",
       "                var nbb_unformatted_code = \"# 1kg metadata\\n\\nwith open(\\\"/scratch/1kg/integrated_call_samples_v3.20130502.ALL.panel\\\", \\\"rt\\\") as f:\\n    next(f)\\n    rows = (line.strip().split(\\\"\\\\t\\\") for line in f)\\n    sample_map = {sample_id: (pop, superpop) for sample_id, pop, superpop, _ in rows}\";\n",
       "                var nbb_formatted_code = \"# 1kg metadata\\n\\nwith open(\\\"/scratch/1kg/integrated_call_samples_v3.20130502.ALL.panel\\\", \\\"rt\\\") as f:\\n    next(f)\\n    rows = (line.strip().split(\\\"\\\\t\\\") for line in f)\\n    sample_map = {sample_id: (pop, superpop) for sample_id, pop, superpop, _ in rows}\";\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": [
    "# 1kg metadata\n",
    "\n",
    "with open(\"/scratch/1kg/integrated_call_samples_v3.20130502.ALL.panel\", \"rt\") as f:\n",
    "    next(f)\n",
    "    rows = (line.strip().split(\"\\t\") for line in f)\n",
    "    sample_map = {sample_id: (pop, superpop) for sample_id, pop, superpop, _ in rows}"
   ]
  },
  {
   "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 = \"# map each 1kg sample id ts nodes\\nimport json\\n\\nsamples_to_nodes = {\\n    json.loads(ind.metadata)[\\\"individual_id\\\"]: ind.nodes\\n    for ind in chroms[0].individuals()\\n}\";\n",
       "                var nbb_formatted_code = \"# map each 1kg sample id ts nodes\\nimport json\\n\\nsamples_to_nodes = {\\n    json.loads(ind.metadata)[\\\"individual_id\\\"]: ind.nodes\\n    for ind in chroms[0].individuals()\\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": [
    "# map each 1kg sample id ts nodes\n",
    "import json\n",
    "\n",
    "samples_to_nodes = {\n",
    "    json.loads(ind.metadata)[\"individual_id\"]: ind.nodes\n",
    "    for ind in chroms[0].individuals()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 9;\n",
       "                var nbb_unformatted_code = \"superpops = {}\\nfor sample_id, (p, sp) in sample_map.items():\\n    superpops.setdefault(sp, [])\\n    superpops[sp].append(sample_id)\";\n",
       "                var nbb_formatted_code = \"superpops = {}\\nfor sample_id, (p, sp) in sample_map.items():\\n    superpops.setdefault(sp, [])\\n    superpops[sp].append(sample_id)\";\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": [
    "superpops = {}\n",
    "for sample_id, (p, sp) in sample_map.items():\n",
    "    superpops.setdefault(sp, [])\n",
    "    superpops[sp].append(sample_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 10;\n",
       "                var nbb_unformatted_code = \"K = 20  # number of samples to take from each superpopulation\\nmu = 1.4e-8  # assumed mutation rate for humans\\n\\n\\ndef process_samples(sample_dict, w, rho_over_theta):\\n    sampled_heights = {}\\n    lines = {}\\n    for sp in sample_dict:\\n        print(sp, flush=True)\\n        xs = []\\n        for sample_id in sample_dict[sp][:K]:\\n            print(\\\"\\\\t%s\\\" % sample_id)\\n            f, p = samples_to_nodes[sample_id]\\n            for data in chroms:\\n                xs.append(\\n                    xsmc.XSMC(\\n                        data, focal=f, panel=[p], w=w, rho_over_theta=rho_over_theta\\n                    )\\n                )\\n        with ThreadPoolExecutor(24) as p:\\n            futs = [\\n                p.submit(\\n                    x.sample_heights,\\n                    j=100,\\n                    k=int(x.ts.get_sequence_length() / 50_000),\\n                    seed=seed(),\\n                )\\n                for x in xs\\n            ]\\n\\n        sampled_heights[sp] = np.concatenate(\\n            [f.result() * 2 * x.theta / (4 * mu) for f, x in zip(futs, xs)], axis=1\\n        )  # rescale each sampled path by 2N0 so that segment heights are in generations\\n    return sampled_heights\";\n",
       "                var nbb_formatted_code = \"K = 20  # number of samples to take from each superpopulation\\nmu = 1.4e-8  # assumed mutation rate for humans\\n\\n\\ndef process_samples(sample_dict, w, rho_over_theta):\\n    sampled_heights = {}\\n    lines = {}\\n    for sp in sample_dict:\\n        print(sp, flush=True)\\n        xs = []\\n        for sample_id in sample_dict[sp][:K]:\\n            print(\\\"\\\\t%s\\\" % sample_id)\\n            f, p = samples_to_nodes[sample_id]\\n            for data in chroms:\\n                xs.append(\\n                    xsmc.XSMC(\\n                        data, focal=f, panel=[p], w=w, rho_over_theta=rho_over_theta\\n                    )\\n                )\\n        with ThreadPoolExecutor(24) as p:\\n            futs = [\\n                p.submit(\\n                    x.sample_heights,\\n                    j=100,\\n                    k=int(x.ts.get_sequence_length() / 50_000),\\n                    seed=seed(),\\n                )\\n                for x in xs\\n            ]\\n\\n        sampled_heights[sp] = np.concatenate(\\n            [f.result() * 2 * x.theta / (4 * mu) for f, x in zip(futs, xs)], axis=1\\n        )  # rescale each sampled path by 2N0 so that segment heights are in generations\\n    return sampled_heights\";\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": [
    "K = 20  # number of samples to take from each superpopulation\n",
    "mu = 1.4e-8  # assumed mutation rate for humans\n",
    "\n",
    "\n",
    "def process_samples(sample_dict, w, rho_over_theta):\n",
    "    sampled_heights = {}\n",
    "    lines = {}\n",
    "    for sp in sample_dict:\n",
    "        print(sp, flush=True)\n",
    "        xs = []\n",
    "        for sample_id in sample_dict[sp][:K]:\n",
    "            print(\"\\t%s\" % sample_id)\n",
    "            f, p = samples_to_nodes[sample_id]\n",
    "            for data in chroms:\n",
    "                xs.append(\n",
    "                    xsmc.XSMC(\n",
    "                        data, focal=f, panel=[p], w=w, rho_over_theta=rho_over_theta\n",
    "                    )\n",
    "                )\n",
    "        with ThreadPoolExecutor(24) as p:\n",
    "            futs = [\n",
    "                p.submit(\n",
    "                    x.sample_heights,\n",
    "                    j=100,\n",
    "                    k=int(x.ts.get_sequence_length() / 50_000),\n",
    "                    seed=seed(),\n",
    "                )\n",
    "                for x in xs\n",
    "            ]\n",
    "\n",
    "        sampled_heights[sp] = np.concatenate(\n",
    "            [f.result() * 2 * x.theta / (4 * mu) for f, x in zip(futs, xs)], axis=1\n",
    "        )  # rescale each sampled path by 2N0 so that segment heights are in generations\n",
    "    return sampled_heights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time to process 1 whole genome"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test\n",
      "\tNA12878\n",
      "CPU times: user 7min 36s, sys: 743 ms, total: 7min 37s\n",
      "Wall time: 37.8 s\n"
     ]
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 11;\n",
       "                var nbb_unformatted_code = \"%%time\\n_ = process_samples({'test': ['NA12878']}, w=500, rho_over_theta=1.)\";\n",
       "                var nbb_formatted_code = \"%%time\\n_ = process_samples({'test': ['NA12878']}, w=500, rho_over_theta=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": [
    "%%time\n",
    "_ = process_samples({'test': ['NA12878']}, w=500, rho_over_theta=1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main event\n",
    "Run pipeline for all 1kg samples split into 5 superpopulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 12;\n",
       "                var nbb_unformatted_code = \"def make_plot(sampled_heights, g, name):\\n    fig, ax = plt.subplots(figsize=(8, 5))\\n    x0 = np.geomspace(1e2, 1e6, 1000)\\n    for sp, col in zip(sampled_heights, TABLEAU):\\n        lines = []\\n        A = np.array(sampled_heights[sp])\\n        with ProcessPoolExecutor() as p:\\n            futs = [p.submit(kde_ne, a.reshape(-1)) for a in A]\\n            for f in as_completed(futs):\\n                x, y = f.result()\\n                lines.append((x * g, y / 2))  # rescale to years, plot diploid Ne\\n        plot_summary(ax, lines, x0, color=col, label=sp)\\n    ax.set_xscale(\\\"log\\\")\\n    ax.set_yscale(\\\"log\\\")\\n    ax.legend(loc=\\\"lower right\\\")\\n    ax.set_xlabel(r\\\"Years ($\\\\mu=1.4\\\\times 10^{-8}, g=%d$)\\\" % g)\\n    ax.set_ylabel(\\\"Effective Population Size ($N_e$)\\\")\\n    fig.tight_layout()\\n    fig.savefig(os.path.join(PAPER_ROOT, \\\"figures\\\", \\\"xsmc_1kg_%s.pdf\\\" % name))\\n    return fig\";\n",
       "                var nbb_formatted_code = \"def make_plot(sampled_heights, g, name):\\n    fig, ax = plt.subplots(figsize=(8, 5))\\n    x0 = np.geomspace(1e2, 1e6, 1000)\\n    for sp, col in zip(sampled_heights, TABLEAU):\\n        lines = []\\n        A = np.array(sampled_heights[sp])\\n        with ProcessPoolExecutor() as p:\\n            futs = [p.submit(kde_ne, a.reshape(-1)) for a in A]\\n            for f in as_completed(futs):\\n                x, y = f.result()\\n                lines.append((x * g, y / 2))  # rescale to years, plot diploid Ne\\n        plot_summary(ax, lines, x0, color=col, label=sp)\\n    ax.set_xscale(\\\"log\\\")\\n    ax.set_yscale(\\\"log\\\")\\n    ax.legend(loc=\\\"lower right\\\")\\n    ax.set_xlabel(r\\\"Years ($\\\\mu=1.4\\\\times 10^{-8}, g=%d$)\\\" % g)\\n    ax.set_ylabel(\\\"Effective Population Size ($N_e$)\\\")\\n    fig.tight_layout()\\n    fig.savefig(os.path.join(PAPER_ROOT, \\\"figures\\\", \\\"xsmc_1kg_%s.pdf\\\" % name))\\n    return fig\";\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 make_plot(sampled_heights, g, name):\n",
    "    fig, ax = plt.subplots(figsize=(8, 5))\n",
    "    x0 = np.geomspace(1e2, 1e6, 1000)\n",
    "    for sp, col in zip(sampled_heights, TABLEAU):\n",
    "        lines = []\n",
    "        A = np.array(sampled_heights[sp])\n",
    "        with ProcessPoolExecutor() as p:\n",
    "            futs = [p.submit(kde_ne, a.reshape(-1)) for a in A]\n",
    "            for f in as_completed(futs):\n",
    "                x, y = f.result()\n",
    "                lines.append((x * g, y / 2))  # rescale to years, plot diploid Ne\n",
    "        plot_summary(ax, lines, x0, color=col, label=sp)\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_yscale(\"log\")\n",
    "    ax.legend(loc=\"lower right\")\n",
    "    ax.set_xlabel(r\"Years ($\\mu=1.4\\times 10^{-8}, g=%d$)\" % g)\n",
    "    ax.set_ylabel(\"Effective Population Size ($N_e$)\")\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(os.path.join(PAPER_ROOT, \"figures\", \"xsmc_1kg_%s.pdf\" % name))\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EUR\n",
      "\tHG00096\n",
      "\tHG00097\n",
      "\tHG00099\n",
      "\tHG00100\n",
      "\tHG00101\n",
      "\tHG00102\n",
      "\tHG00103\n",
      "\tHG00105\n",
      "\tHG00106\n",
      "\tHG00107\n",
      "\tHG00108\n",
      "\tHG00109\n",
      "\tHG00110\n",
      "\tHG00111\n",
      "\tHG00112\n",
      "\tHG00113\n",
      "\tHG00114\n",
      "\tHG00115\n",
      "\tHG00116\n",
      "\tHG00117\n",
      "EAS\n",
      "\tHG00403\n",
      "\tHG00404\n",
      "\tHG00406\n",
      "\tHG00407\n",
      "\tHG00409\n",
      "\tHG00410\n",
      "\tHG00419\n",
      "\tHG00421\n",
      "\tHG00422\n",
      "\tHG00428\n",
      "\tHG00436\n",
      "\tHG00437\n",
      "\tHG00442\n",
      "\tHG00443\n",
      "\tHG00445\n",
      "\tHG00446\n",
      "\tHG00448\n",
      "\tHG00449\n",
      "\tHG00451\n",
      "\tHG00452\n",
      "AMR\n",
      "\tHG00551\n",
      "\tHG00553\n",
      "\tHG00554\n",
      "\tHG00637\n",
      "\tHG00638\n",
      "\tHG00640\n",
      "\tHG00641\n",
      "\tHG00731\n",
      "\tHG00732\n",
      "\tHG00734\n",
      "\tHG00736\n",
      "\tHG00737\n",
      "\tHG00739\n",
      "\tHG00740\n",
      "\tHG00742\n",
      "\tHG00743\n",
      "\tHG01047\n",
      "\tHG01048\n",
      "\tHG01049\n",
      "\tHG01051\n",
      "SAS\n",
      "\tHG01583\n",
      "\tHG01586\n",
      "\tHG01589\n",
      "\tHG01593\n",
      "\tHG02490\n",
      "\tHG02491\n",
      "\tHG02493\n",
      "\tHG02494\n",
      "\tHG02597\n",
      "\tHG02600\n",
      "\tHG02601\n",
      "\tHG02603\n",
      "\tHG02604\n",
      "\tHG02648\n",
      "\tHG02649\n",
      "\tHG02651\n",
      "\tHG02652\n",
      "\tHG02654\n",
      "\tHG02655\n",
      "\tHG02657\n",
      "AFR\n",
      "\tHG01879\n",
      "\tHG01880\n",
      "\tHG01882\n",
      "\tHG01883\n",
      "\tHG01885\n",
      "\tHG01886\n",
      "\tHG01889\n",
      "\tHG01890\n",
      "\tHG01894\n",
      "\tHG01896\n",
      "\tHG01912\n",
      "\tHG01914\n",
      "\tHG01915\n",
      "\tHG01956\n",
      "\tHG01958\n",
      "\tHG01985\n",
      "\tHG01986\n",
      "\tHG01988\n",
      "\tHG01989\n",
      "\tHG01990\n",
      "CPU times: user 15h 50min 50s, sys: 32.1 s, total: 15h 51min 22s\n",
      "Wall time: 41min 10s\n"
     ]
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 13;\n",
       "                var nbb_unformatted_code = \"%%time\\nsampled_heights = process_samples(superpops, w=500, rho_over_theta=1.)\";\n",
       "                var nbb_formatted_code = \"%%time\\nsampled_heights = process_samples(superpops, w=500, rho_over_theta=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": [
    "%%time\n",
    "sampled_heights = process_samples(superpops, w=500, rho_over_theta=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/terhorst/opt/py37/lib/python3.7/site-packages/numpy/lib/nanfunctions.py:1392: RuntimeWarning: All-NaN slice encountered\n",
      "  overwrite_input, interpolation)\n",
      "2020-09-21 16:03:00,023 WARNING matplotlib.font_manager MainThread findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.\n",
      "2020-09-21 16:03:00,225 WARNING matplotlib.font_manager MainThread findfont: Font family ['sans-serif'] not found. Falling back to DejaVu Sans.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49140 characters]\n",
      "text/plain": [
       "<Figure size 576x360 with 1 Axes>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 49140 characters]\n",
      "text/plain": [
       "<Figure size 576x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "            setTimeout(function() {\n",
       "                var nbb_cell_id = 14;\n",
       "                var nbb_unformatted_code = \"make_plot(sampled_heights, 29, \\\"final\\\")\";\n",
       "                var nbb_formatted_code = \"make_plot(sampled_heights, 29, \\\"final\\\")\";\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": [
    "make_plot(sampled_heights, 29, \"final\")"
   ]
  }
 ],
 "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
}
