MackintoshSetter2024.supplementary_python_nb_1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "52c412a4-197c-4f19-a271-9365419491f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import msprime\n",
    "import tskit\n",
    "import numpy as np\n",
    "\n",
    "import os\n",
    "import subprocess\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import math \n",
    "import statistics\n",
    "\n",
    "import seaborn as sns\n",
    "from IPython.display import SVG, set_matplotlib_formats\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01597f69-c167-4ab6-b438-c020e2f884a4",
   "metadata": {},
   "source": [
    "### Estimating expected branch lengths under the IM model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09ee1a59-005a-4919-b464-559d8449feac",
   "metadata": {},
   "source": [
    "##### Set up a branch type dictionary for the unphased unrooted IM model\n",
    "##### Keys are lists of lineages which share a derived mutation, values are mutation types (see Laetsch et al. 2023)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cc1d8009-20e0-44e8-974b-3dd4c5f089da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(0, 1, 2, 3): None,\n",
       " (0,): 'hetA',\n",
       " (1,): 'hetA',\n",
       " (0, 2, 3): 'hetA',\n",
       " (1, 2, 3): 'hetA',\n",
       " (2,): 'hetB',\n",
       " (3,): 'hetB',\n",
       " (0, 1, 2): 'hetB',\n",
       " (0, 1, 3): 'hetB',\n",
       " (0, 2): 'hetAB',\n",
       " (0, 3): 'hetAB',\n",
       " (1, 2): 'hetAB',\n",
       " (1, 3): 'hetAB',\n",
       " (0, 1): 'FD',\n",
       " (2, 3): 'FD'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "HetATypes = [ # when a mutation is ancestral to any of these lists of leaves, it is a hetA mutation\n",
    "    [0],\n",
    "    [1],\n",
    "    [0,2,3], # [0, 2, 3] is equivalent to [1] on an unrooted tree\n",
    "    [1,2,3]\n",
    "]\n",
    "HetBTypes = [\n",
    "    [2],\n",
    "    [3],\n",
    "    [0,1,2],\n",
    "    [0,1,3]\n",
    "]\n",
    "FixedDiffs = [\n",
    "    [0,1],\n",
    "    [2,3]\n",
    "]\n",
    "HetABTypes = [\n",
    "    [0,2],\n",
    "    [0,3],\n",
    "    [1,2],\n",
    "    [1,3]\n",
    "]\n",
    "\n",
    "mutTypeDict = {(0,1,2,3):None};\n",
    "for ls in HetATypes:\n",
    "    mutTypeDict[tuple(sorted(ls))] = \"hetA\"\n",
    "for ls in HetBTypes:\n",
    "    mutTypeDict[tuple(sorted(ls))] = \"hetB\"\n",
    "for ls in HetABTypes:\n",
    "    mutTypeDict[tuple(sorted(ls))] = \"hetAB\"\n",
    "for ls in FixedDiffs:\n",
    "    mutTypeDict[tuple(sorted(ls))] = \"FD\"\n",
    "\n",
    "mutTypeDict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a624371b-a4b8-4910-8783-9b5d9c0a00a1",
   "metadata": {},
   "source": [
    "##### Set up simulations for the IM model with various parameters, add mutations, then show the first tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7d28d198-dbbe-498b-9503-9ab61f505c3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = 1e-8\n",
    "r = 0 # assume no recombination for now\n",
    "Na = 1e5\n",
    "Nb = 1e5\n",
    "Nanc = 1e5\n",
    "M = 1.0 / (2 * Na)\n",
    "Tsplit = 2e5 #generations\n",
    "\n",
    "L = int(1e4)\n",
    "block_length = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0eb17b2c-3179-4528-886b-94ce2322b1fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg baseProfile=\"full\" height=\"200\" version=\"1.1\" width=\"200\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:ev=\"http://www.w3.org/2001/xml-events\" xmlns:xlink=\"http://www.w3.org/1999/xlink\"><defs><style type=\"text/css\"><![CDATA[.background path {fill: #808080; fill-opacity: 0}.background path:nth-child(odd) {fill-opacity: .1}.axes {font-size: 14px}.x-axis .tick .lab {font-weight: bold; dominant-baseline: hanging}.axes, .tree {font-size: 14px; text-anchor: middle}.axes line, .edge {stroke: black; fill: none}.axes .ax-skip {stroke-dasharray: 4}.y-axis .grid {stroke: #FAFAFA}.node > .sym {fill: black; stroke: none}.site > .sym {stroke: black}.mut text {fill: red; font-style: italic}.mut.extra text {fill: hotpink}.mut line {fill: none; stroke: none}.mut .sym {fill: none; stroke: red}.mut.extra .sym {stroke: hotpink}.node .mut .sym {stroke-width: 1.5px}.tree text, .tree-sequence text {dominant-baseline: central}.plotbox .lab.lft {text-anchor: end}.plotbox .lab.rgt {text-anchor: start}]]></style></defs><g class=\"tree t0\"><g class=\"plotbox\"><g class=\"c2 node n6 p2 root\" transform=\"translate(100 26.8)\"><g class=\"a6 c2 m0 node n4 p1 s0\" transform=\"translate(-40 120.704)\"><g class=\"a4 i0 leaf node n0 p0 sample\" transform=\"translate(-20 20.6961)\"><path class=\"edge\" d=\"M 0 0 V -20.6961 H 20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">0</text></g><g class=\"a4 i1 leaf node n3 p1 sample\" transform=\"translate(20 20.6961)\"><path class=\"edge\" d=\"M 0 0 V -20.6961 H -20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">3</text></g><path class=\"edge\" d=\"M 0 0 V -120.704 H 40\" /><g class=\"mut m0 s0\" transform=\"translate(0 -17.9372)\"><line x1=\"0\" x2=\"0\" y1=\"0\" y2=\"17.9372\" /><path class=\"sym\" d=\"M -3,-3 l 6,6 M -3,3 l 6,-6\" /><text class=\"lab lft\" transform=\"translate(-5 0)\">0</text></g><circle class=\"sym\" cx=\"0\" cy=\"0\" r=\"3\" /><text class=\"lab lft\" transform=\"translate(-3 -7.0)\">4</text></g><g class=\"a6 c2 node n5 p1\" transform=\"translate(40 102.435)\"><g class=\"a5 i0 leaf node n1 p0 sample\" transform=\"translate(-20 38.9647)\"><path class=\"edge\" d=\"M 0 0 V -38.9647 H 20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">1</text></g><g class=\"a5 i1 leaf node n2 p1 sample\" transform=\"translate(20 38.9647)\"><path class=\"edge\" d=\"M 0 0 V -38.9647 H -20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">2</text></g><path class=\"edge\" d=\"M 0 0 V -102.435 H -40\" /><circle class=\"sym\" cx=\"0\" cy=\"0\" r=\"3\" /><text class=\"lab rgt\" transform=\"translate(3 -7.0)\">5</text></g><circle class=\"sym\" cx=\"0\" cy=\"0\" r=\"3\" /><text class=\"lab\" transform=\"translate(0 -11)\">6</text></g></g></g></svg>"
      ],
      "text/plain": [
       "'<svg baseProfile=\"full\" height=\"200\" version=\"1.1\" width=\"200\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:ev=\"http://www.w3.org/2001/xml-events\" xmlns:xlink=\"http://www.w3.org/1999/xlink\"><defs><style type=\"text/css\"><![CDATA[.background path {fill: #808080; fill-opacity: 0}.background path:nth-child(odd) {fill-opacity: .1}.axes {font-size: 14px}.x-axis .tick .lab {font-weight: bold; dominant-baseline: hanging}.axes, .tree {font-size: 14px; text-anchor: middle}.axes line, .edge {stroke: black; fill: none}.axes .ax-skip {stroke-dasharray: 4}.y-axis .grid {stroke: #FAFAFA}.node > .sym {fill: black; stroke: none}.site > .sym {stroke: black}.mut text {fill: red; font-style: italic}.mut.extra text {fill: hotpink}.mut line {fill: none; stroke: none}.mut .sym {fill: none; stroke: red}.mut.extra .sym {stroke: hotpink}.node .mut .sym {stroke-width: 1.5px}.tree text, .tree-sequence text {dominant-baseline: central}.plotbox .lab.lft {text-anchor: end}.plotbox .lab.rgt {text-anchor: start}]]></style></defs><g class=\"tree t0\"><g class=\"plotbox\"><g class=\"c2 node n6 p2 root\" transform=\"translate(100 26.8)\"><g class=\"a6 c2 m0 node n4 p1 s0\" transform=\"translate(-40 120.704)\"><g class=\"a4 i0 leaf node n0 p0 sample\" transform=\"translate(-20 20.6961)\"><path class=\"edge\" d=\"M 0 0 V -20.6961 H 20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">0</text></g><g class=\"a4 i1 leaf node n3 p1 sample\" transform=\"translate(20 20.6961)\"><path class=\"edge\" d=\"M 0 0 V -20.6961 H -20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">3</text></g><path class=\"edge\" d=\"M 0 0 V -120.704 H 40\" /><g class=\"mut m0 s0\" transform=\"translate(0 -17.9372)\"><line x1=\"0\" x2=\"0\" y1=\"0\" y2=\"17.9372\" /><path class=\"sym\" d=\"M -3,-3 l 6,6 M -3,3 l 6,-6\" /><text class=\"lab lft\" transform=\"translate(-5 0)\">0</text></g><circle class=\"sym\" cx=\"0\" cy=\"0\" r=\"3\" /><text class=\"lab lft\" transform=\"translate(-3 -7.0)\">4</text></g><g class=\"a6 c2 node n5 p1\" transform=\"translate(40 102.435)\"><g class=\"a5 i0 leaf node n1 p0 sample\" transform=\"translate(-20 38.9647)\"><path class=\"edge\" d=\"M 0 0 V -38.9647 H 20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">1</text></g><g class=\"a5 i1 leaf node n2 p1 sample\" transform=\"translate(20 38.9647)\"><path class=\"edge\" d=\"M 0 0 V -38.9647 H -20\" /><rect class=\"sym\" height=\"6\" width=\"6\" x=\"-3\" y=\"-3\" /><text class=\"lab\" transform=\"translate(0 11)\">2</text></g><path class=\"edge\" d=\"M 0 0 V -102.435 H -40\" /><circle class=\"sym\" cx=\"0\" cy=\"0\" r=\"3\" /><text class=\"lab rgt\" transform=\"translate(3 -7.0)\">5</text></g><circle class=\"sym\" cx=\"0\" cy=\"0\" r=\"3\" /><text class=\"lab\" transform=\"translate(0 -11)\">6</text></g></g></g></svg>'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mts_list = []\n",
    "\n",
    "demography = msprime.Demography()\n",
    "demography.add_population(name=\"A\", initial_size=Na)\n",
    "demography.add_population(name=\"B\", initial_size=Nb)\n",
    "demography.add_population(name=\"C\", initial_size=Nanc)\n",
    "demography.add_population_split(time=Tsplit, derived=[\"A\", \"B\"], ancestral=\"C\")\n",
    "demography.set_migration_rate(source=\"A\",dest=\"B\",rate = M)\n",
    "\n",
    "ts_reps = msprime.sim_ancestry(\n",
    "    samples={\"A\": 1, \"B\": 1}, \n",
    "    demography=demography,\n",
    "    ploidy=2,\n",
    "    recombination_rate=0,\n",
    "    sequence_length=block_length,\n",
    "    discrete_genome=False,\n",
    "    random_seed = 44,\n",
    "    num_replicates = int(L / block_length)\n",
    ")\n",
    "\n",
    "for i, ts in enumerate(ts_reps): \n",
    "    mts_list.append(msprime.sim_mutations(ts, rate=mu, discrete_genome=False, random_seed = 100))\n",
    "    \n",
    "    if i == 0:\n",
    "        first_tree = mts_list[i].first()\n",
    "\n",
    "first_tree.draw_svg()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab707d8a-dfd7-4e85-9c03-afb792373ec8",
   "metadata": {},
   "source": [
    "##### As an example, summarise the mutations on this tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f6d2cedd-5f6f-418e-866f-d7d72912e22b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hetAB mutation? Yes\n",
      "Sums of a branches (sample of tau_a):  294255.05229381763\n",
      "Sums of b branches (sample of tau_b):  294255.05229381763\n"
     ]
    }
   ],
   "source": [
    "tempList = []\n",
    "for node in first_tree.nodes():\n",
    "    leaves = tuple(sorted(list(first_tree.leaves(node))))\n",
    "    mutType = mutTypeDict[leaves]\n",
    "    length = first_tree.get_branch_length(node)\n",
    "    tempList.append([mutType,length,node])\n",
    "if 'hetAB' in [a for (a,b,c) in tempList]:\n",
    "    print(\"hetAB mutation? Yes\")\n",
    "else:\n",
    "    print(\"hetAB mutation? No\")\n",
    "hetA = [[a,b] for a,b,c in tempList if a=='hetA']\n",
    "hetB = [[a,b] for a,b,c in tempList if a=='hetB']\n",
    "branchLengthA = sum([b for a,b in hetA])\n",
    "branchLengthB = sum([b for a,b in hetB])\n",
    "print(\"Sums of a branches (sample of tau_a): \", branchLengthA)\n",
    "print(\"Sums of b branches (sample of tau_b): \", branchLengthB)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ecf1ab8-5f98-4273-b7c9-a0561eaf47b2",
   "metadata": {},
   "source": [
    "#### Now we want to record external branch lengths in three different contexts\n",
    "- all genealogies (idx 0)\n",
    "- incongruent genealogies (idx 1)\n",
    "- incongruent genealogies with hetAB mutations (idx 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "764d81bf-8038-48d6-abcb-75c8009631c8",
   "metadata": {},
   "source": [
    "##### Calculate the store the branch lengths, returning lists for each of the above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "66c03e82-c222-4b89-a9b6-a43de312d757",
   "metadata": {},
   "outputs": [],
   "source": [
    "def record_branch_lengths(mts_list):\n",
    "\n",
    "    a_0_lengths = []\n",
    "    b_0_lengths = []\n",
    "\n",
    "    a_1_lengths = []\n",
    "    b_1_lengths = []\n",
    "\n",
    "    a_2_lengths = []\n",
    "    b_2_lengths = []\n",
    "    \n",
    "    for mts in mts_list:\n",
    "\n",
    "        tree = mts.at(block_length / 2) # assuming no recombination in a block, it doesn't matter where we sample from\n",
    "     \n",
    "        tempList = []\n",
    "        for node in tree.nodes():\n",
    "            leaves = tuple(sorted(list(tree.leaves(node))))\n",
    "            mutType = mutTypeDict[leaves]\n",
    "            length = tree.get_branch_length(node)\n",
    "            tempList.append([mutType,length])\n",
    "    \n",
    "        # calculate the total length of hetA and hetB type branches in this tree\n",
    "        a_length = sum([y for x,y in tempList if x=='hetA'])\n",
    "        b_length = sum([y for x,y in tempList if x=='hetB'])\n",
    "    \n",
    "        # record the hetA and hetB branch lengths for all trees\n",
    "        a_0_lengths.append(a_length)\n",
    "        b_0_lengths.append(b_length)\n",
    "    \n",
    "    \n",
    "        # if the topology is incongruent, \n",
    "        if 'hetAB' in [a for (a,b) in tempList]:\n",
    "            \n",
    "            #store the hetA and hetB branch lengths for this tree conditioned on incongruence\n",
    "            a_1_lengths.append(a_length)\n",
    "            b_1_lengths.append(b_length)\n",
    "    \n",
    "            # given incongruence, check if there is a hetAB mutation present within the block\n",
    "            has_ab_mut = False\n",
    "            muts = tree.mutations()\n",
    "            for mut in muts:\n",
    "                # get the node below the mutation and check the branch type (mutType) of the mutation\n",
    "                node1 = mut.node\n",
    "                leaves = tuple(sorted(list(tree.leaves(node1))))\n",
    "                mutType = mutTypeDict[leaves]\n",
    "                # if the mutation is hetAB\n",
    "                if mutType == 'hetAB':\n",
    "                    has_ab_mut = True\n",
    "                    break\n",
    "                            \n",
    "            # conditioned on a hetAB mutation in the block, record the hetA and hetB branch lengths \n",
    "            if has_ab_mut:\n",
    "                a_2_lengths.append(a_length)\n",
    "                b_2_lengths.append(b_length)\n",
    "\n",
    "    return a_0_lengths, b_0_lengths, a_1_lengths, b_1_lengths, a_2_lengths, b_2_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8ad72c6b-47d7-4cbe-bd92-185ec69820e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_0_lengths, b_0_lengths, a_1_lengths, b_1_lengths, a_2_lengths, b_2_lengths = record_branch_lengths(mts_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8dd21b4-d3a7-411b-876a-081742f39d1c",
   "metadata": {},
   "source": [
    "##### These lists of external branch lengths can be summarised in terms of asymmetry (A_i and A_m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b3c35226-eefb-486a-9db4-d18efaad5ce0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_asymmetry(a_list, b_list):\n",
    "    \n",
    "    a_total = sum(a_list)\n",
    "    b_total = sum(b_list)\n",
    "\n",
    "    asymm = (a_total - b_total) / (a_total + b_total)\n",
    "\n",
    "    return asymm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a474d984-864c-4e57-bd85-27fc75e26763",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A for all genealogies:  \t\t\t\t 0.12565851306521508\n",
      "A_i for incongruent genealogies:  \t\t\t 0.034984076412470325\n",
      "A_m for incongruent genealogies with hetAB mutations:  \t 0.0\n"
     ]
    }
   ],
   "source": [
    "print(\"A for all genealogies: \", \"\\t\\t\\t\\t\", calculate_asymmetry(a_0_lengths, b_0_lengths))\n",
    "print(\"A_i for incongruent genealogies: \", \"\\t\\t\\t\", calculate_asymmetry(a_1_lengths, b_1_lengths))\n",
    "print(\"A_m for incongruent genealogies with hetAB mutations: \", \"\\t\", calculate_asymmetry(a_2_lengths, b_2_lengths))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a46fa524-1b7f-4fb0-b17f-3ecd82ba0986",
   "metadata": {},
   "source": [
    "##### We can use the code above to obtain point estimates of A_m. To obtain the variance of the estimate, and therefore 95%CIs, we can perform block jackknife resampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d5a6f5c1-2c78-43a5-b2d9-3d93c1de1c6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def block_jackknife(a_list, b_list, jk_block):\n",
    "    total_jk_blocks = int((len(a_list) - (len(a_list) % jk_block)) / jk_block)\n",
    "    if total_jk_blocks < 2:\n",
    "        #print(\"not enough blocks for jackknife\")\n",
    "        overall_estimate = calculate_asymmetry(a_list, b_list)\n",
    "        return [-1, overall_estimate, 1]\n",
    "    trimmed_a_list = a_list[:total_jk_blocks * jk_block]\n",
    "    trimmed_b_list = b_list[:total_jk_blocks * jk_block]\n",
    "    blocked_a_list = np.mean(trimmed_a_list.reshape(-1, jk_block), axis=1)\n",
    "    blocked_b_list = np.mean(trimmed_b_list.reshape(-1, jk_block), axis=1)\n",
    "    overall_estimate = calculate_asymmetry(blocked_a_list, blocked_b_list)\n",
    "    n = total_jk_blocks\n",
    "    pseudo_estimates = []\n",
    "    for i in range(0, n):\n",
    "        retained_a_blocks = np.delete(blocked_a_list, i)\n",
    "        retained_b_blocks = np.delete(blocked_b_list, i)\n",
    "        retained_blocks_estimate = calculate_asymmetry(retained_a_blocks, retained_b_blocks)\n",
    "        pseudo_estimate = (n * overall_estimate) - ((n - 1) * retained_blocks_estimate)\n",
    "        pseudo_estimates.append(pseudo_estimate)\n",
    "    pseudovalue_variance = sum([(e - overall_estimate)**2 for e in pseudo_estimates]) / (n - 1)\n",
    "    lower_CI = overall_estimate - (1.96 * math.sqrt(pseudovalue_variance / n))\n",
    "    higher_CI = overall_estimate + (1.96 * math.sqrt(pseudovalue_variance / n))\n",
    "    return [lower_CI, overall_estimate, higher_CI]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac709d2d-e9ae-497d-9df0-a48c942f8029",
   "metadata": {},
   "source": [
    "##### Now we can simulate many blocks (e.g. 500,000) to generate approximate expectations for A, A_i and A_m. Simulations do not need to follow the five parameter IM model and can instead include bidirectional migration or other processes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1fd813c5-2aff-487e-abc9-4ff38978ae8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A for all genealogies in block:  \t\t\t\t 0.16575491363320535\n",
      "A for incongruent genealogies in block:  \t\t\t 0.12537885856889383\n",
      "A for incongruent genealogies with hetAB mutations in block:  \t 0.049175283546627244\n",
      "................................................................\n",
      "................. Running estimates so far .....................\n",
      "[-1, 0.16575491363320535, 1]\n",
      "[-1, 0.12537885856889383, 1]\n",
      "[-1, 0.049175283546627244, 1]\n",
      "----------------------------------------------------------------\n",
      "A for all genealogies in block:  \t\t\t\t 0.15517415536226334\n",
      "A for incongruent genealogies in block:  \t\t\t 0.1053863576755524\n",
      "A for incongruent genealogies with hetAB mutations in block:  \t 0.01340208064969981\n",
      "................................................................\n",
      "................. Running estimates so far .....................\n",
      "[0.15004762120098866, 0.16041717969893562, 0.1707867381968826]\n",
      "[0.09571292923723743, 0.11530615316453796, 0.1348993770918385]\n",
      "[-0.004731331530162596, 0.030372379136541183, 0.06547608980324496]\n",
      "----------------------------------------------------------------\n",
      "A for all genealogies in block:  \t\t\t\t 0.16596557028491848\n",
      "A for incongruent genealogies in block:  \t\t\t 0.12873989814207112\n",
      "A for incongruent genealogies with hetAB mutations in block:  \t 0.05177586247940271\n",
      "................................................................\n",
      "................. Running estimates so far .....................\n",
      "[0.15527650905242235, 0.16228404451756964, 0.16929157998271693]\n",
      "[0.10554636668194273, 0.11987287051632511, 0.1341993743507075]\n",
      "[0.012700379955916035, 0.037740204165414716, 0.0627800283749134]\n",
      "----------------------------------------------------------------\n",
      "A for all genealogies in block:  \t\t\t\t 0.15631735864983626\n",
      "A for incongruent genealogies in block:  \t\t\t 0.1211676513546452\n",
      "A for incongruent genealogies with hetAB mutations in block:  \t 0.036926442513128585\n",
      "................................................................\n",
      "................. Running estimates so far .....................\n",
      "[0.155036379280318, 0.16077613634590407, 0.16651589341149015]\n",
      "[0.11009421834735646, 0.12019942759347783, 0.13030463683959922]\n",
      "[0.01970336742533234, 0.03754198977000136, 0.05538061211467038]\n",
      "----------------------------------------------------------------\n",
      "A for all genealogies in block:  \t\t\t\t 0.15946955849307717\n",
      "A for incongruent genealogies in block:  \t\t\t 0.13807151807275092\n",
      "A for incongruent genealogies with hetAB mutations in block:  \t 0.06161993217438934\n",
      "................................................................\n",
      "................. Running estimates so far .....................\n",
      "[0.15605723873185845, 0.16051087151045193, 0.1649645042890454]\n",
      "[0.11329938851229546, 0.12384033072090596, 0.13438127292951646]\n",
      "[0.025513080587346228, 0.04228932665881584, 0.05906557273028545]\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "mu = 1e-8\n",
    "r = 0\n",
    "Na = 1e5\n",
    "Nb = 1e5\n",
    "Nanc = 1e5\n",
    "Ma = 1 / (2 * Na)\n",
    "Mb = 0 / (2 * Nb)\n",
    "Tsplit = 2e5 #generations\n",
    "\n",
    "L = int(1e6)\n",
    "block_length = 200\n",
    "\n",
    "a_reps_0_list = []\n",
    "b_reps_0_list = []\n",
    "a_reps_1_list = []\n",
    "b_reps_1_list = []\n",
    "a_reps_2_list = []\n",
    "b_reps_2_list = []\n",
    "\n",
    "for i in range(0, 5): # actual sims include 100 tree sequences\n",
    "\n",
    "    mts_list = []\n",
    "\n",
    "    demography = msprime.Demography()\n",
    "    demography.add_population(name=\"A\", initial_size=Na)\n",
    "    demography.add_population(name=\"B\", initial_size=Nb)\n",
    "    demography.add_population(name=\"C\", initial_size=Nanc)\n",
    "    demography.add_population_split(time=Tsplit, derived=[\"A\", \"B\"], ancestral=\"C\")\n",
    "    demography.set_migration_rate(source=\"A\",dest=\"B\",rate = Ma)\n",
    "    demography.set_migration_rate(source=\"B\",dest=\"A\",rate = Mb)\n",
    "    \n",
    "    ts_reps = msprime.sim_ancestry(\n",
    "        samples={\"A\": 1, \"B\": 1}, \n",
    "        demography=demography,\n",
    "        ploidy=2,\n",
    "        recombination_rate=r,\n",
    "        sequence_length=block_length,\n",
    "        discrete_genome=False,\n",
    "        random_seed = 44+i,\n",
    "        num_replicates = int(L / block_length)\n",
    "    )\n",
    "    \n",
    "    for i, ts in enumerate(ts_reps): \n",
    "        mts_list.append(msprime.sim_mutations(ts, rate=mu, discrete_genome=False, random_seed = 44+i))\n",
    "\n",
    "    a_0_lengths, b_0_lengths, a_1_lengths, b_1_lengths, a_2_lengths, b_2_lengths = record_branch_lengths(mts_list)\n",
    "\n",
    "    a_reps_0_list.append(sum(a_0_lengths)) # this groups replicates from the same mts together\n",
    "    b_reps_0_list.append(sum(b_0_lengths)) # the 100 jackknife blocks will be weighted equally, which is only appropriate for large samples\n",
    "    a_reps_1_list.append(sum(a_1_lengths))\n",
    "    b_reps_1_list.append(sum(b_1_lengths))\n",
    "    a_reps_2_list.append(sum(a_2_lengths))\n",
    "    b_reps_2_list.append(sum(b_2_lengths))\n",
    "\n",
    "    print(\"A for all genealogies in block: \", \"\\t\\t\\t\\t\", calculate_asymmetry(a_0_lengths, b_0_lengths))\n",
    "    print(\"A for incongruent genealogies in block: \", \"\\t\\t\\t\", calculate_asymmetry(a_1_lengths, b_1_lengths))\n",
    "    print(\"A for incongruent genealogies with hetAB mutations in block: \", \"\\t\", calculate_asymmetry(a_2_lengths, b_2_lengths))\n",
    "\n",
    "    print(\"................................................................\")\n",
    "\n",
    "    print(\"................. Running estimates so far .....................\")\n",
    "\n",
    "    print(block_jackknife(np.array(a_reps_0_list), np.array(b_reps_0_list), 1)) # ensures 100 (unweighted) jackknife blocks \n",
    "    print(block_jackknife(np.array(a_reps_1_list), np.array(b_reps_1_list), 1))\n",
    "    print(block_jackknife(np.array(a_reps_2_list), np.array(b_reps_2_list), 1))\n",
    "\n",
    "    print(\"----------------------------------------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1ef4679-8191-4b27-a5ec-3212159a5495",
   "metadata": {},
   "source": [
    "##### We also simulate longer sequences with recombination. These sequences can be split into blocks of size 2^n, allowing for (non-independent) estimates of A_m across a range of block sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f493b460-358e-4eb0-a4ae-518147162630",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100\n",
      "jack-knifing...\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS4AAAEcCAYAAAB58lnlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8g+/7EAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAl8UlEQVR4nO3dfVBU1/0/8Dcs+xARtiqGh4grdirbDYTqEgIkfGmmdhFrpqk2gjbUdpppmKYTV9txljAdje2EtOlDmgzECUOj/IFixDSZiRPdzEQkAWNlIJMEI0yJhdjdEmzdFRwXXc/vD8P9ue7ycBfY3bu8XzN3ZvfsOfd8buR+cvbuvefECCEEiIgUJDbcARARycXERUSKw8RFRIrDxEVEisPERUSKw8RFRIrDxEVEisPERUSKw8RFRIrDxEVEihNU4qqrq0NGRgZ0Oh3MZjPa2tomrHv06FF897vfxdKlS5GYmIiCggIcP37cr15LSwtMJhO0Wi1MJhPeeOONGfVLRNFLduJqbm6G1WpFdXU1urq6UFRUhNLSUgwMDASsf+rUKXz3u9/FsWPH0NnZiYcffhiPPPIIurq6pDodHR0oKytDRUUFPvroI1RUVGDz5s348MMPg+6XiKKYkCkvL09UVlb6lBmNRmGz2aa9D5PJJJ599lnp/ebNm8W6det86pSUlIjy8vJZ7ZeIokOcnCQ3NjaGzs5O2Gw2n3KLxYL29vZp7ePmzZu4cuUKFi9eLJV1dHRgx44dPvVKSkrw4osvzqhfj8cDj8fj0/d///tfLFmyBDExMdOKl4hCQwiBK1euIC0tDbGxk38ZlJW4hoeH4fV6kZyc7FOenJwMp9M5rX386U9/wujoKDZv3iyVOZ3OSfcZbL81NTV49tlnpxUXEUWGwcFBLFu2bNI6shLXuDtHK0KIaY1gDh48iD179uDNN9/E3XffLXufcvutqqrCzp07pfculwvLly/H4OAgEhMTp4yXKJqMjo4iLS0NAPDvf/8b8fHxEdMGANxuN9LT05GQkDBlXVmJKykpCSqVym+UMzQ05DcaulNzczN+9rOf4fXXX8fatWt9PktJSZl0n8H2q9VqodVq/coTExOZuGjeUalU0uvExMRpJZRQtbnddAZBsn5V1Gg0MJvNsNvtPuV2ux2FhYUTtjt48CB+8pOfoKmpCd/73vf8Pi8oKPDb54kTJ6R9BtsvEUUpuVfzDx06JNRqtWhoaBA9PT3CarWK+Ph4ceHCBSGEEDabTVRUVEj1m5qaRFxcnKitrRUOh0PaLl++LNX54IMPhEqlEs8//7w4d+6ceP7550VcXJw4ffr0tPudDpfLJQAIl8sl97CJFG9kZEQAEADEyMhIRLURQt75KTtxCSFEbW2tMBgMQqPRiDVr1ojW1lbps23btoni4mLpfXFxsXQQt2/btm3z2efrr78uMjMzhVqtFkajUbS0tMjqdzqYuGg+i6bEFSPE/Fksw+12Q6/Xw+Vy8RoXzTujo6NYuHAhAGBkZGTaF9pD0QaQd37yWUUiUhwmLiIFGh0dRUxMDGJiYjA6OhrucEKOiYuIFIeJi4gUh4mLiBSHiYuIFIeJi4gUh4mLiBSHiYuIFIeJi4gm1NvbK73Ozc3F+fPn56SNXExcRDShrVu3Sq97e3uxcePGOWkjFxMXEQXk9Xrx2WefSe9v3ryJnp4eeL3eWW0TDCYuonlC7lc4lUoFo9EovY+NjYXJZPKZKHA22gSDiYtongjmK1xTU5P0etWqVTh69OictJGL09oQKZDcqWO8Xi/i4vxnar9x48akoyFOa0NEYROqr3ChwsRFNE+E4itcqAS1PBkRza5gv17JsWrVKun12bNn56SPUAlqxFVXV4eMjAzodDqYzWa0tbVNWNfhcGDr1q3IzMxEbGwsrFarX51vf/vb0qRot2+3rwi0Z88ev89TUlKCCZ+IFE524mpubobVakV1dTW6urpQVFSE0tJSDAwMBKzv8XiwdOlSVFdXIycnJ2Cdo0ePwuFwSNsnn3wClUqFxx57zKfevffe61Pv448/lhs+EUUB2V8V//znP+NnP/sZnnjiCQDAiy++iOPHj+OVV15BTU2NX/0VK1bgr3/9KwDgb3/7W8B9Ll682Of9oUOHsGDBAr/EFRcXx1EWEckbcY2NjaGzsxMWi8Wn3GKxoL29fdaCamhoQHl5ud938L6+PqSlpSEjIwPl5eXo7++fdD8ejwdut9tnIyLlk5W4hoeH4fV6/Za9T05OhtPpnJWAzpw5g08++UQa0Y174IEH0NjYiOPHj6O+vh5OpxOFhYW4dOnShPuqqamBXq+XtvT09FmJkYjCK6iL8zExMT7vhRB+ZcFqaGhAVlYW8vLyfMpLS0uxadMmZGdnY+3atXj77bcBAAcOHJhwX1VVVXC5XNI2ODg4KzESUXjJusaVlJQElUrlN7oaGhryG4UF4+rVqzh06BD27t07Zd34+HhkZ2ejr69vwjparRZarXbGcRFRZJE14tJoNDCbzbDb7T7ldrsdhYWFMw7m8OHD8Hg8ePzxx6es6/F4cO7cOaSmps64XyKlCcWcV5FM9q+KO3fuREVFBXJzc1FQUIBXX30VAwMDqKysBHDr69nFixfR2Ngotenu7gZw68a6L7/8Et3d3dBoNDCZTD77bmhowKOPPoolS5b49fvrX/8ajzzyCJYvX46hoSH87ne/g9vtxrZt2+QeApHiBXpg+tNPPw1jRKElO3GVlZXh0qVL2Lt3LxwOB7KysnDs2DEYDAYAt244vfOertWrV0uvOzs70dTUBIPBgAsXLkjlvb29eP/993HixImA/X7xxRfYsmULhoeHsXTpUuTn5+P06dNSv0TzxWRzXin12UO5ODsEUQSQ+8jPN7/5TSl5xcbGwmg0TjniCtVMD5wdgogCiqYHpoPBh6yJFCiaHpgOBkdcRKQ4TFxEpDhMXESkOExcRKQ4TFxEpDhMXESzbHR0VJqld3R0NNzhRCUmLiJSHCYuIlIcJi4iUhwmLiJSHCYuoggQqfNrxcfHQwgBIcS0HysKpo1cTFxEESDQ/Fo0MSYuojCbbH4tCoyJiyjMVCoVjEaj9D42NhYmk2neTAoYDCYuoggQivm1QnHtKVQ4HxdRBJjv82vJFdSIq66uDhkZGdDpdDCbzWhra5uwrsPhwNatW5GZmYnY2FhYrVa/Ovv375cekbh9u3btWtD9ElH0kp24mpubYbVaUV1dja6uLhQVFaG0tNRvgYxxHo8HS5cuRXV1NXJycibcb2JiIhwOh8+m0+mC7peIopiQKS8vT1RWVvqUGY1GYbPZpmxbXFwstm/f7lf+2muvCb1eP+v9Xrt2TbhcLmkbHBwUAITL5ZoyVqJgjYyMCAACgBgZGZmTNsH0EelcLte0z09ZI66xsTF0dnbCYrH4lFssFrS3t88ogY6MjMBgMGDZsmXYsGEDurq6ZtxvTU0N9Hq9tKWnp88oRiKKDLIS1/DwMLxeL5KTk33Kk5OT4XQ6gw7CaDRi//79eOutt3Dw4EHodDo8+OCD6Ovrm1G/VVVVcLlc0jY4OBh0jEQUOYL6VTEmJsbnvRDCr0yO/Px85OfnS+8ffPBBrFmzBi+//DJeeumloPvVarXQarVBx0UUqcZvbZivZI24kpKSoFKp/EY5Q0NDfqOhGQUVG4v7779fGnGFql8iUgZZiUuj0cBsNsNut/uU2+12FBYWzlpQQgh0d3cjNTU1pP0SkTLI/qq4c+dOVFRUIDc3FwUFBXj11VcxMDCAyspKALeuK128eBGNjY1Sm+7ubgC3LsB/+eWX6O7uhkajgclkAgA8++yzyM/Pxze+8Q243W689NJL6O7uRm1t7bT7JaL5Q3biKisrw6VLl7B37144HA5kZWXh2LFjMBgMAG7dcHrnvVWrV6+WXnd2dqKpqQkGgwEXLlwAAFy+fBk///nP4XQ6odfrsXr1apw6dQp5eXnT7peI5o8YMY+u8Lndbuj1erhcLiQmJoY7HIpSo6OjWLhwIYBb3zKm8/hOMG2ijZzzkw9ZE5HiMHERzbJInc00mjBxEc0yzmY695i4iCYhd3FXzmYaGkxcRLOIs5mGBhMX0SwLxWym8x1nQCWaZZzNdO5xxEVEisPERUSKw8RFRIrDxEVEisOL80QRYL5PDCgXR1xEpDhMXESkOExcRKQ4TFxEpDhMXESkOExcRKQ4QSWuuro6ZGRkQKfTwWw2o62tbcK6DocDW7duRWZmJmJjY2G1Wv3q1NfXo6ioCIsWLcKiRYuwdu1anDlzxqfOnj17pOlFxreUlJRgwicihZOduJqbm2G1WlFdXY2uri4UFRWhtLTUb4GMcR6PB0uXLkV1dTVycnIC1jl58iS2bNmC9957Dx0dHVi+fDksFgsuXrzoU+/ee++Fw+GQto8//lhu+EQUDYRMeXl5orKy0qfMaDQKm802Zdvi4mKxffv2KevduHFDJCQkiAMHDkhlu3fvFjk5ObJivXbtmnC5XNI2ODgoAAiXyyVrPzR/jYyMCAACgBgZGZmzNiSEy+Wa9vkpa8Q1NjaGzs5OWCwWn3KLxYL29vZZSqXA1atXcf36dSxevNinvK+vD2lpacjIyEB5eTn6+/sn3U9NTQ30er20paenz1qMRBQ+shLX8PAwvF6v37L3ycnJcDqdsxaUzWbDPffcg7Vr10plDzzwABobG3H8+HHU19fD6XSisLAQly5dmnA/VVVVcLlc0jY4ODhrMRJR+AT1rGJMTIzPeyGEX1mw/vCHP+DgwYM4efIkdDqdVF5aWiq9zs7ORkFBAb7+9a/jwIED2LlzZ8B9abVaaLXaWYmLiCKHrMSVlJQElUrlN7oaGhryG4UF449//COee+45vPvuu7jvvvsmrRsfH4/s7Gz09fXNuF8iUhZZXxU1Gg3MZjPsdrtPud1uR2Fh4YwCeeGFF/Db3/4W77zzDnJzc6es7/F4cO7cOaSmps6oXyJSHtlfFXfu3ImKigrk5uaioKAAr776KgYGBlBZWQng1nWlixcvorGxUWrT3d0N4NbS4l9++SW6u7uh0WhgMpkA3Pp6+Jvf/AZNTU1YsWKFNKJbuHChtCz5r3/9azzyyCNYvnw5hoaG8Lvf/Q5utxvbtm2b0X8AIlKgYH62rK2tFQaDQWg0GrFmzRrR2toqfbZt2zZRXFzsUx9f/TR8+2YwGKTPDQZDwDq7d++W6pSVlYnU1FShVqtFWlqa2Lhxo/j0009lxS3n51YiIXg7RCjJOT9jhJg/s5e53W7o9Xq4XC4kJiaGOxxSgNHRUWnUPzIywhV75pCc85PPKhJNore3V3qdm5uL8+fPhzEaGsfERTSJrVu3Sq97e3uxcePGMEZD45i4iCbg9Xrx2WefSe9v3ryJnp4eeL3eMEZFABMX0YRUKhWMRqP0PjY2FiaTCSqVKoxREcDERTSppqYm6fWqVatw9OjRMEZD47g8GdEkVq1aJb0+e/Ysf1WMEBxxEZHiMHERkeIwcRGR4jBxEZHiMHERkeIwcRGR4jBx0bwxOjoqLW03Ojoa7nBoBpi4iEhxmLiISHGYuIhIcZi4iEhxgkpcdXV1yMjIgE6ng9lsRltb24R1HQ4Htm7diszMTMTGxsJqtQas19LSApPJBK1WC5PJhDfeeGNG/RJR9JKduJqbm2G1WlFdXY2uri4UFRWhtLQUAwMDAet7PB4sXboU1dXVyMnJCVino6MDZWVlqKiowEcffYSKigps3rwZH374YdD9ElEUkzuhfV5enqisrPQpMxqNwmazTdm2uLhYbN++3a988+bNYt26dT5lJSUlory8fFb6HcfFMuY3LnwR2eScn7JGXGNjY+js7ITFYvEpt1gsaG9vDzp5dnR0+O2zpKRE2mew/Xo8Hrjdbp+NiJRPVuIaHh6G1+v1W7U6OTnZb3VrOZxO56T7DLbfmpoa6PV6aUtPTw86RiKKHEFdnI+JifF5L4TwK5uLfcrtt6qqCi6XS9oGBwdnFCMRRQZZM6AmJSVBpVL5jXKGhob8RkNypKSkTLrPYPvVarXQarVBx0VEkUnWiEuj0cBsNsNut/uU2+12FBYWBh1EQUGB3z5PnDgh7XOu+iUiZZI95/zOnTtRUVGB3NxcFBQU4NVXX8XAwAAqKysB3Pp6dvHiRTQ2Nkpturu7AdxaCfjLL79Ed3c3NBoNTCYTAGD79u34v//7P/z+97/H97//fbz55pt499138f7770+7XyKaR4L52bK2tlYYDAah0WjEmjVrRGtrq/TZtm3bRHFxsU99fPVz8u2bwWDwqfP666+LzMxMoVarhdFoFC0tLbL6nQ7eDjG/8XaIyCbn/IwRQogw5cyQc7vd0Ov1cLlcSExMDHc4FGKjo6NYuHAhgFuj/+ms2BNMGwqOnPOTzyoSkeIwcRGR4jBxEZHicCVroknEx8djHl0GVgyOuIhIcZi4iEhxmLiISHGYuGje6O3tlV7n5ubi/PnzYYyGZoKJixQpmDUSt27dKr3u7e3Fxo0b5yo8mmNMXDQveL1efPbZZ9L7mzdvoqenB16vN4xRUbCYuGheUKlUMBqN0vvY2FiYTCaoVKowRkXBYuKieaOpqUl6vWrVKhw9ejSM0dBM8AZUmjdWrVolvT579iwfmFYwjriISHGYuIhIcZi4iEhxmLiISHGYuIhIcYJKXHV1dcjIyIBOp4PZbEZbW9uk9VtbW2E2m6HT6bBy5Urs27fP5/Nvf/vb0l3Qt2/f+973pDp79uzx+zwlJSWY8IlI4WQnrubmZlitVlRXV6OrqwtFRUUoLS3FwMBAwPqff/451q9fj6KiInR1deGZZ57B008/jZaWFqnO0aNH4XA4pO2TTz6BSqXCY4895rOve++916fexx9/LDd8IooGclfiyMvLE5WVlT5lRqNR2Gy2gPV37doljEajT9mTTz4p8vPzJ+zjL3/5i0hISPBZVWX37t0iJydHbrg+uMpP9OCKPdFHzvkpa8Q1NjaGzs5OWCwWn3KLxYL29vaAbTo6Ovzql5SU4OzZs7h+/XrANg0NDSgvL/e7QbCvrw9paWnIyMhAeXk5+vv7J43X4/HA7Xb7bESkfLIS1/DwMLxer9+y98nJyXA6nQHbOJ3OgPVv3LiB4eFhv/pnzpzBJ598gieeeMKn/IEHHkBjYyOOHz+O+vp6OJ1OFBYW4tKlSxPGW1NTA71eL23p6enTPVQiimBBXZyPiYnxeS+E8Cubqn6gcuDWaCsrKwt5eXk+5aWlpdi0aROys7Oxdu1avP322wCAAwcOTNhvVVUVXC6XtA0ODk5+YBQWwUxRQ/ObrGcVk5KSoFKp/EZXQ0NDfqOqcSkpKQHrx8XFYcmSJT7lV69exaFDh7B3794pY4mPj0d2djb6+vomrKPVaqHVaqfcFxEpi6wRl0ajgdlsht1u9ym32+0oLCwM2KagoMCv/okTJ5Cbmwu1Wu1TfvjwYXg8Hjz++ONTxuLxeHDu3DmkpqbKOQQiigZyr/wfOnRIqNVq0dDQIHp6eoTVahXx8fHiwoULQgghbDabqKiokOr39/eLBQsWiB07doienh7R0NAg1Gq1OHLkiN++H3roIVFWVhaw31/96lfi5MmTor+/X5w+fVps2LBBJCQkSP1OB39VjEyh+oWQvypGNjnnp+xpbcrKynDp0iXs3bsXDocDWVlZOHbsGAwGAwDA4XD43NOVkZGBY8eOYceOHaitrUVaWhpeeuklbNq0yWe/vb29eP/993HixImA/X7xxRfYsmULhoeHsXTpUuTn5+P06dNSv0Q0f8QIMX9Wu3S73dDr9XC5XEhMTAx3OPSV0dFRLFy4EAAwMjIyrXmyQtWGQkfO+clnFYlIcZi4iEhxmLiISHGYuEiRuLjr/MbERYrExV3nN67yQ4oz2eKuk62TGB8fj3n0I3pU44iLFIeLuxITFykSF3ed3/hVkRSJi7vObxxx0aziFDUUCkxcRKQ4TFxEpDhMXESkOExcRKQ4TFxEpDhMXESkOExcRKQ4TFwUdpzpgeQKKnHV1dUhIyMDOp0OZrMZbW1tk9ZvbW2F2WyGTqfDypUrsW/fPp/P9+/fL920ePt27dq1GfVLMxOqm0k50wPJJTtxNTc3w2q1orq6Gl1dXSgqKkJpaanPAhm3+/zzz7F+/XoUFRWhq6sLzzzzDJ5++mm0tLT41EtMTITD4fDZdDpd0P2SMkw20wPRhOQuIZSXlycqKyt9yoxGo7DZbAHr79q1SxiNRp+yJ598UuTn50vvX3vtNaHX62e130C4PJk8oVoCzGg0Sm1iY2OFyWSak34ossk5P2WNuMbGxtDZ2QmLxeJTbrFY0N7eHrBNR0eHX/2SkhKcPXsW169fl8pGRkZgMBiwbNkybNiwAV1dXTPqF7i1aKzb7fbZKPJwpgeSS1biGh4ehtfrRXJysk95cnIynE5nwDZOpzNg/Rs3bmB4eBgAYDQasX//frz11ls4ePAgdDodHnzwQfT19QXdLwDU1NRAr9dLW3p6upzDpRC5c6aHzMzMMEZDShDUxfmYmBif90IIv7Kp6t9enp+fj8cffxw5OTkoKirC4cOHsWrVKrz88ssz6reqqgoul0vaBgcHpz44Iop4subjSkpKgkql8hvlDA0N+Y2GxqWkpASsHxcXhyVLlgRsExsbi/vvv18acQXTLwBotVpotdopj2s+4GKoFE1kjbg0Gg3MZjPsdrtPud1uR2FhYcA2BQUFfvVPnDiB3NxcqNXqgG2EEOju7kZqamrQ/UYzznlF857cK/+HDh0SarVaNDQ0iJ6eHmG1WkV8fLy4cOGCEEIIm80mKioqpPr9/f1iwYIFYseOHaKnp0c0NDQItVotjhw5ItXZs2ePeOedd8Q///lP0dXVJX7605+KuLg48eGHH0673+mIxF8VQ/XLHdtQpJNzfsqeurmsrAyXLl3C3r174XA4kJWVhWPHjsFgMAAAHA6Hz71VGRkZOHbsGHbs2IHa2lqkpaXhpZdewqZNm6Q6ly9fxs9//nM4nU7o9XqsXr0ap06dQl5e3rT7JaL5I0aI+bNek9vthl6vh8vlQmJi4qzvP5jrSGwTujYU2eScn3xWcRK8lkQUmZi4aFbxgWkKBSYumlWhemB6fFVqIQS/Js5DTFw0a/jANIUKExfNGpVKBaPRKL2PjY2FyWSCSqUKY1QUjZi4aFbxgWkKBdn3cRFN5s4Hpnn9ieYCR1xEpDhMXDQh3tpAkYqJiybEueApUjFxUUC8tYEiGRMXBcRbGyiSMXHNE8Fcr+KtDRSpmLjCLJiEEkybYK5XcS54ilRMXLMoVAlFbhter6Jow8Q1CbmJKBQJJZg2kX69ig9Mk1xMXJOQk4hClVCCTUK8XkXRhIlrAnITUSgTSjBteL2KoklQiauurg4ZGRnQ6XQwm81oa2ubtH5rayvMZjN0Oh1WrlyJffv2+XxeX1+PoqIiLFq0CIsWLcLatWtx5swZnzp79uyRZiMd31JSUoIJf1qCSUShSihMQjTfyU5czc3NsFqtqK6uRldXF4qKilBaWuqzQMbtPv/8c6xfvx5FRUXo6urCM888g6effhotLS1SnZMnT2LLli1477330NHRgeXLl8NiseDixYs++7r33nvhcDik7eOPP5YbvixyExETClGIyF1CKC8vT1RWVvqUGY1GYbPZAtbftWuXMBqNPmVPPvmkyM/Pn7CPGzduiISEBHHgwAGpbPfu3SInJ0duuD7kLk8mdwmsSF6aK5LbEAkh7/yUNeIaGxtDZ2cnLBaLT7nFYkF7e3vANh0dHX71S0pKcPbsWVy/fj1gm6tXr+L69etYvHixT3lfXx/S0tKQkZGB8vJy9Pf3Txqvx+OB2+322YhI+WQlruHhYXi9Xr9l75OTk+F0OgO2cTqdAevfuHEDw8PDAdvYbDbcc889WLt2rVT2wAMPoLGxEcePH0d9fT2cTicKCwtx6dKlCeOtqamBXq+XtvT09OkeKhFFsKAuzsfExPi8F0L4lU1VP1A5APzhD3/AwYMHcfToUeh0Oqm8tLQUmzZtQnZ2NtauXYu3334bAHDgwIEJ+62qqoLL5ZK2wcHBqQ+OiCKerBlQk5KSoFKp/EZXQ0NDfqOqcSkpKQHrx8XFYcmSJT7lf/zjH/Hcc8/h3XffxX333TdpLPHx8cjOzkZfX9+EdbRaLbRa7aT7ISLlkTXi0mg0MJvNsNvtPuV2ux2FhYUB2xQUFPjVP3HiBHJzc6FWq6WyF154Ab/97W/xzjvvIDc3d8pYPB4Pzp07h9TUVDmHQETRQO6V/0OHDgm1Wi0aGhpET0+PsFqtIj4+Xly4cEEIIYTNZhMVFRVS/f7+frFgwQKxY8cO0dPTIxoaGoRarRZHjhyR6vz+978XGo1GHDlyRDgcDmm7cuWKVOdXv/qVOHnypOjv7xenT58WGzZsEAkJCVK/08FfFfkLIUUuOeen7MQlhBC1tbXCYDAIjUYj1qxZI1pbW6XPtm3bJoqLi33qnzx5UqxevVpoNBqxYsUK8corr/h8bjAYpBPk9m337t1SnbKyMpGamirUarVIS0sTGzduFJ9++qmsuJm4mLgocsk5P2OE+OpK+Tzgdruh1+vhcrmQmJg4Zf3R0VEsXLgQADAyMjLlA8By60djG6JgyTk/+awiESkOExcRKQ4XhKUJjc+TRRRpOOIiIsXhiGsWBTNCCVUbomjCxDVPMNlRNGHimgRPdqLIxGtcRKQ4TFxEpDhMXESkOExcRKQ4TFxEpDhMXESkOExcRKQ4TFxEpDhMXESkOExcRKQ4TFxEpDhBJa66ujpkZGRAp9PBbDajra1t0vqtra0wm83Q6XRYuXIl9u3b51enpaUFJpMJWq0WJpMJb7zxxoz7JaIoJXdC+/FVfurr60VPT4/Yvn27iI+PF//6178C1h9f5Wf79u2ip6dH1NfX+63y097eLlQqlXjuuefEuXPnxHPPPSfi4uLE6dOng+43ELmLZRBR6MzpKj95eXmisrLSp8xoNAqbzRaw/q5du4TRaPQpe/LJJ0V+fr70fvPmzWLdunU+dUpKSkR5eXnQ/QbCxEUUueScn7KmtRkbG0NnZydsNptPucViQXt7e8A2HR0dsFgsPmUlJSVoaGjA9evXoVar0dHRgR07dvjVefHFF4PuF7i1aKzH45Heu1wuALdWEyGiyDJ+XoppTCUlK3ENDw/D6/UiOTnZpzw5ORlOpzNgG6fTGbD+jRs3MDw8jNTU1AnrjO8zmH4BoKamBs8++6xfeXp6+sQHSURhdeXKFej1+knrBDWRYExMjM97IYRf2VT17yyfzj7l9ltVVYWdO3dK7y9fvgyDwYCBgYEp/8NEOrfbjfT0dAwODk5rjchIx+OJXKE6FiEErly5grS0tCnrykpcSUlJUKlUfqOcoaEhv9HQuJSUlID14+LisGTJkknrjO8zmH4BQKvVQqvV+pXr9XrF/zGNS0xMjJpjAXg8kSwUxzLdAYWs2yE0Gg3MZjPsdrtPud1uR2FhYcA2BQUFfvVPnDiB3NxcqNXqSeuM7zOYfokoism98j9+W0JDQ4Po6ekRVqtVxMfHiwsXLgghhLDZbKKiokKqP347xI4dO0RPT49oaGjwux3igw8+ECqVSjz//PPi3Llz4vnnn5/wdoiJ+p2OaPpVMZqORQgeTySLxGORnbiEEKK2tlYYDAah0WjEmjVrRGtrq/TZtm3bRHFxsU/9kydPitWrVwuNRiNWrFghXnnlFb99vv766yIzM1Oo1WphNBpFS0uLrH6n49q1a2L37t3i2rVrstpFomg6FiF4PJEsEo8lRgguY0NEysJnFYlIcZi4iEhxmLiISHGYuIhIcZi4iEhxmLi+8oMf/ACLFi3CD3/4w3CHIsuVK1dw//3341vf+hays7NRX18f7pBm5Pz58/jWt74lbXfddRf+/ve/hzssWSb7W1Li39lkMcfFxUn/Vk888UTIYuLtEF957733MDIyggMHDuDIkSPhDmfavF4vPB4PFixYgKtXryIrKwv/+Mc/pMeplGxkZAQrVqzAv/71L8THx4c7nGmb7G9JiX9nk8WclJSE4eHhkMfEEddXHn74YSQkJIQ7DNlUKhUWLFgAALh27Rq8Xu+0pgVRgrfeegvf+c53FJW0gMn/lpT4dxaJMUdF4qqpqcH999+PhIQE3H333Xj00Udx/vz5cIc1a6Y6vsuXLyMnJwfLli3Drl27kJSUFMZopzbdf6/Dhw+jrKwsDBFOLtr+3mZyPG63G2azGQ899BBaW1vnONL/LyoSV2trK5566imcPn0adrsdN27cgMViwejoKADAbDYjKyvLb/v3v/8d5sinZ6rj+9rXvoaPPvoIn3/+OZqamvCf//wnzBFPbqrjAW6dEB988AHWr18fxkgDm078SjKT47lw4QI6Ozuxb98+/PjHPw7dJJ1hfeBojgwNDQkAsp9lfO+998SmTZvmKKrZM9nxVVZWisOHD4chquAFOp7Gxkbxox/9KIxRTV+g+Cf7W4r0vzO5xzNu3bp14h//+MdchyeEECIqRlx3Gp+iefHixWGOZG7cfnz/+c9/pP/Lud1unDp1CpmZmeEMT7ZA/16R+jUxkGj7e5vu8fzvf/+Tpkb/4osv0NPTg5UrV855fACib8R18+ZN8cgjj4iHHnpIVjuLxSKSkpLEXXfdJe655x5x5syZOYpwZu48vrNnz4qcnBxx3333iezsbFFXVxfmCOUJ9O91+fJlcffddwuPxxPGyKYnUPyT/S1F+t+ZnOP54IMPRFZWlrjvvvtETk6OeOONN0IWZ9Qlrl/84hfCYDCIwcHBcIcyJ6Lt+JR+PEqP/05KOZ6oSly//OUvxbJly0R/f3+4Q5kT0XZ8Sj8epcd/JyUdT1Qkrps3b4qnnnpKpKWlid7e3nCHM+ui7fiUfjxKj/9OSjyeoFb5iTRPPfUUmpqa8OabbyIhIUFaVEOv1+Ouu+4Kc3QzF23Hp/TjUXr8d1Lk8YQ7c84GAAG31157LdyhzYpoOz6lH4/S47+TEo+HzyoSkeJE5X1cRBTdmLiISHGYuIhIcZi4iEhxmLiISHGYuIhIcZi4iEhxmLiISHGYuIhIcZi4iEhxmLiISHGYuIhIcf4fykejGz2FKKkAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: total: 26 s\n",
      "Wall time: 31.8 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import ScalarFormatter\n",
    "\n",
    "mu = 1e-8\n",
    "r = 1e-8\n",
    "Na = 1e5\n",
    "Nb = 1e5\n",
    "Nanc = 1e5\n",
    "M = 1.0 / (2 * Na)\n",
    "Tsplit = 2e5 #generations\n",
    "L = 65_536\n",
    "\n",
    "a_results_dict = {}\n",
    "b_results_dict = {}\n",
    "for dist in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16_384, 32_768, 65_536]:\n",
    "    a_results_dict[dist] = np.array([], dtype=np.float64)\n",
    "    b_results_dict[dist] = np.array([], dtype=np.float64)\n",
    "\n",
    "\n",
    "def a_b_from_tree(tree):\n",
    "    tempList = []\n",
    "    for node in tree.nodes():\n",
    "        leaves = tuple(sorted(list(tree.leaves(node))))\n",
    "        mutType = mutTypeDict[leaves]\n",
    "        length = tree.get_branch_length(node)\n",
    "        tempList.append([mutType,length])\n",
    "\n",
    "    # calculate the total length of hetA and hetB type branches in this tree\n",
    "    a_length = sum([y for x,y in tempList if x=='hetA'])\n",
    "    b_length = sum([y for x,y in tempList if x=='hetB'])\n",
    "\n",
    "    return a_length, b_length\n",
    "\n",
    "\n",
    "for i in range(1, 101): # this would normally be 5001\n",
    "\n",
    "    demography = msprime.Demography()\n",
    "    demography.add_population(name=\"A\", initial_size=Na)\n",
    "    demography.add_population(name=\"B\", initial_size=Nb)\n",
    "    demography.add_population(name=\"C\", initial_size=Nanc)\n",
    "    demography.add_population_split(time=Tsplit, derived=[\"A\", \"B\"], ancestral=\"C\")\n",
    "    demography.set_migration_rate(source=\"A\",dest=\"B\",rate = M)\n",
    "    \n",
    "    ts = msprime.sim_ancestry(\n",
    "        samples={\"A\": 1, \"B\": 1}, \n",
    "        demography=demography,\n",
    "        ploidy=2,\n",
    "        recombination_rate=r,\n",
    "        sequence_length=L,\n",
    "        discrete_genome=True, # means genelogies have defined coords\n",
    "        random_seed = 44+i,\n",
    "        record_migrations=True\n",
    "    )\n",
    "\n",
    "    mt_ts = msprime.sim_mutations(ts, rate=mu, discrete_genome=False, random_seed = 44+i)\n",
    "\n",
    "    seq_a_length = []\n",
    "    seq_b_length = []\n",
    "    hetAB = [0] * L\n",
    "    \n",
    "    for tree in mt_ts.trees():\n",
    "        \n",
    "        start = tree.interval.left\n",
    "        end = tree.interval.right\n",
    "        span = int(end - start)\n",
    "        a_length, b_length = a_b_from_tree(tree)\n",
    "        seq_a_length = seq_a_length + ([a_length] * span)\n",
    "        seq_b_length = seq_b_length + ([b_length] * span)\n",
    "        for site in tree.sites():\n",
    "            mut_pos = int(site.position) # we are still using infinite sites here as we record the mutation rather than the alleles\n",
    "            leaves = tuple(sorted(list(tree.leaves(site.mutations[0].node))))\n",
    "            mutType = mutTypeDict[leaves]\n",
    "            if mutType == \"hetAB\":\n",
    "                hetAB[mut_pos] = 1\n",
    "\n",
    "    \n",
    "    for dist in a_results_dict.keys():\n",
    "        for interval_start in range(0, L, dist):\n",
    "            interval_hetAB = hetAB[interval_start:interval_start+dist]\n",
    "            if sum(interval_hetAB) > 0:\n",
    "                interval_res_a = seq_a_length[interval_start:interval_start+dist]\n",
    "                interval_res_b = seq_b_length[interval_start:interval_start+dist]\n",
    "                interval_a = sum(interval_res_a) / len(interval_res_a)\n",
    "                interval_b = sum(interval_res_b) / len(interval_res_b)\n",
    "                a_results_dict[dist] = np.append(a_results_dict[dist], np.array([interval_a]))\n",
    "                b_results_dict[dist] = np.append(b_results_dict[dist], np.array([interval_b]))\n",
    "\n",
    "\n",
    "    if i % 100 == 0 and i > 0: # this would normally be 5000\n",
    "        print(i)\n",
    "\n",
    "\n",
    "        dist_list = []\n",
    "        asymm_list_l = []\n",
    "        asymm_list_p = []\n",
    "        asymm_list_u = []\n",
    "        topology_list = []\n",
    "        \n",
    "        print(\"jack-knifing...\")\n",
    "        \n",
    "        for dist in a_results_dict.keys():\n",
    "            if len(a_results_dict[dist]) > 0:\n",
    "                jk_block = int(round((len(a_results_dict[dist]) / 25), 0)) # ensures 25 blocks\n",
    "                dist_list.append(dist)\n",
    "                lower, point, upper = block_jackknife(a_results_dict[dist], b_results_dict[dist], jk_block)\n",
    "                asymm_list_l.append(lower)\n",
    "                asymm_list_p.append(point)\n",
    "                asymm_list_u.append(upper)\n",
    "                \n",
    "        fig, ax = plt.subplots()\n",
    "        fig.set_size_inches(3, 3)\n",
    "        ax.scatter(x = dist_list, y = asymm_list_p, s=7, color=\"k\")\n",
    "        ax.errorbar(x = dist_list, y = asymm_list_p, yerr = np.subtract(asymm_list_p, asymm_list_l), color=\"k\", ls='none')\n",
    "        ax.xaxis.set_major_formatter(ScalarFormatter())\n",
    "        ax.set_xscale('log', base=2)\n",
    "        ax.set_xlim(0.5, 2*L)\n",
    "        ax.set_ylim(0, 0.2)\n",
    "        ax.set_xscale('log', base=2)\n",
    "        plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}