notebooks/Hail_tutorial-GWAS-vcf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6252f42-364b-45c7-ad3d-b9cc72873f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "import datetime\n",
    "import hail as hl\n",
    "\n",
    "from hail.plot import show\n",
    "import pandas as pd\n",
    "from pprint import pprint\n",
    "hl.plot.output_notebook()\n",
    "\n",
    "from pyspark.conf import SparkConf\n",
    "from pyspark.context import SparkContext\n",
    "from pyspark.sql.session import SparkSession"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6679e50f-9502-453d-85ba-3b38bbae1e3a",
   "metadata": {},
   "source": [
    "#### Start an [Apache Spark](https://en.wikipedia.org/wiki/Apache_Spark) instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0ebe51c-4ab9-451e-a32f-226e5e477c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_file_name = f\"logs/hail-{datetime.datetime.now():%Y-%m-%d-%H-%M-%S}.log\"\n",
    "# run spark\n",
    "spark_conf = SparkConf().setAppName(\"hail-test\")\n",
    "# .setMaster(\"spark://spark-master:7077\")\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.endpoint\", \"http://lifemap-minio:9000/\")\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.access.key\", \"root\")\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.secret.key\", \"passpass\" )\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.connection.ssl.enabled\", \"false\")\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.path.style.access\", \"true\")\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.impl\", \"org.apache.hadoop.fs.s3a.S3AFileSystem\")\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.connection.maximum\", 1024);\n",
    "spark_conf.set(\"spark.hadoop.fs.s3a.threads.max\", 1024);\n",
    "spark_conf.set(\"spark.hadoop.fs.s3.impl\", \"org.apache.hadoop.fs.s3a.S3AFileSystem\")\n",
    "\n",
    "try:\n",
    "    sc = SparkContext(conf=spark_conf)\n",
    "except:\n",
    "    print (\"Spark session already up\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "996d6260-bf6d-45ba-a0ce-b58c688276d1",
   "metadata": {},
   "source": [
    "#### Create bucket on [Minio](https://min.io/) if it does not exists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddd5137a-45aa-4d0d-8b42-df50d3c97b6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from botocore.exceptions import NoCredentialsError\n",
    "\n",
    "# S3 configuration\n",
    "s3 = boto3.client(\n",
    "    's3',\n",
    "    endpoint_url=\"http://lifemap-minio:9000\",\n",
    "    aws_access_key_id=\"root\",\n",
    "    aws_secret_access_key=\"passpass\",\n",
    ")\n",
    "\n",
    "bucket_name = \"data-hail\"\n",
    "\n",
    "# Check if the bucket exists, if not, create it\n",
    "try:\n",
    "    s3.head_bucket(Bucket=bucket_name)\n",
    "    print(f\"Bucket '{bucket_name}' exists.\")\n",
    "except Exception:\n",
    "    # If the bucket does not exist, create it\n",
    "    s3.create_bucket(Bucket=bucket_name)\n",
    "    print(f\"Bucket '{bucket_name}' created.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7919cb2-8906-44d2-8750-7c359ce58798",
   "metadata": {},
   "source": [
    "### [Hail](https://hail.is/) initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e113ae2b-9754-4e8b-b2d3-2839ac142b4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "hl.init(sc=sc, log=log_file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b11e3c63-e86d-46af-97c5-8651867fbc9c",
   "metadata": {},
   "source": [
    "#### Set filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60272387-1744-4e83-a428-da2df979f722",
   "metadata": {},
   "outputs": [],
   "source": [
    "## VCF\n",
    "vcf_fn = 'data/1kg.vcf'\n",
    "## Annotation file\n",
    "annotations_fn = 'data/1kg_annotations.txt'\n",
    "## Matrix table\n",
    "mt_fn = 's3://data-hail/1kg.mt'\n",
    "\n",
    "print (f\"VCF fn: {vcf_fn}\")\n",
    "print (f\"Annotation file fn: {annotations_fn}\")\n",
    "print (f\"Matrix table fn: {mt_fn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b2f6ab9-715c-4f47-a3bf-cf52e5912bbe",
   "metadata": {},
   "source": [
    "#### Reading vcf with Pandas (N/A if the vcf is stored on s3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cd9a733-da94-4149-b507-c636ec4087e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "vcf_pd = None\n",
    "if \"s3://\" not in vcf_fn: \n",
    "    vcf_pd = pd.read_csv(vcf_fn, sep=\"\\t\", header=109, low_memory=False)\n",
    "vcf_pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55d74bab-ae72-4ebb-a48c-87577ba8b411",
   "metadata": {},
   "source": [
    "### Import VCF to hail as Matrix Table\n",
    "\n",
    "The VCF file must be converted into the [Hail Matrix Table](https://hail.is/docs/0.2/overview/matrix_table.html) data structure.\n",
    "This structure builds upon the [Hail Table](https://hail.is/docs/0.2/hail.Table.html#hail.Table) and comprises four key components:\n",
    "\n",
    "- a **two-dimensional matrix of entry fields** where each entry **is indexed by row key(s) and column key(s)**\n",
    "- a corresponding **rows table** that stores all of the row fields that are constant for every column in the dataset\n",
    "- a corresponding **columns table** that stores all of the column fields that are constant for every row in the dataset\n",
    "- a set of **global fields** that are constant for every entry in the dataset\n",
    "\n",
    "Every Hail Table has a **key** that controls both the order of the rows in the table and the ability to join or annotate one table with the information in another table. \n",
    "Matrix tables have **two keys**: a **row key** and a **column key**. Row fields are indexed by the row key, column fields are indexed by the column key, and entry fields are indexed by the row key and the column key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbc132d8-384b-49ab-aaf6-4bd5228a0525",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Read a vcf file, convert and write it as matrix table\n",
    "_ = hl.import_vcf(vcf_fn).write(mt_fn, overwrite=True) # assign this to a dummy variable to avoid errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd5ebad8-b211-480e-88a0-de4d5a560539",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Read the matrix table from the file and assign it to the mt vaiable\n",
    "mt = hl.read_matrix_table(mt_fn)\n",
    "row_table = mt.rows()\n",
    "col_table = mt.cols()\n",
    "entry_fields = mt.entries()\n",
    "\n",
    "print (f\"Row table: {row_table}\")\n",
    "print (f\"Col table: {col_table}\")\n",
    "print (f\"Entries: {entry_fields}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec96b42d-4377-414e-a2bb-43b45db683cb",
   "metadata": {},
   "source": [
    "![alt text](immagini/vcf_matrix_table.png \"Title\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57fcab9e-7302-49ee-962a-c58ca3d48054",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Summary of the matrix table:\n",
    "mt.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a5c7fd-1b4b-4281-a2a2-c7f5b7e8143d",
   "metadata": {},
   "source": [
    "#### Row table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfee5ad7-e653-4265-a1d3-988937edda43",
   "metadata": {},
   "outputs": [],
   "source": [
    "row_table.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11ccae31-81b3-4dd4-afd8-de014d7a91aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "row_table.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be07b06-9072-47e1-a4f0-a487595f4a5e",
   "metadata": {},
   "source": [
    "#### Column table "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15386e64-1ca6-49fd-ae5c-98535eca8b5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "col_table.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55519c09-f11e-4099-afe7-63cc106f93f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "col_table.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce697e0a-0ca9-4a4b-aa7e-00d5d319740c",
   "metadata": {},
   "source": [
    "#### Counting samples and variants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51f5072c-79d7-4b6c-972c-b3f6fb56c9e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Counts of samples and variants\n",
    "print (\" Printing row and column fields with different methods\")\n",
    "# count row fields\n",
    "print(mt.count_rows())\n",
    "# same as\n",
    "print (row_table.count())\n",
    "\n",
    "# count column fields\n",
    "print(mt.count_cols())\n",
    "# same as\n",
    "print(col_table.count())\n",
    "\n",
    "n_variants = mt.count_rows() # This can be done accessing directy the row table with mt.rows().count()\n",
    "n_samples = mt.count_cols()  # This can be done accessing directy the column table with mt.cols().count()\n",
    "\n",
    "print (f\"\\n\\nTable has {n_variants} variants and {n_samples} samples\") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0044105-9171-45a4-88f2-bbb18be0626b",
   "metadata": {},
   "source": [
    "#### Entry fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67bab8b4-1ec6-4caa-8f91-8f966b0a51a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "entry_fields.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228331a4-92be-4bd0-884d-efda2885425f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "entry_fields.show(n_samples + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3abb72b-1c45-4d19-8799-6cc486304fe6",
   "metadata": {},
   "source": [
    "#### Global values.\n",
    "Common values of the matrix table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dc999fe-c9cc-400b-9d1d-4f3b3b8aa488",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt.globals_table().show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "586b1bdb-3248-47dd-ad21-d0260757f463",
   "metadata": {},
   "source": [
    "### Summary:  \n",
    "In the Hail Matrix Table, VCF data are represented as follows:  \n",
    "\n",
    "- **Row fields:** Each row corresponds to a variant, and row fields contain information common to that variant (e.g., alleles or quality).  \n",
    "- **Column fields:** Each column represents a sample, with column fields specifying sample-related attributes such as the sample ID.  \n",
    "- **Entry fields:** These contain data specific to a (variant, sample) pair, such as the genotype read for a particular sample.  \n",
    "- **Column key:** The column field used for join operations.  \n",
    "- **Row key:** The row field used for join operations.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d44cbb8-dcab-4ebb-80f2-b9e697c0e980",
   "metadata": {},
   "source": [
    "#### Showing rows, cols and entry field data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f2b6deb-8044-46dc-834f-f9b318358e8d",
   "metadata": {},
   "source": [
    "We can use the Hail table[`select`](https://hail.is/docs/0.2/hail.Table.html#hail.Table.select) method to pull out 5 variants. \n",
    "The select method takes either a string refering to a field name in the table, or a [Hail Expression](https://hail.is/docs/0.2/overview/expressions.html). \n",
    "If no arguments are provided only the row key fields, locus and alleles, are selected.\n",
    "\n",
    "The `show` method is then used to display the selected variants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4502bc2-e4c1-42ae-b114-bde9af9068c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "row_table.select().show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c9bcc26-76c8-43c0-9f94-5724804a1af2",
   "metadata": {},
   "outputs": [],
   "source": [
    "## the same as\n",
    "mt.row_key.show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5c96b38-04a6-4e6b-8d93-f402722e10c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Select a field of the row table\n",
    "row_table.select(\"qual\").show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "801acb1f-33bb-4771-bf9f-ea88ce399783",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Select a field of the row table\n",
    "row_table.select(row_table.qual, row_table.info.AC).show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1b6e9c-5113-409d-b16f-50a729351ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## The matrix table has also the colum dimension with the sample IDS. \n",
    "## To peek at the first 5 sample IDs\n",
    "## s is the ID field\n",
    "\n",
    "mt.s.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d50ac85e-e345-4651-9cf4-cc938f816a02",
   "metadata": {},
   "source": [
    "#### Show attributes of entry fields. An example with the genotype field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9304fb17-61dd-47cc-8f96-15d45eb9d679",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Attributes of entry fields\n",
    "entry_structure = mt.entry\n",
    "\n",
    "# To show all entry field names\n",
    "print (list(entry_structure))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb0188e-afda-4f63-93ad-f7a67f9ceaf4",
   "metadata": {},
   "source": [
    "To look at the first few genotype calls, we can use entries along with select and take. The **`take` method collects the first n rows into a Python list**. Alternatively, we can use the `show` method, which prints the first n rows to the console in a table format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5fd16ea-87e7-4acd-ac9b-a7200a0853b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "entry_structure.take(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d392ae7-ef97-422d-a78e-aa130cd2e1ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_expr = mt.GT # Takes the GT entry field for all samples \n",
    "gt_expr.phased.show() # Show the phased attribute of the GT field (It is False for not phased haplotypes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750f9a5b-9050-449e-a5d6-83dbe0002442",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_expr.ploidy.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1257c22-ea80-45b6-a915-65dff49937c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_expr.summarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e4269f6-5616-4f32-bb7f-6ec75231e2a4",
   "metadata": {},
   "source": [
    "### Adding column fields to the matrix table starting from a metadata Table\n",
    "\n",
    "Metadata of samples (like phenotypes or geographical origin) are saved in a separate txt file imported to Hail as a Table.\n",
    "\n",
    "A Hail MatrixTable can have any number of row fields and column fields for storing data associated with each row and column. Annotations are usually a critical part of any genetic study. **Column fields are where you’ll store information about sample phenotypes, ancestry, sex, and covariates. Row fields can be used to store information like gene membership and functional impact for use in QC or analysis**.\n",
    "\n",
    "We'll take the text file and use it to annotate the columns in the MatrixTable.\n",
    "\n",
    "The annotation file contains the sample ID, the population and “super-population” designations, the sample sex, and two simulated phenotypes (one binary, one discrete).\n",
    "\n",
    "This file can be imported into Hail with `import_table`. This function produces a **Table object**. This object functions similarly to a Pandas or R DataFrame but is not constrained by the memory of a single machine and it is distributed across Spark. **The Table object, like the matrix table object, is immutable**. To interact with it locally as a Python datastructure, you should use the `take` method or transform to a Pandas dataframe (using `to_pandas` method).\n",
    "\n",
    "**Table can have global field and row fields**. Global field is common for each element of the table whereas row fields are specific for each row. In this case, each row refers specific attributes of a sample. For example, its ID is specified in the \"Sample\" row field whereas the other fields show its metadata. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1030732-5872-4969-ba6a-17032aea41cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation_table = (hl.import_table(annotations_fn, impute=True)\n",
    "         .key_by('Sample'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4a5dc3f-7ccc-4f4d-b63c-124fb6942d0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation_table.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18151a12-828d-4575-ada3-4e04900dc56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation_table.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57951e46-c273-4de9-94d9-3436701e0b2b",
   "metadata": {},
   "source": [
    "#### Query functions and the [Hail Expression](https://hail.is/docs/0.2/overview/expressions.html) Language\n",
    "\n",
    "Hail has a number of useful query functions that can be used for gathering statistics on our dataset. These query functions take Hail Expressions as arguments. Hail’s expressions are lazy representations of data (each data type in Hail has its own Expression class. For example, an Int32Expression represents a 32-bit integer, and a BooleanExpression represents a boolean value of True or False).\n",
    "\n",
    "We will start by looking at some statistics of the information in our table. The [`aggregate`](https://hail.is/docs/0.2/hail.Table.html#hail.Table.aggregate) method can be used to aggregate over rows of the table (see [Aggregation](https://hail.is/docs/0.2/guides/agg.html) and [Aggretators](https://hail.is/docs/0.2/aggregators.html#sec-aggregators) for details).\n",
    "\n",
    "[`counter`](https://hail.is/docs/0.2/aggregators.html#hail.expr.aggregators.counter) is an aggregation function that counts the number of occurrences of each unique element. We can use this to pull out the population distribution by passing in a Hail Expression for the field that we want to count by.\n",
    "\n",
    "The aggregate method is then used to aggregate something in the table across different rows and the aggregate function like counter, stats, etc... are used to specify how and what to aggregate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84858687-a300-4a2b-ad16-025b2d2d58fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Population distribution\n",
    "## Here counter counts unique geographycal origin label\n",
    "\n",
    "aggregate_expression = hl.agg.counter(annotation_table.SuperPopulation)\n",
    "pprint(annotation_table.aggregate(aggregate_expression))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "335a0664-a1aa-466e-9539-f81faf49796e",
   "metadata": {},
   "source": [
    "[`stats`](https://hail.is/docs/0.2/aggregators.html#hail.expr.aggregators.stats) is an aggregation function that produces some useful statistics about numeric collections. We can use this to see the distribution of the CaffeineConsumption phenotype."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec22783-40dc-4a7c-ba9e-6ed59ae15900",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Stats perform some statistics on the specified field\n",
    "## Here take stats of the caffeine consumption\n",
    "\n",
    "aggregate_expression = hl.agg.stats(annotation_table.CaffeineConsumption)\n",
    "pprint(annotation_table.aggregate(aggregate_expression))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf5cba46-da9d-42a5-a06f-97f25d6702c2",
   "metadata": {},
   "source": [
    "The functionality demonstrated in the last few cells isn’t anything especially new: **it’s certainly not difficult to ask these questions with Pandas or R dataframes, or even Unix tools like awk. But Hail can use the same interfaces and query language to analyze collections that are much larger, like the set of variants.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a1e84e8-197a-434a-87fa-16688546b6cb",
   "metadata": {},
   "source": [
    "#### Grouping function to get some information within superpoulation\n",
    "The Table [`group_by`](https://hail.is/docs/0.2/hail.Table.html#hail.Table.group_by) method can be used to apply aggregation functions to different groups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "134ee1b8-b6d2-4895-b2a3-0dac8f9b56d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "grp = annotation_table.group_by('SuperPopulation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2800eda2-67ac-4d6c-b565-84160726fccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "grp.aggregate(cnt=hl.agg.counter(annotation_table.isFemale)).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac5fcadb-dc3c-4ea5-bec1-c8c04f80ef77",
   "metadata": {},
   "outputs": [],
   "source": [
    "grp.aggregate(stats=hl.agg.stats(annotation_table.CaffeineConsumption)).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57d81a5c-8953-4c6f-b114-b22d14566657",
   "metadata": {},
   "source": [
    "#### Join sample annotations with the matrix table\n",
    "\n",
    "Using the [`annotate_cols`](https://hail.is/docs/0.2/hail.MatrixTable.html#hail.MatrixTable.annotate_cols) method is possible to join the annotation table with the MatrixTable containing our dataset.\n",
    "First, we’ll print the existing column schema using `col`. [`col`](https://hail.is/docs/0.2/hail.MatrixTable.html#hail.MatrixTable.col) is an attribute of the matrix table that return struct expression of all column-indexed fields, including keys.\n",
    "It is different from the [`cols()`](https://hail.is/docs/0.2/hail.MatrixTable.html#hail.MatrixTable.cols) method that returns a table with all column fields in the matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d701a5-7147-4063-936b-ed7fcf4aea1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Column table before adding per sample annotation:\n",
    "mt.col.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd60fff5-90be-438e-bd59-f9bee1b7eac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.annotate_cols(pheno = annotation_table[mt.s])\n",
    "\n",
    "# After the annotation the columns has a new field pheno,\n",
    "# a struct that contains sample metadata\n",
    "mt.col.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f68d8346-71ea-4fce-b4cd-8a2bfe6aefdf",
   "metadata": {},
   "source": [
    "Each column \"name\" now specifies the sample ID and its phenotype.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fd03a59-2f7b-41f5-a530-ed96fd714c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Metadata table samples: {annotation_table.count()}\")\n",
    "print(f\"Matrix table samples: {mt.cols().count()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "324a631c-b39c-479d-bf0a-a8c23463594f",
   "metadata": {},
   "source": [
    "Since there are fewer samples in our dataset than in the full thousand genomes cohort, we need to look at annotations on the dataset. We can use [`aggregate_cols`](https://hail.is/docs/0.2/hail.MatrixTable.html#hail.MatrixTable.aggregate_cols) to get the metrics for only the samples in our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f557a494-4867-4f1c-9dd5-fa33ad703900",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt.aggregate_cols(hl.agg.counter(mt.pheno.SuperPopulation))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfcd510b-667e-4072-b894-77c412d8004f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(mt.aggregate_cols(hl.agg.stats(mt.pheno.CaffeineConsumption)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "692351f7-0268-41ce-80e8-aec4c5a89684",
   "metadata": {},
   "source": [
    "The functionality demonstrated in the last few cells isn’t anything especially new: it’s certainly not difficult to ask these questions with **Pandas** or **R** dataframes, or even Unix tools like `awk`. But **Hail** can use the same interfaces and query language to analyze collections that are much larger, like the set of variants."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "327cdcb6-c9d0-4a5c-b96a-c99ccda9c1da",
   "metadata": {},
   "source": [
    "Here we calculate the counts of each of the 12 possible unique SNPs (4 choices for the reference base * 3 choices for the alternate base).\n",
    "\n",
    "To do this, we need to get the alternate allele of each variant and then count the occurences of each unique ref/alt pair. This can be done with Hail’s **counter function**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e7c516-1305-4761-a9e6-c374aec68f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "struct_expr = hl.Struct(ref=mt.alleles[0], alt=mt.alleles[1])\n",
    "snp_counts = mt.aggregate_rows(hl.agg.counter(struct_expr))\n",
    "pprint(snp_counts)\n",
    "\n",
    "# We can list the counts in descending order using Python’s Counter class.\n",
    "from collections import Counter\n",
    "counts = Counter(snp_counts)\n",
    "counts.most_common()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a051048-2e67-43c9-bbfb-80b99bcf7aed",
   "metadata": {},
   "source": [
    "**The same Python, R, and Unix tools could do this work as well, but we’re starting to hit a wall - the latest gnomAD release publishes about 250 million variants, and that won’t fit in memory on a single computer.**\n",
    "\n",
    "What about genotypes? Hail can query the collection of all genotypes in the dataset, and this is getting large even for our tiny dataset. Our 284 samples and 10,000 variants produce 10 million unique genotypes. The gnomAD dataset has about 5 trillion unique genotypes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e4921ea-b0f9-4c6e-b6ed-13f9d8fd913a",
   "metadata": {},
   "source": [
    "#### Hail plotting functions\n",
    "Hail plotting functions allow Hail fields as arguments, so we can pass in the DP field directly here. If the range and bins arguments are not set, this function will compute the range based on minimum and maximum values of the field and use the default 50 bins."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd7657b-3c38-4ac6-aa6b-a056042b1b06",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.histogram(mt.DP, range=(0,30), bins=30, title='DP Histogram', legend='DP')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1e2454d-c862-4a4b-aa44-837218d89eda",
   "metadata": {},
   "source": [
    "### Quality control\n",
    "\n",
    "QC is where analysts spend most of their time with sequencing datasets. QC is an iterative process, and is different for every project: there is no “push-button” solution for QC. Each time the Broad collects a new group of samples, it finds new batch effects. However, by practicing open science and discussing the QC process and decisions with others, we can establish a set of best practices as a community.\n",
    "\n",
    "QC is entirely based on the ability to understand the properties of a dataset. Hail attempts to make this easier by providing the [`sample_qc`](https://hail.is/docs/0.2/methods/genetics.html#hail.methods.sample_qc) function, which produces a set of useful metrics and stores them in a column field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5425e98-0d63-48b4-b2ef-6382931bf366",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt.col.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b255d1b0-87ca-4d1e-a213-0cc451dcf8c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample_qc is a hail genetic method to compute per-sample metrics useful for quality control.\n",
    "mt = hl.sample_qc(mt)\n",
    "mt.col.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e458ef-a374-4037-ba8c-e583df52910e",
   "metadata": {},
   "outputs": [],
   "source": [
    "##Plotting the QC metrics is a good place to start.\n",
    "p = hl.plot.histogram(mt.sample_qc.call_rate, range=(.88,1), legend='Call Rate')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb13533-e96d-4847-8220-6de6cd16b90d",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.histogram(mt.sample_qc.n_not_called, legend='Not called')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daed4b8e-61bd-40fd-a136-269fbce00908",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Checking corralations between the mean value of dp and the call rate\n",
    "p = hl.plot.scatter(mt.sample_qc.dp_stats.mean, mt.sample_qc.call_rate, xlabel='Mean DP', ylabel='Call Rate')\n",
    "p.line([2,22], [0.97,0.97], color='red', line_width=2)\n",
    "p.line([4,4], [0.878,1.0], color='red', line_width=2)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f9f13dd-b947-420c-9e2b-a308f677bfaf",
   "metadata": {},
   "source": [
    "#### Removing outliers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1db8a553-09e9-49b3-b2da-acdbc107228c",
   "metadata": {},
   "source": [
    "Removing outliers from the dataset will generally improve association results. We can make arbitrary cutoffs and use them to filter.\n",
    "Using matrix table [`filter_cols`](https://hail.is/docs/0.2/hail.MatrixTable.html#hail.MatrixTable.filter_cols) method it is possible to **create a new matrix table considering samples with the DP mean >= 4 and a call rate >= 0.97**. Samples that don't satisfy these criteria are removed.\n",
    "The filtering method does not perform in-place filtering, so the result must be assigned to a variable for the changes to take effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae90eebb-0e0d-4cfc-bfbf-33ad0e79e5aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_cols((mt.sample_qc.dp_stats.mean >= 4) & (mt.sample_qc.call_rate >= 0.97))\n",
    "print('After filter, %d/284 samples remain.' % mt.count_cols())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74c7b1c3-0242-4808-b834-bc39a7526957",
   "metadata": {},
   "source": [
    "Next is genotype QC. It’s a good idea to filter out genotypes where the reads aren’t where they should be: if we find a genotype called homozygous reference with >10% alternate reads, a genotype called homozygous alternate with >10% reference reads, or a genotype called heterozygote without a ref / alt balance near 1:1, it is likely to be an error.\n",
    "\n",
    "In a low-depth dataset like 1KG, it is hard to detect bad genotypes using this metric, since a read ratio of 1 alt to 10 reference can easily be explained by binomial sampling. However, in a high-depth dataset, a read ratio of 10:100 is a sure cause for concern!\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9841dac1-4339-42b6-8c3d-0c50fff2db0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ab = mt.AD[1] / hl.sum(mt.AD)\n",
    "\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",
    "\n",
    "fraction_filtered = mt.aggregate_entries(hl.agg.fraction(~filter_condition_ab))\n",
    "print(f'Filtering {fraction_filtered * 100:.2f}% entries out of downstream analysis.')\n",
    "mt = mt.filter_entries(filter_condition_ab)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94a9a82b-2650-4953-9240-4662e8f1dd79",
   "metadata": {},
   "source": [
    "Variant QC computes per per-variant metric useful for quality control. It is a bit more of the same of sample_qc: we can use the [`variant_qc`](https://hail.is/docs/0.2/methods/genetics.html#hail.methods.variant_qc) function to produce a variety of useful statistics, plot them, and filter. This is made at row level beacause they are stats on variants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff4bb5d8-40bf-41fa-82d6-e8fb6e124b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.variant_qc(mt)\n",
    "mt.row.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "213d4896-abac-4c46-ae94-dcfc79c0d16c",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.histogram(mt.variant_qc.call_rate, legend='Variant QC call rate')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c075f48c-7edc-4dea-b16d-a392c9bff300",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.histogram(mt.variant_qc.dp_stats.mean, legend='Variant QC DP')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7315f777-bc1b-452b-858c-15a700d5e189",
   "metadata": {},
   "source": [
    "## Let’s do a GWAS!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1b07509-4460-47d3-b913-9702ad0dfe24",
   "metadata": {},
   "source": [
    "First, we need to restrict to variants that are :\n",
    "\n",
    "- common (we’ll use a cutoff of 1%)\n",
    "- not so far from Hardy-Weinberg equilibrium as to suggest sequencing error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb857418-924f-426c-b332-8fb14620f714",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_rows(mt.variant_qc.AF[1] > 0.01) # It takes variants for which the alternate allele has a frequency larger than 1%\n",
    "print('Samples: %d  Variants: %d' % (mt.count_cols(), mt.count_rows()))\n",
    "\n",
    "mt = mt.filter_rows(mt.variant_qc.p_value_hwe > 1e-6) # Hardy-Weinberg equilibrium pvalue cut-off\n",
    "print('Samples: %d  Variants: %d' % (mt.count_cols(), mt.count_rows()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a606e349-33af-444c-8747-48505c73c3b9",
   "metadata": {},
   "source": [
    "These filters removed about 15% of sites (we started with a bit over 10,000). This is NOT representative of most sequencing datasets! We have already downsampled the full thousand genomes dataset to include more common variants than we’d expect by chance.\n",
    "\n",
    "In Hail, the association tests accept column fields for the sample phenotype and covariates. Since we’ve already got our phenotype of interest (caffeine consumption) in the dataset, we are good to go:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd29f27-6114-47c8-8843-a6b67bc5d78e",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwas = hl.linear_regression_rows(y=mt.pheno.CaffeineConsumption,\n",
    "                                 x=mt.GT.n_alt_alleles(),\n",
    "                                 covariates=[1.0])\n",
    "gwas.row.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e4a84fa-8071-41c6-b0f4-4b82c02b8199",
   "metadata": {},
   "source": [
    "Looking at the bottom of the above printout, you can see the linear regression adds new row fields for the beta, standard error, t-statistic, and p-value.\n",
    "\n",
    "Hail makes it easy to visualize results! Let’s make a Manhattan plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d70e1d8-7857-499e-bef3-0ed8968beffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.manhattan(gwas.p_value)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ee5873b-cd39-4fc2-9b51-09fa09e924df",
   "metadata": {},
   "source": [
    "This doesn’t look like much of a skyline. Let’s check whether our GWAS was well controlled using a Q-Q (quantile-quantile) plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "324ca913-3b03-4fd1-a0df-4f2a62692017",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.qq(gwas.p_value)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "066a0b77-dfa6-4fac-b34c-e4907fb03893",
   "metadata": {},
   "source": [
    "## Confounded!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d40a4335-b9ae-417a-a4b4-ed1f4c1d6514",
   "metadata": {},
   "source": [
    "The observed p-values drift away from the expectation immediately. Either every SNP in our dataset is causally linked to caffeine consumption (unlikely), or there’s a confounder.\n",
    "\n",
    "We didn’t tell you, but sample ancestry was actually used to simulate this phenotype. This leads to a **stratified distribution** of the phenotype. The solution is to include ancestry as a covariate in our regression.\n",
    "\n",
    "The **linear_regression_rows** function can also take column fields to use as covariates. We already annotated our samples with reported ancestry, but it is good to be skeptical of these labels due to human error. Genomes don’t have that problem! Instead of using reported ancestry, we will use genetic ancestry by including computed principal components in our model.\n",
    "\n",
    "The **pca** function produces eigenvalues as a list and sample PCs as a Table, and can also produce variant loadings when asked. The **hwe_normalized_pca** function does the same, using HWE-normalized genotypes for the PCA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb0b55f0-b8a8-490e-9399-b5d6bc66cd79",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt.GT.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7f1140-1810-4256-8051-ae3f84cba377",
   "metadata": {},
   "outputs": [],
   "source": [
    "eigenvalues, pcs, _ = hl.hwe_normalized_pca(mt.GT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "825e7012-b34a-4bfa-b40c-a74d5a6c1fad",
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(eigenvalues)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1468cb35-c59d-4390-bc8f-bfc4dba6a3ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "pcs.show(5, width=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fccd0ff-fcf6-4629-bd89-dfc2cddf8beb",
   "metadata": {},
   "source": [
    "Now that we’ve got principal components per sample, we may as well plot them! Human history exerts a strong effect in genetic datasets. Even with a 50MB sequencing dataset, we can recover the major human populations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0083fa0e-a2e6-4dc2-aa87-795728976dd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.annotate_cols(scores = pcs[mt.s].scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "543c697b-3b95-4e78-9f5a-50bb720ce056",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcd1de97-4261-401e-9840-6639ee8d823e",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.scatter(mt.scores[0],\n",
    "                    mt.scores[1],\n",
    "                    label=mt.pheno.SuperPopulation,\n",
    "                    title='PCA', xlabel='PC1', ylabel='PC2')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b770b933-7160-4c0f-9d1d-e7f37c038758",
   "metadata": {},
   "source": [
    "Now we can rerun our linear regression, controlling for sample sex and the first few principal components. We’ll do this with input variable the number of alternate alleles as before, and again with input variable the genotype dosage derived from the PL field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1148a65-caf2-4481-951b-5651da095d4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwas = hl.linear_regression_rows(\n",
    "    y=mt.pheno.CaffeineConsumption,\n",
    "    x=mt.GT.n_alt_alleles(),\n",
    "    covariates=[1.0, mt.pheno.isFemale, mt.scores[0], mt.scores[1], mt.scores[2]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cac96bb-3333-412c-b5fd-d25ea1e70df9",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwas.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "571c00cd-3a5c-4f19-86c9-a8f93f4d59c4",
   "metadata": {},
   "source": [
    "We’ll first make a Q-Q plot to assess inflation…"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46816505-170a-4305-9538-9e12b0d7faa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.qq(gwas.p_value)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97ad3912-ccd1-4300-8feb-7a3915dc85b3",
   "metadata": {},
   "source": [
    "That’s more like it! This shape is indicative of a well-controlled (but not especially well-powered) study. And now for the Manhattan plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bec5fb11-a4bc-4cd1-bbf9-6e143718ef7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hl.plot.manhattan(gwas.p_value)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fceeeea8-422d-4521-bf0d-94f4923542ce",
   "metadata": {},
   "source": [
    "We have found a caffeine consumption locus!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ef52005-4d69-4481-a93a-e52bdc11663f",
   "metadata": {},
   "source": [
    "#### How to save the Matrix table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d6016a7-c6e7-4896-89f2-ee4d99da0f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "mt_out_fn = 's3://data-hail/1kg_after_gwas.mt'\n",
    "mt.write(mt_out_fn, overwrite=True)"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}