ATGUworkshop_Population_Unmasking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Practical 3: Unmasking ancestry labels\n",
    "\n",
    "Objectives:\n",
    "- Plot principal components obtained by running PCA on our dataset.\n",
    "- Use principal components to identify population clusters.\n",
    "- Reidentify populations for samples with missing populations using population clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hail as hl\n",
    "import bokeh\n",
    "from hail.plot import show, output_notebook\n",
    "hl.init()\n",
    "output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in QC'ed data and PCA scores\n",
    "\n",
    "First, we'll need to read back in the sample annotations and the PCA scores from the previous practical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pca_scores = hl.read_table('resources/pca_scores.ht')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sa = hl.import_table('resources/1kg_annotations.txt', \n",
    "                     impute=True, \n",
    "                     key='s')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll take the first 4 PCs from the PCA table, and add the population information for each sample from our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ht = pca_scores.select(PC1=pca_scores.scores[0],\n",
    "                       PC2=pca_scores.scores[1],\n",
    "                       PC3=pca_scores.scores[2],\n",
    "                       PC4=pca_scores.scores[3])\n",
    "ht = ht.annotate(pheno = sa[ht.s])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The five populations present in this dataset are `AFR`, `AMR`, `EAS`, `EUR`, and `SAS`. They are three-letter codes from the 1000 Genomes project denoting the [super population of each sample](https://www.internationalgenome.org/category/population/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize!\n",
    "\n",
    "Let's plot several combinations of the first four principal components (PCs) against each other. This will help us visualize the population structure of our dataset, and allow us to try identify our masked samples with different population clusters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1 = hl.plot.scatter(ht.PC1, ht.PC2, xlabel='PC1', ylabel='PC2', label=ht.pheno.super_population, size=6)\n",
    "p2 = hl.plot.scatter(ht.PC1, ht.PC3, xlabel='PC1', ylabel='PC3', label=ht.pheno.super_population, size=6)\n",
    "p3 = hl.plot.scatter(ht.PC2, ht.PC4, xlabel='PC2', ylabel='PC4', label=ht.pheno.super_population, size=6)\n",
    "\n",
    "show(bokeh.layouts.gridplot([[p1], [p2], [p3]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reidentify samples with missing ancestry based on PCA scores\n",
    "\n",
    "Now that we can see how the populations are decomposed by the PCs, let's try to reidentify the masked samples.\n",
    "\n",
    "First, we'll define a grading scheme to check against the true populations of each masked sample. (The `check` function will see how many masked samples you have correctly identified.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_labels = hl.import_table('resources/true_pops.txt', key='s').cache()\n",
    "def check(ht):\n",
    "    ht = ht.annotate(true_pop = true_labels[ht.s].real_super_population)\n",
    "    c = ht.aggregate(hl.agg.filter(hl.is_missing(ht.pheno.super_population), \n",
    "                                   hl.agg.counter((ht.unmasked, ht.true_pop))))\n",
    "    n_correct = sum(count for k, count in c.items() if k[0] == k[1])\n",
    "    n_wrong = sum(count for k, count in c.items() if k[0] != k[1])\n",
    "    print(f'Correctly identified {n_correct} / {n_correct + n_wrong} masked samples.')\n",
    "    print()\n",
    "    \n",
    "    for (unm, true), n in c.items():\n",
    "        if unm != true:\n",
    "            if unm is not None:\n",
    "                print(f'Incorrectly assigned {n} {true} samples as {unm}.')\n",
    "            else:\n",
    "                print(f'Left {n} {true} samples unassigned.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fill in the below\n",
    "\n",
    "Your job is to expand the below code to reidentify the population labels. One of the populations has already been provided as an example.\n",
    "\n",
    "### `case().when()` in Hail\n",
    "\n",
    "The `case` / `when` / `default` motif you see below is a nice way to write `if` / `else if` / `else`. The returned `unmasked` will be equal to the result of the first `when` whose predicate is `True`.\n",
    "\n",
    "### A note on `&` and `|`\n",
    "\n",
    "Python uses `and` and `or` for logical operators. Hail expressions use `&` for 'and' and `|` for or.\n",
    "\n",
    "This can lead to some confusion, especially since `&` and `|` often don't play nicely with expressions involving `>`, `<`, `==`, or `!=`. If both of these operators appear, you will need to wrap the comparison in parentheses.\n",
    "\n",
    "Suppose we want to write code that returns true when \"PC1 is greater than 0.1 or PC2 is less than 0.2\":\n",
    "\n",
    "**correct**:\n",
    "\n",
    "```\n",
    "(ht.PC1 > 0.1) | (ht.PC2 < 0.2)\n",
    "```\n",
    "\n",
    "**incorrect**:\n",
    "```\n",
    "ht.PC1 > 0.1 or ht.PC2 < 0.2\n",
    "ht.PC1 > 0.1 | ht.PC2 < 0.2\n",
    "(ht.PC1 > 0.1) or (ht.PC2 < 0.2)\n",
    "```\n",
    "\n",
    "### To think about\n",
    "\n",
    "Which population is hardest to reidentify? Why?\n",
    "\n",
    "### Extras\n",
    "\n",
    "Try plotting the PCs again with the re-identified population labels.\n",
    "\n",
    "Try plotting the PCs again, highlighting the ones that you missed. (The true population labels are in the table `true_labels`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unmasked = ht.annotate(\n",
    "    unmasked = hl.case()\n",
    "        .when((ht.PC2 > 0.2) & (ht.PC2 > 0.2), 'EAS')\n",
    "#         .when(..., 'AFR')\n",
    "#         .when(..., 'AMR')\n",
    "#         .when(..., 'EUR')\n",
    "#         .when(..., 'SAS')\n",
    "        .default(ht.pheno.super_population)\n",
    ")\n",
    "check(unmasked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}