Python intervaltree.Interval() Examples

def intervalListToIntervalTree(interval_list):
    r"""
    given a dictionary containing tuples of chrom, start, end,
    this is transformed to an interval trees. To each
    interval an id is assigned, this id corresponds to the
    position of the interval in the given array of tuples
    and if needed can be used to identify
    the index of a row/colum in the hic matrix.

    >>> bin_list = [('chrX', 0, 50000), ('chrX', 50000, 100000)]
    >>> res = intervalListToIntervalTree(bin_list)
    >>> sorted(res['chrX'])
    [Interval(0, 50000, 0), Interval(50000, 100000, 1)]
    """
    bin_int_tree = {}

    for intval_id, intval in enumerate(interval_list):
        chrom, start, end = intval[0:3]
        if chrom not in bin_int_tree:
            bin_int_tree[chrom] = IntervalTree()
        bin_int_tree[chrom].add(Interval(start, end, intval_id))

    return bin_int_tree 

def test_merge_overlaps_reducer_wo_initializer():
    def reducer(old, new):
        return "%s, %s" % (old, new)
    # empty tree
    e = IntervalTree()
    e.merge_overlaps(data_reducer=reducer)
    e.verify()
    assert not e

    # One Interval in tree
    o = IntervalTree.from_tuples([(1, 2, 'hello')])
    o.merge_overlaps(data_reducer=reducer)
    o.verify()
    assert len(o) == 1
    assert sorted(o) == [Interval(1, 2, 'hello')]

    # many Intervals in tree, with gap
    t = IntervalTree.from_tuples(data.ivs1.data)
    t.merge_overlaps(data_reducer=reducer)
    t.verify()
    assert len(t) == 2
    assert sorted(t) == [
        Interval(1, 2,'[1,2)'),
        Interval(4, 15, '[4,7), [5,9), [6,10), [8,10), [8,15), [10,12), [12,14), [14,15)')
    ] 

def test_slice():
    t = IntervalTree([Interval(5, 15)])
    t.slice(10)
    assert sorted(t)[0] == Interval(5, 10)
    assert sorted(t)[1] == Interval(10, 15)

    t = IntervalTree([Interval(5, 15)])
    t.slice(5)
    assert sorted(t)[0] == Interval(5, 15)

    t.slice(15)
    assert sorted(t)[0] == Interval(5, 15)

    t.slice(0)
    assert sorted(t)[0] == Interval(5, 15)

    t.slice(20)
    assert sorted(t)[0] == Interval(5, 15) 

def test_interval_sort_interval():
    base = Interval(0, 10)
    ivs = [
        Interval(-10, -5),
        Interval(-10, 0),
        Interval(-10, 5),
        Interval(-10, 10),
        Interval(-10, 20),
        Interval(0, 20),
        Interval(5, 20),
        Interval(10, 20),
        Interval(15, 20),
    ]

    for iv in ivs:
        sort = sorted([base, iv])
        assert sort[0].__cmp__(sort[1]) in (-1, 0)

        sort = sorted([iv, base])
        assert sort[0].__cmp__(sort[1]) in (-1, 0) 

def test_insert_to_filled_tree():
    t = IntervalTree.from_tuples(data.ivs1.data)
    orig = t.print_structure(True)  # original structure record

    assert match.set_data(t[1]) == set(['[1,2)'])
    t.add(Interval(1, 2, '[1,2)'))  # adding duplicate should do nothing
    assert match.set_data(t[1]) == set(['[1,2)'])
    assert orig == t.print_structure(True)

    t[1:2] = '[1,2)'                # adding duplicate should do nothing
    assert match.set_data(t[1]) == set(['[1,2)'])
    assert orig == t.print_structure(True)

    assert Interval(2, 4, '[2,4)') not in t
    t.add(Interval(2, 4, '[2,4)'))
    assert match.set_data(t[2]) == set(['[2,4)'])
    t.verify()

    t[13:15] = '[13,15)'
    assert match.set_data(t[14]) == set(['[8,15)', '[13,15)', '[14,15)'])
    t.verify() 

def test_interval_interval_cmp():
    """
    Test comparisons with other Intervals using __cmp__()
    """
    iv0 = Interval(0, 10)
    iv1 = Interval(-10, -5)
    iv2 = Interval(-10, 0)
    iv3 = Interval(-10, 5)
    iv4 = Interval(-10, 10)
    iv5 = Interval(-10, 20)
    iv6 = Interval(0, 20)
    iv7 = Interval(5, 20)
    iv8 = Interval(10, 20)
    iv9 = Interval(15, 20)

    assert iv0.__cmp__(iv0) == 0
    assert iv0.__cmp__(iv1) == 1
    assert iv0.__cmp__(iv2) == 1
    assert iv0.__cmp__(iv3) == 1
    assert iv0.__cmp__(iv4) == 1
    assert iv0.__cmp__(iv5) == 1
    assert iv0.__cmp__(iv6) == -1
    assert iv0.__cmp__(iv7) == -1
    assert iv0.__cmp__(iv8) == -1
    assert iv0.__cmp__(iv9) == -1 

def test_removei():
    # Empty tree
    e = IntervalTree()
    with pytest.raises(ValueError):
        e.removei(-1000, -999, "Doesn't exist")
    e.verify()
    assert len(e) == 0

    # Non-existent member should raise ValueError
    t = IntervalTree.from_tuples(data.ivs1.data)
    oldlen = len(t)
    with pytest.raises(ValueError):
        t.removei(-1000, -999, "Doesn't exist")
    t.verify()
    assert len(t) == oldlen

    # Should remove existing member
    assert Interval(1, 2, '[1,2)') in t
    t.removei(1, 2, '[1,2)')
    assert len(t) == oldlen - 1
    assert Interval(1, 2, '[1,2)') not in t 

def test_discardi():
    # Empty tree
    e = IntervalTree()
    e.discardi(-1000, -999, "Doesn't exist")
    e.verify()
    assert len(e) == 0

    # Non-existent member should do nothing quietly
    t = IntervalTree.from_tuples(data.ivs1.data)
    oldlen = len(t)
    t.discardi(-1000, -999, "Doesn't exist")
    t.verify()
    assert len(t) == oldlen

    # Should discard existing member
    assert Interval(1, 2, '[1,2)') in t
    t.discardi(1, 2, '[1,2)')
    assert len(t) == oldlen - 1
    assert Interval(1, 2, '[1,2)') not in t 

def __init__(self, data, row_sentence_bounds, window=5, process_all=False):
        """
        Class for managing windowed input data (like TIMIT).

        :param data: Numpy matrix. Each row should be an example data
        :param row_sentence_bounds:  Numpy matrix with bounds for padding. TODO add default NONE
        :param window: half-window size
        :param process_all: (default False) if True adds context to all data at object initialization.
                            Otherwise the windowed data is created in runtime.
        """
        self.window = window
        self.data = data
        base_shape = self.data.shape
        self.shape = (base_shape[0], (2 * self.window + 1) * base_shape[1])
        self.tree = it.IntervalTree([it.Interval(int(e[0]), int(e[1]) + 1) for e in row_sentence_bounds])
        if process_all:
            print('adding context to all the dataset', end='- ')
            self.data = self.generate_all()
            print('DONE')
        self.process_all = process_all 

def intervaltree_prep(h1, h2, ref_seq):
        """Mock up trees as would be done by the `VCFReader` class so we can test medaka.vcf._merge_variants

        :param h1, h2: iterable of variants in first and second haplotype, respectively.
        :param ref_seq: str, reference sequence

        :returns: (`intervaltree.Interval` containing interable of all variants,
                   [`intervaltree.IntervalTree` for each haplotype])
        """
        trees = []
        for variants in h1, h2:
            trees.append(intervaltree.IntervalTree())
            for v in variants:
                trees[-1].add(intervaltree.Interval(v.pos, v.pos + len(v.ref), data=v))
        only_overlapping=True
        comb_tree = intervaltree.IntervalTree(trees[0].all_intervals.union(trees[1].all_intervals))
        # if strict, merge only overlapping intervals (not adjacent ones)
        comb_tree.merge_overlaps(strict=only_overlapping, data_initializer=list(), data_reducer=lambda x,y: x + [y])
        comb_interval = list(comb_tree.all_intervals)[0]
        return comb_interval, trees 

def test_str():
    iv = Interval(0, 1)
    s = str(iv)
    assert s == 'Interval(0, 1)'
    assert repr(iv) == s

    iv = Interval(0, 1, '[0,1)')
    s = str(iv)
    assert s == "Interval(0, 1, '[0,1)')"
    assert repr(iv) == s

    iv = Interval((1,2), (3,4))
    s = str(iv)
    assert s == 'Interval((1, 2), (3, 4))'
    assert repr(iv) == s

    iv = Interval((1,2), (3,4), (5, 6))
    s = str(iv)
    assert s == 'Interval((1, 2), (3, 4), (5, 6))'
    assert repr(iv) == s 

def test_union_operator():
    t = IntervalTree()
    interval = Interval(0, 1)
    s = set([interval])

    # currently runs fine
    # with pytest.raises(TypeError):
    #     t | list(s)
    r = t | IntervalTree(s)
    assert len(r) == 1
    assert sorted(r)[0] == interval

    # also currently runs fine
    # with pytest.raises(TypeError):
    #     t |= s
    t |= IntervalTree(s)
    assert len(t) == 1
    assert sorted(t)[0] == interval 

def test_copy():
    itree = IntervalTree([Interval(0, 1, "x"), Interval(1, 2, ["x"])])
    itree.verify()

    itree2 = IntervalTree(itree)  # Shares Interval objects
    itree2.verify()

    itree3 = itree.copy()         # Shallow copy (same as above, as Intervals are singletons)
    itree3.verify()

    itree4 = pickle.loads(pickle.dumps(itree))  # Deep copy
    itree4.verify()

    list(itree[1])[0].data[0] = "y"
    assert sorted(itree) == [Interval(0, 1, 'x'), Interval(1, 2, ['y'])]
    assert sorted(itree2) == [Interval(0, 1, 'x'), Interval(1, 2, ['y'])]
    assert sorted(itree3) == [Interval(0, 1, 'x'), Interval(1, 2, ['y'])]
    assert sorted(itree4) == [Interval(0, 1, 'x'), Interval(1, 2, ['x'])] 

def test_membership():
    t = IntervalTree.from_tuples(data.ivs1.data)
    assert Interval(1, 2, '[1,2)') in t
    assert t.containsi(1, 2, '[1,2)')
    assert Interval(1, 3, '[1,3)') not in t
    assert not t.containsi(1, 3, '[1,3)')
    assert t.overlaps(4)
    assert t.overlaps(9)
    assert not t.overlaps(15)
    assert t.overlaps(0, 4)
    assert t.overlaps(1, 2)
    assert t.overlaps(1, 3)
    assert t.overlaps(8, 15)
    assert not t.overlaps(15, 16)
    assert not t.overlaps(-1, 0)
    assert not t.overlaps(2, 4) 

def test_overlaps_empty():
    # Empty tree
    t = IntervalTree()
    assert not t.overlaps(-1)
    assert not t.overlaps(0)

    assert not t.overlaps(-1, 1)
    assert not t.overlaps(-1, 0)
    assert not t.overlaps(0, 0)
    assert not t.overlaps(0, 1)
    assert not t.overlaps(1, 0)
    assert not t.overlaps(1, -1)
    assert not t.overlaps(0, -1)

    assert not t.overlaps(Interval(-1, 1))
    assert not t.overlaps(Interval(-1, 0))
    assert not t.overlaps(Interval(0, 0))
    assert not t.overlaps(Interval(0, 1))
    assert not t.overlaps(Interval(1, 0))
    assert not t.overlaps(Interval(1, -1))
    assert not t.overlaps(Interval(0, -1)) 

def combined_cds_introns(self):
        """This property returns the introns which are located between CDS
        segments in the combined CDS."""

        return self._combined_cds_introns

        # if self.number_internal_orfs < 2:
        #     return self.selected_cds_introns
        # if self.number_internal_orfs == 0 or len(self.combined_cds) < 2:
        #     return set()
        #
        # cintrons = []
        # for position in range(len(self.combined_cds) - 1):
        #     former = self.combined_cds[position]
        #     latter = self.combined_cds[position + 1]
        #     junc = intervaltree.Interval(former[1] + 1, latter[0] - 1)
        #     if junc in self.introns:
        #         cintrons.append(junc)
        # cintrons = set(cintrons)
        # return cintrons 

def selected_cds_introns(self):
        """This property returns the introns which are located between
        CDS segments in the selected ORF."""

        return self._selected_cds_introns

        # if len(self.selected_cds) < 2:
        #     return set()
        # if self.number_internal_orfs == 0 or len(self.combined_cds) < 2:
        #     return set()
        #
        # cintrons = []
        # for first, second in zip(self.selected_cds[:-1], self.selected_cds[1:]):
        #     cintrons.append(
        #         intervaltree.Interval(first[1] + 1,
        #                               second[0] - 1)
        #     )
        # cintrons = set(cintrons)
        # assert len(cintrons) > 0
        # return cintrons 

def test_chop():
    t = IntervalTree([Interval(0, 10)])
    t.chop(3, 7)
    assert len(t) == 2
    assert sorted(t)[0] == Interval(0, 3)
    assert sorted(t)[1] == Interval(7, 10)

    t = IntervalTree([Interval(0, 10)])
    t.chop(0, 7)
    assert len(t) == 1
    assert sorted(t)[0] == Interval(7, 10)

    t = IntervalTree([Interval(0, 10)])
    t.chop(5, 10)
    assert len(t) == 1
    assert sorted(t)[0] == Interval(0, 5)

    t = IntervalTree([Interval(0, 10)])
    t.chop(-5, 15)
    assert len(t) == 0

    t = IntervalTree([Interval(0, 10)])
    t.chop(0, 10)
    assert len(t) == 0 

def __init__(self, interval_tuples):
        '''intervals is like [('22', 12321, 12345, 'APOL1'), ...]'''
        self._its = {}
        self._gene_starts = {}
        self._gene_ends = {}
        for interval_tuple in interval_tuples:
            chrom, pos_start, pos_end, gene_name = interval_tuple
            assert isinstance(pos_start, int)
            assert isinstance(pos_end, int)
            if chrom not in self._its:
                self._its[chrom] = intervaltree.IntervalTree()
                self._gene_starts[chrom] = []
                self._gene_ends[chrom] = []
            self._its[chrom].add(intervaltree.Interval(pos_start, pos_end, gene_name))
            self._gene_starts[chrom].append((pos_start, gene_name))
            self._gene_ends[chrom].append((pos_end, gene_name))
        for chrom in self._its:
            self._gene_starts[chrom] = BisectFinder(self._gene_starts[chrom])
            self._gene_ends[chrom] = BisectFinder(self._gene_ends[chrom]) 

def _build_semester_lookup():
    """
    Build data structure to let us easily look up date -> strm.
    """
    all_semesters = Semester.objects.all()
    intervals = ((s.name, Semester.start_end_dates(s)) for s in all_semesters)
    intervals = (
        intervaltree.Interval(st, en+ONE_DAY, name)
        for (name, (st, en)) in intervals)
    return intervaltree.IntervalTree(intervals) 

def find_downstream(self, start, end, offset=0):

        """Method to use the segment tree to find all segments in the transcript downstream of a given interval."""

        return self.segmenttree.downstream_of_interval(Interval(start - offset, end),
                                                       num_intervals=self.exon_num + len(self.introns),
                                                       max_dist=10**9)


    # ###################Class methods##################################### 

def samples_to_bed(args):
    """Write a bed file from samples in a datastore file."""
    logger = medaka.common.get_named_logger('Variants')

    index = medaka.datastore.DataIndex(args.inputs)

    trees = collections.defaultdict(intervaltree.IntervalTree)
    logger.info("Building interval tree")
    for s, f in index.samples:
        d = medaka.common.Sample.decode_sample_name(s)
        # start and end are string repr of floats (major.minor coordinates)
        start, end = int(float(d['start'])), int(float(d['end']))
        # add one to end of interval, as intervaltree intervals and bed file
        # intervals are end-exclusive (i.e. they don't contain the last
        # coordinate), whilst the last position in a sample is included in that
        # sample.
        trees[d['ref_name']].add(intervaltree.Interval(start, end + 1))

    with open(args.output, 'w') as fh:
        for contig, tree in trees.items():
            # strict=False as consecutive samples can start and end on the same
            # major (overlap is in minor) hence if samples are abutting but not
            # overlapping in major coords, merge them
            tree.merge_overlaps(strict=False)
            logger.info("Writing intervals for {}".format(contig))
            for i in sorted(tree.all_intervals):
                fh.write("{}\t{}\t{}\n".format(contig, i.begin, i.end))

    logger.info("All done, bed file written to {}".format(args.output)) 

def __process_loop_file(self):
        interval_tree = {}

        with opener(self.properties['file']) as f:
            for idx, line in enumerate(f):
                line = to_string(line)
                # skip header line
                if idx == 0 and self.__is_header(line):
                    continue

                fields = line.split()
                chr1, x1, x2, chr2, y1, y2, *other = fields
                x1, x2, y1, y2 = list(map(int, [x1, x2, y1, y2]))

                # skip inter-chromosome interaction
                if chr1 != chr2:
                    continue
                chromosome = chr1

                if not chromosome.startswith("chr"):
                    chromosome = change_chrom_names(chromosome)
                if chromosome not in interval_tree:
                    interval_tree[chromosome] = IntervalTree()

                if len(other) == 0:
                    color = PlotHiCPeaks.DEFAULT_COLOR
                else:
                    rgb = other[0].split(",")
                    rgb = list(map(int, rgb))
                    color = rgb2hex(*rgb)

                loop = self.LoopInverval(chr1, x1, x2, chr2, y1, y2, color)
                interval_tree[chromosome].add(Interval(x1, y2, loop))

        return interval_tree 

def __process_bed(self):

        bed_file_h = ReadBed(opener(self.properties['file']))
        self.bed_type = bed_file_h.file_type

        if 'color' in self.properties and self.properties['color'] == 'bed_rgb' and \
                self.bed_type not in ['bed12', 'bed9']:
            log.warning("*WARNING* Color set to 'bed_rgb', but bed file does not have the rgb field. The color has "
                        "been set to {}".format(PlotTADCoverage.DEFAULT_COLOR))
            self.properties['color'] = PlotTADCoverage.DEFAULT_COLOR

        valid_intervals = 0
        interval_tree = {}

        max_score = float('-inf')
        min_score = float('inf')
        for bed in bed_file_h:
            if bed.score < min_score:
                min_score = bed.score
            if bed.score > max_score:
                max_score = bed.score

            interval_tree.setdefault(bed.chromosome, IntervalTree())
            interval_tree[bed.chromosome].add(Interval(bed.start, bed.end, bed))
            valid_intervals += 1

        if valid_intervals == 0:
            log.warning("No valid intervals were found in file {}".format(self.properties['file_name']))

        return interval_tree, min_score, max_score 

def write_interval_info(self, hwname, pclo=None, pchi=None,
                            substage_names=[], substage_entries={}):
        wt = self._get_writestable(hwname)
        if "framac" in hwname:
            return [(r['destlo'], r['desthi']) for r in
                    pytable_utils.get_rows('(%d <= writepclo) & (%d <= writepchi) & (writepclo < %d) & (writepchi <= %d)' % \
                                           (utils.addr_lo(pclo),
                                            utils.addr_hi(pclo),
                                            utils.addr_lo(pchi),
                                            utils.addr_hi(pchi)))]
        else:
            fns = substage_entries
            substages = substage_names
            num = 0
            intervals = {n: intervaltree.IntervalTree() for n in substages}

            for r in wt.read_sorted('index'):
                pc = long(r['pc'])
                if num < len(fns) - 1:
                    # check if we found the entrypoint to the next stage
                    (lopc, hipc) = substage_entries[num + 1]
                    if (lopc <= pc) and (pc < hipc):
                        num += 1
                if num in substages:
                    start = long(r['dest'])
                    end = start + pytable_utils.get_rows(wt, '(pclo == %d) & (pchi == %d)' %
                                                         utils.addr_lo(long(r['pc'])),
                                                         utils.addr_hi(long(r['pc']))[0]['writesize'])
                    intervals[num].add(intervaltree.Interval(start, end))
            return intervals 

def find_upstream(self, start, end, offset=0):

        """Method to use the segment tree to find all segments in the transcript upstream of a given interval."""
        return self.segmenttree.upstream_of_interval(Interval(start - offset, end),
                                                     num_intervals=self.exon_num + len(self.introns),
                                                     max_dist=10**9) 

def make_iv(begin, end, label=False):
    if label:
        return Interval(begin, end, "[{0},{1})".format(begin, end))
    else:
        return Interval(begin, end) 

def test_split_overlap_single_member():
    t = IntervalTree([Interval(0, 1)])
    t.split_overlaps()
    t.verify()
    assert len(t) == 1 

def test_merge_equals_reducer_with_initializer():
    def reducer(old, new):
        return old + [new]
    # empty tree
    e = IntervalTree()
    e.merge_equals(data_reducer=reducer, data_initializer=[])
    e.verify()
    assert not e

    # One Interval in tree, no change
    o = IntervalTree.from_tuples([(1, 2, 'hello')])
    o.merge_equals(data_reducer=reducer, data_initializer=[])
    o.verify()
    assert len(o) == 1
    assert sorted(o) == [Interval(1, 2, ['hello'])]

    # many Intervals in tree, no change
    t = IntervalTree.from_tuples(data.ivs1.data)
    orig = IntervalTree.from_tuples(data.ivs1.data)
    t.merge_equals(data_reducer=reducer, data_initializer=[])
    t.verify()
    assert len(t) == len(orig)
    assert t != orig
    assert sorted(t) == [Interval(b, e, [d]) for b, e, d in sorted(orig)]

    # many Intervals in tree, with change
    t = IntervalTree.from_tuples(data.ivs1.data)
    orig = IntervalTree.from_tuples(data.ivs1.data)
    t.addi(4, 7, 'foo')
    t.merge_equals(data_reducer=reducer, data_initializer=[])
    t.verify()
    assert len(t) == len(orig)
    assert t != orig
    assert not t.containsi(4, 7, 'foo')
    assert not t.containsi(4, 7, '[4,7)')
    assert t.containsi(4, 7, ['[4,7)', 'foo']) 

def test_chop_datafunc():
    def datafunc(iv, islower):
        oldlimit = iv[islower]
        return "oldlimit: {0}, islower: {1}".format(oldlimit, islower)

    t = IntervalTree([Interval(0, 10)])
    t.chop(3, 7, datafunc)
    assert len(t) == 2
    assert sorted(t)[0] == Interval(0, 3, 'oldlimit: 10, islower: True')
    assert sorted(t)[1] == Interval(7, 10, 'oldlimit: 0, islower: False')

    t = IntervalTree([Interval(0, 10)])
    t.chop(0, 7, datafunc)
    assert len(t) == 1
    assert sorted(t)[0] == Interval(7, 10, 'oldlimit: 0, islower: False')

    t = IntervalTree([Interval(0, 10)])
    t.chop(5, 10, datafunc)
    assert len(t) == 1
    assert sorted(t)[0] == Interval(0, 5, 'oldlimit: 10, islower: True')

    t = IntervalTree([Interval(0, 10)])
    t.chop(-5, 15, datafunc)
    assert len(t) == 0

    t = IntervalTree([Interval(0, 10)])
    t.chop(0, 10, datafunc)
    assert len(t) == 0