Ejemplo n.º 1
0
def test_boolean_scalars():
    original_file = uproot.open('tests/boolean_scalars_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['bool_t_branch'].array().dtype == np.dtype('bool')
        assert new_tree['bool_t_branch'].array().tolist() == [False, True, False]
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 2
0
def test_event_selection_vectors():
    original_file = uproot.open('tests/vectors_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename, selection=[True, False, True])
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['int_vector_branch'].array().tolist() == [[], [13]]
        assert abs(new_tree['float_vector_branch'].array() - awkward.fromiter([[], [15.15]])).max().max() < 1e-5
        assert abs(new_tree['double_vector_branch'].array() - awkward.fromiter([[], [16.16]])).max().max() < 1e-5
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 3
0
def test_select_branches():
    original_file = uproot.open('tests/scalars_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        branch_list = ['float_branch', 'bool_branch']
        clone_tree(original_tree, new_filename, branches=branch_list)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert set(map(lambda b: b.decode('utf-8'), new_tree.keys())) == set(branch_list)
        assert np.allclose(new_tree['float_branch'].array(), [0.0, 3.3])
        assert new_tree['bool_branch'].array().tolist() == [False, True]
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 4
0
def test_event_selection_scalars():
    original_file = uproot.open('tests/scalars_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename, selection=[False, True])
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['int_branch'].array().tolist() == [-1]
        assert new_tree['long_branch'].array().tolist() == [-2]
        assert np.allclose(new_tree['float_branch'].array(), [3.3])
        assert np.allclose(new_tree['double_branch'].array(), [4.4])
        assert new_tree['bool_branch'].array().tolist() == [True]
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 5
0
def test_floating_point_scalars():
    original_file = uproot.open('tests/floating_point_scalars_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['float_t_branch'].array().dtype == np.dtype('float32')
        assert np.allclose(new_tree['float_t_branch'].array(), [0.0, 7.7, -19.19])
        # assert np.allclose(new_tree['float16_t_branch'].array(), [0.0, 8.8, -20.20])
        assert new_tree['double_t_branch'].array().dtype == np.dtype('float64')
        assert np.allclose(new_tree['double_t_branch'].array(), [0.0, 9.9, -21.21])
        # assert np.allclose(new_tree['double32_t_branch'].array(), [0.0, 10.10, -22.22])
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 6
0
def test_signed_integer_vectors():
    original_file = uproot.open('tests/signed_integer_vectors_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        # assert new_tree['char_vector_branch'].array().tolist() == [[], [-1, 2, 3], [-37]]
        assert new_tree['short_vector_branch'].array()[0].dtype == np.dtype('int16')
        assert new_tree['short_vector_branch'].array().tolist() == [[], [-7, 8, 9], [-39]]
        assert new_tree['int_vector_branch'].array()[0].dtype == np.dtype('int32')
        assert new_tree['int_vector_branch'].array().tolist() == [[], [-13, 14, 15], [-41]]
        # assert new_tree['long_vector_branch'].array().tolist() == [[], [-19, 20, 21], [-43]]
        # assert new_tree['long64_t_vector_branch'].array().tolist() == [[], [-25, 26, 27], [-45]]
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 7
0
def test_floating_point_vectors():
    original_file = uproot.open('tests/floating_point_vectors_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['float_vector_branch'].array()[0].dtype == np.dtype('float32')
        assert abs(new_tree['float_vector_branch'].array() - awkward.fromiter([[],
                                                                               [-31.31, 32.32, 33.33],
                                                                               [-47.47]])).max().max() < 1e-5
        assert new_tree['double_vector_branch'].array()[0].dtype == np.dtype('float64')
        assert abs(new_tree['double_vector_branch'].array() - awkward.fromiter([[],
                                                                                [-34.34, 35.35, 36.36],
                                                                                [-48.48]])).max().max() < 1e-5
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 8
0
def test_signed_integer_scalars():
    original_file = uproot.open('tests/signed_integer_scalars_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        clone_tree(original_tree, new_filename)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['char_t_branch'].array().dtype == np.dtype('int8')
        assert new_tree['char_t_branch'].array().tolist() == [0, 1, -13]
        assert new_tree['short_t_branch'].array().dtype == np.dtype('int16')
        assert new_tree['short_t_branch'].array().tolist() == [0, 3, -15]
        assert new_tree['int_t_branch'].array().dtype == np.dtype('int32')
        assert new_tree['int_t_branch'].array().tolist() == [0, 5, -17]
        assert new_tree['long64_t_branch'].array().dtype == np.dtype('int64')
        assert new_tree['long64_t_branch'].array().tolist() == [0, 11, -23]
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 9
0
def test_new_vector_branch():
    original_file = uproot.open('tests/vectors_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        new_branch_dictionary = {'new_int16_vector_branch': awkward.fromiter([[-1], [-2, 3], []]).astype(np.dtype('int16')),
                                 'new_int32_vector_branch': awkward.fromiter([[-4], [-5, 6], []]).astype(np.dtype('int32')),
                                 'new_float32_vector_branch': awkward.fromiter([[-7.7],
                                                                                [-8.8, 9.9],
                                                                                []]).astype(np.dtype('float32')),
                                 'new_float64_vector_branch': awkward.fromiter([[-10.10],
                                                                                [-11.11, 12.12],
                                                                                []]).astype(np.dtype('float64'))}
        clone_tree(original_tree, new_filename, new_branches=new_branch_dictionary)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['int_vector_branch'].array().tolist() == [[], [-1, 2, 3], [13]]
        assert abs(new_tree['float_vector_branch'].array() - awkward.fromiter([[],
                                                                               [-7.7, 8.8, 9.9],
                                                                               [15.15]])).max().max() < 1e-5
        assert abs(new_tree['double_vector_branch'].array() - awkward.fromiter([[],
                                                                                [-10.10, 11.11, 12.12],
                                                                                [16.16]])).max().max() < 1e-5
        assert new_tree['new_int16_vector_branch'].array()[0].dtype == np.dtype('int16')
        assert new_tree['new_int16_vector_branch'].array().tolist() == [[-1], [-2, 3], []]
        assert new_tree['new_int32_vector_branch'].array()[0].dtype == np.dtype('int32')
        assert new_tree['new_int32_vector_branch'].array().tolist() == [[-4], [-5, 6], []]
        assert new_tree['new_float32_vector_branch'].array()[0].dtype == np.dtype('float32')
        assert abs(new_tree['new_float32_vector_branch'].array() - awkward.fromiter([[-7.7],
                                                                                     [-8.8, 9.9],
                                                                                     []])).max().max() < 1e-5
        assert new_tree['new_float64_vector_branch'].array()[0].dtype == np.dtype('float64')
        assert abs(new_tree['new_float64_vector_branch'].array() - awkward.fromiter([[-10.10],
                                                                                     [-11.11, 12.12],
                                                                                     []])).max().max() < 1e-5
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)
Ejemplo n.º 10
0
def test_new_scalar_branch():
    original_file = uproot.open('tests/scalars_tree_file.root')
    treename = 'tree'
    original_tree = original_file[treename]
    new_filename = tempfile.mkstemp(suffix='.root', dir=os.getcwd())[1]
    try:
        new_branch_dictionary = {'new_int8_branch': np.array([-5, 6], dtype='int8'),
                                 'new_int16_branch': np.array([-7, 8], dtype='int16'),
                                 'new_int32_branch': np.array([-9, 10], dtype='int32'),
                                 'new_int64_branch': np.array([-11, 12], dtype='int64'),
                                 'new_float32_branch': np.array([-13.13, 14.14], dtype='float32'),
                                 'new_float64_branch': np.array([-15.15, 16.16], dtype='float64'),
                                 'new_bool_branch': np.array([True, False])}
        clone_tree(original_tree, new_filename, new_branches=new_branch_dictionary)
        new_file = uproot.open(new_filename)
        new_tree = new_file[treename]
        assert new_tree['int_branch'].array().tolist() == [0, -1]
        assert new_tree['long_branch'].array().tolist() == [0, -2]
        assert np.allclose(new_tree['float_branch'].array(), [0.0, 3.3])
        assert np.allclose(new_tree['double_branch'].array(), [0.0, 4.4])
        assert new_tree['bool_branch'].array().tolist() == [False, True]
        assert new_tree['new_int8_branch'].array()[0].dtype == np.dtype('int8')
        assert new_tree['new_int8_branch'].array().tolist() == [-5, 6]
        assert new_tree['new_int16_branch'].array()[0].dtype == np.dtype('int16')
        assert new_tree['new_int16_branch'].array().tolist() == [-7, 8]
        assert new_tree['new_int32_branch'].array()[0].dtype == np.dtype('int32')
        assert new_tree['new_int32_branch'].array().tolist() == [-9, 10]
        assert new_tree['new_int64_branch'].array()[0].dtype == np.dtype('int64')
        assert new_tree['new_int64_branch'].array().tolist() == [-11, 12]
        assert new_tree['new_float32_branch'].array()[0].dtype == np.dtype('float32')
        assert np.allclose(new_tree['new_float32_branch'].array(), [-13.13, 14.14])
        assert new_tree['new_float64_branch'].array()[0].dtype == np.dtype('float64')
        assert np.allclose(new_tree['new_float64_branch'].array(), [-15.15, 16.16])
        assert new_tree['new_bool_branch'].array()[0].dtype == np.dtype('bool')
        assert new_tree['new_bool_branch'].array().tolist() == [True, False]
    finally:
        if os.path.isfile(new_filename):
            os.remove(new_filename)