04-inferring-population-labels.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Institute for Behavioral Genetics International Statistical Genetics 2021 Workshop \n",
    "\n",
    "## Inferring Population Labels\n",
    "\n",
    "Learning objectives:\n",
    "\n",
    "1. Plot principal components obtained by running PCA on our dataset.\n",
    "2. Use principal components to identify population clusters.\n",
    "3. Reidentify populations for samples with missing populations using population clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['PYSPARK_SUBMIT_ARGS'] = '--driver-memory 6G pyspark-shell'\n",
    "\n",
    "import hail as hl\n",
    "hl.plot.output_notebook()\n",
    "hl.init()"
   ]
  },
  {
   "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/HGDP_sample_data.tsv', \n",
    "                     impute=True, \n",
    "                     key='sample_id')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's randomly throw away some of our population information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sa = sa.annotate(\n",
    "    continental_pop = hl.if_else(\n",
    "        hl.rand_bool(0.9),\n",
    "        sa.continental_pop,\n",
    "        hl.missing(hl.tstr)\n",
    "    )\n",
    ")\n",
    "sa.write('output/censored_hgdp_data.ht', overwrite=True)\n",
    "sa = hl.read_table('output/censored_hgdp_data.ht')"
   ]
  },
  {
   "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(**sa[ht.s])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ht.aggregate(hl.agg.collect_as_set(ht.continental_pop))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The populations present in this dataset are `afr` `amr`, `eas`, `fin`, `mid`, `nfe`, `oth`, `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 all combinations of the first three principal components (PCs) against one another. Perhaps we can identify clear cluster boundaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import bokeh\n",
    "\n",
    "p12 = hl.plot.scatter(ht.PC1, ht.PC2, xlabel='PC1', ylabel='PC2', label=ht.continental_pop, size=3, width=400, height=400)\n",
    "p13 = hl.plot.scatter(ht.PC1, ht.PC3, xlabel='PC1', ylabel='PC3', label=ht.continental_pop, size=3, width=400, height=400)\n",
    "\n",
    "p23 = hl.plot.scatter(ht.PC2, ht.PC3, xlabel='PC2', ylabel='PC3', label=ht.continental_pop, size=3, width=400, height=400)\n",
    "\n",
    "hl.plot.show(bokeh.layouts.gridplot([[p12],\n",
    "                                     [p13, p23]]))"
   ]
  },
  {
   "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/HGDP_sample_data.tsv', key='sample_id').cache()\n",
    "\n",
    "def check(ht):\n",
    "    ht = ht.annotate(true_pop = true_labels[ht.s].continental_pop)\n",
    "    c = ht.aggregate(hl.agg.filter(hl.is_missing(ht.continental_pop), \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",
    "### `hl.all` and `hl.any`\n",
    "\n",
    "You might also find it easier to use `hl.all` (which is \"and\") and `hl.any` (which is \"or\"). For example, this\n",
    "\n",
    "```\n",
    "(ht.PC1 > 0.1) | (ht.PC2 < 0.2) | (ht.PC3 >= 0.1)\n",
    "```\n",
    "\n",
    "could also be written as\n",
    "\n",
    "```\n",
    "hl.any(ht.PC1 > 0.1,\n",
    "       ht.PC2 < 0.2,\n",
    "       ht.PC3 >= 0.1)\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.continental_pop)\n",
    ")\n",
    "check(unmasked)"
   ]
  },
  {
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}