Belle II Software light-2406-ragdoll
register_pystorearray.py
1#!/usr/bin/env python
2
3
10
11
15
16import math
17
18import basf2
19import ROOT
20from ROOT import Belle2
21
22
23class SillyGeneratorModule(basf2.Module):
24 """Small module to demonstrate how the registration of StoreArrays works from Python"""
25
26 def initialize(self):
27 """ Register a StoreArray on the DataStore"""
28 mcParticles = Belle2.PyStoreArray(Belle2.MCParticle.Class())
29 mcParticles.registerInDataStore()
30
31 def event(self):
32 """ Access and fill the registered StoreArray """
33 # Alternative access by name
34 mcParticles = Belle2.PyStoreArray("MCParticles")
35 mcParticle = mcParticles.appendNew()
36
37 # Fill one test muon for demonstration
38 vertex = ROOT.TVector3(0, 0, 0)
39
40 # Sample a random momentum vector of
41 # gaus distributed length and uniform distibuted direction
42 phi = ROOT.gRandom.Uniform(0.0, 2.0 * math.pi)
43 costheta = ROOT.gRandom.Uniform(-1.0, 1.0)
44 theta = math.acos(costheta)
45 momentum = ROOT.Math.XYZVector(math.sin(theta) * math.cos(phi), math.sin(theta) * math.sin(phi), costheta)
46
47 mcParticle.setPDG(13)
48 mcParticle.setMassFromPDG()
49 mcParticle.addStatus(Belle2.MCParticle.c_PrimaryParticle)
50 mcParticle.addStatus(Belle2.MCParticle.c_StableInGenerator)
51 mcParticle.setProductionVertex(vertex)
52 mcParticle.setMomentum(momentum)
53 m = mcParticle.getMass()
54 mcParticle.setEnergy(math.sqrt(momentum * momentum + m * m))
55 mcParticle.setDecayTime(float("inf"))
56
57
58class ParticleStatisticsModule(basf2.Module):
59 """Collect statistics on particles - also parallel processable"""
60
61 def __init__(self):
62 """set module flags"""
63 super().__init__()
64 self.set_property_flags(
65 basf2.ModulePropFlags.PARALLELPROCESSINGCERTIFIED |
66 basf2.ModulePropFlags.TERMINATEINALLPROCESSES)
67
68 def initialize(self):
69 """init"""
70
71 self.tfile = ROOT.TFile("ParticleStatistics.root", "recreate")
72
73 ntuple = Belle2.PyStoreObj(Belle2.RootMergeable("TNtuple").Class(), "ParticleMomenta", Belle2.DataStore.c_Persistent)
74 ntuple.registerInDataStore()
75 print("IsValid", ntuple.isValid())
76 ntuple.create()
77 print("IsValid", ntuple.isValid())
78 ntuple.obj().assign(ROOT.TNtuple("Particles", # ROOT name
79 "Momentum compontents of the particles", # ROOT title
80 "px:py:pz" # Var list
81 ))
82
83 hist = Belle2.PyStoreObj(Belle2.RootMergeable("TH1D").Class(), "AbsMomentum", Belle2.DataStore.c_Persistent)
84 hist.registerInDataStore()
85 print("IsValid", hist.isValid())
86 hist.create()
87 print("IsValid", hist.isValid())
88 hist.obj().assign(ROOT.TH1D("AbsMomentum", # ROOT name
89 "Absolute momentum of particles", # ROOT title
90 20, # n bins
91 0, # upper bound
92 2 # lower bound
93 ))
94
95 print("IsValid", ntuple.isValid())
96 print("IsValid", hist.isValid())
97
98 def event(self):
99 """actually collect info"""
100 ntuple = Belle2.PyStoreObj("ParticleMomenta", Belle2.DataStore.c_Persistent)
101 hist = Belle2.PyStoreObj("AbsMomentum", Belle2.DataStore.c_Persistent)
102
103 mcParticles = Belle2.PyStoreArray(Belle2.MCParticle.Class())
104
105 for mcParticle in mcParticles:
106 momentum = mcParticle.getMomentum()
107 ntuple.get().Fill(momentum.X(),
108 momentum.Y(),
109 momentum.Z())
110
111 hist.get().Fill(momentum.Mag())
112
113 def terminate(self):
114 """terminate"""
117 print("Writting objects")
118 ntuple = Belle2.PyStoreObj("ParticleMomenta", Belle2.DataStore.c_Persistent)
119 ntuple.write(self.tfile)
120
121 hist = Belle2.PyStoreObj("AbsMomentum", Belle2.DataStore.c_Persistent)
122 hist.write(self.tfile)
123
124 self.tfile.Close()
125
126
127def main():
128 path = basf2.create_path()
129 path.add_module('EventInfoSetter',
130 evtNumList=[1000],
131 runList=[1],
132 expList=[1]
133 )
134
135 path.add_module(SillyGeneratorModule())
136 path.add_module(ParticleStatisticsModule())
137
138 basf2.process(path)
139
140 # Print call statistics
141 print(basf2.statistics)
142
143
144if __name__ == "__main__":
145 main()
static bool isOutputProcess()
Return true if the process is an output process.
Definition: ProcHandler.cc:232
static bool parallelProcessingUsed()
Returns true if multiple processes have been spawned, false in single-core mode.
Definition: ProcHandler.cc:226
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:67
Wrap a root histogram or TNtuple to make it mergeable.
Definition: RootMergeable.h:61
Definition: main.py:1