tests/cookbook-baba-simulations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## (&#x1F4D7;) Cookbook: `ipyrad.analysis simulations` \n",
    "\n",
    "This notebook demonstrates how to use the `baba` module in `ipyrad` to test for admixture and introgression. The code is written in Python and is best implemented in a Jupyter-notebook like this one. Analyses can be split up among many computing cores. Finally, there are some simple plotting methods to accompany tests which I show below. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ipyrad v.0.6.11\n",
      "ipyparallel v.6.0.2\n",
      "numpy v.1.12.0\n"
     ]
    }
   ],
   "source": [
    "## imports\n",
    "import numpy as np\n",
    "import ipyrad as ip\n",
    "import ipyparallel as ipp\n",
    "from ipyrad.analysis import baba\n",
    "\n",
    "## print versions\n",
    "print \"ipyrad v.{}\".format(ip.__version__)\n",
    "print \"ipyparallel v.{}\".format(ipp.__version__)\n",
    "print \"numpy v.{}\".format(np.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### ipcluster setup to run parallel code\n",
    "You can find more details about this in our [ipyparallel tutorial]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "## start 4 engines by running the commented line below in a separate bash terminal. \n",
    "##\n",
    "## ipcluster start --n=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "host compute node: [4 cores] on oud\n"
     ]
    }
   ],
   "source": [
    "## connect to client and print connection summary\n",
    "ipyclient = ipp.Client()\n",
    "print ip.cluster_info(client=ipyclient)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Get default (example) 12 taxon tree \n",
    "The function  `baba.Tree()`  takes a  `newick`  string as an argument, or, if no value is passed, it constructs the default 12 taxon tree shown below. The  `Tree`  Class object holds a representation of the tree, which can be used for plotting, as well as for describing a demographic model for simulating data, as we'll show below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'ipyrad.analysis.baba' from '/home/deren/miniconda2/lib/python2.7/site-packages/ipyrad/analysis/baba.pyc'>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tre.draw(width=400, height=200);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "## tree2dict \n",
    "import ete3\n",
    "newick = \"./analysis_tetrad/pedtest1.full.tre\"\n",
    "## load the newick string as a Tree object\n",
    "tre = ete3.Tree(newick)\n",
    "\n",
    "## set przewalskiis as the outgroup\n",
    "prz = [i for i in tre.get_leaves() if \"prz\" in i.name]\n",
    "out = tre.get_common_ancestor(prz)\n",
    "\n",
    "## set new outgroup and store back as a newick string\n",
    "tre.set_outgroup(out)\n",
    "newick = tre.write()\n",
    "\n",
    "\n",
    "def test_constraint(node, cdict, tip, exact):\n",
    "    names = set(node.get_leaf_names())\n",
    "    const = set(cdict[tip])\n",
    "    if const:\n",
    "        if exact:\n",
    "            if len(names.intersection(const)) == len(const):\n",
    "                return 1\n",
    "            else:\n",
    "                return 0\n",
    "        else:\n",
    "            if len(names.intersection(const)) == len(names):\n",
    "                return 1\n",
    "            else:\n",
    "                return 0        \n",
    "    return 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'p1': [], 'p2': [], 'p3': [], 'p4': ['a']}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## constraints\n",
    "#if not constraint_dict:\n",
    "cdict = {\"p1\":[], \"p2\":[], \"p3\":[], \"p4\":[]}\n",
    "cdict.update({\"p4\":['a']})\n",
    "\n",
    "cdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def tree2tests(newick, constraint_dict=None, constraint_exact=True):\n",
    "    \"\"\"\n",
    "    Returns dict of all possible four-taxon splits in a tree. Assumes\n",
    "    the user has entered a rooted tree. Skips polytomies.\n",
    "    \"\"\"\n",
    "    ## make tree\n",
    "    tree = ete3.Tree(newick)\n",
    "    \n",
    "    ## constraints\n",
    "    #if not constraint_dict:\n",
    "    cdict = {\"p1\":[], \"p2\":[], \"p3\":[], \"p4\":[]}\n",
    "    cdict.update(constraint_dict)\n",
    "    print cdict\n",
    "\n",
    "    ## traverse root to tips. Treat the left as outgroup, then the right.\n",
    "    tests = []\n",
    "    ## topnode must have children\n",
    "    for topnode in tree.traverse(\"levelorder\"):  \n",
    "        ## test onode as either child\n",
    "        for oparent in topnode.children:\n",
    "            ## test outgroup as all descendants on one child branch\n",
    "            for onode in oparent.traverse():\n",
    "                ## put constraints on onode\n",
    "                if test_constraint(onode, cdict, \"p4\", constraint_exact):\n",
    "                    ## p123 parent is sister to oparent\n",
    "                    p123parent = oparent.get_sisters()[0]\n",
    "                    ## test p3 as all descendants of p123parent\n",
    "                    for p3parent in p123parent.children:\n",
    "                        ## p12 parent is sister to p3parent\n",
    "                        p12parent = p3parent.get_sisters()[0]\n",
    "                        for p3node in p3parent.traverse():\n",
    "                            if test_constraint(p3node, cdict, \"p3\", constraint_exact):\n",
    "                                if p12parent.children:\n",
    "                                    p1parent, p2parent = p12parent.children\n",
    "                                    for p2node in p2parent.traverse():\n",
    "                                        if test_constraint(p2node, cdict, \"p2\", constraint_exact):\n",
    "                                            for p1node in p1parent.traverse():\n",
    "                                                if test_constraint(p1node, cdict, \"p1\", constraint_exact):\n",
    "                                                    test = {}\n",
    "                                                    test['p4'] = onode.get_leaf_names()\n",
    "                                                    test['p3'] = p3node.get_leaf_names()\n",
    "                                                    test['p2'] = p2node.get_leaf_names()\n",
    "                                                    test['p1'] = p1node.get_leaf_names()\n",
    "                                                    tests.append(test)\n",
    "    return tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'p2': [], 'p3': ['30686_cyathophylla'], 'p1': [], 'p4': ['32082_przewalskii', '33588_przewalskii']}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "33"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "    \n",
    "tests = tree2tests(newick, \n",
    "                   constraint_dict={\"p4\":['32082_przewalskii', '33588_przewalskii'], \n",
    "                                    \"p3\":['30686_cyathophylla']}, \n",
    "                   constraint_exact=True)\n",
    "len(tests)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'p1': ['38362_rex',\n",
       "   '39618_rex',\n",
       "   '35855_rex',\n",
       "   '40578_rex',\n",
       "   '30556_thamno',\n",
       "   '35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex', '39618_rex', '35855_rex', '40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['30556_thamno', '35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex', '39618_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35855_rex', '40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['30556_thamno'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['39618_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35855_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba',\n",
       "   '30686_cyathophylla',\n",
       "   '41478_cyathophylloides',\n",
       "   '41954_cyathophylloides'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex',\n",
       "   '39618_rex',\n",
       "   '35855_rex',\n",
       "   '40578_rex',\n",
       "   '30556_thamno',\n",
       "   '35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex', '39618_rex', '35855_rex', '40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['30556_thamno', '35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex', '39618_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35855_rex', '40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['30556_thamno'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['39618_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35855_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['29154_superba', '30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex',\n",
       "   '39618_rex',\n",
       "   '35855_rex',\n",
       "   '40578_rex',\n",
       "   '30556_thamno',\n",
       "   '35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex', '39618_rex', '35855_rex', '40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['30556_thamno', '35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex', '39618_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35855_rex', '40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['30556_thamno'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35236_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['38362_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['39618_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['35855_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']},\n",
       " {'p1': ['40578_rex'],\n",
       "  'p2': ['33413_thamno'],\n",
       "  'p3': ['30686_cyathophylla'],\n",
       "  'p4': ['32082_przewalskii', '33588_przewalskii']}]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Or get any arbitrary tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div align=\"center\" class=\"toyplot\" id=\"t441148ce60424ae19347991a725356a4\"><svg class=\"toyplot-canvas-Canvas\" height=\"200.0px\" id=\"t68d81b042a524e5a82d5b3c8efa017b1\" preserveAspectRatio=\"xMidYMid meet\" style=\"background-color:transparent;fill:rgb(16.1%,15.3%,14.1%);fill-opacity:1.0;font-family:Helvetica;font-size:12px;opacity:1.0;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:1.0\" viewBox=\"0 0 200.0 200.0\" width=\"200.0px\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:toyplot=\"http://www.sandia.gov/toyplot\" xmlns:xlink=\"http://www.w3.org/1999/xlink\"><g class=\"toyplot-coordinates-Cartesian\" id=\"te0713cb365ea422fa224a5e9a693bad6\"><clipPath id=\"tf2101a4fe1024ca89d78b79b762240dc\"><rect height=\"200.0\" width=\"200.0\" x=\"0.0\" y=\"0.0\"></rect></clipPath><g clip-path=\"url(#tf2101a4fe1024ca89d78b79b762240dc)\"><g class=\"toyplot-mark-Graph\" id=\"te8238f979c2a49d7817ed43710bbdbf7\"><g class=\"toyplot-Edges\"><path d=\"M 55.0 100.0 L 100.0 190.0\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 55.0 100.0 L 10.0 190.0\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 122.5 10.0 L 55.0 100.0\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 122.5 10.0 L 190.0 190.0\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path></g><g class=\"toyplot-Vertices\"><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"122.5\" cy=\"10.0\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"55.0\" cy=\"100.0\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"100.0\" cy=\"190.0\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"10.0\" cy=\"190.0\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"190.0\" cy=\"190.0\" r=\"10.0\"></circle></g></g><g class=\"toyplot-Labels\"><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(122.5,10.0)\"><tspan style=\"dominant-baseline:inherit\">0</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(55.0,100.0)\"><tspan style=\"dominant-baseline:inherit\">1</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(100.0,190.0)\"><tspan style=\"dominant-baseline:inherit\">a</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(10.0,190.0)\"><tspan style=\"dominant-baseline:inherit\">b</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(190.0,190.0)\"><tspan style=\"dominant-baseline:inherit\">c</tspan></text></g></g></g></g></svg><div class=\"toyplot-interactive\"></div></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## pass a newick string to the Tree class object\n",
    "tre = baba.Tree(newick=\"((a,b),c);\")\n",
    "tre.draw(width=200, height=200);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### get a tree with an admixture edge\n",
    "Shown by an arrow pointing backwards from source to sink (indicating that geneflow should be entered as occurring backwards in time). This is how the program `msprime` will interpret it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div align=\"center\" class=\"toyplot\" id=\"t23633f2d1ffd483bbf18a7012b847b3f\"><svg class=\"toyplot-canvas-Canvas\" height=\"250.0px\" id=\"t5763073eee484e9aa1452371ce24f564\" preserveAspectRatio=\"xMidYMid meet\" style=\"background-color:transparent;fill:rgb(16.1%,15.3%,14.1%);fill-opacity:1.0;font-family:Helvetica;font-size:12px;opacity:1.0;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:1.0\" viewBox=\"0 0 250.0 250.0\" width=\"250.0px\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:toyplot=\"http://www.sandia.gov/toyplot\" xmlns:xlink=\"http://www.w3.org/1999/xlink\"><g class=\"toyplot-coordinates-Cartesian\" id=\"t356a435845a349aeaae713d43ca632fe\"><clipPath id=\"t8a9290e2990b45f6baaba37f592f9be0\"><rect height=\"245.00000000000003\" width=\"245.00000000000003\" x=\"2.5\" y=\"2.5\"></rect></clipPath><g clip-path=\"url(#t8a9290e2990b45f6baaba37f592f9be0)\"><g class=\"toyplot-mark-Graph\" id=\"t44dbac87de104851af21c2f553049ce4\"><g class=\"toyplot-Edges\"><path d=\"M 209.375 162.5 L 181.25 237.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 146.09375 12.5 L 209.375 162.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 82.8125 87.5 L 125.0 237.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 40.625 162.5 L 68.75 237.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 40.625 162.5 L 12.5 237.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 82.8125 87.5 L 40.625 162.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 146.09375 12.5 L 82.8125 87.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path><path d=\"M 209.375 162.5 L 237.5 237.5\" style=\"fill:none;stroke:rgb(16.1%,15.3%,14.1%);stroke-opacity:1.0;stroke-width:3.0\"></path></g><g class=\"toyplot-Vertices\"><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"146.09375000000003\" cy=\"12.5\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"209.37500000000003\" cy=\"162.50000000000003\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"82.812500000000014\" cy=\"87.500000000000014\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"40.625\" cy=\"162.50000000000003\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"125.00000000000001\" cy=\"237.50000000000003\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"68.75\" cy=\"237.50000000000003\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"12.5\" cy=\"237.50000000000003\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"237.50000000000003\" cy=\"237.50000000000003\" r=\"10.0\"></circle></g><g class=\"toyplot-Datum\" style=\"fill:rgb(40%,76.1%,64.7%);fill-opacity:1.0;opacity:1.0;stroke:rgb(40%,76.1%,64.7%);stroke-opacity:1.0\"><circle cx=\"181.25000000000003\" cy=\"237.50000000000003\" r=\"10.0\"></circle></g></g><g class=\"toyplot-Labels\"><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(146.09375000000003,12.5)\"><tspan style=\"dominant-baseline:inherit\">0</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(209.37500000000003,162.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">1</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(82.812500000000014,87.500000000000014)\"><tspan style=\"dominant-baseline:inherit\">2</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(40.625,162.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">3</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(125.00000000000001,237.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">Cow</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(68.75,237.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">a</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(12.5,237.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">b</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(237.50000000000003,237.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">d</tspan></text><text class=\"toyplot-Datum\" style=\"dominant-baseline:middle;font-size:18px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(181.25000000000003,237.50000000000003)\"><tspan style=\"dominant-baseline:inherit\">e</tspan></text></g></g><g class=\"toyplot-mark-Plot\" id=\"tb666e086654c4dc59f6f31a1784d85be\" style=\"fill:none;stroke-dasharray:2, 2;stroke-width:3\"><g class=\"toyplot-Series\"><path d=\"M 195.31250000000003 200.00000000000003 L 103.90625000000001 162.50000000000003\" style=\"stroke:rgb(98.8%,55.3%,38.4%);stroke-dasharray:2, 2;stroke-opacity:1.0;stroke-width:3\"></path></g></g><g class=\"toyplot-mark-Plot\" id=\"t0d5716b3113447ebaccbb399ad64c5af\" style=\"fill:none;stroke-dasharray:2, 2;stroke-width:3\"><g class=\"toyplot-Series\"><path d=\"M 61.718750000000007 125.00000000000001 L 177.73437500000003 87.500000000000014\" style=\"stroke:rgb(98.8%,55.3%,38.4%);stroke-dasharray:2, 2;stroke-opacity:1.0;stroke-width:3\"></path></g></g></g><g class=\"toyplot-coordinates-Axis\" id=\"tb0fd033e2db944288db3e18d0448a556\" transform=\"translate(12.5,237.5)rotate(-90.0)translate(0,-10.0)\"><line style=\"\" x1=\"0\" x2=\"225.00000000000003\" y1=\"0\" y2=\"0\"></line><g><text style=\"dominant-baseline:alphabetic;font-size:10px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(0.0,-6)\"><tspan style=\"dominant-baseline:inherit\">0</tspan></text><text style=\"dominant-baseline:alphabetic;font-size:10px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(75.0,-6)\"><tspan style=\"dominant-baseline:inherit\">1</tspan></text><text style=\"dominant-baseline:alphabetic;font-size:10px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(150.0,-6)\"><tspan style=\"dominant-baseline:inherit\">2</tspan></text><text style=\"dominant-baseline:alphabetic;font-size:10px;font-weight:normal;stroke:none;text-anchor:middle\" transform=\"translate(225.00000000000003,-6)\"><tspan style=\"dominant-baseline:inherit\">3</tspan></text></g><g class=\"toyplot-coordinates-Axis-coordinates\" style=\"visibility:hidden\" transform=\"\"><line style=\"stroke:rgb(43.9%,50.2%,56.5%);stroke-opacity:1.0;stroke-width:1.0\" x1=\"0\" x2=\"0\" y1=\"3.0\" y2=\"-4.5\"></line><text style=\"alignment-baseline:hanging;fill:rgb(43.9%,50.2%,56.5%);fill-opacity:1.0;font-size:10px;font-weight:normal;stroke:none;text-anchor:middle\" x=\"0\" y=\"6\"></text></g></g></g></svg><div class=\"toyplot-interactive\"><ul class=\"toyplot-mark-popup\" onmouseleave=\"this.style.visibility='hidden'\" style=\"background:rgba(0%,0%,0%,0.75);border:0;border-radius:6px;color:white;cursor:default;list-style:none;margin:0;padding:5px;position:fixed;visibility:hidden\">\n",
       "            <li class=\"toyplot-mark-popup-title\" style=\"color:lightgray;cursor:default;padding:5px;list-style:none;margin:0\"></li>\n",
       "            <li class=\"toyplot-mark-popup-save-csv\" onmouseout=\"this.style.color='white';this.style.background='steelblue'\" onmouseover=\"this.style.color='steelblue';this.style.background='white'\" style=\"border-radius:3px;padding:5px;list-style:none;margin:0\">\n",
       "                Save as .csv\n",
       "            </li>\n",
       "        </ul><script>\n",
       "        (function()\n",
       "        {\n",
       "          var data_tables = [{\"title\": \"Plot Data\", \"names\": [\"x\", \"y0\"], \"id\": \"tb666e086654c4dc59f6f31a1784d85be\", \"columns\": [[3.25, 1.625], [0.5, 1.0]], \"filename\": \"toyplot\"}, {\"title\": \"Plot Data\", \"names\": [\"x\", \"y0\"], \"id\": \"t0d5716b3113447ebaccbb399ad64c5af\", \"columns\": [[0.875, 2.9375], [1.5, 2.0]], \"filename\": \"toyplot\"}];\n",
       "\n",
       "          function save_csv(data_table)\n",
       "          {\n",
       "            var uri = \"data:text/csv;charset=utf-8,\";\n",
       "            uri += data_table.names.join(\",\") + \"\\n\";\n",
       "            for(var i = 0; i != data_table.columns[0].length; ++i)\n",
       "            {\n",
       "              for(var j = 0; j != data_table.columns.length; ++j)\n",
       "              {\n",
       "                if(j)\n",
       "                  uri += \",\";\n",
       "                uri += data_table.columns[j][i];\n",
       "              }\n",
       "              uri += \"\\n\";\n",
       "            }\n",
       "            uri = encodeURI(uri);\n",
       "\n",
       "            var link = document.createElement(\"a\");\n",
       "            if(typeof link.download != \"undefined\")\n",
       "            {\n",
       "              link.href = uri;\n",
       "              link.style = \"visibility:hidden\";\n",
       "              link.download = data_table.filename + \".csv\";\n",
       "\n",
       "              document.body.appendChild(link);\n",
       "              link.click();\n",
       "              document.body.removeChild(link);\n",
       "            }\n",
       "            else\n",
       "            {\n",
       "              window.open(uri);\n",
       "            }\n",
       "          }\n",
       "\n",
       "          function open_popup(data_table)\n",
       "          {\n",
       "            return function(e)\n",
       "            {\n",
       "              var popup = document.querySelector(\"#t23633f2d1ffd483bbf18a7012b847b3f .toyplot-mark-popup\");\n",
       "              popup.querySelector(\".toyplot-mark-popup-title\").innerHTML = data_table.title;\n",
       "              popup.querySelector(\".toyplot-mark-popup-save-csv\").onclick = function() { popup.style.visibility = \"hidden\"; save_csv(data_table); }\n",
       "              popup.style.left = (e.clientX - 50) + \"px\";\n",
       "              popup.style.top = (e.clientY - 20) + \"px\";\n",
       "              popup.style.visibility = \"visible\";\n",
       "              e.stopPropagation();\n",
       "              e.preventDefault();\n",
       "            }\n",
       "\n",
       "          }\n",
       "\n",
       "          for(var i = 0; i != data_tables.length; ++i)\n",
       "          {\n",
       "            var data_table = data_tables[i];\n",
       "            var event_target = document.querySelector(\"#\" + data_table.id);\n",
       "            event_target.oncontextmenu = open_popup(data_table);\n",
       "          }\n",
       "        })();\n",
       "        </script><script>\n",
       "        (function()\n",
       "        {\n",
       "            function _sign(x)\n",
       "            {\n",
       "                return x < 0 ? -1 : x > 0 ? 1 : 0;\n",
       "            }\n",
       "\n",
       "            function _mix(a, b, amount)\n",
       "            {\n",
       "                return ((1.0 - amount) * a) + (amount * b);\n",
       "            }\n",
       "\n",
       "            function _log(x, base)\n",
       "            {\n",
       "                return Math.log(Math.abs(x)) / Math.log(base);\n",
       "            }\n",
       "\n",
       "            function _in_range(a, x, b)\n",
       "            {\n",
       "                var left = Math.min(a, b);\n",
       "                var right = Math.max(a, b);\n",
       "                return left <= x && x <= right;\n",
       "            }\n",
       "\n",
       "            function inside(range, projection)\n",
       "            {\n",
       "                for(var i = 0; i != projection.length; ++i)\n",
       "                {\n",
       "                    var segment = projection[i];\n",
       "                    if(_in_range(segment.range.min, range, segment.range.max))\n",
       "                        return true;\n",
       "                }\n",
       "                return false;\n",
       "            }\n",
       "\n",
       "            function to_domain(range, projection)\n",
       "            {\n",
       "                for(var i = 0; i != projection.length; ++i)\n",
       "                {\n",
       "                    var segment = projection[i];\n",
       "                    if(_in_range(segment.range.bounds.min, range, segment.range.bounds.max))\n",
       "                    {\n",
       "                        if(segment.scale == \"linear\")\n",
       "                        {\n",
       "                            var amount = (range - segment.range.min) / (segment.range.max - segment.range.min);\n",
       "                            return _mix(segment.domain.min, segment.domain.max, amount)\n",
       "                        }\n",
       "                        else if(segment.scale[0] == \"log\")\n",
       "                        {\n",
       "                            var amount = (range - segment.range.min) / (segment.range.max - segment.range.min);\n",
       "                            var base = segment.scale[1];\n",
       "                            return _sign(segment.domain.min) * Math.pow(base, _mix(_log(segment.domain.min, base), _log(segment.domain.max, base), amount));\n",
       "                        }\n",
       "                    }\n",
       "                }\n",
       "            }\n",
       "\n",
       "            function display_coordinates(e)\n",
       "            {\n",
       "                var current = svg.createSVGPoint();\n",
       "                current.x = e.clientX;\n",
       "                current.y = e.clientY;\n",
       "\n",
       "                for(var axis_id in axes)\n",
       "                {\n",
       "                    var axis = document.querySelector(\"#\" + axis_id);\n",
       "                    var coordinates = axis.querySelector(\".toyplot-coordinates-Axis-coordinates\");\n",
       "                    if(coordinates)\n",
       "                    {\n",
       "                        var projection = axes[axis_id];\n",
       "                        var local = current.matrixTransform(axis.getScreenCTM().inverse());\n",
       "                        if(inside(local.x, projection))\n",
       "                        {\n",
       "                            var domain = to_domain(local.x, projection);\n",
       "                            coordinates.style.visibility = \"visible\";\n",
       "                            coordinates.setAttribute(\"transform\", \"translate(\" + local.x + \")\");\n",
       "                            var text = coordinates.querySelector(\"text\");\n",
       "                            text.textContent = domain.toFixed(2);\n",
       "                        }\n",
       "                        else\n",
       "                        {\n",
       "                            coordinates.style.visibility= \"hidden\";\n",
       "                        }\n",
       "                    }\n",
       "                }\n",
       "            }\n",
       "\n",
       "            var root_id = \"t23633f2d1ffd483bbf18a7012b847b3f\";\n",
       "            var axes = {\"tb0fd033e2db944288db3e18d0448a556\": [{\"domain\": {\"bounds\": {\"max\": Infinity, \"min\": -Infinity}, \"max\": 3.0, \"min\": 0.0}, \"range\": {\"bounds\": {\"max\": Infinity, \"min\": -Infinity}, \"max\": 225.00000000000003, \"min\": 0.0}, \"scale\": \"linear\"}]};\n",
       "\n",
       "            var svg = document.querySelector(\"#\" + root_id + \" svg\");\n",
       "            svg.addEventListener(\"click\", display_coordinates);\n",
       "        })();\n",
       "        </script></div></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## store newick tree and a migration event list [source, sink, start, stop, rate]\n",
    "## if not edges are entered then it is converted to cladogram\n",
    "newick = \"(((a,b),Cow), (d,e));\"\n",
    "events = [['e', 'Cow', 0, 1, 1e-6], \n",
    "          ['3', '1', 1, 2, 1e-6]]\n",
    "\n",
    "## initiate Tree object with newick and admix args\n",
    "tre = baba.Tree(newick=newick, \n",
    "                admix=events)\n",
    "\n",
    "## show the tree \n",
    "tre.draw(width=250, height=250, yaxis=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4.0, 3.0, 2.0, 1.0, 0.0]\n",
      "['d', 'e', 'Cow', 'a', 'b']\n"
     ]
    }
   ],
   "source": [
    "## a way of finding names xpos from verts\n",
    "tre.verts[tre.verts[:, 1] == 0]\n",
    "tre.tree.search_nodes(name='b')[0].idx\n",
    "\n",
    "tre.verts[6, 0]\n",
    "\n",
    "print [tre.verts[tre.tree.search_nodes(name=name)[0].idx, 0]\n",
    "                for name in tre.tree.get_leaf_names()]\n",
    "\n",
    "print tre.tree.get_leaf_names()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### simulate data on that tree\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "## returns a Sim data object\n",
    "sims = tre.simulate(nreps=10000, Ns=50000, gen=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sims <generator object _replicate_generator at 0x7fa709d3c870>\n",
      "debug 0\n",
      "nreps 10000\n",
      "names ['b', 'a', 'Cow', 'e', 'd']\n"
     ]
    }
   ],
   "source": [
    "## what is in the sims object? 4 attributes:\n",
    "for key, val in sims.__dict__.items():\n",
    "    print key, val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'test' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-3d3bdbaf406e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mbaba\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_msp_to_arr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msims\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'test' is not defined"
     ]
    }
   ],
   "source": [
    "baba._msp_to_arr(sims, test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def _msp_to_arr(Sim, test):\n",
    "    \n",
    "    ## the fixed tree dictionary\n",
    "    #fix = {j: [i, i+1] for j, i in zip(list(\"abcdefghijkl\"), range(0, 24, 2))}\n",
    "    fix = {j: [i, i+1] for j, i in zip(Sim.names, range(0, len(Sim.names)*2, 2))}\n",
    "    \n",
    "    ## fill taxdict by test\n",
    "    keys = ['p1', 'p2', 'p3', 'p4']\n",
    "    arr = np.zeros((Sim.nreps, 4, 100))\n",
    "    \n",
    "    ## unless it's a 5-taxon test\n",
    "    if len(test) == 5:\n",
    "        arr = np.zeros((100000, 6, 100))\n",
    "        keys += ['p5']\n",
    "    \n",
    "    ## create array sampler for taxa\n",
    "    taxs = [test[key] for key in keys]\n",
    "    idxs = [list(itertools.chain(*[fix[j] for j in i])) for i in taxs]\n",
    "    print idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[2, 3], [0, 1], [4, 5], [6, 7, 8, 9]]\n"
     ]
    }
   ],
   "source": [
    "from ipyrad.analysis.baba import *\n",
    "_msp_to_arr(sims, test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### A function to print the variant matrix of one rep.\n",
    "This is used just for demonstration. The matrix will have 2X as many columns as there are \n",
    "tips in the tree, corresponding to two alleles per individual. The number of rows is the \n",
    "number of variant sites simulated. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def print_variant_matrix(tree):\n",
    "    shape = tree.get_num_mutations(), tree.get_sample_size()\n",
    "    arr = np.empty(shape, dtype=\"u1\")\n",
    "    for variant in tree.variants():\n",
    "        arr[variant.index] = variant.genotypes\n",
    "    print(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0 0 0 0 0 0 1 1 0 0]\n",
      " [1 1 1 1 0 0 0 0 0 0]\n",
      " [0 0 1 1 0 0 0 0 0 0]\n",
      " [1 0 1 1 0 0 0 0 0 0]\n",
      " [0 1 0 0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 1 1 1 1]]\n"
     ]
    }
   ],
   "source": [
    "## in order of ladderized tip names (2 copies per tip)\n",
    "print_variant_matrix(sims.sims.next())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Simulate data for ABBA-BABA tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<ipyrad.analysis.baba.Sim at 0x7f92881bbc50>"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "## simulate data on tree\n",
    "#sims = tre.simulate(nreps=10000, Ns=50000, gen=20)\n",
    "\n",
    "## pass sim object to baba\n",
    "test = {\n",
    "    'p4': ['e', 'd'],\n",
    "    'p3': ['Cow'], \n",
    "    'p2': ['b'],\n",
    "    'p1': ['a'],    \n",
    "}\n",
    "\n",
    "#baba.batch(sims, test, ipyclient=ipyclient)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 2",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}