jupyter/common-notebooks/hail-plotting-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AWS Hail on EMR Bokeh Plotting Example\n",
    "\n",
    "This is taken from the [Hail Plotting Tutorial](https://hail.is/docs/0.2/tutorials/08-plotting.html) with adjustments for use with SageMaker Notebook instances and EMR."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List EMR Master Nodes\n",
    "\n",
    "`~/SageMaker/bin/list-clusters` will output the IP of each master node in your account and check Livy connectivity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "~/SageMaker/bin/list-clusters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the Livy Endpoint above and start your session name `-s`, language `-l python`, the livy endpoint `-u`, and authentication type `-t`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext sparkmagic.magics\n",
    "%spark add -s aperry -l python -u http://172.21.20.228:8998 -t None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Retrieval from Master\n",
    "\n",
    "Set a local variable for the EMR Master IP.  http requests will be sent to the Master to pull the remote Hail plots back to this notebook server.  Notebook plots will be stored in a `plots` directory along side to your notebook.  If the plots directory does not exist, it will be created.\n",
    "\n",
    "Once the plots have been pulled locally they will be shown inline.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "emrMasterIp = '172.21.20.228'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "from IPython.display import display, HTML\n",
    "import requests, os\n",
    "\n",
    "def displayRemoteHtml(remoteFile):\n",
    "    url = 'http://' + emrMasterIp + remoteFile\n",
    "    req = requests.get(url)\n",
    "    \n",
    "    # Strip leading /\n",
    "    localPath = remoteFile[1:]\n",
    "    # Drop the filename and create local target directory\n",
    "    targetPath = \"/\".join(localPath.split('/')[0:-1])\n",
    "    if not os.path.exists(targetPath):\n",
    "        os.makedirs(targetPath, exist_ok=True)\n",
    "    open(localPath, 'wb').write(req.content)\n",
    "    \n",
    "    html = HTML(filename=localPath)\n",
    "    display(html)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Tutorial\n",
    "\n",
    "The Hail plot module allows for easy plotting of data. This notebook contains examples of how to use the plotting functions in this module, many of which can also be found in the first tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If using Hail version < 0.2.24, uncomment the following\n",
    "# sc.addPyFile('/opt/hail/hail-python.zip')\n",
    "\n",
    "import hail as hl\n",
    "hl.init(sc)\n",
    "\n",
    "from bokeh.io import show, output_notebook\n",
    "from bokeh.layouts import gridplot\n",
    "output_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hl.utils.get_1kg('data/')\n",
    "mt = hl.read_matrix_table('data/1kg.mt')\n",
    "table = (hl.import_table('data/1kg_annotations.txt', impute=True)\n",
    "         .key_by('Sample'))\n",
    "mt = mt.annotate_cols(**table[mt.s])\n",
    "mt = hl.sample_qc(mt)\n",
    "\n",
    "mt.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Histogram\n",
    "\n",
    "The `histogram()` method takes as an argument an aggregated hist expression, as well as optional arguments for the legend and title of the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, output_file, save\n",
    "\n",
    "dp_hist = mt.aggregate_entries(hl.expr.aggregators.hist(mt.DP, 0, 30, 30))\n",
    "p = hl.plot.histogram(dp_hist, legend='DP', title='DP Histogram')\n",
    "output_file(\"/plots/histogram3.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/histogram3.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This method, like all Hail plotting methods, also allows us to pass in fields of our data set directly. Choosing not to specify the `range` and `bins` arguments would result in a range being computed based on the largest and smallest values in the dataset and a default bins value of 50."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cumulative Histogram\n",
    "\n",
    "The `cumulative_histogram()` method works in a similar way to `histogram()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.cumulative_histogram(mt.DP, range=(0,30), bins=30)\n",
    "output_file(\"/plots/cumulative_histogram.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/cumulative_histogram.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scatter\n",
    "\n",
    "The `scatter()` method can also take in either Python types or Hail fields as arguments for x and y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.scatter(mt.sample_qc.dp_stats.mean, mt.sample_qc.call_rate, xlabel='Mean DP', ylabel='Call Rate')\n",
    "output_file(\"/plots/scatter.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/scatter.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also pass in a Hail field as a `label` argument, which determines how to color the data points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_cols((mt.sample_qc.dp_stats.mean >= 4) & (mt.sample_qc.call_rate >= 0.97))\n",
    "ab = mt.AD[1] / hl.sum(mt.AD)\n",
    "filter_condition_ab = ((mt.GT.is_hom_ref() & (ab <= 0.1)) |\n",
    "                        (mt.GT.is_het() & (ab >= 0.25) & (ab <= 0.75)) |\n",
    "                        (mt.GT.is_hom_var() & (ab >= 0.9)))\n",
    "mt = mt.filter_entries(filter_condition_ab)\n",
    "mt = hl.variant_qc(mt).cache()\n",
    "common_mt = mt.filter_rows(mt.variant_qc.AF[1] > 0.01)\n",
    "gwas = hl.linear_regression_rows(y=common_mt.CaffeineConsumption, x=common_mt.GT.n_alt_alleles(), covariates=[1.0])\n",
    "pca_eigenvalues, pca_scores, _ = hl.hwe_normalized_pca(common_mt.GT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.scatter(pca_scores.scores[0], pca_scores.scores[1],\n",
    "                    label=common_mt.cols()[pca_scores.s].SuperPopulation,\n",
    "                    title='PCA', xlabel='PC1', ylabel='PC2', collect_all=True)\n",
    "output_file(\"/plots/scatter2.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/scatter2.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hail's downsample aggregator is incorporated into the `scatter()`, `qq()`, and `manhattan()` functions. The `collect_all` parameter tells the plot function whether to collect all values or downsample. Choosing not to set this parameter results in downsampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.scatter(pca_scores.scores[0], pca_scores.scores[1],\n",
    "                    label=common_mt.cols()[pca_scores.s].SuperPopulation,\n",
    "                    title='PCA', xlabel='PC1', ylabel='PC2', collect_all=True)\n",
    "p2 = hl.plot.scatter(pca_scores.scores[0], pca_scores.scores[1],\n",
    "                    label=common_mt.cols()[pca_scores.s].SuperPopulation,\n",
    "                    title='PCA (downsampled)', xlabel='PC1', ylabel='PC2', collect_all=False, n_divisions=50)\n",
    "\n",
    "# show(gridplot([p, p2], ncols=2, plot_width=400, plot_height=400))\n",
    "grid = gridplot([p, p2], ncols=2, plot_width=400, plot_height=400)\n",
    "output_file(\"/plots/gridplot.html\")\n",
    "save(grid)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/gridplot.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2-D histogram\n",
    "\n",
    "For visualizing relationships between variables in large datasets (where scatter plots may be less informative since they highlight outliers), the `histogram_2d()` function will create a heatmap with the number of observations in each section of a 2-d grid based on two variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.histogram2d(pca_scores.scores[0], pca_scores.scores[1])\n",
    "output_file(\"/plots/2d-histogram.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/2d-histogram.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Q-Q (Quantile-Quantile)\n",
    "\n",
    "The `qq()` function requires either a Python type or a Hail field containing p-values to be plotted. This function also allows for downsampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.qq(gwas.p_value, collect_all=True)\n",
    "p2 = hl.plot.qq(gwas.p_value, n_divisions=75)\n",
    "\n",
    "# show(gridplot([p, p2], ncols=2, plot_width=400, plot_height=400))\n",
    "qq = gridplot([p, p2], ncols=2, plot_width=400, plot_height=400)\n",
    "output_file(\"/plots/qq.html\")\n",
    "save(qq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/qq.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manhattan\n",
    "\n",
    "The `manhattan()` function requires a Hail field containing p-values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.manhattan(gwas.p_value)\n",
    "output_file(\"/plots/manhattan.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/manhattan.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also pass in a dictionary of fields that we would like to show up as we hover over a data point, and choose not to downsample if the dataset is relatively small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hover_fields = dict([('alleles', gwas.alleles)])\n",
    "p = hl.plot.manhattan(gwas.p_value, hover_fields=hover_fields, collect_all=True)\n",
    "output_file(\"/plots/manhattan_hover.html\")\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%local\n",
    "displayRemoteHtml('/plots/manhattan_hover.html')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove the Livy notebook session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%spark cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Sparkmagic (PySpark)",
   "language": "",
   "name": "pysparkkernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "python",
    "version": 3
   },
   "mimetype": "text/x-python",
   "name": "pyspark",
   "pygments_lexer": "python3"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}