Mercurial > repos > bcclaywell > argo_navis

comparison bin/pact_wrapper.py @ 0:d67268158946 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload commit a3f181f5f126803c654b3a66dd4e83a48f7e203b
author bcclaywell
date Mon, 12 Oct 2015 17:43:33 -0400
parents
children
comparison
equal deleted inserted replaced
-1:000000000000 0:d67268158946
1 #!/usr/bin/env python
2
3 import argparse
4 from os import path
5 import re
6 import os
7 import csv
8 import shutil
9 import subprocess
10
11
12 setting_template = """
13 # TREE MANIPULATION
14 {prune}
15
16 # TREE OUTPUT
17 print rule tree # print out the tree...
18
19 # SUMMARY STATS
20 summary tmrca # time to mrca
21 summary proportions # proportions on trunk of geneology
22 summary coal rates # coalescent rates; separate for each label
23 summary mig rates # migration rates; separate for each label pair
24 #summary fst # diversity between vs within labels
25
26 skyline settings -{sky_start} -{sky_end} {sky_interval}
27 skyline proportions
28 """
29
30
31 def tips_from_label(meta_reader, label, deme_col="deme"):
32 return [row['sequence'] for row in meta_reader if row[deme_col] == label]
33
34 def tips_from_labels(meta_reader, labels, deme_col="deme"):
35 return [row['sequence'] for row in meta_reader if row[deme_col] in labels]
36
37 def translate_tips(tips, translation):
38 result = list()
39 for tip in tips:
40 try:
41 result.append(translation[tip])
42 except KeyError:
43 pass
44 return result
45
46
47 def translation_from_nexus(handle):
48 """This reads the seqname -> id translation from the trees file. Should probably eventually move this out
49 and into a separate file so that we can use the results in plotting. Not now though..."""
50 in_trans = False
51 trans = dict()
52 for line in handle:
53 if re.search("Translate", line):
54 # This is the demarkator for the translation section of the nexus file
55 in_trans = True
56 continue
57 if in_trans:
58 if re.match(";", line):
59 # If we hit this, we've reached the end of the translation region
60 break
61 else:
62 # Otherwise we're still in the Translate region, so populate trans
63 index, name = line.strip().strip(',').split()
64 trans[name] = index
65 return trans
66
67
68 def prune_tips_string(tips, args):
69 if args.translate_trees:
70 trans = translation_from_nexus(file(args.trees_in))
71 tips = translate_tips(tips, trans)
72 return "prune to tips " + " ".join(tips)
73
74
75 def prune_string(args):
76 if args.labels:
77 if args.metadata:
78 labels = args.labels.replace("\"", "").replace("'", "").split()
79 print "Pruning to labels:", labels
80 tips = tips_from_labels(csv.DictReader(args.metadata), labels, deme_col=args.deme_col)
81 print "Tips for labels are:", tips
82 return prune_tips_string(tips, args)
83 else:
84 return "prune to label " + args.labels
85
86 elif args.tip_file or args.tips:
87 if args.tip_file:
88 tips = args.tip_file.read().split()
89 else:
90 tips = args.tips.replace("'", "").replace("\"", "").split()
91 print "Tips specified are:", tips
92 return prune_tips_string(tips, args)
93 else:
94 return ""
95
96
97 def get_args():
98 parser = argparse.ArgumentParser()
99 parser.add_argument('trees_in')
100 parser.add_argument('-t', '--tips', help="List of tips in quotes")
101 parser.add_argument('-T', '--tip-file', help="File of tips sep by space", type=argparse.FileType('r'))
102 parser.add_argument('-r', '--translate-trees', action="store_true", default=False,
103 help="""If flagged, the trees_in file will be used for translating tip names from indices to actual
104 tip names; this is necessary for BEAST runs only""")
105 parser.add_argument('-m', '--metadata', type=argparse.FileType('r'),
106 help="""Required for filtering by deme/label with the beast method""")
107 parser.add_argument('-l', '--labels')
108 parser.add_argument('-s', '--sky-start', default=2.0, type=float,
109 help="How far back to compute skyline statistics (should be a potive number)")
110 parser.add_argument('-S', '--trim-start', type=float,
111 help="""Trim the tree from 0 back to the specified time; overrides sky-start for skyline plots.
112 (should be a positive number)""")
113 parser.add_argument('-e', '--trim-end', default=0.0, type=float,
114 help="Most recent time to trim to in time interval; will also be used for sky (should be non-negative)")
115 parser.add_argument('-o', '--out-dir')
116 parser.add_argument('-p', '--prune-to-trunk', action="store_true", default=False)
117 parser.add_argument('-d', '--deme-col', default="deme", help="Deme column in metadata file")
118 return parser.parse_args()
119
120
121 def main():
122 args = get_args()
123
124 # Create the param file in the proper directory
125 outfile = file(path.join(args.out_dir, 'in.param'), 'w')
126
127 prune = prune_string(args)
128 if args.prune_to_trunk:
129 prune += "\nprune to trunk"
130
131 if args.trim_start:
132 prune += "\ntrim ends -%s -%s" % (args.trim_start, args.trim_end)
133
134 sky_start = args.trim_start if args.trim_start else args.sky_start
135 sky_end = args.trim_end
136 sky_interval = abs(sky_end - sky_start) / 100
137 outfile.write(setting_template.format(
138 prune=prune,
139 sky_end=sky_end,
140 sky_start=sky_start,
141 sky_interval=sky_interval))
142 outfile.close()
143
144 # Copy the tree file over to the directory
145 shutil.copyfile(args.trees_in, path.join(args.out_dir, 'in.trees'))
146
147 # Actually run PACT
148 os.chdir(args.out_dir)
149 subprocess.check_call(['pact'])
150
151
152 if __name__ == '__main__':
153 main()
154
155