Pandas-for-Population-Structure-Barplots.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7cd6189",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This Jupyter notebook will demonstrate how to use python and pandas to make a publication-quality \n",
    "# barplot figure from an Admixture analysis. The required input files are the ped \n",
    "# file used to run admixture and an output Q file from the analysis.\n",
    "\n",
    "# To run this notebook, you will need to specify the paths to these files. \n",
    "# You will also need to have the python packages pandas (https://pandas.pydata.org/)\n",
    "# and seaborn (https://seaborn.pydata.org/) available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beea2df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64d83d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################################################\n",
    "# First step is to read in the PED and Q files as pandas dataframes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f39aefb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# full path to the PED file used for Admixture\n",
    "fileped = '/Users/portik/Dropbox/Trachylepis/Trachylepis_recoded.ped'\n",
    "\n",
    "# full path to the Q file to use for plotting\n",
    "fileq = '/Users/portik/Dropbox/Trachylepis/Trachylepis_recoded.4.Q'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2bcbab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in ped file as dataframe, this file is delimited by whitespace\n",
    "df_ped = pd.read_csv(fileped, sep=' ', header=None)\n",
    "\n",
    "# read in the Q file, which is also whitespace delimited \n",
    "df_q = pd.read_csv(fileq, sep=' ', header=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1d38e7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is what the PED dataframe looks like\n",
    "df_ped\n",
    "\n",
    "# the sample names in a PED file should be in the first column:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c33bfbd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# and this is what a Q file dataframe looks like\n",
    "df_q\n",
    "\n",
    "# notice there are no sample names here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffae965a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################################################\n",
    "# We now need to make some edits to the Q dataframe to prepare it for plotting. \n",
    "# First, it needs column labels and the sample names. After that, we need \n",
    "# to assign each sample to a population, and add that as another \n",
    "# column. With all of that information, we can make the high quality barplots!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "683b1094",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we are adding labels to the population proportion columns, and \n",
    "# adding a new column with the sample names from the ped file.\n",
    "\n",
    "# automatically generate column names based on number of columns in Q (pop1, pop2, pop3, etc.)\n",
    "names = [\"pop{}\".format(i) for i in range(1, df_q.shape[1]+1)]\n",
    "\n",
    "# add column names to dataframe\n",
    "df_q.columns = names\n",
    "\n",
    "# insert the sample names into the first column position\n",
    "df_q.insert(0, 'Sample', df_ped[0])\n",
    "\n",
    "# now  set the dataframe index to the sample names (e.g., the 'Sample' column)\n",
    "df_q.set_index('Sample', inplace=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b062c589",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check what this did to the q dataframe:\n",
    "df_q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908cd325",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next we want to assign each individual to a population, based on highest proportion of ancestry.\n",
    "# For each row, we find the column with the greatest value, then assign it that column label value\n",
    "# in a new column called 'assignment'.\n",
    "df_q['assignment'] = df_q.idxmax(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b7bca9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can see this new column now, with the assignments provided\n",
    "df_q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7564170f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################################################\n",
    "# Now it is time to make the plots. First up is choosing a color scheme, followed \n",
    "# by some stacked barplots. We will first create a plot from an unsorted Q \n",
    "# dataframe. Then, we will use a custom sorting function to produce better looking \n",
    "# plots, and save them as PDF files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51dd9029",
   "metadata": {},
   "outputs": [],
   "source": [
    "# To make plotting easier we can set the color palette here. \n",
    "\n",
    "# For palette ideas and color schemes, please refer to the seaborn \n",
    "# documentation: https://seaborn.pydata.org/tutorial/color_palettes.html\n",
    "# I really like the Paired, tab10, tab20, and tab20c palettes. To set those \n",
    "# instead, you could use:\n",
    "# pal = sns.color_palette(\"Paired\", 4)\n",
    "# The number following the palette name specifies the number of colors to select.\n",
    "\n",
    "# I chose a color palette based on what I found using: https://colorbrewer2.org/\n",
    "\n",
    "pal = sns.color_palette(['#ef8a62','#92c5de','#fddbc7','#0571b0'])\n",
    "pal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2882e119",
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################################################\n",
    "# We will start by creating a stacked barplot for the unsorted Q dataframe.\n",
    "# The plotting arguments are all explained below, and some values may need to be adjusted \n",
    "# based on your dataset qualities (for example, the width and height of plot, and \n",
    "# sample label sizes).\n",
    "\n",
    "\n",
    "\n",
    "# Create the stacked bar plot using the pandas .plot.bar() function with the following args:\n",
    "\n",
    "# stacked: controls whether or not to create a stacked barplot\n",
    "# figsize: the dimensions of the figure, in (width, height)\n",
    "# width: this controls the width of the bars, 1 = no separation between them, <1 puts spacing in\n",
    "# color: sets the colors, we refer to pal which was set above\n",
    "# fontsize: the size of sample labels (values = xx-small, x-small, small, medium, large, x-large, xx-large)\n",
    "# edgecolor: controls the color of lines separating colors within bars, as well as bar outline\n",
    "# linewidth: the line size for the edgecolor\n",
    "\n",
    "ax = df_q.plot.bar(stacked=True, \n",
    "                    figsize=(25,5), \n",
    "                    width=1,\n",
    "                    color=pal, \n",
    "                    fontsize='x-small',\n",
    "                    edgecolor='black', \n",
    "                    linewidth=0.5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b43728d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can clean up some aspects of this plot using some additional commands: \n",
    "# 1) remove the bounding box around the figure,\n",
    "# 2) rotate the sample labels,\n",
    "# 3) change legend placement and look.\n",
    "\n",
    "ax = df_q.plot.bar(stacked=True, \n",
    "                    figsize=(25,5), \n",
    "                    width=1,\n",
    "                    color=pal, \n",
    "                    fontsize='x-small',\n",
    "                    edgecolor='black', \n",
    "                    linewidth=0.5)\n",
    "\n",
    "# 1) these commands eliminate the bounding box for the barplot\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "\n",
    "# 2) this controls rotation of the sample names along the x-axis\n",
    "ax.set_xticklabels(df_q.index, rotation=45, ha='right')\n",
    "\n",
    "# 3) this controls the placement of the legend, as well as \n",
    "# font (fontsize), spacing (labelspacing), and bounding box (frameon)\n",
    "ax.legend(bbox_to_anchor=(1,1), fontsize='medium', labelspacing=0.5, frameon=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ffbaab",
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################################################\n",
    "# That's looking better! However the main issue is that the populations need to be sorted. \n",
    "# This can be attempted using a built-in sort method with pandas. It sorts in the order of \n",
    "# the columns provided in a list.\n",
    "\n",
    "df_q_sorted = df_q.sort_values(['assignment', 'pop1','pop2','pop3','pop4'], ascending=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89968b1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see how well this method was able to sort:\n",
    "\n",
    "ax = df_q_sorted.plot.bar(stacked=True, \n",
    "                    figsize=(25,5), \n",
    "                    width=1,\n",
    "                    color=pal, \n",
    "                    fontsize='x-small',\n",
    "                    edgecolor='black', \n",
    "                    linewidth=0.5)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "ax.set_xticklabels(df_q_sorted.index, rotation=45, ha='right')\n",
    "ax.legend(bbox_to_anchor=(1,1), fontsize='medium', labelspacing=0.5, frameon=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14f63bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# It's looking better, but the plot isn't quite right. There are a few individuals \n",
    "# that seem out of place. \n",
    "\n",
    "\n",
    "#############################################################################\n",
    "# A better way to sort is to pull out each population subset from the dataframe,\n",
    "# and sort them individually. After each subset is sorted, they can be merged \n",
    "# into a complete dataframe again.\n",
    "\n",
    "# Here is a function for better sorting of the dataframes.\n",
    "def sort_df_by_pops(df):\n",
    "    temp_dfs = []\n",
    "    for pop in sorted(df['assignment'].unique()):\n",
    "        temp = df.loc[df['assignment'] == pop].sort_values(by=[pop], ascending=False)\n",
    "        temp_dfs.append(temp)\n",
    "    return pd.concat(temp_dfs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c62054b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can now use the function to sort the Q dataframe:\n",
    "\n",
    "df_sorted_q = sort_df_by_pops(df_q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00f2dfc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see what this did:\n",
    "df_sorted_q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcd1d164",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we can make the barplot for the sorted Q dataframe.\n",
    "\n",
    "# make the stacked barplot\n",
    "ax = df_sorted_q.plot.bar(stacked=True, \n",
    "                          figsize=(25,5), \n",
    "                          width=1,\n",
    "                          color=pal, \n",
    "                          fontsize='x-small',\n",
    "                          edgecolor='black', \n",
    "                          linewidth=0.5)\n",
    "\n",
    "# these commands eliminate the bounding box for the barplot\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "\n",
    "# this controls rotation of the sample names along the x-axis\n",
    "ax.set_xticklabels(df_sorted_q.index, rotation=45, ha='right')\n",
    "\n",
    "# this controls the placement of the legend, as well as \n",
    "# font (fontsize), spacing (labelspacing), and bounding box (frameon)\n",
    "ax.legend(bbox_to_anchor=(1,1), fontsize='medium', labelspacing=0.5, frameon=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32c81f99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Looks really good!\n",
    "\n",
    "# We can save this plot using the command below:\n",
    "ax.figure.savefig('Admixture-K4.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eedd47bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can also save the Q dataframe we created. \n",
    "\n",
    "df_sorted_q.to_csv('Sorted-Q4-data.txt', sep=\",\", index=True, float_format='%.6f')\n",
    "\n",
    "# Changing the sep= value to ' ' or '\\t' will produce white-space or tab-delimited output, \n",
    "# rather than a comma-delimited csv output. The float_format='%.6f' ensures that the same \n",
    "# number of decimals will be written for the floats (ancestry proportions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e2bb8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "###########################################################################################\n",
    "# A final advanced topic - changing the order of populations on the figure!\n",
    "\n",
    "# We may wish to shift the order of the populations, and this is also easy to do.\n",
    "\n",
    "# It takes a slight edit to the sort function. Here it will return a list of the sorted\n",
    "# subdataframes, rather than the concatenated sorted subdataframes:\n",
    "\n",
    "def sort_df_by_pops_nocat(df):\n",
    "    temp_dfs = []\n",
    "    for pop in sorted(df['assignment'].unique()):\n",
    "        temp = df.loc[df['assignment'] == pop].sort_values(by=[pop], ascending=False)\n",
    "        temp_dfs.append(temp)\n",
    "    return temp_dfs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d758903d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A list of the subdataframes is returned:\n",
    "sub_dfs = sort_df_by_pops_nocat(df_q)\n",
    "\n",
    "# There are four subdataframes in this list, and we can merge them in any order \n",
    "# by using the index of the list with the concatenate function. Let's say\n",
    "# we want to move population one to the end of the plot, and shift population two\n",
    "# to be in between three and four:\n",
    "\n",
    "df_custom_sort = pd.concat([sub_dfs[2], sub_dfs[1], sub_dfs[3], sub_dfs[0]])\n",
    "\n",
    "# Remember python indices are zero-based!\n",
    "\n",
    "# If your K value is higher or lower, make sure to include all the subdataframes in \n",
    "# the list for pd.concat() . For example with K=6, you'll need sub_dfs[0] to sub_dfs[5]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79bd0b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we can make the barplot for the custom sorted Q dataframe.\n",
    "\n",
    "ax = df_custom_sort.plot.bar(stacked=True, \n",
    "                             figsize=(25,5), \n",
    "                             width=1,\n",
    "                             color=pal, \n",
    "                             fontsize='x-small',\n",
    "                             edgecolor='black', \n",
    "                             linewidth=0.5)\n",
    "\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "ax.set_xticklabels(df_sorted_q.index, rotation=45, ha='right')\n",
    "ax.legend(bbox_to_anchor=(1,1), fontsize='medium', labelspacing=0.5, frameon=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61772906",
   "metadata": {},
   "outputs": [],
   "source": [
    "# That looks even better than the automated sorting function, definitely going to save this!\n",
    "ax.figure.savefig('Admixture-K4-Perfection.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37a838c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And that's it! Hopefully this will help save some time and effort for an upcoming\n",
    "# phylogeography figure. Pandas is extremely useful for manipulating dataframes and \n",
    "# creating plots, and this is one example of a useful application."
   ]
  }
 ],
 "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}