Week3_solutions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trees\n",
    "\n",
    "This week we'll implement the Sankoff algorithm from lecture. This is a good introduction to programming with phylogenetic trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from math import inf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the substitution matrix\n",
    "First, we will create the $c_{ij}$ matrix that contains the substitution scores. As we saw in lecture, the score for no change is 0, the score for a trasition is 1.0 and the score for a transversion is 2.5.\n",
    "\n",
    "We could hard code this since it's only 4x4 but here is some code that will build the matrix as a dictionary of dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "purines = ['A', 'G']\n",
    "pyrimidines = ['C', 'T']\n",
    "nucleotides = purines + pyrimidines\n",
    "\n",
    "c = dict()\n",
    "for nt1 in nucleotides:\n",
    "    c[nt1] = dict()\n",
    "    for nt2 in nucleotides:\n",
    "        # no change\n",
    "        if nt1 == nt2:\n",
    "            c[nt1][nt2] = 0.0\n",
    "        # transition\n",
    "        elif (nt1 in purines and nt2 in purines) or (nt1 in pyrimidines and nt2 in pyrimidines):\n",
    "            c[nt1][nt2] = 1.0\n",
    "        # transversion\n",
    "        else:\n",
    "            c[nt1][nt2] = 2.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see below, the dictionary has nucleotides as keys and another dictionary as values. This second level of dictionaries has nucleotides as keys and the $c_{ij}$ entries as values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'A': {'A': 0.0, 'C': 2.5, 'G': 1.0, 'T': 2.5},\n",
       " 'C': {'A': 2.5, 'C': 0.0, 'G': 2.5, 'T': 1.0},\n",
       " 'G': {'A': 1.0, 'C': 2.5, 'G': 0.0, 'T': 2.5},\n",
       " 'T': {'A': 2.5, 'C': 1.0, 'G': 2.5, 'T': 0.0}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In practical terms, this allows us to easily get a specific $c_{ij}$ value using dictionary lookups, which improves code readability for minimal performance cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c['A']['G']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing the tree\n",
    "We will also implement the tree as a dictionary of dictionaries. The tree dictionary keys are node names and the values are node dictionaries. Thwe node dictionaries have the node's data, scores and left/right node names, as the key-value pairs.\n",
    "\n",
    "The following function creates an empty tree dictionary suitable for running the Sankoff algorithm on DNA data with the correct number of nodes given a number of tips."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def init_tree(num_tips):\n",
    "    \"\"\"\n",
    "    Create an empty tree dictionary with the correct number of nodes.\n",
    "    This function does not link the nodes together, so there is no tree structure yet.\n",
    "    \"\"\"\n",
    "    tree = dict()\n",
    "    \n",
    "    # create all the nodes\n",
    "    tree['root'] = dict()\n",
    "    for x in range(1, num_tips + 1):\n",
    "        tree['tip{}'.format(x)] = dict()\n",
    "    for x in range(1, num_tips + 1 - 2):\n",
    "        tree['int{}'.format(x)] = dict()\n",
    "    \n",
    "    # initialize node data\n",
    "    for node in tree:\n",
    "        for child in ['left', 'right']:\n",
    "            tree[node][child] = None\n",
    "        for nt in nucleotides:\n",
    "            tree[node][nt] = None\n",
    "    \n",
    "    return tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 1: Initializing tip values\n",
    "The following function should initialize a tip node for the Sankoff algorithm to 0 for the observed nucleotide and infinity (`inf`) for others. Try to complete it. Note that the function modifies the tree in-place and does not return a value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_tip(tree, tip_name, observed_nt):\n",
    "    \"\"\"\n",
    "    Intialise a tip of a tree with the correct costs.\n",
    "    At this tip node, only the observed nucleotides are possible:\n",
    "    the cost of the observed nucleotide is zero, and\n",
    "    the cost of any other nucleotide is infinity.\n",
    "    \"\"\"\n",
    "    for nt in nucleotides:\n",
    "        if nt == observed_nt:\n",
    "            tree[tip_name][nt] = 0.0\n",
    "        else:\n",
    "            tree[tip_name][nt] = inf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code sets the sequence at the tips (from left to right) and initializes the tree using the function above. Set the sequence for your tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tip_seq = 'CACAG'\n",
    "my_tree = init_tree(len(tip_seq))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code explicitly sets the tree topology. We can change the tree's structure by modifying these lines.\n",
    "\n",
    "Note that the left and right nodes for the tips are left as `None` because they have no descendents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "my_tree['root']['left'] = 'int1'\n",
    "my_tree['root']['right'] = 'int3'\n",
    "my_tree['int1']['left'] = 'tip1'\n",
    "my_tree['int1']['right'] = 'tip2'\n",
    "my_tree['int3']['left'] = 'tip3'\n",
    "my_tree['int3']['right'] = 'int2'\n",
    "my_tree['int2']['left'] = 'tip4'\n",
    "my_tree['int2']['right'] = 'tip5'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This loop uses the helpful `enumerate` built in function to initialize the tips of the tree using the function you wrote above. Double check that it is working correctly for the tip sequence you entered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tip1 {'left': None, 'right': None, 'A': inf, 'G': inf, 'C': 0.0, 'T': inf}\n",
      "tip2 {'left': None, 'right': None, 'A': 0.0, 'G': inf, 'C': inf, 'T': inf}\n",
      "tip3 {'left': None, 'right': None, 'A': inf, 'G': inf, 'C': 0.0, 'T': inf}\n",
      "tip4 {'left': None, 'right': None, 'A': 0.0, 'G': inf, 'C': inf, 'T': inf}\n",
      "tip5 {'left': None, 'right': None, 'A': inf, 'G': 0.0, 'C': inf, 'T': inf}\n"
     ]
    }
   ],
   "source": [
    "for i, nt in enumerate(tip_seq):\n",
    "    tip_name = 'tip{}'.format(i + 1)\n",
    "    init_tip(my_tree, tip_name, nt)\n",
    "    print(tip_name, my_tree[tip_name])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 2: Sankoff function core calculation\n",
    "This function calculates the $S$ arrays for each node. It assumes that the left and right nodes have already been filled.\n",
    "\n",
    "The outer loop is looping over the four entries (ACGT) in the $S$ array. Complete the function using the variables initialized and the calculation from lecture:\n",
    "\n",
    "$S_a(i) = \\min_{j}[c_{ij} + S_L(j)] + \\min_{k}[c_{ik} + S_R(k)]$\n",
    "\n",
    "As before, this function modifies the tree in place and doesn't return a value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def sankoff_calculate(c_matrix, tree, node_name):\n",
    "    \"\"\"\n",
    "    For the specified node of the tree, calculate the minimum possible cost \n",
    "    for each nucleotide. \n",
    "    \"\"\"\n",
    "    for nt1 in nucleotides:\n",
    "        left_costs = list()\n",
    "        right_costs = list()\n",
    "        left_node = tree[tree[node_name]['left']]\n",
    "        right_node = tree[tree[node_name]['right']]\n",
    "\n",
    "        for nt2 in nucleotides:\n",
    "            left_costs.append(c_matrix[nt2][nt1] + left_node[nt2])\n",
    "            right_costs.append(c_matrix[nt2][nt1] + right_node[nt2])\n",
    "    \n",
    "        tree[node_name][nt1] = min(left_costs) + min(right_costs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the `sankoff_calculate` function requires that the left and right nodes be populated, we will perform a recursive traversal. Try to follow the print statements while looking at the tree on your handout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sankoff_traverse(c_matrix, tree, node_name):\n",
    "    \"\"\"\n",
    "    Traverse the tree recursively, calculating costs for each node.\n",
    "    Calculate a node's children before the node itself.\n",
    "    If a child node already has its costs calculated, don't try to calculate it.\n",
    "    This function depends on correctly initialising the costs in the tips\n",
    "    of the tree before traversing the tree.\n",
    "    \"\"\"\n",
    "    print(\"Examining node {}\".format(node_name))\n",
    "    if tree[tree[node_name]['left']]['A'] is None:\n",
    "        sankoff_traverse(c_matrix, tree, tree[node_name]['left'])\n",
    "    else:\n",
    "        print(\"Node {} values already known\".format(tree[node_name]['left']))\n",
    "    \n",
    "    if tree[tree[node_name]['right']]['A'] is None:\n",
    "        sankoff_traverse(c_matrix, tree, tree[node_name]['right'])\n",
    "    else:\n",
    "        print(\"Node {} values already known\".format(tree[node_name]['right']))   \n",
    "    \n",
    "    print(\"Calculating node {}\".format(node_name))\n",
    "    sankoff_calculate(c_matrix, tree, node_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Examining node root\n",
      "Examining node int1\n",
      "Node tip1 values already known\n",
      "Node tip2 values already known\n",
      "Calculating node int1\n",
      "Examining node int3\n",
      "Node tip3 values already known\n",
      "Examining node int2\n",
      "Node tip4 values already known\n",
      "Node tip5 values already known\n",
      "Calculating node int2\n",
      "Calculating node int3\n",
      "Calculating node root\n"
     ]
    }
   ],
   "source": [
    "sankoff_traverse(c, my_tree, 'root')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can print the tree dictionary and see the $S$ array values for each node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'int1': {'A': 2.5,\n",
       "  'C': 2.5,\n",
       "  'G': 3.5,\n",
       "  'T': 3.5,\n",
       "  'left': 'tip1',\n",
       "  'right': 'tip2'},\n",
       " 'int2': {'A': 1.0,\n",
       "  'C': 5.0,\n",
       "  'G': 1.0,\n",
       "  'T': 5.0,\n",
       "  'left': 'tip4',\n",
       "  'right': 'tip5'},\n",
       " 'int3': {'A': 3.5,\n",
       "  'C': 3.5,\n",
       "  'G': 3.5,\n",
       "  'T': 4.5,\n",
       "  'left': 'tip3',\n",
       "  'right': 'int2'},\n",
       " 'root': {'A': 6.0,\n",
       "  'C': 6.0,\n",
       "  'G': 7.0,\n",
       "  'T': 8.0,\n",
       "  'left': 'int1',\n",
       "  'right': 'int3'},\n",
       " 'tip1': {'A': inf, 'C': 0.0, 'G': inf, 'T': inf, 'left': None, 'right': None},\n",
       " 'tip2': {'A': 0.0, 'C': inf, 'G': inf, 'T': inf, 'left': None, 'right': None},\n",
       " 'tip3': {'A': inf, 'C': 0.0, 'G': inf, 'T': inf, 'left': None, 'right': None},\n",
       " 'tip4': {'A': 0.0, 'C': inf, 'G': inf, 'T': inf, 'left': None, 'right': None},\n",
       " 'tip5': {'A': inf, 'C': inf, 'G': 0.0, 'T': inf, 'left': None, 'right': None}}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 3: Finding the minimum cost\n",
    "By examining the root specifically, we can determine the minimum cost of the tree and infer the ancestral state at this position according to this tree. Write some code that prints out the minimum cost at the root and all nucleotides with that cost.\n",
    "\n",
    "Real reconstructions of ancestral sequences will incorporate lots of additional information and not be so simple as this!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Minimum cost is 6.0 for A and C\n"
     ]
    }
   ],
   "source": [
    "min_nts = list()\n",
    "min_cost = inf\n",
    "\n",
    "for nt in nucleotides:\n",
    "    if my_tree['root'][nt] < min_cost:\n",
    "        min_nts = [nt]\n",
    "        min_cost = my_tree['root'][nt]\n",
    "    elif my_tree['root'][nt] == min_cost:\n",
    "        min_nts.append(nt)\n",
    "\n",
    "print(\"Minimum cost is {} for {}\".format(min_cost, \" and \".join(min_nts)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}