added gtr argument to ancestral reconstruction script
#!/usr/bin/env python
from __future__ import print_function, division
import numpy as np
-from treetime import TreeAnc
+from treetime import TreeAnc, GTR
from Bio import Phylo, AlignIO
if __name__=="__main__":
@@ -16,6 +16,16 @@ if __name__=="__main__":
' like _A45G_... The inferred GTR model is written to stdout')
parser.add_argument('--aln', required = True, type = str, help ="fasta file with input sequences")
parser.add_argument('--tree', required = True, type = str, help ="newick file with tree")
+
+ 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('--prot', default = False, action="store_true", help ="protein alignment")
parser.add_argument('--marginal', default = False, action='store_true', help='marginal instead of joint ML reconstruction')
parser.add_argument('--infer_gtr', default = False, action='store_true', help='infer substitution model')
@@ -25,14 +35,34 @@ if __name__=="__main__":
###########################################################################
### ANCESTRAL RECONSTRUCTION
###########################################################################
- model = 'aa' if params.prot else 'Jukes-Cantor'
- treeanc = TreeAnc(params.tree, aln=params.aln, gtr=model, verbose=4, fill_overhangs=not params.keep_overhangs)
- treeanc.infer_ancestral_sequences('ml', infer_gtr=params.infer_gtr,
- marginal=params.marginal)
+
+ # creata GTR model
+ 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
+
+ treeanc = TreeAnc(params.tree, aln=params.aln, gtr=gtr, verbose=4, fill_overhangs=not params.keep_overhangs)
+ treeanc.infer_ancestral_sequences('ml', infer_gtr=infer_gtr,
+ marginal=params.3)
###########################################################################
### OUTPUT and saving of results
###########################################################################
+
+ model = 'aa' if params.prot else 'Jukes-Cantor'
if params.infer_gtr:
print('\nInferred GTR model:')
print(treeanc.gtr)
