Log In Sign Up
Models:
Joel Williams
JoelW/
NeurIPS2018-AIforProsthet
0
Downloads: 1
[f9c9f2]: / baselines / common / tests / test_segment_tree.py

Download this file

104 lines (76 with data), 2.8 kB 

  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
import numpy as np



from baselines.common.segment_tree import SumSegmentTree, MinSegmentTree





def test_tree_set():

    tree = SumSegmentTree(4)



    tree[2] = 1.0

    tree[3] = 3.0



    assert np.isclose(tree.sum(), 4.0)

    assert np.isclose(tree.sum(0, 2), 0.0)

    assert np.isclose(tree.sum(0, 3), 1.0)

    assert np.isclose(tree.sum(2, 3), 1.0)

    assert np.isclose(tree.sum(2, -1), 1.0)

    assert np.isclose(tree.sum(2, 4), 4.0)





def test_tree_set_overlap():

    tree = SumSegmentTree(4)



    tree[2] = 1.0

    tree[2] = 3.0



    assert np.isclose(tree.sum(), 3.0)

    assert np.isclose(tree.sum(2, 3), 3.0)

    assert np.isclose(tree.sum(2, -1), 3.0)

    assert np.isclose(tree.sum(2, 4), 3.0)

    assert np.isclose(tree.sum(1, 2), 0.0)





def test_prefixsum_idx():

    tree = SumSegmentTree(4)



    tree[2] = 1.0

    tree[3] = 3.0



    assert tree.find_prefixsum_idx(0.0) == 2

    assert tree.find_prefixsum_idx(0.5) == 2

    assert tree.find_prefixsum_idx(0.99) == 2

    assert tree.find_prefixsum_idx(1.01) == 3

    assert tree.find_prefixsum_idx(3.00) == 3

    assert tree.find_prefixsum_idx(4.00) == 3





def test_prefixsum_idx2():

    tree = SumSegmentTree(4)



    tree[0] = 0.5

    tree[1] = 1.0

    tree[2] = 1.0

    tree[3] = 3.0



    assert tree.find_prefixsum_idx(0.00) == 0

    assert tree.find_prefixsum_idx(0.55) == 1

    assert tree.find_prefixsum_idx(0.99) == 1

    assert tree.find_prefixsum_idx(1.51) == 2

    assert tree.find_prefixsum_idx(3.00) == 3

    assert tree.find_prefixsum_idx(5.50) == 3





def test_max_interval_tree():

    tree = MinSegmentTree(4)



    tree[0] = 1.0

    tree[2] = 0.5

    tree[3] = 3.0



    assert np.isclose(tree.min(), 0.5)

    assert np.isclose(tree.min(0, 2), 1.0)

    assert np.isclose(tree.min(0, 3), 0.5)

    assert np.isclose(tree.min(0, -1), 0.5)

    assert np.isclose(tree.min(2, 4), 0.5)

    assert np.isclose(tree.min(3, 4), 3.0)



    tree[2] = 0.7



    assert np.isclose(tree.min(), 0.7)

    assert np.isclose(tree.min(0, 2), 1.0)

    assert np.isclose(tree.min(0, 3), 0.7)

    assert np.isclose(tree.min(0, -1), 0.7)

    assert np.isclose(tree.min(2, 4), 0.7)

    assert np.isclose(tree.min(3, 4), 3.0)



    tree[2] = 4.0



    assert np.isclose(tree.min(), 1.0)

    assert np.isclose(tree.min(0, 2), 1.0)

    assert np.isclose(tree.min(0, 3), 1.0)

    assert np.isclose(tree.min(0, -1), 1.0)

    assert np.isclose(tree.min(2, 4), 3.0)

    assert np.isclose(tree.min(2, 3), 4.0)

    assert np.isclose(tree.min(2, -1), 4.0)

    assert np.isclose(tree.min(3, 4), 3.0)





if __name__ == '__main__':

    test_tree_set()

    test_tree_set_overlap()

    test_prefixsum_idx()

    test_prefixsum_idx2()

    test_max_interval_tree()