PBSTajD_acrossChrom.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the modules\n",
    "import allel\n",
    "import zarr\n",
    "import numcodecs\n",
    "import numpy as np\n",
    "import sys\n",
    "import pandas as pd\n",
    "import itertools\n",
    "import ipytree\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_zarr1kgp(chrom, outfold):\n",
    "#     chrom = \"1\"\n",
    "#     outfold=\"data\"\n",
    "    # VCF Path\n",
    "    vcf_path= outfold + '/ALL.chr' + chrom + '.phase3_shapeit2_mvncall_integrated_v5a.20130502.genotypes.vcf.gz'\n",
    "    callset = allel.read_vcf(vcf_path, fields=['numalt'], log=sys.stdout)\n",
    "    zarr_path = outfold + '/ALL.chr' + chrom + '.phase3_shapeit2_mvncall_integrated_v5a.20130502.genotypes.zarr'\n",
    "    print(\"converting to zarr chrom \"+ chrom+ \" take a coffee...\")\n",
    "    allel.vcf_to_zarr(vcf_path, zarr_path, group=chrom,fields='*', alt_number=1, log=sys.stdout,\n",
    "                  overwrite=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_PBSTaj(chrom, outfold, taj_window):\n",
    "    zarr_path = outfold + '/ALL.chr' + chrom + '.phase3_shapeit2_mvncall_integrated_v5a.20130502.genotypes.zarr'\n",
    "    callset = zarr.open_group(zarr_path, mode='r')\n",
    "    pos = allel.SortedIndex(callset[ chrom + '/variants/POS'])\n",
    "    rs = callset[ chrom + '/variants/ID'][:]\n",
    "    gt_dask = allel.GenotypeDaskArray(callset[chrom + '/calldata/GT'])\n",
    "    multiall_only = callset[chrom + '/variants/MULTI_ALLELIC'][:]\n",
    "    SNPs_only = callset[chrom + '/variants/is_snp'][:]\n",
    "    loc_variants = ~multiall_only & SNPs_only\n",
    "    pos_forTab = pos[loc_variants]\n",
    "    rs_forTab = rs[loc_variants]\n",
    "    pos_forTabv2 = [chrom + \"_\" + str(s) for s in pos_forTab]\n",
    "    gt_variant_selection = gt_dask.compress(loc_variants, axis=0)\n",
    "    panel_path='data/integrated_call_samples_v3.20130502.ALL.panel'\n",
    "    panel = pd.read_csv(panel_path, sep='\\t', usecols=['super_pop', 'pop', 'sample'],encoding='utf-8-sig')\n",
    "    samples = callset[chrom + '/samples'][:]\n",
    "    # Test if the samples in the genotype have the same order in the panel loaded\n",
    "    if np.all(samples == panel['sample'].values):\n",
    "        samples_list = list(samples)\n",
    "        samples_callset_index = [samples_list.index(s) for s in panel['sample']]\n",
    "        panel['callset_index'] = samples_callset_index\n",
    "\n",
    "        loc_AFR=panel[panel.super_pop == 'AFR'].callset_index.values\n",
    "        loc_YRI=panel[panel['pop'] == 'YRI'].callset_index.values\n",
    "        loc_GBR=panel[panel['pop'] == 'GBR'].callset_index.values\n",
    "        loc_IBS=panel[panel['pop'] == 'IBS'].callset_index.values\n",
    "        loc_TSI=panel[panel['pop'] == 'TSI'].callset_index.values\n",
    "        loc_CHB=panel[panel['pop'] == 'CHB'].callset_index.values\n",
    "\n",
    "        print(\"converting to gt single pops.. in chrom\" + chrom)\n",
    "        gt_IBS= gt_variant_selection.take(loc_IBS, axis=1).compute()\n",
    "        gt_TSI= gt_variant_selection.take(loc_TSI, axis=1).compute()\n",
    "        gt_CHB= gt_variant_selection.take(loc_CHB, axis=1).compute()\n",
    "        gt_YRI= gt_variant_selection.take(loc_YRI, axis=1).compute()\n",
    "\n",
    "        print(\"converting to allele counts... in chrom\" + chrom)\n",
    "        ac_IBS_du = gt_IBS.count_alleles()\n",
    "        ac_TSI_du = gt_TSI.count_alleles()\n",
    "        ac_CHB_du = gt_CHB.count_alleles()\n",
    "        ac_YRI_du = gt_YRI.count_alleles()\n",
    "\n",
    "        # Make the PBS\n",
    "        print(\"Make the PBS of chrom \" + chrom + \"TSI vs CHB...\")\n",
    "        PBS_TSI_CHB_YRI_2all_noNormed = allel.pbs(ac_TSI_du, ac_CHB_du, ac_YRI_du, window_size=1, normed=False)\n",
    "\n",
    "        # Make the PBS\n",
    "        print(\"Make the PBS of chrom \" + chrom + \"IBS vs CHB...\")\n",
    "        PBS_IBS_CHB_YRI_2all_noNormed = allel.pbs(ac_IBS_du ,ac_CHB_du, ac_YRI_du, window_size=1, normed=False)\n",
    "\n",
    "        # Make the PBS\n",
    "        print(\"Make the PBS of chrom \" + chrom + \"CHB vs TSI...\")\n",
    "        PBS_CHB_TSI_YRI_2all_noNormed = allel.pbs(ac_CHB_du, ac_TSI_du, ac_YRI_du, window_size=1, normed=False)\n",
    "\n",
    "        # Make the PBS\n",
    "        print(\"Make the PBS of chrom \" + chrom + \"CHB vs IBS...\")\n",
    "        PBS_CHB_IBS_YRI_2all_noNormed = allel.pbs(ac_CHB_du, ac_IBS_du, ac_YRI_du, window_size=1, normed=False)\n",
    "        \n",
    "        # Make the Tajima\n",
    "        print(\"Making the tajima\")\n",
    "        TajD_TSI_YRI_2all = allel.moving_delta_tajima_d(ac_TSI_du, ac_YRI_du, size=taj_window)\n",
    "\n",
    "\n",
    "        print(\"Saving the output of chrom \"+ chrom)\n",
    "        # Save the results in a table to be viewed \n",
    "        res_df= pd.DataFrame( {'PBS': PBS_TSI_CHB_YRI_2all_noNormed, 'rs' : rs_forTab , 'pos': pos_forTabv2})\n",
    "        res_df.to_csv('PBS_scikit_'+chrom+ 'TSI' + 'CHB' + 'YRI' +'v2.txt', sep= \"\\t\", na_rep= 'NaN')\n",
    "\n",
    "        res_df= pd.DataFrame( {'PBS': PBS_IBS_CHB_YRI_2all_noNormed, 'rs' : rs_forTab , 'pos': pos_forTabv2})\n",
    "        res_df.to_csv('PBS_scikit_'+chrom + 'IBS' + 'CHB' + 'YRI' +'v2.txt', sep= \"\\t\", na_rep= 'NaN')\n",
    "\n",
    "        res_df= pd.DataFrame( {'PBS': PBS_CHB_TSI_YRI_2all_noNormed, 'rs' : rs_forTab , 'pos': pos_forTabv2})\n",
    "        res_df.to_csv('PBS_scikit_'+chrom + 'CHB' + 'TSI' + 'YRI' +'v2.txt', sep= \"\\t\", na_rep= 'NaN')\n",
    "\n",
    "        res_df= pd.DataFrame( {'PBS': PBS_CHB_IBS_YRI_2all_noNormed, 'rs' : rs_forTab , 'pos': pos_forTabv2})\n",
    "        res_df.to_csv('PBS_scikit_'+chrom + 'CHB' + 'IBS' + 'YRI' +'v2.txt', sep= \"\\t\", na_rep= 'NaN')\n",
    "        \n",
    "        # Save Tajima D\n",
    "        start_every5K = pos_forTab[::taj_window][:-1]\n",
    "        end_every5K = (pos_forTab[::taj_window][1:]-1)\n",
    "        \n",
    "        pos_taj = [str(s) + \"_\" + str(e) for s,e in zip(start_every5K, end_every5K)]\n",
    "        \n",
    "         # Save the TajD Table\n",
    "        res_df= pd.DataFrame(TajD_TSI_YRI_2all, index= pos_taj, columns= ['TajD'])\n",
    "        res_df['chromosome'] = chrom\n",
    "        res_df = res_df.rename_axis(\"regions\").reset_index()\n",
    "\n",
    "        res_df.to_csv('TajD_scikit_' + chrom + 'TSI' + 'ogYRI' + 'v2.txt', sep= \"\\t\", na_rep= 'NaN')\n",
    "\n",
    "\n",
    "\n",
    "        print(\"Job finished, the dinner is on the Tables!!\")\n",
    "    else:\n",
    "        print(\"Problem with the samples name! Look the input\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for c in range(1,23):\n",
    "    convert_zarr1kgp(str(c), \"data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for c in range(1,23):\n",
    "    calculate_PBSTaj(str(c), \"data\", taj_window= 5000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PBS",
   "language": "python",
   "name": "pbs"
  },
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}