#!/usr/bin/env python
# coding: utf-8

# # <span style="color:gray">ipyrad-analysis toolkit:</span> RAxML
# 
# RAxML is one of the most popular tools for inferring phylogenetic trees using maximum likelihood. It is fast even for very large data sets. The documentation for raxml is huge, and there are many options. However, I tend to use the same small number of options very frequently, which motivated me to write the `ipa.raxml()` tool to automate the process of generating RAxml command line strings, running them, and accessing the resulting tree files. The simplicity of this tool makes it easy to incorporate into other more complex tools, for example, to infer tress in sliding windows along the genome using the `ipa.treeslider` tool.
# 
# 

# ### Required software

# In[1]:


# conda install ipyrad -c conda-forge -c bioconda
# conda install raxml -c conda-forge -c bioconda
# conda install ipcoal -c conda-forge


# In[2]:


import ipyrad.analysis as ipa
import toytree
import ipcoal


# ### Simulate example dataset

# In[3]:


# the true tree
tree = toytree.rtree.imbtree(ntips=10, treeheight=1e7)
tree.draw(ts='p');

# setup simulator
subst = {
    "state_frequencies": [0.3, 0.2, 0.3, 0.2],
    "kappa": 0.25, 
    "gamma": 0.20,
    "gamma_categories": 4,
}
mod = ipcoal.Model(tree=tree, Ne=1e5, nsamples=2, mut=1e-8, substitution_model=subst)
mod.sim_loci(nloci=1, nsites=10000)
mod.write_concat_to_phylip(name="raxtest", outdir="/tmp", diploid=True)


# ### Infer a ML tree

# In[4]:


# init raxml object with input data and (optional) parameter options
rax = ipa.raxml(data="/tmp/raxtest.phy", T=4, N=100)

# print the raxml command string for prosperity
print(rax.command)

# run the command, (options: block until finishes; overwrite existing)
rax.run(block=True, force=True)


# ### Draw the inferred tree
# After inferring a tree you can then visualize it in a notebook using `toytree`. 

# In[12]:


# load from the .trees attribute of the raxml object, or from the saved tree file
tre = toytree.tree(rax.trees.bipartitions)

# draw the tree
rtre = tre.root("r9")
rtre.draw(tip_labels_align=True, node_sizes=18, node_labels="support");


# ### Setting parameters
# 
# By default several parameters are pre-set in the raxml object. To remove those parameters from the command string you can set them to None. Additionally, you can build complex raxml command line strings by adding almost any parameter to the raxml object init, like below. You probably can't do everythin in raxml using this tool, it's only meant as a convenience. You can always of course just write the raxml command line string by hand instead.

# In[13]:


# init raxml object
rax = ipa.raxml(data="/tmp/raxtest.phy", T=4, N=10)

# parameter dictionary for a raxml object
rax.params


# In[14]:


# paths to output files produced by raxml inference
rax.trees


# ### Cookbook
# 
# Most frequently used: perform 100 rapid bootstrap analyses followed by 10 rapid hill-climbing ML searches from random starting trees under the GTRGAMMA substitution model. 

# In[15]:


rax = ipa.raxml(
    data="/tmp/raxtest.phy",
    name="test-1",
    workdir="analysis-raxml",
    m="GTRGAMMA",
    T=8,
    f="a",
    N=50,
)
print(rax.command)
rax.run(force=True)


# Another common option: Perform N rapid hill-climbing ML analyses from random starting trees, with no bootstrap replicates. Be sure to use the `BestTree` output from this analysis since it does not produce a `bipartitions` output file. 

# In[16]:


rax = ipa.raxml(
    data="/tmp/raxtest.phy",
    name="test-2",
    workdir="analysis-raxml",
    m="GTRGAMMA",
    T=8,
    f="d",
    N=10,
    x=None,
)
print(rax.command)
rax.run(force=True)


# ### Check your files
# The .info and related log files will be stored in the `workdir`. Be sure to look at these for further details of your analyses. 

# In[18]:


get_ipython().system(' cat ./analysis-raxml/RAxML_info.test-1')


# In[19]:


get_ipython().system(' cat ./analysis-raxml/RAxML_info.test-2')