| 15 | |
| 16 | |
| 17 | class BedIntervalTree(object): |
| 18 | def __init__(self): |
| 19 | """reads in a BED file and converts it to an interval tree for searching""" |
| 20 | self.tree = defaultdict(IntervalTree) |
| 21 | self.intCount = 0 |
| 22 | |
| 23 | def mkzero(): |
| 24 | return int(0) |
| 25 | self.count_by_label = defaultdict(mkzero) |
| 26 | self.nt_count_by_label = defaultdict(mkzero) |
| 27 | |
| 28 | def __str__(self): |
| 29 | return str(self.intCount) + ' intervals' |
| 30 | |
| 31 | def __repr__(self): |
| 32 | return str(self) |
| 33 | |
| 34 | def _addEntryToTree(self, bedentry, label): |
| 35 | """ Add a BED entry to the tree |
| 36 | :param bedentry: BED entry [chr, start, stop(, optional extra fields)] |
| 37 | :type bedentry: list of int and string |
| 38 | :param label: the label for the entry |
| 39 | :type label: str |
| 40 | """ |
| 41 | chrom = bedentry[0] |
| 42 | start = int(bedentry[1]) + 1 |
| 43 | end = int(bedentry[2]) + 1 |
| 44 | lbl = [label] + bedentry[3:] |
| 45 | currInt = Interval(start, end, value=lbl, chrom=chrom) |
| 46 | self.tree[chrom].add_interval(currInt) |
| 47 | self.count_by_label[label] += 1 |
| 48 | self.nt_count_by_label[label] += end - start |
| 49 | self.intCount += 1 |
| 50 | |
| 51 | def intersect(self, chrom, start, end): |
| 52 | """ Return all overlapping intervals in chr:[start,end) |
| 53 | :param chrom: Chromosome |
| 54 | :param start: start (1-based) |
| 55 | :param end: end |
| 56 | :rtype: list of Interval |
| 57 | |
| 58 | Intervals have a value associated, this value is an array -- the first column will be |
| 59 | the label, followed by the bed columns |
| 60 | |
| 61 | """ |
| 62 | return self.tree[chrom].find(start, end) |
| 63 | |
| 64 | def countbases(self, chrom=None, start=0, end=0, label=None): |
| 65 | """ Return the number of bases covered by intervals in chr:[start,end) |
| 66 | :param chrom: Chromosome |
| 67 | :param start: start (1-based) |
| 68 | :param end: end |
| 69 | :param label: label |
| 70 | :rtype: int |
| 71 | |
| 72 | Intervals have a value associated, this value is an array -- the first column will be |
| 73 | the label, followed by the bed columns |
| 74 | |