Belle II Software development
lhereader_nominal.py
1#!/usr/bin/env python3
2
3
10
11# Test LHEReader and LHEInputModule
12
13import os
14import sys
15import math
16from ROOT import TFile
17from pdg import add_particle
18from basf2 import create_path, register_module, process, print_params, find_file
19from modularAnalysis import fillParticleListsFromMC, printVariableValues, variablesToNtuple
20from variables import variables as vm
21from tempfile import TemporaryDirectory
22
23# check that the file exists, if not: skip the test
24inputfile = find_file('generators/tests/event_nominal.lhe')
25if len(inputfile) == 0:
26 sys.stderr.write(
27 'TEST SKIPPED: input file ' +
28 inputfile +
29 ' not found.')
30 sys.exit(-1)
31
32
33# add the A' particle to the pdg database so we can specify it later
34add_particle('A', 9000008, 5.5, 0.1329, 0, 0)
35
36# configure the LHE reade
37lhereader = register_module('LHEInput')
38lhereader.param('createEventMetaData', True)
39lhereader.param('runNum', 1337)
40lhereader.param('expNum', 0)
41lhereader.param('inputFileList', [inputfile])
42lhereader.param('useWeights', False)
43lhereader.param('nInitialParticles', 2)
44lhereader.param('nVirtualParticles', 0)
45lhereader.param('wrongSignPz', False) # Using BII convention where e- goes in +z
46print_params(lhereader)
47
48# prepare the path
49testpath = create_path()
50testpath.add_module('Progress')
51testpath.add_module(lhereader)
52testpath.add_module('BoostMCParticles')
53fillParticleListsFromMC([('gamma:gen', ''), ('A:gen', '')], path=testpath)
54
55# dump information from basf2
56vm.addAlias('pxcms', 'useCMSFrame(px)')
57vm.addAlias('pycms', 'useCMSFrame(py)')
58vm.addAlias('pzcms', 'useCMSFrame(pz)')
59vm.addAlias('pecms', 'useCMSFrame(E)')
60variables = ['M', 'px', 'py', 'pz', 'E',
61 'pxcms', 'pycms', 'pzcms', 'pecms']
62
63printVariableValues('A:gen', variables, path=testpath)
64printVariableValues('gamma:gen', variables, path=testpath)
65
66variablesToNtuple('A:gen', variables, 'darkphoton', 'test.root', path=testpath)
67variablesToNtuple('gamma:gen', variables, 'gammas', 'test.root', path=testpath)
68
69# temporary directory to keep cwd clean
70with TemporaryDirectory() as tmp:
71
72 # process the basf2 path
73 os.chdir(tmp)
74 process(testpath)
75
76 # open output and check the momenta are what is expected
77 fi = TFile('test.root')
78 t1 = fi.Get('darkphoton')
79 t2 = fi.Get('gammas')
80
81 assert t1.GetEntries() == 1, f'Output contains {int(t1.GetEntries())} entries'
82 assert t2.GetEntries() == 1, f'Output contains {int(t1.GetEntries())} entries'
83
84 t1.GetEntry(0)
85 t2.GetEntry(0)
86
87 assert t1.__run__ == 1337, 'Run number not set correctly'
88 assert t1.__experiment__ == 0, 'Experiment number not set correctly'
89 assert t2.M == 0, 'Photon is not as expected'
90
91 # precision is low since MCParticle class is using floats to store particle momenta
92 # when changed to double precision was much better
93 eps = 2e-5
94
95 # test momenta balance in CMS
96 assert math.isclose(t1.pxcms, -t2.pxcms, rel_tol=eps), 'Momenta don\'t balance'
97 assert math.isclose(t1.pycms, -t2.pycms, rel_tol=eps), 'Momenta don\'t balance'
98 assert math.isclose(t1.pzcms, -t2.pzcms, rel_tol=eps), 'Momenta don\'t balance'
99
100 # test that dark photon CM momentum is close to LHE input generators/tests/event.lhe
101 assert math.isclose(t1.pxcms, +1.4572035746e-03, rel_tol=eps), 'CMS momenta are not as expected'
102 assert math.isclose(t1.pycms, +2.5198484375e-03, rel_tol=eps), 'CMS momenta are not as expected'
103 assert math.isclose(t1.pzcms, -3.8826254795e+00, rel_tol=eps), 'CMS momenta are not as expected'
104 assert math.isclose(t1.pecms, 6.6973734293e+00, rel_tol=eps), 'CMS momenta are not as expected'
105
106 # test that dark photon mass is as in LHE
107 assert math.isclose(t1.M, 5.4571074540e+00, rel_tol=eps), 'Mass is not as expected'
108
109 # test dark photon momentum in LAB
110 assert math.isclose(t1.px, 0.24481175912036506892, rel_tol=eps), 'Boosted momenta are not as expected'
111 assert math.isclose(t1.py, 0.002519848437500000083, rel_tol=eps), 'Boosted momenta are not as expected'
112 assert math.isclose(t1.pz, -2.1368737390471825854, rel_tol=eps), 'Boosted momenta are not as expected'
113
114 # test photon momentum in LAB
115 assert math.isclose(t2.px, 0.21182522798639694117, rel_tol=eps), 'Boosted momenta are not as expected'
116 assert math.isclose(t2.py, -0.002519848437500000083, rel_tol=eps), 'Boosted momenta are not as expected'
117 assert math.isclose(t2.pz, 5.1364143846516343572, rel_tol=eps), 'Boosted momenta are not as expected'
118
119 fi.Close()