Cookbook: TreeMix

The program TreeMix by Pickrell & Pritchard (2012) is used to infer population splits and admixture from allele frequency data. From the TreeMix documentation: "In the underlying model, the modern-day populations in a species are related to a common ancestor via a graph of ancestral populations. We use the allele frequencies in the modern populations to infer the structure of this graph."

Required software

In [1]:
## conda install -c ipyrad ipyrad
## conda install -c ipyrad treemix
## conda install -c eaton-lab toytree
In [2]:
import ipyrad.analysis as ipa
import toytree
import toyplot
import numpy as np
In [3]:
print 'ipyrad', ipa.__version__
print 'toytree', toytree.__version__
print 'toyplot', toyplot.__version__
! treemix --version | grep 'TreeMix v. '
ipyrad 0.7.1
toytree 0.1.3
toyplot 0.15.0-dev
TreeMix v. 1.12

Define the populations

These two dictionaries will be used to parse out SNPs from the ipyrad SNPs output file (.snps.phy), to filter the SNPs for inclusion in the data set based on missing data (minmap), and to group them into populations (imap).

In [4]:
## a dictionary mapping sample names to 'species' names
imap = {
    "prz": ["32082_przewalskii", "33588_przewalskii"],
    "cys": ["41478_cyathophylloides", "41954_cyathophylloides"],
    "cya": ["30686_cyathophylla"],
    "sup": ["29154_superba"],
    "cup": ["33413_thamno"],
    "tha": ["30556_thamno"],
    "rck": ["35236_rex"],
    "rex": ["35855_rex", "40578_rex"],
    "lip": ["39618_rex", "38362_rex"],  
    }

## optional: loci will be filtered if they do not have data for at
## least N samples in each species. Minimums cannot be <1.
minmap = {
    "prz": 2,
    "cys": 2,
    "cya": 1,
    "sup": 1,
    "cup": 1,
    "tha": 1, 
    "rck": 1,
    "rex": 2,
    "lip": 2,
    }

Create a Treemix object

The format for this object is similar to many of the other ipyrad.analysis objects.

In [5]:
t = ipa.treemix(
    name="test",
    data="analysis-ipyrad/pedicularis_outfiles/pedicularis.snps.phy",
    imap=imap,
    minmap=minmap,
    )
In [6]:
## you can set additional parameter args here
t.params.m = 1
t.params.root = "prz"
t.params.global_ = 1
#t.params.se = 1
t.params
Out[6]:
binary      treemix             
bootstrap   0                   
climb       0                   
cormig      0                   
g           (None, None)        
global_     1                   
k           0                   
m           1                   
root        prz                 
se          0                   

Generate the treemix input file

In [7]:
## write treemix input files so you can call treemix from the command line
s = t.write_output_file()
ntaxa 13; nSNPs total 194653; nSNPs written 12247

The command string

This shows the command string that corresponds to the parameter settings in the Treemix object. You can see that the input file (-i) is the string we enetered in the data field above, and the output prefix (-o) corresponds to the default working directory and the name field that we provided above. In addition, the argument (-m 1) is added because we added that to the params dictionary.

In [8]:
## the command string
print t.command
treemix -i /home/deren/Documents/ipyrad/tests/analysis-treemix/test.treemix.in.gz -o /home/deren/Documents/ipyrad/tests/analysis-treemix/test -m 1 -root prz -global
In [9]:
## you can run the command in a notebook by using bash one-liners (!)
! $t.command > analysis-treemix/treemix.log

Run treemix jobs

Alternatively, you can use the .run() command of the treemix object to run treemix. This is more convenient because the results will automatically be parsed by the treemix object so that they are easily accessible for downstream plotting. In the loop below we run treemix over a range of migration parameters (-m) and with 4 replicates per setting.

In [10]:
## a dictionary for storing treemix objects
tdict = {}

## iterate over values of m
for rep in xrange(4):
    for mig in xrange(4):
        
        ## create new treemix object copy
        name = "mig-{}-rep-{}".format(mig, rep)
        tmp = t.copy(name)

        ## set params on new object
        tmp.params.m = mig
    
        ## run treemix analysis
        tmp.run()
        
        ## store the treemix object
        tdict[name] = tmp

Accessible results

In [11]:
## choose a treemix object from the above analysis
t = tdict['mig-2-rep-3']
In [12]:
## access output files produced by treemix
t.files
Out[12]:
cov        ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.cov.gz
covse      ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.covse.gz
edges      ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.edges.gz
llik       ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.llik
modelcov   ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.modelcov.gz
treeout    ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.treeout.gz
vertices   ~/Documents/ipyrad/tests/analysis-treemix/mig-2-rep-3.vertices.gz
In [13]:
## access the newick string representation of the tree
t.results.tree
Out[13]:
'((((tha:0.0528677,(lip:0.0360458,(rex:0.015358,rck:0.0400938):0.0013585):0.000729372):0.0061529,cup:0.0402529):0.0346747,((sup:0.0493299,cya:0.0485717):0.0144389,cys:0.0693095):0.016622):0.00587823,prz:0.309518);'
In [14]:
## access a list of admixture edges
t.results.admixture
Out[14]:
[(12, 0.0402529, 13, 0.0528677, 0.087311),
 (8, 0.309518, 16, 0.0400938, 0.0459799)]
In [15]:
## access the covariance matrix
t.results.modelcov
Out[15]:
array([[ 0.267703  , -0.0258312 , -0.0268198 , -0.0267356 , -0.0364621 ,
        -0.0393228 , -0.038564  , -0.0365603 , -0.037407  ],
       [-0.0258312 ,  0.0819627 ,  0.0116645 ,  0.0117488 , -0.0145997 ,
        -0.0174604 , -0.0167016 , -0.0146979 , -0.0160852 ],
       [-0.0268198 ,  0.0116645 ,  0.0744446 ,  0.025199  , -0.0155884 ,
        -0.0184491 , -0.0176903 , -0.0156866 , -0.0170739 ],
       [-0.0267356 ,  0.0117488 ,  0.025199  ,  0.0738549 , -0.0155041 ,
        -0.0183649 , -0.0176061 , -0.0156024 , -0.0169897 ],
       [-0.0364621 , -0.0145997 , -0.0155884 , -0.0155041 ,  0.0496971 ,
         0.00940488,  0.00734218,  0.00934589,  0.00636426],
       [-0.0393228 , -0.0174604 , -0.0184491 , -0.0183649 ,  0.00940488,
         0.0531333 ,  0.0100971 ,  0.0121008 ,  0.00886101],
       [-0.038564  , -0.0167016 , -0.0176903 , -0.0176061 ,  0.00734218,
         0.0100971 ,  0.0481683 ,  0.0141262 ,  0.0108281 ],
       [-0.0365603 , -0.0146979 , -0.0156866 , -0.0156024 ,  0.00934589,
         0.0121008 ,  0.0141262 ,  0.0328464 ,  0.0141279 ],
       [-0.037407  , -0.0160852 , -0.0170739 , -0.0169897 ,  0.00636426,
         0.00886101,  0.0108281 ,  0.0141279 ,  0.0473745 ]])

You can generate plots in R

Follow the directions in the Treemix tutorial for plotting results.

Or, you can generate plots in Python using Toytree

The code to produce tree plots with admixture edges in Toytree is still in development and will change to be more user friendly in the near future.

In [16]:
def _get_admix_point(tre, idx, dist):
    ## parent coordinates
    px, py = tre.verts[idx]
    ## child coordinates
    cx, cy = tre.verts[tre.tree.search_nodes(idx=idx)[0].up.idx]
    ## angle of hypotenuse
    theta = np.arctan((px-cx) / (py-cy))
    ## new coords along the hypot angle
    horz = np.sin(theta) * dist
    vert = np.cos(theta) * dist
    
    ## change x
    a = tre.verts[idx, 0]
    b = tre.verts[idx, 1] 
    a -= abs(horz)
    if py < cy:
        b += abs(vert)
    else:
        b -= abs(vert)
    return a, b
In [17]:
def treemix_plot(tmp, axes):
        
    ## create a toytree object from the treemix tree result
    tre = toytree.tree(newick=tmp.results.tree)
    tre.draw(
        axes=axes,
        use_edge_lengths=True,
        tree_style='c',
        tip_labels_align=True,
        edge_align_style={"stroke-width": 1}
    );

    ## get coords 
    for admix in tmp.results.admixture:
        ## parse admix event
        pidx, pdist, cidx, cdist, weight = admix
        a = _get_admix_point(tre, pidx, pdist)
        b = _get_admix_point(tre, cidx, cdist)

        ## add line for admixture edge
        mark = axes.plot(
            a = (a[0], b[0]),
            b = (a[1], b[1]),
            style={"stroke-width": 10*weight,
                   "stroke-opacity": 0.95, 
                   "stroke-linecap": "round"}
        )

        ## add points at admixture sink
        axes.scatterplot(
            a = (b[0]),
            b = (b[1]),
            size=8,
            title="weight: {}".format(weight),
        )

    ## add scale bar for edge lengths
    axes.y.show=False
    axes.x.ticks.show=True
    axes.x.label.text = "Drift parameter"
    return axes

Draw a single result

In [18]:
toyplot.__version__
Out[18]:
'0.15.0-dev'
In [24]:
## select a result
tmp = tdict["mig-1-rep-0"]

## draw the tree similar to the Treemix plotting R code
canvas = toyplot.Canvas(width=350, height=350)
axes = canvas.cartesian(padding=25, margin=75)
axes = treemix_plot(tmp, axes)
przcyssupcyathacupliprexrckweight: 0.137873-0.3-0.2-0.10.0Drift parameter

Draw replicate runs

In [25]:
canvas = toyplot.Canvas(width=800, height=1200)
idx = 0
for mig in range(4):
    for rep in range(4):
        tmp = tdict["mig-{}-rep-{}".format(mig, rep)]
        ax = canvas.cartesian(grid=(4, 4, idx), padding=25, margin=(25, 50, 100, 25))
        ax = treemix_plot(tmp, ax)
        idx += 1
przcyssupcyacupthaliprckrex-0.2-0.10.0Drift parameterprzcyssupcyacupthaliprckrex-0.2-0.10.0Drift parameterprzcyssupcyacupthaliprexrck-0.2-0.10.0Drift parameterprzcyscyasupcupthaliprexrck-0.2-0.10.0Drift parameterprzcyssupcyathacupliprexrckweight: 0.137873-0.3-0.2-0.10.0Drift parameterprzcyscyasupthacupliprexrckweight: 0.137961-0.3-0.2-0.10.0Drift parameterprzcyscyasupthacupliprckrexweight: 0.137873-0.3-0.2-0.10.0Drift parameterprzcyscyasupthacupliprckrexweight: 0.137873-0.3-0.2-0.10.0Drift parameterprzcyscyasupcupthaliprexrckweight: 0.0460636weight: 0.139281-0.3-0.2-0.10.0Drift parameterprzcyssupcyacupthaliprexrckweight: 0.0963147weight: 0.0459799-0.3-0.2-0.10.0Drift parameterprzcyscyasupcupthaliprexrckweight: 0.0460636weight: 0.139281-0.3-0.2-0.10.0Drift parameterprzcyssupcyacupthaliprexrckweight: 0.087311weight: 0.0459799-0.3-0.2-0.10.0Drift parameterprzsupcyacyscupthaliprexrckweight: 0.141451weight: 0.457816weight: 0.0459704-0.3-0.2-0.10.0Drift parameterprzsupcyacysthacupliprexrckweight: 0.141451weight: 0.457816weight: 0.0459704-0.3-0.2-0.10.0Drift parameterprzsupcyacysthacupliprexrckweight: 0.141451weight: 0.457816weight: 0.0459704-0.3-0.2-0.10.0Drift parameterprzsupcyacyscupthaliprexrckweight: 0.141451weight: 0.457816weight: 0.0459704-0.3-0.2-0.10.0Drift parameter
In [26]:
## grab names from the tree
tre = toytree.tree(tmp.results.tree)

lnames = toyplot.locator.Explicit(
    locations=range(len(tre.get_tip_labels())),
    labels=tre.get_tip_labels(),
)

## get a colormap and plot the matrix
cmap = toyplot.color.diverging.map("BlueRed", tmp.results.cov.min(), tmp.results.cov.max())
canvas, table = toyplot.matrix(
              (tmp.results.cov, cmap),
               width=400, 
               height=400, 
               bshow=True,
               tshow=False,
               llocator=lnames,
               blocator=lnames,      
              );

## add a color scale
tlocs = np.linspace(tmp.results.cov.min(), tmp.results.cov.max(), 5)
tlabs = ["{:.2f}".format(i) for i in tlocs]
canvas.color_scale(cmap, 
                   x1=100, x2=300, y1=50, y2=50, 
                   ticklocator=toyplot.locator.Explicit(
                       locations=tlocs, 
                       labels=tlabs,
                   ));
                       
prz0.267703-0.0270258-0.0280964-0.0242687-0.0364397-0.0391614-0.038536-0.0367678-0.037407sup-0.02702580.07444460.0251990.0111396-0.0156012-0.0181699-0.0177945-0.015368-0.0168238cya-0.02809640.0251990.07385490.0122737-0.0152837-0.0181586-0.0174872-0.0153056-0.0169961cys-0.02426870.01113960.01227370.0819627-0.0148089-0.018092-0.0170124-0.0149584-0.0162356cup-0.0364397-0.0156012-0.0152837-0.01480890.04969710.009404880.007483780.009634720.0059131tha-0.0391614-0.0181699-0.0181586-0.0180920.009404880.05313330.01019190.01232250.00852938lip-0.038536-0.0177945-0.0174872-0.01701240.007483780.01019190.04816830.01346840.0115176rex-0.0367678-0.015368-0.0153056-0.01495840.009634720.01232250.01346840.03284640.0141279rck-0.037407-0.0168238-0.0169961-0.01623560.00591310.008529380.01151760.01412790.0473745przsupcyacyscupthaliprexrck-0.040.040.110.190.27