[3] | 1 | import sys, os |
---|
| 2 | import unittest |
---|
| 3 | import tempfile |
---|
| 4 | try: |
---|
| 5 | sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) |
---|
| 6 | except: |
---|
| 7 | sys.path.insert(0, os.path.dirname(os.path.abspath("."))) |
---|
| 8 | |
---|
| 9 | from bx.arrays.array_tree import ArrayTree, FileArrayTree, FileArrayTreeDict, array_tree_dict_from_reader |
---|
| 10 | from bx.arrays.bed import BedReader |
---|
| 11 | from bx.arrays.wiggle import WiggleReader |
---|
| 12 | |
---|
| 13 | class TestArrayTree(unittest.TestCase): |
---|
| 14 | def setUp(self): |
---|
| 15 | tree = ArrayTree(10000, 10) # max value of 10000, each block has 10 numbers |
---|
| 16 | for i in range(5000): |
---|
| 17 | tree[i] = i |
---|
| 18 | |
---|
| 19 | # Insert extra copies to test frequency |
---|
| 20 | for i in range(3000): |
---|
| 21 | tree[i] = i |
---|
| 22 | |
---|
| 23 | tree.set_range(5000, 9001, 100) |
---|
| 24 | tree.root.build_summary() |
---|
| 25 | |
---|
| 26 | d = {'test': tree} |
---|
| 27 | f = tempfile.TemporaryFile() |
---|
| 28 | FileArrayTreeDict.dict_to_file( d, f ) |
---|
| 29 | f.seek(0) |
---|
| 30 | self.filearraytreedict = FileArrayTreeDict(f) |
---|
| 31 | self.filearraytree = self.filearraytreedict['test'] |
---|
| 32 | |
---|
| 33 | def test_get_summary(self): |
---|
| 34 | f = self.filearraytree |
---|
| 35 | lvl1 = f.get_summary(0, 1) |
---|
| 36 | self.assertEqual( map(float, lvl1.sums/lvl1.counts), [4.5, 14.5, 24.5, 34.5, 44.5, 54.5, 64.5, 74.5, 84.5, 94.5]) |
---|
| 37 | lvl2 = f.get_summary(0, 2) |
---|
| 38 | self.assertEqual( map(float, lvl2.sums/lvl2.counts), [49.5, 149.5, 249.5, 349.5, 449.5, 549.5, 649.5, 749.5, 849.5, 949.5]) |
---|
| 39 | lvl3 = f.get_summary(0, 3) |
---|
| 40 | self.assertEqual( map(float, lvl3.sums/lvl3.counts), [499.5, 1499.5, 2499.5, 3499.5, 4499.5, 100.0, 100.0, 100.0, 100.0, 100.0]) |
---|
| 41 | lvl2_2 = f.get_summary(3000, 2) |
---|
| 42 | self.assertEqual( map(float, lvl2_2.sums/lvl2_2.counts), [3049.5, 3149.5, 3249.5, 3349.5, 3449.5, 3549.5, 3649.5, 3749.5, 3849.5, 3949.5]) |
---|
| 43 | |
---|
| 44 | def test_get_leaf(self): |
---|
| 45 | f = self.filearraytree |
---|
| 46 | from_start = [int(i) for i in f.get_leaf(0)] |
---|
| 47 | from_middle = [int(i) for i in f.get_leaf(5)] |
---|
| 48 | self.assertEqual(from_start, from_middle) |
---|
| 49 | self.assertEqual(from_start, range(10)) |
---|
| 50 | |
---|
| 51 | from_start = [int(i) for i in f.get_leaf(4999)] |
---|
| 52 | self.assertEqual(from_start, range(4990, 5000)) |
---|
| 53 | |
---|
| 54 | from_start = [int(i) for i in f.get_leaf(9600)] |
---|
| 55 | self.assertEqual(from_start, []) |
---|
| 56 | |
---|
| 57 | def test_big(self): |
---|
| 58 | tree = ArrayTree(2147483647, 1000) # What we use for tracks |
---|
| 59 | for i in range(5000): |
---|
| 60 | tree[i] = i |
---|
| 61 | |
---|
| 62 | # Insert extra copies to test frequency |
---|
| 63 | for i in range(3000): |
---|
| 64 | tree[i] = i |
---|
| 65 | |
---|
| 66 | tree.set_range(5000, 9001, 100) |
---|
| 67 | tree.set_range(14000000, 15000000, 200) |
---|
| 68 | tree.root.build_summary() |
---|
| 69 | |
---|
| 70 | d = {'test': tree} |
---|
| 71 | f = tempfile.TemporaryFile() |
---|
| 72 | FileArrayTreeDict.dict_to_file( d, f ) |
---|
| 73 | f.seek(0) |
---|
| 74 | at = FileArrayTreeDict(f)['test'] |
---|
| 75 | |
---|
| 76 | lvl1 = at.get_summary(14000000, 1) |
---|
| 77 | avgs = map(float, lvl1.sums/lvl1.counts) |
---|
| 78 | self.assertEqual( len(avgs), 1000 ) |
---|
| 79 | self.assertEqual( avgs, [ 200 for i in range(0, 1000)] ) |
---|
| 80 | |
---|
| 81 | |
---|
| 82 | # def create_bed(self): |
---|
| 83 | # reader = BedReader( open( "22.bed.txt" ) ) |
---|
| 84 | # temp = tempfile.TemporaryFile() |
---|
| 85 | # |
---|
| 86 | # d = array_tree_dict_from_reader( reader, {}, block_size = 1000 ) |
---|
| 87 | # |
---|
| 88 | # for array_tree in d.itervalues(): |
---|
| 89 | # array_tree.root.build_summary() |
---|
| 90 | # |
---|
| 91 | # FileArrayTreeDict.dict_to_file( d, open("tree.at", "w"), no_leaves=True ) # just summaries |
---|
| 92 | # |
---|
| 93 | # def test_bed(self): |
---|
| 94 | # # self.create_bed() |
---|
| 95 | # print "bed" |
---|
| 96 | # at = FileArrayTreeDict( open( "tree.at" ) )['chr22'] |
---|
| 97 | # print map(, at.get_summary(14000000, 1).frequencies) |
---|
| 98 | |
---|
| 99 | |
---|
| 100 | def test_get_frequencies(self): |
---|
| 101 | f = self.filearraytree |
---|
| 102 | self.assertEqual( map(float, f.get_summary(0, 1).frequencies), ([20] * 10) ) |
---|
| 103 | self.assertEqual( map(float, f.get_summary(4000, 1).frequencies), ([10] * 10) ) |
---|
| 104 | self.assertEqual( map(float, f.get_summary(0, 2).frequencies), ([200] * 10) ) |
---|
| 105 | self.assertEqual( map(int, f.get_summary(0, 3).frequencies), [2000, 2000, 2000, 1000, 1000, 1000, 1000, 1000, 1000, 1] ) |
---|
| 106 | |
---|
| 107 | def test_wrong_dictkey(self): |
---|
| 108 | self.assertRaises(KeyError, self.filearraytreedict.__getitem__, "non-existing") |
---|
| 109 | |
---|
| 110 | def test_higher_level_than_tree(self): |
---|
| 111 | f = self.filearraytree |
---|
| 112 | self.assertEqual(3, f.levels) |
---|
| 113 | self.assertRaises(ValueError, f.get_summary, 0, 4) |
---|
| 114 | |
---|
| 115 | |
---|
| 116 | if __name__ == '__main__': |
---|
| 117 | unittest.main() |
---|