-
Notifications
You must be signed in to change notification settings - Fork 55
/
test_treetime.py
273 lines (226 loc) · 11.7 KB
/
test_treetime.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
from __future__ import print_function
from io import StringIO
# Tests
def test_import_short():
print("testing short imports")
from treetime import GTR
from treetime import TreeTime
from treetime import TreeAnc
from treetime import seq_utils
def test_assign_gamma(root_dir=None):
print("testing assign gamma")
import os
from treetime import TreeTime
from treetime.utils import parse_dates
if root_dir is None:
root_dir = os.path.dirname(os.path.realpath(__file__))
fasta = root_dir + "/treetime_examples/data/h3n2_na/h3n2_na_20.fasta"
nwk = root_dir + "/treetime_examples/data/h3n2_na/h3n2_na_20.nwk"
dates = parse_dates(root_dir + "/treetime_examples/data/h3n2_na/h3n2_na_20.metadata.csv")
seq_kwargs = {"marginal_sequences":True,
"branch_length_mode": 'input',
"reconstruct_tip_states":False}
tt_kwargs = {'clock_rate': 0.0001,
'time_marginal':'assign'}
myTree = TreeTime(gtr='Jukes-Cantor', tree = nwk, use_fft=False,
aln = fasta, verbose = 1, dates = dates, precision=3, debug=True)
def assign_gamma(tree):
return tree
success = myTree.run(infer_gtr=False, assign_gamma=assign_gamma, max_iter=1, verbose=3, **seq_kwargs, **tt_kwargs)
assert success
def test_GTR():
from treetime import GTR
import numpy as np
for model in ['Jukes-Cantor']:
print('testing GTR, model:',model)
myGTR = GTR.standard(model, alphabet='nuc')
print('Frequency sum:', myGTR.Pi.sum())
assert (myGTR.Pi.sum() - 1.0)**2<1e-14
# the matrix is the rate matrix
assert abs(myGTR.Q.sum(0)).sum() < 1e-14
# eigendecomposition is made correctly
n_states = myGTR.v.shape[0]
assert abs((myGTR.v.dot(myGTR.v_inv) - np.identity(n_states)).sum() < 1e-10)
assert np.abs(myGTR.v.sum()) > 1e-10 # **and** v is not zero
def test_reconstruct_discrete_traits():
from Bio import Phylo
from treetime.wrappers import reconstruct_discrete_traits
# Create a minimal tree with traits to reconstruct.
tiny_tree = Phylo.read(StringIO("((A:0.60100000009,B:0.3010000009):0.1,C:0.2):0.001;"), 'newick')
traits = {
"A": "?",
"B": "North America",
"C": "West Asia",
}
# Reconstruct traits with "?" as missing data.
mugration, letter_to_state, reverse_alphabet = reconstruct_discrete_traits(
tiny_tree,
traits,
missing_data="?",
)
# With two known states, the letters "A" and "B" should be in the alphabet
# mapping to those states.
assert "A" in letter_to_state
assert "B" in letter_to_state
# The letter for missing data should be the next letter in the alphabet,
# following the two known state letters.
assert letter_to_state["C"] == "?"
def test_ancestral(root_dir=None):
import os
from Bio import AlignIO
import numpy as np
from treetime import TreeAnc, GTR
if root_dir is None:
root_dir = os.path.dirname(os.path.realpath(__file__))
fasta = str(os.path.join(root_dir, 'treetime_examples/data/h3n2_na/h3n2_na_20.fasta'))
nwk = str(os.path.join(root_dir, 'treetime_examples/data/h3n2_na/h3n2_na_20.nwk'))
for marginal in [True, False]:
print('loading flu example')
t = TreeAnc(gtr='Jukes-Cantor', tree=nwk, aln=fasta)
print('ancestral reconstruction' + ("marginal" if marginal else "joint"))
t.reconstruct_anc(method='ml', marginal=marginal)
assert t.data.compressed_to_full_sequence(t.tree.root.cseq, as_string=True) == 'ATGAATCCAAATCAAAAGATAATAACGATTGGCTCTGTTTCTCTCACCATTTCCACAATATGCTTCTTCATGCAAATTGCCATCTTGATAACTACTGTAACATTGCATTTCAAGCAATATGAATTCAACTCCCCCCCAAACAACCAAGTGATGCTGTGTGAACCAACAATAATAGAAAGAAACATAACAGAGATAGTGTATCTGACCAACACCACCATAGAGAAGGAAATATGCCCCAAACCAGCAGAATACAGAAATTGGTCAAAACCGCAATGTGGCATTACAGGATTTGCACCTTTCTCTAAGGACAATTCGATTAGGCTTTCCGCTGGTGGGGACATCTGGGTGACAAGAGAACCTTATGTGTCATGCGATCCTGACAAGTGTTATCAATTTGCCCTTGGACAGGGAACAACACTAAACAACGTGCATTCAAATAACACAGTACGTGATAGGACCCCTTATCGGACTCTATTGATGAATGAGTTGGGTGTTCCTTTTCATCTGGGGACCAAGCAAGTGTGCATAGCATGGTCCAGCTCAAGTTGTCACGATGGAAAAGCATGGCTGCATGTTTGTATAACGGGGGATGATAAAAATGCAACTGCTAGCTTCATTTACAATGGGAGGCTTGTAGATAGTGTTGTTTCATGGTCCAAAGAAATTCTCAGGACCCAGGAGTCAGAATGCGTTTGTATCAATGGAACTTGTACAGTAGTAATGACTGATGGAAGTGCTTCAGGAAAAGCTGATACTAAAATACTATTCATTGAGGAGGGGAAAATCGTTCATACTAGCACATTGTCAGGAAGTGCTCAGCATGTCGAAGAGTGCTCTTGCTATCCTCGATATCCTGGTGTCAGATGTGTCTGCAGAGACAACTGGAAAGGCTCCAATCGGCCCATCGTAGATATAAACATAAAGGATCATAGCATTGTTTCCAGTTATGTGTGTTCAGGACTTGTTGGAGACACACCCAGAAAAAACGACAGCTCCAGCAGTAGCCATTGTTTGGATCCTAACAATGAAGAAGGTGGTCATGGAGTGAAAGGCTGGGCCTTTGATGATGGAAATGACGTGTGGATGGGAAGAACAATCAACGAGACGTCACGCTTAGGGTATGAAACCTTCAAAGTCATTGAAGGCTGGTCCAACCCTAAGTCCAAATTGCAGATAAATAGGCAAGTCATAGTTGACAGAGGTGATAGGTCCGGTTATTCTGGTATTTTCTCTGTTGAAGGCAAAAGCTGCATCAATCGGTGCTTTTATGTGGAGTTGATTAGGGGAAGAAAAGAGGAAACTGAAGTCTTGTGGACCTCAAACAGTATTGTTGTGTTTTGTGGCACCTCAGGTACATATGGAACAGGCTCATGGCCTGATGGGGCGGACCTCAATCTCATGCCTATA'
print('testing LH normalization')
from Bio import Phylo,AlignIO
tiny_tree = Phylo.read(StringIO("((A:0.60100000009,B:0.3010000009):0.1,C:0.2):0.001;"), 'newick')
tiny_aln = AlignIO.read(StringIO(">A\nAAAAAAAAAAAAAAAACCCCCCCCCCCCCCCCGGGGGGGGGGGGGGGGTTTTTTTTTTTTTTTT\n"
">B\nAAAACCCCGGGGTTTTAAAACCCCGGGGTTTTAAAACCCCGGGGTTTTAAAACCCCGGGGTTTT\n"
">C\nACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGTACGT\n"), 'fasta')
mygtr = GTR.custom(alphabet = np.array(['A', 'C', 'G', 'T']), pi = np.array([0.9, 0.06, 0.02, 0.02]), W=np.ones((4,4)))
t = TreeAnc(gtr=mygtr, tree=tiny_tree, aln=tiny_aln)
t.reconstruct_anc('ml', marginal=True, debug=True)
lhsum = np.exp(t.sequence_LH(pos=np.arange(4**3))).sum()
print (lhsum)
assert(np.abs(lhsum-1.0)<1e-6)
t.optimize_branch_len()
def test_seq_joint_reconstruction_correct():
"""
evolve the random sequence, get the alignment at the leaf nodes.
Reconstruct the sequences of the internal nodes (joint)
and prove the reconstruction is correct.
In addition, compute the likelihood of the particular realization of the
sequences on the tree and prove that this likelihood is exactly the same
as calculated in the joint reconstruction
"""
from treetime import TreeAnc, GTR
from treetime import seq_utils
from Bio import Phylo, AlignIO
import numpy as np
from collections import defaultdict
def exclusion(a, b):
"""
Intersection of two lists
"""
return list(set(a) - set(b))
tiny_tree = Phylo.read(StringIO("((A:.060,B:.01200)C:.020,D:.0050)E:.004;"), 'newick')
mygtr = GTR.custom(alphabet = np.array(['A', 'C', 'G', 'T']),
pi = np.array([0.15, 0.95, 0.05, 0.3]),
W=np.ones((4,4)))
seq = np.random.choice(mygtr.alphabet, p=mygtr.Pi, size=400)
myTree = TreeAnc(gtr=mygtr, tree=tiny_tree, aln=None, verbose=4)
# simulate evolution, set resulting sequence as ref_seq
tree = myTree.tree
seq_len = 400
tree.root.ref_seq = np.random.choice(mygtr.alphabet, p=mygtr.Pi, size=seq_len)
print ("Root sequence: " + ''.join(tree.root.ref_seq.astype('U')))
mutation_list = defaultdict(list)
for node in tree.find_clades():
for c in node.clades:
c.up = node
if hasattr(node, 'ref_seq'):
continue
t = node.branch_length
p = mygtr.evolve( seq_utils.seq2prof(node.up.ref_seq, mygtr.profile_map), t)
# normalize profile
p=(p.T/p.sum(axis=1)).T
# sample mutations randomly
ref_seq_idxs = np.array([int(np.random.choice(np.arange(p.shape[1]), p=p[k])) for k in np.arange(p.shape[0])])
node.ref_seq = np.array([mygtr.alphabet[k] for k in ref_seq_idxs])
node.ref_mutations = [(anc, pos, der) for pos, (anc, der) in
enumerate(zip(node.up.ref_seq, node.ref_seq)) if anc!=der]
for anc, pos, der in node.ref_mutations:
mutation_list[pos].append((node.name, anc, der))
print (node.name, len(node.ref_mutations), node.ref_mutations)
# set as the starting sequences to the terminal nodes:
alnstr = ""
i = 1
for leaf in tree.get_terminals():
alnstr += ">" + leaf.name + "\n" + ''.join(leaf.ref_seq.astype('U')) + '\n'
i += 1
print (alnstr)
myTree.aln = AlignIO.read(StringIO(alnstr), 'fasta')
# reconstruct ancestral sequences:
myTree.infer_ancestral_sequences(final=True, debug=True, reconstruct_leaves=True)
diff_count = 0
mut_count = 0
for node in myTree.tree.find_clades():
if node.up is not None:
mut_count += len(node.ref_mutations)
diff_count += np.sum(node.sequence != node.ref_seq)
if np.sum(node.sequence != node.ref_seq):
print("%s: True sequence does not equal inferred sequence. parent %s"%(node.name, node.up.name))
else:
print("%s: True sequence equals inferred sequence. parent %s"%(node.name, node.up.name))
# the assignment of mutations to the root node is probabilistic. Hence some differences are expected
assert diff_count/seq_len<2*(1.0*mut_count/seq_len)**2
# prove the likelihood value calculation is correct
LH = myTree.ancestral_likelihood()
LH_p = (myTree.tree.sequence_LH)
print ("Difference between reference and inferred LH:", (LH - LH_p).sum())
assert ((LH - LH_p).sum())<1e-9
return myTree
def test_seq_joint_lh_is_max():
"""
For a single-char sequence, perform joint ancestral sequence reconstruction
and prove that this reconstruction is the most likely one by comparing to all
possible reconstruction variants (brute-force).
"""
from treetime import TreeAnc, GTR
from treetime import seq_utils
from Bio import Phylo, AlignIO
import numpy as np
mygtr = GTR.custom(alphabet = np.array(['A', 'C', 'G', 'T']), pi = np.array([0.91, 0.05, 0.02, 0.02]), W=np.ones((4,4)))
tiny_tree = Phylo.read(StringIO("((A:.0060,B:.30)C:.030,D:.020)E:.004;"), 'newick')
#terminal node sequences (single nuc)
A_char = 'A'
B_char = 'C'
D_char = 'G'
# for brute-force, expand them to the strings
A_seq = ''.join(np.repeat(A_char,16))
B_seq = ''.join(np.repeat(B_char,16))
D_seq = ''.join(np.repeat(D_char,16))
#
def ref_lh():
"""
reference likelihood - LH values for all possible variants
of the internal node sequences
"""
tiny_aln = AlignIO.read(StringIO(">A\n" + A_seq + "\n"
">B\n" + B_seq + "\n"
">D\n" + D_seq + "\n"
">C\nAAAACCCCGGGGTTTT\n"
">E\nACGTACGTACGTACGT\n"), 'fasta')
myTree = TreeAnc(gtr=mygtr, tree = tiny_tree,
aln =tiny_aln, verbose = 4)
logLH_ref = myTree.ancestral_likelihood()
return logLH_ref
#
def real_lh():
"""
Likelihood of the sequences calculated by the joint ancestral
sequence reconstruction
"""
tiny_aln_1 = AlignIO.read(StringIO(">A\n"+A_char+"\n"
">B\n"+B_char+"\n"
">D\n"+D_char+"\n"), 'fasta')
myTree_1 = TreeAnc(gtr=mygtr, tree = tiny_tree,
aln=tiny_aln_1, verbose = 4)
myTree_1.reconstruct_anc(method='ml', marginal=False, debug=True)
logLH = myTree_1.tree.sequence_LH
return logLH
ref = ref_lh()
real = real_lh()
print(abs(ref.max() - real) )
# joint chooses the most likely realization of the tree
assert(abs(ref.max() - real) < 1e-10)
return ref, real