{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import dask.array as da\n",
    "from LDtools.liftover import *\n",
    "from LDtools.genodata import *\n",
    "from LDtools.sumstat import *\n",
    "from LDtools.ldmatrix import *\n",
    "from LDtools.utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "# Scalable pipeline for computing LD matrix in big sample phenotype"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "kernel": "SoS"
   },
   "source": [
    "### 4 modules\n",
    "- Genodata\n",
    "- Sumstats\n",
    "- Liftover\n",
    "- LDmatrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pip install LDtoolsets`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "lf = Liftover('hg38','hg19')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "vcf ='/home/yh3455/Github/SEQLinkage/MWE/small_sample_ii_coding.vcf.gz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "lf.vcf_liftover(vcf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_Genodata.ipynb.\n",
      "Converted 01_Sumstat.ipynb.\n",
      "Converted 02_Liftover.ipynb.\n",
      "Converted 03_LDmatrix.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "from nbdev.export import notebook2script; notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/yh3455/miniconda3/bin/python\n"
     ]
    }
   ],
   "source": [
    "!which python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "kernel": "SoS"
   },
   "outputs": [],
   "source": [
    "region = [5,272741,1213528-900000]\n",
    "geno_path = 'MWE_region_extraction/ukb23156_c5.merged.filtered.5_272741_1213528.bed'\n",
    "sumstats_path = 'MWE_region_extraction/090321_UKBB_Hearing_aid_f3393_expandedwhite_6436cases_96601ctrl_PC1_2_f3393.regenie.snp_stats'\n",
    "pheno_path = None\n",
    "unr_path = 'MWE_region_extraction/UKB_genotypedatadownloaded083019.090221_sample_variant_qc_final_callrate90.filtered.extracted.white_europeans.filtered.092821_ldprun_unrelated.filtered.prune.txt'\n",
    "imp_geno_path = 'MWE_region_extraction/ukb_imp_chr5_v3_05_272856_1213643.bgen'\n",
    "imp_sumstats_path = 'MWE_region_extraction/100521_UKBB_Hearing_aid_f3393_expandedwhite_15601cases_237318ctrl_500k_PC1_PC2_f3393.regenie.snp_stats.gz'\n",
    "imp_ref = 'hg19'\n",
    "bgen_sample_path = 'MWE_region_extraction/ukb_imp_chr5_v3_05_272856_1213643.sample'\n",
    "output_sumstats = 'test.snp_stats.gz'\n",
    "output_LD = 'test_corr.csv.gz'\n",
    "\n",
    "#main(region,geno_path,sumstats_path,pheno_path,unr_path,imp_geno_path,imp_sumstats_path,imp_ref,output_sumstats,output_LD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "    def main(region,geno_path,sumstats_path,pheno_path,unr_path,imp_geno_path,imp_sumstats_path,imp_ref,output_sumstats,output_LD,bgen_sample_path):\n",
    "\n",
    "        print('1. Preprocess sumstats (regenie format) and extract it from a region')\n",
    "        if pheno_path is not None:\n",
    "            # Load phenotype file\n",
    "            pheno = pd.read_csv(pheno_path, header=0, delim_whitespace=True, quotechar='\"')\n",
    "        if unr_path is not None:\n",
    "            # Load unrelated sample file\n",
    "            unr = pd.read_csv(unr_path, header=0, delim_whitespace=True, quotechar='\"')  \n",
    "        # Load the file of summary statistics and standardize it.\n",
    "        exome_sumstats = Sumstat(sumstats_path)\n",
    "        exome_geno = Genodata(geno_path,bgen_sample_path)\n",
    "\n",
    "        print('1.1. Region extraction')\n",
    "        exome_sumstats.extractbyregion(region)\n",
    "        exome_geno.extractbyregion(region)\n",
    "        exome_sumstats.match_ss(exome_geno.bim)\n",
    "        exome_geno.geno_in_stat(exome_sumstats.ss)\n",
    "\n",
    "        if imp_geno_path is not None:\n",
    "            #two genotype data\n",
    "            imput_sumstats = Sumstat(imp_sumstats_path)\n",
    "            imput_geno = Genodata(imp_geno_path,bgen_sample_path)\n",
    "\n",
    "            if imp_ref is None:\n",
    "                imput_sumstats.extractbyregion(region)\n",
    "                imput_geno.extractbyregion(region)\n",
    "                imput_sumstats.match_ss(imput_geno.bim)\n",
    "                imput_geno.geno_in_stat(imput_sumstats.ss)\n",
    "            else:\n",
    "                print('1.2. LiftOver the region')\n",
    "                hg38toimpref = Liftover('hg38',imp_ref)\n",
    "                imp_region = hg38toimpref.region_liftover(region)\n",
    "                imput_sumstats.extractbyregion(imp_region)\n",
    "                imput_geno.extractbyregion(imp_region)\n",
    "                imput_sumstats.match_ss(imput_geno.bim)\n",
    "                imput_geno.geno_in_stat(imput_sumstats.ss)\n",
    "\n",
    "                print('1.3. Regional SNPs Liftover')\n",
    "                impreftohg38 = Liftover(imp_ref,'hg38') #oppsite with hg38toimpref\n",
    "                imput_geno.bim = impreftohg38.bim_liftover(imput_geno.bim)\n",
    "                imput_sumstats.ss.POS = list(imput_geno.bim.pos)\n",
    "                imput_sumstats.ss.SNP = list(imput_geno.bim.snp)\n",
    "\n",
    "            print('1.1.1 Get exome unique sumstats and geno and Combine sumstats')\n",
    "            snp_match = compare_snps(exome_sumstats.ss,imput_sumstats.ss)\n",
    "            exome_sumstats.ss = exome_sumstats.ss.loc[snp_match.qidx[snp_match.exact==False].drop_duplicates()] #remove by exact match. can be improved.\n",
    "            exome_sumstats.extractbyvariants(imput_sumstats.ss.SNP,notin=True)\n",
    "            exome_geno.geno_in_stat(exome_sumstats.ss)\n",
    "            sumstats = pd.concat([exome_sumstats.ss,imput_sumstats.ss])\n",
    "        else:\n",
    "            #one genotype data\n",
    "            sumstats = exome_sumstats\n",
    "\n",
    "        print('2. Remove relative samples')\n",
    "        if unr_path is not None:\n",
    "            exome_geno.geno_in_unr(unr)\n",
    "            if imp_geno_path is not None:\n",
    "                imput_geno.geno_in_unr(unr)\n",
    "        else:\n",
    "            print('Warning:There is no file of relative sample. All sample are included in computing LD matrix')\n",
    "\n",
    "        if pheno_path is not None:\n",
    "            print('Warning: This function has been implementd yet.')\n",
    "            pass #sld and pld\n",
    "\n",
    "        print('3. Calculate LD matrix')\n",
    "        if imp_geno_path is None:\n",
    "            cor_da = geno_corr(exome_geno.bed.T)\n",
    "        else:\n",
    "            xx = geno_corr(exome_geno.bed.T)\n",
    "            yy = geno_corr(imput_geno.bed.T,step=500)\n",
    "\n",
    "            imput_fam = imput_geno.fam.copy()\n",
    "            imput_fam.index = list(imput_fam.iid.astype(str))\n",
    "            imput_fam['i'] = list(range(imput_fam.shape[0]))\n",
    "            imput_fam_comm = imput_fam.loc[list(exome_geno.fam.iid.astype(str))]\n",
    "            imput_geno.extractbyidx(list(imput_fam_comm.i),row=False)\n",
    "            xy = geno_corr(exome_geno.bed.T,imput_geno.bed.T,step=500)\n",
    "            cor_da = da.concatenate([da.concatenate([xx,xy],axis=1),da.concatenate([xy.T,yy],axis=1)],axis=0)\n",
    "\n",
    "        print('4. Output sumstats and LD matrix')\n",
    "        index = list(sumstats.SNP.apply(shorten_id))\n",
    "        sumstats.SNP = index\n",
    "        sumstats.index = list(range(sumstats.shape[0]))\n",
    "        sumstats.to_csv(output_sumstats, sep = \"\\t\", header = True, index = False,compression='gzip')\n",
    "\n",
    "        corr = cor_da.compute()\n",
    "        np.fill_diagonal(corr, 1)\n",
    "        corr = pd.DataFrame(corr, columns=index)\n",
    "        corr.to_csv(output_LD, sep = \"\\t\", header = True, index = False,compression='gzip')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. Preprocess sumstats (regenie format) and extract it from a region\n",
      "1.1. Region extraction\n",
      "this region [5, 272741, 313528] has 498 SNPs\n",
      "Total SNPs 119 . Flip SNPs 118\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yh3455/miniconda3/lib/python3.8/site-packages/pandas/core/generic.py:5516: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  self[name] = value\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.2. LiftOver the region\n",
      "this region (5, 272856, 313643) has 1736 SNPs\n",
      "Total SNPs 385 . Flip SNPs 0\n",
      "1.3. Regional SNPs Liftover\n",
      "1.1.1 Get exome unique sumstats and geno and Combine sumstats\n",
      "keep   exact  flip   reverse  both   complement\n",
      "False  False  False  False    False  False         106\n",
      "True   False  True   False    False  False          12\n",
      "       True   False  False    False  False           1\n",
      "dtype: int64\n",
      "2. Remove relative samples\n",
      "3. Calculate LD matrix\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/mfs/statgen/yin/Github/LDtools/LDtools/ldmatrix.py:29: RuntimeWarning: invalid value encountered in true_divide\n",
      "  geno_i = (geno_i - np.nanmean(geno_i,axis=0)[None,:])/np.nanstd(geno_i,axis=0)[None,:]\n",
      "/mnt/mfs/statgen/yin/Github/LDtools/LDtools/genodata.py:74: PerformanceWarning: Slicing with an out-of-order index is generating 22098 times more chunks\n",
      "  geno = geno[:,idx]\n",
      "/mnt/mfs/statgen/yin/Github/LDtools/LDtools/ldmatrix.py:67: RuntimeWarning: invalid value encountered in true_divide\n",
      "  geno_i = (geno_i - np.nanmean(geno_i,axis=0)[None,:])/np.nanstd(geno_i,axis=0)[None,:]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4. Output sumstats and LD matrix\n"
     ]
    }
   ],
   "source": [
    "main(region,geno_path,sumstats_path,pheno_path,unr_path,imp_geno_path,imp_sumstats_path,imp_ref,output_sumstats,output_LD,bgen_sample_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
