14
14
# See the License for the specific language governing permissions and
15
15
# limitations under the License.
16
16
17
+ from collections import namedtuple
17
18
import unittest
18
19
19
20
import numpy
@@ -26,8 +27,12 @@ def runTest(self):
26
27
27
28
def test_branch_array (self ):
28
29
file = uproot .open ("tests/simple.root" )
30
+ repr (file )
29
31
30
32
tree = file ["tree" ]
33
+ repr (tree )
34
+ repr (tree .branch ("one" ))
35
+ repr (tree .branch ("one" ))
31
36
self .assertEqual (tree .branch ("one" ).array ().tolist (), [1 , 2 , 3 , 4 ])
32
37
self .assertEqual (tree .branch ("two" ).array ().tolist (), numpy .array ([1.1 , 2.2 , 3.3 , 4.4 ], dtype = numpy .float32 ).tolist ())
33
38
self .assertEqual (tree .branch ("three" ).array ().tolist (), [b"uno" , b"dos" , b"tres" , b"quatro" ])
@@ -243,3 +248,71 @@ def test_directories(self):
243
248
self .assertEqual (dict ((name , array .tolist ()) for name , array in file ["one/tree;1" ].arrays (["one" , "two" , "three" ]).items ()), {b"one" : [1 , 2 , 3 , 4 ], b"two" : [1.100000023841858 , 2.200000047683716 , 3.299999952316284 , 4.400000095367432 ], b"three" : [b"uno" , b"dos" , b"tres" , b"quatro" ]})
244
249
self .assertEqual (file ["one/two/tree;1" ].array ("Int32" ).shape , (100 ,))
245
250
self .assertEqual (file ["three/tree;1" ].array ("I32" ).shape , (100 ,))
251
+
252
+ def test_cast (self ):
253
+ tree = uproot .open ("tests/Zmumu.root" )["events" ]
254
+ one = numpy .cast [numpy .int32 ](numpy .floor (tree .array ("M" )))
255
+ two = tree .array ("M" , numpy .int32 )
256
+ self .assertEqual (one .dtype , two .dtype )
257
+ self .assertEqual (one .shape , two .shape )
258
+ self .assertTrue (numpy .array_equal (one , two ))
259
+
260
+ for (one ,) in tree .iterator (10000 , "M" , outputtype = tuple ):
261
+ one = numpy .cast [numpy .int32 ](numpy .floor (one ))
262
+ for (two ,) in tree .iterator (10000 , {"M" : numpy .int32 }, outputtype = tuple ):
263
+ pass
264
+ self .assertEqual (one .dtype , two .dtype )
265
+ self .assertEqual (one .shape , two .shape )
266
+ self .assertTrue (numpy .array_equal (one , two ))
267
+
268
+ def test_pass_array (self ):
269
+ tree = uproot .open ("tests/Zmumu.root" )["events" ]
270
+ one = numpy .cast [numpy .int32 ](numpy .floor (tree .array ("M" )))
271
+ two = numpy .zeros (one .shape , dtype = one .dtype )
272
+ tree .array ("M" , two )
273
+ self .assertTrue (numpy .array_equal (one , two ))
274
+
275
+ for (one ,) in tree .iterator (10000 , "M" , outputtype = tuple ):
276
+ one = numpy .cast [numpy .int32 ](numpy .floor (one ))
277
+ two = numpy .zeros (one .shape , dtype = one .dtype )
278
+ for (two ,) in tree .iterator (10000 , {"M" : numpy .int32 }, outputtype = tuple ):
279
+ self .assertTrue (numpy .array_equal (one , two ))
280
+
281
+ def test_outputtype (self ):
282
+ tree = uproot .open ("tests/simple.root" )["tree" ]
283
+
284
+ arrays = tree .arrays (["three" , "two" , "one" ], outputtype = dict )
285
+ self .assertTrue (isinstance (arrays , dict ))
286
+ self .assertEqual (arrays [b"one" ].tolist (), [1 , 2 , 3 , 4 ])
287
+ self .assertEqual (arrays [b"three" ].tolist (), [b"uno" , b"dos" , b"tres" , b"quatro" ])
288
+
289
+ arrays = tree .arrays (["three" , "two" , "one" ], outputtype = tuple )
290
+ self .assertTrue (isinstance (arrays , tuple ))
291
+ self .assertEqual (arrays [2 ].tolist (), [1 , 2 , 3 , 4 ])
292
+ self .assertEqual (arrays [0 ].tolist (), [b"uno" , b"dos" , b"tres" , b"quatro" ])
293
+
294
+ arrays = tree .arrays (["three" , "two" , "one" ], outputtype = namedtuple )
295
+ self .assertEqual (arrays .one .tolist (), [1 , 2 , 3 , 4 ])
296
+ self .assertEqual (arrays .three .tolist (), [b"uno" , b"dos" , b"tres" , b"quatro" ])
297
+
298
+ arrays = tree .arrays (["three" , "two" , "one" ], outputtype = list )
299
+ self .assertTrue (isinstance (arrays , list ))
300
+ self .assertEqual (arrays [2 ].tolist (), [1 , 2 , 3 , 4 ])
301
+ self .assertEqual (arrays [0 ].tolist (), [b"uno" , b"dos" , b"tres" , b"quatro" ])
302
+
303
+ class Awesome (object ):
304
+ def __init__ (self , one , two , three ):
305
+ self .one = one
306
+ self .two = two
307
+ self .three = three
308
+
309
+ arrays = tree .arrays (["one" , "two" , "three" ], outputtype = Awesome )
310
+ self .assertTrue (isinstance (arrays , Awesome ))
311
+ self .assertEqual (arrays .one .tolist (), [1 , 2 , 3 , 4 ])
312
+ self .assertEqual (arrays .three .tolist (), [b"uno" , b"dos" , b"tres" , b"quatro" ])
313
+
314
+ def test_informational (self ):
315
+ self .assertEqual (uproot .open ("tests/simple.root" )["tree" ].branchnames , [b"one" , b"two" , b"three" ])
316
+ self .assertEqual (uproot .open ("tests/simple.root" )["tree" ].branchtypes , {b"two" : numpy .dtype (">f4" ), b"one" : numpy .dtype (">i4" ), b"three" : numpy .dtype ("O" )})
317
+ self .assertEqual (uproot .open ("tests/small-evnt-tree-fullsplit.root" )["tree" ].allbranchnames , ["evt" , b"Beg" , b"I16" , b"I32" , b"I64" , b"U16" , b"U32" , b"U64" , b"F32" , b"F64" , b"Str" , b"P3.Px" , b"P3.Py" , b"P3.Pz" , b"ArrayI16[10]" , b"ArrayI32[10]" , b"ArrayI64[10]" , b"ArrayU16[10]" , b"ArrayU32[10]" , b"ArrayU64[10]" , b"ArrayF32[10]" , b"ArrayF64[10]" , b"N" , b"SliceI16" , b"SliceI32" , b"SliceI64" , b"SliceU16" , b"SliceU32" , b"SliceU64" , b"SliceF32" , b"SliceF64" , b"End" ])
318
+ self .assertEqual (uproot .open ("tests/small-evnt-tree-fullsplit.root" )["tree" ].allbranchtypes , {b"Str" : numpy .dtype ("O" ), b"P3.Px" : numpy .dtype (">i4" ), b"I64" : numpy .dtype (">i8" ), b"U64" : numpy .dtype (">u8" ), b"ArrayF32[10]" : numpy .dtype (">f4" ), b"SliceI16" : numpy .dtype (">i2" ), b"ArrayI64[10]" : numpy .dtype (">i8" ), b"evt" : numpy .dtype (">i4" ), b"SliceF64" : numpy .dtype (">f8" ), b"End" : numpy .dtype ("O" ), b"U32" : numpy .dtype (">u4" ), b"Beg" : numpy .dtype ("O" ), b"I32" : numpy .dtype (">i4" ), b"N" : numpy .dtype (">i4" ), b"SliceI32" : numpy .dtype (">i4" ), b"P3.Py" : numpy .dtype (">f8" ), b"U16" : numpy .dtype (">u2" ), b"SliceU32" : numpy .dtype (">u4" ), b"P3.Pz" : numpy .dtype (">i4" ), b"ArrayI32[10]" : numpy .dtype (">i4" ), b"ArrayF64[10]" : numpy .dtype (">f8" ), b"I16" : numpy .dtype (">i2" ), b"SliceU64" : numpy .dtype (">u8" ), b"F64" : numpy .dtype (">f8" ), b"ArrayI16[10]" : numpy .dtype (">i2" ), b"ArrayU16[10]" : numpy .dtype (">u2" ), b"ArrayU32[10]" : numpy .dtype (">u4" ), b"F32" : numpy .dtype (">f4" ), b"SliceF32" : numpy .dtype (">f4" ), b"ArrayU64[10]" : numpy .dtype (">u8" ), b"SliceU16" : numpy .dtype (">u2" ), b"SliceI64" : numpy .dtype (">i8" )})
0 commit comments