added gtr argument to timetree inference script
#!/usr/bin/env python
from __future__ import print_function, division
import numpy as np
-from treetime import TreeTime
+from treetime import TreeTime, GTR
from Bio import Phylo, AlignIO
if __name__=="__main__":
@@ -20,7 +20,15 @@ if __name__=="__main__":
parser.add_argument('--tree', required = True, type = str, help ="newick file with tree")
parser.add_argument('--dates', required = True, type = str,
help ="csv with dates for nodes with 'node_name, date' where date is float (as in 2012.15)")
- parser.add_argument('--infer_gtr', default = True, action='store_true', help='infer substitution model')
+ # parser.add_argument('--infer_gtr', default = True, action='store_true', help='infer substitution model')
+ parser.add_argument('--gtr', required=True, type = str, help="GTR model to use. "
+ " Type 'infer' to infer the model from the data. Or, specify the model type. "
+ "Optionally, feed the arguments with the '--gtr_args' option")
+ parser.add_argument('--gtr_params', type=str, nargs='+', help="GTR parameters for the model "
+ "specified by the --gtr argument. The parameters should be feed as 'key=value' list of parameters. "
+ "Example: '--gtr K80 --gtr_params kappa=0.2 pis=0.25,0.25,0.25,0.25'. See the exact definitions of "
+ " the parameters in the GTR creation methods.")
+
parser.add_argument('--reroot', required = False, type = str, default='best',
help ="reroot the tree. Valid arguments are 'best', 'midpoint', or a node name")
parser.add_argument('--resolve_polytomies', default = True, action='store_true',
@@ -54,14 +62,35 @@ if __name__=="__main__":
except:
continue
+
+ ###########################################################################
+ ### GTR SET-UP
+ ###########################################################################
+ model = params.gtr
+ gtr_params = params.gtr_params
+ if model == 'infer':
+ gtr = GTR.standard('jc')
+ infer_gtr = True
+ else:
+ kwargs = {}
+ for param in gtr_params:
+ keyval = param.split('=')
+ if len(keyval)!=2: continue
+ if keyval[0] in ['pis', 'pi', 'Pi', 'Pis']:
+ keyval[1] = map(int, keyval[1].split(','))
+ kwargs[keyval[0]] = keyval[1]
+
+ gtr = GTR.standard(model, **kwargs)
+ infer_gtr = False
+
###########################################################################
### ANCESTRAL RECONSTRUCTION AND SET-UP
###########################################################################
myTree = TreeTime(dates=dates, tree=params.tree,
- aln=params.aln, gtr='JC69', verbose=params.verbose)
+ aln=params.aln, gtr=gtr, verbose=params.verbose)
myTree.run(root=params.reroot, relaxed_clock=params.relax,
resolve_polytomies=params.resolve_polytomies,
- Tc=params.Tc, max_iter=params.max_iter)
+ Tc=params.Tc, max_iter=params.max_iter, infer_gtr=infer_gtr)
###########################################################################
### OUTPUT and saving of results
