8 from ROOT
import Belle2, TFile
9 from pdg
import add_particle
10 from basf2
import create_path, register_module, process, print_params
11 from modularAnalysis
import fillParticleListsFromMC
12 from modularAnalysis
import variablesToNtuple
as v2nt
13 from variables
import variables
as vm
14 from beamparameters
import add_beamparameters
15 from tempfile
import TemporaryDirectory
19 if len(inputfile) == 0:
21 'TEST SKIPPED: input file ' +
28 add_particle(
'A', 9000008, 5.5, 0.1329, 0, 0)
31 lhereader = register_module(
'LHEInput')
32 lhereader.param(
'makeMaster',
True)
33 lhereader.param(
'runNum', 1337)
34 lhereader.param(
'expNum', 0)
35 lhereader.param(
'inputFileList', [inputfile])
36 lhereader.param(
'useWeights',
False)
37 lhereader.param(
'nInitialParticles', 2)
38 lhereader.param(
'nVirtualParticles', 0)
39 lhereader.param(
'boost2Lab',
True)
40 lhereader.param(
'wrongSignPz',
True)
41 print_params(lhereader)
44 testpath = create_path()
45 testpath.add_module(
'Progress')
46 add_beamparameters(testpath,
'Y4S')
47 testpath.add_module(lhereader)
48 fillParticleListsFromMC([(
'gamma:gen',
''), (
'A:gen',
'')], path=testpath)
51 vm.addAlias(
'pxcms',
'useCMSFrame(px)')
52 vm.addAlias(
'pycms',
'useCMSFrame(py)')
53 vm.addAlias(
'pzcms',
'useCMSFrame(pz)')
54 vm.addAlias(
'pecms',
'useCMSFrame(E)')
55 variables = [
'M',
'px',
'py',
'pz',
'E',
56 'pxcms',
'pycms',
'pzcms',
'pecms']
57 v2nt(
'A:gen', variables,
'darkphoton',
'test.root', path=testpath)
58 v2nt(
'gamma:gen', variables,
'gammas',
'test.root', path=testpath)
61 with TemporaryDirectory()
as tmp:
68 fi = TFile(
'test.root')
69 t1 = fi.Get(
'darkphoton')
72 assert t1.GetEntries() == 1,
'Output contains %i entries' % t1.GetEntries()
73 assert t2.GetEntries() == 1,
'Output contains %i entries' % t1.GetEntries()
78 assert t1.__run__ == 1337,
'Run number not set correctly'
79 assert t1.__experiment__ == 0,
'Experiment number not set correctly'
80 assert t2.M == 0,
'Photon is not as expected'
85 assert math.isclose(t1.M, 5.4571074540e+00, rel_tol=eps),
'Mass is not as expected'
87 assert math.isclose(t1.pxcms, -t2.pxcms, rel_tol=eps),
'Momenta don\'t balance'
88 assert math.isclose(t1.pycms, -t2.pycms, rel_tol=eps),
'Momenta don\'t balance'
89 assert math.isclose(t1.pzcms, -t2.pzcms, rel_tol=eps),
'Momenta don\'t balance'
91 assert math.isclose(t1.pxcms, -0.023859674786793544, rel_tol=eps),
'CMS momenta are not as expected'
92 assert math.isclose(t1.pycms, 0.0025198485236614943, rel_tol=eps),
'CMS momenta are not as expected'
93 assert math.isclose(t1.pzcms, -3.882552444880825, rel_tol=eps),
'CMS momenta are not as expected'
94 assert math.isclose(t1.pecms, 6.697373505125157, rel_tol=eps),
'CMS momenta are not as expected'
96 assert t1.px == 0.24189484119415283,
'Boosted momenta are not as expected'
97 assert t1.py == 0.0025198485236614943,
'Boosted momenta are not as expected'
98 assert t1.pz == -2.136873722076416,
'Boosted momenta are not as expected'
100 assert t2.px == 0.21474215388298035,
'Boosted momenta are not as expected'
101 assert t2.py == -0.0025198485236614943,
'Boosted momenta are not as expected'
102 assert t2.pz == 5.136414527893066,
'Boosted momenta are not as expected'