MackintoshSetter2024.supplementary_python_nb_2.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": [
    "### Testing the expected branch length 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 = [\n",
    "    [0],\n",
    "    [1],\n",
    "    [0,2,3],\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": "code",
   "execution_count": 3,
   "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 non-recombining blocks\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": 4,
   "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": 5,
   "id": "d5a6f5c1-2c78-43a5-b2d9-3d93c1de1c6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def block_jackknife(a_list, b_list, jk_block):\n",
    "    #print(\"trimming {} sims...\".format(len(a_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",
    "    #print(\"total of {} jk blocks\".format(total_jk_blocks))\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": "41aa5fd7-b89d-4584-92f5-52b60aa1f2c9",
   "metadata": {},
   "source": [
    "##### Here we simulate sequences with recombination and calculate A_m across a range of block lengths. The demography is that of Heliconius melpomene and H. cydno, inferred by Laetsch et al (2023) using the mutlilocus inference method gIMble. Their demographic model is the five parameter IM model, where H. melpomene is population A and H. cydno is population B."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f493b460-358e-4eb0-a4ae-518147162630",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jack-knifing...\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASUAAAEXCAYAAADr1VLRAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8g+/7EAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAdCUlEQVR4nO3df0zU9/0H8Ccc90MRqJV6olIEEzkJ/jzqxMlIt+5Ql646F3HtqMnqIpVFgTQWZIuzTaXdusU0E4iOrjWbP1rQ1m1s45ooY4K1/NAsxR9ZQGHqlWLqHaeRH8f7+4fjvl7vBD6nd/e+u+cj+STem/f7w+td7p593+c+n89FCCEEiIgkERnoAoiI7sdQIiKpMJSISCoMJSKSCkOJiKTCUCIiqTCUiEgqDCUikkpUoAt4VEZGRnD9+nXExMQgIiIi0OUQ0X2EEOjv78fMmTMRGTn2WihkQun69etITEwMdBlENIaenh7Mnj17zD4hE0oxMTEA7k06NjY2wNUQ0f1sNhsSExOdr9OxhEwojb5li42NZSgRSWoih1Z4oJuIpMJQIiKpMJSISCoMJSKSCkOJiKTCUCIiqTCUiEgqDCXyqdu3byMiIgIRERG4fft2oMuhIMBQIiKpMJTClMwrGJlrI99jKBGRVBhKRCQVhhIRSYWhRERSYSgRkVQYSkQkFYYSEUmFoUREUmEoEZFUGEpEJBWvQqmiogLJycnQ6XQwGo1obGx8YN8bN27g+eefR2pqKiIjI1FYWOixX21tLdLS0qDVapGWlobjx497UxoRBTnFoXT06FEUFhairKwM7e3tyMrKwurVq9Hd3e2x/8DAAJ544gmUlZVh0aJFHvs0NzcjNzcXeXl5OH/+PPLy8rBhwwZ8+umnSssjomAnFFq2bJnIz893aTMYDKKkpGTcsdnZ2WL79u1u7Rs2bBCrVq1yacvJyREbN26ccF1Wq1UAEFardcJjwpndbhcABABht9uDfgzJTcnrU9FKaXBwEK2trTCZTC7tJpMJTU1NXgdjc3Oz2z5zcnLG3OfAwABsNpvLRkTBT1Eo9fX1weFwQK/Xu7Tr9XpYLBavi7BYLIr3WV5ejri4OOfGr+wmCg1eHej++rdcCiEm9M2Xj3KfpaWlsFqtzq2np+ehfj8RyUHR13bHx8dDpVK5rWB6e3vdVjpKzJgxQ/E+tVottFqt17+TiOSkaKWk0WhgNBphNptd2s1mM1asWOF1EZmZmW77rK+vf6h9ElFwUrRSAoDi4mLk5eUhIyMDmZmZ2L9/P7q7u5Gfnw/g3tuqa9eu4eDBg84x586dAwDY7XZ8+eWXOHfuHDQaDdLS0gAA27dvx7e+9S289dZbeO655/Dxxx/jk08+wb/+9a9HMEUiCirefLy3b98+kZSUJDQajVi6dKloaGhw/mzTpk0iOzvbpT/+9/Hu/VtSUpJLnw8//FCkpqYKtVotDAaDqK2tVVQTTwlQRuaP93lKQOhR8vqMEEKIAGThI2ez2RAXFwer1YrY2NhAlyO927dvY8qUKQDurWCjo6ODegzJTcnrk9e+EZFUGEpEJBWGEhFJhaFERFJhKFFI4Lfqhg6GEhFJhaFERFJhKBGRVBhKRCQVhhIRSYWhRERSYSgRkVQYSkQkFYYS+dTly5ed/87IyMClS5cCWA0FA4YS+dTzzz/v/Pfly5fxgx/8IIDVUDBgKJHPOBwOXLx40fl4ZGQEHR0dcDgcAayKZMdQIp9RqVQwGAzOx5GRkUhLS4NKpQpgVSQ7hhL51KFDh5z/njdvHo4dOxbAaigYKP7iACIl5s2b5/x3S0sLb21L4+JKiYikwlAiIqkwlIhIKgwlIpIKQ4mIpMJQIiKpMJSISCoMJZIOL+INbwwlkg4v4g1vDCWSCi/iJa9CqaKiAsnJydDpdDAajWhsbByzf0NDA4xGI3Q6HVJSUlBVVeXWZ+/evUhNTcWkSZOQmJiIoqIi3L1715vywk4ofREjL+IlCIWOHDki1Gq1OHDggOjo6BDbt28X0dHR4urVqx77d3Z2ismTJ4vt27eLjo4OceDAAaFWq0VNTY2zzx//+Eeh1WrFn/70J9HV1SX+8Y9/iISEBFFYWDjhuqxWqwAgrFar0ikFPbvdLgAIAMJutwf9mLa2NucYg8EgLl686JPfQ/6j5PWpOJSWLVsm8vPzXdoMBoMoKSnx2H/Hjh3CYDC4tG3ZskUsX77c+bigoEB8+9vfdulTXFwsVq5cOeG6GEpyBozMY8h/lLw+Fb19GxwcRGtrK0wmk0u7yWRCU1OTxzHNzc1u/XNyctDS0oKhoSEAwMqVK9Ha2oqzZ88CADo7O1FXV4fvfe97SsojBfgJF8lK0a1L+vr64HA4oNfrXdr1ej0sFovHMRaLxWP/4eFh9PX1ISEhARs3bsSXX36JlStXQgiB4eFhvPzyyygpKXlgLQMDAxgYGHA+ttlsSqYS9jx9wvX5558HsCKie7w60B0REeHyWAjh1jZe//vbT506hTfeeAMVFRVoa2vDsWPH8Je//AWvv/76A/dZXl6OuLg455aYmOjNVMISP+EimSkKpfj4eKhUKrdVUW9vr9tqaNSMGTM89o+KisK0adMAAL/4xS+Ql5eHzZs3Y8GCBVi3bh327NmD8vJyjIyMeNxvaWkprFarc+vp6VEylbDm7SdcfMtH/qAolDQaDYxGI8xms0u72WzGihUrPI7JzMx0619fX4+MjAyo1WoAwJ07dxAZ6VqKSqWCuHcg3uN+tVotYmNjXTaaOG9uU8uTGskvlB5FHz0loLq6WnR0dIjCwkIRHR0trly5IoQQoqSkROTl5Tn7j54SUFRUJDo6OkR1dbXbKQG7du0SMTEx4vDhw6Kzs1PU19eLuXPnig0bNky4Ln765ttPq4aHh53979+Gh4cDXpu3Y8h/fHpKgBBC7Nu3TyQlJQmNRiOWLl0qGhoanD/btGmTyM7Odul/6tQpsWTJEqHRaMScOXNEZWWly8+HhobEL3/5SzF37lyh0+lEYmKi2Lp1q/jqq68mXBNDyfcvYoPB4BwTGRkp0tLSpKmNoSQ3Ja/PCCEe8P4oyNhsNsTFxcFqtYbdW7nbt29jypQpAAC73T6hm/N7M6a9vR1Lly4FABgMBnz00UdITU2VojZvxpD/KHl98ttMaML4zSTkD7wgl4ikwlAiIqkwlIhIKgwlIpIKQ0kyoXRvJCJvMJSISCoMJSKSCkOJiKTCUCIiqTCUKGzxQwU5MZSISCoMJSKSCkOJiKTCuwSQT0VHRz/w7qFEnnClRERSYSgRkVQYSj7Ej5yJlGMoEZFUGEpEJBWGEhFJhaFEIYHf3hs6GEoUEvjtvaGDoURBz+Fw4OLFi87HIyMj6OjogMPhCGBV5C2GEgU9lUoFg8HgfBwZGYm0tDSoVKoAVkXeYihRSDh06JDz3/PmzcOxY8cCWA09DF77RiGB394bOrhSIiKpMJRCAD8Op1DCUAoB/DicQolXoVRRUYHk5GTodDoYjUY0NjaO2b+hoQFGoxE6nQ4pKSmoqqpy63Pr1i0UFBQgISEBOp0O8+fPR11dnTflhRV+HE6hRnEoHT16FIWFhSgrK0N7ezuysrKwevVqdHd3e+zf1dWFNWvWICsrC+3t7di5cye2bduG2tpaZ5/BwUF897vfxZUrV1BTU4NLly7hwIEDmDVrlvczCxP8OJxCjlBo2bJlIj8/36XNYDCIkpISj/137NghDAaDS9uWLVvE8uXLnY8rKytFSkqKGBwcVFqOk9VqFQCE1Wr1eh+Pmt1uFwAEAGG32302pq2tzTnGYDCIixcvSlObN/xVm7/mQ8pen4pWSoODg2htbYXJZHJpN5lMaGpq8jimubnZrX9OTg5aWlowNDQEADhx4gQyMzNRUFAAvV6P9PR07NmzZ8y3IAMDA7DZbC6br8l6f6SvfxyempoawGqIHo6iUOrr64PD4YBer3dp1+v1sFgsHsdYLBaP/YeHh9HX1wcA6OzsRE1NDRwOB+rq6vDzn/8cv/nNb/DGG288sJby8nLExcU5t8TERCVTISJJeXWgOyIiwuWxEMKtbbz+97ePjIxg+vTp2L9/P4xGIzZu3IiysjJUVlY+cJ+lpaWwWq3Oraenx5upEElJ6arcm1W8v8YopeiM7vj4eKhUKrdVUW9vr9tqaNSMGTM89o+KisK0adMAAAkJCVCr1S4HZ+fPnw+LxYLBwUFoNBq3/Wq1Wmi1WiXlE1EQULRS0mg0MBqNMJvNLu1msxkrVqzwOCYzM9Otf319PTIyMqBWqwEA3/zmN/Gf//wHIyMjzj6XL19GQkKCx0AiotCl+O1bcXExfv/73+Pdd9/FhQsXUFRUhO7ubuTn5wO497bqxRdfdPbPz8/H1atXUVxcjAsXLuDdd99FdXU1XnnlFWefl19+GTdv3sT27dtx+fJl/PWvf8WePXtQUFDwCKZIwWb0u+KEELyGLQwpviA3NzcXN2/exGuvvYYbN24gPT0ddXV1SEpKAgDcuHHD5Zyl5ORk1NXVoaioCPv27cPMmTPxzjvvYP369c4+iYmJqK+vR1FRERYuXIhZs2Zh+/btePXVVx/BFIkomHh1l4CtW7di69atHn/23nvvubVlZ2ejra1tzH1mZmbizJkz3pRDRCGE175JhhfXUrhjKEmGF9dSuONN3iQy1sW1MlzLNnoAmsiXuFKSCC+uJWIoSYf3mpabrNc/hhK+fZMM7zVN4Y4rJSKSCkOJyMf4lk8ZhhIRSYXHlMIUP94nWXGlRERSCdtQ4vt8IjmFbSgRkZwYSkQkFYYShS3ekUFODCUKW7wjg5wYShSW+HXn8mIoUVjiHRnkxVCisMU7MsiJoeRDPJAqN37duZwYSj7EA6lEyjGUfIQHUom8w1DyER5IJfIOQ8mHeCCVSDneusSHeGtbIuW4UiIiqTCUiEgqDCUikgpDiYikwlBSgGdoE/meV6FUUVGB5ORk6HQ6GI1GNDY2jtm/oaEBRqMROp0OKSkpqKqqemDfI0eOICIiAmvXrvWmNJ+S9Qzt0S8BEELwEz4KeopD6ejRoygsLERZWRna29uRlZWF1atXo7u722P/rq4urFmzBllZWWhvb8fOnTuxbds21NbWuvW9evUqXnnlFWRlZSmfiY/xDG0i/1AcSr/97W/x0ksvYfPmzZg/fz727t2LxMREVFZWeuxfVVWFJ598Env37sX8+fOxefNm/OQnP8Hbb7/t0s/hcOCFF17A7t27kZKS4t1sfIhnaBP5h6JQGhwcRGtrK0wmk0u7yWRCU1OTxzHNzc1u/XNyctDS0oKhoSFn22uvvYYnnngCL7300oRqGRgYgM1mc9l8jWdoE/meojO6+/r64HA4oNfrXdr1ej0sFovHMRaLxWP/4eFh9PX1ISEhAadPn0Z1dTXOnTs34VrKy8uxe/duJeU/NJ6hTeR7Xh3ojoiIcHkshHBrG6//aHt/fz9+/OMf48CBA4iPj59wDaWlpbBarc6tp6dHwQyISFaKVkrx8fFQqVRuq6Le3l631dCoGTNmeOwfFRWFadOm4fPPP8eVK1fw7LPPOn8+MjJyr7ioKFy6dAlz5851269Wq4VWq1VSPhEFAUUrJY1GA6PRCLPZ7NJuNpuxYsUKj2MyMzPd+tfX1yMjIwNqtRoGgwH//ve/ce7cOef2/e9/H08//TTOnTuHxMREhVMiomCm+C4BxcXFyMvLQ0ZGBjIzM7F//350d3cjPz8fwL23VdeuXcPBgwcBAPn5+fjd736H4uJi/PSnP0VzczOqq6tx+PBhAIBOp0N6errL73jssccAwK2diEKf4lDKzc3FzZs38dprr+HGjRtIT09HXV0dkpKSAAA3btxwOWcpOTkZdXV1KCoqwr59+zBz5ky88847WL9+/aObBRGFDK/up7R161Zs3brV48/ee+89t7bs7Gy0tbVNeP+e9hEuRs/OJgpXvMkbhQSGeejgBblEJBWGEhFJhaFERFJhKBGRVMI2lHjDNiI5hW0oyXrDNgo9/B+gMmEZSrxhG3nLm4Dx5n+ASn+PN3X5a4xiIkRYrVYBQFit1gn1NxgMAoAAICIjI0VaWtq4Y+x2u3OM3W5/5P3Jv7z5+yh93gwPDzv7378NDw8/0t/jzfPZX2OEUPb6DNtQamtrc/7HNRgM4uLFi+OOYSiFFqV/H38FjNLf401d/hozSsnrMyzfvgHuN2xLTU0NYDUUDLy9JbLSO5Yq/T3e1OWvMV4ZN7aChNKVkjerGK6UQos3fx9/rLC9+T3e1OWvMUIoe33y2jciBfx1S2Slv8ebuvw1RqmwfftGRHJiKBGRVBhKRCQVhhIRSYWhRERSYSgRkVQYSkQkFYYSEUmFoUREUmEoEZFUeJkJhS1+LZOcuFIiIqlwpeRD/D8xkXJcKRGRVBhKRCQVhhIRScWrUKqoqEBycjJ0Oh2MRiMaGxvH7N/Q0ACj0QidToeUlBRUVVW5/PzAgQPIysrC1KlTMXXqVDzzzDM4e/asN6X51OgxIiGEz27uRRTuFIfS0aNHUVhYiLKyMrS3tyMrKwurV69Gd3e3x/5dXV1Ys2YNsrKy0N7ejp07d2Lbtm2ora119jl16hR+9KMf4eTJk2hubsaTTz4Jk8mEa9eueT8zIgpOE7rB7n2WLVsm8vPzXdoMBoMoKSnx2H/Hjh3CYDC4tG3ZskUsX778gb9jeHhYxMTEiPfff3/CdfnjHt1E/ri3uzdjZK1rlM++zWRwcBCtra0wmUwu7SaTCU1NTR7HNDc3u/XPyclBS0sLhoaGPI65c+cOhoaG8Pjjjz+wloGBAdhsNpeNiIKfolDq6+uDw+GAXq93adfr9bBYLB7HWCwWj/2Hh4fR19fncUxJSQlmzZqFZ5555oG1lJeXIy4uzrklJiYqmQoRScqrA90REREuj4UQbm3j9ffUDgC/+tWvcPjwYRw7dgw6ne6B+ywtLYXVanVuPT09SqZARJJSdEZ3fHw8VCqV26qot7fXbTU0asaMGR77R0VFYdq0aS7tb7/9Nvbs2YNPPvkECxcuHLMWrVYLrVarpHwiCgKKVkoajQZGoxFms9ml3Ww2Y8WKFR7HZGZmuvWvr69HRkYG1Gq1s+3Xv/41Xn/9dfz9739HRkaGkrKIKJRM+PD5/xw5ckSo1WpRXV0tOjo6RGFhoYiOjhZXrlwRQghRUlIi8vLynP07OzvF5MmTRVFRkejo6BDV1dVCrVaLmpoaZ5+33npLaDQaUVNTI27cuOHc+vv7J1wXP30jf5D1Uy5Z6xql5PXp1dd279u3TyQlJQmNRiOWLl0qGhoanD/btGmTyM7Odul/6tQpsWTJEqHRaMScOXNEZWWly8+TkpKcE71/27Vr14RrYiiRP8j64pe1rlFKXp8RQoTGZew2mw1xcXGwWq2IjY0dt//t27cxZcoUAIDdbucZ2jQh3jxv/DFG1rpGKXl98to3IpIKQ4mIpMJQIiKphO2dJ3lXSCI5caVERFJhKBGRVBhKRCQVhhIRSYWhRERSYSgRkVTC9pQAIn/h6SfKcKVERFJhKBGRVBhKRCQVHlMiUoDHh3yPKyUikgpDiYikwlAiIqkwlIhIKgwlIpIKQ4mIpMJQIiKpMJSISCoMJSKSCs/oJpKQ0jPHvTnT3F9jlOJKiYikwlAiIqkwlIhIKgwlIpIKQ4mIpOJVKFVUVCA5ORk6nQ5GoxGNjY1j9m9oaIDRaIROp0NKSgqqqqrc+tTW1iItLQ1arRZpaWk4fvy4N6URUZBTHEpHjx5FYWEhysrK0N7ejqysLKxevRrd3d0e+3d1dWHNmjXIyspCe3s7du7ciW3btqG2ttbZp7m5Gbm5ucjLy8P58+eRl5eHDRs24NNPP/V+ZkQUlCKEwpMOvvGNb2Dp0qWorKx0ts2fPx9r165FeXm5W/9XX30VJ06cwIULF5xt+fn5OH/+PJqbmwEAubm5sNls+Nvf/ubss2rVKkydOhWHDx+eUF02mw1xcXGwWq2IjY1VMiUi8jElr09FK6XBwUG0trbCZDK5tJtMJjQ1NXkc09zc7NY/JycHLS0tGBoaGrPPg/YJAAMDA7DZbC4bEQU/RWd09/X1weFwQK/Xu7Tr9XpYLBaPYywWi8f+w8PD6OvrQ0JCwgP7PGifAFBeXo7du3e7tTOciOQz+rqcyBszry4ziYiIcHkshHBrG6//19uV7rO0tBTFxcXOx11dXVi8eDESExPHnwARBUR/fz/i4uLG7KMolOLj46FSqdxWML29vW4rnVEzZszw2D8qKgrTpk0bs8+D9gkAWq0WWq3W+TgpKQkA0N3dPe6kZWez2ZCYmIienp6QOD7G+cjLX3MRQqC/vx8zZ84ct6+iUNJoNDAajTCbzVi3bp2z3Ww247nnnvM4JjMzE3/+859d2urr65GRkQG1Wu3sYzabUVRU5NJnxYoVE64tMvLe4bG4uLigf6KMio2NDZm5AJyPzPwxlwkvFoRCR44cEWq1WlRXV4uOjg5RWFgooqOjxZUrV4QQQpSUlIi8vDxn/87OTjF58mRRVFQkOjo6RHV1tVCr1aKmpsbZ5/Tp00KlUok333xTXLhwQbz55psiKipKnDlzZsJ1Wa1WAUBYrValU5JOKM1FCM5HZjLORXEoCSHEvn37RFJSktBoNGLp0qWioaHB+bNNmzaJ7Oxsl/6nTp0SS5YsERqNRsyZM0dUVla67fPDDz8UqampQq1WC4PBIGpraxXVJON/XG+F0lyE4HxkJuNcvAolGd29e1fs2rVL3L17N9ClPLRQmosQnI/MZJyL4pMniYh8iRfkEpFUGEpEJBWGEhFJhaFERFIJm1Bat24dpk6dih/+8IeBLmXC+vv78dRTT2Hx4sVYsGABDhw4EOiSHsqlS5ewePFi5zZp0iR89NFHgS5LsbGeS8H2PBur3qioKOffavPmzX6rKWw+fTt58iTsdjvef/991NTUBLqcCXE4HBgYGMDkyZNx584dpKen47PPPnNenhPM7HY75syZg6tXryI6OjrQ5Sgy1nMp2J5nY9UbHx+Pvr4+v9cUNiulp59+GjExMYEuQxGVSoXJkycDAO7evQuHw+Hz79zylxMnTuA73/lO0AUSMPZzKdieZzLWGxShVF5ejqeeegoxMTGYPn061q5di0uXLgW6rEdivLndunULixYtwuzZs7Fjxw7Ex8cHsNrxTfRv9cEHHyA3NzcAFY4vlJ5vDzMXm80Go9GIlStXoqGhwceV/r+gCKWGhgYUFBTgzJkzMJvNGB4ehslkwu3btwEARqMR6enpbtv169cDXPn4xpvbY489hvPnz6OrqwuHDh3CF198EeCKxzbefIB7T/bTp09jzZo1Aaz0wSYyh2DxMHO5cuUKWltbUVVVhRdffNF/9yoL6PnkXurt7RUAXK65m4iTJ0+K9evX+6iqR2OsueXn54sPPvggAFV5z9N8Dh48KF544YUAVqWMpzmM9VyS+XmmdC6jVq1aJT777DNflyeEECIoVkpfZ7VaAQCPP/54gCt59O6f2xdffOH8v5PNZsM///lPpKamBrI8xTz9rWR+6+ZJKD3fJjqXr776CgMDAwCA//73v+jo6EBKSorP6wMQfCulkZER8eyzz4qVK1cqGmcymUR8fLyYNGmSmDVrljh79qyPKvTe1+fW0tIiFi1aJBYuXCgWLFggKioqAlyhMp7+Vrdu3RLTp08XAwMDAaxs4jzNYaznkszPMyVzOX36tEhPTxcLFy4UixYtEsePH/dbnUEXSlu3bhVJSUmip6cn0KU8cqE2t1CYTyjMYVSwzCWoQulnP/uZmD17tujs7Ax0KY9cqM0tFOYTCnMYFUxzCYpQGhkZEQUFBWLmzJni8uXLgS7nkQq1uYXCfEJhDqOCcS5efZuJvxUUFODQoUP4+OOPERMT4/ySgbi4OEyaNCnA1T2cUJtbKMwnFOYwKijnEuhUnAgAHrc//OEPgS7toYXa3EJhPqEwh1HBOJewufaNiIJDUJ6nREShi6FERFJhKBGRVBhKRCQVhhIRSYWhRERSYSgRkVQYSkQkFYYSEUmFoUREUmEoEZFUGEpEJBWGEhFJ5f8A6We9B+ZUrBgAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n",
      "CPU times: total: 32.3 s\n",
      "Wall time: 39.9 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 = 2.9e-9\n",
    "r = 1.89e-8\n",
    "\n",
    "Tmya = 1.054\n",
    "Tyears = Tmya*1e6\n",
    "Tgens = Tyears*4\n",
    "\n",
    "Nmel = 5.49e5\n",
    "Ncyd = 1.415e6\n",
    "Nanc = 9.279e5\n",
    "\n",
    "me = 7.4e-7  #msprime now uses me directly, it seems\n",
    "M = 4*Nmel*me\n",
    "\n",
    "\n",
    "L = 65_536\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, 11): # this would normally be 1001\n",
    "\n",
    "    demography = msprime.Demography()\n",
    "    demography.add_population(name=\"MEL\", initial_size=Nmel)\n",
    "    demography.add_population(name=\"CYD\", initial_size=Ncyd)\n",
    "    demography.add_population(name=\"ANC\", initial_size = Nanc) # this was an approximation based on info in the table, and I assumed it was for ultiamte ancestor\n",
    "\n",
    "    demography.set_migration_rate(source=\"MEL\",dest=\"CYD\", rate=me)\n",
    "\n",
    "    demography.add_population_split( time = Tgens, derived = [\"MEL\",\"CYD\"], ancestral = \"ANC\")\n",
    "    demography.add_migration_rate_change(time = Tgens, source=\"CYD\", dest = \"MEL\", rate = 0.0)\n",
    "    \n",
    "    \n",
    "    ts = msprime.sim_ancestry(\n",
    "        samples={\"MEL\": 1, \"CYD\": 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) # effectively assumes infinite sites\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",
    "    if i % 10 == 0 and i > 0:\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]) / 10), 0)) # this would normally be 25\n",
    "                #print(jk_block)\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.25)\n",
    "        ax.set_xscale('log', base=2)\n",
    "        plt.show()\n",
    "\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d3db9c5b-2454-4b67-a46a-4d427743b8dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]\n",
      "\n",
      "[0.015150545279164644, 0.022360492096428836, 0.023912013144455537, 0.0370315045344794, 0.04178326042741269, 0.052407356517363335, 0.055556463279089155, 0.06985143670282702, 0.05830995470235376, 0.03639193059113682, 0.025155977211910395, 0.009606519589526432, 0.011811394584269917, 0.004778706703041476, 0.004778706703041379, 0.004778706703041379, 0.004778706703041379]\n",
      "\n",
      "[0.03308698479273315, 0.04152779985230648, 0.04610084295327104, 0.05563329178565965, 0.061958284393398774, 0.07988399876245739, 0.07838290874113336, 0.08906374848865159, 0.0828970410361777, 0.06555730633445316, 0.04864633842444908, 0.03205956679488518, 0.032406037246984806, 0.0321582651490545, 0.0321582651490545, 0.0321582651490545, 0.0321582651490545]\n",
      "\n",
      "[0.051023424306301655, 0.060695107608184126, 0.06828967276208654, 0.0742350790368399, 0.08213330835938486, 0.10736064100755144, 0.10120935420317757, 0.10827606027447616, 0.10748412737000165, 0.09472268207776949, 0.07213669963698777, 0.05451261400024393, 0.0530006799096997, 0.05953782359506753, 0.059537823595067624, 0.059537823595067624, 0.059537823595067624]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# print results\n",
    "print(dist_list)\n",
    "print()\n",
    "print(asymm_list_l)\n",
    "print()\n",
    "print(asymm_list_p)\n",
    "print()\n",
    "print(asymm_list_u)\n",
    "print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cf856921-d366-43e9-aa8d-b3e09fbfdcaa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0, 0.03308698479273315}, {1, 0.04152779985230648}, {2, 0.04610084295327104}, {3, 0.05563329178565965}, {4, 0.061958284393398774}, {5, 0.07988399876245739}, {6, 0.07838290874113336}, {7, 0.08906374848865159}, {8, 0.0828970410361777}, {9, 0.06555730633445316}, {10, 0.04864633842444908}, {11, 0.03205956679488518}, {12, 0.032406037246984806}, {13, 0.0321582651490545}, {14, 0.0321582651490545}, {15, 0.0321582651490545}, {16, 0.0321582651490545}, "
     ]
    }
   ],
   "source": [
    "# print results for plotting in Mathematica\n",
    "for i, dist in enumerate(dist_list):\n",
    "    print(\"{}{}, {}{}, \".format(\"{\", i, asymm_list_p[i], \"}\"), end=\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "21cef1df-66a8-4d57-9150-7406881ee918",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Line[{{0, 0.015150545279164644}, {0, 0.051023424306301655}}], Line[{{1, 0.022360492096428836}, {1, 0.060695107608184126}}], Line[{{2, 0.023912013144455537}, {2, 0.06828967276208654}}], Line[{{3, 0.0370315045344794}, {3, 0.0742350790368399}}], Line[{{4, 0.04178326042741269}, {4, 0.08213330835938486}}], Line[{{5, 0.052407356517363335}, {5, 0.10736064100755144}}], Line[{{6, 0.055556463279089155}, {6, 0.10120935420317757}}], Line[{{7, 0.06985143670282702}, {7, 0.10827606027447616}}], Line[{{8, 0.05830995470235376}, {8, 0.10748412737000165}}], Line[{{9, 0.03639193059113682}, {9, 0.09472268207776949}}], Line[{{10, 0.025155977211910395}, {10, 0.07213669963698777}}], Line[{{11, 0.009606519589526432}, {11, 0.05451261400024393}}], Line[{{12, 0.011811394584269917}, {12, 0.0530006799096997}}], Line[{{13, 0.004778706703041476}, {13, 0.05953782359506753}}], Line[{{14, 0.004778706703041379}, {14, 0.059537823595067624}}], Line[{{15, 0.004778706703041379}, {15, 0.059537823595067624}}], Line[{{16, 0.004778706703041379}, {16, 0.059537823595067624}}], "
     ]
    }
   ],
   "source": [
    "# print results for plotting in Mathematica\n",
    "for i, l, u in zip([k for k in range(0, 17)], asymm_list_l, asymm_list_u):\n",
    "   print(\"Line{}{}, {}{}{}, {}{}\".format(\"[{{\", i, l, \"}, {\", i, u, \"}}], \"), end=\"\")"
   ]
  }
 ],
 "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
}