{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp ldmatrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# LDmatrix module\n",
    "\n",
    "> Calculate LDmatrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import dask.array as da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def geno_corr(x,y=None,step=100):\n",
    "    '''functions to calculate LD matrix'''\n",
    "    if y is None:\n",
    "        dd = dask_corr(x,step)\n",
    "        return(dict2mat(dd))\n",
    "    else:\n",
    "        dd = dask_corr_pair(x,y,step)\n",
    "        return(dict2mat_pair(dd))\n",
    "\n",
    "def dask_corr(genos,step=100):\n",
    "    #sample by snps (normalized)\n",
    "    nsample = genos.shape[0]\n",
    "    nsnp = genos.shape[1]\n",
    "    da_corr = {}\n",
    "    for i in range(0,nsnp,step):\n",
    "        da_corr[i] = {}\n",
    "        geno_i = genos[:,i:min(i+step,nsnp)].compute().astype(np.float64)\n",
    "        geno_i = (geno_i - np.nanmean(geno_i,axis=0)[None,:])/np.nanstd(geno_i,axis=0)[None,:]\n",
    "        geno_i = torch.from_numpy(geno_i)\n",
    "        geno_i[torch.isnan(geno_i)] = 0\n",
    "        chunk_i = da.from_array((torch.matmul(geno_i.T,geno_i)/nsample).numpy())\n",
    "        da_corr[i][i]=chunk_i\n",
    "        for j in range(i+step,nsnp,step):\n",
    "            geno_j = genos[:,j:min(j+step,nsnp)].compute().astype(np.float64)\n",
    "            geno_j = (geno_j - np.nanmean(geno_j,axis=0)[None,:])/np.nanstd(geno_j,axis=0)[None,:]\n",
    "            geno_j = torch.from_numpy(geno_j)\n",
    "            geno_j[torch.isnan(geno_j)] = 0\n",
    "            cor_ij = da.from_array((torch.matmul(geno_i.T,geno_j)/nsample).numpy())\n",
    "            da_corr[i][j]=cor_ij\n",
    "    return da_corr\n",
    "\n",
    "def dict2mat(dd):\n",
    "    da_mat=[]\n",
    "    for i in dd.keys():\n",
    "        rowi = []\n",
    "        for j in dd.keys():\n",
    "            if i>j:\n",
    "                rowi.append(dd[j][i].T)\n",
    "            else:\n",
    "                rowi.append(dd[i][j])\n",
    "        rowi = da.concatenate(rowi,axis=1)\n",
    "        da_mat.append(rowi)\n",
    "    return(da.concatenate(da_mat,axis=0))\n",
    "\n",
    "def dask_corr_pair(genos,pgenos,step=100):\n",
    "    #sample by snps (normalized)\n",
    "    nsample = genos.shape[0]\n",
    "    nsnp = genos.shape[1]\n",
    "    psample = pgenos.shape[0]\n",
    "    psnp = pgenos.shape[1]\n",
    "    if nsample != psample: print(\"error: sample not match\")\n",
    "    da_corr = {}\n",
    "    for i in range(0,nsnp,step):\n",
    "        da_corr[i] = {}\n",
    "        geno_i = genos[:,i:min(i+step,nsnp)].compute().astype(np.float64)\n",
    "        geno_i = (geno_i - np.nanmean(geno_i,axis=0)[None,:])/np.nanstd(geno_i,axis=0)[None,:]\n",
    "        geno_i = torch.from_numpy(geno_i)\n",
    "        geno_i[torch.isnan(geno_i)] = 0\n",
    "        for j in range(0,psnp,step):\n",
    "            geno_j = pgenos[:,j:min(j+step,psnp)].compute().astype(np.float64)\n",
    "            geno_j = (geno_j - np.nanmean(geno_j,axis=0)[None,:])/np.nanstd(geno_j,axis=0)[None,:]\n",
    "            geno_j = torch.from_numpy(geno_j)\n",
    "            geno_j[torch.isnan(geno_j)] = 0\n",
    "            cor_ij = da.from_array((torch.matmul(geno_i.T,geno_j)/nsample).numpy())\n",
    "            da_corr[i][j]=cor_ij\n",
    "    return da_corr\n",
    "\n",
    "def dict2mat_pair(dd):\n",
    "    da_mat=[]\n",
    "    for i in dd.keys():\n",
    "        rowi = []\n",
    "        for j in dd[0].keys():\n",
    "            rowi.append(dd[i][j])\n",
    "        rowi = da.concatenate(rowi,axis=1)\n",
    "        da_mat.append(rowi)\n",
    "    return(da.concatenate(da_mat,axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
