test_traversal.py 文件源码

def test_prune_from_top():
    d, traversal = get_four_level_data_traversal()
    reduced_storm_id = [1,]
    d = traversal.reduce_to_entities('storm_id', reduced_storm_id)
    reduced_stroke_id = np.asarray([])
    reduced_flash_id = np.asarray([])
    reduced_trig_id = np.asarray([])
    assert_equal(d['storm_id'], reduced_storm_id)
    assert_equal(d['flash_id'], reduced_flash_id)
    assert_equal(d['stroke_id'], reduced_stroke_id)
    assert_equal(d['trig_id'], reduced_trig_id)

    reduced_storm_id = [2,]
    d = traversal.reduce_to_entities('storm_id', reduced_storm_id)
    reduced_flash_id = [4,5,6,7,8]
    reduced_stroke_id = [13,14,15,19,20,23,46]
    reduced_trig_id = [18,19,20,22,23,25,26,30,31,32]    
    assert_equal(d['storm_id'].data, reduced_storm_id)
    assert_equal(d['flash_id'].data, reduced_flash_id)
    assert_equal(d['stroke_id'].data, reduced_stroke_id)
    assert_equal(d['trig_id'].data, reduced_trig_id)