s03_alignments_haplotypes_Ace1_2020-02-24.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Haplotype alignments\n",
    "\n",
    "Alignment of haplotypes around the *Ace1* duplication for further phylogenetic analysis.\n",
    "\n",
    "## Input\n",
    "\n",
    "Input files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input data\n",
    "outdir      = \"results_admixture_phylo/\"\n",
    "metasam_fn  = \"metadata/samples.meta_phenotypes_acegenotype.simple.txt\"\n",
    "callset_fn  = \"/home/xavi/Documents/VariationAg1k/data/phase2.AR1/variation/main/zarr2/ag1000g.phase2.ar1.pass/\"\n",
    "accessi_fn  = \"/home/xavi/Documents/VariationAg1k/data/phase2.AR1/accessibility/accessibility.h5\"\n",
    "haploty_fn  = \"/home/xavi/Documents/VariationAg1k/data/phase2.AR1/haplotypes/zarr2/ag1000g.phase2.ar1.samples/\"\n",
    "snpeff_fn   = \"/home/xavi/Documents/VariationAg1k/data/phase2.AR1/snpeff/zarr2/\"\n",
    "gffann_fn   = \"metadata/Anopheles-gambiae-PEST_BASEFEATURES_AgamP4.9.gff3\"\n",
    "\n",
    "# define populations\n",
    "popl    = [\"BFcol\",\"BFgam\",\"CIcol\",\"GHcol\",\"GHgam\",\"GNgam\"]\n",
    "popc    = \"population\"\n",
    "sub1l   = popl\n",
    "sub1c   = popc\n",
    "chrom   = \"2R\"\n",
    "\n",
    "# exclude these samples\n",
    "excludec   = \"ox_code\"\n",
    "excludel   = list((\"NO RES\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import zarr\n",
    "import pandas as pd\n",
    "import allel\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data\n",
    "\n",
    "### Genotypes, haplotypes, variants & samples\n",
    "\n",
    "Load data for all variants & genotypes. Population and sample structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data:\n",
      "* Samples     =  345\n",
      "* Populations =  {'GHcol', 'BFgam', 'CIcol', 'GHgam', 'GNgam', 'BFcol'}\n",
      "population\n",
      "BFcol    75\n",
      "BFgam    92\n",
      "CIcol    71\n",
      "GHcol    55\n",
      "GHgam    12\n",
      "GNgam    40\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# load samples list with sample code, groupings, locations etc.\n",
    "samples_df   = pd.read_csv(metasam_fn, sep='\\t')\n",
    "samples_bool = (\n",
    "    samples_df[popc].isin(popl).values & \n",
    "    samples_df[sub1c].isin(sub1l).values &\n",
    "    ~samples_df[excludec].isin(excludel).values\n",
    ")\n",
    "samples_sub  = samples_df[samples_bool]\n",
    "samples_sub.reset_index(drop=True, inplace=True)\n",
    "\n",
    "# indexed dictionary of populations\n",
    "popdict = dict()\n",
    "for popi in popl: \n",
    "    popdict[popi]  = samples_sub[samples_sub[popc] == popi].index.tolist()\n",
    "\n",
    "# add an extra population composed of all other locations\n",
    "popdict[\"all\"] = []\n",
    "for popi in popl:\n",
    "    popdict[\"all\"] = popdict[\"all\"] + popdict[popi]\n",
    "    \n",
    "# report\n",
    "print(\"Data:\")\n",
    "print(\"* Samples     = \", samples_sub.shape[0])\n",
    "print(\"* Populations = \", set(samples_sub[popc]))\n",
    "print(samples_sub.groupby((\"population\")).size())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load variant, genotypes, haplotypes, accessibility, etc. data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load accessibility array...\n"
     ]
    }
   ],
   "source": [
    "# Accessibility\n",
    "import h5py\n",
    "print(\"Load accessibility array...\")\n",
    "accessi_df  = h5py.File(accessi_fn,mode=\"r\")\n",
    "accessi_arr = accessi_df[chrom][\"is_accessible\"][:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Divergence\n",
    "\n",
    "Calculate divergence of genotype frequencies between species and duplicated sequences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "has_dup = np.array([any(b in s for b in [\"TRUE\"]) for s in samples_sub[\"pop_dup\"].values])\n",
    "\n",
    "# dictionary with species (nondup) and dups\n",
    "popdict_div = dict()\n",
    "popdict_div[\"dup\"] = np.where(has_dup)[0].tolist()\n",
    "popdict_div[\"col\"] = samples_sub[  np.logical_and( samples_sub[\"m_s\"] == \"M\" , np.logical_not(has_dup) )  ].index.tolist()\n",
    "popdict_div[\"gam\"] = samples_sub[  np.logical_and( samples_sub[\"m_s\"] == \"S\" , np.logical_not(has_dup) )  ].index.tolist()\n",
    "popdict_div[\"cod\"] = samples_sub[  np.logical_and( samples_sub[\"m_s\"] == \"M\" , has_dup )  ].index.tolist()\n",
    "popdict_div[\"gad\"] = samples_sub[  np.logical_and( samples_sub[\"m_s\"] == \"S\" , has_dup )  ].index.tolist()\n",
    "popdict_div[\"all\"] = popdict_div[\"dup\"] + popdict_div[\"gam\"] + popdict_div[\"col\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Region of interest:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ace_start = 3489213 # start gene\n",
    "ace_end   = 3493788 # end gene\n",
    "ace_119S  = 3492074 # resistant variant\n",
    "ace_dups  = 3436800 # start duplication\n",
    "ace_dupe  = 3639600 # end duplication"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load data from region of interest:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Subset haps...\n"
     ]
    }
   ],
   "source": [
    "export_start = ace_dups\n",
    "export_end   = ace_dupe\n",
    "export_name  = \"duplication\"\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# calculate population allele counts\n",
    "hapalco_sub = haploty_sub.count_alleles_subpops(subpops=popdict_div)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pairwise differences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dif gam to dup = 0.0901244 +/- 0.0027261 \n",
      "dif col to dup = 0.0919380 +/- 0.0027750 \n"
     ]
    }
   ],
   "source": [
    "# pairwise difference\n",
    "mpd_dg = allel.mean_pairwise_difference_between(ac1=hapalco_seg[\"dup\"], ac2=hapalco_seg[\"gam\"])\n",
    "mpd_dc = allel.mean_pairwise_difference_between(ac1=hapalco_seg[\"dup\"], ac2=hapalco_seg[\"col\"])\n",
    "# jack-knife per-SNP estimates\n",
    "mpd_dg_est = allel.stats.misc.jackknife(mpd_dg , np.mean)\n",
    "mpd_dc_est = allel.stats.misc.jackknife(mpd_dc , np.mean)\n",
    "# report\n",
    "print(\"dif gam to dup = %.7f +/- %.7f \" % (mpd_dg_est[0], mpd_dg_est[1]))\n",
    "print(\"dif col to dup = %.7f +/- %.7f \" % (mpd_dc_est[0], mpd_dc_est[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dxy divergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dxy gam to dup = 0.0105591 +/- 0.0010111 \n",
      "dxy col to dup = 0.0108962 +/- 0.0010787 \n"
     ]
    }
   ],
   "source": [
    "# dxy per window\n",
    "dxy_dg = allel.windowed_divergence(ac1=hapalco_seg[\"dup\"], ac2=hapalco_seg[\"gam\"], pos = hapvars_seg[\"POS\"], size=100, is_accessible = accessi_arr)\n",
    "dxy_dc = allel.windowed_divergence(ac1=hapalco_seg[\"dup\"], ac2=hapalco_seg[\"col\"], pos = hapvars_seg[\"POS\"], size=100, is_accessible = accessi_arr)\n",
    "\n",
    "# jack-knife per-SNP estimate\n",
    "dxy_dg_est = allel.stats.misc.jackknife(dxy_dg[0] , np.nanmean)\n",
    "dxy_dc_est = allel.stats.misc.jackknife(dxy_dc[0] , np.nanmean)\n",
    "\n",
    "# report\n",
    "print(\"dxy gam to dup = %.7f +/- %.7f \" % (dxy_dg_est[0], dxy_dg_est[1]))\n",
    "print(\"dxy col to dup = %.7f +/- %.7f \" % (dxy_dc_est[0], dxy_dc_est[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PBS relative to duplicated sequences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PBS col = 0.0258508 +/- 0.0027189\n",
      "PBS gam = 0.0080164 +/- 0.0022532\n"
     ]
    }
   ],
   "source": [
    "pbs_col = allel.pbs(ac1=hapalco_seg[\"col\"], ac2=hapalco_seg[\"gam\"], ac3=hapalco_seg[\"dup\"], window_size=100)\n",
    "pbs_gam = allel.pbs(ac1=hapalco_seg[\"gam\"], ac2=hapalco_seg[\"col\"], ac3=hapalco_seg[\"dup\"], window_size=100)\n",
    "\n",
    "# jack-knifing\n",
    "pbs_col_est = allel.stats.misc.jackknife(pbs_col , np.nanmean)\n",
    "pbs_gam_est = allel.stats.misc.jackknife(pbs_gam , np.nanmean)\n",
    "\n",
    "# report\n",
    "print(\"PBS col = %.7f +/- %.7f\" % (pbs_col_est[0], pbs_col_est[1]))\n",
    "print(\"PBS gam = %.7f +/- %.7f\" % (pbs_gam_est[0], pbs_gam_est[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same, using only duplicated sequences from *col*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PBS col = 0.0206925 +/- 0.0024442\n",
      "PBS gam = 0.0134438 +/- 0.0025498\n"
     ]
    }
   ],
   "source": [
    "pbs_col = allel.pbs(ac1=hapalco_seg[\"col\"], ac2=hapalco_seg[\"gam\"], ac3=hapalco_seg[\"cod\"], window_size=100)\n",
    "pbs_gam = allel.pbs(ac1=hapalco_seg[\"gam\"], ac2=hapalco_seg[\"col\"], ac3=hapalco_seg[\"cod\"], window_size=100)\n",
    "\n",
    "# jack-knifing\n",
    "pbs_col_est = allel.stats.misc.jackknife(pbs_col , np.nanmean)\n",
    "pbs_gam_est = allel.stats.misc.jackknife(pbs_gam , np.nanmean)\n",
    "\n",
    "# report\n",
    "print(\"PBS col = %.7f +/- %.7f\" % (pbs_col_est[0], pbs_col_est[1]))\n",
    "print(\"PBS gam = %.7f +/- %.7f\" % (pbs_gam_est[0], pbs_gam_est[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same, using only duplicated sequences from *gam*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PBS col = 0.0332122 +/- 0.0032000\n",
      "PBS gam = -0.0000764 +/- 0.0020300\n"
     ]
    }
   ],
   "source": [
    "pbs_col = allel.pbs(ac1=hapalco_seg[\"col\"], ac2=hapalco_seg[\"gam\"], ac3=hapalco_seg[\"gad\"], window_size=100)\n",
    "pbs_gam = allel.pbs(ac1=hapalco_seg[\"gam\"], ac2=hapalco_seg[\"col\"], ac3=hapalco_seg[\"gad\"], window_size=100)\n",
    "\n",
    "# jack-knifing\n",
    "pbs_col_est = allel.stats.misc.jackknife(pbs_col , np.nanmean)\n",
    "pbs_gam_est = allel.stats.misc.jackknife(pbs_gam , np.nanmean)\n",
    "\n",
    "# report\n",
    "print(\"PBS col = %.7f +/- %.7f\" % (pbs_col_est[0], pbs_col_est[1]))\n",
    "print(\"PBS gam = %.7f +/- %.7f\" % (pbs_gam_est[0], pbs_gam_est[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot, with a bit of a rolling window and including flanking regions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Subset haps...\n"
     ]
    }
   ],
   "source": [
    "# reload data to include flanks\n",
    "export_start = ace_dups-1e6\n",
    "export_end   = ace_dupe+1e6\n",
    "export_name  = \"duplication\"\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# calculate population allele counts\n",
    "hapalco_sub = haploty_sub.count_alleles_subpops(subpops=popdict_div)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "winsize=5000\n",
    "stepsize=500\n",
    "\n",
    "# window bp\n",
    "windows_pos = allel.moving_statistic(\n",
    "    hapvars_seg[\"POS\"][:],\n",
    "    statistic=lambda v: v[0],\n",
    "    size=winsize,step=stepsize)\n",
    "\n",
    "# accessibility?\n",
    "#inaccessible_pos = np.where(np.invert(accessi_arr))[0]\n",
    "#inaccessible_pos_in_region = inaccessible_pos[np.where(np.logical_and(inaccessible_pos>export_start, inaccessible_pos<export_end))[0]]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# based on all duplications\n",
    "pbs_col = allel.pbs(ac1=hapalco_seg[\"col\"], ac2=hapalco_seg[\"gam\"], ac3=hapalco_seg[\"dup\"], window_size=winsize, window_step=stepsize)\n",
    "pbs_gam = allel.pbs(ac1=hapalco_seg[\"gam\"], ac2=hapalco_seg[\"col\"], ac3=hapalco_seg[\"dup\"], window_size=winsize, window_step=stepsize)\n",
    "\n",
    "plt.figure(figsize=(5,1.5))\n",
    "plt.step(windows_pos/1e6, pbs_gam, color=\"orangered\", label=\"gam to dup\", where=\"post\")\n",
    "plt.step(windows_pos/1e6, pbs_col, color=\"blue\", label=\"col to dup\", where=\"post\")\n",
    "plt.title(\"Distance from duplication genotypes\")\n",
    "plt.ylim([-0.01,0.15])\n",
    "plt.ylabel(\"Distance\")\n",
    "plt.xlabel(\"Mb\")\n",
    "plt.axhline(0, color='k',linestyle=\"--\",label=\"\")\n",
    "plt.axvline(ace_dups/1e6, color='r',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_dupe/1e6, color='r',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_start/1e6, color='orange',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_end/1e6, color='orange',linestyle=\":\",label=\"\")\n",
    "plt.legend()\n",
    "plt.savefig(\"%s/dist_3pops_alldup.pdf\" % (outdir), bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# based on col duplications\n",
    "pbs_col = allel.pbs(ac1=hapalco_seg[\"col\"], ac2=hapalco_seg[\"gam\"], ac3=hapalco_seg[\"cod\"], window_size=winsize, window_step=stepsize)\n",
    "pbs_gam = allel.pbs(ac1=hapalco_seg[\"gam\"], ac2=hapalco_seg[\"col\"], ac3=hapalco_seg[\"cod\"], window_size=winsize, window_step=stepsize)\n",
    "\n",
    "plt.figure(figsize=(5,1.5))\n",
    "plt.step(windows_pos/1e6, pbs_gam, color=\"orangered\", label=\"gam to dup\", where=\"post\")\n",
    "plt.step(windows_pos/1e6, pbs_col, color=\"blue\", label=\"col to dup\", where=\"post\")\n",
    "plt.title(\"Distance from coluzzii duplication genotypes\")\n",
    "plt.ylim([-0.01,0.15])\n",
    "plt.ylabel(\"Distance\")\n",
    "plt.xlabel(\"Mb\")\n",
    "plt.axhline(0, color='k',linestyle=\"--\",label=\"\")\n",
    "plt.axvline(ace_dups/1e6, color='r',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_dupe/1e6, color='r',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_start/1e6, color='orange',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_end/1e6, color='orange',linestyle=\":\",label=\"\")\n",
    "plt.legend()\n",
    "plt.savefig(\"%s/dist_3pops_coldup.pdf\" % (outdir), bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# based on gam duplications\n",
    "pbs_col = allel.pbs(ac1=hapalco_seg[\"col\"], ac2=hapalco_seg[\"gam\"], ac3=hapalco_seg[\"gad\"], window_size=winsize, window_step=stepsize)\n",
    "pbs_gam = allel.pbs(ac1=hapalco_seg[\"gam\"], ac2=hapalco_seg[\"col\"], ac3=hapalco_seg[\"gad\"], window_size=winsize, window_step=stepsize)\n",
    "\n",
    "plt.figure(figsize=(5,1.5))\n",
    "plt.step(windows_pos/1e6, pbs_gam, color=\"orangered\", label=\"dup to gam\", where=\"post\")\n",
    "plt.step(windows_pos/1e6, pbs_col, color=\"blue\", label=\"dup to col\", where=\"post\")\n",
    "plt.title(\"distance from gambiae duplication genotypes\")\n",
    "plt.ylim([-0.01,0.15])\n",
    "plt.ylabel(\"Distance\")\n",
    "plt.xlabel(\"Mb\")\n",
    "plt.axhline(0, color='k',linestyle=\"--\",label=\"\")\n",
    "plt.axvline(ace_dups/1e6, color='r',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_dupe/1e6, color='r',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_start/1e6, color='orange',linestyle=\":\",label=\"\")\n",
    "plt.axvline(ace_end/1e6, color='orange',linestyle=\":\",label=\"\")\n",
    "plt.legend()\n",
    "plt.savefig(\"%s/dist_3pops_gamdup.pdf\" % (outdir), bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export alignment\n",
    "\n",
    "Export `Phylip` alignment of haplotypes.\n",
    "\n",
    "First, the entire duplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dups\n",
    "export_end   = ace_dupe\n",
    "export_name  = \"duplication\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, a region **just downstream** of the duplication breakpoint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dupe\n",
    "export_end   = ace_dupe+5e4\n",
    "export_name  = \"breakdodu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dupe\n",
    "export_end   = ace_dupe+5e3\n",
    "export_name  = \"breakdodu5k\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dupe\n",
    "export_end   = ace_dupe+1e3\n",
    "export_name  = \"breakdodu1k\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, a region **just upstream** of the duplication breakpoint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dups-5e4\n",
    "export_end   = ace_dups\n",
    "export_name  = \"breakupdu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dups-5e3\n",
    "export_end   = ace_dups\n",
    "export_name  = \"breakupdu5k\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dups-1e3\n",
    "export_end   = ace_dups\n",
    "export_name  = \"breakupdu1k\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, an unadmixed region upstream of the duplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dups - 1e6\n",
    "export_end   = export_start + 5e4\n",
    "export_name  = \"upstream\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, a region downstream of the duplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "export_start = ace_dupe + 1e6\n",
    "export_end   = export_start + 5e4\n",
    "export_name  = \"downstream\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load haplotype variants...\n",
      "Load haplotype haplotypes...\n",
      "Drop ploidy haplotypes...\n",
      "Samples dictionary for haps...\n",
      "Allele counts haplotypes...\n",
      "Subset haps...\n",
      "FASTA...\n"
     ]
    }
   ],
   "source": [
    "poplhap = popl\n",
    "\n",
    "# haplotypes: variants\n",
    "hapcall     = zarr.open(haploty_fn)\n",
    "print(\"Load haplotype variants...\")\n",
    "hapcall_var = hapcall[chrom][\"variants\"]\n",
    "hapvars     = allel.VariantChunkedTable(hapcall_var,names=[\"POS\",\"REF\",\"ALT\"],index=\"POS\")\n",
    "hap_bool    = np.logical_and(hapvars[\"POS\"][:] >= export_start, hapvars[\"POS\"][:] <= export_end)\n",
    "hapvars_sub = hapvars.compress(hap_bool)\n",
    "\n",
    "# haplotypes: phased genotypes\n",
    "print(\"Load haplotype haplotypes...\")\n",
    "hapcall_gen = hapcall[chrom][\"calldata/genotype\"]\n",
    "haploty_gen = allel.GenotypeChunkedArray(hapcall_gen)\n",
    "# find samples in haplotype dataset that coincide with genotypes\n",
    "haploty_sam = hapcall[chrom][\"samples\"][:].astype(str)\n",
    "hapsam_bool = np.isin(haploty_sam, np.array(samples_sub[\"ox_code\"]))\n",
    "haploty_sub = haploty_gen.subset(sel0=hap_bool,sel1=hapsam_bool)\n",
    "\n",
    "# recast haplotypes: drop ploidy\n",
    "print(\"Drop ploidy haplotypes...\")\n",
    "haploty_sub_hap = haploty_sub.to_haplotypes()\n",
    "\n",
    "# haplotype dicts\n",
    "# arrays of hap ids and populations of each hap (double the size of genotype arryays: 2 haps per individual except in X chromosome)\n",
    "print(\"Samples dictionary for haps...\")\n",
    "is_samp_in_hap = np.isin(np.array(samples_sub[\"ox_code\"]),haploty_sam)\n",
    "hap_ids        = np.array(list(itertools.chain(*[[s + 'a', s + 'b'] for s in haploty_sam[hapsam_bool]])))\n",
    "hap_pops       = np.array(list(itertools.chain(*[[s, s] for s in np.array(samples_sub[popc][is_samp_in_hap])])))\n",
    "hap_pops_df    = pd.DataFrame(data={ popc : hap_pops , \"ids\" : hap_ids})\n",
    "\n",
    "# pop dicts for haplotype data\n",
    "popdicthap = dict()\n",
    "for popi in poplhap: \n",
    "    popdicthap[popi]  = hap_pops_df[hap_pops_df[popc] == popi].index.tolist()\n",
    "\n",
    "popdicthap[\"all\"] = []\n",
    "for popi in poplhap:\n",
    "    popdicthap[\"all\"] = popdicthap[\"all\"] + popdicthap[popi]\n",
    "\n",
    "# haplotypes: allele counts\n",
    "print(\"Allele counts haplotypes...\")\n",
    "hapalco_sub = haploty_sub_hap.count_alleles_subpops(subpops=popdicthap)\n",
    "\n",
    "# filter haplotypes: segregating alleles, no singletons\n",
    "is_hapseg   = hapalco_sub[\"all\"].is_segregating()[:] # segregating\n",
    "is_hapnosing= hapalco_sub[\"all\"][:,:2].min(axis=1)>2 # no singletons\n",
    "filhap_bool = (is_hapseg[:] & is_hapnosing[:])\n",
    "\n",
    "# subset\n",
    "print(\"Subset haps...\")\n",
    "haploty_seg = haploty_sub_hap.compress(filhap_bool)\n",
    "hapvars_seg = hapvars_sub.compress(filhap_bool)\n",
    "hapalco_seg = hapalco_sub.compress(filhap_bool)\n",
    "\n",
    "# refs and alts\n",
    "hapvars_seg_REF = hapvars_seg[\"REF\"][:].astype(str)\n",
    "hapvars_seg_ALT = hapvars_seg[\"ALT\"][:].astype(str)\n",
    "\n",
    "\n",
    "# output\n",
    "print(\"FASTA...\")\n",
    "happhy = pd.DataFrame({\n",
    "    \"hap\": \">\"+hap_pops_df[\"ids\"]+\"_\"+hap_pops_df[\"population\"],\n",
    "    \"seq\": np.nan},    \n",
    "    columns=[\"hap\", \"seq\"])\n",
    "\n",
    "for pn,popi in enumerate(hap_pops_df[\"ids\"]):\n",
    "    \n",
    "    popi_gen = np.ndarray.tolist(haploty_seg[:,pn])\n",
    "    popi_seq = [hapvars_seg_REF[gn] if gei == 0 else hapvars_seg_ALT[gn] for gn,gei in enumerate(popi_gen)]\n",
    "    happhy[\"seq\"][pn] = ''.join(str(e) for e in popi_seq)\n",
    "\n",
    "happhy.to_csv(\"%s/hapalignment_%s.fasta\" % (outdir,export_name),sep=\"\\n\",index=False, header=False)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}