{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65947629",
   "metadata": {},
   "source": [
    "This script is used for plotting heatmap of FR family kinship\n",
    "\n",
    "\"an estimated kinship coefficient range >0.354, [0.177, 0.354], [0.0884, 0.177] and [0.0442, 0.0884] corresponds to duplicate/MZ twin, 1st-degree, 2nd-degree, and 3rd-degree relationships respectively. \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cdf08e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0a43b0fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GetHMdata(file, labelFile):\n",
    "    '''\n",
    "    # output of KING software\n",
    "    FID1    ID1     FID2    ID2     N_SNP   HetHet  IBS0    Kinship\n",
    "    P363    P363    P474    P474    6047111 0.2361  0.0057  0.2204\n",
    "    P363    P363    P475    P475    6046933 0.2244  0.0056  0.2238\n",
    "    P363    P363    P478    P478    6050679 0.2324  0.0055  0.2274\n",
    "    '''\n",
    "    # pandas read txt\n",
    "    data = pd.read_table(file, delimiter = '\\t')\n",
    "    labels = pd.read_csv(labelFile, header=None)\n",
    "    xylabels = list(labels[0])\n",
    "    # get xy labels\n",
    "    #xylabels = list(set(list(data[\"FID1\"]) + list(data[\"FID2\"])))\n",
    "    #xylabels.sort()\n",
    "    #xylabels = ['P363', 'P481', 'P474', 'P475', 'P478', 'P479', 'P484', 'P485', 'P487', 'P490', 'P491', 'P492']\n",
    "    # initial kinCoeff Matrix\n",
    "    kinCoeff = np.zeros([len(xylabels),len(xylabels)]) + 0.500\n",
    "    for i in range(len(data['FID1'])):\n",
    "        # Matirx X rows\n",
    "        x = xylabels.index(data['FID1'][i])\n",
    "        # Matrix Y column\n",
    "        y = xylabels.index(data['FID2'][i])\n",
    "        # change KinCoeff in corresponding position, Matrix[x,y] = newvalue\n",
    "        kinCoeff[x,y] = data['Kinship'][i]\n",
    "        kinCoeff[y,x] = data['Kinship'][i]\n",
    "    return xylabels, kinCoeff\n",
    "\n",
    "\n",
    "def heatmap(data, row_labels, col_labels, ax=None,\n",
    "            cbar_kw={}, cbarlabel=\"\", **kwargs):\n",
    "    \"\"\"\n",
    "    Create a heatmap from a numpy array and two lists of labels.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    data\n",
    "        A 2D numpy array of shape (M, N).\n",
    "    row_labels\n",
    "        A list or array of length M with the labels for the rows.\n",
    "    col_labels\n",
    "        A list or array of length N with the labels for the columns.\n",
    "    ax\n",
    "        A `matplotlib.axes.Axes` instance to which the heatmap is plotted.  If\n",
    "        not provided, use current axes or create a new one.  Optional.\n",
    "    cbar_kw\n",
    "        A dictionary with arguments to `matplotlib.Figure.colorbar`.  Optional.\n",
    "    cbarlabel\n",
    "        The label for the colorbar.  Optional.\n",
    "    **kwargs\n",
    "        All other arguments are forwarded to `imshow`.\n",
    "        \n",
    "    Note: modified from online python script resource.\n",
    "    \"\"\"\n",
    "\n",
    "    if not ax:\n",
    "        ax = plt.gca()\n",
    "\n",
    "    # Plot the heatmap\n",
    "    im = ax.imshow(data, **kwargs)\n",
    "\n",
    "    # Create colorbar\n",
    "    cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw, shrink=0.5)\n",
    "    cbar.ax.set_ylabel(cbarlabel, rotation=-90, va=\"bottom\").set_fontsize(20)\n",
    "\n",
    "    # Show all ticks and label them with the respective list entries.\n",
    "    ax.set_xticks(np.arange(data.shape[1]))\n",
    "    ax.set_yticks(np.arange(data.shape[0]))\n",
    "    ax.set_xticklabels(col_labels)\n",
    "    ax.set_yticklabels(row_labels)\n",
    "\n",
    "    # Let the horizontal axes labeling appear on top.\n",
    "    ax.tick_params(top=True, bottom=False,\n",
    "                   labeltop=True, labelbottom=False, labelsize=20)\n",
    "\n",
    "    # Rotate the tick labels and set their alignment.\n",
    "    plt.setp(ax.get_xticklabels(), rotation=-30, ha=\"right\",\n",
    "             rotation_mode=\"anchor\")\n",
    "\n",
    "    # Turn spines off and create white grid.\n",
    "    # unhashed it if necessary\n",
    "    # ax.spines[:].set_visible(False)\n",
    "\n",
    "    ax.set_xticks(np.arange(data.shape[1]+1)-.5, minor=True)\n",
    "    ax.set_yticks(np.arange(data.shape[0]+1)-.5, minor=True)\n",
    "    \n",
    "    ax.grid(which=\"minor\", color=\"w\", linestyle='-', linewidth=3)\n",
    "    ax.tick_params(which=\"minor\", bottom=False, left=False, labelsize=20)\n",
    "\n",
    "    return im, cbar\n",
    "\n",
    "\n",
    "def annotate_heatmap(im, data=None, valfmt=\"{x:.2f}\",\n",
    "                     textcolors=(\"black\", \"black\"),\n",
    "                     threshold=None, **textkw):\n",
    "    \"\"\"\n",
    "    A function to annotate a heatmap.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    im\n",
    "        The AxesImage to be labeled.\n",
    "    data\n",
    "        Data used to annotate.  If None, the image's data is used.  Optional.\n",
    "    valfmt\n",
    "        The format of the annotations inside the heatmap.  This should either\n",
    "        use the string format method, e.g. \"$ {x:.2f}\", or be a\n",
    "        `matplotlib.ticker.Formatter`.  Optional.\n",
    "    textcolors\n",
    "        A pair of colors.  The first is used for values below a threshold,\n",
    "        the second for those above.  Optional.\n",
    "    threshold\n",
    "        Value in data units according to which the colors from textcolors are\n",
    "        applied.  If None (the default) uses the middle of the colormap as\n",
    "        separation.  Optional.\n",
    "    **kwargs\n",
    "        All other arguments are forwarded to each call to `text` used to create\n",
    "        the text labels.\n",
    "        \n",
    "    Note: modified from online python script resource.\n",
    "    \"\"\"\n",
    "\n",
    "    if not isinstance(data, (list, np.ndarray)):\n",
    "        data = im.get_array()\n",
    "\n",
    "    # Normalize the threshold to the images color range.\n",
    "    if threshold is not None:\n",
    "        threshold = im.norm(threshold)\n",
    "    else:\n",
    "        threshold = im.norm(data.max())/2.\n",
    "\n",
    "    # Set default alignment to center, but allow it to be\n",
    "    # overwritten by textkw.\n",
    "    kw = dict(horizontalalignment=\"center\",\n",
    "              verticalalignment=\"center\")\n",
    "    kw.update(textkw)\n",
    "\n",
    "    # Get the formatter in case a string is supplied\n",
    "    if isinstance(valfmt, str):\n",
    "        valfmt = matplotlib.ticker.StrMethodFormatter(valfmt)\n",
    "\n",
    "    # Loop over the data and create a `Text` for each \"pixel\".\n",
    "    # Change the text's color depending on the data.\n",
    "    texts = []\n",
    "    for i in range(data.shape[0]):\n",
    "        for j in range(data.shape[1]):\n",
    "            kw.update(color=textcolors[int(im.norm(data[i, j]) > threshold)])\n",
    "            text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)\n",
    "            texts.append(text)\n",
    "\n",
    "    return texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c0acbdbb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 197340 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "kinshipFile = 'Fam1.kin0'\n",
    "labelFile = 'Fam1.samples.txt'\n",
    "xylabels, kinCoeff = GetHMdata(kinshipFile,labelFile)\n",
    "filename = kinshipFile.split('.')[0]\n",
    "# set up subplot in figure\n",
    "fig, ax = plt.subplots(figsize=(16,12))\n",
    "# image and colorbar\n",
    "im, cbar = heatmap(kinCoeff, xylabels, xylabels, ax=ax, cmap=\"YlGn\", cbarlabel=\"Kinship Coeffient\")\n",
    "texts = annotate_heatmap(im, valfmt=\"{x:.3f}\",fontsize=15)\n",
    "fig.tight_layout()\n",
    "# save figures\n",
    "plt.savefig('SupplFig1_Fam1HeatMap.jpg', format = 'jpeg', dpi=200, pad_inches=0)\n",
    "# also show here\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "61787863",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 208288 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "kinshipFile = 'Fam2.kin0'\n",
    "labelFile = 'Fam2.samples.txt'\n",
    "xylabels, kinCoeff = GetHMdata(kinshipFile,labelFile)\n",
    "filename = kinshipFile.split('.')[0]\n",
    "# set up subplot in figure\n",
    "fig, ax = plt.subplots(figsize=(16,12))\n",
    "# image and colorbar\n",
    "im, cbar = heatmap(kinCoeff, xylabels, xylabels, ax=ax, cmap=\"YlGn\", cbarlabel=\"Kinship Coeffient\")\n",
    "texts = annotate_heatmap(im, valfmt=\"{x:.3f}\",fontsize=15)\n",
    "fig.tight_layout()\n",
    "# save figures\n",
    "plt.savefig('SupplFig2_Fam2HeatMap.jpg', format = 'jpeg', dpi=200, pad_inches=0)\n",
    "# also show here\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ccc3e445",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 199168 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "kinshipFile = 'Fam3.kin0'\n",
    "labelFile = 'Fam3.samples.txt'\n",
    "xylabels, kinCoeff = GetHMdata(kinshipFile,labelFile)\n",
    "filename = kinshipFile.split('.')[0]\n",
    "# set up subplot in figure\n",
    "fig, ax = plt.subplots(figsize=(16,12))\n",
    "# image and colorbar\n",
    "im, cbar = heatmap(kinCoeff, xylabels, xylabels, ax=ax, cmap=\"YlGn\", cbarlabel=\"Kinship Coeffient\")\n",
    "texts = annotate_heatmap(im, valfmt=\"{x:.3f}\",fontsize=15)\n",
    "fig.tight_layout()\n",
    "# save figures\n",
    "plt.savefig('SupplFig3_Fam3HeatMap.jpg', format = 'jpeg', dpi=200, pad_inches=0)\n",
    "# also show here\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7a7f0751",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 92256 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "kinshipFile = 'parents.kin0'\n",
    "labelFile = 'parents.txt'\n",
    "xylabels, kinCoeff = GetHMdata(kinshipFile,labelFile)\n",
    "filename = kinshipFile.split('.')[0]\n",
    "# set up subplot in figure\n",
    "fig, ax = plt.subplots(figsize=(16,12))\n",
    "# image and colorbar\n",
    "im, cbar = heatmap(kinCoeff, xylabels, xylabels, ax=ax, cmap=\"YlGn\", cbarlabel=\"Kinship Coeffient\")\n",
    "texts = annotate_heatmap(im, valfmt=\"{x:.3f}\",fontsize=20)\n",
    "fig.tight_layout()\n",
    "# save figures\n",
    "plt.savefig('SupplFig4_ParentsHeatMap.jpg', format = 'jpeg', dpi=200, pad_inches=0)\n",
    "# also show here\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54eabb9e",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
