Belle II Software development
test_b2file-mix.py
1
8
9'''
10Test the tool b2file-mix and make sure it works.
11'''
12
13import subprocess
14
15import basf2 as b2
16import b2test_utils as b2tu
17
18
19class PopulateEvents(b2.Module):
20 '''Class for populating the events with some random numbers.'''
21
22 def initialize(self):
23 '''Initialize.'''
24 import ROOT
25
26 self.evt = ROOT.Belle2.PyStoreObj('EventMetaData')
27 self.evt.isRequired()
28
29 def event(self):
30 '''Event.'''
31 import ROOT
32 # Let's store a random number in the slot for the generator weight
33 self.evt.setGeneratedWeight(ROOT.gRandom.Rndm())
34
35
36def path_for_test_file(events, output_name, seed):
37 b2.set_random_seed(seed)
38 main = b2.Path()
39 main.add_module('EventInfoSetter', expList=0, runList=0, evtNumList=events)
40 main.add_module(PopulateEvents())
41 main.add_module('RootOutput', outputFileName=output_name)
42 b2.process(main)
43
44
45if __name__ == '__main__':
46
47 with b2tu.clean_working_directory():
48
49 events = 1000
50 file_name_a = 'test_a.root'
51 file_name_b = 'test_b.root'
52 file_name_ab = 'test_ab.root'
53
54 # Create a test file in the current directory
55 b2tu.run_in_subprocess(target=path_for_test_file, events=events, output_name=file_name_a, seed='aaa')
56 b2tu.run_in_subprocess(target=path_for_test_file, events=events, output_name=file_name_b, seed='bbb')
57 subprocess.check_call(['b2file-merge', file_name_ab, file_name_a, file_name_b])
58
59 # Now mix the file and run a simple check
60 subprocess.check_call(['b2file-mix', file_name_a, '-o', 'test1.root'])
61 metadata = b2.get_file_metadata('test1.root')
62 assert (metadata.getNEvents() == events)
63
64 # Mix again the file, but keeping less events
65 subprocess.check_call(['b2file-mix', file_name_a, '-o', 'test2.root', '-n', '100'])
66 metadata = b2.get_file_metadata('test2.root')
67 assert (metadata.getNEvents() == 100)
68
69 # Mix again the file, but changing the exp. number
70 subprocess.check_call(['b2file-mix', file_name_a, '-o', 'test3.root', '--exp', '114'])
71 metadata = b2.get_file_metadata('test3.root')
72 assert (metadata.getExperimentLow() == 114)
73
74 # Mix two files, this time passing the seed. Few times: twice with the same seed and once with a different one.
75 subprocess.check_call(['b2file-mix', file_name_a, file_name_b, '-o', 'test4.root', '--seed', 'abc'])
76 metadata = b2.get_file_metadata('test4.root')
77 assert (metadata.getRandomSeed() == 'abc')
78 subprocess.check_call(['b2file-mix', file_name_a, file_name_b, '-o', 'test5.root', '--seed', 'abc'])
79 subprocess.check_call(['b2file-mix', file_name_a, file_name_b, '-o', 'test6.root', '--seed', 'def'])
80
81 # Compare the content of test4, test5 and test6 and make sure it's different from test_ab.
82 # Then, test4 and test5 must be identical, while test6 must be different than test4 and test5.
83 # Since b2file-mix keep the relative ordering between the events in a file: For checking test6
84 # is different, let's count how many times the entries differ and ask this number is large enough.
85 # Because of the same reason, let's skip the first and last few events.
86 import ROOT # noqa
87 counts = 0
88 with ROOT.TFile.Open(file_name_ab) as fileab:
89 treeab = fileab.Get('tree')
90 with ROOT.TFile.Open('test4.root') as file4:
91 tree4 = file4.Get('tree')
92 with ROOT.TFile.Open('test5.root') as file5:
93 tree5 = file5.Get('tree')
94 with ROOT.TFile.Open('test6.root') as file6:
95 tree6 = file6.Get('tree')
96 for i, (eventab, event4, event5, event6) in enumerate(zip(treeab, tree4, tree5, tree6)):
97 if i < 10 or i > 1990:
98 continue
99 weightab = eventab.EventMetaData.getGeneratedWeight()
100 weight4 = event4.EventMetaData.getGeneratedWeight()
101 weight5 = event5.EventMetaData.getGeneratedWeight()
102 weight6 = event6.EventMetaData.getGeneratedWeight()
103 assert (weightab != weight4)
104 assert (weightab != weight6)
105 assert (weight4 == weight5)
106 counts += (weight4 != weight6)
107 # It should be above 1900: let's be generous here.
108 assert (counts > 1800)