{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to the Expression Language\n",
    "\n",
    "This notebook starts with the basics of the Hail expression language, and builds up practical experience with the type system, syntax, and functionality. By the end of this notebook, we hope that you will be comfortable enough to start using the expression language to slice, dice, filter, and query genetic data. These are covered in the next notebook!\n",
    "\n",
    "The best part about a Jupyter Notebook is that you don't just have to run what we've written - you can and **should** change the code and see what happens!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Every Hail practical notebook starts the same: import the necessary modules, and construct a [HailContext](https://hail.is/hail/hail.HailContext.html#hail.HailContext). This is the entry point for Hail functionality. This object also wraps a SparkContext, which can be accessed with `hc.sc`.\n",
    "\n",
    "As always, visit the [documentation](https://hail.is/hail/api.html) on the Hail website for full reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from hail import *\n",
    "hc = HailContext()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hail Expression Language\n",
    "\n",
    "The Hail expression language is used everywhere in Hail: filtering conditions, describing covariates and phenotypes, storing summary statistics about variants and samples, generating synthetic data, plotting, exporting, and more. The Hail expression language takes the form of Python strings passed into various Hail methods like [filter_variants_expr](https://hail.is/hail/hail.VariantDataset.html#hail.VariantDataset.annotate_variants_expr) and [linear regression](https://hail.is/hail/hail.VariantDataset.html#hail.VariantDataset.linreg).\n",
    "\n",
    "The expression language is a programming language just like Python or R or Scala. While the syntax is different, programming experience will certainly translate. We have built the expression language with the hope that even people new to programming are able to use it to explore genetic data, even if this means copying motifs and expressions found on places like [Hail discussion forum](http://discuss.hail.is/).\n",
    "\n",
    "For learning purposes, `HailContext` contains the method [eval_expr_typed](https://hail.is/hail/hail.HailContext.html#hail.HailContext.eval_expr_typed). This method takes a Python string of Hail expr code, evaluates it, and returns a tuple with the result and the type.  We'll be using this method throughout the expression language tutorial.\n",
    "\n",
    "## Hail Types\n",
    "\n",
    "\n",
    "The Hail expression language is strongly typed, meaning that every _expression_ has an associated type.\n",
    "\n",
    "Hail defines the following types:\n",
    "\n",
    "Primitives:\n",
    " - [Int](https://hail.is/hail/types.html#int)\n",
    " - [Double](https://hail.is/hail/types.html#double)\n",
    " - [Float](https://hail.is/hail/types.html#float)\n",
    " - [Long](https://hail.is/hail/types.html#long)\n",
    " - [Boolean](https://hail.is/hail/types.html#boolean)\n",
    " - [String](https://hail.is/hail/types.html#string)\n",
    " \n",
    "Compound Types:\n",
    " - [Array[T]](https://hail.is/hail/types.html#array)\n",
    " - [Set[T]](https://hail.is/hail/types.html#set)\n",
    " - [Dict[K, V]](https://hail.is/hail/types.html#dict)\n",
    " - [Aggregable[T]](https://hail.is/hail/types.html#aggregable)\n",
    " - [Struct](https://hail.is/hail/types.html#struct)\n",
    " \n",
    "Genetic Types:\n",
    " - [Variant](https://hail.is/hail/types.html#variant)\n",
    " - [Locus](https://hail.is/hail/types.html#locus)\n",
    " - [AltAllele](https://hail.is/hail/types.html#altallele)\n",
    " - [Interval](https://hail.is/hail/types.html#interval)\n",
    " - [Genotype](https://hail.is/hail/types.html#genotype)\n",
    " - [Call](https://hail.is/hail/types.html#call)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Primitive Types\n",
    "\n",
    "Let's start with simple primitive types. Primitive types are a basic building block for any programming language - these are things like numbers and strings and boolean values. \n",
    "\n",
    "Hail expressions are passed as Python strings to Hail methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# the Boolean literals are 'true' and 'false'\n",
    "hc.eval_expr_typed('true') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The return value is `True`, not `true`.  Why?  When values are returned by Hail methods, they are returned as the corresponding Python value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('123')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('123.45')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "String literals are denoted with double-quotes. The 'u' preceding the printed result denotes a unicode string, and is safe to ignore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('\"Hello, world\"')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Primitive types support all the usual operations you'd expect.  For details, refer to the documentation on [operators](https://hail.is/hail/operators.html) and [types](https://hail.is/hail/types.html). Here are some examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('3 + 8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('3.2 * 0.5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('3 ** 3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('25 ** 0.5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('true || false')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('true && false')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Missingness\n",
    "\n",
    "Like R, all values in Hail can be missing.  Most operations, like addition, return missing if any of their inputs is missing.  There are a few special operations for manipulating missing values.  There is also a missing literal, but you have to specify it's type.  Missing Hail values are converted to `None` in Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('NA: Int') # missing Int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('NA: Dict[String, Int]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('1 + NA: Int')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can test missingness with `isDefined` and `isMissing`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('isDefined(1)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('isDefined(NA: Int)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('isMissing(NA: Double)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`orElse` lets you convert missing to a default value and `orMissing` lets you turn a value into missing based on a condtion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('orElse(5, 2)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('orElse(NA: Int, 2)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('orMissing(true, 5)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('orMissing(false, 5)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let\n",
    "\n",
    "You can assign a value to a variable with a `let` expression.  Here is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let a = 5 in a + 1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The variable, here `a` is only visible in the body of the let, the expression following `in`.  You can assign multiple variables.  Variable assignments are separated by `and`.  Each variable is visible in the right hand side of the following variables as well as the body of the let.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('''\n",
    "let a = 5\n",
    "and b = a + 1\n",
    " in a * b\n",
    "''')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conditionals\n",
    "\n",
    "Unlike other languages, conditionals in Hail return a value.  The arms of the conditional must have the same type.  The predicate must be of type Boolean.  If the predicate is missing, the value of the entire conditional is missing.  Here are some simple examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('if (true) 1 else 2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('if (false) 1 else 2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('if (NA: Boolean) 1 else 2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `if` and `else` branches need to return the same type. The below expression is invalid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Uncomment and run the below code to see the error message\n",
    "\n",
    "# hc.eval_expr_typed('if (true) 1 else \"two\"')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compound Types\n",
    "\n",
    "Hail has several compound types:\n",
    " - [Array[T]](https://hail.is/hail/types.html#array)\n",
    " - [Set[T]](https://hail.is/hail/types.html#set)\n",
    " - [Dict[K, V]](https://hail.is/hail/types.html#dict)\n",
    " - [Aggregable[T]](https://hail.is/hail/types.html#aggregable)\n",
    " - [Struct](https://hail.is/hail/types.html#struct)\n",
    " \n",
    "`T`, `K` and `V` here mean any type, including other compound types.  Hail's `Array[T]` objects are similar to Python's lists, except they must be homogenous: that is, each element must be of the same type.  Arrays are 0-indexed.  Here are some examples of simple array expressions.\n",
    "\n",
    "Array literals are constructed with square brackets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('[1, 2, 3, 4, 5]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arrays are indexed with square brackets and support Python's slice syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let a = [1, 2, 3, 4, 5] in a[0]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let a = [1, 2, 3, 4, 5] in a[1:3]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let a = [1, 2, 3, 4, 5] in a[1:]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let a = [1, 2, 3, 4, 5] in a.length()')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arrays can be transformed with functional operators `filter` and `map`.  These operations return a new array, never modify the original."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# keep the elements that are less than 10\n",
    "hc.eval_expr_typed('let a = [1, 2, 22, 7, 10, 11] in a.filter(x => x < 10)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# square the elements of an array\n",
    "hc.eval_expr_typed('let a = [1, 2, 22, 7, 10, 11] in a.map(x => x * x)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# combine the two: keep elements less than 10 and then square them\n",
    "hc.eval_expr_typed('let a = [1, 2, 22, 7, 10, 11] in a.filter(x => x < 10).map(x => x * x)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above filter / map expressions, you can see a strange syntax:\n",
    "\n",
    "`x => x < 10`\n",
    "\n",
    "This syntax is a [lambda function](https://en.wikipedia.org/wiki/Anonymous_function). The functions `filter` and `map` take _functions_ as arguments! A Hail lambda function takes the form:\n",
    "\n",
    "`binding => expression`\n",
    "\n",
    "That we named the binding 'x' in every example above is a point of preference, and no more. We can name the bindings anything we want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# use 'foo' and 'bar' as bindings\n",
    "hc.eval_expr_typed('let a = [1, 2, 22, 7, 10, 11] in a.filter(foo => foo < 10).map(bar => bar * bar)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The full list of methods on arrays can be found [here](https://hail.is/hail/types.html#array-t)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numeric Arrays\n",
    "\n",
    "Numeric arrays, like `Array[Int]` and `Array[Double]` have additional operations like `max`, `mean`, `median`, `sort`.  For a full list, see, for example, [Array[Int]](https://hail.is/hail/types.html#array-int).  Here are a few examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('[1, 2, 22, 7, 10, 11].sum()')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('[1, 2, 22, 7, 10, 11].max()')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('[1, 2, 22, 7, 10, 11].mean()')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# take the square root of each element\n",
    "hc.eval_expr_typed('let a = [1, 2, 22, 7, 10, 11] in a.map(x => x ** 0.5)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "Write an expression that calculates the sum of the squared residuals (x - mean) of an array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Uncomment the below code by deleting the triple-quotes and write an expression to calculate the residuals.\n",
    "\n",
    "\"\"\"\n",
    "result, t = hc.eval_expr_typed('''\n",
    "let a = [1, -2, 11, 3, -2]\n",
    "and mean = <FILL IN>\n",
    "in a.map(x => <FILL IN> ).sum() \n",
    "''')\n",
    "\"\"\"\n",
    "\n",
    "try:\n",
    "    print('Your result: %s (%s)' % (result, t))\n",
    "    print('Expected answer:  114.8 (Double)')\n",
    "except NameError:\n",
    "    print('### Remove the triple quotes around the above code to start the exercise ### ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if `a` contains a missing value NA: Int? Will your code still work?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structs\n",
    "\n",
    "`Struct`s are a collection of named values known as fields.  Hail does not have tuples like Python.  Unlike arrays, the values can be heterogenous.  Unlike `Dict`s, the set of names are part of the type and must be known statically.  `Struct`s are constructed with a syntax similar to Python's `dict` syntax.  `Struct` fields are accessed using the `.` syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(hc.eval_expr_typed('{gene: \"ACBD\", function: \"LOF\", nHet: 12}'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let s = {gene: \"ACBD\", function: \"LOF\", nHet: 12} in s.gene')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('let s = NA: Struct { gene: String, function: String, nHet: Int} in s.gene')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Genetic Types\n",
    "\n",
    "Hail contains several genetic types:\n",
    " - [Variant](https://hail.is/hail/types.html#variant)\n",
    " - [Locus](https://hail.is/hail/types.html#locus)\n",
    " - [AltAllele](https://hail.is/hail/types.html#altallele)\n",
    " - [Interval](https://hail.is/hail/types.html#interval)\n",
    " - [Genotype](https://hail.is/hail/types.html#genotype)\n",
    " - [Call](https://hail.is/hail/types.html#call)\n",
    " \n",
    "These are designed to make it easy to manipulate genetic data. There are many built-in functions for asking common questions about these data types, like whether an alternate allele is a SNP, or the fraction of reads a called genotype that belong to the reference allele."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo variables\n",
    "\n",
    "To explore these types and constructs, we have defined five representative variables which you can access in `eval_expr`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 'v' is used to indicate 'Variant' in Hail\n",
    "hc.eval_expr_typed('v')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 's' is used to refer to sample ID in Hail\n",
    "hc.eval_expr_typed('s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 'g' is used to refer to the genotype in Hail\n",
    "hc.eval_expr_typed('g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 'sa' is used to refer to sample annotations\n",
    "hc.eval_expr_typed('sa')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above output is a bit wordy. Let's try `'va'`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 'va' is used to refer to variant annotations\n",
    "hc.eval_expr_typed('va')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**This is totally illegible.** `pprint` **can solve our problems!**\n",
    "\n",
    "`pprint` is a Python standard library module that tries to print objects legibly. Let's try it out here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 'va' is used to refer to variant annotations\n",
    "pprint(hc.eval_expr_typed('va'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll rarely need to construct a `Variant` or `Genotype` object inside the Hail expression language. More commonly, these objects will be provided to you as variables. In the remainder of this notebook, we will explore how to to manipulate the demo variables. In the next notebook, we start using the expression langauge to annotate and filter a dataset.\n",
    "\n",
    "First, a short demonstration of some of the methods accessible on `Variant` and `Genotype` objects:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v.contig')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v.start')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v.ref')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v.altAlleles')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v.altAlleles.map(aa => aa.isSNP())')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('v.altAlleles.map(aa => aa.isInsertion())')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('g.dp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('g.ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('g.fractionReadsRef()')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('g.isHet()')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrangling complex nested types\n",
    "\n",
    "Structs and Arrays allow arbitrarily deep grouping and nesting of values.\n",
    "\n",
    "Remember the type of `sa`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pprint(hc.eval_expr_typed('sa')[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select elements of a `Struct` with a `'.'`. If we want to select `PC1` from the above type, we first index into the top-level struct with `covariates`, then select the field with `PC1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('sa.covariates.PC1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can construct an array from the struct elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('[sa.covariates.PC1, sa.covariates.PC2, sa.covariates.PC3]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll use `va`. Here's its type of `va`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pprint(hc.eval_expr_typed('va')[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This schema is somewhat representative of typical variant annotations: `AC`, `AN`, and `AF` are typically included in the `INFO` field of a VCF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.info.AF')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.info.AF[1]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AC and AF mean \"allele count\" and \"allele frequency\" and are \"A-indexed\", which means that there is one element per alternate allele. Perhaps we want to construct an array which contains each alternate allele and its count and frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pprint(hc.eval_expr_typed('''range(v.altAlleles.length()).map(i => \n",
    "                      {allele: v.altAlleles[i], \n",
    "                       count: va.info.AC[i], \n",
    "                       frequency: va.info.AF[i]})'''))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's manipulate the `va.transcripts` array. Here's what it looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.transcripts')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start by pulling out just the gene field. Our result will be an `Array[String]`. We need to do this with the `map` function, to map each struct element of the array to its field `gene`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.transcripts.map(t => t.gene)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perhaps we just want the set of unique genes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.transcripts.map(t => t.gene).toSet()')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can find the canonical transcript with `find`, which returns the first element where the predicate is true:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.transcripts.find(t => t.canonical)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, `find` returns `None` if there isn't an element where the predicate is true:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.transcripts.find(t => t.gene == \"GENE5\")')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll pull out all transcripts marked \"MIS\" (missense):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('va.transcripts.filter(t => t.consequence == \"MIS\")')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a bit of a complex motif - we can sort the transcripts by an arbitrary function. Here we'll sort so that `\"LOF\"` comes before `\"MIS\"`, and `\"MIS\"` comes before `\"SYN\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('''va.transcripts.sortBy(t => \n",
    "                        if (t.consequence == \"LOF\") 1 \n",
    "                        else if (t.consequence == \"MIS\") 2 \n",
    "                        else 3)''')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we are interested in pulling out the worst-consequence transcript, we can use this sorting motif and then take the first element:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hc.eval_expr_typed('''va.transcripts.sortBy(t => \n",
    "                        if (t.consequence == \"LOF\") 1 \n",
    "                        else if (t.consequence == \"MIS\") 2 \n",
    "                        else 3)[0]''')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Learn more!\n",
    "\n",
    "- [Basic language constructs](https://hail.is/hail/language_constructs.html)\n",
    "\n",
    "- [Operators](https://hail.is/hail/operators.html)\n",
    "\n",
    "- [Functions](https://hail.is/hail/functions.html)\n",
    "\n",
    "- [Types](https://hail.is/hail/types.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Exercises\n",
    "\n",
    "Uncomment the code blocks, fill them in, and run each block to check your answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def check(answer, answer_key):\n",
    "    print('Your answer / type:')\n",
    "    pprint(answer)\n",
    "    print('')\n",
    "    if (answer == answer_key):\n",
    "        print('Correct!')\n",
    "    else:\n",
    "        print('Incorrect. Expected:')\n",
    "        pprint(answer_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 1: using** `filter` **and** `map` **to pull out the gene isoform for synonymous transcripts**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "result_1 = hc.eval_expr_typed(\n",
    "'''\n",
    "va.transcripts.filter(t => <FILL IN>)\n",
    "  .map(t => <FILL IN>)\n",
    "''')\n",
    "\"\"\"\n",
    "# check the answer\n",
    "try:\n",
    "    answer_key = [u'GENE1.1', u'GENE3.1', u'GENE3.2'], TArray(TString())\n",
    "    check(result_1, answer_key)\n",
    "except NameError:\n",
    "    print('### Remove the triple quotes around the above code to start the exercise ### ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 2: using** `groupBy` **and** `mapValues` **to produce a mapping from gene to all observed consequences**\n",
    "\n",
    "Remember: `<array>.toSet()` converts an array to a Set, the desired type of the dictionary value.\n",
    "\n",
    "Hint: Once you've grouped by gene, you can fill in the `mapValues` step with `ts => ts` to see the type of `ts`. It's an `Array[Struct{...}]`. How do we pull just one field out?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "result_2 = hc.eval_expr_typed(\n",
    "'''\n",
    "  va.transcripts.groupBy(t => <FILL IN>)\n",
    "    .mapValues(ts => <FILL IN>)\n",
    "''')\n",
    "\"\"\"\n",
    "\n",
    "# check the answer\n",
    "try:\n",
    "    answer_key = {u'GENE1': {u'LOF', u'SYN'}, u'GENE2': {u'MIS'}, u'GENE3': {u'SYN'}}, TDict(TString(), TSet(TString()))\n",
    "    check(result_2, answer_key)\n",
    "except NameError:\n",
    "    print('### Remove the triple quotes around the above code to start the exercise ### ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 3: Do the reverse: group** `va.transcripts` **by consequence, and produce a mapping from consequence to all genes with that consequence**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "result_3 = hc.eval_expr_typed(\n",
    "'''\n",
    "<FILL IN>\n",
    "''')\n",
    "\"\"\"\n",
    "\n",
    "# check the answer\n",
    "try:\n",
    "    answer_key = {u'LOF': {u'GENE1'}, u'MIS': {u'GENE2'}, u'SYN': {u'GENE1', u'GENE3'}}, TDict(TString(), TSet(TString()))\n",
    "    check(result_3, answer_key)\n",
    "except NameError:\n",
    "    print('### Remove the triple quotes around the above code to start the exercise ### ')"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}