def test_union_samples(self):
s1 = argbook.Segment()
s1.samples.update({k: k
for k in [0, 1, 2]})
s1.left = 0
s1.right = 10
s2 = argbook.Segment()
s2.samples.update({k: k
for k in [3, 4]})
s2.left = 0
s2.right = 10
self.assertEqual(sorted(s1.union_samples(s2)), [0, 1, 2, 3, 4])
self.assertEqual(sorted(s1.samples), [0, 1, 2])
def test_num_links(self):
node = argbook.Node(2)
s1 = argbook.Segment()
s2 = argbook.Segment()
s1.left = 10
s1.right = 20
s1.node = node
s2.left = 0
s2.right = 5
s2.node = node
s1.prev = s2
s2.next = s1
node.first_segment = s2
self.assertEqual(node.num_links(), 19)
self.assertNotEqual(node.num_links(), 20)
def test_get_first_segment(self):
s1 = argbook.Segment()
s2 = argbook.Segment()
s1.samples.update({k: k for k in [0, 1, 2]})
s1.left = 10
s1.right = 20
s1.node = argbook.Node(2)
s2.samples.update({k: k for k in [0, 1]})
s2.left = 0
s2.right = 4
s2.node = argbook.Node(2)
s1.prev = s2
s2.next = s1
self.assertEqual(s1.get_first_segment().left, 0)
self.assertEqual(s1.get_first_segment().right, 4)
self.assertEqual(sorted(s1.get_first_segment().samples), [0, 1])
def test_defrag_segment_chain(self):
s1 = argbook.Segment()
s2 = argbook.Segment()
s1.samples.update({k: k for k in [0, 1, 2]})
s1.left = 10
s1.right = 20
s1.node = argbook.Node(2)
s2.samples.update({k: k for k in [0, 1, 2]})
s2.left = 0
s2.right = 10
s2.node = argbook.Node(2)
s1.prev = s2
s2.next = s1
s1.defrag_segment_chain()
self.assertEqual(s2.right, 20)
self.assertEqual(s2.left, 0)
self.assertEqual(sorted(s2.samples), [0, 1, 2])
def test_contains(self):
s1 = argbook.Segment()
s1.samples.update({k: k
for k in [0, 1, 2]})
s1.left = 0
s1.right = 10
self.assertFalse(s1.contains(10), False)
self.assertTrue(s1.contains(0), True)
self.assertTrue(s1.contains(7), True)
def test_is_mrca(self):
sample_size = 5
s1 = argbook.Segment()
s1.samples.update({k: k
for k in [0, 1, 2, 3, 4]})
s1.left = 0
s1.right = 10
self.assertTrue(s1.is_mrca(sample_size), True)
self.assertEqual(len(set(s1.samples)), len(s1.samples))
self.assertTrue(len(s1.samples) <= sample_size)
def test_equal(self):
s1 = argbook.Segment()
s1.samples = [0, 1, 2]
s1.left = 0
s1.right = 10
s1.node = argbook.Node(2)
s2 = argbook.Segment()
s2.samples = [0, 1, 2]
s2.left = 0
s2.right = 10
s2.node = argbook.Node(2)
self.assertTrue(s1.equal(s2))
#----
s3 = argbook.Segment()
s3.samples.update({k: k
for k in [0, 1, 2]})
s3.left = 0
s3.right = 10
s3.node = argbook.Node(1)
self.assertFalse(s3.equal(s1))
def test_equal_samples(self):
#equal
s1 = argbook.Segment()
s1.samples.update({k: k
for k in [0, 1, 2]})
s1.left = 0
s1.right = 10
s2 = argbook.Segment()
s2.samples.update({k: k
for k in [0, 1, 2]})
s2.left = 0
s2.right = 10
self.assertTrue(s1.equal_samples(s2), True)
#------ not equal
s3 = argbook.Segment()
s3.samples.update({k: k
for k in [0, 1]})
s3.left = 1
s3.right = 10
self.assertFalse(s1.equal_samples(s3), False)
def test_contains(self):
node = argbook.Node(2)
s1 = argbook.Segment()
s2 = argbook.Segment()
s1.left = 10
s1.right = 20
s1.node = node
s2.left = 0
s2.right = 10
s2.node = node
s1.prev = s2
s2.next = s1
node.first_segment = s2
self.assertFalse(node.contains(20))
self.assertTrue(node.contains(10))
self.assertTrue(node.contains(0))
self.assertEqual(node.first_segment.prev, None)
#---- testx_segment here
self.assertEqual(node.x_segment(5).left, 0)
self.assertEqual(node.x_segment(10).left, 10)
self.assertRaises(ValueError, node.x_segment, 20)
self.assertRaises(ValueError, node.x_segment, 25)
def test_detach_update(self):
'''detach a node and then update ancestral material
RE: (2)b=5--> (3, 4), t= 1.5
CA: (3, 4)--> 6, t= 2.5
CA: (0, 1)--> 5, t= 3.5
CA: (5, 6)--> 7, t= 4.5
'''
arg = argbook.ARG()
arg.nodes[0] = argbook.Node(0)
arg.nodes[0].first_segment = argbook.Segment()
arg.nodes[1] = argbook.Node(1)
arg.nodes[1].first_segment = argbook.Segment()
arg.nodes[2] = argbook.Node(2)
arg.nodes[2].first_segment = argbook.Segment()
arg.nodes[0].time = 0
arg.nodes[1].time = 0
arg.nodes[2].time = 0
#--------- n = 3 , m=0 , seq length =10
arg.nodes[0].first_segment.left = 0
arg.nodes[0].first_segment.right = 10
arg.nodes[1].first_segment.left = 0
arg.nodes[1].first_segment.right = 10
arg.nodes[2].first_segment.left = 0
arg.nodes[2].first_segment.right = 10
arg.nodes[0].first_segment.samples.__setitem__(0, 0)
arg.nodes[1].first_segment.samples.__setitem__(1, 1)
arg.nodes[2].first_segment.samples.__setitem__(2, 2)
#----rec on 2 at b= 5
arg.nodes[3] = argbook.Node(3)
arg.nodes[3].first_segment = argbook.Segment()
arg.nodes[3].first_segment.left = 0
arg.nodes[3].first_segment.right = 5
arg.nodes[3].left_child = arg.nodes[2]
arg.nodes[3].right_child = arg.nodes[2]
arg.nodes[3].first_segment.samples.__setitem__(2, 2)
arg.nodes[4] = argbook.Node(4)
arg.nodes[4].first_segment = argbook.Segment()
arg.nodes[4].first_segment.left = 5
arg.nodes[4].first_segment.right = 10
arg.nodes[4].left_child = arg.nodes[2]
arg.nodes[4].right_child = arg.nodes[2]
arg.nodes[4].first_segment.samples.__setitem__(2, 2)
arg.nodes[2].left_parent = arg.nodes[3]
arg.nodes[2].right_parent = arg.nodes[4]
arg.nodes[2].breakpoint = 5
arg.nodes[3].time = 1.5
arg.nodes[4].time = 1.5
#------ CA (3,4 ) ---6
arg.nodes[6] = argbook.Node(6)
arg.nodes[6].first_segment = argbook.Segment()
arg.nodes[6].first_segment.left = 0
arg.nodes[6].first_segment.right = 10
arg.nodes[6].left_child = arg.nodes[3]
arg.nodes[6].right_child = arg.nodes[4]
arg.nodes[6].first_segment.samples.__setitem__(2, 2)
arg.nodes[6].time = 2.5
arg.nodes[3].left_parent = arg.nodes[6]
arg.nodes[3].right_parent = arg.nodes[6]
arg.nodes[4].left_parent = arg.nodes[6]
arg.nodes[4].right_parent = arg.nodes[6]
#-------- CA (0,1 ) ---> 5
arg.nodes[5] = argbook.Node(5)
arg.nodes[5].first_segment = argbook.Segment()
arg.nodes[5].first_segment.left = 0
arg.nodes[5].first_segment.right = 10
arg.nodes[5].left_child = arg.nodes[0]
arg.nodes[5].right_child = arg.nodes[1]
arg.nodes[5].first_segment.samples.update({0: 0, 1: 1})
arg.nodes[5].time = 3.5
arg.nodes[0].left_parent = arg.nodes[5]
arg.nodes[0].right_parent = arg.nodes[5]
arg.nodes[1].left_parent = arg.nodes[5]
arg.nodes[1].right_parent = arg.nodes[5]
#--- CA (5, 6) ---7
arg.nodes[7] = argbook.Node(7)
arg.nodes[7].first_segment = argbook.Segment()
arg.nodes[7].left_child = arg.nodes[5]
arg.nodes[7].right_child = arg.nodes[6]
arg.nodes[7].time = 4.5
arg.nodes[5].left_parent = arg.nodes[7]
arg.nodes[5].right_parent = arg.nodes[7]
arg.nodes[6].left_parent = arg.nodes[7]
arg.nodes[6].right_parent = arg.nodes[7]
#-------
mcmc = MC.MCMC(ts_full='test')
mcmc.arg = arg
mcmc.arg.coal.update({5: 5, 6: 6, 7: 7})
mcmc.arg.rec.update({3: 3, 4: 4})
mcmc.data = {}
mcmc.n = 3
mcmc.m = 0
mcmc.seq_length = 10
# mcmc.print_state()
# -- - - - -
detach = mcmc.arg.nodes[6]
old_merger_time = detach.left_parent.time
sib = mcmc.arg.nodes[6].sibling()
mcmc.arg.detach(detach, sib)
#--- parent is root --> add both to floating
mcmc.floatings[detach.time] = detach.index
mcmc.floatings[detach.left_parent.time] = sib.index
assert sib.left_parent is None
if sib.left_parent is not None:
mcmc.update_all_ancestral_material(sib)
mcmc.spr_reattach_floatings(detach, sib, old_merger_time)
# print("coal", mcmc.arg.coal)
# mcmc.print_state()
mcmc.spr()
def test_equal(self):
node1 = argbook.Node(2)
s1 = argbook.Segment()
s2 = argbook.Segment()
s1.samples.update({k: k for k in [0, 1, 2]})
s1.left = 10
s1.right = 20
s1.node = node1
s2.samples.update({k: k for k in [0, 1]})
s2.left = 0
s2.right = 5
s2.node = node1
s1.prev = s2
s2.next = s1
node1.first_segment = s2
node2 = argbook.Node(2)
s3 = argbook.Segment()
s4 = argbook.Segment()
s3.samples.update({k: k for k in [0, 1, 2]})
s3.left = 10
s3.right = 20
s3.node = node2
s4.samples.update({k: k for k in [0, 1]})
s4.left = 0
s4.right = 5
s4.node = node2
s3.prev = s4
s4.next = s3
node2.first_segment = s4
self.assertTrue(node2.equal(node1))
self.assertTrue(node1.equal(node2))
#----- not equal
s5 = argbook.Segment()
s5.samples.update({k: k for k in [0, 1]})
s5.left = 21
s5.right = 25
s5.node = node2
s3.next = s5
s5.prev = s3
self.assertFalse(node2.equal(node1))
self.assertFalse(node1.equal(node2))
#---- raise
node3 = None
self.assertRaises(ValueError, node2.equal, node3)
#-------snps
node4 = argbook.Node(2)
s3 = argbook.Segment()
s4 = argbook.Segment()
s3.samples.update({k: k for k in [0, 1, 2]})
s3.left = 10
s3.right = 20
s3.node = node4
s4.samples.update({k: k for k in [0, 1]})
s4.left = 0
s4.right = 5
s4.node = node4
s3.prev = s4
s4.next = s3
node4.first_segment = s4
node4.snps.update({12: 12, 23: 23})
self.assertFalse(node1.equal(node4))
node1.snps.update({23: 23, 12: 12})
self.assertTrue(node1.equal(node4))